What is the softmax precision cliff?

The logit margin at which 1 + e^-m rounds to exactly 1: about 5.5 in bf16, 16.6 in fp32 and 36.7 in fp64. Past it, standard implementations return a zero loss and a zero target gradient.

Why does a zero target gradient cause loss spikes?

The remaining gradients no longer sum to zero, and Adam's second-moment estimate decays while gradients are tiny. When an example falls back below the cliff, its restored gradient is divided by a very small denominator, producing a large step.

Softmax and Cross-Entropy Numerics

Subtracting the max prevents overflow, not rounding. See exactly where the target gradient falls to zero, and the rewrite that keeps it.

9 min readReviewed October 2026By Oliver Perrin · Machine Learning Engineer

Learn with a guided session

Jump to a section

Why this matters

Softmax turns logits into probabilities, and cross-entropy scores those probabilities against a target. Every framework subtracts the maximum logit first so that exp cannot overflow. That fix is necessary but, as Barron, Gnanapragasam and Tschannen showed in October 2026, not sufficient: the standard implementations in JAX, PyTorch and textbooks still lose the target gradient once the model is confident enough.

The cause is ordinary floating point. Adding a tiny number to 1.0 can return exactly 1.0. When the denominator of softmax becomes exactly 1, the target probability becomes exactly 1, the loss becomes exactly 0, and the target gradient becomes exactly 0, while the other gradients stay positive. The gradients stop summing to zero.

Build the mental model

Measure confidence with the logit margin $m$ : how far the target logit sits above the log-sum-exp of the others. As $m$ grows, the true target gradient shrinks like $-e^{-m}$ , smoothly, all the way to underflow. The standard computation forms $1 + e^{-m}$ first. Once $e^{-m}$ is smaller than half the spacing between floating-point numbers near 1, that sum rounds to 1 and the gradient falls off a precision cliff.

Explore the ideaRelated Studio experiment

Keep the target gradient past the precision cliff

Standard bf16 and fp32 gradients end abruptly; the fixed fp32 expression keeps the small numerator.

Exact inputs and equations appear with the figure.

Preparing the scientific figure…

Scalar margin arithmetic with fp32 rounding and emulated round-to-nearest-even bf16. Zero gradients are omitted on this logarithmic axis, so the standard curves end. This does not reproduce a particular fused kernel or mixed-precision training run.

Barron et al., A Fix for Softmax and Cross-Entropy

Section 2, Equation 3

Experiment with related Studio inputs

Follow each line to the right. The standard computation steps down, then stops where 1 + e^-m rounds to 1: about m = 5.5 in bf16 and 16.6 in fp32. The fixed computation keeps going.

Computed with real fp32 rounding and emulated bf16 rounding for a single target logit. Points where the gradient is exactly zero are omitted, which is why the standard lines end.

Inspect the related Studio values

