Source-linked AI summary
A Fix for Softmax and Cross-Entropy
Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen
TL;DR
Standard softmax and cross-entropy have two numerical bugs that distort gradients at high confidence and near-equal student–teacher logits. The paper rewrites these functions to preserve accurate gradients, eliminating Slingshot spikes and preventing student drift during distillation.
Problem
Standard implementations can zero the target gradient at high confidence or produce nonzero gradients when student logits equal teacher logits.
Source: §1 · pp. 1–2, §7 · p. 8Method
The paper rewrites log-softmax, softmax, and cross-entropy in log space using stable formulas to preserve accurate gradients without higher-precision casting.
Source: §2 · p. 3, §3 · p. 4, §8 · pp. 8–9Results
Takeaways & Limitations
The corrected functions are drop-in replacements that can be implemented without meaningful slowdown, and AdamW margins with weight decay converge to log(1/ϵ).
Source: §5 · p. 6, §8 · p. 9The experiments probe extreme margins and near-zero student–teacher gaps that are uncommon in standard large-scale training.
Source: §8 · p. 9
Abstract
from arXiv · showhide
Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, rounds the target logit's gradient to zero if that one logit is much larger than the others, resulting in gradients that no longer sum to zero. The second bug, Softmax Cancellation, occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal. We present simple implementations of log-softmax, softmax, and cross-entropy that fix both bugs. Our fix eliminates the previously-observed "Slingshot" phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks from drifting away from the teacher during distillation, and shows that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's $\epsilon$ hyperparameter.
1 INTRODUCTION
Max-shifting logits prevents overflow, but standard softmax and cross-entropy still have two numerical bugs. Softmax Collapse breaks gradients at high confidence, while Softmax Cancellation leaves residual gradients when student logits match teacher logits.
1 INTRODUCTION
- Max-shifting logits prevents overflow, but two numerical bugs remain in standard softmax and cross-entropy.For target distribution p, the exact cross-entropy gradient is q − p.Source: §1 · p. 1
- Softmax Collapse occurs at high confidence when the target logit dominates the non-target logits, causing standard implementations to round its gradient to zero.The resulting target and non-target gradients no longer sum to zero after the non-target terms remain nonzero beyond the precision cliff.Source: §1 · p. 2
- Softmax Cancellation occurs in distillation and KL regularization when separately computed student and teacher probabilities lose precision upon subtraction, leaving a nonzero gradient even at equal logits.The paper fixes this by computing q − p directly from the student–teacher logit difference in log space.Source: §1 · p. 2
2 SOFTMAX COLLAPSE
Softmax Collapse arises because adding a tiny non-target exponential to 1 can erase it, zeroing the target gradient while leaving non-target gradients nonzero. The proposed log-space formulas avoid that addition and preserve zero-sum gradients.
2 SOFTMAX COLLAPSE
- The fixed target gradient keeps the small exponential in its numerator, preserving the target gradient down to exponent underflow instead of rounding it to zero at the precision cliff.Standard cross-entropy rounds the denominator 1 + e^-m to 1, causing the loss and target gradient to become zero.Source: §2 · p. 2
- For standard cross-entropy, non-target gradients remain positive after the target gradient rounds to zero, breaking the zero-sum gradient property.The same collapse affects standalone softmax, whose largest-logit gradient becomes zero while the remaining gradients stay nonzero.Source: §2 · pp. 2–3
- The fix excludes the largest logit from the log-sum-exp margin calculation and combines the margin with stable softplus, avoiding direct addition of 1 and e^-m.For hard-target cross-entropy, the target can replace the argmax, and autodiff computes its gradient as the negative sum of non-target gradients.Source: §2 · p. 3
- For z = [9.0, −9.0, −9.0] in fp32, the standard target loss and gradient are zero, whereas the fixed loss is 3.046e-8 and its gradients sum to zero.The fixed target gradient is −3.046e-8 and each non-target gradient is 1.523e-8.Source: §2 · p. 3
3 SOFTMAX CANCELLATION
Standard soft-target cross-entropy loses precision when student and teacher logits are close, producing a nonzero gradient even when they match. The proposed method computes their probability difference in log space to preserve accurate gradients near equality.
3 SOFTMAX CANCELLATION
- The method computes q − p from student–teacher logit differences in log space using expm1 and log1p, avoiding subtraction of nearly equal probabilities.It centers the logit difference under p and corrects residual drift so the gradient sums to zero.Source: §3 · p. 4
- When student logits equal teacher logits, the fixed gradient is zero, unlike standard implementations that retain a nonzero gradient from mismatched forward and backward computations.The true gradient shrinks toward zero as student logits approach teacher logits.Source: §3 · p. 4
- For a small fp32 student-logit perturbation, the fixed implementation matches the exact fp64 gradient, while the standard gradient shrinks one entry by 3× and flips another entry's sign.The standard gradient also fails to sum to zero in this example.Source: §3 · p. 4
4 THE SLINGSHOT MECHANISM
Softmax Collapse can trigger Slingshot spikes after memorization, with effects that depend on how weight norms are controlled. Fixed cross-entropy prevents the spikes, speeding grokking with weight decay and avoiding accuracy collapse without it.
4 THE SLINGSHOT MECHANISM
- Slingshot was observed after memorization as alternating near-zero loss plateaus and sudden spikes that coincide with test-accuracy jumps.The cited work observed this mechanism during grokking on modular addition with Adam.Source: §4 · p. 4
- After memorization, margins can exceed the precision cliff, zeroing target gradients while non-target gradients remain positive and nearly vanishing parameter gradients let Adam’s second moment decay.When gradients turn back on below the cliff, Adam’s small second-moment denominator can produce a large parameter step and loss spike.Source: §4 · pp. 4–5
- Fixed cross-entropy keeps the target-logit gradient from rounding to zero at the cliff, so Adam’s moving averages decay smoothly and Slingshot spikes disappear.Figure 2 reports faster grokking on all five seeds with weight decay and smooth continued learning without it.Source: §4 · p. 5, Figure 2 · p. 5
- 94 ± 5% vs. 100 ± 0% test accuracy on a + b mod 31, while on a2 + b2 mod 97 standard CE falls from 61% to 23% and fixed CE reaches 85%.With weight decay, Slingshot spikes slow and hurt grokking; without weight decay or LayerNorm, a spike causes a sharp accuracy drop.Source: §4 · p. 5, Figure 2 · p. 5
5 MARGIN EQUILIBRIA
After memorization, weight decay can balance Adam’s normalized gradient step and bound logit-margin growth. With positive ε, the predicted equilibrium is m⋆(ε) = log(1/ε); without weight decay, margins continue to rise.
5 MARGIN EQUILIBRIA
- For ε = 0, the normalized gradient step remains magnitude 1 as gradients shrink, so margins grow until gradients underflow.In the reported fp32 setting, the margin reaches approximately 81.2 before gradient underflow.Source: §5 · p. 5, Figure 4 · p. 6
- When |gi| ≪ ε, the normalized step shrinks as gi/ε; balancing it against weight decay requires |gi| ≈ ελ|θi|.The resulting regularization acts like an L2 regularizer with weight ελ.Source: §5 · p. 6
- At ε = 10−8, the normalized Adam step stays near 1 and decays around log(1/ε) ≈ 18.4, while standard CE zeroes the target gradient at c16 ≈ 5.5 and c32 ≈ 16.6.Since the equilibrium lies just above the fp32 cliff, standard CE reaches the cliff before the equilibrium; for ε ≥ 10−7, the cited passage says it settles without spiking.Source: Figure 3 · p. 6, §5 · p. 6
- With weight decay and ε > 0, margins level off at the predicted equilibrium m⋆(ε) = log(1/ε); without weight decay, they continue drifting upward.Figure 4 shows this pattern across ε settings using fixed cross-entropy and LaProp on a + b mod 31.Source: Figure 4 · p. 6, §5 · p. 6
6 GROKKING ACROSS PRECISIONS AND BENCHMARKS
Fixed cross-entropy crosses precision cliffs that stall or disrupt standard cross-entropy, reaching high test accuracy across floating-point precisions and grokking benchmarks. Comparisons with alternative fixes show that they do not match its results across the tested settings.
6 GROKKING ACROSS PRECISIONS AND BENCHMARKS
- At ϵ = 10−8, fixed cross-entropy crosses the bf16 and fp32 precision cliffs smoothly, reaches the predicted equilibrium margin of about 18.4, and achieves 100% test accuracy in all three precisions.Standard cross-entropy stalls at c16 ≈ 5.5 in bf16 and repeatedly spikes at c32 ≈ 16.6 in fp32.Source: §6 · p. 6, Figure 5 · p. 7
- At ϵ = 10−16, fixed cross-entropy converges smoothly even though standard cross-entropy fails at the fp64 precision cliff c64.The predicted equilibrium margin, log(1/ϵ) ≈ 36.8, exceeds c64 ≈ 36.7.Source: §6 · p. 7, Figure 5 · p. 7
- Zero-sum projection stalls near margin 18 and drops below 50% test accuracy for ϵ ≤ 10−12; StableMax has lower test accuracy than fixed cross-entropy across tested ϵ values.StableMax's margins grow to 10^3–10^4.Source: §6 · p. 7
- Across ϵ ∈ {10−8, 10−12, 10−16}, fixed cross-entropy reaches margins of about 18.4, 27.6, and 36.8 and 99–100% test accuracy across precisions, while standard cross-entropy stalls near its precision cliff.The comparison covers modular addition a + b mod 31 with five seeds.Source: Table 1 · p. 7
- Across four grokking tasks, fp32 fixed cross-entropy matches fp64 standard cross-entropy and reaches 100% test accuracy, while fp32 standard cross-entropy spikes at c32.The tasks include 1-layer Transformers and 2-layer MLPs.Source: §6 · p. 7, Figure 6 · p. 8
7 SOFT-TARGET DISTILLATION
The authors test soft-target distillation with a student initialized from the teacher's weights, where the true gradient should be zero. Standard cross-entropy can nevertheless drive drift or a KL spike, while the fixed version preserves the initial match across precisions.
7 SOFT-TARGET DISTILLATION
- The experiment uses distillation temperature T = 4, keeping the scaled margin m/T below c16 so the teacher probabilities do not round to one-hot.The student is initialized with weights from a teacher trained to 100% test accuracy on a + b mod 53.Source: §7 · p. 8
- With the student initialized from the teacher, fixed soft-target cross-entropy keeps KL = 0 and 100% test accuracy across precisions, while standard cross-entropy causes bf16 drift and a brief fp32 KL spike.In bf16, the residual gradient is about 10−3; in fp32 it is about 10−8, against Adam ϵ = 10−8.Source: §7 · p. 8
- At initialization, zS = zT makes the true gradient zero, but standard soft-target cross-entropy produces a residual because its forward softmax and backward log-softmax differ.When the residual exceeds Adam's ϵ, Adam normalizes it into steps of size η.Source: §7 · p. 8
8 CONCLUSION
The paper presents fixes for Softmax Collapse and Softmax Cancellation by rewriting softmax-related functions in log space. The authors report accurate zero-sum gradients across three precisions and improvements in grokking and distillation behavior.
8 CONCLUSION
- The fixes address Softmax Collapse, which zeros the dominant logit's gradient while leaving non-dominant gradients nonzero, and Softmax Cancellation, which loses precision as student logits approach teacher logits.Rewriting the functions in log space with softplus, expm1, and log1p restores accurate zero-sum gradients across bf16, fp32, and fp64 without internally casting to higher precision.Source: §8 · pp. 8–9
- The authors report that the fixes eliminate Slingshot loss spikes during grokking, prevent student drift during distillation, and yield logit margins that converge to log(1/ϵ) under AdamW with weight decay.The authors note that the extreme margins and near-zero student–teacher gaps studied are uncommon in standard large-scale training, and that the functions are easy drop-in replacements without meaningful slowdown.Source: §8 · p. 9