Google DeepMind finds softmax flaw: gradient drops to zero at a logit gap of 16.6

On 22 October 2026, Jonathan Barron, Dan Gnanapragasam and Michael Tschannen of Google DeepMind posted the preprint A Fix for Softmax and Cross-Entropy. The authors analyze two numerical flaws in the standard softmax and cross-entropy implementations used by PyTorch and JAX. The first, Softmax Collapse, zeroes the gradient of the correct class in float32 when its lead over the competitors exceeds roughly 16.6. The second, Softmax Cancellation, adds a nonzero gradient in distillation even when the student has exactly reproduced the teacher. Both flaws are fixed by rearranging the terms of the formula, without changing the math.

Where float32 loses a small number

The standard scheme is already protected against overflow: the maximum is subtracted from the logits, so the exponential does not blow up. But overflow protection does not guarantee accuracy for small differences. Take the logits [20, 0] in float32. The probability of the second class comes out at 2.06·10⁻⁹, and that number is stored perfectly well. The probability of the first class rounds to exactly 1, because an increment of 10⁻⁹ added to one does not fit into seven significant digits. After one is subtracted, the gradient of the correct class comes out as zero, although it should be about −2.06·10⁻⁹. The competitors' gradients remain, so the derivatives no longer sum to zero.

Let m be the correct answer's lead and r = e^(−m). The usual formula computes 1/(1+r) − 1. The corrected one computes −r/(1+r). It is the same function, but the small quantity sits in the numerator and is not rounded away together with the one. Check in float32: at m = 17 the usual scheme gives 0, while the corrected one gives −4.14·10⁻⁸. At m = 20 the result is −2.06·10⁻⁹, and at m = 30 it is −9.36·10⁻¹⁴.

The thresholds depend on the number of fraction bits in the computation format:

• bfloat16, 7 bits: lead threshold ≈ 5.54;
• float16, 10 bits: ≈ 7.62;
• float32, 23 bits: ≈ 16.63;
• float64, 52 bits: ≈ 36.74.

These are thresholds of the arithmetic of a specific operation. Having weights in bfloat16 does not mean that softmax and the loss are also computed in bfloat16: under mixed precision, sensitive operations are often run in float32.

The student matched the teacher, but the gradient remained

The second flaw shows up in distillation. The student learns to reproduce the probabilities of a fixed teacher, and when they match perfectly, the gradient with respect to the logits should be zero. Cross-entropy then stays equal to the teacher's entropy, while the gradient and the KL divergence go to zero.

In practice, the teacher's and student's probabilities are often computed by different routes: one through a division of exponentials, the other through the exponential of log_softmax. For the logits [2, 1, 0] in float32, the difference is about 1.49·10⁻⁸ in two components, with a noise norm of 2.1·10⁻⁸. The logits are identical, but the gradient is not.

The authors propose comparing the logits first and moving to probabilities afterward. The method computes the difference between the student and teacher logits, subtracts from it the mean weighted by the teacher's probabilities, and then evaluates the small correction with expm1 and log1p. These functions preserve accuracy for small arguments. When the logits match, the result is a strict zero: expm1(0) returns 0.

Why 10⁻⁸ reaches training

A gradient on the order of 10⁻⁹ looks harmless, but Adam normalizes its step by an estimate of the gradient scale, so the absolute smallness of the signal does not protect against the consequences. In Softmax Collapse, the gradient of the target logit vanished while the competitors' gradients remained. The distorted signal passes through the network and reaches the adaptive optimizer. The authors link this asymmetry to the Slingshot effect: after a period of low loss, sharp spikes appear.

The experiments in the preprint were run in float32, with five runs using different seeds. The first is addition modulo 31 with weight decay λ = 1 and Adam ε = 10⁻⁸. The second is (a² + b²) mod 97 with λ = 0 and ε = 10⁻¹⁶. In both cases, fixing the numerical scheme removes the spikes and changes the subsequent training dynamics. This is a result on specific models and tasks. It does not prove that every loss jump in an LLM is caused by softmax.

To compute the hard-label loss stably, the authors use logsumexp over the competitors and softplus instead of directly adding one. The illustrative NumPy example in the write-up works only when the correct class is in the lead, and it does not replace autograd.

What to check in your own code

The fix changes the order of operations, not the architecture: it requires no new parameters and no extra memory. For teams training models to low loss, this is a cheap check.

• Feed the loss kernel logits with a lead of 17–30 in float32 and see whether the gradient of the correct class goes to zero.
• Make sure that in your mixed-precision setup, softmax and the loss are actually computed in float32, not in bfloat16 or float16.
• In distillation, feed the student the teacher's logits. The gradient should be exactly zero, not on the order of 10⁻⁸.
• Compare behavior at different values of ε in Adam: this parameter determines how much of the noise turns into a step.
• If you see loss spikes after a long plateau in training, check the numerical scheme before changing hyperparameters.

Take the full implementation for large logit differences from the authors' preprint: the short demonstrations in the write-up suit small differences and do not replace a ready-made function.