x,y,series
0,-0.3010299956639812,Standard (bf16)
0.25,-0.359021942641668,Standard (bf16)
0.5,-0.4214682310456047,Standard (bf16)
0.75,-0.4944261129281329,Standard (bf16)
1,-0.5693908745745943,Standard (bf16)
1.25,-0.6523651096393581,Standard (bf16)
1.5,-0.7454821336302755,Standard (bf16)
1.75,-0.8284563686950394,Standard (bf16)
2,-0.9311187105921871,Standard (bf16)
2.25,-1.0102999566398119,Standard (bf16)
2.5,-1.1294863643590205,Standard (bf16)
2.75,-1.2321487062561682,Standard (bf16)
3,-1.3668472801536244,Standard (bf16)
3.25,-1.4082399653118496,Standard (bf16)
3.5,-1.505149978319906,Standard (bf16)
3.75,-1.630088714928206,Standard (bf16)
4,-1.806179973983887,Standard (bf16)
4.25,-1.806179973983887,Standard (bf16)
4.5,-2.1072099696478683,Standard (bf16)
4.75,-2.1072099696478683,Standard (bf16)
5,-2.1072099696478683,Standard (bf16)
5.25,-2.1072099696478683,Standard (bf16)
5.5,-2.1072099696478683,Standard (bf16)
0,-0.3010299956639812,Standard (fp32)
0.25,-0.3587009569476138,Standard (fp32)
0.5,-0.4230362831219311,Standard (fp32)
0.75,-0.49373682567775584,Standard (fp32)
1,-0.5703423347604135,Standard (fp32)
1.25,-0.6522795508640861,Standard (fp32)
1.5,-0.7389143153112003,Standard (fp32)
1.75,-0.8295996013365485,Standard (fp32)
2,-0.9237129184890965,Standard (fp32)
2.25,-1.0206816226150572,Standard (fp32)
2.5,-1.1199976054673955,Standard (fp32)
2.75,-1.2212218816936253,Standard (fp32)
3,-1.3239847289527833,Standard (fp32)
3.25,-1.4279777049693034,Standard (fp32)
3.5,-1.532950799536737,Standard (fp32)
3.75,-1.638700019870078,Standard (fp32)
4,-1.7450596640856364,Standard (fp32)
4.25,-1.8519035955138177,Standard (fp32)
4.5,-1.959123872847957,Standard (fp32)
4.75,-2.066637976304475,Standard (fp32)
5,-2.174390417101289,Standard (fp32)
5.25,-2.2823248728561216,Standard (fp32)
5.5,-2.390394582190755,Standard (fp32)
5.75,-2.4985715886686997,Standard (fp32)
6,-2.6068497394352015,Standard (fp32)
6.25,-2.715167927340134,Standard (fp32)
6.5,-2.8235606969117595,Standard (fp32)
6.75,-2.9319987581610056,Standard (fp32)
7,-3.0404828669191772,Standard (fp32)
7.25,-3.1489540177198143,Standard (fp32)
7.5,-3.257405977648465,Standard (fp32)
7.75,-3.3660022810752572,Standard (fp32)
8,-3.4745201681063667,Standard (fp32)
8.25,-3.5829484252815895,Standard (fp32)
8.5,-3.6915836076569097,Standard (fp32)
8.75,-3.8001649193288354,Standard (fp32)
9,-3.908749550478631,Standard (fp32)
9.25,-4.017354858466477,Standard (fp32)
9.5,-4.125730256534371,Standard (fp32)
9.75,-4.234381041147947,Standard (fp32)
10,-4.342764924595948,Standard (fp32)
10.25,-4.450933450954355,Standard (fp32)
10.5,-4.560077920379423,Standard (fp32)
10.75,-4.6684173951682615,Standard (fp32)
11,-4.77756186459333,Standard (fp32)
11.25,-4.886263402330944,Standard (fp32)
11.5,-4.994270974557275,Standard (fp32)
11.75,-5.104145964729699,Standard (fp32)
12,-5.2076865566367685,Standard (fp32)
12.25,-5.321629908943605,Standard (fp32)
12.5,-5.432328206437295,Standard (fp32)
12.75,-5.543478658559962,Standard (fp32)
13,-5.644936299318738,Standard (fp32)
13.25,-5.747598641215887,Standard (fp32)
13.5,-5.844508654223943,Standard (fp32)
13.75,-5.969447390832243,Standard (fp32)
14,-6.078591860257311,Standard (fp32)
14.25,-6.224719895935548,Standard (fp32)
14.5,-6.321629908943605,Standard (fp32)
14.75,-6.446568645551905,Standard (fp32)
15,-6.446568645551905,Standard (fp32)
15.25,-6.622659904607587,Standard (fp32)
15.5,-6.622659904607587,Standard (fp32)
15.75,-6.923689900271567,Standard (fp32)
16,-6.923689900271567,Standard (fp32)
16.25,-6.923689900271567,Standard (fp32)
16.5,-6.923689900271567,Standard (fp32)
0,-0.3010299956639812,Fixed (fp32)
0.25,-0.35870092738550835,Fixed (fp32)
0.5,-0.42303624883957314,Fixed (fp32)
0.75,-0.49373682567775584,Fixed (fp32)
1,-0.5703422866347501,Fixed (fp32)
1.25,-0.652279492745647,Fixed (fp32)
1.5,-0.7389144217353153,Fixed (fp32)
1.75,-0.8295998198982882,Fixed (fp32)
2,-0.9237131085030447,Fixed (fp32)
2.25,-1.0206817583576406,Fixed (fp32)
2.5,-1.1199976054673955,Fixed (fp32)
2.75,-1.2212220432475993,Fixed (fp32)
3,-1.3239846607253358,Fixed (fp32)
3.25,-1.4279782684327815,Fixed (fp32)
3.5,-1.5329511307034278,Fixed (fp32)
3.75,-1.6386996678120267,Fixed (fp32)
4,-1.7450603387157104,Fixed (fp32)
4.25,-1.8519026464367008,Fixed (fp32)
4.5,-1.9591231365765387,Fixed (fp32)
4.75,-2.0666400511039607,Fixed (fp32)
5,-2.1743888458435823,Fixed (fp32)
5.25,-2.2823190228901367,Fixed (fp32)
5.5,-2.390390855628095,Fixed (fp32)
5.75,-2.4985733415786435,Fixed (fp32)
6,-2.6068420511728503,Fixed (fp32)
6.25,-2.715178082304978,Fixed (fp32)
6.5,-2.823566590721971,Fixed (fp32)
6.75,-2.931995991399,Fixed (fp32)
7,-3.0404571998533285,Fixed (fp32)
7.25,-3.1489432956257986,Fixed (fp32)
7.5,-3.257448780005576,Fixed (fp32)
7.75,-3.3659692537994776,Fixed (fp32)
8,-3.474501510733053,Fixed (fp32)
8.25,-3.5830429254061245,Fixed (fp32)
8.5,-3.6915914674755683,Fixed (fp32)
8.75,-3.800145532982693,Fixed (fp32)
9,-3.908703914338644,Fixed (fp32)
9.25,-4.017265677038108,Fixed (fp32)
9.5,-4.125830092216643,Fixed (fp32)
9.75,-4.234396517527198,Fixed (fp32)
10,-4.34296454000104,Fixed (fp32)
10.25,-4.45153380379883,Fixed (fp32)
10.5,-4.5601040267709045,Fixed (fp32)
10.75,-4.668674996668983,Fixed (fp32)
11,-4.777246543035809,Fixed (fp32)
11.25,-4.88581857015582,Fixed (fp32)
11.5,-4.994390936346622,Fixed (fp32)
11.75,-5.102963586044208,Fixed (fp32)
12,-5.211536478114386,Fixed (fp32)
12.25,-5.32010948452557,Fixed (fp32)
12.5,-5.428682639654707,Fixed (fp32)
12.75,-5.537255908152309,Fixed (fp32)
13,-5.645829263530775,Fixed (fp32)
13.25,-5.754402668113415,Fixed (fp32)
13.5,-5.862976105450714,Fixed (fp32)
13.75,-5.971549578999756,Fixed (fp32)
14,-6.080123100761703,Fixed (fp32)
14.25,-6.188696616782098,Fixed (fp32)
14.5,-6.297270179030691,Fixed (fp32)
14.75,-6.4058437730501785,Fixed (fp32)
15,-6.514417392670421,Fixed (fp32)
15.25,-6.622990956971233,Fixed (fp32)
15.5,-6.73156456977471,Fixed (fp32)
15.75,-6.840138112528796,Fixed (fp32)
16,-6.9487117609513405,Fixed (fp32)
16.25,-7.057285359703564,Fixed (fp32)
16.5,-7.165858985625292,Fixed (fp32)
16.75,-7.274432580786749,Fixed (fp32)
17,-7.383006182284615,Fixed (fp32)
17.25,-7.491579806504776,Fixed (fp32)
17.5,-7.600153420331542,Fixed (fp32)
17.75,-7.708727044217298,Fixed (fp32)
18,-7.817300693906956,Fixed (fp32)
18.25,-7.925874297122894,Fixed (fp32)
18.5,-8.034447903262397,Fixed (fp32)
18.75,-8.143021546873614,Fixed (fp32)
19,-8.251595155272389,Fixed (fp32)
19.25,-8.360168798216145,Fixed (fp32)
19.5,-8.468742388191194,Fixed (fp32)
19.75,-8.5773160195568,Fixed (fp32)
20,-8.685889623372777,Fixed (fp32)
20.25,-8.794463256809427,Fixed (fp32)
20.5,-8.90303687772196,Fixed (fp32)
20.75,-9.011610480690402,Fixed (fp32)
21,-9.120184104931619,Fixed (fp32)
21.25,-9.228757738634,Fixed (fp32)
21.5,-9.337331368039683,Fixed (fp32)
21.75,-9.445904978546114,Fixed (fp32)
22,-9.554478600613402,Fixed (fp32)
22.25,-9.663052235979555,Fixed (fp32)
22.5,-9.77162585586622,Fixed (fp32)
22.75,-9.8801994706676,Fixed (fp32)
23,-9.988773087631898,Fixed (fp32)
23.25,-10.097346705407361,Fixed (fp32)
23.5,-10.205920311625682,Fixed (fp32)
23.75,-10.31449394692874,Fixed (fp32)
24,-10.42306758552304,Fixed (fp32)

