A Fix for Softmax and Cross-Entropy

The decomposition of Slingshot as a direct consequence of Softmax Collapse is convincing. The authors show precisely that once $m > c_d$, standard cross-entropy rounds the target gradient $\tilde{g}_t$ to zero while non-target gradients persist, violating the zero-sum constraint and causing Adam's second moment to decay until a restored gradient triggers a catastrophic step: "Because $\sqrt{\hat{v}_i}$ decayed while those examples were above $c_d$, Adam divides the restored gradients by a very small denominator $\sqrt{\hat{v}_i}+\epsilon$, producing a large parameter step that spikes the training loss" (Barron et al., Section 4). The log-space derivation for soft-target gradients using expm1 and log1p (Section 3, Equation 9) is mathematically sound and avoids subtracting nearly-equal probabilities.

The primary concern is narrow experimental scope: every demonstration uses small synthetic grokking benchmarks or tiny distillation setups, with no ImageNet, language-model, or large-scale validation where softmax dominates. The margin equilibrium analysis (Section 5) relies on a stationary-point approximation that does not strictly hold during transient training: "At a stationary point where the parameters stop changing, the moving averages converge to $\hat{m}_i=g_i$ and $\hat{v}_i=g_i^2$" (Barron et al., Section 5). In practice, Adam's biased moving averages deviate from this ideal during Slingshot cycles. Moreover, the soft-target fix requires a custom backward pass with jax.lax.stop_gradient tricks, which adds integration complexity not present in standard framework kernels. The paper also omits any runtime benchmark of the argmax-and-exclude logsumexp versus standard two-pass softmax.

Within the studied grokking regime the evidence strongly supports the core claims: fixed cross-entropy eliminates precision-cliff spikes across all tested precisions, reaches predicted margin equilibria, and outperforms prior remedies. Table 1 shows fixed CE achieving 99–100% test accuracy where standard CE stalls at roughly 94% and fp32 StableMax yields about 97% with margins ballooning to thousands (Barron et al., Table 1). However, comparisons to Zero-sum Projection and StableMax are restricted to modular arithmetic; their broader failure modes are not tested on other task families. The claim that StableMax "remains lower than with our fixed cross-entropy" is only validated on $a+b \bmod 31$, leaving open how the methods compare in realistic high-entropy classification settings where margins rarely approach $c_d$.

The paper provides complete inline JAX implementations for all three fixed functions and reports hyperparameters for grokking and distillation, including learning rate, betas, weight decay, and $\epsilon$ values. Random seeds are mentioned but not listed, and no public repository is linked. The modular arithmetic datasets are standard in the grokking literature, yet exact train/validation splits and model initializations are not fully specified. The core gradient identities can be verified immediately from the concrete numerical examples provided: "Standard implementations add 1 and $e^{-m}$ to form the denominator $1+e^{-m}$, which rounds to 1 when $m > c_d$ ... In contrast, our fixed expression $-e^{-m}/(1+e^{-m})$ keeps $e^{-m}$ in the numerator" (Barron et al., Section 2). Reproducing the full training curves, however, would require re-implementing data loaders and initialization protocols absent from the main text.