A fix for softmax and cross-entropy: the target gradient that rounds to zero

A Fix for Softmax and Cross-Entropy (Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen, Google DeepMind, arXiv 22 October 2026) makes a claim that sounds like it cannot be true: the softmax and cross-entropy in JAX, PyTorch and the textbooks are wrong in two places. Not unstable in some exotic regime. Wrong, in the sense that they return a different gradient from the one the formula defines, in a regime that ordinary training reaches.

Barron announced it on X on 23 October: "all existing implementations of softmax and cross-entropy (that's JAX, PyTorch, textbooks, etc) have a subtle bug --- two bugs, actually. A simple rewrite fixes both, and eliminates 'slingshots' when grokking." The HTML version of the paper carries the code, and there is an interactive demo.

I wanted to see the numbers myself, so this page has three kinds of evidence, each labelled. Measured: I ran it, in PyTorch 2.14.1 (CPU build) and NumPy 2.5.4 on a shared 4-core cloud VM, with no GPU. Reported: the paper's figure, not re-run. Reasoned: my arithmetic on the other two.

Softmax as everyone writes it

Take logits $\mathbf{z} \in \mathbb{R}^K$. Computing $e^{z_k} / \sum_j e^{z_j}$ directly overflows once a logit passes about 88 in fp32, so every framework first subtracts the largest logit:

Cross-entropy against a target distribution $\mathbf{p}$ is $\mathcal{L} = -\sum_k p_k \log\operatorname{softmax}(\mathbf{z})_k$, and its gradient is the cleanest result in deep learning: $\mathbf{g} = \mathbf{q} - \mathbf{p}$. $\mathbf{q}$ and $\mathbf{p}$ each sum to one, so the gradient sums to zero. Raising one logit has to lower the others. Keep that invariant in mind; both bugs break it.

Bug one: Softmax Collapse

Call the target class $t$ and suppose the model is right and confident, so $z_t = z_{\max}$. The paper measures confidence with the logit margin, how far the target logit sits above the log-sum-exp of all the others:

After the max shift, the target's own term in the denominator is $e^0 = 1$ and the other $K-1$ terms add up to exactly $e^{-m}$. So the denominator is $1 + e^{-m}$, and the hard-target loss and gradient are

The two forms of $g_t$ are equal on paper. In floating point they are not. Floats are dense near zero and sparse near one: fp32 can hold numbers as small as about $10^{-38}$, but the gap between 1.0 and the next float is $2^{-23}$. Add anything at or below $2^{-24}$ to 1.0 and the sum rounds back to exactly 1.0. That happens once $e^{-m} \le 2^{-24}$, which is $m \ge 24 \ln 2 = 16.64$ (reasoned). The paper calls this the precision cliff $c_d$ and gives $c_{16} \approx 5.5$ for bf16, $c_{32} \approx 16.6$ for fp32 and $c_{64} \approx 36.7$ for fp64 (reported; $8 \ln 2 = 5.55$ and $53 \ln 2 = 36.74$ match, reasoned).

Past the cliff the standard path computes $q_t = 1/1.0 = 1.0$, subtracts 1, and gets a target gradient of exactly 0.0. The loss is $\log(1.0) = 0$ as well. But the non-target gradients are each $e^{z_j - z_t}$ divided by 1.0. That numerator was never added to 1, so it stays a perfectly good small positive number all the way down to exponent underflow near $m \approx 88$. The row of gradients now sums to $e^{-m}/(1 + e^{-m}) > 0$. The invariant is gone.

What PyTorch prints

Here is the paper's own example, $\mathbf{z} = [9, -9, -9]$ with target 0, so $m = 18 - \ln 2 = 17.31$, run through torch.nn.functional.cross_entropy and through my PyTorch port of the fix (measured):

