A Fix for Softmax and Cross-Entropy
Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen
cs.LG
2026-10-23
DeepMind 重写 softmax 与交叉熵的实现,修掉两个浮点 bug:Slingshot 尖峰消失,三种精度下 grokking 测试精度达 99–100%。
JAX、PyTorch、TensorFlow 里的 softmax 和交叉熵,为防指数溢出都会先把 logits 减去最大值。这一步防住了溢出,却留下两个舍入 bug,而且都出在梯度上。第一个叫 Softmax Collapse:模型对正确类别已经很有信心、目标 logit 远大于其余时,标准实现会把目标 logit 的梯度舍入成 0,而其余 K−1 个梯度照常为正,梯度向量的加和不再为零。第二个叫 Softmax Cancellation:软目标蒸馏里学生 logits 已经贴住教师时,标准实现仍给出非零梯度,推着学生离开教师。
这不只是精度洁癖。grokking(网络先记住训练集、很久之后才泛化)实验里反复出现的 Slingshot 现象,训练损失在近零平台和突发尖峰之间来回,一直被当作自适应优化器的固有怪癖。这篇证明它就是 Softmax Collapse 在 Adam 下发作的样子。此前的工作用代理损失、梯度投影或强制 fp64 绕,这篇直接把函数修对。
两个 bug 同源:把很小的数加到 1.0 上。浮点数在 0 附近精度极高,fp32 能表示到约 10⁻³⁸,但加到 1.0 上,小于约 10⁻⁷ 的部分会被舍入吃掉。标准实现 shift 之后目标项是 e⁰=1,分母是 1+e⁻ᵐ,其中 m 是 logit margin,即目标 logit 领先其余 logit 的 log-sum-exp 量。margin 一旦越过精度悬崖(bf16 约 5.5、fp32 约 16.6、fp64 约 36.7),e⁻ᵐ 加进 1.0 就消失,损失和目标梯度双双归零;非目标梯度用负差值直接算,不经过加 1,能一直活到指数下溢(m≈88)。
修法是让 e⁻ᵐ 永远不站在被加到 1.0 的位置:
每个修法十几行 JAX,内部不升高精度,论文称没有明显减速。
主实验是 mod 加法 grokking(1 层 width-64 Transformer,a+b mod 31,50% 训练集,AdamW,λ=1),覆盖三种精度、三档 ε:
| 设置 | 标准 CE | 修复版 CE |
| fp32,ε=10⁻⁸ | 94.3±5.4% | 100.0±0.0% |
| bf16,ε=10⁻⁸ | 93.3±7.4% | 100.0±0.1% |
| fp64,ε=10⁻¹⁶ | 97.2±3.2% | 99–100% |
9 组设置里,修复版全部把 margin 送到预测平衡点 log(1/ε),测试精度 99–100%;标准版的 margin 全部停在各自精度悬崖附近。ε=10⁻¹⁶ 时平衡点 36.8 恰好越过 fp64 悬崖 36.7,连 fp64 标准实现都会撞上。无 weight decay 时更直接:a²+b² mod 97 上,标准 CE 一路跟到 61%,约第 1250 步撞悬崖,一个尖峰把精度打到 23%(a³+b mod 97 是 60% 掉到 21%),修复版平滑走到 85%。
对比两个近期修补:Zero-sum Proj(把梯度投影回零和)在 ε≤10⁻¹² 时精度跌破 50%;StableMax(把 exp 换成分段有理函数)的 margin 爆到 10³–10⁴,精度全面更低。蒸馏这边,学生直接用教师权重初始化,真实梯度应为 0,标准软目标 CE 的残余梯度在 bf16 下约 10⁻³,超过 ε=10⁻⁸ 就被 Adam 归一化成正常步长的更新,学生被推离教师;修复版 KL 恒为 0,精度保持 100%。
这三个函数每一步分类训练都在跑。修好的版本是 drop-in 替换,十几行代码、不升精度,框架层可以直接采纳。对用 bf16 做蒸馏的团队,10⁻³ 量级的伪梯度是量化过的实害;对 grokking 和优化器稳定性研究,Slingshot 有了完整因果链:目标梯度归零,Adam 二阶矩随之衰减,梯度恢复后被极小的分母放大成大步长。顺带一个洞察:margin 的平衡点 m≈log(1/ε),默认 ε=10⁻⁸ 对应 18.4,恰好压在 fp32 悬崖 16.6 上方,这解释了 Slingshot 为什么在 fp32 里常见;ε≥10⁻⁷ 时平衡点降到悬崖以下,标准实现也不会尖峰。ε 与 weight decay 组合起来,效果像一个权重为 ελ 的隐式 L2 正则。
对常规大规模预训练,直接影响有限,margin 很少被推过 16。但小数据记忆、长尾类别、蒸馏到收敛这类高置信场景,这是白捡的保险。
作者在结论里自己承认,极端 margin 和近乎重合的师生 logits 在常规大规模训练里不常见。实验也全部小规模:模运算 grokking、学生=教师初始化的玩具蒸馏,没有真实教师-学生的 LLM 蒸馏,没有大规模训练验证,「没有明显减速」只有定性说法,没有计时数字。
论文同时展示,有 LayerNorm、无 weight decay 时,Slingshot 尖峰撑大权重范数,像一次意外的学习率衰减,可能有益。消除它并非在所有设置下都更好,收益在于行为可预测。蒸馏危害的证据也偏构造性:专挑学生=教师的极端设置来暴露噪声,真实蒸馏里师生差距大,伪梯度占比多少没有测。