研究指出 PyTorch、JAX、TensorFlow 等框架的 Softmax 與交叉熵標準實作存在兩項數值缺陷

深度学习中最基础、使用最广泛的两个函数，竟然藏着两个隐蔽的 Bug。

昨日，Google DeepMind 首席研究科学家 Jonathan T. Barron 在 X 上公布了一项新研究。

他指出，包括 PyTorch、JAX、TensorFlow 在内的主流深度学习框架，普遍采用的 Softmax 和交叉熵（Cross-Entropy）标准实现中，存在两个容易被忽视的数值精度 Bug。

这些 bug 平时很难察觉，却可能在模型训练的特定阶段悄悄改变梯度，造成损失突然飙升，甚至让模型已经学到的能力出现退化。

更有意思的是，研究团队发现，只需调整计算方式，就能修复这两个 bug，无须提高浮点精度，也不需要重新设计损失函数。

相关论文已于 10 月 22 日提交至 arXiv，除了 Jonathan T. Barron，还有另外两位研究工程师与研究科学家 Dan Gnanapragasam 和 Michael Tschannen。

• 论文标题：A Fix for Softmax and Cross-Entropy

• 论文地址：https://arxiv.org/pdf/2610.31642

并且按照 Barron 的说法，从发现问题、开展实验到撰写论文，Gemini 在这项研究中承担了相当多的工作。

模型越自信

梯度反而算错了

第一个 bug 叫做 Softmax Collapse（Softmax 坍缩）。

要理解它，先得看看 Softmax 在做什么。

训练分类模型时，神经网络通常会为每个候选类别输出一个原始分数，也就是 Logit。Softmax 负责将这些分数转换成概率，再由交叉熵衡量模型的预测与正确答案之间的差距。比如，一张图片中有一只猫，模型给猫的分数远高于狗和鸟。经过 Softmax 转换，猫的预测概率可能达到 99.99%。

随着训练继续，模型对正确答案越来越有把握，正确类别的 Logit 与其他类别之间的差距也越来越大。问题就出在这里。

现有实现通常会先从所有 Logit 中减去最大值，再进行指数计算。这是一种经典的数值稳定性处理，可以防止指数运算发生溢出。但减去最大值后，最大 Logit 对应的指数值变成了 1，其他类别则是一些很小的正数。在正确类别分数最高的情况下，Softmax 的分母实际上变成了：

1 + 一个很小的数。

浮点数的精度有限。当这个小数足够小时，计算机在执行加法时就会直接将其舍去，结果被四舍五入成 1。于是，正确类别的预测概率被计算为精确的 1，交叉熵损失变成 0，对应梯度也被算成 0。然而，其他错误类别的梯度仍然保留着微小的非零值。

按照交叉熵的数学定义，所有类别对应的 Logit 梯度之和应该为 0。现在，正确类别的梯度提前归零，其他类别的梯度却还存在，原本应当成立的关系被破坏。论文将这种突然失去精度的位置称为精度悬崖（Precision Cliff）。

标准 Softmax 和交叉熵中的两个数值计算缺陷。（a）Softmax Collapse（Softmax 坍缩），（b）Softmax Cancellation（Softmax 相消误差）。

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

• 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 函数完成计算。对于单一正确标签的交叉熵，修复后的公式还让自动微分直接得到：正确类别的梯度等于其他类别梯度之和的相反数。这样，极小的梯度就能继续保留下来，直到浮点数确实无法表示它为止。

论文同时给出了 JAX 参考实现，覆盖 Softmax、Log-Softmax 和交叉熵。

一个小小的精度问题

如何引发训练中的「弹弓」

这个 Bug 最有意思的地方，是它可能解释一种长期被观察到的神经网络训练现象：Slingshot Mechanism（弹弓机制）。

2022 年，Thilak 等研究者在研究 Grokking 时观察到了这种现象。Grokking 通常被翻译为「顿悟」：模型很早就记住了训练数据，却要经过相当长时间的训练，才能突然在测试数据上获得较好的泛化能力。在这一过程中，研究者发现训练损失有时会长时间维持在接近 0 的位置，随后突然出现尖峰，再次回落。这样的变化可能反复发生，并伴随测试准确率的跃升。

为什么模型明明已经记住了训练数据，损失还会突然暴涨？论文给出了一种具体解释。

当模型记住训练集后，正确类别与其他类别的 Logit 差距不断扩大，逐渐接近精度悬崖。一旦越过临界值，标准交叉熵就会错误地将正确类别的梯度归零。此时，Adam 这类自适应优化器也会受到影响。Adam 会根据历史梯度及其平方的移动平均值，动态调整参数更新的幅度。

当训练梯度长期接近 0 时，优化器记录的梯度二阶矩估计也会逐渐衰减。随后，权重衰减或剩余梯度带来的参数变化，可能让一部分样本重新跌回精度悬崖以内。这些样本原本被舍去的梯度突然恢复。由于 Adam 用来归一化梯度的分母已经变得很小，一次正常的梯度恢复就可能造成异常大的参数更新。训练损失随之突然飙升。

这正是论文提出的 Slingshot 形成机制。

论文在模运算任务中验证了这一点。在一个使用单层 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 实现。不过，Slingshot 与泛化能力之间的关系仍需要结合具体训练配置理解。

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

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

研究团队还解释了另一个值得注意的细节。Adam 优化器中有一个常见的超参数 ε，通常设置为 10⁻⁸。它被加在梯度归一化的分母中，用于控制数值行为。