-- torch.float32
torch.log_softmax(z)             [ 0.0000e+00, -1.8000e+01, -1.8000e+01]
log_softmax_fixed(z)             [-3.0460e-08, -1.8000e+01, -1.8000e+01]
F.cross_entropy     loss=0.0000e+00 grad=[ 0.0000e+00, 1.5230e-08, 1.5230e-08] sum=3.0460e-08
cross_entropy_fixed loss=3.0460e-08 grad=[-3.0460e-08, 1.5230e-08, 1.5230e-08] sum=0.0000e+00
-- torch.float64
F.cross_entropy     loss=3.0460e-08 grad=[-3.0460e-08, 1.5230e-08, 1.5230e-08]

The fp32 standard gradient is the correct gradient with its first entry deleted. In fp64 the same call is right, because 17.31 is below fp64's cliff. The paper reports the identical numbers from jax.nn and optax (reported), and a NumPy textbook version, q = e / e.sum(); g = q; g[t] -= 1, prints the same [0.000e+00 1.523e-08 1.523e-08] (measured).

bf16 hits it at a margin anyone reaches in minutes. With $\mathbf{z} = [3.5, -3.5, -3.5]$, $m = 6.31$ (measured):

F.cross_entropy bf16   loss=0.0000e+00 grad=[ 0.0000e+00, 9.1171e-04, 9.1171e-04]
fixed bf16             loss=1.8082e-03 grad=[-1.8234e-03, 9.1171e-04, 9.1171e-04]
exact fp64 grad                         [-1.8204e-03, 9.1022e-04, 9.1022e-04]

The widget below does the same arithmetic in your browser, one rounding at a time, for three classes. Drag the margin across the cliff and watch the standard target gradient fall off a wall while the fixed one keeps decaying as $e^{-m}$.

The cliff is a little earlier than 16.6

One thing the paper's single number hides (measured). I swept $m$ in steps of 0.05 with one target logit and $K-1$ equal others, and recorded where F.cross_entropy's target gradient first became exactly zero:

With two classes it is the textbook $24 \ln 2$. With more, the $K-1$ small terms are not summed before they meet the 1.0; each one is added to a running total near 1.0 and rounded on its own, so each is lost once it alone falls below half a step. With $K = 3$ each term is $e^{-m}/2$, and the cliff moves to $16.64 - \ln 2 = 15.95$, exactly what PyTorch shows (reasoned from measured). Where it lands for larger $K$ depends on how the kernel vectorises the sum. bf16 does not move, because PyTorch's CPU kernel accumulates in fp32 and the cliff comes from rounding $q_t$ to bf16 afterwards. The paper's claim is unaffected; the cliff is, if anything, closer than it says.

The fix

Never add $e^{-m}$ to 1. Take the max logit out of the sum entirely, compute the margin from the others with a log-sum-exp, and let a numerically stable softplus, $\operatorname{softplus}(x) = \log(1 + e^{x})$, do the rest:

Because the target slot is the constant 0 rather than $z_t - z_t$, $z_t$ reaches the loss only through the $K-1$ differences, so autodiff gives each non-target $g_j = e^{z_j - z_t - \mathcal{L}}$ and the target exactly minus their sum. The paper ships it as nine lines of JAX. My PyTorch port of the hard-target version:

import torch, torch.nn.functional as F
 
def cross_entropy_fixed(z, t):
    z_t = z.gather(-1, t[..., None])                      # target logit
    is_t = F.one_hot(t, z.shape[-1]).bool()               # target mask
    others = (z - z_t).masked_fill(is_t, float("-inf"))   # exclude target from the sum
    m = -torch.logsumexp(others, -1)                      # logit margin
    return F.softplus(-m).mean()                          # = -log_softmax(z)_t

log_softmax gets the same treatment with the argmax standing in for $t$. Across the whole sweep above, the port's gradient row summed to zero to within one rounding step in every format, and its target entry matched fp64 (measured).

Why a missing $10^{-8}$ matters: Adam

A gradient of $3 \times 10^{-8}$ that should have been there is not obviously a problem. It becomes one because Adam does not care how big a gradient is. It divides the running mean of each parameter's gradient by the square root of the running mean of its square:

While $|g| \gg \epsilon$, that normalised step is about 1 whatever the size of $g$. So the paper's account of the slingshot (Thilak et al., 2022), the recurring loss spikes after a grokking network memorises its training set, runs like this (reported):

