A Fix for Softmax and Cross-Entropy
This presentation examines two numerical pathologies in standard softmax and cross-entropy implementations that can alter training dynamics under adaptive optimization. We explore how precision-dependent gradient cliffs cause apparent loss instabilities and parameter drift, and present algebraic fixes that preserve exact gradient structure in low-precision arithmetic. The talk demonstrates these effects on grokking benchmarks and distillation tasks, revealing an optimizer equilibrium that controls margin growth.Script
Standard softmax implementations hide a precision trap. When one logit dominates the others by a large enough margin, the gradient of the target class vanishes completely, even though gradients for non-target classes remain positive. This violates a fundamental property: the gradients should sum to exactly zero.
The precision cliff appears at different margins depending on your datatype: around 5.5 in bf16, 16.6 in fp32, and 36.7 in fp64. Once the margin crosses these thresholds, the standard computation rounds the target probability to exactly one, zeroing its gradient while leaving a residual that breaks the zero-sum conservation law.
This discontinuity produces the Slingshot phenomenon in grokking experiments. After memorization drives margins past the cliff, target gradients disappear and Adam's second moments decay. Weight decay then pulls some margins back below the threshold, suddenly restoring large gradients that get divided by the decayed adaptive denominator, triggering a loss spike and accuracy drop.
The second pathology, Softmax Cancellation, appears during distillation. When student and teacher logits are nearly identical, computing their probabilities separately and then subtracting them discards the information in the small logit difference. The corrected approach works entirely in log space, using expm1 and log1p to preserve relative accuracy and projecting out the residual to guarantee an exact zero sum.
The fixes eliminate precision-dependent failures completely. With corrected cross-entropy, bf16, fp32, and fp64 all reach 100 percent test accuracy on modular arithmetic tasks, crossing their respective precision cliffs smoothly and reaching margins predicted by the optimizer equilibrium law: approximately log of 1 over epsilon.
These aren't just cosmetic rounding errors. They reshape training dynamics, break distillation near convergence, and create optimizer equilibria that depend on your floating-point format. The algebraic corrections restore the exact geometry of the loss landscape. You can explore the full derivations and experiments at EmergentMind.com, where you can also create video presentations of research that interests you.
