# Two hidden bugs in softmax, the maths at the heart of every language model

URL: https://ai-pulse.aglaganova.com/blog/two-hidden-bugs-in-softmax-the-maths-at-the-heart-of-every-language-model
Published: 2026-10-26
Author: Roman Aglaganov, Full-stack developer
Category: Research
Region: Global
Tags: science

> A new paper says the standard code for softmax and cross-entropy, used to train almost every neural network, quietly gets gradients wrong in two situations. The authors offer a short fix and say it explains some odd training behaviour.

![Illustration: a smooth black curve rising across a grid of grey dots, with one red dot knocked off the curve at its peak](https://ai-pulse.aglaganova.com/api/media/file/softmax-fix-1.png)

Almost every modern AI model ends with the same few lines of maths. Softmax turns a model’s raw scores, called logits, into probabilities, and cross-entropy measures how far those probabilities are from the right answer. Training follows the gradient of that number. A paper posted to arXiv on 22 October by Jonathan T. Barron, Dan Gnanapragasam and Michael Tschannen argues that the standard implementations of these functions have two numerical bugs.

## The two bugs

**Softmax Collapse.** When one logit is much larger than all the others, the standard code rounds that logit’s gradient to zero. The gradients then no longer add up to zero, as the maths says they should, so the model gets a slightly wrong signal exactly when it is most confident.

**Softmax Cancellation.** In distillation, where a smaller “student” model learns to copy a larger “teacher”, the student’s logits end up very close to the teacher’s. There the standard code produces gradients that are not zero even when the two sets of logits are equal, so the student keeps being pushed when it should be left alone.

Both are problems of finite precision: computers store numbers with a limited number of digits, and the usual formulas lose the small differences that matter in these cases. The authors present rewritten versions of log-softmax, softmax and cross-entropy that avoid both.

## What the fix changes

According to the paper, the fix removes the “Slingshot” effect, a known pattern of sudden spikes during long training runs; speeds up grokking, where a model suddenly generalises long after it has memorised its training data; prevents accuracy collapsing when no weight decay is used; and stops distilled students drifting away from their teachers. The authors also report that, after memorisation, the growth of the logits is limited by a balance between weight decay and a small constant in the Adam optimiser.

## Why it matters

This is the most heavily used code in machine learning, written into every major framework, which is why a claim like this is worth watching. A feature request asking PyTorch to adopt the approach has already been opened on its GitHub. The results are the authors’ own and the paper has not yet been peer reviewed, so the next step is for others to test whether the fix changes anything at the scale of real language models.

## Sources

- [arXiv: A Fix for Softmax and Cross-Entropy (Barron, Gnanapragasam, Tschannen)](https://arxiv.org/abs/2610.31642)
- [GitHub: PyTorch issue #201986, “A Fix for Softmax and Cross-Entropy”](https://github.com/pytorch/pytorch/issues/201986)