A Fix for Softmax and Cross-Entropy: Two Numerical Bugs and Their Consequences

Forum topic · 小凯 · 2026-10-24

Summary

A paper by Jonathan T. Barron, Dan Gnanapragasam, and Michael Tschannen (arXiv:2610.31642) identifies two numerical bugs in standard softmax, log-softmax, and cross-entropy implementations. The first, Softmax Collapse, rounds a target logit's gradient to zero when it is much larger than the others, breaking the property that gradients sum to zero. The second, Softmax Cancellation, produces non-zero gradients during distillation even when student logits exactly equal teacher logits. The authors propose simple implementation fixes for both. The fix eliminates the previously observed "Slingshot" phenomenon, accelerates grokking with weight decay, prevents accuracy collapse when weight decay is absent, stops student networks from drifting away from their teacher during distillation, and reveals that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's epsilon hyperparameter. The paper was shared on the zhichai.net forum with a Chinese abstract alongside the original English abstract.

Paper Overview

Field: Machine Learning
Authors: Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen
Published: 2026-10-22
arXiv: 2610.31642

Abstract (translated)

Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs:

1. Softmax Collapse: the target logit's gradient is rounded to zero if that one logit is much larger than the others, resulting in gradients that no longer sum to zero. 2. Softmax Cancellation: occurs when student logits are close to teacher logits, producing non-zero gradients even when they are equal.

We present simple implementations that fix both bugs. Our fix eliminates the previously-observed "Slingshot" phenomenon, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, stops student networks from drifting away from the teacher during distillation, and shows that after memorization, logit growth is bounded by an equilibrium between weight decay and Adam's ε hyperparameter.

Key Findings

Two distinct numerical defects exist in softmax/log-softmax/cross-entropy implementations: Softmax Collapse and Softmax Cancellation.
Both can be fixed with simple implementation changes.
Fixing the bugs removes the "Slingshot" instability observed in prior work.
Grokking with weight decay is accelerated; accuracy collapse without weight decay is prevented.
In distillation, student networks no longer drift from the teacher.
Post-memorization logit growth is bounded by an equilibrium between weight decay and Adam's ε.

---

*Auto-collected on 2026-10-24*