DeepMind paper finds numerical bugs in softmax and cross-entropy; fix lifts grokking accuracy to 100%

JFPuget · x · 2026-10-24

A new DeepMind paper by Jonathan Barron, Dan Gnanapragasam, and Michael Tschannen shows standard implementations of softmax, log-softmax, and cross-entropy contain two numerical bugs: Softmax Collapse (a dominant logit's gradient rounds to zero, so gradients no longer sum to zero) and Softmax Cancellation (non-zero gradients persist in distillation even when student and teacher logits match). These errors can trigger sudden loss spikes and disrupt knowledge distillation.
The fix computes the tiny residual quantity directly instead of deriving it by subtracting nearly equal rounded numbers — no higher numerical precision needed.
On small-scale grokking benchmarks, the fix eliminates "Slingshot" loss spikes and raises test accuracy from 94.3% to 100% in one setting; it also stops student networks from drifting away from teachers during distillation, and shows post-memorization logit growth is bounded by an equilibrium between weight decay and Adam's epsilon.
Caveat: results are from toy benchmarks, not large-scale training. Commenters are already asking when PyTorch will adopt the fix.

More from Research