谷歌DeepMind披露Softmax与交叉熵实现存在舍入Bug

谷歌DeepMind主任研究科学家Jon Barron在X平台披露，PyTorch、JAX及教科书中Softmax和交叉熵的标准实现存在舍入Bug，可能导致模型训练时梯度消失、训练曲线剧烈震荡。Barron称，该问题最初由Gemini Flash 3.7发现并完成修复，团队已在arXiv提交论文《Softmax和交叉熵的修复》，提出重写方案。

Softmax和交叉熵是AI模型训练中最基础的运算：前者将模型输出的原始得分转换为概率，后者计算预测与真实标签之间的误差。为防止计算溢出，现有标准实现会先减去最大得分，使分母变成类似“1+极小数字”的结构。在浮点运算中，当这个极小数字足够小时会被直接舍入为0，导致正确的梯度瞬间归零。DeepMind论文将这一现象称为“Softmax崩溃”（Softmax Collapse），并指出主流优化器Adam会进一步放大该问题，造成损失函数突然飙升，即此前被部分研究者视为“顿悟”特征的“弹弓效应”。

论文还描述了另一类“Softmax抵消”问题：在知识蒸馏场景中，标准代码分别计算概率再相减，舍入误差会留下“幽灵梯度”，使学生模型偏离正确答案。修复的核心原则是“永远不要把一个极小的数字加到1上”，并对交叉熵采用log-sum-exp与softplus的等价形式，在蒸馏场景改用expm1和log1p计算。

测试数据显示，在模加法（a+b mod 31）任务中，bfloat16精度下标准代码的准确率为93.3%且不稳定，改用修复代码后准确率升至100.0%。Barron在评论区补充，该Bug最初由Gemini Flash 3.7发现并修复，论文中大部分实验由Flash 3.8完成，Gemini 4 Argon撰写了论文和图表初稿。10月24日，独立工程师Satyajit Ghana在PyTorch中复现了这两个Bug，并证实了论文在fp32精度下的结果。