The cliff sits at about $m = 5.5$ in bf16, $16.6$ in fp32 and $36.7$ in fp64. The chart computes both versions with real fp32 and bf16 rounding. Before the cliff, the standard curve already moves in visible steps; past it, the line simply ends.

Work through the math

With target $t$ , define the margin

m = z_t - \log\sum_{j \ne t} e^{z_j}

The loss and target gradient are

L = \log\left(1 + e^{-m}\right), \qquad g_t = \underbrace{\frac{1}{1+e^{-m}} - 1}_{\text{standard}} = \underbrace{-\frac{e^{-m}}{1+e^{-m}}}_{\text{fixed}}

The two expressions are equal in exact arithmetic. In floating point, the standard one subtracts 1 from a number that has already rounded to 1. The fixed one keeps $e^{-m}$ in the numerator, so dividing by a denominator that rounded to 1 still returns about $-e^{-m}$ . The non-target gradients $g_j = e^{z_j - z_t}/(1+e^{-m})$ never added 1 to their numerators, which is why they survive while $g_t$ does not.

The fix computes the loss as $\operatorname{softplus}(-m)$ from the $K-1$ differences $z_j - z_t$ , and sets $g_t = -\sum_{j\ne t} g_j$ . The gradients then sum to zero by construction.

Key takeaway: Subtracting the max prevents overflow. It does not prevent rounding. Never form $1 + \text{tiny}$ and then subtract 1.

