Jon Barron: every softmax and cross-entropy implementation has two subtle numerical bugs
jon_barron · x · 2026-10-27
Jon Barron (Google DeepMind) reports that all existing implementations of softmax and cross-entropy — JAX, PyTorch, textbooks included — harbor two subtle numerical bugs that produce false or vanishing gradients when logits nearly agree, causing "slingshots" during grokking. A simple rewrite fixes both.
Brandon Damos extends the analysis to distillation objectives:
- Forward KL / soft cross-entropy: naive q−p gradients round away differences or leave false residuals; fix by reconstructing from logits and correcting gradient-sum roundoff.
- Reverse KL: subtracting rounded log-probabilities erases small log-ratios; fix by centering logit differences under q to cancel the log-normalizer before rounding.
- Generalized JSD: log-ratio subtraction and cancellation between KL terms lose small gradients; recover τᵢ from logit differences, compute hᵢ with log1p/expm1, and center under q.
Experiments show naive gradient error plateaus below logit gaps of 1e-5 (off by 5+ orders of magnitude), while the fixed versions track a high-precision reference down to 1e-12.
More from Research
- FlexGlue: sparse-to-dense matching that beats dense accuracy at a fraction of the cost — zhenjun_zhao · 2026-10-27
- PixGR: pixel-space generative decoding for fast feed-forward 3D Gaussian reconstruction — zhenjun_zhao · 2026-10-27
- Mundo turns a few room photos into an explorable, simulatable 3D world — at 1.5 FPS — mark_k · 2026-10-27
- Mathematics as an experimental science: Marc Lelarge to talk LLM-driven math experiments at ENS — marc_lelarge · 2026-10-27
- DOE Genesis Mission Phase II: AI for particle accelerators and rare isotope separation — prescott · 2026-10-27
- Tenstorrent runs Proteina protein model at high throughput, claims no HBM shortage — gajesh · 2026-10-27