Standard softmax and cross-entropy carry two rounding bugs — a DeepMind fix kills the grokking "slingshot"

Jonathan T. Barron, Dan Gnanapragasam, and Michael Tschannen (Google DeepMind) posted a paper showing that the softmax, log-softmax, and cross-entropy implementations used by every major framework carry two floating-point bugs: past a "precision cliff," the target class's gradient rounds to exactly zero while every other gradient stays live, and in distillation the gradient never fully cancels even when student logits equal teacher logits. Their drop-in fix — a few lines per function, with JAX code in the paper — eliminates the "slingshot" loss spikes in grokking experiments, speeds up grokking under weight decay, and stops distillation drift.

What shipped

The paper (arXiv:2610.31642, posted October 22) names the frameworks directly: the max-subtraction scheme used "in all deep learning frameworks (JAX, PyTorch, TensorFlow, etc.)" leaves two bugs behind. The authors publish corrected versions of log_softmax, softmax, and hard- and soft-target cross-entropy, plus an interactive demo. Independent code is not required — the paper's JAX listing is the reference implementation.

Bug one, Softmax Collapse. To avoid overflow, standard softmax shifts logits by their max, so the denominator becomes 1 + e^{-m}, where m is the logit margin — how far the target logit sits above the log-sum-exp of the rest. Floating point is coarse near 1.0: in fp32, anything below roughly 10⁻⁷ vanishes when added to 1.0.

Once e^{-m} drops below that threshold, 1 + e^{-m} rounds to exactly 1.0, the loss evaluates to exactly 0.0, and the target gradient q_t − 1 rounds to exactly 0.0. The non-target gradients are computed near zero, where precision is fine, so they survive all the way to exponent underflow (m ≈ 88).

The gradient vector therefore stops summing to zero, and under Adam/AdamW — which normalize each parameter's gradient by its own magnitude — that asymmetry produces the recurring "slingshot" spikes observed by Thilak et al. (2022) and the grokking failures studied by Prieto et al. (2025), who first identified the collapse and proposed a surrogate, StableMax. The cliff sits at margin ≈ 5.5 in bf16, 16.6 in fp32, and 36.7 in fp64.

The paper's worked example makes the failure concrete. With logits [9, -9, -9] and target 0 in fp32, standard cross-entropy reports loss 0.0 and gradient [0.0, 1.523e-8, 1.523e-8]; the correct values are loss 3.046e-8 and gradient [-3.046e-8, 1.523e-8, 1.523e-8]. The target entry — the argmax's own update — has been deleted.

Bug two, Softmax Cancellation. In distillation and KL regularization, the target p = softmax(z_T) comes from teacher logits. Standard implementations compute q = softmax(z_S) and p separately and subtract them. When z_S ≈ z_T, subtracting two nearly equal rounded probability vectors loses the small difference, and because the forward softmax and backward log-softmax run different op sequences, the gradient is non-zero even at z_S = z_T. In fp32 the gradient norm plateaus at a noise floor; independent replication by Frank Chu measures the floor around 1e-8–5e-8, and at a student–teacher gap of δ=1e-7, rounding error is 74% of the gradient's own norm.

The fix. For hard targets, never form 1 + e^{-m}: compute the margin by excluding the target from the log-sum-exp, combine with a stable softplus, and define the target gradient as the negated sum of the non-target gradients. Written that way, the target gradient cannot round away on its own, and gradients sum to zero by construction down to exponent underflow. For soft targets, the paper passes z_T into the loss and computes q − p directly from the logit gap δ = z_S − z_T in log space using expm1/log1p, returning exactly zero gradient when the logits coincide.

A side result, separate from the bug fix: after memorization, logit margins plateau at an equilibrium between weight decay and Adam's ε, predicted near log(1/ε) ≈ 18.4 for λ=1, ε=1e-8 — a phenomenon the alphaXiv overview distinguishes from the numerical bug.

