Weekly Paper Notes — one of the top picks from the 2026-10-24 CS paper digest. Area: AI / ML.

Authors: Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen (Google DeepMind) arXiv: 2610.31642 · PDF

TL;DR

Every deep learning framework computes softmax and cross-entropy the same way: subtract the largest logit so exp can’t overflow, then normalise. That trick fixes overflow and leaves two quieter bugs in place.

The first, which the authors call Softmax Collapse, hits when the model is very confident. The denominator becomes 1 + e^{-m} for logit margin m, and once e^{-m} is too small to change 1.0 it’s rounded away. The target-class gradient goes to exactly zero while the non-target gradients stay positive, so the gradient vector no longer sums to zero. The cliff sits at a margin of about 5.5 in bf16, 16.6 in fp32 and 36.7 in fp64.

The second, Softmax Cancellation, hits in distillation. Standard code computes student and teacher probabilities separately and subtracts them, which throws away precision when the two are close and leaves a non-zero gradient even when student and teacher logits are identical.

The fixes are a few lines each. They rewrite the computation so 1 is never added to a tiny number and nearly-equal probabilities are never subtracted, using softplus, expm1 and log1p. On the paper’s experiments the fix removes the “Slingshot” loss spikes seen in grokking, speeds grokking up, keeps a distilled student sitting on its teacher, and exposes a clean equilibrium: under AdamW with weight decay, logit margins stop growing at log(1/ε).

What problem is the paper actually attacking?

The textbook numerically stable softmax shifts logits by their max:

z̃_k = z_k − z_max
logsoftmax(z)_k = z̃_k − log Σ_j exp(z̃_j)

Higham’s Accuracy and Stability of Numerical Algorithms covers this, and JAX, PyTorch and TensorFlow all do it. Cross-entropy against a target distribution p then has the exact gradient q − p, where q is the softmax output.

The authors’ point is that the shift only protects against overflow. It does nothing about the other way floating point loses information, which is adding a very small number to 1. In fp32 a value near zero can be as small as about 10⁻³⁸, but anything below about 10⁻⁷ added to 1.0 disappears.

Two research threads had already run into the symptoms without a clean diagnosis. Thilak et al. (2022) described the Slingshot Mechanism: after a network memorises its training set, the loss sits near zero for a while, then spikes, then drops again, and the spikes line up with jumps in test accuracy. Prieto et al. (ICLR 2025) linked grokking failures to numerical stability at the edge of float precision and proposed StableMax, a replacement for exp. Liu et al. (2026) argued the Slingshot spikes come from low precision and proposed projecting the gradient so it sums to zero. Others sidestepped the problem by casting to fp64.

Each of those is a workaround. This paper fixes the functions.

Two panels. Left: target-logit gradient magnitude falling exponentially with margin, with standard implementations dropping to zero at bf16, fp32 and fp64 cliffs. Right: distillation gradient norm shrinking as student logits approach the teacher’s, with standard implementations stuck at a noise floor. Left: the true target gradient shrinks smoothly as confidence grows, but standard code rounds it to zero at a format-specific cliff. Right: in distillation the standard gradient bottoms out at a noise floor instead of reaching zero. Source: Barron et al., arXiv:2610.31642, Figure 1.

The mechanism, part 1: never add 1 to a tiny number

Take hard-label cross-entropy with target t, and suppose the model is right and confident, so z_t is the max. Define the logit margin as how far the target logit sits above the log-sum-exp of everything else:

m = z_t − log Σ_{j≠t} exp(z_j)

After shifting by z_t, the denominator is 1 + e^{-m}. Mathematically the target gradient is

g_t = 1/(1 + e^{-m}) − 1        (standard form)
    = −e^{-m}/(1 + e^{-m})       (fixed form)

Those are equal on paper. In floats they aren’t. The standard form computes q_t = 1/(1+e^{-m}), which becomes exactly 1.0 once m passes the cliff, and subtracting 1 gives exactly zero. The fixed form keeps e^{-m} in the numerator. The denominator still rounds to 1, but dividing a small number by 1 leaves it small, so the gradient survives until e^{-m} itself underflows at m ≈ 88.

