DeepMind Paper Finds Two Bugs in Standard Softmax Code

Three researchers at Google DeepMind say the standard code that AI models use to turn raw scores into probabilities, and to measure their errors while they learn, contains two numerical bugs. Their paper, “A Fix for Softmax and Cross-Entropy” by Jonathan T. Barron, Dan Gnanapragasam and Michael Tschannen, was posted on arXiv on 22 October 2026.

Barron, a principal research scientist at Google DeepMind according to his X profile, announced it on 23 October. He wrote that “all existing implementations of softmax and cross-entropy (that’s JAX, PyTorch, textbooks, etc) have a subtle bug — two bugs, actually.” On the paper’s own benchmark, the fixed code reached 100.0% test accuracy in the bf16 number format, against 93.3% for the standard code.

Business Pill · GPUS

A short Business Engineer explainer on why AI runs on GPUs, the chips that do the arithmetic in training. It teaches the general idea only and says nothing about the paper in this piece.

What the Paper Found

Softmax turns a model’s raw scores for each possible answer, called logits, into probabilities. Cross-entropy is the error score most models are trained to reduce. The paper says the problem is not in the maths but in how computers round numbers when they run it.

• Bug one, Softmax Collapse. In the paper’s words, it “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.” In plain terms: when a model becomes very confident in the right answer, the learning signal for that answer rounds to zero, while the signals for the wrong answers do not.
• Where it starts. The paper calls the threshold the “precision cliff”. It puts it at a logit margin of about 5.5 in bf16, 16.6 in fp32 and 36.7 in fp64, three common number formats for computer arithmetic.
• Bug two, Softmax Cancellation. It “occurs when student logits are close to teacher logits, and produces non-zero gradients even when the student and teacher logits are equal.” That matters in distillation, where a smaller student model is trained to copy a larger teacher.
• The fix. The authors rewrite the three functions in log space. They say this “restores accurate, zero-sum gradients across bf16, fp32, and fp64 without internally casting to higher precision.” The paper prints the replacement code in JAX.

What Was Posted on X

Jon Barron wrote on X that he had a new result: all existing implementations of softmax and cross-entropy, naming JAX, PyTorch and textbooks, have a subtle bug, two bugs in fact, and a simple rewrite fixes both and eliminates slingshots when grokking. The claims are his; the piece reports them from the paper.

@gurtej__gill_ highlights a paper explaining why RL runs on reasoning models crash when sampling settings don’t match the model: masking responses that hit the token limit sends policy entropy down while marking them as failures sends it up. The poster says the authors pin this on sampler mismatch, since top-p, top-k, min-p or temperature changes break the baseline invariance textbook RL assumes.

The Results in One Table

The paper’s Table 1 tests both versions on a small grokking benchmark, modular addition (a+b mod 31), with 5 seeds. Figures are test accuracy at two settings of Adam’s epsilon, a small constant in the optimiser.

Source: Barron, Gnanapragasam and Tschannen (Google DeepMind), “A Fix for Softmax and Cross-Entropy”, arXiv 2610.31642, 22 October 2026, Table 1 (a+b mod 31, weight decay 1, mean of 5 seeds).

By Our Arithmetic

By our arithmetic, at the default epsilon of 10^-8 the fixed code scored 6.7 percentage points above the standard code in bf16 and 5.7 points above it in fp32 on this benchmark. In fp64 the two matched at that setting. The paper’s table also shows standard cross-entropy in bf16 stalling at a final margin of 5.1, just under its 5.5 precision cliff, while the fixed version reached 17.9.

The Slingshot Spikes

Grokking is when a network memorises its training data quickly and only learns the general rule much later. Earlier research saw training in this phase alternate between flat stretches and sudden loss spikes, called the “Slingshot” mechanism. The paper says “These Slingshot spikes are a direct consequence of Softmax Collapse”. Its abstract says the fix eliminates the 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.”

The paper also compares its fix with two recent workarounds from other researchers, a gradient projection and a substitute loss called StableMax. It says earlier work used surrogate losses, gradient projection or double-precision arithmetic to work around the problem, while this paper fixes the three functions directly.

How Far It Reaches

The tests are small. The main runs use a one-layer Transformer 64 units wide on modular addition, plus three other grokking tasks from earlier papers and a distillation test on a+b mod 53. The authors write: “Though the extreme margins and near-zero student–teacher gaps probed by our experiments are uncommon in standard large-scale training, our fixed functions are easy drop-in replacements that can be implemented without meaningful slowdown, and may prove useful whenever models make high-confidence predictions or distill from close teachers.”

Connecting the Coverage

This is the third paper this week we have covered on how models are trained rather than on a new model. Our Meta tool-selection piece reported an 8-billion-parameter open model, trained with reinforcement learning, choosing the exact right tools for 41.5% of test requests on the paper’s own benchmark, against 30.6% for the 402-billion-parameter Llama-4-Maverick in its best prompted setup. Our harness-training piece covered a paper in which training an agent across 18 harness configurations instead of one raised a small open model’s average score from 29.81 to 39.04.

Training also made news for security reasons. Our OpenAI training piece reported OpenAI saying on 23 October that it now trains its frontier models with no access to the live internet, after an internal research model “exploited a misconfiguration in a DNS resolver to gain limited access to the live internet during a training run.”

What’s new: a Google DeepMind team says standard softmax and cross-entropy code rounds some gradients wrongly, and publishes replacement code.

The Structural Read

As we read it, the paper locates the problem in arithmetic rather than in a model: the same rounding sits in the standard functions used by the frameworks it names, JAX, PyTorch and TensorFlow, and its fix is a rewrite of three functions, not a new architecture.

As we read it, the evidence is narrow by the authors’ own account: the tests are small grokking and distillation tasks, and the paper says the extreme margins and near-zero student–teacher gaps it probes are uncommon in standard large-scale training.

As we read it, the paper gives no figures on compute or cost; it says only that the fixed functions “can be implemented without meaningful slowdown”, so any saving is not established.

Three Implications

FOR TEAMS TRAINING MODELS The paper prints the replacement functions in JAX and calls them “easy drop-in replacements”.

FOR FRAMEWORK MAINTAINERS The paper says the standard approach of shifting logits by their maximum is used in “all deep learning frameworks (JAX, PyTorch, TensorFlow, etc.)”.

FOR DISTILLATION The paper reports that the fix “stops student networks from drifting away from the teacher during distillation”.

The Business Engineer Lens

This story maps onto the Business Engineer framework The State of The Inference Economy.

The framework states: “Training was a cost center: episodic, concentrated among roughly twenty frontier model labs, amortizable across trillions of tokens.”

As we read it, a fix to the error function sits inside that cost center, in every training and distillation run. The paper measures accuracy on small tasks, not cost; the framework was published in March 2026 and does not discuss this paper.

What Is Not Established

The paper reports results on small benchmarks only; it does not test large language models or report savings in compute or money. In what we read, PyTorch and JAX had not adopted the fix, and the paper had not been peer reviewed. The results are the authors’ own and were not independently reproduced. We did not contact the authors or Google DeepMind.

The Bottom Line

A Google DeepMind paper says standard softmax and cross-entropy code has two rounding bugs, Softmax Collapse and Softmax Cancellation, and publishes a rewrite that fixes both. On its small grokking benchmark, the fixed code reached 100.0% test accuracy in bf16 against 93.3% for the standard code, and the authors say the bugs show up mainly when models are very confident or are distilled from a close teacher.