arXiv Machine Learning
Oct 23

A Fix for Softmax and Cross-Entropy

The paper identifies two numerical bugs in standard softmax, log‑softmax, and cross‑entropy implementations: Softmax Collapse, where a dominant logit causes its gradient to be rounded to zero, and Softmax Cancellation, where student and teacher logits that are equal still produce non‑zero gradients. The authors provide straightforward implementations that correct both issues. Their fix removes the previously observed "Slingshot" phenomenon, accelerates grokking with weight decay, prevents accuracy collapse without weight decay, keeps student networks aligned with teachers during distillation, and demonstrates that post‑memorization logit growth is bounded by an equilibrium between weight decay and Adam’s epsilon hyperparameter.

By Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen