A Dozen Lines of JAX Fix Softmax's Hidden Gradient-Killing Bug
Two overlooked floating-point bugs in softmax and cross-entropy have been quietly corrupting gradients. A new paper fixes them in a few lines of JAX.
- A new paper from Google DeepMind researchers identifies two numerical bugs in standard softmax and cross-entropy implementations.
- Softmax Collapse zeros out target gradients once logit margin exceeds precision cliffs: 5.5 (bf16), 16.6 (fp32), 36.7 (fp64).
- Softmax Cancellation produces non-zero gradients in distillation even when student and teacher logits are identical.
- The fix rewrites log-softmax using
softplus,expm1, andlog1pto avoid adding tiny numbers to 1. - Fixed version eliminates Slingshot loss spikes during grokking and reveals logit margins equilibrate at
log(1/eps). - Available as drop-in JAX code with no meaningful slowdown; most useful for grokking, interpretability, and distillation.
Softmax and cross-entropy are foundational operations in JAX, PyTorch, TensorFlow, and other deep learning frameworks. An arXiv paper by Jonathan Barron, Dan Gnanapragasam, and Michael Tschannen identifies two floating-point failure modes in standard implementations. At extreme confidence, they can erase target gradients, produce loss spikes during grokking, and move a distilled student away from an identical teacher. The authors propose algebraically equivalent replacements that require roughly a dozen lines of JAX.
Two ways floating point bends softmax
Softmax converts a vector of logits into probabilities, while cross-entropy measures the probability assigned to the target. For a one-hot target, the gradient is the predicted probability vector minus the target vector. Those components should sum to zero because adding the same constant to every logit leaves the probabilities unchanged.
Most implementations prevent overflow by subtracting the largest logit before exponentiation. When the target has the largest logit, its probability can be written as 1 / (1 + exp(-m)), where m is the gap between the target logit and the log-sum-exp of the remaining logits. Once exp(-m) becomes too small to change the floating-point value 1.0, the denominator rounds to exactly 1.0. The loss and target gradient then round to zero.
| Format | Approximate margin cliff |
|---|---|
bf16 |
5.5 |
fp32 |
16.6 |
fp64 |
36.7 |
At this “Softmax Collapse” threshold, the non-target gradients remain representable because they sit near zero and never require adding a tiny value to 1.0. The resulting gradient components no longer sum to zero, breaking the shift-invariance property of cross-entropy.
Distillation exposes a second failure mode called “Softmax Cancellation.” A soft-target loss computes student and teacher probabilities separately, then forms their difference. As the models converge, subtracting these nearly equal values discards significant digits. Slight differences between the forward softmax and backward log-softmax operation sequences can even produce a nonzero gradient when the student and teacher logits are identical.
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
