A New Fix for a Training Bug Is Real. The Accuracy Gain Mostly Isn't.
A recent paper found a rounding bug inside the loss function that nearly every classifier and language model trains with, and proposed a fix that it said makes models learn better. We rebuilt the fix in PyTorch and re-ran the experiments: 480 training runs plus a distillation test, with our predictions written down before we started. The bug is exactly as described. The accuracy gain was much smaller than reported, and at some settings it went the other way. Along the way we found three things that matter more for everyday training.
01 — the bug
When a model is very sure, its learning signal rounds to zero
Training nudges the model towards the right answer by an amount that shrinks as the model grows more confident. Call that confidence the margin: how far the right answer's score is ahead of the rest. The standard loss works the nudge out as "1 minus a probability". Once the probability is so close to 1 that the computer stores it as exactly 1, the nudge becomes exactly zero, while the push away from the wrong answers carries on.
Where that happens depends on the number format. We swept the margin and found the cutoffs close to where the paper put them: 5.75 in 16-bit (bf16), 16 in 32-bit and 36 in 64-bit. The fixed loss does the same calculation in a way that never forms "1 minus a probability", so its signal stays correct all the way out.
02 — what we ran
480 small training runs, predictions written first
We followed the paper's main experiment: a small transformer learning to add numbers on a clock with 31 positions, trained on half the possible sums and tested on the other half. The paper didn't describe its model, so we used a well-known one-layer design from earlier work on this task. We compared three versions of the loss:
Standard, PyTorch's built-in cross-entropy. Fixed, a line-by-line port of the paper's code. Upcast, our addition: the standard loss with one line that converts the model's scores to 64-bit first, to see whether the cheap workaround is enough.
Each version ran in 16-, 32- and 64-bit, at three values of the Adam optimiser's epsilon setting (the tiny number it adds to avoid dividing by zero), from 20 random starting points. Every version began from identical weights on identical data for a given starting point, so each comparison is like for like. Sixty-seven unit tests check the losses against a high-precision reference. The whole grid took 3.5 hours on a DGX Spark and a Mac.
Before any full run we wrote down what we expected, in two tiers: how the numbers should behave, which ought to hold for any model, and how accuracy should move, which might not. We report every prediction, including the ones that failed.
03 — what held up
In 32- and 64-bit, the numbers behaved as predicted
In 32-bit, the standard loss's margin climbed to the cutoff, stalled, slid back and climbed again, over and over. At the top of each cycle up to 74% of training examples were getting exactly zero signal for the right answer. Under the fixed loss that share was zero in every run.
The paper also predicts where the fixed loss settles: at a margin set by Adam's epsilon, the same in every format. That held too. Fixed margins came out at about 17, 26 and 36 for the three epsilon settings in both 32- and 64-bit, about one unit below the predicted value each time. The standard loss in 32-bit was capped near 11 to 12 by its cutoff.
The one-line upcast matched the fixed loss wherever the 64-bit cutoff wasn't reached, which in our grid meant every setting except the smallest epsilon. The 16-bit runs are a separate story (section 05).
04 — what didn't
The accuracy gain was small, and it reversed
At the default epsilon, the fixed loss beat the standard one by 1.5 points of test accuracy, a real but small gain, against 5.7 in the paper. At the next setting it made no clear difference. At the smallest it cost points: 5.9 behind at the end of training (an interval that only just includes zero), and 10 behind averaged over the last 6,000 steps. In 64-bit, where the cutoff is far away and the two losses should agree, they did.
Our reading, which we haven't tested directly: the bug works as an accidental cap on how confident the model can get, and in this task that cap helped it generalise. Removing it let margins grow larger, and larger margins went with later or no learning of the rule. Our model learned about twice as slowly as the paper's, so the gap may come down to model design.
05 — the 16-bit surprises
"16-bit" can mean three different things
16-bit weights froze training. In our first 16-bit runs, the weights and optimiser state were stored in 16-bit too. None of the 20 runs learned the task at any setting. The cause is plain arithmetic: AdamW shrinks each weight by multiplying it by 0.999 every step, and that change is smaller than the gap between neighbouring 16-bit numbers, so the result rounds back to the same weight. Most of Adam's small updates vanish the same way. We set those runs aside, as our written rules required.
16-bit only at the loss reproduced the paper. With 32-bit weights and only the model's scores converted to 16-bit before the loss, the standard loss stalled at a margin of 4.9 (the paper: 5.1) and the fix gained 2.2 points on average across five runs.
16-bit on an NVIDIA GPU didn't collapse early at all. Under PyTorch's usual mixed-precision mode, the standard loss settled at a margin of 15.3, the 32-bit value. The scores reaching the loss were 16-bit; we checked. But PyTorch's GPU cross-entropy does its internal arithmetic in 32-bit and only rounds the outputs, so the cutoff sits at 16, not 5.75. Our loss-only runs were on a CPU, where the 16-bit maths rounds inside the calculation. The same "16-bit" setting moves the cutoff by a factor of three depending on where the arithmetic runs.
06 — distillation
A copy of the teacher drifted away, even in 64-bit
The paper's second bug affects distillation, where a student model learns to copy a teacher's outputs. If the student already matches the teacher, the signal should be exactly zero. The standard loss computes the two sides separately and subtracts them, so a tiny leftover remains. We trained three teachers, made each student an exact copy, and ran 2,000 steps of Adam.
One setup trap first: in PyTorch, a model in evaluation mode can take a faster attention path whose outputs differ in the fifth decimal place from training mode. That was enough to make the "exact copy" not exact. We kept both models in training mode, after which their outputs matched to the last bit.
The fixed loss gave an exactly zero signal and the student never moved. The standard loss drifted in every format, including 64-bit, where the paper expected the leftover (around 10⁻¹⁷) to be too small to matter.
The reason is the optimiser, not the size of the leftover. When the signal is far below Adam's epsilon, Adam behaves like plain gradient descent with a very large step, and at the default settings that step is too large to stay at the teacher. Any nonzero signal, however small, gets multiplied every step until the student breaks away; a fully 64-bit student escaped within 20 steps. We measured how sharply curved the loss was at each teacher, predicted the epsilon at which the drift should stop (1.46 × 10⁻⁴ for the first teacher), and checked it. The real boundary fell between 1.2 and 1.8 × 10⁻⁴. The fix works because its signal is exactly zero, not because it is small.
Accuracy never dropped in any distillation run, so we did not reproduce the paper's accuracy loss. From a more realistic start, a student near but not equal to its teacher, all three losses reached similar error in 32- and 64-bit; in 16-bit the fix reached about half the error of the standard loss.
07 — what this does not show
Limits we want to be clear about
One task and one model design. The paper's model wasn't specified, and ours learned the task more slowly and less reliably. The findings about the numbers hold for any model; the accuracy findings may not.
Some parts are small. The 16-bit follow-ups used 5 runs each and the distillation tests 3 teachers. They show direction and cause, not exact sizes. The 64-bit drift explanation and the epsilon sweep were worked out after we saw the drift, though the near-threshold prediction was written down before its runs.
08 — practical lessons
Check where your precision actually applies
The bug is real, but in our tests, where the 16-bit arithmetic happened mattered more than which loss function we used.
Keep 32-bit master weights. With weights stored in 16-bit, AdamW's weight decay silently rounds away and training can stall. Standard mixed-precision setups already keep a 32-bit copy; hand-rolled 16-bit training may not.
Know where your loss is computed. On an NVIDIA GPU, PyTorch's cross-entropy already does its internals in 32-bit, so the early 16-bit collapse doesn't happen there. On a CPU it did, and other hardware may differ.
If in doubt, convert the scores before the loss. logits.float() or .double() is one line, and in our grid it matched the fix everywhere but the smallest epsilon.
Don't count on an accuracy gain. The collapse only bites on examples the model already gets right with near-total confidence. In language-model fine-tuning those tokens usually add little either way; it matters most when you train many passes to near-zero loss. Measure on your own task.
Near zero isn't zero for Adam. When fine-tuning or distilling from a model that is already at a good solution, a raised epsilon, a decaying learning rate or a loss that returns an exact zero keeps it there. A tiny leftover signal won't.