99.999997% Confident, and the Gradient Coming Back from Backprop Is Zero

The model is 99.999997% sure of one answer. The loss is near zero. Backpropagated, the gradient for the target logit is 0.0.

This is what standard softmax plus cross-entropy gives under fp32. Jonathan T. Barron of Google DeepMind and two colleagues wrote a paper with a blunt title: "A Fix for Softmax and Cross-Entropy." The conclusion is equally blunt: the standard implementations in JAX, PyTorch, TensorFlow, and textbooks all have two numerical bugs.

Jia-Bin Huang made a half-hour explainer video, going from the softmax derivation to floating point; if you don't want to read the paper, you can watch this first.

The First Bug: Softmax Collapse

The standard implementation first subtracts the maximum to prevent overflow, something Higham put in textbooks back in 2002:

z_k - z_max

When the target happens to be the largest logit, after the shift it becomes 0, and the denominator becomes 1 + sum_{j≠t} exp(z_j - z_t), i.e. 1 + e^{-m}. m is the logit margin, how much higher the target class is than all other classes.

The problem is this 1 +. fp32 has roughly 1e-7 precision near 1.0. Once e^{-m} is smaller than that magnitude, adding it is the same as not adding it, and the denominator is 1.0. q_t = 1/(1+e^{-m}) also becomes 1.0. The gradient g_t = q_t - 1 = 0.

The gradients for non-target classes are computed near 0, where precision is much finer and can go all the way down to exponential underflow. The target gradient is zeroed out, while the others remain. The sum of gradients is no longer zero.

The paper calls the critical point the precision cliff: 5.5 for bf16, 16.6 for fp32, and 36.7 for fp64. Once the margin crosses this line, the gradient falls off the cliff.

The fix is to avoid the 1 + small quantity addition. Keep e^{-m} in the numerator:

g_t = -e^{-m} / (1 + e^{-m})

Use softplus and logsumexp to exclude the argmax from the sum. Barron gives a JAX implementation, twenty-nine lines. The second bug is handled similarly, using expm1 and log1p to avoid directly subtracting two nearly equal probabilities.

Between the Two Bugs Lies a Slingshot

When Thilak et al. studied grokking in 2022, they saw a phenomenon: after the model memorizes the training set, the loss bounces back and forth between a plateau near zero and sudden spikes, with the spikes coinciding with jumps in test accuracy. They called this a slingshot.

This spike is a direct consequence of Softmax Collapse. AdamW's second moment tracks the square of the gradient. If all samples in a batch have margins past the precision cliff, the target gradients all go to zero, the parameter gradients nearly vanish, and the second moment decays exponentially to a tiny value. Weight decay slowly pulls the parameters back; a few samples fall back below the line, the gradient suddenly recovers, and dividing by a denominator that has already decayed to a tiny value makes the parameters take a big step, causing a loss spike. Then it repeats.

With the fixed version, g_t does not go to zero at the cliff; the first and second moments decay smoothly together, the spikes disappear, and grokking gets faster.

The Default eps Sits Just Outside the Cliff

There is one detail in this paper worth calling out separately. AdamW's weight decay and Adam's eps form an equilibrium point. The gradient decays exponentially with the margin; once eps is larger than the gradient, the normalized step size starts shrinking as g/eps. The equilibrium works out to:

m* ≈ log(1/eps)

With the default eps = 1e-8, log(1/eps) ≈ 18.4. The fp32 cliff is 16.6. The default configuration pushes the margin just past the cliff. This may be why some people find small eps dangerous.

The paper also tried eps = 1e-16. The equilibrium rises to 36.8, a bit higher than the fp64 cliff at 36.7. At this point even the standard implementation cannot hold up under fp64.

The Second Bug: Softmax Cancellation

In knowledge distillation, the student network's logits approach the teacher's. When z_S equals z_T, q - p should be exactly zero. But the forward pass uses softmax and the backward pass uses log-softmax, so the two paths round in different orders, and the residual after subtraction is not zero. In fp32 this residual is about 1e-8; in bf16 it is 1e-3.

As soon as Adam sees that the residual is larger than eps, it normalizes it into an update of normal step size. The student thus drifts away from the teacher. In Barron's experiments, when the student is initialized with the teacher's weights, test accuracy drops noticeably under bf16, and under fp32 the KL first spikes and then settles onto a floor of 1e-6.

The fixed version returns an exact zero gradient when z_S = z_T. KL stays 0 throughout, in all three precisions.

The Ugliest Set of Numbers in the Experiments

a+b mod 31, a one-layer Transformer, weight decay 1. Standard CE's test accuracy wobbles around 94%, while the fixed version is between 99% and 100%. The gap is not large, but the trend is visible.

a²+b² mod 97, no weight decay, eps 1e-16, and the gap becomes obvious. Standard CE runs to around step 1250, the margin hits the fp32 cliff, and one spike knocks test accuracy from 61% down to 23%. The fixed version goes all the way to 85%.

The paper also compares two existing remedies. Liu et al.'s 2026 Zero-sum Proj subtracts the mean from the gradient to ensure the sum is zero, but once the target gradient has already been zeroed out, subtracting the mean can only restore an amount of -1/K; the margin gets stuck around 18, and when eps reaches 1e-12 accuracy drops to 49%. Prieto et al.'s 2025 StableMax replaces exp with a piecewise rational function; the margin shoots up to the 10^4 range, but accuracy still lags.

A Side Note

In a tweet, Barron said this bug was first found and fixed by Gemini Flash 3.7, most of the paper's experiments were run with Flash 3.8, and the first draft was written by Gemini 4 Argon. He nitpicked until the content was fine but the writing was not, then threw his own first-author papers from previous years into the context window and let the model learn his style; the effect was unexpectedly good, and finally he revised it by hand. He said he really likes this workflow.

Whether that counts as self-deprecation is up to interpretation.

Why It's Worth Reading

The extreme margins in the paper are uncommon in routine large-scale training. But the fix is a drop-in replacement: no internal conversion to higher precision, and no obvious slowdown. Any model making high-confidence predictions, and any distillation task where student and teacher are very close, will encounter these two bugs.

Back to the number from the beginning. 99.999997% confidence, and the gradient is zero. 1 + a very small number = 1; this floating-point rule has lain in textbooks for twenty years.