模型越自信梯度越归零，Softmax藏了两个Bug

Google DeepMind 首席研究科学家 Jonathan T. Barron 在 X 上公布了一项新研究：PyTorch、JAX、TensorFlow 等主流深度学习框架普遍采用的 Softmax 和交叉熵（Cross-Entropy）标准实现里，藏着两个容易被忽视的数值精度 Bug。

它们平时很难察觉，却可能在模型训练的特定阶段悄悄改变梯度，造成损失突然飙升，甚至让模型已经学到的能力出现退化。论文已于 10 月 22 日提交至 arXiv，作者除 Barron 外，还有研究工程师与研究科学家 Dan Gnanapragasam 和 Michael Tschannen，标题为《A Fix for Softmax and Cross-Entropy》。

按 Barron 的说法，从发现问题、开展实验到撰写论文，Gemini 承担了相当多的工作。

第一个 Bug：模型越自信，梯度反而算错

这个 Bug 叫 Softmax Collapse（Softmax 坍缩）。训练分类模型时，网络为每个候选类别输出原始分数，也就是 Logit；Softmax 把分数转成概率，交叉熵再衡量预测与正确答案的差距。一张图里有只猫，模型给猫的分数远高于狗和鸟，转换后猫的概率可能达到 99.99%。

训练继续，正确类别的 Logit 与其他类别的差距越拉越大。现有实现通常先从所有 Logit 中减去最大值再做指数计算，这是经典的数值稳定性处理，能防止指数运算溢出。但减去最大值后，最大 Logit 对应的指数值变成 1，其他类别只剩很小的正数，Softmax 的分母实际变成「1 + 一个很小的数」。

浮点精度有限。当这个小数足够小，加法时会直接被舍去，结果四舍五入成 1。正确类别的预测概率被算成精确的 1，交叉熵损失变成 0，对应梯度也被算成 0；而其他错误类别的梯度仍保留微小的非零值。按交叉熵的数学定义，所有类别的 Logit 梯度之和应为 0，现在这个关系被破坏了。论文把突然失去精度的位置称为精度悬崖（Precision Cliff）。

论文给出不同浮点格式下的临界值：

• bf16：Logit Margin 约为 5.5
• fp32：Logit Margin 约为 16.6
• fp64：Logit Margin 约为 36.7

Logit Margin 衡量正确类别的 Logit 相对其他类别合计得分的优势。以 fp32 为例，差距超过约 16.6 就会撞上精度悬崖，此时对应的理论预测概率已高达约 99.9999938%。论文举了个直观例子：三个 Logit 为 [9, -9, -9]，第一个是正确答案，标准 JAX 和 Optax 实现算出的交叉熵损失为 0、正确类别梯度为 0；按修复后的方式，损失约为 3.046×10⁻⁸，正确类别梯度约为 -3.046×10⁻⁸，另两个类别梯度分别约为 1.523×10⁻⁸。

这一现象并非完全无人发现。2025 年 Prieto 等研究者已识别出相关问题，并提出 StableMax，通过替换 Softmax 中的指数函数来缓解。Barron 团队选择更直接的办法：保留原本的 Softmax 和交叉熵定义，重新组织计算过程，把最大 Logit 对应的项从指数求和中单独分离，再用数值稳定的 softplus 函数完成计算。论文同时给出了覆盖 Softmax、Log-Softmax 和交叉熵的 JAX 参考实现。

第二个 Bug：学生已经等于老师，训练却还在继续

第一个 Bug 发生在模型对答案过于自信时，第二个 Bug 叫 Softmax Cancellation（Softmax 相消误差），出现在两个模型输出非常接近的时候，典型场景是知识蒸馏。教师模型提供目标概率分布，学生模型学习这些输出；如果两者 Logit 完全相同，学生已达到目标，梯度在数学上应为 0。

但标准实现可能算出非零梯度。标准软目标交叉熵通常分别计算教师和学生的 Softmax 概率再相减得到梯度，两组概率非常接近时，浮点运算会丢失微小差异。更麻烦的是，前向传播的 Softmax 与反向传播涉及的 Log-Softmax 可能使用略有不同的运算顺序，即使 Logit 完全相同，两条路径的舍入误差也可能不同，留下一个本不该存在的梯度，还可能被 Adam 放大。

研究团队先训练一个教师模型，使其在模 53 加法任务上达到 100% 测试准确率，再直接复制教师权重初始化学生，用蒸馏温度 4、ε=10⁻⁸ 的 Adam、不使用权重衰减继续训练 2000 步。理论上学生没有任何偏离教师的必要，实验结果却是：fp64 残余梯度非常小，学生基本不变；fp32 出现短暂的 KL 散度尖峰；bf16 残余梯度更大，学生明显偏离教师，测试准确率随之下降。