1. After memorisation the margins grow, because the loss keeps pushing $q_t$ toward 1.
2. Once a batch sits past the cliff, the target gradients are zero and what is left is the small, lopsided push from the non-targets. The gradient almost vanishes, and $\hat v_i$ decays toward zero with it.
3. Weight decay and that residual push drag a few examples back below the cliff. Their target gradients switch back on at about $e^{-m}$.
4. Adam divides those restored gradients by a $\sqrt{\hat v_i}$ that has just decayed. The step is enormous, the loss spikes, $\hat v_i$ resets, and the cycle starts again.

With the fixed loss, $g_t$ fades smoothly through the cliff, $\hat m_i$ and $\sqrt{\hat v_i}$ decay together, and there is nothing to snap back.

The paper adds a second result that I find the more useful one. At a stationary point the update is zero, so $g_i / (|g_i| + \epsilon) = -\lambda\theta_i$. Once $|g_i|$ drops below $\epsilon$ the normalised step shrinks like $g_i / \epsilon$, and since $|g_i| \propto e^{-m}$, weight decay stops the margin near

For PyTorch's default $\epsilon = 10^{-8}$ that is 18.42, just past fp32's 16.64 (reasoned). That is the whole reason the default triggers slingshots: training wants to settle at a margin the standard loss cannot represent. At $\epsilon = 10^{-7}$ the equilibrium is 16.12, below the cliff, and the paper notes standard cross-entropy then settles without spiking (reported). Barron's thread puts it as "the default eps=1e-8 happens to be juuust past where this bug kicks in for fp32 (this may be why some think a small eps is dangerous)."

expand

The paper's Table 1 puts numbers on the top row: on $a + b \bmod 31$ at $\epsilon = 10^{-8}$, fp32 fixed cross-entropy reaches 100.0 ± 0.0% test accuracy against 94.3 ± 5.4% for standard (reported). The bottom row is the dramatic one: with no weight decay and $\epsilon = 10^{-16}$, one spike at around step 1,250 drops standard cross-entropy from 61% to 23%, while the fixed loss continues to 85% (reported).

I ran the grokking experiment

The paper's Section 6 setup is a 1-layer, width-64 ReLU Transformer on $a + b \bmod 31$ with a 50% training split, full-batch AdamW at learning rate $10^{-3}$, betas 0.9 and 0.98, weight decay 1, $\epsilon = 10^{-8}$, 10,000 epochs. The paper does not give head count, MLP width, initialisation or normalisation, so I chose: 4 heads, an MLP of width 256, PyTorch default initialisation, no LayerNorm, inputs a b =, read out at the last position. That leaves 480 training pairs and 481 held out. Each run is one process, one thread, nice -n 19, and took 2.5 to 11 minutes on the shared CPU, depending on load (measured). The only difference between a standard and a fixed run is the loss function: same seed, same split, same initialisation.

What I saw (measured, three seeds at learning rate $10^{-3}$):

• The slingshots are real and the fix removes them. Standard cross-entropy produced 6, 6 and 6 slingshots in 10,000 epochs (a mean-margin drop of more than 3 within 100 epochs); the fixed loss produced 0, 0 and 0. The standard runs climb to a mean margin of 16.05 to 16.24, the moment a quarter to a third of training pairs are past the 16.64 cliff, and fall back to between 4.75 and 8.14.
• The fixed margin settles below the predicted equilibrium. It flattens at 16.51 to 16.83 by epoch 10,000, not at the 18.42 that $\log(1/\epsilon)$ predicts or the 17.9 the paper's Table 1 reports for fp32. My guess is that the approximation drops constants, such as the weight norm, that differ in my model (reasoned).
• Neither loss grokked in 10,000 epochs. Final test accuracy was 46.2, 47.6 and 50.7% for standard against 43.7, 46.2 and 48.6% for fixed. My model groks slower than the paper's, so 10,000 epochs is not enough to test the speed claim.

So I ran the same seeds at a learning rate of $3 \times 10^{-3}$, which makes weight decay act three times faster per epoch (switch to it in the widget).

