5.5这道坎：softmax里藏了个bug，把梯度悄悄清零

训练一个模型，softmax把原始分数变成概率，交叉熵再给这些概率打分。这两步几乎每个模型都要走一遍。三位谷歌DeepMind研究者——Jonathan T. Barron、Dan Gnanapragasam、Michael Tschannen——在2026年10月22日挂出一篇论文，认为这两个函数在JAX、PyTorch、TensorFlow以及教科书里的标准实现，都带着两个数值bug。修法不长，用softplus、expm1、log1p重写几行。在论文的grokking基准上，bf16测试准确率从93.3%升到100%。

这不是新架构的故事，是算术的故事：当你把一个小到不能再小的数加到1.0上，计算机安静地回你一句"那就是1.0"。

想象一块有一万个格子的记分牌，你只能写大约7位有效数字。模型赢得很彻底，正确答案的分数远高于其他。要报告它有多自信，你得写下类似0.00000004这样的数——判断错误的那点残余概率。标准做法把它算成1 + 0.00000004，再减1。在7位精度的记分牌上，1 + 0.00000004直接舍入成1，1 - 1 = 0。模型本来有99.999996%的自信，账本上却记成100%，一点学习空间都不剩。

更麻烦的是，记分牌上其他格子没有先加1，那些微小值被保留了下来。梯度不再求和为零——这是梯度下降依赖的基本恒等式——优化器开始朝一个特定、有结构的方向走错步。

这个临界点在论文里有名字：精度悬崖。bf16下大约在logit差距5.5处，fp32是16.6，fp64是36.7。越过悬崖，正确答案上的梯度精确为零，错误答案上的梯度却不是。作者管这叫Softmax Collapse。

softmax是exp(z_k)除以sum(exp(z_j))，直接算会溢出——exp(88)已经超出fp32范围。所以每个框架都先减去最大logit，即max-shift技巧，自Higham 2002年起就是标准做法。移位后最大logit变成0，它的指数精确为1，分母变成1 + e^-m，m是领先幅度：赢家logit比其余部分的log-sum-exp高出多少。

不对称就在这里。浮点数在0附近比在1附近密得多：fp32里0附近的值可以小到约10^-38，但任何加到1.0上的东西，低于约10^-7就消失了。e^-m在输家类别的分子里活得好好的——那些是直接算的，没加1——可一旦在分母里被加到1.0上，m越过悬崖它就舍入没了。损失本身舍入成精确的0（log(1) = 0），目标梯度舍入成精确的0（q_t - 1 = 1 - 1），只有输家保住了梯度。梯度向量求和变成一个正数，不再是零。

一个错得很小但确实错的梯度，为什么要紧？因为Adam。自适应优化器用梯度幅度的滑动平均来归一化每个梯度，一个本该是-4e-8却读成0、而同伴读成+1.5e-8的梯度，不是Adam能平均掉的舍入误差，而是每一步都重复的系统性单边失真。论文把两个此前说不清的现象追到了这里：grokking训练里反复出现的"弹弓"式损失尖峰（Thilak等，2022），以及权重衰减下的grokking失败（Prieto等，2025）。此前的绕法——替代损失、梯度投影、转成fp64——治的都是症状，这篇论文修的是函数。

修法是永远不要直接把1加到e^-m上。对硬目标交叉熵，把目标那一格换成常数0，只用其他logit算出m，返回softplus(-m) = log(1 + e^-m)，算法上不形成1 + e^-m这个舍入和。自动微分随后给出目标梯度，等于非目标梯度之和取负——梯度按构造求和为零，精确到margin约88处的指数下溢。论文自己的例子很具体：logits为[9, -9, -9]、目标为0时，标准JAX/optax返回损失0.0、目标梯度0.0；修好的代码返回损失3.046e-8、梯度[-3.046e-8, 1.523e-8, 1.523e-8]——小、精确、求和为零。

第二个bug叫Softmax Cancellation，打在知识蒸馏上——大多数小而可部署的模型都靠它，学生网络学着匹配老师的概率。标准代码用softmax算学生的概率，用另一条log-softmax路径算老师的，再相减。学生逼近老师时，真实梯度本该缩到零，实际却停在一个非零噪声底上，因为两个公式在fp32的最后几位上不一致。

配套notebook直接复现了这点：老师logits为[2.0, 0.5, -1.0]时，前向softmax与log-softmax取指数后的结果，即使学生和老师logits完全相同，梯度范数也差约6.2e-8。一个已经收敛的学生持续收到幽灵梯度，从它早已匹配上的老师身边漂走。修法是在对数空间里直接从logit差z_S - z_T算出q - p（借助expm1/log1p），logit相等时按构造精确为零。

论文的头号实验是个小型grokking基准——模加法，a + b mod 31，5个种子取平均——一张表六个数字就是全部论证：

• bf16：93.3% → 100.0%（+6.7pp）；fp32：94.3% → 100.0%（+5.7pp），均在Adam默认epsilon 1e-8下。标准bf16跑到最后margin停在5.1，刚好卡在它5.5的悬崖下面，正是一个在边界上磨到停的run的签名。
• 弹弓消失。修好的代码去掉了困扰grokking训练的损失尖峰，加快带权重衰减的grokking，并防止不带权重衰减时的准确率崩塌。
• 蒸馏不再漂移。用修好的损失训练的学生网络停在老师那里，而不是收敛后四处游走。
• 新的平衡，而非失控。记忆之后，论文显示logit增长被权重衰减与Adam epsilon之间的平衡所约束——一个干净、可检验的预测，取代了旧图景里logit无界增长直到某处崩掉。

该留意的有三点。其一，这是2026年10月22日的预印本，尚未同行评审，头号数字来自小基准（模加法、5个种子），不是前沿大模型训练——作者的主张是关于算术的，算术在notebook里可查，但还没有人展示过某个100B参数训练因为这一修法改变了结果。其二，还没有框架采纳：JAX、PyTorch、TensorFlow仍在发布标准代码，所以今天的修法是一段粘贴进去的代码，不是一个升级就能拿到的版本；论文自己的说法是它"对grokking、可解释性和蒸馏最有用"，而不是对每种工作负载都保证加速。其三，悬崖是低精度的故事：在fp64和默认Adam epsilon下，标准版和修复版都到100%，全程满精度训练的团队本来就没暴露。暴露的是所有用bf16/fp32配Adam训练的人——也就是说，今天训练前沿模型的几乎所有人。

bug在分母里，不在模型里。1 + e^-m越过margin 5.5（bf16）/ 16.6（fp32）/ 36.7（fp64）就舍入成精确的1.0，这一次舍入把目标梯度清零，却留下输家的梯度。Adam把微小的不对称放大成可见的损伤：逐坐标归一化把单边误差变成弹弓尖峰和grokking失败。修的地方在损失函数，不在优化器。

蒸馏有它自己的地板。两个分别舍入的概率向量永远不能完全抵消（fp32里约6e-8），收敛的学生因此漂移；从logit差直接算差值，按构造就修好了。整套修法是三个函数，不是新框架：log-softmax、softmax、交叉熵的softplus/expm1/log1p重写，论文里给了可直接替换的JAX代码，据称没有明显变慢，梯度求和为零可撑到margin约88。

想自己验证也不难。配套notebook在NumPy里复现两个bug：悬崖穿越、论文的[9, -9, -9]例子（损失3.046e-8对0.0）、以及蒸馏地板。如果你的训练栈跑bf16配Adam，这大概是今年最便宜的一次审计。