Google DeepMind researchers find and fix two bugs in softmax and cross-entropy

Three Google DeepMind researchers, Jonathan T. Barron, Dan Gnanapragasam and Michael Tschannen, have published a paper describing two numerical bugs in the standard softmax, log-softmax and cross-entropy implementations used across deep learning frameworks like JAX, PyTorch and TensorFlow. The first bug, which they call Softmax Collapse, rounds the gradient on a correctly predicted target to exactly zero once the model becomes confident enough, which breaks the mathematical property that all gradients should sum to zero. The second bug, Softmax Cancellation, shows up in knowledge distillation, where a student model's logits are nearly equal to a teacher's, and produces a non-zero gradient even when the two sets of logits are identical.

The researchers trace both bugs to how floating-point numbers lose precision near 1.0 and to the order of operations used when subtracting the maximum logit before exponentiating. They present an alternative formulation, including a short JAX implementation, that keeps gradients accurate down to the point where floating-point numbers underflow to zero, rather than rounding them away earlier.

According to the paper, the fix eliminates a known "slingshot" phenomenon of sudden loss spikes during training, speeds up a behavior called grokking when weight decay is used, prevents accuracy from collapsing when it isn't, and stops student models from drifting away from their teacher during distillation.

Why it matters

A fix at this level sits underneath nearly every classifier and language model, so correcting it could quietly improve training stability across the field.

Source: arXiv