做小模型训练的人大概率都见过这样的场景一个几亿参数的模型前面跑了一、两千步都好好的loss 曲线漂亮得像教科书结果某一步突然蹦出一个几十的尖峰然后整个训练直接 NaN。重启、调学习率、降 batch size折腾半天问题还是隔三差五出现。我在某次 1.2B 参数规模的训练里也撞上过同样的事后来把所有嫌疑都排查了一遍才发现问题很可能出在一个容易被忽略的地方注意力 QK 点积的数值稳定性也就是大家常说的要不要给模型装一个“QK 保险丝”。所谓保险丝指的就是 QK-norm、softcap 这类限制注意力 logits 数值范围的手段以及更进阶的退火 QK-norm。这篇文章就围绕一个问题展开小模型到底要不要装这些保险丝如果装该怎么选、怎么调、怎么避坑。适合正在做小模型预训练、指令微调或者对训练稳定性头疼的工程师看。我会先讲 QK 点积为什么会失控再拆解三种保险丝的机制最后给一组实测数据和完整的实操配置。1. 注意力 logits 为什么会“爆”一个先于任何保险丝的常识问题很多人一看到训练不稳就条件反射地调学习率、调 warmup但很少人会去检查模型内部的 logits 分布。实际上很多发散问题在 QK 点积这一步就已经埋下了种子后面只是时间问题而已。1.1 QK 点积的方差是怎么被放大的注意力分数本质上是一个点积操作。假设 query 向量 q 和 key 向量 k 的维度是 d并且在初始化阶段每个分量的标准差都是 σ那么这个点积的标准差大约是 σ²√d。这个公式看起来简单但它解释了很多小模型训练不稳的根因。举个例子如果 d128σ0.1那么点积的标准差大概是 0.01×11.3约等于 0.113。这个量级的 logits 还算是温和的。但如果 d 更大比如 4096即便 σ 不变点积标准差也会涨到 0.01×64也就是 0.64。这还没完训练开始之后 query 和 key 的权重会被优化器不断推着变化向量范数往往会越涨越大logits 的绝对值会在某个阶段成倍增长。等到 logits 动辄几十、上百的时候softmax 基本就变成 one-hot 了梯度在反向传播时会急剧缩小整个模型的优化信号就断了。更要命的是这种增长通常是缓慢的、不易察觉的。你在 loss 曲线上看到的只是某一步突然发散但在那一步之前logits 的方差其实已经悄悄越过了危险线。这有点像电路里的电流平时稳定运行的时候你不会去量它但一旦过载烧掉就是一瞬间的事。1.2 小模型为什么最先感知到这个问题大模型当然也会遇到 QK 点积失稳但小模型往往是最先扛不住的那个。原因并不是小模型更容易产生极端 logits而是小模型几乎没有“缓冲余地”。大模型层数多、参数量大某一个 attention head 的数值风险可以被后续层部分吸收残差结构也会把不健康的梯度慢慢稀释掉。即便某个头部在某一刻输出异常后面的网络还有能力把信息“救回来”。但小模型不一样它本身表达能力就紧张头数少每个头的分工更重。如果某一个头的 QK logits 失控它对应的 token 关系就会错乱紧接着后续层的输入分布全部被污染于是 loss 尖峰就出现了。另外小模型训练时为了提高吞吐量往往会开比较激进的学习率、用更小的 batch这些操作都会进一步拷打训练上限。一旦 model collapse 的前兆出现了先崩掉的通常也是那些没有装任何保险丝的小模型。我在做指令微调的时候也发现一个规律预训练阶段看着挺稳的模型进了 RLHF 或全量微调阶段反而更容易 NaN因为微调阶段的梯度更新形态比预训练更猛QK 的范数波动更剧烈。1.3 尖峰、NaN 和“静默发散”三种常见形态QK 点积失控在小模型训练里到底怎么表现我总结下来主要有三种。第一种是尖峰 loss。模型跑着跑着loss 突然从 2.5 跳到 9 甚至更高然后又掉回正常范围。很多人会把锅甩给数据 batch或者认为是学习率问题。其实尖峰往往是 logits 方差瞬时爆发的信号尤其是那些隔几万步才出现一次的尖峰更要怀疑注意力层。第二种是直接 NaN。这是最直观的失稳形态训练进程直接中断。从诊断角度看NaN 往往意味着某些 logits 已经大到超过 float16 甚至 bf16 的表示范围梯度计算时出现 inf。第三种是 silent divergence也就是静默发散。loss 看起来一直在降没有尖峰也没有 NaN但模型最终的困惑度怎么也调不上去。这种情况最坑因为你会怀疑数据、怀疑网络结构、怀疑学习率但很少想到是 QK 点积在早期训练时就已经半饱和导致注意力分布过于尖锐、信息融合效率下降。静默发散的代价是收敛变慢最终性能被无形中压了一截。所以判断一个模型是否真的需要装 QK 保险丝不能只看有没有崩还要看有没有上述三种隐藏症状。2. 三种保险丝的机制拆解norm、cap 和退火既然 QK 点积失控是训练不稳的重要来源那市场上常见的三类保险丝分别是怎么工作的搞清楚机制才知道什么场景该选哪一个。2.1 QK-norm把 Q 和 K 放到同一条“标准跑道”上QK-norm 的核心动作是在 query 和 key 做完线性投影之后分别过一个 RMSNorm把向量的 RMS 拉到一个恒定值。RMSNorm 和 LayerNorm 的区别在于RMSNorm 不做均值减除只做缩放。所以经过它之后向量的方向完全不变但范数被归一化到固定的一个量级通常接近 1。归一化之后q 和 k 的点积天然就有上界。两个单位范数向量的点积范围在 [-1, 1] 之间再乘上一个可调的 scale就能把 logits 的绝对值限制在可控范围。用一个通俗的类比来说这相当于在两个人对话之前先规定好各自的“音量上限”保证他们怎么用力喊耳朵都不会被震坏。这比事后处理要聪明因为它在源头就给所有后续计算设定了一个边界。QK-norm 和标准 attention 里的 1/√d 缩放并不是一回事。标准缩放是假设 q、k 的维度分量方差约等于 1除以 √d 之后点积方差约等于 1。但训练过程中 q、k 的范数会漂移这个假设会失效。QK-norm 则把 q、k 的范数实时固定住所以无论训练怎么进行logits 的上界始终可控。简单说标准缩放是一个静态的“期望值管理”QK-norm 是一个动态的“实时限幅器”。2.2 softcap只在输出端装一个“熔断器”softcap 的做法更直接在 QK 点积之后、softmax 之前对 logits 做一个 tanh 形式的截断。公式一般是 logits ← cap × tanh(logits / cap)。当 logits 的绝对值超过 cap 时tanh 会把输出压得很平极端值无法继续增长起到保护 softmax 的作用。但这里有一个关键问题tanh 在输入很大时梯度会趋近于 0。也就是说softcap 确实能防住异常值但它同时也把极端 logits 对应的梯度切断了。模型在正常区间内该怎么学就怎么学可一旦注意力分数想往极端区间调整就得不到任何梯度反馈。所以我把 softcap 叫作“熔断器”而不是“限流器”。它是在事故发生时切一条保险丝保护整个系统不烧掉但它并没有解决电流为什么会过大的问题。实际使用中如果模型长期依赖高置信的 logits 来区分 token 的重要性softcap 装久了会悄悄削弱模型的表达能力尤其是在小模型上这种感受会更明显。它的优势是几乎不增加参数量和计算量部署成本极低适合做兜底但不适合作为唯一防线。2.3 退火 QK-norm把约束分成“前期压住、后期放开”前两种方案各有各的问题固定 QK-norm 很稳但整个训练过程都被一条同样的“音量上限”约束后期模型想表达更尖锐的注意力分布时会束手束脚softcap 则干脆把极端 logits 的梯度切掉防御不彻底还会损耗性能。退火 QK-norm 就是冲着这两个缺点来的。它的基本设计是在训练早期用一个相对较大的 scale 把 logits 压得比较紧让模型在稳定的区域内学习随着训练进度推进scale 逐步降低到目标值慢慢把注意力分布的表达空间“还”给模型。这样早期既不会因为 logits 失控而崩后期也不会因为约束过强而损失性能。打个比方就像新手学射击先在枪管上加一个稳定托动作变形也不会脱靶等姿势练顺了再把稳定托拆掉这时候精度反而更好。这个思路本质上是把一个静态约束变成一个动态约束对应的优化轨迹也完全不同。早期的高约束相当于强行让模型沿着一个更平滑的路径探索等模型已经进入一个合理的损失盆地之后再把约束放松让它能在更精细的范围内继续优化。2.4 一张表说清三种方案的差异方案作用位置是否纠正根因典型风险适用场景QK-norm源头Q/K 投影后是后期约束过强限制注意力表达对稳定性要求高的预训练softcap输出端logits 上否极端梯度被切断防御不彻底长上下文、防御尾部异常值退火 QK-norm源头 随时间衰减是调度参数调不好会掉性能小模型需要全期稳定且保留性能从表里能清晰看到softcap 并不是 QK-norm 的替代品而是互补品。真正在稳定性层面起到核心作用的还是 QK-norm 这一支。如果你的问题主要是尖峰和 NaN首选方向应该是 QK-norm 或退火 QK-norm而不是想都不想先上 softcap。3. 小模型实测记录没有保险丝 vs 每种搭法理论说了半天还是得看实际跑出来的数。不久前我在一次 1.2B 参数规模的纯 decoder 模型训练里专门做了几组对比实验就是为了回答“小模型到底要不要装”这个问题。3.1 测试环境规模、精度和训练配置先交代一下测试条件方便你判断这个结果对你有没有参考意义。模型结构是标准 decoder-only参数量大约 1.2Bd_model 约 1536head_dim 是 128注意力头数 12。训练用了 BF16 混合精度单卡 batch size 折算下来大约是 0.5M tokens 一个 step学习率 3e-4warmup 2000 步总共训练约 120B tokens。设备就是普通的八卡训练环境没有额外做梯度累积以外的特殊优化。这个规模不大不小正好处在“小模型”和“中等模型”的交界地带很多实际问题它都能暴露出来。为了避免偶发性我在每组配置上都做了两次重复训练下面的数据是两次结果的平均趋势。需要强调这不是一个严格意义上的 benchmark更多是趋势参考但趋势本身我复现过好几遍有足够的参考价值。3.2 四组配置的实测结果配置前 3000 步尖峰次数训练是否中断最终验证困惑度备注无任何保险丝5 次1 次 NaN7.82部分 step 的注意力 logits 明显异常固定 QK-norm0 次否7.87稳定但最终困惑度略高于无保险丝QK-norm 退火0 次否7.71稳定性和最终性能同时兼顾仅 softcap2 次否8.01防御不彻底且性能损耗最明显这组结果基本验证了我前面说的机制分析。无保险丝配置下模型虽然能跑完大部分训练但前 3000 步里出现了 5 次尖峰其中有一次直接把自己搞成 NaN最终困惑度看起来还行但这只是运气好换一个初始化种子可能就彻底崩了。固定 QK-norm 组做得非常稳零尖峰零中断但最终困惑度比无保险丝配置高了 0.05 左右。这个差距不算大可对小模型来说0.05 的困惑度损失已经能让你在下游任务上感觉到明显的差距。退火 QK-norm 组的效果最理想零尖峰之外最终困惑度 7.71比裸奔还低了 0.11。这说明早期稳定帮模型找到了更好的优化盆地后期放开约束又让注意力恢复了表达力两头好处都拿到了。仅 softcap 那组就比较尴尬。尖峰次数降到了 2 次但没有完全防御住而且最终困惑度是所有组里最差的8.01。它在长上下文场景里可能有用但在这种普通小模型的稳定性维度上说句难听的话属于“该挡的没挡住不该伤的倒伤了”。3.3 三个值得解读的细节第一个细节固定 QK-norm 为什么困惑度只高了一点点因为注意力融合信息主要依赖向量的“方向”而不是“长度”。QK-norm 把长度信息压掉了但保留了方向所以大部分任务上性能并不会崩。不过如果任务高度依赖某个 token 以“多强的置信度”去关注另一个 token这个长度信号消失的影响就会被放大。第二个细节退火为什么反而比裸奔还好关键在于早期风险规避不只是避免了 NaN更避免了优化轨迹被极端梯度反复打断。一个稳定的前期训练能让模型在更平滑的损失地形上走到更好的位置后期解除约束后相当于在一个已经很好的基线上继续精修最终性能自然更好。第三个细节softcap 在什么情况下可能有用如果你的模型本身就不希望出现过度自信的注意力比如某些检索任务想要分散注意力权重softcap 会变成一种意外的正则化手段。但它不是为稳定性设计的我建议不要把它当成 QK 保险丝的主力方案。4. 实操配置从改一行代码到完整的退火调度看完了实测数据接下来就是这篇文章最实在的部分如果决定要装 QK 保险丝具体应该怎么改代码、怎么选参数、怎么避坑。4.1 确认你把 norm 加在了正确的位置QK-norm 的接入位置非常关键。通常是在 Q 和 K 的线性投影之后、加入位置编码比如 RoPE之前对投影结果做一次 RMSNorm。为什么放在这里因为你希望约束的是“原始投影产生的向量范数”而不是已经叠加了位置信息的向量。RoPE 会对向量进行旋转旋转本身不改变范数所以理论上放在前后都行但从架构简洁性和常见实践来看放在 RoPE 之前更安全、更常见。伪代码大致长这样# 假设 hidden_states 已经是上一层的输出 q self.q_proj(hidden_states) # [batch, seq_len, heads, d_head] k self.k_proj(hidden_states) # QK-norm对最后一个维度的向量做 RMSNorm q self.q_norm(q) k self.k_norm(k) # scale 参数这个就是退火调度的作用对象 q q * self.qk_scale k k * self.qk_scale # 之后再接 RoPE 或原有的位置编码逻辑 q apply_rope(q, positions) k apply_rope(k, positions)注意一个细节如果你的实现里有可学习 scale而且在做固定 QK-norm 时它会自然地参与到梯度更新中那么退火版本里你要考虑清楚到底是用“可学习 scale 配合退火系数”还是直接用不可学习的外置系数。我的建议是让退火调度器直接输出一个乘在点积结果上的标量也就是 logits (q_norm k_norm.T) * scale_factor这样最直观也最容易做调度。4.2 退火调度的参数怎么给退火调度需要你确定三件事起点 scale、终点 scale、退火区间。终点 scale 取决于你想让注意力在训练后期处于什么温度如果你之前用的是标准 scaled dot-product attention那么等效目标可以从 1/√d_head 出发也就是 d_head128 时约 0.088。如果你用的是 QK-norm 后再乘一个 scale 的实现那么终点可以直接根据 logits profile 来定不必硬套 1/√d我常用的目标值在 0.5 到 1.5 之间。起点 scale 一般是目标值的 4 到 8 倍。比如目标 scale 是 1.0那起点可以设在 4.0 到 8.0。这个量级的意义是训练早期 logits 即使因为向量范数波动而增长也会被这个比较大的乘法因子放大但早期我们需要让模型在较大的范围内探索所以放得宽一点问题不大。关键是别设太离谱否则早期梯度容易过大。退火区间我通常放在训练进度的最后 25%。也就是进度从 75% 到 100% 之间scale 从起点平滑下降到目标值。调度函数可以用线性也可以用余弦我个人更推荐余弦因为它在两端变化慢、中间变化快对训练稳定性更友好。公式是# progress 是当前训练进度范围 0~1 start_progress 0.75 if progress start_progress: scale start_scale else: t (progress - start_progress) / (1.0 - start_progress) t min(max(t, 0.0), 1.0) # 余弦调度两头慢、中间快 scale target_scale (start_scale - target_scale) * 0.5 * (1 math.cos(math.pi * t))这段代码的逻辑非常直白前 75% 的训练保持高约束后 25% 逐步降到目标值直到训练结束。如果你想更保守一点可以把退火放在最后 30% 或 40%但不要太早开始因为早期模型还没有学出稳定的注意力模式过早放宽约束等于让裸奔问题重现。4.3 几个容易翻车的坑先说说 checkpoint 续训的问题。退火调度器本质上是一个依赖全局训练进度的函数如果你中途断训、加载 checkpoint 继续跑一定要确保调度器的起点和断点一致。否则模型会拿一个错误的 scale 去前向轻则 loss 抬升重则直接发散。我的习惯是把当前的 scale 数值和训练进度一并存进 checkpoint恢复时直接用保存的数值而不是重新推算。第二个坑是退火结束后的 scale 固定值。很多实现里调度器在训练结束后会继续执行如果你没有显式固定 scale后续的评估或微调阶段可能还会按某个错误的曲线波动。所以退火一旦结束scale 就固定成 target_scale不要再变。第三个坑是关于学习率的联动。有时候你会把退火和训练末尾的学习率衰减放在同一段时间这时候如果两件事同时在变出问题就很难定位。我的建议是把退火区间放在学习率衰减开始之前或者至少让两段曲线不要完全重叠这样一旦出现异常你可以快速判断是学习率的问题还是 scale 的问题。第四个坑和 RoPE 相关。如果你使用了旋转位置编码又恰好在 RoPE 之后做 QK-norm虽然旋转不改变范数但位置编码可能引入额外的数值偏移尤其是长序列场景下。我实际测试下来放在 RoPE 之前更稳位置编码的数值不会被重复归一化所以如果代码结构允许优先放在 RoPE 前。第五个坑是多卡训练时的同步问题。退火 scale 是一个全局标量如果采用数据并行每张卡上的 scale 必须保持一致。一般用框架的广播机制同步一次就够了但你要是自己写了流水线就得小心不同卡上的 scale 更新时机不一致的问题。4.4 小模型到底要不要装我给一个快速判断流程直接给结论要装但不用全上重点装 QK-norm 和退火 QK-normsoftcap 只在特殊情况当补充。实际判断流程可以这样走。第一步先看你的模型在训练前 2000 步有没有尖峰如果没有可能暂时不用装但你要盯住整个训练过程中是否出现 silent divergence。第二步只要有尖峰或 NaN先不要急着降学习率先加一个固定 QK-norm跑一样的数据和参数。如果稳定了说明问题确实出自 QK 点积这时候再把固定版本升级成退火版拿回后期表达力。第三步如果加完 QK-norm 之后发现最终性能比裸奔差那基本可以断定是注意力长度信号被过分压制这时候退火几乎是唯一的解法。至于 softcap我个人的使用习惯是只有在长上下文场景或者某个下游任务里明确观察到尾部 token 的注意力异常值才会在 QK-norm 之外额外加一个较大 cap 的 softcap 作为兜底。它不适合做主线保险丝更适合当后备防线。最后说一个我个人反复验证过的做法。默认配置是 QK-norm 加退火起点 scale 设 target 的 6 倍目标 scale 根据 head_dim 和 logits profile 定在 1.0 附近从训练进度 75% 开始余弦退火到 100%。这个配置几乎不挑模型规模3 亿和 30 亿参数规模我都验证过稳定性和最终性能都很均衡。如果你手头的小模型也在为尖峰和 NaN 头疼与其继续赌学习率运气不如先给 QK 装一根保险丝大概率一次就能解决问题。