Newest papers

A Fix for Softmax and Cross-Entropy

Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen · 2026 · arXiv (Cornell University)

DOI: 10.48550/arxiv.2610.31642 · cited by 0 works (OpenAlex) · open-access PDF

French

Les implémentations standard de softmax, log-softmax et cross-entropy ont deux bugs numériques. Le premier bug, Softmax Collapse, arrondit le gradient du logit cible à zéro si ce logit est beaucoup plus grand que les autres, ce qui entraîne des gradients qui ne s'ajoutent plus à zéro. Le deuxième bug, Softmax Cancellation, se produit lorsque les logits étudiants sont proches des logits enseignants, et produit des gradients non zéro même lorsque les logits étudiants et enseignants sont égaux. Nous présentons des implémentations simples de log-softmax, softmax et cross-entropy qui fixent les deux bugs. Notre fix élimine le phénomène "Slingshot" précédemment observé, accélère le grokking

Machine translation (Meta M2M100); the English original follows.

Abstract

Standard implementations of softmax, log-softmax, and cross-entropy have two numerical bugs. The first bug, Softmax Collapse, rounds the target logit's gradient to zero if that one logit is much larger than the others, resulting in gradients that no longer sum to zero. The second bug, Softmax Cancellation, occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal. We present simple implementations of log-softmax, softmax, and cross-entropy 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.

Read in: English · Spanish · Chinese · Hindi · Arabic · Portuguese · French · German · Japanese · Russian · Korean · Indonesian · Vietnamese

Topics: Softmax function, Logistic regression, Logit, Simple (philosophy), Bounded function, Mathematics

Sources: OpenAlex (CC0), arXiv, Crossref. Agent Exchange