Google DeepMind researchers fix two numerical bugs in softmax and cross-entropy

A fix for a function every model uses

Three Google DeepMind researchers say the standard implementations of softmax, log-softmax and cross-entropy, used in essentially every classifier and language model, contain two numerical bugs. In a paper posted to arXiv on 22 October 2026, titled simply A Fix for Softmax and Cross-Entropy, they describe both bugs, give short replacement implementations, and show that the fix removes a training oddity researchers have puzzled over for years.

The functions are among the most heavily used pieces of code in machine learning. Softmax turns a model's raw scores, called logits, into probabilities, and cross-entropy is the loss most models are trained on. Every major framework, including JAX, PyTorch and TensorFlow, computes them the same way, so a bug here is not a quirk of one library.

Who wrote it and where

The authors are Jonathan T. Barron, Dan Gnanapragasam and Michael Tschannen, all listed with Google DeepMind. The paper is released under a CC BY 4.0 licence, and Barron has posted an HTML version on his site that includes JAX code for the fixed functions. It is a preprint and has not been peer reviewed.

It is a small, precise piece of work rather than a new model, the kind of numerical housekeeping that matters more as training moves to lower precision. Our coverage of mixed-precision reinforcement learning samplers touched the same tension between speed and accuracy in bf16 arithmetic.

Bug one: Softmax Collapse

Frameworks avoid overflow by subtracting the largest logit before exponentiating. That is standard practice, but it creates the first bug, which the authors call Softmax Collapse. When the model is confident and the correct class has by far the largest logit, the shifted target becomes zero, so its exponential is exactly one, and the other classes add a tiny number to it.

Floating-point numbers are very precise near zero but coarse near one. Once that tiny number is smaller than the rounding step around one, it disappears. The denominator becomes exactly 1.0, the predicted probability becomes exactly 1.0, and the loss and the gradient on the correct class round to zero.

The paper measures confidence with a logit margin and finds the cliff at a margin of about 5.5 in bf16, 16.6 in fp32 and 36.7 in fp64. The trouble is that the gradients for the other classes do not round to zero, because they are computed near zero, where precision is fine. They stay positive until the margin reaches about 88. So past the cliff the gradients no longer sum to zero, which mathematically they always should.

How that causes Slingshot spikes

That asymmetry explains the Slingshot mechanism, first described by Thilak and colleagues in 2022. In grokking experiments, where a small network memorises a task and generalises much later, training loss alternates between flat stretches near zero and sudden spikes.

The paper traces each spike to the cliff. When every example in a batch is above it, the gradients nearly vanish and Adam's running estimate of gradient size decays toward zero. Eventually weight decay pulls a few examples back below the cliff and their gradients switch back on. Adam divides those restored gradients by its now tiny estimate, takes a very large step, the loss spikes, and the cycle repeats. With the fixed loss, the target gradient shrinks smoothly instead of snapping to zero, and the spikes disappear.

Bug two: Softmax Cancellation

The second bug, Softmax Cancellation, affects knowledge distillation and KL regularisation, where a student model learns to match a teacher's probabilities. Standard code computes the student's softmax and the teacher's softmax separately and subtracts them. When the two are nearly equal, that subtraction loses most of its precision.

Worse, small differences between the softmax used in the forward pass and the log-softmax used in the backward pass leave a non-zero gradient even when the student's logits exactly equal the teacher's. The true gradient in that case is zero. The paper fixes this by computing the difference directly from the gap between student and teacher logits, in log space.

The fix

The fix rewrites the three functions in log space using softplus, expm1 and log1p, standard functions designed to stay accurate for very small inputs. The authors say the new versions give accurate, zero-sum gradients in bf16, fp32 and fp64 without internally casting to higher precision, and can be implemented without meaningful slowdown.

That distinguishes it from earlier workarounds. Previous papers cast to fp64, used surrogate losses, or projected gradients so they sum to zero. The authors compare against two of them. A zero-sum projection method from Liu and colleagues in 2026 stalls the margin near 18 and drops test accuracy below 50 percent at small Adam epsilon values. StableMax, from Prieto and colleagues in 2025, replaces the exponential with a slower-growing function, which lets margins blow up to around 10,000 and lags the fix on accuracy.

What the experiments show

Most experiments use small grokking tasks such as modular addition. On a one-layer Transformer learning a plus b mod 31, standard cross-entropy stalled at the bf16 cliff and spiked repeatedly at the fp32 cliff, while the fixed version crossed both and reached 100 percent test accuracy in all three precisions. In the paper's table, fixed cross-entropy in fp32 averaged 100 percent test accuracy at the default Adam epsilon of 1e-8, against 94.3 percent for the standard version.

Without weight decay the effect was starker. On a task learning a squared plus b squared mod 97, hitting the fp32 cliff around step 1,250 triggered a spike that dropped standard test accuracy from 61 percent to 23 percent, while the fixed loss continued smoothly to 85 percent. Across four grokking benchmarks, fp32 with the fix matched fp64 with the standard loss.

For distillation, the authors started a student with the teacher's exact weights, so it should never move. With standard code in bf16, the residual gradient pushed the student away and hurt test accuracy; in fp32 it caused a brief KL spike. The fixed loss kept the student at zero divergence and 100 percent accuracy in every precision.

A neat result about Adam

Removing the cliff exposed something else. With weight decay, the fixed loss's margin still levelled off, just above the fp32 cliff. The authors show this is an equilibrium between weight decay and Adam's epsilon setting, and that the margin settles near the logarithm of one over epsilon. With the default epsilon of 1e-8, that is about 18.4, which sits just above the fp32 cliff of 16.6. That, they argue, is why default Adam settings make Slingshot spikes appear in fp32 at all. It is the kind of result that sits well alongside our explainer on how training choices shape model outcomes.

Their LaProp experiments confirmed the prediction across many orders of magnitude of epsilon; with epsilon set to zero, margins grew until fp32 gradients underflowed at about 81.2.

How much this matters in practice

The authors are explicit that the extreme margins and near-zero student-teacher gaps in their experiments are uncommon in standard large-scale training. Nothing in the paper shows that today's frontier language models are being damaged by these bugs, and the experiments are on small networks and toy tasks.

The relevance grows in two places: when models make very confident predictions, and when they are distilled from close teachers, which is now routine for small models. Lower precision also brings the cliff much closer, from 16.6 in fp32 to 5.5 in bf16. Work that automates optimiser tuning, like MIT's Opal project, may find such numerical artefacts easier to spot once the baseline is clean.

What to do with it

The fixed functions are short drop-in replacements, and the JAX reference code is in the paper and on Barron's page. Teams that see unexplained loss spikes late in training, run distillation in bf16, or study grokking have the clearest reason to try them. It also adds to a run of Google DeepMind research this month, including its Lehmer polynomial search.

Framework maintainers are the other audience. No change to PyTorch, JAX or TensorFlow defaults has been announced, and any such change would need wider testing than one preprint. Still, the paper makes a clear case that one of the most familiar lines of code in deep learning has been rounding away information it did not need to lose.