Why a zero-sum violation becomes a loss spike

Adam divides each update by a running estimate of the gradient's magnitude. When every example in a batch sits past the cliff, the target gradients vanish and that estimate decays toward zero. When an example falls back below the cliff, its gradient returns and is divided by a tiny denominator: a large step. The paper shows this explains the "Slingshot" loss spikes reported during grokking, and that the fixed loss removes them.

Distillation has a second bug

With a soft target $p = \operatorname{softmax}(z_T)$ , the gradient is $q - p$ . Computing $q$ and $p$ separately and subtracting nearly equal probabilities loses precision, so the gradient can stay non-zero even when student and teacher logits are identical. The paper computes $q - p$ directly from $z_S - z_T$ in log space with expm1 and log1p.

The same identity inside attention

The softmax Jacobian always sums to zero across a row. FlashAttention's backward pass relies on it: for one query row, $dS = P \circ (dP - D)$ with $D = \operatorname{rowsum}(dO \circ O)$ , and the subtraction of $D$ is what makes each row of $dS$ sum to zero. Low-precision arithmetic can break that identity just as it breaks the cross-entropy one, so check it explicitly when testing an attention kernel.

Read the implementation

A PyTorch version of the fixed hard-target loss for one batch. Autograd produces the correct gradient because the target slot enters only through the differences.