The non-target gradients e^{z_j − z_t}/(1 + e^{-m}) never had the problem, because their numerators are computed directly. So past the cliff the standard implementation reports a target gradient of 0 and a set of small positive non-target gradients. The forward loss reports 0 too.

The fix is structural. Set the shifted target slot to the constant 0 instead of computing z_t − z_t, compute the margin with a log-sum-exp that excludes the max, and combine with a stable softplus:

def cross_entropy(z, t):
    z_t = jnp.take_along_axis(z, t[..., None], axis=-1)
    is_t = jax.nn.one_hot(t, z.shape[-1], dtype=bool)
    others = jnp.where(is_t, -jnp.inf, z - z_t)
    m = -jax.nn.logsumexp(others, axis=-1)
    return jax.nn.softplus(-m)

Because z_t now only enters through the differences z_j − z_t, autodiff produces the target gradient as minus the sum of the others, and the gradients sum to zero by construction. The paper gives the same treatment to log_softmax and softmax using the argmax in place of t.

Their worked example is z = [9, −9, −9] in fp32 with target 0, a margin of about 17.3, just past the fp32 cliff. optax.softmax_cross_entropy_with_integer_labels returns a loss of 0.0 and gradients [0.0, 1.523e-8, 1.523e-8]. The fixed version returns 3.046e-8 and [-3.046e-8, 1.523e-8, 1.523e-8].

The mechanism, part 2: don’t subtract nearly equal probabilities

Distillation and KL regularisation use a soft target p = softmax(z_T) from a teacher. The gradient is still q − p. Standard code computes q and p as separate probability vectors and subtracts them. As the student approaches the teacher, the information about their difference lives in the logit difference δ = z_S − z_T, and it’s lost once both have been rounded to probabilities. Worse, the forward pass computes p with softmax and the backward pass computes q as exp(logsoftmax), which are slightly different sequences of operations. So even at z_S = z_T exactly, q − p isn’t zero.

The fix works from δ directly in log space:

d_k = δ_k − Σ_i p_i δ_i
s_k = d_k − log1p( Σ_i p_i expm1(d_i) )     # log(q_k / p_k)
u_k = p_k expm1(s_k)                        # q_k − p_k
g_k = u_k − q_k Σ_i u_i                     # force Σ g = 0

Nothing here adds or subtracts 1 explicitly, and at δ = 0 every term is zero. The implementation returns the ordinary forward loss and substitutes g in the backward pass with a stop-gradient trick. In their fp32 example, with the student perturbed from the teacher by ±1e-7, the standard gradient shrinks one entry threefold and flips the sign of another. The fixed gradient matches the fp64 reference to every printed digit.

Why it matters for training: Slingshot is the cliff

Section 4 is the part I’d point people to. It explains the Slingshot Mechanism as a direct consequence of Softmax Collapse interacting with Adam.

After memorisation, margins grow. Once every example in a batch crosses the cliff, the target gradients vanish and the parameter gradients nearly vanish with them. Adam’s second-moment estimate v̂ decays toward zero. Eventually weight decay, or the leftover non-target gradients, pulls a few examples back under the cliff. Their target gradients switch back on and get divided by √v̂ + ε, which is now tiny. The result is a huge step, a loss spike, a reset of v̂, and the cycle repeats.

Three columns of training curves (mean logit margin, training loss, test accuracy) for standard and fixed cross-entropy, with and without weight decay Top row, with weight decay: standard cross-entropy repeatedly hits the fp32 cliff near margin 16.6 and rebounds; the fixed loss crosses it smoothly and settles near 18.4. Bottom row, without weight decay: one spike at the cliff drops standard test accuracy from 61% to 23%, while the fixed version keeps climbing to 85%. Source: Barron et al., arXiv:2610.31642, Figure 2.

Whether the spikes look helpful depends on the setup. With LayerNorm before the output and no weight decay, each spike inflates the weight norm and acts like an accidental learning-rate decay, which is how Thilak et al. saw them coinciding with generalisation. With weight decay, they slow grokking down: 94 ± 5% versus 100 ± 0% test accuracy on a + b mod 31. With neither, a single spike can wreck a run.

