A Fix for Softmax and Cross-Entropy

论文 A Fix for Softmax and Cross-Entropy
作者 Jonathan T. Barron, Dan Gnanapragasam, Michael Tschannen
版本 arXiv v3，2026 年 10 月 29 日
项目与实现 Softmax Fix

减去最大值之后，softmax 还会有什么问题？

通常为了避免指数溢出，会先把 logits 减去最大值，再计算 softmax。这个处理很常见，但论文指出，当模型极度确信某个答案，或者学生和教师的输出已经非常接近时，仍然可能出现精度损失，并影响反向传播。

文章主要分析 Softmax Collapse 和 Softmax Cancellation 两种情况。前者发生在概率接近 1 的区域，后者涉及两个很接近的概率相减。

对正确类别 $t$，普通交叉熵的梯度为

\[g_t=q_t-1,\qquad g_j=q_j\quad(j\ne t).\]

当 $q_t$ 非常接近 1，浮点数可能直接把它舍入为 1，于是 $g_t$ 变成 0。但是其他类别的梯度仍然是小的正数，梯度之和也就偏离了数学上应该满足的零和性质。

例如论文给出 fp32 下的 logits $[9,-9,-9]$。交叉熵实际约为 $3.046\times10^{-8}$，标准实现却可能返回 0，目标类别的梯度也变成 0，另外两个梯度仍约为 $1.523\times10^{-8}$。

把正确类别相对于其余类别的 margin 写成

\[m=z_t-\log\sum_{j\ne t}\exp(z_j),\]

就可以将损失和目标梯度改写成

\[\mathcal{L}=\operatorname{softplus}(-m),\qquad g_t=-\frac{\exp(-m)}{1+\exp(-m)}.\]

实现时结合 log1p 等运算保留小量，并通过非目标类别梯度之和计算目标梯度，避免先把概率舍入成 1 再减去 1。数学上等价的表达式，在有限精度下可能有不同的行为。

对于 soft target，梯度涉及学生概率和教师概率的差。如果两者几乎相同，分别算出两个 softmax 再相减，就可能丢失差值的信息。前向和反向采用不同数值路径时，甚至可能在师生 logits 完全一致的情况下留下残余梯度。

这种微小误差经过 Adam 等优化器后，可能造成持续漂移。作者的修复从 logits 的差异出发，使用 expm1、log1p 等稳定计算方式，并让梯度满足应有的零和约束。

文章用了一个很直接的检查。把教师参数完整复制给学生，再让学生蒸馏教师，理想情况应该一直保持不变。标准实现中可以观察到漂移，修复后则在测试的精度配置下保持零 KL 和一致预测。

在模加法的 grokking 实验中，部分训练会出现 loss 长时间停滞后突然尖峰，也就是 slingshot 现象。数值修复能够消除论文设置中的这类不稳定。表 1 的一个 fp32 配置下，标准实现准确率为 $94.3\%\pm5.4\%$，修复后为 $100\%\pm0\%$。

这说明训练曲线中的某些现象可能来自实现的数值性质。解释优化行为之前，需要先检查梯度是否真的符合理论公式。

作者也指出，极大的 margin 和极小的师生差异，在常规大规模训练中未必经常出现。实际收益需要结合具体模型、精度、优化器和 kernel 实现测试。

感觉这篇文章的问题非常具体，几个简单例子就能看出错误发生在哪里。平时遇到 loss spike，容易先去改学习率、梯度裁剪和正则化，这篇提醒我也要检查最基础的运算实现。

尤其喜欢“复制教师给学生之后还会不会变化”这个实验。它有一个很明确的预期结果，可以直接暴露数值路径中的问题。以后写训练或蒸馏代码，除了检查 loss 下降，也值得检查这种应该保持不动的情况，以及梯度和是否接近零。

真正集成到训练系统时，还需要验证 fused kernel、自定义 backward 和混合精度是否保留这些性质，并测量吞吐开销。论文里的公式写对了，实现里的每一步仍然需要单独确认。