The result was the opposite of the paper's (measured, three seeds):

Standard cross-entropy grokked on all three seeds; the fixed loss grokked on one, later than standard did on the same seed, and on seed 0 it was just starting to climb when the run ended. In the standard runs, test accuracy climbs between slingshots: on seed 0 it goes from 50.5% at epoch 4,800, just after one, to 90.9% at epoch 6,000, just after the next. The fixed margins sat at 16.60 to 17.17 at the end, again under 18.42.

That is not what the paper reports, and I do not think it refutes it: three seeds of one architecture I had to partly guess, against five seeds and four benchmarks. The paper itself says that whether spikes help or hurt depends on how the weight norm is controlled, and that with LayerNorm and no weight decay each spike "acts as an accidental learning-rate decay" (reported). My model has weight decay and no LayerNorm, the case where the paper says spikes only hurt. What my runs do support without reservation is the mechanism: every sawtooth tops out just as a quarter or more of the training pairs cross the cliff, and the fixed loss has no sawtooth at all.

expand

Bug two: Softmax Cancellation

The second bug lives in distillation, where the target is itself a softmax, $\mathbf{p} = \operatorname{softmax}(\mathbf{z}_T)$ of teacher logits. The standard code computes $\mathbf{q}$ and $\mathbf{p}$ separately, rounds both to probabilities, and subtracts. When the student is close to the teacher, that is the subtraction of two nearly equal numbers, and the small logit gap $\boldsymbol{\delta} = \mathbf{z}_S - \mathbf{z}_T$ is gone before the subtraction happens. Worse, $\mathbf{p}$ comes from softmax in the forward pass while $\mathbf{q}$ comes from exp(log_softmax) in the backward pass, two different sequences of roundings, so even at $\mathbf{z}_S = \mathbf{z}_T$, where the true gradient is zero, they do not cancel (reported).

The fix computes $\mathbf{q} - \mathbf{p}$ from $\boldsymbol{\delta}$ directly, in log space, with expm1 and log1p so nothing is ever added to or subtracted from 1:

with $s_k = \log(q_k / p_k)$ and $u_k = q_k - p_k$, a fallback to $q_k - p_k$ when $|s_k| \ge 1$, and a final projection $g_k = u_k - q_k \sum_i u_i$ so the row sums to zero.

My PyTorch reproduction of the paper's example, $\mathbf{z}_T = [2, 1, 0]$ and $\mathbf{z}_S = \mathbf{z}_T + [-10^{-7}, 10^{-7}, 0]$ in fp32 (measured). In fp32 that perturbation rounds to $\pm 1.19 \times 10^{-7}$, so the exact gradient is of those rounded inputs:

exact (fp64 of the rounded inputs)  [-4.5955e-08,  4.1442e-08,  4.5131e-09]
fp32 standard (PyTorch)             [-5.9605e-08,  4.4703e-08,  7.4506e-09]  sum=-7.451e-09
fp32 fixed                          [-4.5955e-08,  4.1442e-08,  4.5131e-09]  sum= 1.776e-15

PyTorch's standard gradient is wrong by 30% in the first entry and 65% in the third. The paper's JAX run of the same thing is wrong differently, [-5.960e-8, 1.452e-8, -1.123e-9], with the third sign flipped (reported). Different frameworks, different rounding sequences, different wrong answers, which is the point.

Then the clean test, suggested by more than one reply to Barron's thread: set the student equal to the teacher, where the gradient must be zero. 64 rows of 1,000 random logits each (measured):

That residual is harmless in plain SGD. Under Adam it is not, for the same reason as before: any gradient larger than $\epsilon$ becomes a full-sized step. The paper initialises a student with its teacher's weights on $a + b \bmod 53$ and runs 2,000 Adam steps: the fp64 residual of about $10^{-17}$ sits below $\epsilon = 10^{-8}$ and nothing moves, fp32's $10^{-8}$ residual causes an initial KL spike and a floor near $10^{-6}$, and bf16's $10^{-3}$ drives the student away from the teacher and costs test accuracy, while the fixed loss stays at KL 0 and 100% in every format (reported, Section 7). I did not re-run that experiment.

