最近和几个做预训练的同行聊天，话题总绕不开一个小模型稳定性的问题：训练一版 0.5B 规模的模型，loss 曲线一开始挺正常，跑到十几万步突然出现尖峰，甚至直接发散。查下来发现是注意力 logits 的数值分布出了问题。然后就会有一个经典争论冒出来——要不要给这个小模型装上 QK 保险丝（QK-norm、softcap、退火 QK-norm）？很多人的第一反应是“那是大模型才需要的东西，我们小模型加它干嘛”，但我实测下来的感受是：这个问题的答案没那么简单，小模型反而更需要想清楚“什么时候装、装哪种、装到什么强度”，而不是一刀切地“装”或者“不装”。
这篇文章就当是我自己踩坑后的复盘。我会从 QK-norm 的原理讲起，对比 softcap 和退火 QK-norm 各自解决的问题，然后重点落在小模型这个场景上：哪些信号说明你需要保险丝，哪些情况加了纯属白费，以及加了之后怎么调、怎么监控、怎么排查。适合正在做 LLM 预训练、微调、蒸馏或者做长上下文工作的工程师和研究者参考。
1. 注意力“烧起来”之前发生了什么
1.1 logits 失控的底层机制
注意力机制的核心公式是 softmax(QK^T / sqrt(d))。理论上 d 开根号是个天然的缩放因子，把方差拉回 1 附近，但实际训练中这个假设经常不成立，因为 Q 和 K 的分布本身就在不断变化。随着训练步数增加，如果 Q、K 投影矩阵的谱范数持续增大，QK^T 的绝对值会跟着变大，logits 出现高方差、长尾分布，softmax 的输出就容易退化：要么集中在极少数 token 上（尖锐化），要么彻底均匀（熵坍塌）。这两种情况都不是我们希望看到的注意力行为，前者放大梯度噪声，后者让信息无法有效传递。
我观测过一个小规模模型在正常训练下的 QK 点积标准差变化：前 5000 步大概是 8 左右，比较平稳；跑到 30000 步左右，标准差开始缓慢爬升到 20；然后在某个独特样本（长文本块）上出现了一次 80 以上的 spike，训练直接崩了。这个演进过程很典型，说明 logits 数值本身存在“渐进失控”的趋势，只是很多时候被 optimizer 和 loss 的平滑性暂时掩盖了。
1.2 三类保险丝的设计哲学
QK-norm、softcap、退火 QK-norm，名字听起来都像在埋保险丝，但思路本质不同：
- QK-norm：直接约束 Q 和 K 向量的模长或分布，让点积的尺度被限制在可控范围。它的思路是“我不让源头产生大火”，属于结构性防范。
- softcap：不对 Q、K 做任何约束，而是在 logits 计算完成后用
logits = tanh(logits) * cap做一次非线性压缩。它的思路是“就算你烧起来，我也给你按回去”，属于输出端的熔断保护。 - 退火 QK-norm：动态调节 QK-norm 作用强度，一般是训练早期强约束，后期逐渐放松；或者反过来，训练早期完全不约束，后期等注意力出现坍缩趋势时再逐步收紧。它的思路是“保险丝应该按需接入，而不是全程固定”。
理解这三类机制之间的区别之后，才能回答“小模型装哪个”的问题。很多开源代码里默认不加任何保险丝，不是因为不需要，而是因为再小的模型都有概率遇到训练不稳定，只是触发概率和使用场景不同。
2. QK-norm 到底在约束什么
2.1 从 RMSNorm 到 QK-norm 的落地细节
大家常说“QK-norm”，其实多数实现是给 Q 和 K 分别过一个 RMSNorm（或者 LayerNorm），然后在后面加一个可学习的缩放参数。RMSNorm 的好处是计算量小、无均值偏移，对显存和算力的消耗几乎可以忽略。具体实现时，在 Q、K 投影之后插入：
import torch import torch.nn as nn import torch.nn.functional as F class QKNorm(nn.Module): def __init__(self, dim, scale=True): super().__init__() self.norm = nn.RMSNorm(dim) # PyTorch 2.4+ 直接支持 self.scale = nn.Parameter(torch.ones(1)) if scale else 1.0 def forward(self, q, k): q = self.norm(q) * self.scale k = self.norm(k) return q, k这个 scale 参数即“软保险丝”的软化程度。如果你把它初始化成 0.5，前期 QK 点积尺度就被压到 sqrt(d)*0.5 左右；如果初始化为 1.0，则接近标准缩放。实际训练中 scale 会被梯度推动着变化，某种程度上它学习到的就是“该给这个注意力头多少信任”。我个人的习惯是 scale 初始化为 1.0 或者稍微小于 1，然后给学习率设置相对主网络稍小的系数，避免它在训练初期变化太剧烈。
2.2 QK-norm 的边界：它管不住所有异常
QK-norm 虽然能把每个 head 的 logits 尺度压住，但注意一个细节：它约束的是单头内部 Q、K 的全局尺度，并不能直接解决跨样本的极端离群点问题。有些瞬间异常来自某个 token 的 embedding 本身数值巨大或 norm 异常，这时候即便过了 RMSNorm，点积也可能因为少数维度上的极端值而出现虚高。
所以在大型模型实践中，通常把 QK-norm 和别的策略（如额外的 attention dropout、参数裁剪）一起用。在小模型上，遇到这类少量极端值导致的 spike，你可以临时验证一下：把训练数据里几个疑似异常样本打印出来，算一下每个 token 的 norm，如果和历史分布的差距超过 5 到 8 倍，那问题根本不在 QK 维度，而在 embedding 的数值稳定性。这个判断很关键，因为很多人花了大力气调 QK-norm，最后发现是数据侧的问题。
3. softcap：输出端的熔断保护
3.1 tanh 压缩的数值行为
softcap 的做法很简单：计算出 logits 之后，先除以 cap，经过 tanh，再乘回 cap。公式可以写成logits = cap * tanh(logits / cap)。
这个操作的好处有两个：第一，logits 的绝对值永远不会超过 cap，从数学上彻底杜绝了极大值穿透 softmax；第二，tanh 在零附近近似线性，意味着正常数值范围的 logits 几乎不受影响，只有超出 cap 的值才会被非线性压缩。实践里 cap 常见取值为 30、50 或者 100，具体看任务和模型尺寸。cap 太大会失去保护意义，太小则会让注意力在训练后期变得“僵硬”，头部之间的辨识度下降。
举个例子：一个 0.8B 参数模型，在前 20 万步用 cap=50 的 softcap，loss 曲线非常平稳；后来我把 cap 改到 30，曲线出现轻微震荡，但模型在评测集上的困惑度反而微降。这说明 softcap 的“手感”是非常依赖任务阶段的，需要在稳定性和表达力之间找平衡。
3.2 和 QK-norm 组合使用时的注意点
很多开源模型其实是“QK-norm + softcap”一起用的，比如一些采用规范化注意力的大模型。两者并不冲突：QK-norm 压缩常规尺度，softcap 兜底极端异常。但小模型加这套组合的时候要特别留意训练动态：如果 Q、K 在 RMSNorm 之后又接了一个可学习 scale，这个 scale 已经被 push 到比较大的值，那么 softcap 相当于在“scale 刚大起来”的边缘做压制，效果可能会互相抵消。
我踩过的一个坑就是：把 scale 初始化为 1.0，加了 softcap=30，训练 2 万步之后发现 scale 涨到了 4.0 以上，QK 点积标准差却又回落到 5 左右。等于模型为了规避 softcap 的限制，主动去压缩了 QK 的表达，这种做法让注意力变得更加同质化，信息利用效率反而变差。后来我把 scale 的初始值改成 0.6，同时把 scale 的权重衰减调大，情况就正常了。
所以如果决定组合使用，我建议：
- 监控 QK 点积的实际标准差，而不是只关注 loss。
- 给 QK-norm 的 scale 加一个合理的上界，比如设成 2.0，防止它无限增长。
- softcap 的 cap 一开始可以较大（50 以上），后期再根据稳定情况尝试下调。
- 小模型如果想节省调参成本，更推荐先单独上 softcap 或者单独上 QK-norm，组合方案留给大模型或者不缺算力的场景。
4. 退火 QK-norm：让保险丝跟着训练节奏走
4.1 为什么要“退火”而不是“固定”
固定强度的 QK-norm 有一个天然矛盾：训练初期模型参数混乱，logits 动态范围大，确实需要强约束；训练后期模型已经形成一定的注意力模式，如果仍然强力约束，就限制了 QK 表达能力的上限。说白了，大 logits 不等于坏事，在查询和键高度匹配的时候，稍微大一点的 logits 能让注意力更果断。
退火 QK-norm 的核心思想，就是让约束强度随训练进度变化。最常用的方式是做“严格到宽松”的退火：训练早期系数 λ 接近 1，QK-norm 的约束完整接入；训练后期系数逐渐衰减到 0，让模型恢复原始的自由度。也有反向退火的方案，即先不加，后期加，这在我试下来是“问题已经出现了才做补救”，效果不如预期好，因为注意力坍缩一旦发生，梯度信号已经污染，同修复基础损伤。
4.2 动态强度调整的实现与观测
具体实现不需要改模型结构，只需要在 forward 里让 QK-norm 的 scale 变成一个动态系数，随 step 线性或者余弦衰减：
def anneal_coef(current_step, total_steps, start=1.0, end=0.0, warmup_ratio=0.05): # 前 5% 步维持完整约束，之后按剩余步数线性退火 if current_step < int(total_steps * warmup_ratio): return start progress = (current_step - int(total_steps * warmup_ratio)) / (total_steps - int(total_steps * warmup_ratio)) return start + (end - start) * progress把返回的系数乘到 QK-norm 的输出上即可。要注意的是，退火末尾的 scale 不要直接降到 0，否则会造成一步陡变，等价于训练末期突然改变模型架构。我一般给它留一个最低值，比如 0.1 或者 0.2，让 QK-norm 始终保留一点兜底作用。对比实测下来，平滑退火比简单线性要稳定不少，尤其在小模型上，陡变引发的 loss spike 概率非常高。
实际操作上还有一个更省心的做法：让 QK-norm 的 scale 直接变成一个可学习参数，然后给它的损失函数里加一个小的 L2 正则，形成一个“逐渐退火”的隐式效果。这样做的好处是不用设计退火 schedule，坏处是可控性差，尤其是在训练中断继续跑的时候，scale 状态可能和你预期不一致。我建议训练任务时间充足、想少写代码的时候用隐式退火，想精细控制跑实验的时候用显式退火。
4.3 小模型场景下的退火效果
在我的小模型实验里，退火 QK-norm 比固定 QK-norm 的最终困惑度平均好 1.5% 到 2%，比不加保险丝好 3% 以上。这个数字看着不大，但对小模型来说已经算显著提升。原因可能在于：小模型本身容量有限，如果全程用强力 QK-norm 约束，它很难学习到一些需要在注意力上“用力”的模式；退火则给了模型先稳定后放开的空间。
另一个有意思的观察是：退火周期和总训练步数的比例会影响最终效果。比如一个训练 10 万步的模型，我在第 6 万步退火到位和在第 8 万步退火到位，表现差别很明显。提前退火给模型更多时间在“自由状态”下打磨注意力分布，效果更好；延迟退火则把过多时间花在受限状态，后期来不及做精细调整。不过如果退火太早，前期的稳定红利又不够，容易在训练中段就开始出现轻微 spikes。这个平衡点目前更多靠经验和验证集 loss 来调，还没有一个通用公式。
5. 小模型到底要不要装：决策框架与配置建议
5.1 小模型和大模型在稳定性上的本质差异
很多人认为小模型训练不容易爆炸，直觉上“参数越少越稳定”。这个观点有一定道理：参数规模小，优化轨迹更容易被约束住，极端权重的产生概率低于大模型。但小模型有一个独特的劣势——它的容量小，容易在训练后期进入“过拟合式锐化”状态，注意力分布逐渐变得保守且尖锐，尤其是在训练数据不够多样或者学习率偏大的时候。
也就是说，小模型和大模型的不稳定来源不太一样：大模型的发散多来自优化动态的混沌和大批量并行导致的有效学习率过高；小模型的发散多来自单一任务模式过拟合和数值退化。这意味着保险丝的必要性是存在的，但设计目标要调整。
5.2 什么情况下加，什么情况下别加
我整理了一个比较实用的决策清单，基于我对多个 0.1B 到 1B 规模模型的经验：
- 训练数据比较干净、领域单一（比如只做代码或者只做百科类数据）：不加保险丝的概率大一些，但建议准备一个训练监控脚本，时刻盯着 QK 点积的标准差，如果连续几千步高于某个阈值再停训加保险丝。
- 训练数据来源复杂、混合比例高、文本长度变化剧烈：建议直接加，宁可杀错不可放过。小模型对数据分布偏移更敏感，保险丝是低成本高保障。
- 训练总步数超过 50万步或者多轮重复数据：强烈建议加退火 QK-norm，因为长期训练很容易在后期出现累积性数值问题。
- 用于继续预训练或做长上下文扩展：必加 softcap，长序列下 QK 点积的新鲜性高，不加容易踩坑。
- 训练周期极短，比如几天内的微调任务：不加，因为模型初始化权重已经经过了预训练，QK 动态范围相对可控，加了反而可能拖慢收敛。
我给过一个“高风险”配置实际测试：0.3B 模型，混合数据训练 80 万步，不使用任何保险丝。第 30 万步出现一次小 spike，手动恢复后第 37 万步又出现一次，且一次比一次严重。加了 softcap 之后，整个训练平稳走完，最终指标反而比不加的稍好。这印证了一个点：注意力保险丝对最终效果的影响，很多时候不是负数，而是一种正则作用。
5.3 推荐的三套可选配置
针对小模型，我根据实际经验推荐三档配置，直接抄作业即可：
- 保守稳定型：
QK-norm(scale=0.8) + softcap(cap=50)，适合数据极杂、训练步数超长。 - 平衡高效型：
QK-norm(scale=1.0) + ln_anneal（线性退火到 0.2），适合常规规模的预训练。 - 资源最少介入型：只加
softcap(cap=80)，适合在已有架构上做最小改动快速稳定。
这三档配置都经过至少一个实际小模型项目的检验，直接复用问题不大。但要提醒一点，任何保险丝方案的超参数都必须配合你的训练时长、学习率、batch size 来调整，不存在全局最佳的固定值。
6. 实操中的坑与排查技巧
6.1 常见问题的快速定位方法
我把自己做训练稳定性排查时用到的判断路径整理了一下，按出现频率排序：
| 现象 | 可能原因 | 排查方式 | 建议处置 |
|---|---|---|---|
| 训练中后期出现单点 loss spike | QK 点积突然变大，或 embedding norm 异常 | 打印qk_std和token_norm的历史曲线 | 先看是否数据异常，再考虑加 softcap |
| 加了 QK-norm 后 loss 反而升高 | scale 过小或 softcap cap 过小，限制表达 | 观察 scale 和学习率的关系 | 调大 scale 初始值或提高 cap |
| 注意力分布变得极其均匀（熵高） | 约束太强，模型“不敢”聚焦 | 计算 attention entropy 分布 | 用退火 QK-norm 或降低约束强度 |
| 训练早期就发散 | optimizer 或学习率本身不合理 | 检查 loss 曲线和 grad norm | 先排除基础问题，再考虑保险丝 |
| 退火后 loss 出现台阶式下降 | 退火造成架构等效变化，残差重分配 | 观察退火系数平滑度 | 改成余弦退火或加长过渡段 |
表里第五个问题是我自己踩过最深的坑。线性退火到 0.0 的那一步，模型等于突然少了一层约束，权重瞬间重新调整，loss 曲线出现类似阶梯的形态。虽然最终可能收敛到更好的平原，过程中却容易产生几个小 spikes。后来我改成余弦退火，并让退火末端系数停在 0.2，这个现象基本消失。
6.2 训练过程中的监控指标
加了保险丝之后，判断它有没有生效不能只看 loss。我通常会在训练 loop 里额外记录四个指标：
qk_std：QK 点积的标准差。这个值应该在 QK-norm 或 softcap 的约束下保持在一个稳定范围内，如果持续攀升，说明保险丝没起作用或者强度不够。qk_norm_scale：QK-norm 的可学习 scale 值。如果它持续增长，说明模型在试图跨越约束边界。attention_entropy：每个注意力头的平均熵。如果整体熵过高，保险丝限制了表达；如果某个头熵特别低，可能该头已经坍缩。max_logits：QK 点积的最大值。softcap 生效时这个值应严格小于 cap；如果接近 cap，说明有异常 token 正在冲击上限。
这四个指标配合 original loss 一起看，基本能把问题定位到“是否保险丝问题”还是“优化器问题”还是“数据问题”。定位清楚之后，再动手改参数，效率会高很多。
6.3 终极建议：先加最便宜的保险
如果要给一个不熟悉这个领域人建议，我的结论很直接：小模型加上 QK 保险丝的成本很低，但收益可能比想象高。计算开销上，两个 RMSNorm 对整体训练时间影响可忽略，显存增加也微乎其微。调试成本才是主要成本，好在有 softcap 这种几乎零参数的方案兜底。
所以我的建议是：当你准备训练一个 1B 以下的小模型，与其赌它不会出问题，不如先在最不关键的位置加一个最便宜的保险（single softcap），把训练跑稳再说。等模型和数据成熟了，再考虑是否需要上退火 QK-norm 精细调优。小模型很容易出现“训练过程稳定，评测结果不行”的问题，而这类问题往往和注意力分布过度平均有关，退火 QK-norm 比传统调整学习率的办法更直接。
我在实际操作中还会做一件小事：把 QK-norm 从一种“保险”思路理解成“初始化偏置”。它其实是在训练早期给注意力提供一个先验，告诉模型“你暂时不需要用大 logits 来表达关注”，慢慢再教它学会用。这个视角上的转换帮我省了不少调参时间——我不再把它当成一个防御性组件，而是当成一个可设计的先验模块。这样无论在哪个训练阶段，我都知道它应该起什么作用、什么时候该松手。
最后再分享一个小技巧：不管选哪类保险丝，记得固定住随机种子对比实验，否则很难判断效果是来自保险丝还是来自随机噪声。小模型训练的随机性方差比想象中大，同样配置跑两次可能差 2% 到 3%，如果评估指标只有 1% 的提升，无法下结论。多跑几个种子，统计稳定性再下判断，这才是科学的方法。