python

1import torch
2import torch.nn.functional as F
3
4def cross_entropy_fixed(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
5    # logits: [B, K]; target: [B] class indices
6    z_t = logits.gather(1, target[:, None])                     # [B, 1]
7    diffs = logits - z_t                                          # z_j - z_t
8    diffs = diffs.scatter(1, target[:, None], float("-inf"))      # drop the target slot
9    margin = -torch.logsumexp(diffs, dim=1)                       # m
10    return F.softplus(-margin).mean()                             # log(1 + e^-m)
11
12z = torch.tensor([[9.0, -9.0, -9.0]], requires_grad=True)
13t = torch.tensor([0])
14cross_entropy_fixed(z, t).backward()
15print(z.grad)            # target entry is small and negative, not 0; row sums to 0

In the paper's fp32 JAX example, the standard computation on this row returns a loss and target gradient of exactly 0, because its margin of about 17.3 is past the fp32 cliff. The fixed version returns a loss of about 3.05e-8 and a target gradient of about -3.05e-8. The PyTorch code above implements the same hard-target identity; exact rounding depends on the kernel and accumulation precision.

Questions and trade-offs

Conceptual: Why do frameworks subtract the max logit? (To keep exp from overflowing. It shifts every logit by the same amount, so probabilities are unchanged.)
Applied: Why does the target gradient vanish while the others do not? (The target gradient is computed as q_t minus 1 after 1 + e^-m has rounded to 1; the others are e^(z_j - z_t) divided by that same denominator and never lost their small numerators.)
Implementation: How do you compute log(1 + e^x) safely? (Use softplus or log1p: for x > 0, x + log1p(e^-x); otherwise log1p(e^x).)
Systems-level: Why does this matter more in bf16? (Its cliff is at a margin of about 5.5, which ordinary confident predictions reach routinely.)
Failure modes: How would you detect this bug in a training run? (Check that per-example logit gradients sum to zero, and log the fraction of examples with margins past the cliff for the training dtype.)

Check your understanding

From memory: define the logit margin, write the standard and fixed target gradient, and state the precision cliff for fp32. Check against Stage 3.

Sources and further reading

Barron, Gnanapragasam & Tschannen, A Fix for Softmax and Cross-Entropy: Softmax Collapse, Softmax Cancellation, and the Slingshot mechanism.
Thilak et al., The Slingshot Mechanism: the loss spikes observed during grokking.
Dao et al., FlashAttention: the tiled forward pass and the rowsum term in the backward pass.
Higham, Accuracy and Stability of Numerical Algorithms: rounding and stable summation.

Oliver Perrin

Machine Learning Engineer · Founder, LiminalML

LiminalML brings together concept explanations, mathematical examples, and working visuals. Find more writing by Oliver on the LiminalML Substack.

Keep exploring the idea.

Ask a follow-up in a guided session, work through the self-check, or open the visual in Studio to build on it.

Ask a follow-up Try a recall question See pricing

Related concepts

Deep Learning

FlashAttention

Follow tiled exact attention and the online-softmax invariant that avoids a full intermediate score matrix.

LLMs

Reasoning Distillation

Separate imitating a teacher's text from matching its distribution, and see why distillation is cheap but capped by the teacher.

Optimization

Gradient Descent (SGD, Momentum, Adam)

Follow updates across a loss surface, then connect learning rate, momentum, and Adam to the update equations.