DEV Community

Cover image for I measured softmax's precision cliff. In fp32, cross-entropy returns exactly zero at margin 16.
frank chu
frank chu

Posted on

I measured softmax's precision cliff. In fp32, cross-entropy returns exactly zero at margin 16.

Jon Barron posted a paper today — "A Fix for Softmax and Cross-Entropy" — claiming that the standard implementations of softmax and cross-entropy in every major framework have two numerical bugs. Not "imprecision." Bugs, with specific, falsifiable numbers attached: a precision cliff at margin ≈5.5 in bf16, ≈16.6 in fp32, ≈36.7 in fp64, past which the loss and the target gradient round to exactly zero while the other gradients don't.

Precise claims deserve precise checking. I reproduced both bugs in an afternoon with nothing but NumPy 1.26.4 — no framework, computing the gradient analytically as q − p, which is exactly what the frameworks' backward passes compute. Everything below is numbers I ran today.

Bug 1: Softmax Collapse, and the cliff is where he says it is

The mechanism is embarrassingly simple. For a hard target, the loss is log(1 + e^(−m)), where m is the logit margin — how far the target logit sits above the log-sum-exp of the rest. Floating point near 1.0 is coarse: in fp32, anything smaller than half an ulp of 1.0 (≈5.96e−8) vanishes when added to 1.0. Once e^(−m) drops below that, the denominator 1 + e^(−m) rounds to exactly 1.0, the loss rounds to exactly 0, and the target gradient q_t − 1 rounds to exactly 0.

The non-target gradients are computed near 0, where floating point is fine-grained, so they don't round away. Result: the gradient vector stops summing to zero — it sums to exactly the target gradient that got deleted.

I swept the margin for a 3-class problem in three dtypes. First margin (0.1 steps) where the standard loss comes out as exactly 0.0:

dtype     measured cliff   theory −log(eps/2)
float16        7.7              7.62
float32       16.0             16.64
float64       36.1             36.74
Enter fullscreen mode Exit fullscreen mode

Close to the paper's numbers, but not identical — and the difference is the interesting part. The paper's cliff assumes the small terms meet the 1.0 in one addition. In a real reduction, e^(−m) arrives split across the non-target classes, and each piece is judged against half-ulp separately. I checked the dependence on class count in fp32:

K=2:   cliff at margin 16.7
K=3:   cliff at margin 16.0
K=11:  cliff at margin 15.8
K=101: cliff at margin 16.0
Enter fullscreen mode Exit fullscreen mode

And the raw rounding experiment explains the scatter. At margin 16, each of the two non-target exponentials is 5.63e−08 — just under half an ulp of 1.0:

sum([1.0, a, a])  = 1.0         -> log = 0.0
sum([a, a, 1.0])  = 1.0000001   -> log = 1.19e-07
Enter fullscreen mode Exit fullscreen mode

Same three numbers, different order, one survives and one doesn't. So "the fp32 cliff" isn't a constant of nature; it's a property of your reduction tree — sequential vs. pairwise summation, how many classes, how the kernel tiles. It lives somewhere in the 15.8–16.7 band for fp32, and the paper's 16.6 is the clean two-class idealization. Frameworks with parallel tree reductions will each have their own.

The paper's worked example reproduces digit for digit. With z = [9.0, −9.0, −9.0], target 0, fp32:

standard:  loss = -0.0        grad = [ 0.0,       1.523e-08,  1.523e-08]  sum = 3.046e-08
fixed:     loss =  3.046e-08  grad = [-3.046e-08, 1.523e-08,  1.523e-08]  sum = 0.0
fp64 ref:  loss =  3.046e-08  grad = [-3.046e-08, 1.523e-08,  1.523e-08]
Enter fullscreen mode Exit fullscreen mode

The margin sweep (fp32, K=3) shows the damage isn't a clean step either — the loss is already quantized garbage a full nat before it hits zero:

m      std loss     std g_target  grad sum   | fixed loss   fixed g_target
5.0    6.7153e-03   -6.693e-03    8.4e-08    | 6.7153e-03   -6.693e-03
10.0   4.5299e-05   -4.530e-05    9.8e-08    | 4.5399e-05   -4.540e-05
15.0   2.3842e-07   -2.980e-07    7.9e-09    | 3.0590e-07   -3.059e-07
16.0   0.0           0.0          1.1e-07    | 1.1254e-07   -1.1254e-07
20.0   0.0           0.0          2.1e-09    | 2.0612e-09   -2.0612e-09
30.0   0.0           0.0          9.4e-14    | 9.3576e-14   -9.3576e-14
Enter fullscreen mode Exit fullscreen mode

At m=15 the standard loss reads 2.38e−07 — that's exactly 2^−22, one ulp — against a true value of 3.06e−07. A 22% error, one nat before the cliff edge. And note the gradient sum is never exactly zero even at m=5, because the softmax is evaluated as exp(log_softmax): the rounding residue is ~1e−7 in fp32 at every margin. Past the cliff that residue stops being noise relative to the signal and becomes the whole gradient.

