DeepMind paper points out rounding bugs in Softmax and cross-entropy implementations; Gemini participated in discovery and fix

PyTorch got it wrong. JAX got it wrong. Even the standard formulas in textbooks are wrong.

Jon Barron, a Principal Research Scientist at Google DeepMind, dropped this bombshell!

This week, Google DeepMind showed a problem: in PyTorch, JAX, and various textbooks, every standard version of the algorithm has been quietly rounding a certain number to zero.

The more confident the model, the more easily the correct training signal vanishes into thin air. The optimizer thus loses control, and the training curve oscillates violently.

But the most chilling and also the most exciting part is this: this underlying logic flaw, hidden right under the noses of all humanity, was caught by AI (Gemini) itself.

This time, humans were only responsible for the post-hoc review and typesetting.

The vanishing number, the collapsing model

This bug is hidden in the most fundamental operations of AI: Softmax and cross-entropy.

If you have ever touched machine learning, these two words are certainly no strangers to you.

Softmax is responsible for converting the model's raw output scores (logits) into probabilities. Cross-entropy is responsible for computing the error, telling the model how badly it got things wrong.

All classifiers, all large language models, call this pair of functions in huge numbers every day, hundreds of millions of times.

How could they go wrong?

The problem lies in "limits" and "rounding."

To prevent computational overflow, existing standard frameworks (whether PyTorch or JAX) all perform a standard action:

First subtract the maximum score.

With this subtraction, the denominator becomes a structure like "1 + a tiny number."

This "tiny number" depends on how confident the model is.

The more confident the model, the more the score of the correct answer exceeds the other answers, and the smaller this number becomes.

But here lies a key detail: mathematical equivalence does not mean the results are completely equivalent in computer floating-point arithmetic.

In the world of computer floating-point numbers, if you add a sufficiently small number to 1.0, the computer will think "this is too small, no need to compute it" and directly round the result to exactly 1.0.

In the bfloat16 format, this cliff appears at a value of around 5.5. In fp32, it appears at around 16.6.

Once you cross this cliff, what happens?

The correct gradient instantly becomes absolute zero.

The model is clearly still learning, but the feedback signal about the correct answer is directly erased by the computer's rounding mechanism.

DeepMind's paper calls this Softmax Collapse.

And the optimizer makes things worse.

The principle of the mainstream optimizer Adam is to divide each update step by the average of recent gradients. When the gradient becomes zero for a long time due to the bug, this average also approaches zero infinitely.

When the real gradient finally reappears, Adam divides this gradient by a tiny number.

The result is that the model lunges forward a huge step, and the loss function (Loss) instantly spikes.

The academic community previously observed this phenomenon and gave it a nice name, "Slingshots," thinking it was some advanced feature of the model during the "Grokking" process.

After all that, it's not a feature, it's a bug.

Besides collapse, there is also Softmax Cancellation.

During knowledge distillation, if the student model completely copies the teacher, the gradient should theoretically be zero.

But because the standard code computes the probabilities separately and then subtracts them, rounding error leaves behind a "ghost gradient." The student model inexplicably deviates from the correct answer.

All of this is because a tiny number was lost.

Jonathan T. Barron's team at DeepMind submitted a paper on arXiv: "Fixing Softmax and Cross-Entropy."

The next day, Barron dropped a bombshell on the X platform:

It turns out that all existing implementations of Softmax and cross-entropy (including JAX, PyTorch, textbooks, etc.) have a subtle bug—actually two.

A simple rewrite can fix them and eliminate the "slingshot effect" during grokking.

For cross-entropy (fixing collapse): Take the target out of the sum, directly compute the margin using log-sum-exp, then return the softplus of the negative margin. Core principle: never add a tiny number to 1.

For distillation (fixing cancellation): Use expm1 and log1p for computation in log space. This way, when the student perfectly matches the teacher, the gradient is perfectly absolute zero, with no ghost residue.

The paper provides intuitive data tests:

In the modular addition (a + b mod 31) task, using the standard code, the test accuracy under bfloat16 precision is only 93.3% (and extremely unstable).

With the fixed code, the accuracy directly soars to 100.0%, rock solid.

And in the distillation test, the fixed code keeps the divergence between the student model and the teacher model at zero, maintaining a perfect 100% accuracy.

But what really set the entire tech world abuzz was a sentence Barron added in the comments:

This bug was initially discovered and fixed by Gemini Flash 3.7. Most of the experiments in the paper were completed by Flash 3.8. Gemini 4 Argon wrote the first draft of the paper and the figures.

I merely nitpicked with extreme rigor until the content was solid.

On October 24, independent engineer Satyajit Ghana carried out a comprehensive reproduction in PyTorch.

He reproduced both bugs on the CPU using printed values, ported the fixes, and confirmed the results in the paper at fp32 precision.

His conclusion was resounding: "Regardless of other details, the mechanism has been confirmed without reservation."

The tragedy of history: machines are always dropping digits

Digits being quietly dropped is not a first. Every time it happens, the cost is extremely painful.

In 1965, Professor William Kahan of the University of Toronto issued a warning in Communications of the ACM.

He pointed out that when a computer accumulates tens of thousands of numbers, each addition produces a rounding error. Add a million numbers, and you may lose 6 significant digits.

His solution was extremely simple: add one more variable, store the error lost last time, and add it back next time. This is the famous "Kahan summation algorithm."

In the article, Professor Kahan left a mocking remark full of prophetic meaning: "Double precision will soon be universally accepted as a replacement for intelligence in solving numerical problems."

61 years later today, in pursuit of faster chip speeds, AI has been downgraded all the way to the 16-bit bfloat16 format. Precision has become lower, and intelligence must be brought back to the table.

In 1991, it cost people their lives.

During the Gulf War, a U.S. Patriot missile failed to intercept an incoming Scud missile. The missile hit a barracks, and 28 American soldiers died.

The subsequent investigation report by the U.S. General Accounting Office (GAO) revealed the cause of the tragedy. The system's internal clock kept time in units of one-tenth of a second. But when converting this time into a real number, the register was only 24 bits long.

As the system ran, conversion errors continued to accumulate.

The missile battery had been running continuously for more than 100 hours. After 8 hours of operation, the error was 0.02 seconds, with a distance deviation of 55 meters. After 100 hours of operation, the error was 0.34 seconds, with a distance deviation as high as 687 meters. The target had long since flown out of the tracking range set by the radar.

An invisible rounding turned into a fatal deviation of 687 meters four days later.

In 1994, it cost Intel a huge sum of money.

In October 1994, Thomas Nicely, a mathematics professor at Lynchburg College, discovered that his two computers produced different results.

He was calculating the sum of the reciprocals of twin primes. With the same program, the Pentium chip and an old 486 produced divergent results. A few days later, he pinpointed the culprit: when calculating 1 divided by 824,633,702,441, the Pentium processor's floating-point unit calculated it incorrectly.

The reason was that the Pentium's divider used a lookup table, and the table was missing 5 entries.

What uncovered the problem was not only Professor Nicely's comparison of two computers, but also rigorous cross-validation.

Ultimately, Intel was forced to agree to unconditionally replace chips for all users. In January 1995, Intel set aside a huge loss of $475 million for this.

And the most inspiring part is: this time, it was the machine that found its own defect, and humans were responsible for the final check.

References:

This article comes from the WeChat public account "Xinzhiyuan", author: ASI Revelation, published by 36Kr with authorization.