谷歌研究员称 JAX/PyTorch 的 Softmax 与交叉熵实现全有 bug
Google 研究员 Jon Barron 等人于 10 月 24 日发布 arXiv 论文《A Fix for Softmax and Cross-Entropy》，宣称所有现存的 softmax 与交叉熵实现——包括 JAX、PyTorch 以及教科书上的标准写法——都存在两个微妙的数值 bug，一个简单的重写即可同时修复，且修复效果在 grokking 训练动态上得到验证。
已确认
- 第一个 bug 与 grokking 过程中的「slingshot」现象相关：Barron 展示训练曲线对比，标准交叉熵的训练损失在 grokking 阶段反复震荡，而修复版平稳收敛；修复后不仅能消除 slingshot 尖峰，还能加速 grokking。
- 第二个 bug 被称为「Softmax Cancellation」：当模型输出的 logits 非常接近软目标 logits 时（例如蒸馏场景中 student 初始化为 teacher），现有实现即使 student 与 teacher 完全一致仍会产生非零梯度。
- Barron 还补充，Adam 优化器默认的 eps=1e-8 恰好踩中 bug 相关的阈值，这一巧合值得注意。
为什么重要
- softmax 与交叉熵是深度学习最基础、使用最广泛的算子之一，若实现层面普遍存在数值缺陷，影响面覆盖几乎所有训练与蒸馏流程。
- 修复可能改变学界对 grokking、slingshot 等现象的解释——这些「神秘」动态或部分源于数值误差而非算法本质，对相关研究结论的可信度有直接影响。
2026-10-24 ~ 2026-10-24 · 6 条相关
一手来源
- Jon Barron：JAX、PyTorch 的 softmax 与交叉熵实现都藏着两个数值 bug — jon_barron ·
- Softmax 两大 bug 之二：蒸馏时 student=teacher 仍产生非零梯度 — jon_barron ·
- 修复交叉熵可消除 grokking 尖峰,Adam eps=1e-8 恰好踩中 bug 阈值 — jon_barron ·
- 【源头】Jon Barron：JAX、PyTorch 的 softmax 与交叉熵实现都藏着两个数值 bug — jon_barron · 2026-10-24
- 【源头】修复交叉熵可消除 grokking 尖峰,Adam eps=1e-8 恰好踩中 bug 阈值 — jon_barron · 2026-10-24
- 【源头】Softmax 两大 bug 之二：蒸馏时 student=teacher 仍产生非零梯度 — jon_barron · 2026-10-24
- 谷歌研究员发现 Softmax 与交叉熵两个数值 bug，修复后消除 Slingshot 现象 — jon_barron · 2026-10-24
另有 2 条近重复转述:jon_barron · ZeYanjie