expand

Who should care

The paper is candid that "the extreme margins and near-zero student–teacher gaps probed by our experiments are uncommon in standard large-scale training" (reported). I would not expect language-model pretraining over a 100,000-token vocabulary to sit at a margin of 16 on most tokens (reasoned). But four places do reach it.

• bf16 losses. The bf16 cliff is a margin of 5.55, a probability of about 0.996 on the right answer (reasoned). Most training code upcasts logits to fp32 before the loss for this reason among others; code that does not is collapsing on every easy token.
• Small-output problems trained long. Grokking, classifiers with a few classes, anything past memorisation with weight decay and default $\epsilon$.
• Distillation from a close teacher, and KL penalties that start the policy as a copy of the reference, which is how most RL fine-tuning begins. One reply points at continual-learning penalties that start the new model as an exact copy of the old one.
• Attention. Attention is a softmax too. A reply to the thread from the authors of Broken Symmetry in BF16 Attention argues the same broken zero-sum shows up in FlashAttention's bf16 backward. That claim is their own; a second group has since made the stochastic-rounding version of it, and I measured that one in the 25 October section below.

Prior fixes treated symptoms: Liu et al. project the gradient to zero sum, StableMax (Prieto et al., 2025) replaces the exponential, and others cast to fp64. The paper's Table 1 finds zero-sum projection stalls the margin near 18 and drops test accuracy below 50% at $\epsilon \le 10^{-12}$, and StableMax lets margins run to about $10^4$ (reported). Liu replied that their main contribution is the mechanism, "Numerical Feature Inflation", and that the projection was a causal test.

What the fix costs: an argmax, a masked log-sum-exp and a softplus instead of one fused kernel. The paper says it runs "without meaningful slowdown" (reported). I did not benchmark it, and a reply asking about fused vocabulary-scale kernels has no public answer yet. My unfused port allocates a few extra $B \times K$ temporaries, which at a large vocabulary is the memory a fused cross-entropy exists to avoid.

Two notes on provenance. Barron writes that the bug was "first found and fixed by Gemini Flash 3.7", that most experiments were run by Flash 3.8 and the first draft written by a Gemini model before he rewrote it (reported, from his thread); a reply points out the paper itself does not say so. And no code repository is published: the code is the listings on the project page, which is what I ported.

25 October: the same bug in FlashAttention's backward pass

A day after this page went up, Alexander Johansen (@AlexRoseJo) posted that "FlashAttention's backward pass has the same zero-sum bug Jon Barron just fixed in softmax. With FP4 stochastic rounding it's wrong on every single draw." The fix is in his ICLR 2027 submission, CroissAttn, a complete FP4 attention backward. The repository holds one file, the anonymous submission PDF (reported; no kernels are published). Everything below about the paper is from that PDF and his thread.

What the backward computes. For one query row, attention gives probabilities $P$ and output $O = PV$. The backward forms $dP = dO\,V^\top$ and then the softmax derivative

Because $P$ sums to one, the row sum of $dS$ is $\sum_j P_j dP_j - D$, which is zero exactly when $D$ equals the $P$-weighted mean of the row it is subtracted from. FlashAttention never stores $P$, so it takes the second form of $D$, the dot product of $dO$ with the stored output $O$. With rounded operands it is a different number from $\langle P, \widehat{dP}\rangle$ computed from the row actually used, and the gap is the row sum of $dS$ (reasoned).

Where the property is lost. Two places. If $O$ is stored in bf16, its rounding error is a fixed offset per row. With stochastic rounding of $dO$ and $V$, $\widehat{dP}$ is a random draw, while $D$ taken from $O$ is the average draw. The expected gradient is right, but each draw has a non-zero row sum, and the error lands on the key that holds most of the row: the attention sink. The submitted CroissAttn kernels use a high-precision auxiliary product $O' = PV_f$ for $D$, which fixes the mean and not the individual draw (reported, PDF Eq. 4 and Appendix A.3).