Why it matters is Adam. Adam normalizes each gradient by its running magnitude, so a target gradient that blinks off at the cliff while the non-target gradients stay on produces exactly the asymmetric updates the paper blames for the "slingshot" loss spikes in grokking experiments. I didn't retrain a grokking model to check that part — the paper's training curves are their evidence. What I can confirm is the arithmetic underneath them: the blinking is real, and it starts where they say it does, give or take your reduction tree.

Bug 2: Softmax Cancellation has a hard error floor

The second bug is about distillation: gradient = q − p, student softmax minus teacher softmax. When the student is close to the teacher, you're subtracting two nearly equal probability vectors, and the small true difference drowns in the rounding of each side.

The paper's example: teacher [2, 1, 0], student perturbed by [−1e−7, +1e−7, 0] in fp32 (which fp32 stores as [−1.19e−07, +1.19e−07, 0] — already a hint about the precision regime). My NumPy run:

standard grad:  [ 0.0,        4.470e-08,  2.235e-08]   sum = 6.7e-08
fixed grad:     [-4.596e-08,  4.144e-08,  4.513e-09]   sum = -2.2e-15
exact (fp64):   [-4.596e-08,  4.144e-08,  4.513e-09]
Enter fullscreen mode Exit fullscreen mode

The standard gradient's first component is exactly 0.0 where the true value is −4.6e−08 — 100% error — and the third component is off by 395%. The fixed version (compute q − p in log space with expm1/log1p, per the paper's Eq. 9) matches the fp64 reference to 9.2e−16. Worst component error: standard 4.6e−08, fixed 9.2e−16. That's eight orders of magnitude from rearranging the arithmetic.

Sweeping the perturbation size shows it's a floor, not a proportional error:

delta    exact ‖g‖    std max err   fixed max err
1e-03    5.21e-04     5.27e-08      2.27e-11
1e-05    5.21e-06     4.55e-08      1.61e-13
1e-06    4.96e-07     1.12e-08      1.95e-14
1e-07    6.21e-08     4.60e-08      9.16e-16
5e-08    0 (unrepresentable in fp32 — the perturbation itself rounds away)
Enter fullscreen mode Exit fullscreen mode

The standard implementation's error sits at ~1e−8–5e−8 no matter how small the true gradient gets. At δ=1e−7 the error is 74% of the gradient's own norm. Distillation spends its late training exactly in this regime — student near teacher, true gradients tiny — which means a standard soft-target loss is, in that regime, mostly returning rounding noise with the right sign sometimes.

One honest negative from my run: the paper also says the gradient is non-zero even when student equals teacher exactly, because forward and backward evaluate softmax through different op sequences. In my NumPy setup the two paths happen to be the same operations, and identical student/teacher gave me exactly 0.0. That sub-claim depends on your framework's kernels; the floor above does not — it appears the moment δ is small and non-zero.

The fix is embarrassingly cheap

For hard targets, never form 1 + e^(−m). Compute the margin excluding the target and use softplus:

def fixed_ce_and_grad(z, t):
    diffs = np.array([z[j] - z[t] for j in range(len(z)) if j != t],
                     dtype=z.dtype)
    mx = np.max(diffs)
    m = -(mx + np.log(np.sum(np.exp(diffs - mx))))  # margin, target excluded
    loss = np.logaddexp(np.zeros((), dtype=z.dtype), -m)  # softplus(-m)
    g = np.zeros_like(z)
    gj = np.exp(diffs - loss)                      # non-target gradients
    idx = 0
    for j in range(len(z)):
        if j != t:
            g[j] = gj[idx]; idx += 1
    g[t] = -np.sum(gj)                             # target grad by construction
    return loss, g
Enter fullscreen mode Exit fullscreen mode

The last line is the whole trick: the target gradient is defined as the negated sum of the others, so it cannot round away independently, and the gradient sums to zero by construction.

Every fixed number above came out of ~40 lines like this. No fp64 casting, no surrogate loss.

What I take from this

Two things. First, both bugs are real, at the magnitudes claimed. If you train with cross-entropy in fp32 — everyone — your loss has been reporting exactly 0.0 on your most confident examples while still pushing on the non-target logits, and your distillation gradients have had a ~5e−8 noise floor that no learning-rate schedule can fix.

Second, the cliff isn't a constant. It moved by almost a full nat depending on how many classes I had and how the summation was ordered. Anyone who "checks" their setup by testing whether the loss hits zero at exactly 16.6 will measure their reduction tree, not the phenomenon. The phenomenon is the band, and your kernel is somewhere in it.

Sources: Jon Barron, "A Fix for Softmax and Cross-Entropy," https://jonbarron.info/softmaxfix/ (posted October 2026; cliff values and the slingshot/grokking experiments are his). All measurements above are my own, run today with NumPy 1.26.4 on the analytic gradient q − p; framework kernels will differ in detail, not in mechanism.

Top comments (0)