softmax 藏两个数值 Bug

2026 年 10 月 22 日，Google DeepMind 研究员 Jonathan Barron、Dan Gnanapragasam 和 Michael Tschannen 提交 arXiv 论文《A Fix for Softmax and Cross-Entropy》（编号 2610.31642）：PyTorch、JAX、TensorFlow 与教科书通用的标准 softmax、log-softmax、交叉熵实现带着两个浮点 Bug。Softmax Collapse 在目标 logit 独大时把它的梯度舍入成零；Softmax Cancellation 让蒸馏里学生与教师 logits 完全相等时梯度仍残留非零值。作者项目页附上修复用的 JAX 参考实现和交互式演示。

要点

论文 arXiv:2610.31642（v1，2026-10-22 提交，cs.LG）由三位 Google DeepMind 研究员署名，主张标准 softmax、log-softmax、交叉熵实现存在两个数值 Bug。
按作者项目页，Softmax Collapse 的精度悬崖在 margin ≈5.5（bf16）、≈16.6（fp32）、≈36.7（fp64）：目标 logit 领先越过悬崖后，它的梯度被舍入为零，梯度不再求和为零。
项目页内嵌两段 JAX 实现（硬目标版约 15 行、软目标版约 12 行）；一个独立 NumPy 复现证实了两个 Bug 的量级，同时测出 fp32 悬崖随归约顺序在 15.8–16.7 之间浮动。

背景

softmax 加交叉熵是所有分类与蒸馏训练的地基，从教材到框架教的都是同一套写法：logits 先减去最大值防上溢，再做指数与归一化。这套惯例挡住了上溢；论文盯上的是减完之后的舍入行为。目标 logit 把其余 logits 甩开一个 margin m 之后，分母里的 1+e⁻ᵐ 在浮点下舍入成 1.0，目标梯度 qₜ−1 跟着归零。项目页把这条失真路径与已知的 Slingshot loss 突刺（由 Adam/AdamW 驱动）和 grokking 现象连在一起。数值精度与训练互相牵制的话题本月已两次出现在本站：BF16 张量核心算力解锁与 Humming 量化内核，处理的都是同一层的细节。

两个 Bug 的机制

两个 Bug 的机制在项目页里写得最细。Softmax Collapse：目标 logit 独大时，分母 1+e⁻ᵐ 舍入为 1.0，目标梯度 gₜ = qₜ−1 舍入为 0.0，其余 K−1 个非目标梯度仍为正，梯度整体不再求和为零。Softmax Cancellation：蒸馏与 KL 散度里学生、教师各算一遍 softmax 再相减，两组 logits 几乎相等时，相近概率相减丢失精度，即使两者完全相等也残留非零梯度（项目页口径）。

修复走对数空间：把目标（argmax）从 log-sum-exp 里摘出来得到 margin m，log-softmax 写成 ẑ_k − softplus(−m)；目标梯度改用 −e⁻ᵐ/(1+e⁻ᵐ)，避开「先加 1 再减」的写法；软目标一侧用 expm1 与 log1p 从 z_S−z_T 直接算 q−p（项目页公式 9）。项目页内嵌的两段 JAX 实现分别约 15 行与 12 行，另附交互式 demo.html。

悬崖与复现

悬崖位置是全部论证里最可核对的部分，也是两个来源出现分歧的地方。项目页口径如下（margin m 指目标 logit 与其余 logits 的对数和之差）：

精度	悬崖（项目页）	NumPy 实测（dev.to）
bf16	≈5.5	未单独实测
fp16	页面未给	7.7（理论 7.62）
fp32	≈16.6	16.0，随类别数在 15.8–16.7 浮动
fp64	≈36.7	36.1（理论 36.74）

frank chu 10 月 23 日在 dev.to 发出的复现不走框架，用 NumPy 1.26.4 直接算解析梯度 q−p，也就是框架反向传播算的那一项。复现结果：悬崖之后 loss 与目标梯度恰好为 0.0；悬崖之前已有量化损伤，m=15 时标准损失读数 2.38e−7、真值 3.06e−7，偏差约 22%；蒸馏一侧有约 1e−8 到 5e−8 的梯度误差地板，δ=1e−7 时占梯度范数 74%，论文的修复版与 fp64 真值只差 9.2e−16。

与论文的出入有两处。悬崖位置：他从 K=2 扫到 K=101，fp32 悬崖在 15.8–16.7 之间随归约顺序移动，16.6 是两类问题的理想化数字。一个子声明未复现：论文称学生与教师 logits 完全相等时梯度仍非零，他的 NumPy 环境里两条路径用同样的算子，结果恰好为 0.0；他把差异归给框架 kernel 的实现方式，同时说明那个误差地板与 kernel 无关。他也没有重训 grokking 模型，slingshot 与 Adam 的关联仍依赖论文自己的训练曲线。

各方说法

dev.to 的复现站在「Bug 属实」一侧：量级与论文一致，但悬崖位置测的是 kernel 的归约顺序，别把 16.6 当常数写进断言。

SamCodeMan 的批评笔记（Samir Sengupta，10 月 23 日）把质疑停在摘要层面：v1 摘要没有给出触发阈值、伪梯度的大小、fp32/bf16/fp16 下哪些框架中招，蒸馏漂移幅度、grokking 加速幅度、Adam ε 的公式与测试值一概空缺，也没有独立代码包和生产开销数据。他同时认为 Bug 描述具体到可以立刻对着现有训练栈写测试，并给出五步清单：单元测试蒸馏损失，相同 logits 时断言梯度恰为零；测一行单主导 logit，查目标梯度是否为零、梯度求和是否为零；任一失败就换论文实现并重跑两个测试；重查建立在 Slingshot 类行为、无 weight decay 精度崩塌、蒸馏学生漂移上的结论；把 Adam ε 连同 weight decay 视为共同封顶记忆后 logit 增长的设置。

两边的分歧在证据强度：复现方证明算术层面成立，批评方指出训练层面的收益仍是作者的说法。

怎么试

最短路径在作者项目页 jonbarron.info/softmaxfix：页内嵌两段 JAX 实现和交互式 demo，可直接抄进自己的损失函数。动手前先跑两个单元测试（Sengupta 清单前两条）：学生与教师 logits 完全相同时，蒸馏损失梯度应恰为零；一行里放一个独大的 logit，目标梯度应为零、梯度求和也应为零。哪个测试失败，再考虑换实现重跑。

编辑判断

自己写训练循环、做蒸馏、或训练过度自信分类器（大 margin 场景）的团队最该花这十分钟。bf16 的悬崖在 ≈5.5，比 fp32 的 ≈16.6 低得多，而作者给出的平衡点 m*(ε) ≈ log(1/ε) 意味着记忆完成后的 margin 本来就会往上爬，ε 取常见的 10⁻⁸ 时平衡点约 18.4，已经越过 fp32 悬崖。

诚实的缺口：这是 v1 论文，没有独立代码仓库，生产规模收益没有第三方验证；悬崖数字随归约顺序浮动，直接抄进断言会误判；一个子声明在纯 NumPy 下未复现。各框架 kernel 的实际行为、修复的运行开销（论文称 drop-in 替换、无可感知减速，作者口径）都还欠第三方测试。