修复方式是从学生与教师的 Logit 差值出发，直接在对数空间中计算所需的概率差，具体实现使用 expm1 和 log1p 两个专门处理接近零的小数值的函数，并进一步修正剩余浮点误差。这样当两者 Logit 完全相同时，梯度能精确归零。同样的蒸馏实验中，修复版本在 bf16、fp32、fp64 三种精度下都保持与教师一致，KL 散度维持在 0，测试准确率保持 100%。

精度问题如何引发训练中的「弹弓」

这个 Bug 最有意思的地方，是它可能解释一种长期被观察到的训练现象：Slingshot Mechanism（弹弓机制）。2022 年 Thilak 等研究者在研究 Grokking（「顿悟」）时观察到，训练损失有时长时间维持在接近 0 的位置，随后突然出现尖峰，再次回落，并伴随测试准确率跃升。

论文给出的解释是：模型记住训练集后，正确类别与其他类别的 Logit 差距不断扩大，逐渐接近精度悬崖；一旦越过临界值，标准交叉熵错误地把正确类别梯度归零。Adam 这类自适应优化器会根据历史梯度及其平方的移动平均值动态调整更新幅度，梯度长期接近 0 时，二阶矩估计逐渐衰减。随后权重衰减或剩余梯度带来的参数变化，可能让一部分样本重新跌回悬崖以内，被舍去的梯度突然恢复；由于 Adam 归一化梯度的分母已经很小，一次正常的梯度恢复就可能造成异常大的参数更新，损失随之飙升。

论文在模运算任务中验证了这一点。单层 Transformer 学习模 31 加法，采用 fp32、AdamW、权重衰减系数 1、ε=10⁻⁸：标准交叉熵反复出现损失尖峰，5 次随机实验平均测试准确率 94.3%±5.4%；修复版本训练更平稳，5 次实验均达到 100%。bf16 下同一任务，标准交叉熵平均 93.3%，修复版本约 100%。

另一组实验更严重：让 MLP 学习模 97 的平方和运算，不使用权重衰减，Adam 的 ε 设为 10⁻¹⁶。标准交叉熵在训练约 1250 步时触及 fp32 精度悬崖，测试准确率从 61% 骤降至 23%；修复版本没有发生这次崩塌，准确率继续提升至约 85%。论文还在另外几种 Grokking 任务上观察到类似现象，相应实验条件下修复后的 fp32 交叉熵表现接近标准 fp64 实现。

论文也指出，在使用 LayerNorm 且没有权重衰减的某些网络中，Slingshot 带来的参数范数增长可能意外起到降低有效学习率的作用，因此过去部分实验中看似有利的 Slingshot，并不意味着这种数值问题本身对模型有益。

Adam 的默认参数，恰好踩在临界点附近

Adam 优化器有个常见超参数 ε，通常设为 10⁻⁸，加在梯度归一化的分母中控制数值行为。论文发现，在模型已记住训练集、并使用权重衰减的特定条件下，Logit Margin 的增长受权重衰减与 Adam 更新之间平衡关系的限制，平衡值近似满足 m ≈ ln(1/ε)。ε=10⁻⁸ 时这个值约为 18.4，而 fp32 的精度悬崖恰好位于 16.6——模型在接近自然形成的平衡位置之前，就可能先触发标准交叉熵的数值问题。这也为「Adam 的 ε 太小可能导致不稳定」的观察提供了新解释。在相应实验设置中，把 ε 提高到 10⁻⁷ 或更大，可以让平衡位置落在悬崖以内，避免这类尖峰；但调整 ε 只是改变优化器行为，修复交叉熵的计算方式才能直接消除这处数值错误。

Gemini 参与发现 Bug，还承担了大量实验和论文写作

按 Barron 的介绍，Softmax Cancellation 最早由 Gemini Flash 3.7 发现并给出修复方案，论文中的大部分实验由 Gemini Flash 3.8 完成，写作阶段又使用了 Gemini 4 Argon。Barron 表示，Argon 完成了论文及图表的初稿，随后他进行了严格审查和反复修改，直到技术内容足够扎实，但对最终的文字表达仍不太满意。AI 已参与到数值问题发现、修复方案设计、实验执行和论文撰写等多个环节，研究人员则持续承担技术核查与质量把关。

回到技术结果，论文主要通过模运算 Grokking 任务和特定知识蒸馏实验验证修复效果。研究团队也承认，实验中出现的极大 Logit Margin，以及学生与教师几乎完全相同的输出，在常规大规模训练中并不常见。因此目前还不能认定主流大模型训练普遍受这两个问题严重影响，更无法直接推断修复后就能让大模型训练整体提速。

但研究确实揭示了一个容易被忽视的事实：即使数学表达式完全正确，计算机实际执行的浮点运算，也可能在特定条件下改变梯度的性质。论文给出的修复方案保留了原有函数的数学定义，不需要提高浮点精度，作者认为这些函数有望直接替换现有实现且不会带来明显性能开销，其在大规模模型训练中的实际效果仍需进一步验证。