PyTorch、JAX都被点名！Softmax曝出隐藏Bug，Gemini立功
2026.10.25 · 事件 · 论文 · 论文预印本
主流框架的 Softmax 与交叉熵标准实现存在两处数值精度缺陷，可致梯度算错。
Google DeepMind 首席研究科学家 Jonathan T. Barron 公布新研究，指出 PyTorch、JAX、TensorFlow 等主流框架普遍采用的 Softmax 与交叉熵标准实现存在两个数值精度 Bug，可能在训练特定阶段悄悄改变梯度、造成损失飙升。论文 10 月 22 日提交至 arXiv，提出不提高浮点精度、不重设损失函数的修复方式。
- 两个 Bug 分别是 Softmax Collapse（Softmax 坍缩）与 Softmax Cancellation（Softmax 相消误差）。
- 精度悬崖临界值：bf16 约 5.5、fp32 约 16.6、fp64 约 36.7（Logit Margin）。
- 修复思路是保留原定义、重组计算，用 softplus 处理最大 Logit 项，并给出 JAX 参考实现。
- 论文用模 31 加法、模 97 平方和等实验验证，修复版训练更平稳、测试准确率更高。
- Adam 默认 ε=10⁻⁸ 时平衡值约 18.4，落在 fp32 精度悬崖 16.6 之外，可能先触发数值问题。
为什么值得看
两个最基础的函数可能长期影响训练稳定性，值得开发者关注。
适合谁关注
深度学习框架开发者、模型训练工程师、AI 研究者。