专题 · FULL STORY

softmax 与交叉熵的两大数值 bug:从揭露到论文

Google DeepMind 研究员 Barron 等人揭露主流框架中 softmax 与交叉熵实现里的两个数值 bug,引发社区热议,随后团队发布 arXiv 论文《A Fix for Softmax and Cross-Entropy》系统阐述问题与修复方案。

2026-10-24 ~ 2026-10-28 · 2 集 · 24 条

第 1 集 · Barron 揭 softmax 与交叉熵实现两大数值 bug(2026-10-24,21 条)

Google/DeepMind 研究员 Jonathan T. Barron、Dan Gnanapragasam 与 Michael Tschannen 于 10 月 24 日发布 arXiv 论文《A Fix for Softmax and Cross-Entropy》，指出包括 JAX、PyTorch 以及教科书标准写法在内的所有现存 softmax、log-softmax 与交叉熵实现都存在两个微妙的数值 bug，一个简单的重写即可同时修复。修复效果在 grokking 训练动态上得到验证，且 FlashAttention 的反向传播同样中招。消息由 Barron 本人连续发帖公布，经 ZeYanjie、codestar、IgorCarron、khademinori、chrisemoody 等人转发扩散，burkov、lmoroney、JFPuget、kalyankpl 亦做了转述。

已确认

  • 第一个 bug 与 grokking 过程中的「slingshot」现象相关。据 lmoroney 转述，该 bug 被称为「Softmax Collapse」，在某个 logit 远大于其他 logit 时触发；burkov 补充其对高置信度的正确预测会把梯度完全抹掉，而其他类别却残留小梯度。Barron 展示训练曲线对比：标准交叉熵的训练损失在 grokking 阶段反复震荡，而修复版平稳收敛；修复不仅消除 slingshot 尖峰，还能加速 grokking。
  • 第二个 bug 被称为「Softmax Cancellation」：当模型输出的 logits 非常接近软目标 logits 时（例如蒸馏场景中 student 初始化为 teacher），现有实现即使 student 与 teacher 完全一致仍会产生非零梯度。
  • Barron 补充，Adam 优化器默认的 eps=1e-8 恰好踩中 bug 相关的阈值，这一巧合值得注意。
  • 据转发信息，FlashAttention 的反向传播同样受该 bug 影响，波及面不止于直接的 softmax/交叉熵调用。
  • Barron 同时公开了论文的 AI 协作工作流：bug 最早由 Gemini Flash 3.7 发现并修复，论文中大部分实验由 Flash 3.8 完成，Gemini 4 Argon 在喂入其旧论文学习文风后写出初稿，最后由他本人人工重写。

为什么重要

  • softmax 与交叉熵是深度学习最基础、使用最广泛的算子之一，若实现层面普遍存在数值缺陷，影响面覆盖几乎所有训练与蒸馏流程，FlashAttention 反向传播中招进一步扩大波及范围。
  • 修复可能改变学界对 grokking、slingshot 等现象的解释——这些「神秘」动态或部分源于数值误差而非算法本质，对相关研究结论的可信度有直接影响。

还有 1 条相关讨论 →

第 2 集 · Google 论文揭秘 softmax 与交叉熵实现中的两个数值 bug(2026-10-27,3 条)

Google DeepMind 研究员 Jon Barron 等人发布 arXiv 论文《A Fix for Softmax and Cross-Entropy》，指出包括 JAX、PyTorch 乃至教科书在内的所有现有 softmax 与 cross-entropy 标准实现都藏着两个数值 bug，并引发关于形式验证的讨论。随后学者 Brandon Damos 补充称，这一 "近一致态" 的数值问题并不限于蒸馏中的软交叉熵损失，凡是使用 KL 散度的场景（如强化学习中的 KL 正则项）同样可能受到影响。