DeepMind Uncovers Hidden Bugs in Softmax and Cross-Entropy Implementations
Google DeepMind researchers identified numerical precision bugs in standard Softmax and Cross-Entropy implementations across major frameworks like PyTorch and JAX, which can cause training instability. They proposed a fix involving computational adjustments without increasing floating-point precision, with Gemini playing a key role in discovering the issue and validating experiments.
Full Report
Contents2 sections
Gemini discovers bugs and handles a large portion of the experiments.
Jiqizhixin Editorial Team
Two of the most fundamental and widely used functions in deep learning have been found to contain two subtle bugs.
Yesterday, Jonathan T. Barron, Chief Research Scientist at Google DeepMind, announced new research on X.
He pointed out that standard implementations of Softmax and Cross-Entropy (CE), which are universally adopted by mainstream deep learning frameworks such as PyTorch, JAX, and TensorFlow, suffer from two easily overlooked numerical precision bugs.
These bugs are difficult to detect under normal circumstances but can silently alter gradients during specific stages of model training, causing sudden spikes in loss or even degrading capabilities the model has already learned.
Interestingly, the research team discovered that these bugs can be fixed simply by adjusting the computation method—without increasing floating-point precision or redesigning the loss function.
The related paper was submitted to arXiv on October 22. In addition to Jonathan T. Barron, the authors include research engineers and scientists Dan Gnanapragasam and Michael Tschannen.
- Paper Title: A Fix for Softmax and Cross-Entropy
- Paper Link: https://arxiv.org/pdf/2610.31642
According to Barron, Gemini handled a significant portion of the work in this study, ranging from problem identification and experiment execution to paper writing.
The More Confident the Model, The More Incorrect the Gradient Calculation
The first bug is called Softmax Collapse.
To understand it, we must first look at what Softmax does.
When training classification models, neural networks typically output raw scores, known as logits, for each candidate class. Softmax converts these scores into probabilities, and Cross-Entropy then measures the discrepancy between the model's predictions and the correct answers. For example, if an image contains a cat, the model assigns a much higher score to "cat" than to "dog" or "bird." After Softmax transformation, the predicted probability for "cat" might reach 99.99%.
As training progresses, the model becomes increasingly confident in the correct answer, and the gap between the logit of the correct class and those of other classes widens. This is where the problem arises.
Existing implementations typically subtract the maximum value from all logits before performing exponential calculations. This is a classic technique for ensuring numerical stability, preventing overflow during exponentiation. However, after subtracting the maximum, the exponential value corresponding to the maximum logit becomes 1, while values for other classes become very small positive numbers. When the correct class has the highest score, the denominator of the Softmax calculation effectively becomes:
1 + a very small number.
Floating-point numbers have limited precision. When this small number is sufficiently tiny, the computer rounds it off during addition, resulting in exactly 1. Consequently, the predicted probability for the correct class is calculated as precisely 1, the Cross-Entropy loss becomes 0, and the corresponding gradient is also computed as 0. However, gradients for other incorrect classes retain small non-zero values.
Mathematically, according to the definition of Cross-Entropy, the sum of gradients with respect to the logits for all classes should be zero. But now, the gradient for the correct class prematurely zeroes out, while gradients for other classes persist, breaking the relationship that should hold true. The paper refers to this point of sudden precision loss as the Precision Cliff.
Two numerical computation defects in standard Softmax and Cross-Entropy. (a) Softmax Collapse, (b) Softmax Cancellation Error.
The paper provides critical thresholds for different floating-point formats:
- bf16: Logit Margin ≈ 5.5
- fp32: Logit Margin ≈ 16.6
- fp64: Logit Margin ≈ 36.7
Here, Logit Margin measures the advantage of the correct class's logit over the combined scores of other classes. Taking fp32 as an example, when this margin exceeds approximately 16.6, the standard implementation encounters the Precision Cliff, erroneously zeroing out the gradient for the correct class. At this point, the theoretical predicted probability is already as high as approximately 99.9999938%.
The paper offers a very intuitive example. Assume three logits are [9, -9, -9], with the first class being the correct answer.
The cross-entropy loss computed by standard JAX and Optax implementations is 0, and the gradient for the correct class is also 0. However, with the corrected calculation method, the loss is actually approximately , the gradient for the correct class is approximately , and the gradients for the other two classes are approximately each. Although these numbers are small, they directly affect whether the gradients satisfy mathematical constraints.
This issue has not gone entirely unnoticed. In 2025, researchers such as Prieto et al. identified related phenomena and proposed StableMax to mitigate the problem by replacing the exponential function in Softmax.
Barron’s team chose a more direct approach this time: retaining the original definitions of Softmax and cross-entropy while reorganizing the computation process. They separated the term corresponding to the maximum logit from the sum of exponentials and used a numerically stable softplus function to complete the calculation. For cross-entropy with a single correct label, the corrected formula also allows automatic differentiation to directly yield that the gradient of the correct class equals the negative sum of the gradients of the other classes. This ensures that extremely small gradients are preserved until floating-point representation can no longer hold them.
The paper also provides a reference implementation in JAX, covering Softmax, Log-Softmax, and cross-entropy.
How a Small Precision Issue Triggers the "Slingshot" Mechanism in Training
The most interesting aspect of this bug is that it may explain a long-observed phenomenon in neural network training: the Slingshot Mechanism.
In 2022, Thilak et al. observed this phenomenon while studying Grokking. Grokking is often translated as "sudden insight": the model memorizes the training data early on but requires a significant amount of additional training before suddenly achieving good generalization performance on test data. During this process, researchers found that the training loss sometimes remains near zero for an extended period, then suddenly spikes, and subsequently drops back down. These fluctuations can occur repeatedly and are accompanied by jumps in test accuracy.
Why does the loss suddenly spike when the model has already memorized the training data? The paper offers a specific explanation.
Once the model memorizes the training set, the gap between the logits of the correct class and other classes continues to widen, gradually approaching the precision cliff. Once this critical threshold is crossed, standard cross-entropy incorrectly zeroes out the gradient for the correct class. At this point, adaptive optimizers like Adam are also affected. Adam dynamically adjusts the magnitude of parameter updates based on moving averages of historical gradients and their squares.
When training gradients remain close to zero for a long time, the optimizer's estimate of the second moment of the gradients also decays. Subsequently, parameter changes caused by weight decay or residual gradients may cause some samples to fall back within the precision cliff. The gradients for these samples, which were previously discarded, suddenly recover. Since the denominator used by Adam to normalize gradients has become very small, a normal gradient recovery can result in anomalously large parameter updates. Consequently, the training loss suddenly spikes.
This is precisely the formation mechanism of the Slingshot effect proposed in the paper.
The paper verified this in modular arithmetic tasks. In an experiment where a single-layer Transformer learned addition modulo 31, the research team used fp32 precision, the AdamW optimizer, a weight decay coefficient of 1, and , testing both standard cross-entropy and the corrected version.
The results were quite clear. Standard cross-entropy exhibited repeated loss spikes during training, with an average test accuracy of across 5 random experiments. The corrected cross-entropy resulted in smoother training, achieving 100% test accuracy in all 5 experiments. Under bf16 precision, the average test accuracy for standard cross-entropy in the same task was , whereas the corrected version achieved approximately 100%.
Another set of experiments revealed even more severe issues.
The research team trained an MLP to learn the sum of squares operation modulo 97, without using weight decay and setting Adam's to . Standard cross-entropy hit the fp32 precision cliff at around step 1250, causing test accuracy to plummet from 61% to 23%. The corrected version did not experience this accuracy collapse, and test accuracy continued to improve to approximately 85%.
The paper also observed similar phenomena in several other Grokking tasks. Under the corresponding experimental conditions, the corrected fp32 cross-entropy performed comparably to the standard fp64 implementation. However, the relationship between the Slingshot effect and generalization capability still needs to be understood in the context of specific training configurations.
Previous studies have observed a correlation between the Slingshot effect and the occurrence of Grokking. The new paper points out that in certain networks using LayerNorm without weight decay, the increase in parameter norm caused by the Slingshot effect may inadvertently act to reduce the effective learning rate. Therefore, the seemingly beneficial Slingshot effects observed in past experiments do not imply that this numerical issue itself is beneficial for the model.
Adam's Default Parameters Sit Right on the Edge of the Critical Point
The research team also explained another noteworthy detail. Adam optimizer includes a common hyperparameter , typically set to . It is added to the denominator of gradient normalization to control numerical behavior.
The paper finds that under specific conditions—where the model has memorized the training set and weight decay is applied—the growth of the Logit Margin is constrained by the balance between weight decay and Adam updates. This equilibrium value approximately satisfies:
m ≈ ln (1/ε)
When ε=10⁻⁸, this value is approximately 18.4. However, the precision cliff for fp32 lies exactly at 16.6. This means the model may trigger numerical issues in standard cross-entropy before reaching its naturally formed equilibrium point. This provides a new explanation for observations in some training experiments where "an excessively small ε in Adam can lead to instability."
Researchers found that in corresponding experimental setups, increasing ε to 10⁻⁷ or larger shifts the equilibrium point within the fp32 precision cliff, thereby avoiding these spikes. However, adjusting ε only alters optimizer behavior. Fixing the computation method of cross-entropy is required to directly eliminate this numerical error.
Student Equals Teacher, Yet Training Continues
The first bug occurs when the model becomes overly confident in its answers. The second bug, Softmax Cancellation, appears when the outputs of two models are very close.
A typical scenario is knowledge distillation. In knowledge distillation, the teacher model provides the target probability distribution, while the student model attempts to learn these outputs. If the logits of the student and teacher are identical, their corresponding probability distributions should also be identical. At this point, the student has achieved the distillation objective, and its gradient should mathematically be zero.
However, existing standard implementations may compute a non-zero gradient. The reason is that standard soft-target cross-entropy typically calculates the Softmax probabilities for the teacher and student separately, then derives the gradient by subtracting one from the other. When the two sets of probabilities are very close, floating-point operations lose the minute differences between them.
More troublesome is that the Softmax used in forward propagation and the Log-Softmax involved in backpropagation may use slightly different operation orders. Even if the logits of the teacher and student are completely identical, the rounding errors produced by the two computational paths may differ, ultimately leaving a residual gradient that should not exist. This residual gradient appears negligible but can be amplified by Adam.
To test its impact, the research team first trained a teacher model to achieve 100% test accuracy on a modulo-53 addition task, then initialized the student model by directly copying the teacher's weights. In other words, at the start of training, the student and teacher were identical.
The paper subsequently continued training the student for 2000 steps using a distillation temperature of 4, an Adam optimizer with ε=10⁻⁸, and no weight decay.
Theoretically, there was no need for the student to deviate from the teacher. However, experimental results showed that when using standard soft-target cross-entropy:
- With fp64, the residual gradient was extremely small, and the student remained largely unchanged.
- With fp32, a brief KL divergence spike occurred, indicating that the student's output temporarily deviated from the teacher.
- With bf16, the residual gradient was larger; the student clearly deviated from the teacher during continuous training, and test accuracy declined accordingly.
Researchers adopted a different fix for the second bug.
Since calculating two sets of probabilities separately and then subtracting them easily loses precision, they started from the difference between the student and teacher logits and directly computed the required probability difference in log-space. The specific implementation used expm1 and log1p, functions designed to handle small values near zero, and further corrected remaining floating-point errors.
This way, when the student and teacher logits are completely identical, the fixed gradient precisely returns to zero.
In the same distillation experiment, after adopting the fixed version, the student remained consistent with the teacher across all three precisions (bf16, fp32, fp64). The KL divergence stayed at 0, and test accuracy remained at 100%. These results indicate that in distillation scenarios where model outputs are very close, minor errors introduced by low-precision computations can indeed interfere with the optimization process.
Gemini Participated in Bug Discovery, Conducted Extensive Experiments, and Assisted with Paper Writing
Beyond the two numerical issues, Barron shared insights into the research process on social media that are worth noting.
According to his introduction, Softmax Cancellation was first discovered and fixed by Gemini Flash 3.7. Most of the experiments in the paper were conducted by Gemini Flash 3.8.
During the paper writing phase, the research team used Gemini 4 Argon.
Barron stated that Argon completed the initial drafts of the paper and figures. He then rigorously reviewed and repeatedly revised the drafts until the technical content was sufficiently solid, though he remained somewhat dissatisfied with the final textual expression.
In this research workflow, AI participated in multiple stages, including discovering numerical issues, designing fixes, executing experiments, and writing the paper. Researchers continued to bear the responsibility of technical verification and quality control.
Summary
Returning to the technical results, the value of this work needs to be understood within practical application contexts.
The paper primarily validated the fixes through Grokking tasks involving modular arithmetic and specific knowledge distillation experiments. The research team acknowledged that the extremely large Logit Margins observed in the experiments, as well as cases where student and teacher outputs were nearly identical, are not common in conventional large-scale training.
Therefore, it cannot yet be concluded that mainstream large model training is severely affected by these two issues, nor can it be directly inferred that fixing them will accelerate overall large model training.
However, the study does reveal an often-overlooked fact: even if mathematical expressions are completely correct, the floating-point operations actually executed by computers may change the nature of gradients under specific conditions.
Softmax and cross-entropy are foundational components of the deep learning toolchain. For years, researchers have optimized model architectures, training algorithms, and computational efficiency. This discovery indicates that even these most basic computational implementations still have room for improvement.
The proposed fix preserves the original mathematical definition of the functions and does not require increasing floating-point precision to resolve the issue. The authors believe these functions can directly replace existing implementations without introducing significant performance overhead. However, their actual effectiveness in large-scale model training still requires further validation.
References: