Bug one: softmax collapse
When a model is very confident, the score of the correct class far exceeds the others. The softmax denominator becomes one plus a tiny sum, and in floating-point arithmetic that sum vanishes in rounding. The gradient for the correct class is computed as a probability close to one minus one — and turns into an exact zero.
The gradients for the other classes stay positive, so the gradients no longer sum to zero as they should. The threshold at which this happens depends on the precision of the computation.
Authors' data. In bf16, popular for training large models, the threshold is the lowest.
Bug two: mutual cancellation
In distillation the student learns to reproduce the teacher's distribution. The standard implementation computes both probabilities separately and subtracts them. When the student's outputs almost match the teacher's, the difference is lost to rounding — and the gradient is non-zero even where student and teacher match exactly.
How it was fixed
- For ordinary training. The sum of exponents over the wrong classes is computed separately and carefully combined with the correct one. The gradient of the correct class is built as minus the sum of the others, not by subtracting one from a rounded probability.
- For distillation. The difference is computed directly from the logit difference, using functions that keep precision near zero. When student and teacher match, the gradient becomes exactly zero.
What it gives
On modular arithmetic tasks the "slingshots" disappear — recurring spikes of the loss late in training with Adam. Without weight regularization the effect is sharper: with standard cross-entropy test accuracy drops from 61 to 23%, while with the fixed one it keeps growing to 85%.
Authors' data; the task is the square of a sum modulo 97.
In distillation where the student starts from the teacher's weights, the fixed version keeps the divergence at zero and accuracy at 100%. The standard one in bf16 drives the student away from the teacher. A curious detail: the bug was found and fixed by a Gemini model, and most of the experiments were also run by AI.
What it means in practice
If you distill models in bf16 or train to very high confidence, it is worth checking your loss implementation — the fix is a drop-in replacement for the standard one. The authors are upfront that such regimes are rare in ordinary large-scale training, the tests were run on small tasks, and the paper has no speed measurements.