Jon Barron: all softmax and cross-entropy implementations (JAX, PyTorch, textbooks) harbor two subtle bugs
agarwl_ · x · 2026-10-24
Google researcher Jon Barron shares a new result: every existing implementation of softmax and cross-entropy — in JAX, PyTorch, and even textbooks — contains two subtle bugs. A simple rewrite fixes both and eliminates 'slingshots' during grokking.
Retweeter petrenkoai notes his team spotted this softmax 'bug' 1.5 years ago in a paper, but it's hard to construct training scenarios where it matters since gradients on highly likely tokens are tiny anyway. Still, he recommends applying the fix — one less thing to worry about.
More from Research
- Mistral's Timestral tops time-series benchmarks with 209M params, built mostly by agents — hkproj · 2026-10-24
- Svante Pääbo's team uses AlphaGenome as a time machine to reconstruct human gene regulatory evolution — gleech · 2026-10-24
- KAN author Ziming Liu: LLMs learn alchemy, not the science of AI — ZimingLiu11 · 2026-10-24
- Training-Physics Intuition Compressed Into a Decision Tree Hits 78.2% Out-of-Fold — ZimingLiu11 · 2026-10-24
- LLMs Rank Optimizers at 85% but Are Data-Blind on Architectures, New Benchmark Shows — ZimingLiu11 · 2026-10-24
- NVIDIA releases 43K instruction-hierarchy rubrics for judging system vs user prompts — kalyan_kpl · 2026-10-24