DeepMind发现Softmax与交叉熵数值Bug，Gemini参与研究
深度学习领域最基础、应用最广泛的两个计算函数——Softmax与交叉熵，近期被发现存在两处不易察觉的数值精度缺陷。Google DeepMind的首席研究科学家Jonathan T. Barron在社交平台公布了这项研究成果，指出包括PyTorch、JAX、TensorFlow在内的主流深度学习框架的标准实现均存在此类问题。
这两处缺陷平时很难被察觉，却可能在模型训练的特定阶段悄悄改变梯度，造成训练损失突然飙升，甚至导致模型已学到的能力出现退化。更关键的是，研究团队找到了无需提升浮点精度、也无需重新设计损失函数的修复方案。
第一个缺陷被称为Softmax坍缩，与模型对答案的自信程度直接相关。训练分类模型时，神经网络会为每个候选类别输出原始分数（即对数几率Logit），Softmax负责将这些分数转换为概率，交叉熵则用来衡量模型预测与正确答案的差距。随着训练推进，模型对正确答案的把握逐渐提升，正确类别的Logit与其他类别之间的差距会不断扩大。
现有实现中，通常会先从所有Logit中减去最大值再进行指数计算，这是防止指数溢出的经典数值稳定性处理。但减去最大值后，正确类别对应的指数值变为1，其他类别则是极小的正数。当正确类别分数足够高时，Softmax的分母就近似为“1+极小正数”。由于浮点数精度有限，当这个极小正数足够小时，计算机会直接将其舍去，导致分母被四舍五入为1，正确类别的预测概率被判定为精确的1，交叉熵损失变为0，对应梯度也被算成0。
按照交叉熵的数学定义，所有类别对应的Logit梯度之和应为0，此时正确类别的梯度提前归零，其他类别的梯度却仍存在，这就破坏了原本的数学约束，论文将这种现象称为精度悬崖。
不同浮点格式下的精度临界值不同：bf16的Logit差距临界值约为5.5，fp32约为16.6，fp64约为36.7。这里的Logit差距指正确类别Logit相对于其他类别合计得分的优势，以fp32为例，当这一差距超过16.6时，就会触发精度悬崖，此时模型对正确答案的预测概率已高达约99.9999938%。
论文通过具体例子说明：假设三个Logit为[9,-9,-9]，标准实现计算出的交叉熵损失为0，正确类别的梯度也为0；而修复后的计算方式中，损失约为3.046×10⁻⁸，正确类别的梯度约为-3.046×10⁻⁸，另外两个类别的梯度各约为1.523×10⁻⁸，这些微小的梯度对数学约束的保持至关重要。
2025年已有研究者识别出类似现象并提出StableMax，通过替换Softmax的指数函数缓解问题，但Barron团队选择了更直接的方式：保留原有Softmax与交叉熵的定义，重新组织计算流程，将最大Logit对应的项从指数求和中分离，用数值稳定的softplus函数完成计算，修复后的公式能让自动微分直接得到正确类别的梯度等于其他类别梯度之和的相反数，确保微小梯度不会丢失。
第二个缺陷被称为Softmax相消误差，出现在两个模型输出非常接近的场景中，典型代表是知识蒸馏任务。在知识蒸馏中，教师模型输出概率分布，学生模型学习这些输出；如果学生与教师的Logit完全相同，两者的概率分布应一致，此时学生的梯度数学上应为0。
但现有标准实现中，通常会分别计算教师与学生的Softmax概率，再通过两者相减得到梯度，当两组概率非常接近时，浮点运算会丢失它们之间的微小差异，导致算出非零梯度。此外，前向传播的Softmax与反向传播的Log-Softmax运算顺序略有不同，即使Logit完全相同，两条路径的舍入误差也可能留下不该存在的梯度，这类微小梯度还可能被Adam优化器放大。
研究团队通过蒸馏实验验证了这一点：先训练教师模型在模53加法任务上达到100%准确率，再复制教师权重初始化学生模型，用蒸馏温度4、ε=10⁻⁸的Adam优化器训练2000步。结果显示，标准实现中fp64的残余梯度很小，学生基本不变；fp32出现短暂的KL散度尖峰，学生输出偏离教师；bf16的残余梯度更大，学生明显偏离，准确率下降。
针对这个缺陷，研究团队采用了不同的修复方式：从学生与教师的Logit差值出发，直接在对数空间计算概率差，用expm1和log1p处理接近零的小数值，修正剩余浮点误差，这样当Logit完全相同时，修复后的梯度能精确归零。在同样的蒸馏实验中，修复版本在三种精度下都保持与教师一致，KL散度为0，准确率100%。
这两个缺陷还能解释神经网络训练中的弹弓机制现象：2022年研究者在研究“顿悟（Grokking）”时发现，模型记住训练数据后，损失会长时间接近0，随后突然飙升再回落，伴随测试准确率跃升。论文指出，这正是Softmax坍缩导致的：当模型记住训练集后，Logit差距扩大到精度悬崖，标准实现将正确类别的梯度归零，而Adam优化器的二阶矩估计会因长期梯度为0而衰减，后续参数变化让样本回到精度悬崖内，恢复的梯度因Adam分母过小导致参数更新异常，损失随之飙升。
研究团队通过模运算实验验证了这一机制：用单层Transformer学习模31加法，fp32精度下，标准交叉熵训练中反复出现损失尖峰，5次实验平均准确率为94.3%±5.4%；修复版本训练平稳，准确率均达100%。在bf16下，标准实现平均准确率93.3%，修复版本约100%。另一组MLP学习模97平方和的实验中，标准实现在约1250步触发精度悬崖，准确率从61%骤降至23%，修复版本则提升至约85%。
研究团队还发现Adam优化器的默认参数恰好踩在临界点附近：Adam的ε通常设为10⁻⁸，论文指出，在模型记住训练集且使用权重衰减的条件下，Logit差距的平衡值约为ln(1/ε)，当ε=10⁻⁸时约为18.4，而fp32的精度悬崖在16.6，这意味着模型在达到自然平衡前就可能触发数值问题。调整ε至10⁻⁷以上可避免尖峰，但修复交叉熵的计算方式才是根本解决方法。
值得关注的是，Gemini在这项研究中承担了大量工作：Softmax相消误差由Gemini Flash 3.7发现并提出修复方案，多数实验由Gemini Flash 3.8完成，论文初稿由Gemini 4 Argon生成，后续由研究人员完成技术核查与修改。这套流程中，AI参与了从问题发现到论文撰写的全环节，研究人员则负责技术把关。
这项研究的价值需要结合实际应用理解：论文主要通过模运算顿悟任务和知识蒸馏实验验证修复效果，而实验中出现的极大Logit差距、学生与教师完全相同的输出，在常规大规模训练中并不常见，因此目前无法认定主流大模型训练普遍受这两个问题影响，也不能直接推断修复后能大幅提升大模型训练效率。
但研究揭示了一个易被忽视的事实：即使数学表达式正确，计算机实际执行的浮点运算也可能在特定条件下改变梯度性质。Softmax与交叉熵作为深度学习的基础组件，多年来研究者优化了架构、算法与算力，却没想到最基础的计算实现仍有改进空间。修复方案保留了原有函数的数学定义，无需提升精度，有望直接替换现有实现且无明显性能开销，其在大规模模型中的实际效果仍需进一步验证。
