A Fix for Softmax and Cross-Entropy

TL;DR

Standard implementations of softmax, log-softmax, and cross-entropy in JAX, PyTorch, and textbooks contain two subtle numerical bugs that corrupt gradients at high confidence and during distillation. Softmax Collapse rounds the target logit’s gradient to zero when one logit dominates, so gradients no longer sum to zero. Softmax Cancellation produces non-zero gradients between student and teacher logits that are already equal (relevant for warm-starting distillation from the teacher). A simple drop-in rewrite fixes both and empirically eliminates the “Slingshot” grokking artifact, speeds up grokking with weight decay, prevents accuracy collapse without weight decay, and keeps distilled students from drifting away from the teacher.

Key claims

Softmax Collapse: when one logit is much larger than the others, the target-logit gradient term silently rounds to zero in standard implementations, so the per-example gradient no longer sums to zero and the “excess” leaks into non-target logits [Abstract, §3].
Softmax Cancellation: when student logits are close to teacher logits (as at the start of distillation from a copy of the teacher), standard implementations produce non-zero gradients even when the two sets of logits are equal [Abstract, §3].
A single rewrite of softmax / log-softmax / cross-entropy fixes both bugs; the fix is a few lines and drop-in in both JAX and PyTorch [Abstract, §3].
With the fix, Slingshot loss spikes during grokking disappear and grokking proceeds smoothly past the fp32 cliff at logit margin m ≈ c₃₂ ≈ 16.6, saturating at m ≈ 18.4 across all five tested seeds [§4, Fig. 2].
Without weight decay, the fixed cross-entropy avoids the accuracy collapse that standard CE exhibits once logits grow past the Softmax Collapse threshold [Abstract].
After memorization, logit growth is bounded by an equilibrium between weight decay λ and Adam’s ε hyperparameter — the default eps=1e-8 sits just past the regime where Softmax Collapse begins in fp32, consistent with the folklore that very small ε is “dangerous” [Abstract].
In distillation, the Cancellation fix prevents student logits from drifting away from the teacher during warm start, which standard CE does not [Abstract].

Method

The authors derive the two bugs analytically from the standard log-softmax identity log σ(z)_t = z_t - logsumexp(z) and show both arise from catastrophic cancellation at specific operating points — one driven by the magnitude of the top logit (Collapse), the other by the magnitude of the gap between student and teacher logits (Cancellation). The proposed rewrite restructures the arithmetic so neither cancellation triggers, and the resulting gradient satisfies the “sums-to-zero” property exactly in finite precision. Experimentally, the authors instrument grokking runs on modular-arithmetic tasks following Prieto et al. 2025 (who identified but did not directly fix Softmax Collapse), swap in fixed CE, and show the Slingshot spikes and grokking delays vanish without any other change to the training recipe. A separate distillation experiment initializes a student from a copy of the teacher and shows standard CE produces drift while fixed CE keeps the student pinned.

Notably, the paper credits Gemini Flash 3.7 with originally finding and fixing the bug, with most experiments run by Flash 3.8 and Gemini 4 Argon writing the first draft — a public example of a frontier model identifying a numerical bug in a universally deployed primitive.

Results

Fixed CE eliminates Slingshot spikes and speeds up grokking under weight decay; crosses the fp32 numerical cliff at c₃₂ ≈ 16.6 smoothly rather than repeatedly dropping and rebounding [§4, Fig. 2].
Fixed CE prevents accuracy collapse when training without weight decay — a regime where standard CE eventually fails as logits grow past the Collapse threshold [Abstract].
Logit magnitude at saturation is set by the equilibrium between weight decay and Adam’s ε: with default eps=1e-8, saturation sits at m ≈ 18.4 under fixed CE, just above the fp32 cliff c₃₂ [§5].
Fixed CE keeps distillation students from drifting away from the teacher under warm-start initialization; standard CE does not [Abstract].

Why it’s interesting

This is a stability-locus paper that lands cleanly inside the ongoing wiki argument on Training stability at scale about where the load-bearing stability primitives actually live — architecture (QK-Norm / Gated Attention / mHC), optimizer (Muon / SSO / AdamH), parametrization (µP / µA), or precision (FP16 / NVFP4). Barron et al. add a fifth locus: the loss function’s finite-precision implementation. The mechanism connects directly to A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training‘s “outlier-driven implicit rescaling” story — logit outliers at the LM head are what Softmax Collapse interacts with — and to Lost in Backpropagation: The LM Head is a Gradient Bottleneck‘s claim that the LM head is a gradient bottleneck. It also complements Defeating the Training-Inference Mismatch via FP16‘s precision-side framing: both papers argue that a finite-precision implementation choice at the output layer is doing more work than the field assumed. For the distillation side, the Cancellation fix is directly relevant to the on-policy distillation line (Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe, On-Policy Distillation, On-policy distillation isn't a free lunch — four failure modes of OPD/OPSD in current post-training recipes (Bhavin Jawade)) — “student drifts from teacher” is a reported failure mode and this paper identifies a baseline numerical cause.