Numerical bugs in standard softmax and cross-entropy cause gradient collapse and distillation drift

Barron, Gnanapragasam, and Tschannen uncover two pervasive numerical bugs in mainstream softmax implementations, presenting stable fixes that eliminate the slingshot effect and prevent distillation drift.

Primary source: arXiv:2610.31642: https://arxiv.org/abs/2610.31642

What’s new: Softmax, log-softmax, and categorical cross-entropy are foundational building blocks of deep learning, yet standard implementations across major deep learning frameworks contain two subtle numerical bugs. In A Fix for Softmax and Cross-Entropy (arXiv:2610.31642, 22 Oct 2026), Google DeepMind researchers Jonathan T. Barron, Dan Gnanapragasam, and Michael Tschannen identify and resolve both flaws. The first bug, termed Softmax Collapse, occurs when a target logit becomes much larger than non-target logits; floating-point rounding rounds the target logit gradient to zero, causing class gradients to no longer sum to zero. The second bug, termed Softmax Cancellation, arises when student logits approach teacher logits during knowledge distillation or soft-target training, where catastrophic cancellation generates spurious non-zero gradients even when student and teacher outputs match.

The authors present numerically robust implementations of log-softmax, softmax, and cross-entropy that eliminate both failure modes. In empirical evaluations across synthetic optimization tasks and neural network training, the corrected primitives eliminate the previously documented slingshot phenomenon, where models exhibit sudden parameter and loss oscillations during training. The fix also accelerates grokking under weight decay, prevents catastrophic accuracy collapse during training runs without weight decay, and prevents student networks from drifting away from their teachers during distillation. Finally, the analysis proves that post-memorization logit growth is governed by an equilibrium between weight decay and Adam’s epsilon hyperparameter.

Why it matters: Practitioners routinely attribute training instability, late-stage loss spikes, and distillation divergence to optimizer dynamics, data ordering, or architecture quirks. Demonstrating that these phenomena stem directly from low-level floating-point rounding in ubiquitous loss primitives shifts how engineers should diagnose optimization failures. Adopting stable softmax and cross-entropy implementations removes artificial gradient noise, providing cleaner optimization trajectories for large language models and knowledge distillation pipelines.

Caveats: The proposed implementations require modifying low-level autograd kernels, which may introduce minor computational overhead compared to naive implementations that rely on fused hardware operations. Furthermore, in massive-scale pre-training runs where logit growth is already constrained by logit capping, z-loss regularization, or aggressive weight decay, the practical performance gains may be less pronounced than in distillation or grokking regimes.