1. 从一次训练发散说起QK 保险丝到底在保什么很多人第一次听到“QK 保险丝”这个说法会以为是某种硬件保护机制其实它说的是注意力计算里对 query 和 key 点积结果做约束的一类手段。核心场景很具体当序列变长、头数变多、训练步数拉长之后注意力 logit 的数值会逐渐失控轻则 loss 抖动重则直接发散成 NaN。所谓“保险丝”就是在 logit 进入 softmax 之前给它加一道限制让数值别飙到不可收拾。我最初接触这个问题是在一个中等规模的语言模型预训练任务上。模型参数量不算大单卡能跑序列长度 2048头数 16。前几千步一切正常到某个节点之后梯度范数突然跳了一个数量级紧接着 loss 直接变成 nan。排查了很久从学习率、数据、初始化一路查到注意力模块最后定位到 attention logit 的极值已经涨到了几百的量级。softmax 在这种输入下会退化成近似 one-hot梯度几乎全部分配给单个位置反向传播时数值极不稳定。这就是 QK 保险丝要解决的问题。它不是一个单点技术而是一族思路QK-norm、softcap、以及把两者结合并随训练进程调整的退火 QK-norm。三者解决的是同一个病根但用药方式和副作用各不相同。小模型要不要装装哪一种什么时候装是这篇要拆开讲清楚的事。先把结论方向摆出来方便你带着判断往下读小模型并非“参数量小所以更稳”恰恰相反小模型在注意力数值稳定性上往往更脆弱因为它的容量小、冗余少一旦 logit 分布畸形没有足够的其他维度去补偿。所以我的基本立场是——小模型值得装保险丝但装法要讲究盲目照搬大模型配置反而会拖慢收敛。2. attention logit 为什么会失控数值层面的病根拆解2.1 点积的方差随维度线性增长注意力 logit 的计算是 query 和 key 的点积再除以 sqrt(d_k)。问题出在点积本身如果 query 和 key 的每个分量都是独立、均值 0、方差 1 的随机变量那么 d_k 维点积的方差就是 d_k。也就是说头维度越大logit 的数值范围天然越宽。除以 sqrt(d_k) 这个缩放操作本意就是把方差拉回到 1 附近。理论上没问题但前提是 query 和 key 的分布真的接近标准正态。实际训练中随着参数更新query 和 key 的分布会漂移方差可能远大于 1缩放因子就压不住了。头维度 64 的时候logit 标准差偏个两三倍极值就能到几十头维度 128情况更明显。2.2 长序列放大了极值出现的概率logit 是一个矩阵形状是序列长度乘序列长度。序列越长这个矩阵里的元素越多出现极端值的概率就越高。这跟“样本越多越容易见到极端事件”是一个道理。2048 长度的序列注意力矩阵有四百多万个元素只要分布尾部稍微厚一点就一定会冒出几个数值离谱的 logit。这些极端 logit 经过 softmax 之后会把对应位置的注意力权重推到接近 1其他位置接近 0。前向看好像只是“注意力集中”但反向传播时softmax 的雅可比矩阵在这种饱和区几乎为零梯度传不回去。更糟的是如果极端值出现在多个头、多个位置梯度会出现剧烈的稀疏化优化过程直接失稳。2.3 小模型为什么反而更容易中招这里有个反直觉的点。大模型头数多、维度大按理说 logit 更容易失控但实际训练中大模型往往有更多的正则手段、更精细的学习率调度、更充分的数据而且它的容量足以让注意力分布学得比较“平滑”。小模型没这些条件。小模型的隐藏维度小query 和 key 的投影矩阵表达能力有限容易出现某些维度被过度激活、某些维度长期沉默的情况。这种不均衡会让点积的方差估计偏离理论值缩放失效。再加上小模型通常训练数据量也小分布覆盖不足遇到没见过的模式时 logit 更容易爆。所以“小模型不用装保险丝”这个判断从数值稳定性角度看是站不住的。2.4 一个可复现的观察方法想确认自己的模型有没有这个问题不用等它发散。训练过程中定期记录 attention logit 的绝对值最大值和分位数画一条曲线。如果最大值随训练步数持续攀升或者 99.9 分位数开始明显抬头那就是预警信号。我一般会在验证阶段顺手统计一下成本很低但能提前几百步发现问题。提示统计 logit 时记得在缩放之后、softmax 之前取这才是真正进入 softmax 的数值。缩放之前的点积极值没有直接意义。3. QK-norm 的工作方式与它的隐藏代价3.1 QK-norm 到底 norm 了什么QK-norm 的做法很直接在算出 query 和 key 之后分别对它们做归一化通常是用 RMSNorm 或 LayerNorm然后再做点积。这样一来query 和 key 的每个向量都被约束到单位尺度附近点积的数值范围就被压住了不再随头维度和参数漂移而失控。它的逻辑跟“把两个向量都归一化到单位球面上再算相似度”是一样的。余弦相似度的取值范围天然是负一到一乘上温度系数之后也不会爆。QK-norm 相当于把注意力相似度计算从“无界点积”改成了“有界余弦型”从根上掐掉了 logit 失控的可能。3.2 为什么它比单纯调缩放因子更彻底有人会问既然问题是缩放不够那把 sqrt(d_k) 换成一个更大的常数不就行了。这个思路在小范围内有效但它治标不治本。因为 query 和 key 的方差是随训练动态变化的你固定一个缩放因子只能适配某个阶段的分布。训练早期合适后期可能就压不住了反过来早期又可能压得太狠导致注意力过于均匀、学不到东西。QK-norm 是自适应的。不管 query 和 key 的原始尺度怎么变归一化之后都回到统一尺度。它把“缩放多少”这个需要手工调的超参数变成了一个由数据自身决定的归一化操作。这是它比固定缩放更稳的根本原因。3.3 隐藏代价表达能力的损失但 QK-norm 不是没有代价。归一化会丢掉向量的模长信息。原本 query 和 key 的模长可以携带“这个位置有多重要”的信号归一化之后这个信号没了所有位置的向量都被拉到同一尺度。对于需要精细区分注意力强度的任务这可能是个损失。我在一个小规模检索式任务上对比过加 QK-norm 之后训练确实更稳但最终指标比不加的时候低了大概一个百分点。分析下来就是模型失去了用模长表达重要性的能力只能靠方向来区分表达空间变窄了。这个代价在小模型上更明显因为小模型本来容量就紧张再砍掉一个自由度影响会被放大。3.4 什么情况下 QK-norm 是划算的判断标准其实不复杂。如果你的训练经常出现 loss 尖峰、梯度爆炸或者序列长度超过 1024、头维度超过 64那 QK-norm 带来的稳定性收益通常大于表达损失。反过来如果训练一直很稳序列也不长那加它可能只是徒增约束。还有一个折中做法只对 query 做 normkey 不动或者反过来。这样能保留一半的模长信息稳定性收益也能拿到大部分。我试过只 norm query 的版本在几个任务上表现介于全 norm 和不 norm 之间算是个可选的中间档。4. softcap给 logit 加一道软天花板4.1 softcap 的数学形式与直觉softcap 的思路跟 QK-norm 完全不同。它不动 query 和 key而是在 logit 算出来之后用一个平滑函数把它压到一个固定范围内。最常见的形式是 tanh 型logit 经过 c 乘以 tanh(logit 除以 c)其中 c 是 cap 值。当 logit 远小于 c 时这个变换近似恒等当 logit 接近或超过 c 时输出被平滑地压到正负 c 之间。用生活化的说法QK-norm 是“把原料标准化再加工”softcap 是“加工完之后加个限流阀”。前者从源头控制后者在末端兜底。两者可以单独用也可以叠加。4.2 softcap 保住了什么softcap 最大的好处是它不丢模长信息。logit 在 cap 范围内基本原样通过模型仍然可以用 logit 的大小表达注意力强度。只有在极端值出现时才被压制。这相当于给模型留了完整的表达空间只在越界时才干预。对于小模型这个特性很有吸引力。因为小模型经不起表达能力的额外损失softcap 的“平时不管、越界才管”正好契合。我在几个小模型上对比过softcap 版本的下游指标基本和不加约束的版本持平但训练稳定性明显更好loss 尖峰少了很多。4.3 cap 值怎么选一个容易踩的坑cap 值 c 的选择是 softcap 的关键。选太大等于没加极端 logit 照样能到几十选太小正常范围内的 logit 也被压注意力变得过于均匀模型学不到尖锐的分布。我的经验是先在不加 cap 的情况下跑几百步统计 logit 的 99 分位数把 c 设在这个分位数的 1.5 到 2 倍左右。这样大部分 logit 不受影响只有尾部极端值被压。如果训练中途 logit 分布整体上移可以定期重新统计、调整 c或者干脆设一个偏大的固定值靠 tanh 的平滑性兜底。注意cap 值不要设成整数后就不管了。不同任务、不同序列长度下logit 的合理范围差别很大。我见过直接照搬别人配置设 c 等于 30 的结果在自己的任务上把正常 logit 也压了收敛慢了一半。4.4 softcap 与温度系数的关系softcap 和 softmax 的温度系数在效果上有重叠都是调节注意力分布的尖锐程度。但机制不同温度系数是全局线性缩放softcap 是尾部非线性压缩。温度系数调小能让分布整体变平但压不住极端值softcap 专门压极端值对中间部分几乎无影响。实践中可以两者配合温度系数负责整体调节softcap 负责兜底。不过要注意别过度调节两个都往平的方向调注意力会退化成近似均匀分布模型等于没在选信息。5. 退火 QK-norm把约束强度交给训练进程5.1 为什么需要“退火”QK-norm 和 softcap 都有一个共同问题约束强度是固定的。但训练的不同阶段对约束的需求是不一样的。训练早期参数随机、分布混乱logit 容易爆这时候需要强约束训练后期参数趋于稳定分布收敛过强的约束反而限制了模型的精细表达能力。退火 QK-norm 就是针对这个矛盾设计的。它的核心思想是让约束强度随训练进程从强到弱平滑过渡。早期用强 norm 或小 cap 把数值压住后期逐渐放松把表达自由还给模型。这跟学习率退火是同一个哲学——先粗后细先稳后精。5.2 退火的具体实现方式实现上有几种常见做法。一种是插值系数设一个从 1 衰减到 0 的系数 alpha最终 query 等于 alpha 乘以归一化后的 query 加上 (1 减 alpha) 乘以原始 query。训练早期 alpha 接近 1几乎完全归一化后期 alpha 接近 0回归原始 query。另一种是 cap 值退火softcap 的 c 从一个较小的值逐渐增大到目标值。早期强压后期放松。两种方式可以结合也可以单独用。退火曲线一般跟学习率曲线保持一致用余弦或线性衰减都行。5.3 退火带来的实际收益我在一个序列长度 4096 的小模型上做过对比。固定 QK-norm 的版本训练很稳但最终验证 loss 比不加约束的版本高了一点退火 QK-norm 的版本既保持了训练稳定最终 loss 又追平甚至略优于不加约束的版本。原因就是后期约束放松后模型重新获得了表达自由度。这个收益在小模型上尤其值得关注。小模型容量紧张任何长期的硬约束都是净损失。退火相当于把约束变成“临时脚手架”用完就拆不占用最终的表达预算。5.4 退火调度的注意事项退火不是越慢越好。如果衰减太慢模型长期处于强约束下等于没退火衰减太快早期还没稳住就放松了可能中途发散。我的经验是让退火周期覆盖训练的前 60% 到 70%后 30% 保持弱约束或零约束给模型充分的精细调整时间。另外退火过程中要监控 logit 的极值。如果放松之后极值重新抬头说明退得太快需要把曲线拉长。这个监控跟前面说的 logit 统计是同一套工具复用即可。6. 三种方案横向对比与选型建议6.1 关键维度对照维度QK-normsoftcap退火 QK-norm作用位置query/key 归一化logit 末端压缩归一化强度动态调整模长信息丢失保留后期恢复稳定性收益高中高高表达损失中低低实现复杂度低低中超参数敏感度低高cap 值中适合场景长序列、易发散想保表达、轻度不稳长训练、追求最终指标6.2 小模型的选型逻辑如果训练经常发散、序列长、头维度大优先考虑退火 QK-norm它兼顾稳定和最终指标。如果训练基本稳、只是偶尔有尖峰softcap 是性价比最高的选择实现简单、表达损失小。如果任务对注意力强度的精细区分要求不高固定 QK-norm 最省事。我的默认推荐是小模型先上 softcapcap 值按 logit 分位数定如果还压不住再叠加退火 QK-norm。这个组合的稳定性上限最高且表达损失可控。6.3 一个容易忽略的交互与位置编码的关系QK 保险丝的效果跟位置编码方式有交互。旋转位置编码会把位置信息编码进 query 和 key 的相对关系里归一化之后这种相对关系可能被削弱。我在用旋转位置编码的模型上加 QK-norm 时发现需要把 norm 的 epsilon 调大一点否则位置信号的区分度会下降。这个细节在文档里很少提但实测有影响。7. 实操落地从监控到上线的完整链路7.1 第一步建立 logit 监控在注意力模块里加一个统计钩子记录缩放后 logit 的绝对值最大值、99 分位数和标准差。频率不用太高每几百步记一次就够。这些数据是后续所有决策的依据。# 伪代码示意在注意力前向中插入统计 logit torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) with torch.no_grad(): stats[logit_max].append(logit.abs().max().item()) stats[logit_p999].append(torch.quantile(logit.abs().flatten(), 0.999).item())7.2 第二步根据监控结果决定方案如果 logit 最大值长期稳定在个位数到十几不用加保险丝。如果持续攀升到几十先上 softcap。如果 softcap 之后仍然有尖峰或者序列长度超过 2048上退火 QK-norm。7.3 第三步参数标定softcap 的 cap 值按前面说的分位数法标定。退火 QK-norm 的退火周期按训练总步数的 60% 到 70% 设定。归一化用 RMSNorm 还是 LayerNorm小模型上差别不大RMSNorm 计算略省可以优先。7.4 第四步上线后的回归验证加了保险丝之后除了看稳定性一定要对比下游指标。稳定性提升但指标下降说明约束过强需要放松。我一般会跑一个不加约束的基线和一个加约束的版本对比验证 loss 曲线和几个关键任务的指标确认收益是净正的。提示对比时注意控制变量学习率、数据顺序、随机种子尽量一致否则指标差异可能来自其他因素误导判断。8. 几个我踩过的坑和对应的处理第一个坑是 cap 值照搬。早期我直接用了别人论文里的 cap 值结果在自己的任务上把正常 logit 也压了模型收敛慢指标还降了。后来改成按自己的 logit 分布标定问题解决。教训是任何跟数值范围相关的超参数都必须基于自己的数据重新标定。第二个坑是退火退得太快。有一次把退火周期设成了总步数的 30%结果中期 logit 重新抬头训练在后期发散。把周期拉长到 65% 之后稳定了。退火速度要跟训练总长度匹配短训练可以快一点长训练必须慢。第三个坑是只监控 loss 不监控 logit。loss 是滞后指标等 loss 出问题的时候logit 已经失控很久了。加上 logit 监控之后能在发散前几百步就发现苗头提前干预。这个习惯我后来在所有注意力相关的训练里都保留了。第四个坑是忽略了位置编码的交互。用旋转位置编码时加 QK-norm没调 epsilon导致位置区分度下降模型在长距离依赖任务上表现变差。调大 epsilon 之后恢复。这个细节不显眼但影响实在。9. 关于“要不要装”的最终判断回到标题的问题。小模型要不要装 QK 保险丝我的答案是看你的训练稳不稳而不是看模型大小。稳就不装不稳就装装的时候优先选表达损失小的方案并且尽量用退火把约束变成临时的。小模型因为容量紧张对表达损失更敏感所以更不适合长期硬约束。但这不意味着不该装而是说装的方式要更精细。softcap 加退火 QK-norm 的组合是我目前在小模型上找到的平衡点稳定性够用表达损失可控实现也不复杂。最后分享一个我自己的习惯每次开新训练任务先跑一个短程的探针实验几百步就够专门看 logit 的走势。如果走势平稳直接上正式训练如果有抬头趋势先加保险丝再上。这个探针实验花不了多少时间但能省下大量排查发散的时间。踩过几次发散的坑之后我现在已经把这个步骤固化成流程的一部分了。