论文发现，在模型已经记住训练集、并使用权重衰减的特定条件下，Logit Margin 的增长会受到权重衰减与 Adam 更新之间平衡关系的限制。其平衡值近似满足：

m ≈ ln (1/ε)

当 ε=10⁻⁸ 时，这个值约为 18.4。而 fp32 的精度悬崖恰好位于 16.6。这意味着，模型在接近自然形成的平衡位置之前，就可能先触发标准交叉熵的数值问题。这也为一些训练实验中「Adam 的 ε 太小可能导致不稳定」的观察，提供了新的解释。

研究者发现，在相应实验设置中，将 ε 提高至 10⁻⁷ 或更大，可以使平衡位置落在 fp32 精度悬崖以内，从而避免这类尖峰。不过，调整 ε 只是改变优化器的行为。修复交叉熵的计算方式，才能直接消除这处数值错误。

学生已经等于老师

训练却还在继续

第一个 Bug 发生在模型对答案过于自信的时候，第二个 bug Softmax Cancellation（Softmax 相消误差），则出现在两个模型的输出非常接近的时候。

典型场景就是知识蒸馏。在知识蒸馏中，教师模型负责提供目标概率分布，学生模型尝试学习这些输出。如果学生与教师的 Logit 完全相同，那么两者对应的概率分布也应该完全一致。此时，学生已经达到了这个蒸馏目标，其梯度在数学上应当为 0。

但现有标准实现可能算出一个非零梯度。原因在于，标准软目标交叉熵通常会分别计算教师和学生的 Softmax 概率，再通过两者相减得到梯度。当两组概率非常接近时，浮点运算会丢失它们之间的微小差异。

更麻烦的是，前向传播中的 Softmax 与反向传播涉及的 Log-Softmax，可能使用略有不同的运算顺序。即使教师和学生的 Logit 完全相同，两条计算路径产生的舍入误差也可能不同，最终留下一个本不应该存在的梯度。这个残余梯度看起来微不足道，却可能被 Adam 放大。

为了测试其影响，研究团队先训练一个教师模型，使其在模 53 加法任务上达到 100% 的测试准确率，再直接复制教师的权重初始化学生模型。也就是说，训练开始时，学生和教师完全相同。

论文随后使用蒸馏温度 4、ε=10⁻⁸ 的 Adam 优化器，在不使用权重衰减的条件下继续训练学生 2000 步。

理论上，学生没有任何偏离教师的必要。实验结果却显示，使用标准软目标交叉熵时：

• fp64 的残余梯度非常小，学生基本保持不变。

• fp32 出现短暂的 KL 散度尖峰，说明学生输出一度偏离教师。

• bf16 的残余梯度更大，学生在持续训练中明显偏离教师，测试准确率也随之下降。

研究者对第二个 bug 采用了另一种修复方式。

既然分别计算两组概率再相减容易丢失精度，就从学生与教师的 Logit 差值出发，直接在对数空间中计算所需的概率差。具体实现使用了 expm1 和 log1p 两个专门处理接近零的小数值运算的函数，并进一步修正剩余的浮点误差。

这样，当学生与教师的 Logit 完全相同时，修复后的梯度就能精确归零。

在同样的蒸馏实验中，采用修复版本后，学生在 bf16、fp32、fp64 三种精度下都保持与教师一致，KL 散度维持在 0，测试准确率保持 100%。这一结果表明，在模型输出非常接近的蒸馏场景中，低精度计算带来的微小误差确实可能干扰优化过程。

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

除了两处数值问题，Barron 在社交平台分享的研究过程也值得关注。

按照他的介绍，Softmax Cancellation 最早由 Gemini Flash 3.7 发现并给出修复方案，论文中的大部分实验则由 Gemini Flash 3.8 完成。

到了论文写作阶段，研究团队又使用了 Gemini 4 Argon。

Barron 表示，Argon 完成了论文及图表的初稿。随后，他对初稿进行了严格审查和反复修改，直到技术内容足够扎实，但对最终的文字表达仍然不太满意。

这套研究流程中，AI 已经参与到数值问题发现、修复方案设计、实验执行和论文撰写等多个环节，研究人员则持续承担技术核查与质量把关的工作。

总结

回到技术结果，这项工作的价值还需要放在实际应用场景中理解。

论文主要通过模运算 Grokking 任务和特定知识蒸馏实验验证修复效果。研究团队也承认，实验中出现的极大 Logit Margin，以及学生与教师几乎完全相同的输出，在常规大规模训练中并不常见。

因此，目前还不能据此认定，主流大模型训练普遍受到这两个问题的严重影响，更无法直接推断修复后就能让大模型训练整体提速。

但研究确实揭示了一个容易被忽视的事实：即使数学表达式完全正确，计算机实际执行的浮点运算，也可能在特定条件下改变梯度的性质。

Softmax 和交叉熵已经是深度学习工具链中的基础组件。多年来，研究者一直在优化模型架构、训练算法与算力效率，而这次发现表明，连这些最基础的计算实现，仍然存在可以进一步改进的空间。

论文给出的修复方案保留了原有函数的数学定义，也不需要通过提高浮点精度来解决问题。作者认为，这些函数有望直接替换现有实现，且不会带来明显的性能开销。不过，其在大规模模型训练中的实际效果，还需要进一步验证。