Google DeepMind Researcher Discloses Rounding Bug in Softmax and Cross-Entropy Implementations

Jon Barron, a principal research scientist at Google DeepMind, said on X that standard implementations of Softmax and cross-entropy in PyTorch, JAX, and textbooks contain a rounding bug that can zero out gradients and destabilize training. He said the issue was originally discovered and fixed by Gemini Flash 3.7, and the team has submitted a paper titled "Fixing Softmax and Cross-Entropy" to arXiv.

Softmax converts raw logits into probabilities, while cross-entropy measures the difference between predictions and targets. To avoid overflow, standard implementations subtract the maximum logit, creating a denominator shaped like "1 + a very small number." In floating-point arithmetic, if that small number falls below the precision threshold, it is rounded to zero, causing the gradient for the correct answer to vanish. The DeepMind paper calls this "Softmax collapse" and notes that optimizers such as Adam can amplify the issue, producing sudden loss spikes that were previously described as "slingshots" during grokking.

The paper also identifies "Softmax cancellation" in knowledge distillation, where rounding leaves a spurious gradient that pushes student models away from the correct answer. The proposed fix avoids adding a very small number to 1, computes cross-entropy via log-sum-exp and softplus, and uses expm1 and log1p in log space for distillation.

In tests on a modular addition task (a + b mod 31), the standard implementation reached only 93.3% accuracy under bfloat16 and was unstable, while the fixed version reached 100.0%. Barron added that Gemini Flash 3.7 initially found and fixed the bug, most of the experiments in the paper were done by Flash 3.8, and Gemini 4 Argon drafted the paper and figures. On October 24, independent engineer Satyajit Ghana reproduced both bugs in PyTorch and confirmed the paper's fp32 results.