先说一个可能让你意外的事实。
在做低精度大模型训练这件事上，业界早就把矩阵乘法、权重、激活值都压缩到了很窄的数值格式里，FP8、FP4满天飞。但唯独softmax这一步，几乎所有主流方案都留了一手，继续用高精度浮点数计算。英伟达自己的MXFP8方案文档里甚至直接写着，把softmax精度降下来这件事,"留给未来的工作"。
这不是疏忽,这是心虚。
因为softmax这个操作,表面上看只是把一堆分数变成加起来等于1的概率,实际上它牵扯着整个注意力机制的学习过程。你把它简化了,模型的前向计算会变,但更麻烦的是,反向传播的梯度也会跟着变。而梯度这东西,直接决定了模型往哪个方向学习。如果这一步出了偏差,你可能要等到模型训练了几千万个token之后才会发现问题——那时候已经晚了。
这篇来自伊利诺伊大学香槟分校的论文,做的就是把这个"心虚"的地方掰开揉碎,看看到底会出什么问题,以及怎么才能避免。
一个被忽略的假设:"求导的时候能不能偷懒"
要理解这篇论文在做什么,得先明白softmax本身在计算时有一个小技巧。
Softmax
一种把任意实数向量转换成概率分布的函数,常用于注意力机制里给不同的"关注对象"分配权重,所有权重加起来等于1。
标准softmax计算时,为了避免数值爆炸(指数函数在输入很大时会瞬间冲到无穷大),工程实现上都会先减去这一行分数里的最大值,再做指数运算。这个操作在数学上是"安全的",因为softmax函数有一个性质叫平移不变性,你把所有输入都加上或减去同一个常数,输出的概率分布完全不变。既然如此,减去最大值这一步,在反向传播时它的梯度贡献会精确抵消为零,所以工程师们干脆把这一步的梯度直接掐断("detach"),图省事,也没什么坏处。
这个习惯用了很多年,大家都觉得理所当然。
论文提出的核心问题是,当你不再用标准softmax,而是用一种基于查找表的近似方法时,这个"偷懒掐断梯度"的习惯还安全吗?
答案是,不安全,而且后果很隐蔽。
因为在近似方案里,这一行分数的最大值和最小值不再只是"用来防止数值爆炸的临时工具",它们本身变成了决定查找表精度的关键参数。掐断它们的梯度,相当于告诉模型"这两个数是常量,你别管它们",但实际上模型的学习行为会持续调整这两个数,导致查找表越来越粗糙,而这个偏差却完全没有反映在梯度里。
这就好比一个工厂的质检流程。原本质检员只是抽查一下最后成品,顺手记录一下,这个记录不影响生产线的任何决策,因为无论怎么记录,产品该怎么造还是怎么造。但现在工厂换了一套新流程,质检员记录的这两个数字,直接决定了下一批原材料要按什么规格采购。如果车间主任还按老习惯,不把这个记录反馈给采购部门,采购部门就会一直按旧标准买料,原材料规格和实际生产需求渐行渐远,而且这个偏差在账面上完全看不出来,因为产品出厂检验的时候表面上还合格。如果不把这个反馈链接上,工厂运转个把月看着一切正常,但库存积压的问题会越滚越大,直到某天彻底压垮生产线。
论文管这套"能不能求导"的完整数学推导叫全量校准梯度,而只保留了老习惯的那种叫detach。实验里,用LERP方法在K=32的设置下,detach版本的模型训练了大约2500万到3000万个token时还和标准softmax差不多,然后突然开始跑偏,最后跑出的验证损失比正常版本高出0.65到3.07纳特(这是衡量语言模型预测好坏的单位,数值越低说明模型预测得越准)。更让人意外的是,研究者试着做了一个"补救",把梯度硬投影到一个数学上"更干净"的子空间(让梯度满足零和的性质),结果这个补救版本反而比什么都不做的detach版本更差。这说明,问题不只是梯度的某个数学性质被破坏了那么简单,背后还有别的机制在起作用。
三个旋钮:K区间注意力到底怎么设计的
要研究这个问题,论文构造了一整套叫K区间注意力的操作族。
K区间注意力
用K+1个预先算好的查找表数值,通过插值或者四舍五入的方式,来近似指数函数exp(z),从而替代标准softmax里代价较高的指数运算。
这套操作族里有三个可以独立调节的设计维度,论文把它们叫做三个轴。
第一个轴是校准方式,也就是这张查找表该怎么划分区间。论文比较了两种做法。一种叫MinMax,顾名思义,就是拿这一行分数里真实的最大值和最小值来划定查找表的范围,这个范围会随着模型训练不断变化。另一种叫固定窗口锚定,简写为FWM,做法是只盯着最大值,往下切出一个固定宽度(论文里定为6纳特)的窗口,窗口之外的分数一律当作零权重处理。
这两种做法的区别,类比一下就像两种不同的相机曝光策略。MinMax相当于每次拍照都根据画面里最亮和最暗的地方自动调整曝光范围,画面对比度越大,曝光的"刻度"就要拉得越宽,每一格代表的亮度差异也就越粗糙。而FWM则像是固定死了曝光范围,只根据画面里最亮的点往下切一段固定宽度,不管画面整体对比度有多剧烈,曝光刻度的精细程度始终不变。如果你的场景对比度会随时间剧烈波动(这正是模型训练过程中注意力分数分布会发生的事情),用MinMax的相机会随着对比度变大而越拍越糊,而用固定窗口的相机分辨率始终如一。这正是论文最后发现FWM往往比MinMax更稳的根本原因,自适应看似聪明,但如果这个"自适应"本身失控了,后果反而更严重。
第二个轴是重建方式,也就是拿到查找表之后,怎么把一个具体的分数变成一个权重。一种叫LERP,是分段线性插值,分数落在两个查找表数值之间时,按比例插值出一个连续的权重,这种方式的好处是处处可导,反向传播时可以算出精确的梯度。另一种叫Nearest,直接把分数四舍五入到最近的一个查找表格点,这样权重就变成了一个个离散的台阶,这种方式在数学上几乎处处梯度为零,因为四舍五入这个操作本身就是不连续的。
这就带来一个麻烦。Nearest方式的前向计算虽然可以做,但反向传播没法直接用,因为你没法对一个"四舍五入"求导数。论文的解决办法是引入第三个轴,叫直通估计器的放置位置。
直通估计器
英文缩写STE,是一种应对不可导操作的训练技巧,前向计算时用真实的离散结果,反向传播时假装这个操作是"透明"的,直接把梯度传过去,不做任何修正。
论文比较了两个放置位置。一个叫Weight-STE,在权重还没有归一化(也就是还没有除以总和变成概率)之前就做这个直通替换。另一个叫Prob-STE,等到权重已经归一化成概率之后再做替换。两者的前向输出在数学上是完全一样的,区别只在于反向传播时,那个"用于计算梯度的雅可比矩阵"是在权重归一化之前的点上算的,还是在归一化之后的点上算的。
这个区别听起来很技术,但后果差异巨大。论文发现,在K=4这样比较粗糙的设置下,把直通估计器放在归一化之前(Weight-STE)配合MinMax校准,会导致训练损失比标准softmax高出0.39纳特(在2.5亿token时)甚至0.89纳特(在25亿token时)。换成归一化之后放置(Prob-STE),或者换用FWM校准,这个差距几乎都能大幅收窄。而如果两个改动同时做,效果并不会比单独做一个改动好太多,说明这两个因素之间存在很强的相互替代关系。
这就好比一场接力赛里交接棒的时机。Weight-STE相当于在选手还没跑到交接区之前就把棒交出去,这时候前一位选手的速度和后一位选手的速度还没对齐,交接的动作生硬,容易掉棒。Prob-STE则是等两位选手速度基本同步之后再交棒,整个动作更平滑。区间越宽(也就是MinMax校准下查找表格子越粗),两位选手的速度差就越大,交接的风险也就越大;而如果区间本身就窄(比如FWM固定宽度),不管在哪个时机交接,风险都不高。这也解释了论文观察到的另一个现象:随着K(区间数量)增大,两种直通估计器放置方式之间的差距会缩小,而且在FWM设置下缩小得比MinMax更快,因为FWM的区间宽度从一开始就被死死锁定,不会随着训练变大。
崩溃的真实模样:范围失控的连锁反应
光看数字可能还不够直观,论文里还专门追踪了模型训练过程中,注意力分数本身的几何形态是怎么变化的。
论文发现,在MinMax校准配合Weight-STE的这个组合里,模型学到的注意力分数范围(也就是一行分数里最大值减最小值)会持续膨胀。到25亿token训练结束时,这个范围达到了标准softmax下的大约7倍,查找表的每个格子宽度达到53纳特,而模型自己的原生概率分布和标准softmax概率分布之间的总变差达到了0.86(总变差是衡量两个概率分布差异的指标,取值在0到1之间,越接近1说明两个分布差异越大)。相当于模型自己越训练,分数越锋利,但近似算子的分辨率却跟不上,最后变成了近乎"一刀切"的粗糙判断,注意力权重的顶格那一档几乎独占了99.5%的概率质量。
论文把这个现象叫做范围-分辨率反馈假说:区间越宽,直通估计器传回来的梯度信号就越失真,而失真的梯度又会进一步鼓励模型把分数拉得更极端,分数越极端,区间又变得更宽,如此循环,最后失控。
这就像一台老式收音机,天线越拉越长本来是为了收到更远的信号,但如果接收电路本身跟不上,拉得越长反而引入越多噪音,噪音又让你以为信号不够强,于是继续拉天线,最后天线拉到屋顶那么长,收到的却全是杂音,一句清楚的话都听不见。如果不设法把天线长度锁定在一个合理范围(这正是FWM在做的事情),这个恶性循环会一直持续下去,直到整个系统彻底失灵。
不过论文很坦诚地指出,这个假说不是无懈可击的。他们发现K=64的MinMax-Weight设置下,分数范围同样在持续膨胀翻倍,但最终的损失差距却一直稳定在0.02到0.04纳特之间,并没有跟着崩溃。而且在2.5亿token的实验里,损失差距在训练早期(5000万token)就已经打开了,那时候两种直通估计器放置方式对应的分数几何形态几乎还没分化。这说明因果链条没有那么干净,更准确的说法是,这是一个描述性的关联,不是被实验直接证实的机制。
不同规模、不同预算下,这个问题有多普遍
论文没有只在一个模型规模、一个训练预算下做实验,而是设计了四组正式的实验套件,分别在1.24亿参数模型上训练2.5亿或25亿token,在10亿参数模型上训练1亿token,以及在1.24亿参数模型上用五个不同的随机种子重复训练。这样设计的目的是检验前面发现的这些效应,是不是只在某个特定条件下才会出现的巧合。
结果是相当一致的。
在K=4的设置下,校准方式的选择、直通估计器的放置方式,这两个因素各自的影响方向,以及它们的相互作用方向,在10亿参数规模和五个随机种子上都保持不变。用Kendall's W这个指标衡量五个种子之间条件排序的一致性,得到的数值是0.91,这是一个相当高的一致性水平。平均下来,把校准方式从MinMax换成FWM能让损失降低0.067纳特,把直通估计器从Weight-STE换成Prob-STE能降低0.021纳特,两者叠加还有额外的0.031纳特收益,说明这两个改动确实存在替代关系而不是简单叠加关系。
同时,论文也观察到,随着K增大,这些效应本身的绝对大小都在缩小。当K达到16以上时,FWM系列的条件已经落在了不同随机种子之间的自然波动范围之内,换句话说,这时候操作方式的选择已经不太重要了,精细的查找表本身就足够好用了。
论文还专门测了一组更极端的对照,把指数运算直接换成BF16或者FP8这种更低精度的浮点格式,不用任何K区间查找表,看看单纯降低数值精度是不是也会出问题。结果这两个对照组的损失差距都在0.01纳特以内,远远小于粗糙查找表方案的差距。这说明问题的根源不是"精度低"这么简单,而是"查找表的分辨率会不会跟着训练动态跑偏"这件事,固定格式的浮点数不会像MinMax那样越训练越粗糙,所以它稳。
好消息是,只要选对了组合,近似方案完全可以逼近标准softmax的表现。论文里表现最好的几组配置,比如LERP方法在K=4、16、32这几个设置下,在25亿token训练结束时和标准softmax的损失差距都控制在0.005纳特以内,FWM配合Weight-STE在K=16时差距是0.004纳特。这些数字小到什么程度呢,换算成困惑度(一种更直观衡量语言模型好坏的指标),差距只有大约0.4%到0.5%,基本可以认为训练效果和标准做法没有实质区别。
下游任务:小的损失差距不等于处处都没事
论文没有止步于验证损失这一个指标,还专门做了一批下游任务评测,涵盖BLiMP语法判断、HellaSwag常识推理、PIQA物理常识、ARC-Easy科学问答、LAMBADA长文本理解、SciQ科学问答等七种配置。
结果印证了一个很朴素但常被忽略的道理:验证损失差距大的条件,几乎在所有下游任务上都会掉分。论文统计了35组"大差距条件对比标准softmax"的比较,全部呈现负向差距(也就是表现变差),其中27组差距在统计上是显著的。表现最差的MinMax-Weight在K=4时,准确率损失最多达到19.4个百分点,这是一个相当惊人的跌幅,相当于每5道题里,这个模型比标准softmax多错一道。
但反过来,验证损失差距很小的那些"看似安全"的条件,在下游任务上也不是完全没有波动。论文统计了35组"近乎持平条件"的比较,有8组达到了统计显著的水平,其中6组是变差(集中在BLiMP和LAMBADA上),2组反而是变好(集中在SciQ的两个配置上)。这说明,验证损失这一个数字,并不能完全代表模型在所有具体任务上的真实能力,近乎持平不代表处处等价。
论文还做了一个很值得玩味的对照实验。他们把用近似算子训练出来的模型权重,在评测时强行换回标准softmax来计算,结果发现42组对照里有37组差距反而变大了,有的甚至从原本几乎持平的0.005纳特一下跳到1.814纳特。这说明模型的权重本身已经"适应"了训练时使用的那套近似算子,你没法简单地在训练完之后再换回标准算子,期待获得更好的效果,模型已经和它的训练环境绑定在了一起。
这就好比一个人常年在高海拔地区生活,身体已经适应了那里稀薄的氧气,你把他突然空运到海平面,以为氧气充足了应该表现更好,结果他反而因为身体调节机制紊乱而更难受。适应是双向的,模型学到的不只是"任务本身",还包括"在什么样的计算环境里完成任务",硬生生换掉这个环境,收益不但没有,反而可能是负的。
写在后面
读完这篇论文,最让我意外的其实不是那些数字,而是论文反复强调的一句话:同一个前向计算,配上不同的反向传播规则,能训练出完全不同的模型。
这句话乍一听像是废话,梯度不同当然结果不同。但真正让人后背发凉的地方在于,detach这个操作根本不改变前向计算的任何一行代码,它改变的只是反向传播时一个几乎看不见的开关。如果你只检查模型的前向输出对不对,你永远发现不了问题,因为前向计算真的完全一样。这种"表面正常,内在已经出问题"的状态,可能持续训练两三千万个token都毫无异样,直到某个临界点突然爆发。这种延迟暴露的特性,对任何做工程优化的人来说都是个警钟:你以为的"无害简化",可能正在悄悄改写整个学习过程的方向,而你手头的监控指标暂时还看不出来。
另一个让我反复琢磨的细节是那个"零和投影"的补救实验。研究者发现detach破坏了梯度的一个数学性质(所有位置的梯度加起来应该等于零),于是很自然地想,那我把梯度强行投影回这个性质成立的子空间,是不是就能修好?结果这个补救版本反而比什么都不做更差。这提醒我们,一个看似优雅的数学修补,不一定对应着真实世界里坏掉的那个机制。有时候诊断对了症状,开的药方却可能是错的,因为症状背后可能有不止一个原因在起作用。
论文自己也很诚实地承认,这个范围-分辨率反馈假说存在反例,没有一个干净利落的因果证明。这种坦诚在学术论文里其实不算常见,大部分研究会倾向于把故事讲圆。但正是这种不完美,让人更愿意相信这些数据是真实的,而不是被精心挑选出来配合某个漂亮结论的。
这篇论文留下的问题,可能比它回答的还多。比如,如果换成更极端的4比特查找表会怎样?如果换一种完全不同的自适应校准策略,是不是还会出现类似的范围膨胀现象?这个问题会不会在更大的模型、更长的训练周期里表现出完全不同的模式?这些都还没有答案。
Q&A
Q1：K区间注意力是什么？
A：K区间注意力是论文提出的一套近似softmax的操作族,用K+1个预先计算好的查找表数值,通过插值或四舍五入来近似指数运算,从而在注意力计算中替代标准softmax,主要用于研究不同校准和重建方式对模型训练效果的影响。
Q2：为什么detach行梯度反而会让训练变差？
A：因为在近似算子里,一行分数的最大值和最小值不再只是数值稳定性技巧,而是直接决定查找表精度的关键参数。掐断它们的梯度会让模型持续调整这两个值却得不到正确反馈,导致查找表逐渐失真,训练在两三千万token后突然出现验证损失大幅上升。
Q3：用粗糙的softmax近似训练大模型安全吗？
A：要看具体配置。论文发现固定窗口锚定配合归一化后放置直通估计器的组合,在多种规模和随机种子下都能把损失差距控制在0.005纳特左右,接近标准softmax;但校准方式选错、直通估计器放置位置不当,损失差距可能高达0.89纳特,下游任务准确率最多掉19个百分点。
