DeepMind's Barron flags two Softmax bugs in standard ML code

The paper targets a piece of numerical plumbing that most practitioners never inspect: the handoff between Softmax and Cross-Entropy, where the gradient simplifies to "probability minus the correct-class indicator." The authors show that the standard trick of subtracting the maximum logit before exponentiating, while numerically safe on its own, creates a path for rounding to erase the correct-class gradient once confidence passes a precision-dependent threshold. Their fix avoids ever forming "1 plus a tiny number" by computing the log-sum-exp only over non-correct logits and wrapping the loss in softplus.

The Mac reproduction adds independent evidence that the bugs are not theoretical. In PyTorch on CPU and Apple GPU (MPS), the standard cross-entropy produced a gradient sum of +3.0e-08 instead of zero for a confident example, while the fixed version returned exactly zero. In a distillation setup where student and teacher started with identical weights, the standard implementation drifted apart within three steps even in fp64, whereas the fixed version stayed put.

The practical takeaway hinges on where the loss is computed. Because Hugging Face transformers casts logits to fp32 before the loss, the reproduction found only a small share of tokens beyond the cliff in pretrained LLMs; but when bf16 logits go straight into the loss, up to 47% of tokens for an 8B model on code data sat beyond the cliff, meaning their correct-class gradients rounded to zero. That distinction, rather than the existence of the bugs alone, is likely to determine whether a given training pipeline is exposed.