Results: a margin equilibrium you can predict

With the cliff gone, the fixed loss still stops growing its margin at about 18.4 in the weight-decay runs. The authors show this is an optimiser equilibrium, not a precision effect. At a stationary point of AdamW, the update g/(|g| + ε) + λθ is zero. After memorisation |g| shrinks like e^{-m}. While |g| ≫ ε the normalised step has magnitude about 1; once |g| ≪ ε it shrinks like g/ε. Balancing that against weight decay pins the margin at

m*(ε) ≈ log(1/ε)

For PyTorch’s and Optax’s default ε = 1e-8, that’s 18.4, which sits just above the fp32 cliff at 16.6. That’s why standard fp32 training with default Adam hits Slingshot: the optimiser is trying to settle 1.8 nats past the point where the loss function stops reporting gradients. With ε ≥ 1e-7 the equilibrium (≤ 16.1) lands below the cliff and standard cross-entropy settles quietly. The authors confirm log(1/ε) across many orders of magnitude of ε using LaProp, down to ε = 0, where the margin grows until fp32 gradients underflow at about 81.

The comparison table on a + b mod 31 covers bf16, fp32 and fp64 at three values of ε. Fixed cross-entropy reaches the predicted margins (17.9, 27.4, 36.8) and 99–100% test accuracy in every precision. Standard cross-entropy stalls near its format’s cliff: margins of about 5–6 in bf16 and 13–15 in fp32. In fp64 it’s fine at ε = 1e-8 and 1e-12, but at ε = 1e-16 the equilibrium of 36.8 lies past the fp64 cliff of 36.7 and even fp64 breaks.

The two competing fixes don’t hold up. Zero-sum gradient projection hands the target logit only 1/K of the missing gradient, stalls the margin near 18, and drops test accuracy below 50% at ε ≤ 1e-12. StableMax grows linearly instead of exponentially, so margins run into the thousands, and its accuracy trails the fixed loss throughout. The same pattern holds across four grokking benchmarks from Nanda et al., Varma et al. and Doshi et al.: fp32 with the fixed loss matches fp64 with the standard one.

For distillation, the authors initialise a student with the teacher’s exact weights, so the true gradient is zero and nothing should move. With standard soft-target cross-entropy, the residual gradient is about 1e-17 in fp64 (harmless, below ε), about 1e-8 in fp32 (a brief KL spike), and about 1e-3 in bf16, where Adam normalises it into full-size steps that push the student away from the teacher and cost test accuracy. The fixed loss returns exactly zero in all three.

Why this matters

The authors are upfront that “the extreme margins and near-zero student–teacher gaps probed by our experiments are uncommon in standard large-scale training.” Grokking on modular arithmetic is a small, adversarial setting. Nobody should read this as a hidden cause of instability in frontier pretraining.

Still, I think there are three reasons to care.

First, bf16’s cliff is at a margin of 5.5. That’s p ≈ 0.996. Plenty of real classifiers, reward models and token predictors hit that confidence on easy examples, and anything computing loss or softmax gradients in bf16 is silently zeroing the target gradient on them while keeping the others. Most stacks upcast the loss to fp32, which pushes the cliff to 16.6, but that’s a convention, not a guarantee.

Second, distillation and KL-regularised RL (policy versus reference model) are exactly the setting where Softmax Cancellation bites, because the two distributions start identical. A spurious gradient at initialisation, amplified by Adam, is a strange thing to be optimising against.

Third, the log(1/ε) result is useful on its own. It gives a closed-form account of where post-memorisation margins end up under AdamW, and it turns ε from a “don’t touch” constant into a knob with a clear meaning.

The cost is a handful of extra ops, no casting, and no change to the forward value. That’s a cheap trade, and it’d be a reasonable default for Optax and PyTorch to adopt.

Read alongside

📄 arXiv abstract · 📄 PDF

Part of the Weekly CS Paper Digest series. Summary written from a close read of the preprint; figures cropped from the arXiv PDF and reproduced here under fair use for educational commentary.