A fix for softmax and cross-entropy — Softmax Collapse and the precision cliff — What does it mean?

The news. On October 22, 2026, a Google DeepMind team posted "A Fix for Softmax and Cross-Entropy" (arXiv 2610.31642). The paper names two numerical bugs in standard softmax-family code, Softmax Collapse and Softmax Cancellation, and gives short replacement implementations of log-softmax, softmax and cross-entropy. It reports that the fix removes Slingshot loss spikes, speeds up grokking with weight decay, and stops a student network from drifting away from its teacher during distillation. Read the paper →

Picture the softmax denominator as a bank balance. Standard code first shifts every logit so the largest becomes 0, which makes its exponential exactly 1 and prevents overflow. That shift is sound, but it parks the biggest term at 1 and then adds every other class on top of it. When the model is very confident, the other classes together contribute only e^-m, where m is the logit margin. In the ledger picture the balance is $1.00, the other classes are $0.00000003 of interest, and the ledger keeps only seven decimal places, much as float32 keeps about seven digits.

A float32 number near 1 moves in steps of about 1.2e-7, so once e^-m drops below about 6e-8, the sum 1 + e^-m is stored as exactly 1 and the interest is gone. The forward pass then reports a loss of exactly 0. The backward pass forms the target gradient as 1 minus the rounded probability, which is also exactly 0. Meanwhile the other classes' gradients are computed straight from small exponentials that floating point stores well near zero, so they stay positive.

The gradients of softmax cross-entropy must sum to zero across classes, and after the cliff the offsetting entry has vanished while the deposits remain, so the books no longer balance. This is the first bug, Softmax Collapse. The paper shows the same failure in standalone softmax, which is also what turns attention scores into weights; its experiments, however, train with cross-entropy, and the source does not claim attention is affected in practice. For how logits become probabilities in generation, see From logits to probabilities.

Take the paper's own example: logits [9, -9, -9], target class 0, float32. The margin is about 17.3, because the log-sum-exp of the two -9 logits is -9 + ln 2 ≈ -8.31, and 9 + 8.31 ≈ 17.3. After the shift each other class contributes e^-18 ≈ 1.52e-8, so the extra term in the denominator is about 3.05e-8, below the roughly 6e-8 threshold, and 1 + 3.05e-8 is stored as 1. The paper reports that the standard jax.nn and optax code returns loss 0 and gradient [0, 1.52e-8, 1.52e-8], which sums to 3.05e-8 instead of 0. The fixed code returns loss 3.05e-8 and gradient [-3.05e-8, 1.52e-8, 1.52e-8], which sums to zero.

The numbers are tiny, so why does it matter? Because AdamW is the accountant who divides every adjustment by the recent average size of adjustments. After memorization the margins keep growing; once every example in a batch is past the cliff, the parameter gradients nearly vanish and Adam's running estimate of gradient size decays toward zero while weight decay slowly shrinks the parameters. Then a few examples fall back below the cliff and their target gradients switch on again. Adam divides those restored gradients by a very small denominator, takes a large step, spikes the loss, and the cycle repeats. The paper identifies this as the Slingshot. With the fixed cross-entropy the target gradient never rounds to zero, so the gradient and Adam's estimate decay smoothly together and the spikes disappear.

The effect depends on the setup. With weight decay, the paper reports spikes only slow grokking, and the fixed loss groks faster on all 5 seeds shown. Without weight decay, one spike can drop test accuracy sharply while the fixed loss continues smoothly. The paper also finds that after memorization the margin settles at an equilibrium between weight decay and Adam's epsilon, and that the default epsilon places that equilibrium just above the float32 cliff, which is why float32 runs hit it.

The fix is to never add e^-m to 1. The paper replaces the largest logit's slot with the constant 0, computes the margin from the other logits only, and returns the shifted logits minus softplus(-m). Because the target logit now enters only through differences with the other logits, automatic differentiation makes the target gradient the negated sum of the other gradients, so the books balance by construction, down to exponent underflow.

The second bug, Softmax Cancellation, appears in distillation. A student learns from a teacher's probabilities by minimizing a soft-target cross-entropy. Standard code computes the teacher's probabilities and the student's log-probabilities separately and subtracts, so when the two sets of logits are nearly equal it subtracts nearly equal numbers, and because its forward and backward formulas differ slightly it keeps a non-zero residual gradient even when the logits match. The paper's example uses teacher logits [2, 1, 0] and a student nudged by 1e-7 in float32: standard code gives [-5.96e-8, 1.45e-8, -1.12e-9] against an exact [-4.60e-8, 4.14e-8, 4.51e-9], shrinking the second entry, flipping the sign of the third, and failing to sum to zero. The fix computes the probability difference directly in log space.

The honest limit: the paper itself says the extreme margins and near-zero student-teacher gaps in its experiments are uncommon in standard large-scale training, and its grokking tests use small models. It presents the functions as easy drop-in replacements that may help whenever models make high-confidence predictions or distill from close teachers.

Goes deeper in: LLM Internals → Text Generation → From Logits to Probabilities