A Fix for Softmax and Cross-Entropy
Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen
cs.LG
2026-10-23
DeepMind rewrites softmax and cross-entropy in log space to fix two rounding bugs; Slingshot spikes vanish and grokking reaches 99-100% test accuracy in bf16/fp32/fp64.
Every deep learning framework computes softmax and cross-entropy by first subtracting the max logit to avoid overflow. That shift prevents overflow but leaves two rounding bugs, and both land in the gradients. Softmax Collapse: when a model is highly confident in the correct class, with the target logit far above the rest, standard implementations round the target logit's gradient to zero while the other K−1 gradients stay positive, so the gradient vector no longer sums to zero. Softmax Cancellation: in soft-target distillation, when student logits sit right on top of teacher logits, the standard loss still emits a non-zero gradient that pushes the student away from its teacher.
These bugs are not cosmetic. The Slingshot phenomenon in grokking experiments, where training loss alternates between flat near-zero plateaus and sudden spikes, had been treated as an intrinsic quirk of adaptive optimizers. This paper shows it is Softmax Collapse playing out under Adam. Earlier work dodged the issue with surrogate losses, gradient projection, or fp64 casting; this one fixes the functions themselves.
Both bugs share one root cause: adding a very small number to 1.0. Floating point is extremely fine near zero, fp32 reaches about 1e-38, but anything below roughly 1e-7 vanishes when added to 1.0. After the standard max-shift the target term is e^0 = 1 and the denominator is 1 + e^-m, where m is the logit margin, how far the target logit leads the log-sum-exp of the rest. Once m crosses a precision cliff (about 5.5 in bf16, 16.6 in fp32, 36.7 in fp64), e^-m is rounded away: the loss and the target gradient both become zero, while the non-target gradients, computed from negative differences without ever adding 1, survive down to exponent underflow (m ≈ 88).
The fix keeps e^-m out of any add-to-1 position:
Each fix is roughly ten lines of JAX, with no internal upcasting and, per the authors, no meaningful slowdown.
Main experiment: modular-addition grokking (1-layer width-64 Transformer, a+b mod 31, 50% train split, full-batch AdamW, weight decay λ=1) across three precisions and three values of ε:
| Setting | Standard CE | Fixed CE |
| fp32, ε=1e-8 | 94.3±5.4% | 100.0±0.0% |
| bf16, ε=1e-8 | 93.3±7.4% | 100.0±0.1% |
| fp64, ε=1e-16 | 97.2±3.2% | 99–100% |
In all nine settings the fixed loss drives the margin to the predicted equilibrium log(1/ε) and lands 99–100% test accuracy; the standard loss stalls at the precision cliff of its format. At ε=1e-16 the equilibrium (36.8) sits just above the fp64 cliff (36.7), and even fp64 standard cross-entropy breaks. Without weight decay the damage is sharper: on a²+b² mod 97, standard CE tracks the fix up to 61% accuracy, then one spike at the cliff around step 1250 drops it to 23% (60% to 21% on a³+b mod 97), while the fixed version glides to 85%.
Two recent patches fare worse: Zero-sum Projection, which projects gradients back to zero sum, falls below 50% accuracy for ε ≤ 1e-12; StableMax, which replaces exp with a piecewise rational function, blows margins up to 10³–10⁴ and stays below the fix at every ε. In the distillation test, a student initialized with the teacher's weights should not move at all: the true gradient is zero. The standard loss's residual gradient is about 1e-3 in bf16, above ε=1e-8, so Adam normalizes it into full-size steps and the student drifts off; the fix stays at KL=0 and 100% accuracy in every precision.
These three functions run in every classification step of every framework, and the fixed versions are drop-in: ten-odd lines, no upcasting, no claimed slowdown. Teams distilling in bf16 get quantified evidence of a 1e-3 noise floor actively pushing students away from teachers. Grokking and optimizer-stability research gets a complete causal chain for Slingshot: target gradients zero out, Adam's second moment decays, gradients snap back, and the shrunken denominator turns them into oversized steps.
A clean side result: under AdamW with weight decay, logit margins converge to m ≈ log(1/ε). The default ε=1e-8 puts that equilibrium at 18.4, just above the fp32 cliff of 16.6, which is why Slingshot is common in fp32; at ε ≥ 1e-7 the equilibrium drops below the cliff and standard CE settles without spiking. The pair of ε and weight decay acts like an implicit L2 regularizer with weight ελ.
For standard large-scale pretraining the direct impact is limited; margins rarely get pushed past 16. In high-confidence regimes, memorization on small datasets, long-tail classes, distillation run to convergence, this is cheap insurance.
The authors state it themselves: extreme margins and near-zero student-teacher gaps are uncommon in standard large-scale training. All experiments are small-scale, modular arithmetic grokking plus a toy distillation with the student initialized at the teacher. There is no real teacher-student LLM distillation, no large-scale training run, and the no-slowdown claim is qualitative with no timings.
The paper also shows Slingshot spikes can help when LayerNorm and no weight decay let each spike inflate the weight norm, an accidental learning-rate decay. Removing them is not uniformly better; the gain is predictability. And the distillation harm is demonstrated in a deliberately extreme setup; how much the spurious gradient matters in real distillation, where student and teacher differ, goes untested.