The fixed softmax still needs something to stop the logits

Softmax turns logits—unbounded class scores—into probabilities. Cross-entropy then penalizes the negative log-probability of the target class. The standard max-shift implementation prevents exponential overflow, but Barron, Gnanapragasam, and Tschannen show that it still breaks a conservation law: once the winning logit is confident enough, its gradient becomes exactly zero while the losing gradients remain positive.

Let the target logit exceed the log-sum-exp of all other logits by margin m. The target gradient is −exp(−m)/(1+exp(−m)). Standard code instead forms the target probability near one and subtracts one. In fp32, values smaller than roughly 10⁻⁷ cannot change 1.0, so that subtraction returns zero around m = 16.6.

At the paper's three-class example [9, −9, −9], the standard fp32 target gradient is 0.0, but the two non-target entries are still 1.523e-8 each. The true loss is not zero; it is 3.046e-8.

def fixed_gradient(margin, non_targets):
    tail = exp(-margin)
    each = tail / (non_targets * (1.0 + tail))
    others = [each] * non_targets
    target = -sum(others)
    return [target] + others

The running TypeScript uses exactly this rearrangement. Python and C++ are faithful translations: compute the tiny tail first, then make the target gradient the negative sum of all non-target gradients. No higher-precision cast is required.

Adam divides a moving average of the gradient by a moving scale plus epsilon. When collapsed gradients make that scale decay, a later non-zero gradient is divided by a tiny denominator and produces the sudden loss spike called Slingshot. Correct cross-entropy makes the gradient decay smoothly instead of switching off at a format boundary.

At a stationary AdamW point, the normalized scalar step is |g|/(|g|+epsilon). After memorization, |g| is proportional to exp(−m). With weight decay opposing the step, the balance occurs around m* ≈ log(1/epsilon). The learning rate and moment coefficients cancel from this stationary approximation.

On modular addition, the fixed loss reaches the epsilon-predicted margin across bf16, fp32, and fp64. Standard cross-entropy stalls near each format's cliff. The table below reproduces selected means from the paper's five-seed Table 1.

Knowledge distillation trains a student against a teacher's probability distribution. Computing both softmaxes separately and subtracting nearly equal probabilities loses significant digits. The paper instead works with the student–teacher logit difference using expm1 and log1p, then explicitly removes residual zero-sum drift. Starting student and teacher from identical weights, the corrected loss stays at zero KL and 100% accuracy across all tested precisions; standard bf16 drifts.

The paper is deliberately strongest in extreme regimes: very large margins and nearly identical teacher–student logits, which the authors note are uncommon in ordinary large-scale training. Next, benchmark wall-clock cost in fused GPU kernels, test label smoothing and huge vocabularies, and separate improvements caused by restored arithmetic from those caused by the changed long-run margin trajectory.