DeepMind 指出 PyTorch、JAX 的 Softmax 实现存在数值精度缺陷
最常用的 Softmax 和交叉熵在大 Logit 差距下会把梯度舍入成 0，这可能解释训练损失为何突然飙升；修复只需改计算顺序。
Google DeepMind 团队发现，PyTorch、JAX、TensorFlow 等主流框架的 Softmax 与交叉熵标准实现中存在两个数值精度 Bug：当模型对正确答案非常自信时，正确类别的梯度会被错误地算成 0，破坏梯度之和为零的数学约束；另一个是相消误差。论文给出只需重排计算方式、无需提高浮点精度的修复方案，并解释它可能与训练中损失突然飙升的「弹弓」现象有关。
正文摘录
.jpg) 深度学习中最基础、使用最广泛的两个函数，竟然藏着两个隐蔽的 Bug。 昨日，Google DeepMind 首席研究科学家 Jonathan T. Barron 在 X 上公布了一项新研究。 他指出，包括 PyTorch、JAX、TensorFlow 在内的主流深度学习框架，普遍采用的 Softmax 和交叉熵（Cross-Entropy）标准实现中，存在两个容易被忽视的数值精度 Bug。 这些 bug 平时很难察觉，却可能在模型训练的特定阶段悄悄改变梯度，造成损失突然飙升，甚至让模型已经学到的能力出现退化。 更有意思的是，研究团队发现，只需调整计算方式，就能修复这两个 bug，无须提高浮点精度，也不需要重新设计损失函数。