Google 研究员称 JAX/PyTorch 的 softmax 全有 bug,给出重写修复

jon_barron · x · 2026-10-24

Google 研究员 Jon Barron 宣布新结果:所有现有 softmax 与交叉熵实现(包括 JAX、PyTorch、教科书)都存在两个细微 bug,一个简单重写即可修复,并能消除 grokking 过程中的 "slingshot" 尖峰。

第一个 bug 是 "Softmax Collapse":当一个 logit 远大于其他 logit 时,其梯度会静默下溢为 0,而其他 logit 的梯度不会,导致梯度之和不再为零。配图显示标准 CE 在不同精度(float64/32/bfloat16)下目标 logit 梯度会在特定 logit margin 处突降为零,而修复版曲线平滑。帖子附上了完整的 JAX 修复代码(logsoftmax 与 crossentropy 的重写实现)。