谷歌研究员：主流框架softmax都有两个bug
一篇新上传至arXiv的论文指出，从PyTorch、TensorFlow到JAX，所有主流深度学习框架的标准softmax和交叉熵实现都带着两个数值bug。论文编号2610.31642，作者是谷歌研究员Jon Barron等人，标题直白——《修复softmax与交叉熵》。
第一个bug叫softmax坍缩。当目标类别的logit远大于其他类别时，交叉熵损失中e的负指数次项小到在浮点加法里直接被吃掉，损失与目标梯度被舍入为零，梯度向量不再满足求和为零的性质。论文给出了悬崖的精确位置：fp32下当目标logit领先约16.6时精度崩塌，bf16约5.5，fp64约36.7——不是慢慢变不准，而是在阈值上一步归零。
第二个bug出现在知识蒸馏场景：当学生网络的logit与教师完全一致时，梯度本应为零，标准实现却会算出非零梯度，把学生从教师身边推走，论文称之为softmax抵消。
修复本身简单得可疑：只需在计算中保留目标logit的独立通路，几行代码即可同时解决两个问题。论文报告的效果立竿见影——消除训练中著名的弹弓（Slingshot）异常、在有权重衰减时加快grokking、无权重衰减时阻止精度崩塌、蒸馏中学生不再漂移；论文还据此推出记忆化之后logit增长受权重衰减与Adam的epsilon之间平衡约束的新结论。
论文发布后，有开发者仅用NumPy花一个下午就独立复现了两个bug和理论悬崖位置（fp32实测约16，随类别数与求和顺序在15.8至16.7之间浮动），并指出悬崖位置其实取决于框架的归约树实现——每个框架各有自己的悬崖。对绝大多数用户来说，这未必是急需修复的故障，但它再次提醒：深度学习的地基里，浮点才是第一性的。
编译自 arXiv，原标题："A Fix for Softmax and Cross-Entropy"
