Google DeepMind's Jon Barron says every standard softmax and cross entropy implementation has two subtle numerical bugs, and a short rewrite fixes both and removes grokking slingshots
Hi all, I have a new result I'm excited to share. It turns out that *all existing implementations* of softmax and cross-entropy (that's JAX, PyTorch, textbooks, etc) have a subtle bug --- two bugs, actually. A simple rewrite fixes both, and eliminates "slingshots" when grokking:
Context
Softmax and cross entropy sit inside nearly every classifier and language model. If the paper holds up, a small code change could make training and distillation more stable, though the effects shown so far are in research settings.
The paper, A Fix for Softmax and Cross Entropy, is by Jonathan T. Barron, Dan Gnanapragasam and Michael Tschannen of Google DeepMind and is posted on arXiv. Grokking is when a network memorizes its training data and only generalizes much later, and slingshots are sudden loss spikes seen during that phase. The post was published at 3:03 PM Eastern on October 23.
Deep dive
Jon Barron posted at 3:03 PM Eastern on October 23 that all existing implementations of softmax and cross entropy, including JAX, PyTorch and textbooks, have two subtle bugs, and that a simple rewrite fixes both. The claim that every implementation is affected is the authors' own and has not yet been independently reviewed.
The paper names the bugs. Softmax Collapse happens when one logit is much larger than the others: its gradient rounds to zero in floating point while the other gradients do not, so the gradients no longer sum to zero. Softmax Cancellation happens when student logits are very close to teacher logits, as in distillation, and produces non zero gradients even when the two are equal.
The authors tie the first bug to the slingshot effect seen in grokking. After memorization, gradients nearly vanish, Adam's running estimate of gradient size shrinks, and when gradients switch back on the optimizer takes an oversized step that spikes the loss. With their rewritten cross entropy, they report, the gradient no longer rounds to zero, the spikes disappear, grokking with weight decay speeds up, and accuracy collapse without weight decay is prevented. They also report that student networks stop drifting from the teacher in distillation.
In a later post in the thread, Barron writes that the bug was first found and fixed by a Gemini Flash model, that most experiments were run by Gemini Flash 3.8, and that Gemini 4 Argon wrote a first draft that he then heavily revised.
What to expect next
The practical test is whether the fixed functions make a difference at scale, since the paper's evidence centers on grokking and distillation experiments. Watch for framework maintainers in PyTorch and JAX responding, and for others reproducing the results on larger training runs. Barron's account of Gemini models finding the bug and running most experiments will also draw attention as an example of AI assisted research, which other labs may try to replicate.