Google DeepMind 的首席研究科学家、NeRF 的原作者 Jon Barron 他带团队在 arXiv 上发了一篇题为A Fix for Sof… · @snwiki238337 · TwKit

@snwiki238337

Google DeepMind 的首席研究科学家、NeRF 的原作者 Jon Barron 他带团队在 arXiv 上发了一篇题为A Fix for Softmax and Cross-Entropy，直接点名包括 PyTorch、JAX、TensorFlow 在内的所有主流深度学习框架。 我们在 ai 训练中每天在用的这个 Softmax 和交叉熵这两个最基础的算子，标准实现里是有精度的 bug 的。 其实现有框架为了防止指数运算溢出，表面上都会先从所有 Logit 中减去最大值，但这在浮点数精度有限的条件下会引发两个 bug。 第一个叫 Softmax 坍缩，简单来说的话其实就是当模型对正确答案越来越自信、得分差距拉大时，分母里的 1 加上一个极小数会被计算机直接四舍五入成 1，导致正确类别的梯度提前归零，而错误类别的微小梯度还在。 直接破坏了所有类别梯度之和必须等于 0 的规则。 这顺手解释了模型顿悟训练中的弹弓效应，也就是梯度归零后 Adam 优化器的分母不断衰减，一旦样本退回精度悬崖内梯度恢复，就会瞬间出现一次巨大的参数更新。 其实这从日常的训练、调参都能感受到，这点我相信很多人都有所感受，有时候明明训练集其实已经拟合好了，但是Loss 贴着地面走却突然莫名其妙出一个尖峰。 第二个叫 Softmax 相消误差，主要出现在知识蒸馏里，目前的话标准实现是分别算出老师和学生的概率再相减，当学生已经学得跟老师一模一样时，前反向不同的舍入误差会留下无法归零的噪声梯度，在 bf16 低精度下及其容易把已经对齐的学生模型训练出问题。 针对这两个底层缺陷，Barron 团队在项目主页给出的解法只需把最大 Logit 项单独剥离出来用 softplus 重组计算，并在蒸馏时直接从老师和学生的 Logit 差值出发在对数空间里算概率差。 在模运算测试里，原版算子跑到中途准确率从 61% 下降到 23%，而换上修复代码后平稳升到了 85%。 但实际上Barron 透露出第二个相消误差 Bug 最早就是由 Gemini Flash 3.7 发现并给出修复方案的，论文里的大部分实验由 Gemini Flash 3.8 跑完，初稿和图表是由刚发布的 Gemini 4 Argon 生成。 目前的话呢团队已经在论文里直接放出了可直接替换的 JAX 参考实现，几乎不增加任何额外计算开销。 arxiv：https://arxiv.org/pdf/2610.31642

曝光 81 · 评论 1 · 点赞 0 · 书签 0