What CroissAttn changes. Compute $D$ from the same rounded $\widehat{dP}$ row: $D = \langle P, \widehat{dP}\rangle$. Then $\sum_j P_j(\widehat{dP}_j - D) = 0$ on every draw, and the mean is unchanged because the rule is linear in $\widehat{dP}$. Johansen's thread says the per-draw rule was derived after the submitted kernels were built, and training results with it are "to follow" (reported). All the paper's training numbers, such as the 0.0439-nat gap at 500M parameters and the 1.5 to 2.2 times backward speedup on B200, were produced with the auxiliary version, so they do not test this claim.

Their evidence. In the PDF, an exact enumeration of E2M1 rounding on a two-key row gives a total squared error of 0.21037027 with the auxiliary rule and 0.01670725 with per-draw centering, and 0 against 0.03609975 on a one-key row (reported, Table 8). The variance ratio falls from 1.98 on peaky rows to 1.14 on diffuse ones (reported, Table 9). The thread adds that on one sink row, 150 draws land within 1 times the true gradient 1 time in 150 with the shortcut and 150 times in 150 with CroissAttn; that experiment is not in the PDF, so I could not check it against a table. The paper notes a catch: for some downstream maps, sampled centering gives a larger error than the auxiliary rule (reported, Appendix C.1, a toy with 1/16 against 0).

I ran the row sum

(measured) I rebuilt the claim in PyTorch 2.14.1 on the shared 4-core CPU, two threads. One attention row, 256 keys, head size 64, Gaussian tensors; a sink row adds +6 or +12 to the first key's logit and shrinks its value vector to 0.01 of the others. Gradient operands $dO$ and $V$ are rounded to bf16, fp8 (E4M3) or fp4 (E2M1, blocks of 16 with an E4M3 scale), by round to nearest or by stochastic rounding I emulated myself; the stored $O$ is bf16. Each cell is 150 draws. Three choices of $D$ are compared, and the target is the exact-arithmetic gradient of the rounded forward. I did not run the authors' code, because none exists, so this is my reading of the rules, not their kernel.

• Stored output, fp4 with stochastic rounding, +6 sink row: the row sum of $dS$ averaged $3.2 \times 10^{-2}$ in absolute value, from a mean of $-5.7 \times 10^{-3}$. That is zero on average and wrong on each draw, which is what the thread says.
• Per-draw rule, same cell: $1.6 \times 10^{-8}$, which is fp32 rounding of $P$ and nothing else. It is $1.4 \times 10^{-9}$ or lower on diffuse rows and $3.9 \times 10^{-10}$ on the +12 sink row.
• bf16 with round to nearest gives a deterministic offset, not noise: $1.3 \times 10^{-3}$ on a diffuse row from the stored $O$.
• The sink amplifies it. On the +12 row the sink holds 99.9% of the probability and its true $dS$ is only $-1.3 \times 10^{-3}$. With fp4 and stochastic rounding the shortcut's sink gradient stayed within 1 times the true value in 16 of 150 draws and the per-draw rule in 150 of 150; with fp8 and stochastic rounding, 29 of 150 against 150 of 150. Mean relative error of $dS$ was 7.4 for the shortcut against 0.11 for per-draw (fp4, stochastic).
• On diffuse rows the three rules have the same error in $dS$ (0.138, 0.138 and 0.137 in fp4 with stochastic rounding), matching the thread's "about 1 times on diffuse rows".

Limits: one row and one seed per cell, Gaussian tensors rather than a trained model's, my own rounding emulation, no training. I did not check the speedup, the training gaps or a reply's claim that GProj fixes the bf16 side with a rank-one correction.

The takeaway

The bug is one addition: $1 + e^{-m}$. Everything else follows from where floats are dense and where they are sparse. If a quantity you care about is tiny, compute it directly; never compute it as the difference between 1 and something close to 1. The same lesson is why log1p and expm1 exist, and it is the whole of both fixes. For a different, older problem with the same output layer, see the softmax bottleneck; for softmax inside attention, the attention explainer; for training natively in low precision, A.X-K2; for loss spikes caught by an optimizer that watches itself, reactive optimizers; and for distillation measured on a reasoning model, Inkling-Small.