当 margin 超 5.5：Softmax 梯度静默死亡

两个 Bug 的定量画像

上个月我在调一个混合精度训练的分类模型，loss 在验证集上已经收敛了，但训练集 accuracy 还在缓慢爬升，logit margin 一路飙到 60 以上。我去翻了 PyTorch 的 cross-entropy 源码，又看了 JAX 的实现，发现所有主流框架的标准版本都存在两个静默 Bug 来源：Barron。Jonathan Barron 在 Google DeepMind 的这篇论文把这两个问题讲清楚了，也给了修复方案。

第一个 Bug 叫 Softmax Collapse。当一个模型的 logit margin 超过某个精度阈值后，目标 logit 的真实梯度会以指数级衰减，但标准实现会把这个趋近于零的值直接舍入为零 来源：Barron。结果就是梯度不再求和为零，优化器实际上在接收一个错误的信号。这个阈值取决于数值精度，bf16 下大约是 5.5，fp32 下是 16.6，fp64 下是 36.7 来源：Barron。

第二个 Bug 叫 Softmax Cancellation，出现在知识蒸馏场景。当学生网络的 logits 趋近教师网络的 logits 时，真实梯度应该趋近于零，但标准实现由于精度损失，会在非零噪声底处停滞 来源：Barron。换句话说，即使学生 logits 和教师 logits 完全相等，标准实现仍然会产生非零梯度。两个 Bug 同时存在于 softmax、log-softmax 和 cross-entropy 的所有标准实现中。

修复消除的五类现象与一个反直觉发现

修复之后，之前社区观察到但不知道根因的 slingshot 现象消失了 来源：Barron。slingshot 是指训练后期 logits 出现剧烈震荡的现象，Barron 等人指出它的成因就是 Softmax Cancellation 在梯度中引入的系统性偏差。配合 weight decay 时，grokking 过程加速了 来源：Barron。grokking 是模型在训练大量 epoch 后突然理解泛化的现象，修复后这个时间窗口缩短。

没有 weight decay 时，修复能防止 accuracy collapse 来源：Barron。我猜这里的机制是梯度被错误归零后，优化器的动量或自适应学习率仍然在更新权重，但这些更新不再对应正确的梯度方向，最终导致精度从高位掉落。在知识蒸馏中，修复阻止了学生网络从教师网络漂移，也就是 Softmax Cancellation 不再让梯度持续推动学生偏离教师 来源：Barron。

论文里最让我意外的发现是，记忆化之后 logit 增长实际上是有界的 来源：Barron。过去学界普遍认为模型完全记住训练数据后，logit 会无界增长。但修复精度问题后，实际观察到 logit 增长会收敛到一个平衡点，这个平衡点由 weight decay 的强度与 Adam 优化器的 ε 超参数共同决定。我不确定这个平衡点的理论证明是否完全严格，但实验结果相当一致。

三种精度的悬崖横向对比与修复复杂度

bf16 的悬崖只有 5.5，这意味着混合精度训练中，模型只要对一个样本达到中等置信度就会触发梯度归零 来源：Barron。对于 100 类分类任务，logit margin 到 5.5 其实相当容易，尤其是训练后期。fp32 的悬崖是 16.6，fp64 的悬崖在 36.7 来源：Barron。fp64 理论空间更充裕，但现代 GPU 训练几乎不直接用 fp64。

修复的复杂度远低于预期。Barron 等人提出的是 simple implementations，核心思路是对 softmax 的前向和反向传播做数值稳定改写，让梯度在小 margin 处不再因舍入而突变 来源：Barron。改动量小、可移植性强，理论上可以无缝替换现有框架中的对应函数。不过我没有逐行对照 PyTorch 的源码去验证改动细节，论文里的实现代码应该是最权威的参考。

从应用角度看，如果你的训练 pipeline 跑在 bf16 上，或者你的知识蒸馏训练在 loss 完全收敛后出现了奇怪的 accuracy 回落，值得检查一下是否在悬崖之上。5.5 这个阈值不算高，很多场景下已经在用了。

原文：

A Fix for Softmax and Cross-Entropy · 来源：Hacker News

作者：itfan123 值班编辑 · 发布于 2026年10月24日