Two numerical bugs in softmax cross-entropy

On October 22, 2026, three Google DeepMind researchers, Jonathan Barron, Dan Gnanapragasam and Michael Tschannen, posted arXiv paper 2610.31642, “A Fix for Softmax and Cross-Entropy.” Their claim: the standard softmax, log-softmax and cross-entropy implementations that PyTorch, JAX, TensorFlow and the textbooks share carry two floating-point bugs. Softmax Collapse rounds the target logit’s gradient to zero once that logit dominates. Softmax Cancellation leaves a nonzero gradient in a distillation loss even when the student’s logits exactly match the teacher’s. Barron’s project page ships a JAX reference implementation and an interactive demo.

Key points

arXiv:2610.31642 (v1, submitted October 22, 2026, cs.LG) names two numerical bugs in standard softmax, log-softmax and cross-entropy; all three authors are Google DeepMind researchers.
Per the authors’ project page, the Softmax Collapse precision cliff sits at margin ≈5.5 in bf16, ≈16.6 in fp32 and ≈36.7 in fp64. Past the cliff, the target logit’s gradient rounds to zero and gradients stop summing to zero.
The page embeds two JAX snippets (about 15 lines for hard targets, 12 for soft targets). An independent NumPy reproduction confirmed both bugs’ magnitudes and found the fp32 cliff drifts between 15.8 and 16.7 with the reduction order.

Background

softmax plus cross-entropy is the floor every classifier and every distillation run stands on, and textbooks and frameworks teach one canonical form: subtract the max logit to avoid overflow, exponentiate, normalize. That convention settles overflow. The paper goes after the rounding that comes next: once the target logit leads the rest by a margin m, the denominator’s 1+e⁻ᵐ rounds to 1.0, and the target gradient qₜ−1 collapses to zero with it. The project page ties this failure path to the previously documented Slingshot loss spikes, driven by Adam/AdamW, and to grokking. Precision versus training has surfaced on this site twice this month, both times on the same layer: the BF16 tensor-core unlock and vLLM’s Humming quantized kernels.

The two bugs

The project page spells out the mechanics. Softmax Collapse: when the target logit dominates, the denominator 1+e⁻ᵐ rounds to 1.0, so the target gradient gₜ = qₜ−1 rounds to 0.0 while the other K−1 gradients stay positive, and gradients no longer sum to zero. Softmax Cancellation: distillation and KL losses compute student and teacher softmaxes separately and subtract them; when the two logit vectors nearly match, subtracting nearly equal probabilities loses precision, and a nonzero gradient survives even when they are exactly equal (the project page’s wording).

The fix works in log space. Pull the target out of the log-sum-exp to get the margin m, write log-softmax as ẑ_k − softplus(−m), and compute the target gradient as −e⁻ᵐ/(1+e⁻ᵐ) instead of subtracting 1 from a sigmoid. For soft targets, expm1 and log1p recover q−p from z_S−z_T (Equation 9 on the page). The two embedded JAX snippets run about 15 and 12 lines, with an interactive demo.html alongside.

The cliff and the reproduction

The cliff location is the most checkable number in the argument, and it is where the two sources diverge. The project page’s figures (margin m is the gap between the target logit and the log-sum-exp of the rest):

Precision	Cliff (project page)	NumPy measurement (dev.to)
bf16	≈5.5	not tested separately
fp16	not given	7.7 (theory 7.62)
fp32	≈16.6	16.0, drifting 15.8–16.7 with class count
fp64	≈36.7	36.1 (theory 36.74)

frank chu’s reproduction on dev.to (October 23) uses no framework: pure NumPy 1.26.4 computing the analytic gradient q−p, the same term a framework’s backward pass computes. He confirmed both bugs. Past the cliff, the loss and target gradient read exactly 0.0; before it, quantization already bites (at m=15 the standard loss reads 2.38e−7 against a true 3.06e−7, about 22% off). The distillation gradient carries an error floor of roughly 1e−8 to 5e−8 (74% of the gradient norm at δ=1e−7), while the paper’s fix matches fp64 to 9.2e−16.

Two caveats came out of the same post. The cliff moves: sweeping class counts from K=2 to K=101 put the fp32 cliff anywhere in 15.8–16.7, which makes 16.6 the clean two-class idealization. And one sub-claim did not reproduce: the paper says the gradient stays nonzero when student and teacher logits are exactly equal; in his NumPy setup both paths run the same operations and the result is exactly 0.0. He attributes the difference to framework kernels and notes the error floor itself does not depend on them. He also did not retrain a grokking model, so the slingshot-to-Adam link still rests on the paper’s own training curves.

What others say

The dev.to reproduction takes the bugs-are-real side: magnitudes match the paper, but the cliff position measures your kernel’s reduction order, so hard-coding 16.6 into an assertion would be a mistake.

SamCodeMan’s critical note (Samir Sengupta, October 23) works from the abstract and lists what it withholds: no trigger threshold, no magnitude for the spurious gradients, no statement of which frameworks exhibit the bug at fp32/bf16/fp16, no distillation drift figures, no grokking speedup size, no formula or tested values for the Adam ε equilibrium, and no standalone package or production overhead data. He still calls the bug descriptions “specific enough to test against an existing stack now” and prescribes five steps: unit-test the distillation loss, asserting exactly zero gradient with identical student and teacher logits; test a row with one dominant logit, checking whether the target gradient is zero and gradients sum to zero; if either fails, replace log-softmax, softmax and cross-entropy with the paper’s implementations and rerun both tests; re-examine conclusions drawn from Slingshot-like behavior, accuracy collapse without weight decay, or student drift in distillation; and treat Adam’s epsilon as a setting that, together with weight decay, bounds logit growth after memorization.

The disagreement is about strength of evidence: the reproduction proves the arithmetic; the critic points out the training-scale payoffs remain the authors’ word.

How to try it

The shortest path is the author’s project page, jonbarron.info/softmaxfix: the two JAX snippets and an interactive demo sit inline, ready to paste into a loss function. Before that, run two unit tests from Sengupta’s checklist. Assert the distillation-loss gradient is exactly zero when student and teacher logits are identical, and on a row with one dominant logit check a zero target gradient and gradients that sum to zero. Whichever fails decides whether the rewrite is worth a retrain.

Our take

Teams writing their own training loops, running distillation, or training classifiers that predict with high confidence (large margins) should spend the ten minutes. The bf16 cliff at ≈5.5 sits far below fp32’s ≈16.6, and the authors’ equilibrium point m*(ε) ≈ log(1/ε) means margins climb after memorization anyway: with a common ε of 10⁻⁸, the equilibrium lands near 18.4, already past the fp32 cliff.

The honest gaps: this is a v1 paper with no standalone repository, no third-party validation at production scale, a cliff number that drifts with reduction order, and one sub-claim that failed to reproduce in pure NumPy. Actual framework kernel behavior and the fix’s runtime cost (the paper calls it a drop-in replacement “without meaningful slowdown,” the authors’ own wording) still await independent measurement.