The training results are author-measured, all on modular-arithmetic grokking tasks in fp32 (five seeds). On addition mod 31 with weight decay λ=1 and ε=1e-8, standard cross-entropy repeatedly drops and rebounds after hitting the fp32 cliff, while the fixed version crosses it smoothly and levels near margin 18.4; on squared addition mod 97 with λ=0 and ε=1e-16, standard-loss test accuracy falls from 61% to 23% around step 1,250 while fixed runs continue to 85%. Against StableMax — the previous remedy, which swaps softmax for a different function — fixed CE reaches 99.6±0.7% test accuracy with final margin 27.4 versus 95.6±2.0% and margin 5,661 for StableMax (mod 31, λ=1, ε=1e-12). In distillation (addition mod 53, temperature 4, student initialized at the teacher), the fixed loss holds KL≈0 and 100% accuracy across bf16, fp32, and fp64, while standard bf16 drifts away from the teacher and loses accuracy.

Why it matters

The headline shift is attribution. The slingshot phenomenon has been studied as an optimizer-dynamics effect — Thilak et al. even tied the spikes to improved generalization. In these experiments it is a floating-point artifact: a target gradient that blinks off at the precision cliff while every non-target gradient stays on. That changes what an engineer should do about spiky late-training curves — test the fixed cross-entropy before hunting for exotic LR schedules or optimizer settings.

The bf16 cliff at margin ≈5.5 is the number to internalize. It sits far closer to ordinary training dynamics than fp32's 16.6 — any bf16 model whose confident predictions cross margin 5.5 enters the same zeroed-gradient regime the paper demonstrates on toy tasks. The authors themselves hedge that "the extreme margins and near-zero student–teacher gaps probed by our experiments are uncommon in standard large-scale training" (quoted in full on Hacker News). The most exposed pipelines are low-precision training, late-stage confident classifiers, and distillation runs where the student approaches the teacher.

For anyone who wants to check their own framework today: the whole fix is algebraic rearrangement, not a surrogate loss, fp64 casting, or gradient projection. Chu reproduced both bugs in pure NumPy in an afternoon, matching the paper's worked example digit for digit — and added a caveat the paper glosses: the fp32 cliff is not a constant. Across class counts and reduction orderings he measured the cliff in a band from 15.8 to 16.7; the paper's 16.6 is the clean two-class idealization. Different frameworks' kernels will each have their own cliff, so a "did it hit zero at exactly 16.6" test measures your reduction tree, not the phenomenon.

The Adam connection is the final piece of folklore on the table. Barron notes that the default eps=1e-8 sits "juuust past" where the fp32 bug kicks in, which "may be why some think a small eps is dangerous" (thread). With the fix, logits are free to grow until weight decay and Adam's ε reach equilibrium — margin ≈ log(1/ε).

Caveats

• All training claims are author-reported, unpublished outside arXiv, and demonstrated only on small grokking tasks (mod 31 / mod 97 / mod 53). No LLM- or vision-scale run is shown, and the paper reports no runtime benchmark — "without meaningful slowdown" is a claim, not a measurement.
• Chu's replication covers the arithmetic of both bugs; he did not retrain models, so the slingshot/grokking curves rest on the authors' evidence alone.
• The "non-zero gradient at exact equality" sub-claim depends on op sequence: Chu measured exactly zero when student equaled teacher in his NumPy setup. Test per framework.
• The distillation experiment is deliberately extreme (student initialized at the teacher), not ordinary distillation from a distant student.
• Several experiment figures above (mod 31 / mod 97 / StableMax / distillation) could only be read via the alphaXiv overview of the paper rather than directly in the full text during research; they are author-reported and not independently reproduced.
• Process note from Barron's announcement thread: the bug was first found and fixed by Gemini Flash 3.7, with most experiments run by Flash 3.8 and the first paper draft written by Gemini 4 Argon (disclosure) — an unusually explicit record of how much of this was LLM-assisted.

Sources

• arXiv abstract — A Fix for Softmax and Cross-Entropy (2610.31642v1)
• Full text (HTML) on arXiv
• Project page, jonbarron.info/softmaxfix
• Interactive demo
• alphaXiv AI overview of the paper (experiment details)
• Frank Chu, "I measured softmax's precision cliff…" (independent NumPy replication)
• Barron's announcement thread (mirror)
• Hacker News discussion
• Thilak et al., "The Slingshot Mechanism" (arXiv 2206.04817)
• Prieto et al., "Grokking at the Edge of Numerical Stability" (arXiv 2501.04697)