A seemingly basic mathematical computation issue may affect the stability of large‑scale model training. Google DeepMind researchers have found that long‑used implementations of Softmax and cross‑entropy can produce incorrect gradients under certain floating‑point precision conditions. Notably, Gemini not only helped identify the problem but also participated in experimental validation and paper writing.
In late October 2026, Google DeepMind researchers Jonathan T. Barron, Dan Gnanapragasam and Michael Tschannen posted a paper on arXiv titled "Fixing Softmax and Cross‑Entropy". The paper points out two subtle numerical precision issues in common implementations of two almost ubiquitous operations in deep learning: Softmax and cross‑entropy.
These issues are not errors in the mathematical formulas themselves but arise because mathematically equivalent expressions can yield different results in limited‑precision floating‑point arithmetic. They may be hard to detect during ordinary training, but when model predictions become extremely confident or when a student model closely matches a teacher model, tiny rounding errors can change gradients and thus affect optimization.
1. Why do two fundamental functions show numerical flaws?
To understand the study, first recall the roles of Softmax and cross‑entropy in neural network training.
Softmax converts model output logits into a probability distribution, and cross‑entropy measures the discrepancy between the predicted distribution and the target. Together they form one of the most common loss computation paths in classification models and large language model training.
In theory, Softmax is computed as:
where Zi are the model logits and Pi are the predicted probabilities for each class. To avoid overflow from exponentials, engineering implementations commonly subtract the maximum logit before exponentiating.
Mathematically equivalent and usually numerically more stable, this trick nevertheless can still lead to two problems in finite‑precision floating‑point arithmetic.
1. Softmax Collapse: the correct class gradient unexpectedly becomes zero
When training progresses, the model's prediction for the correct class may become extremely confident, with its probability approaching 1 while other classes approach 0.
Floating‑point representations cannot represent values arbitrarily close to 1 with infinite precision. When the difference from 1 is smaller than the representable precision of the data type, that difference can be rounded away.
The paper shows that under certain computation conditions, both bfloat16 and float32 can exhibit this phenomenon once logit differences reach a particular range.
The result can be that the gradient for the correct class is computed as exactly 0 while gradients for incorrect classes remain nonzero. This violates mathematical properties the Softmax gradient should satisfy and can deliver inconsistent update signals to the optimizer.
Such anomalies may relate to observed sudden loss spikes and optimization instability during training. The authors discuss a so‑called Slingshot Effect in this context.
It is important to emphasize this does not mean every model using bfloat16 or float32 will encounter the problem. Occurrence depends on specific logit distributions, computation paths, and numeric precision.
2. Softmax Cancellation: the model is nearly correct yet still produces spurious gradients
The second issue appears mainly in knowledge distillation.
Knowledge distillation trains a smaller student model to match the teacher model's output probabilities rather than only the final hard labels.
Ideally, when the student output matches the teacher exactly, the loss should be minimal and gradients should be zero.
However, the researchers found that standard computation can produce nonzero gradients due to floating‑point rounding errors in some cases.
This means even when the student closely matches the teacher, the optimizer may receive an incorrect update direction that pushes the student away from the teacher.
This type of error is especially concerning for systems that rely on distillation for model compression, capability transfer, or post‑training adjustments.
2. Root cause: mathematical equivalence does not guarantee identical results in computer arithmetic
The most valuable contribution of this study is its reminder of a fact often overlooked in engineering practice: correct mathematical formulas do not guarantee correct numerical results in finite‑precision computing environments.
Floating‑point numbers use a finite number of bits to represent real numbers and therefore cannot represent all values exactly. Especially near 1.0, two mathematically distinct values may be rounded to the same representable number in low‑precision formats.
This leads to two consequences:
First, very small but real gradients may be rounded to zero.
Second, two values that should cancel each other may produce a residual error depending on rounding order.
This explains why identical mathematical formulas can produce different training outcomes under different data types, computation orders, and hardware environments.
As large language models increasingly use bfloat16, FP16, and sometimes FP8, the importance of these issues rises.
3. Gemini did more than help write code — it helped find the problem
Another notable aspect of this study is Gemini's role in the research process.
According to public statements by paper author Jonathan T. Barron:
- Problem discovery: Softmax Cancellation was initially found by Gemini Flash 3.7, which also proposed a fix.
- Experimental validation: Most experiments in the paper were carried out by Gemini Flash 3.8.
- Paper drafting: Gemini 4 Argon contributed to writing the initial draft and producing figures.
- Research oversight: Barron primarily reviewed, questioned and validated the research conclusions to ensure they withstand rigorous scrutiny.
If these attributions are recorded accurately, this is a representative case where AI not only helps people use machine learning but begins to participate in checking and improving machine learning's low‑level computational methods.
It is important to note Gemini did not replace researchers' mathematical judgment and experimental verification. The meaningful outcome is a new collaborative research workflow: a model proposes issues, explores fixes and runs experiments, and human researchers review and confirm the results.
4. Fixes: changing computation, not increasing precision
The team did not opt to simply increase floating‑point precision, nor did they redefine Softmax or cross‑entropy mathematically; instead they adjusted computation expressions and operation order to reduce numerical error.
The paper uses numerically stable functions such as softplus, expm1, and log1p.
These functions are optimized for specific numeric ranges and avoid computations that easily lose precision.
For example:
- expm1(x) computes e^x ‑ 1 and is more accurate than computing e^x and subtracting 1 when x is small.
- log1p(x) computes log(1+x) and reduces loss of significant digits when x is small.
- softplus(x) provides a more stable form for computations related to Softmax.
The team validated fixes on tasks such as modular addition and Grokking testbeds. Under certain bfloat16 test conditions reported in the paper, the standard implementation achieved about 93.3% accuracy while the fixed implementation reached 100% and produced more stable training.
In knowledge distillation experiments, the corrected computations avoid spurious residual gradients when the student matches the teacher.
Independent engineers have also reproduced the issues in PyTorch and validated the fixes.
5. What this means for large language models
Softmax and cross‑entropy are not specific to any single model; they are fundamental computational primitives in modern deep learning systems. In large language model training, Softmax appears not only in classification tasks but directly in computing output probability distributions.
Thus the study's significance goes beyond fixing two functions: it signals to the AI engineering community that model improvements depend not only on architecture, data, and scale but also on the reliability of low‑level numerical computations.
However, the scope of impact should be judged objectively. Everyday model inference will not usually be directly affected by these two issues; the primary concern lies in large‑scale pretraining, low‑precision training, knowledge distillation and certain post‑training tasks. Whether mainstream frameworks will adopt the paper's fixes as default implementations remains to be seen in updates to PyTorch, JAX, TensorFlow and others.
6. More important than fixing two bugs: AI begins to improve AI
I believe the most noteworthy outcome is not merely that Softmax and cross‑entropy have numerical flaws, but that the AI R&D model itself may be changing.
Historically, deep learning development relied on human researchers to propose theory, write code and design experiments, with models acting as tools or subjects.
Now, advanced models can participate in exploring mathematical problems, analyzing code, designing experiments, comparing results, and even discovering low‑level issues that researchers have overlooked.
This creates a new closed loop:
AI models discover problems → AI suggests fixes → automated experiments validate → human researchers review → training infrastructure is improved → AI capabilities advance further.
This does not mean AI can yet conduct fully reliable scientific research independently, but it shows AI is shifting from a mere tool to an active collaborator and important assistant in research and engineering.
If in future AI can continuously identify issues in algorithms, numeric computations and hardware utilization, its contribution to its own development could outpace mere increases in parameter counts or training data.
From this perspective, Gemini's role in fixing Softmax and cross‑entropy is not just about patching numerical computations but is a notable sign that AI is beginning to help improve its own technical foundation.
▶️ Enjoyed this article? Creating great content takes effort!
Subscribe to our YouTube channel:
AI Rosetta
Your likes, subscriptions and shares mean the world to us! 🙏
For more AI news, visit:
https://www.airosetta.com