掘金· lucas_AI·· 20 天前精选AI 评分71
DeepMind 论文指出 PyTorch、JAX 等框架 softmax 与交叉熵存在两个浮点 Bug
PyTorch、JAX 全中招：Softmax 藏了两个数值 Bug，DeepMind 用十余行代码修好了
AI 导读
Google DeepMind 首席研究科学家 Jonathan T. Barron 与同事在 arXiv 论文《A Fix for Softmax and Cross-Entropy》中指出，PyTorch、JAX、TensorFlow 和教科书共用的 softmax 与交叉熵标准实现存在两个浮点 Bug。
借鉴点
可照抄论文的 softplus/expm1/log1p 重写 log-softmax 十余行实现，替换训练循环里的交叉熵，代价是需自行核对 bf16 与 fp32 下 margin 悬崖位置，论文未给统一阈值。
来源：掘金 · juejin.cn