上个月我在训一个 0.5B 的小模型，大概跑到 20B token 的时候，loss 开始规律性地出现尖峰。查了一圈，不是数据问题也不是学习率问题，最后把 attention 里的 logits 打出来看了一眼，一个 head 的注意力分数直接冲到了四位数的量级。这要是再往上点，bf16 下直接就 inf 了。当时我脑子里冒出来的第一个东西，就是大模型训练里常说的那几根“保险丝”：QK-norm、softcap，还有个最近才被频繁提到的退火 QK-norm。
这篇文章就聊聊我自己的实践和判断：小模型到底要不要装这些保险丝，什么时候装，怎么装，以及装完之后有哪些坑。
先说结论，免得你往下看半天：“要不要装”不看参数量，看训练规模和稳定性指标。模型再小，训练 token 一旦冲上去，注意力机制该出的事故一件都不会少。
1. 保险丝防的到底是什么事故
在讨论“小模型要不要装”之前，得先把这三根保险丝各自防的是什么事故讲清楚。它们不是同一个东西，出问题的场景也有差别。
1.1 QK-norm：把无界变有界
标准的注意力分数是 q 和 k 做点积，然后除以 sqrt(head_dim)。问题在于，q 和 k 是经过深层残差网络输出的向量，它们的范数完全不受控制。训练到某个阶段，某些 head 会走上“捷径”，让一部分 q_k 向量的范数诡异地膨胀，点积结果从几十跳到几千甚至上万。softmax 对这种极端输入非常敏感，一个巨大分数直接把整个注意力分布压成 one-hot，梯度穿过之后也开始剧烈震荡，loss 曲线就疯了。
QK-norm 的思路很朴素：在 q 和 k 做点积之前，先对每个 head 各自的 head_dim 维度做一次归一化。常见做法是 RMSNorm，把向量拉回一个稳定的尺度范围，这样点积结果不再依赖 q 和 k 的绝对值大小，注意力分布不会走极端。
你可以在 Gemma、Chinchilla 这些模型里看到它的影子，社区里老玩家跑几百 B 模型也基本默认带一个。这个操作对显存影响很小，就是一组 norm 权重，浅层另算一点计算量。
具体实现上有个容易被忽略的细节：QK-norm 通常不做在 hidden_size 上，而是先 reshape 成B, heads, T, head_dim，然后在 head_dim 上逐 head 做 norm。你要是图省事对整个 hidden 做了 norm，等于把所有 head 混在一起拉尺度，注意力分布的特性会被破坏。
1.2 softcap：防止单个 logits 失控
softcap 是另一个维度的问题。QK-norm 保证的是 q,k 的范数稳定，但即便范数稳定，随着训练推进，个别位置的 logits 依然可能偏大。softcap 直接在 logits 层面做对称压缩：
logits = beta * tanh(logits / beta)
beta 一般是 30 或 50。它的含义是：在正常范围内，tanh(logits / beta) 近似等于 logits / beta，相当于不做什么处理；但一旦某个 logits 偏大，tanh 会把它温和地压住，最高不超过 beta，从而给所有 logits 加了一个软上限。
这个东西和 QK-norm 是互补的。QK-norm 管的是“源头”（q、k 的范数），softcap 管的是“出口”（最终的 logits 分数）。PaLM 和 Gemma 都用过，Gemma 2 的 config 里甚至可以分别设置 attention_logit_softcapping 和 final_logit_softcapping。
注意一点：softcap 不是归一化，它不改变 logits 的相对大小关系。它只解决“局部极端值”的问题，不会像 norm 那样重新分配相对尺度。所以它和 QK-norm 并不冲突，可以同时开。
1.3 退火 QK-norm：给模型一个“断奶”期
退火 QK-norm 这个名字在社区里还没有完全统一的标准定义。我理解它的核心诉求是：QK-norm 在训练期是稳定器，但模型训练完进入推理或微调阶段时，如果干脆把 QK-norm 去掉，模型往往会掉点，因为它已经依赖 norm 后的 q、k 尺度来维持注意力分布了。
这就好比孩子一直扶着栏杆走路，栏杆突然撤了，他步子是稳不住的。退火 QK-norm 的思路是在训练末期留一段“断奶”时间，把 QK-norm 的作用强度从一个固定值逐渐降下来，让模型慢慢学会在没有它的情况下也能保持正常注意力分布。最终推理时可以选择不启用 QK-norm，省掉这部分计算，减少和训练环境的差异。
退火方式大致有两种流派：
- 一种是直接对 RMSNorm 的输出做线性混合：
q = tau * rms_norm(q) + (1-tau) * q，让 tau 从 1 逐渐降到 0。 - 另一种是保持 norm 结构，但把 QK-norm 里附加的缩放参数从一个设定退火到另一个更低的值。
这两种我没有看到官方严格对照实验，也没法告诉你哪个包治百病。我自己实践更喜欢后者，因为它稳定性更高，模型的注意力分布不会发生突变。
2. 小模型会不会犯同一个病
聊完原理，回到核心问题：小模型会犯这个病吗？我的答案是会，但触发条件比很多人想得更晚。
2.1 训练规模是真正的门槛
决定要不要保险丝的不是“参数量”这个单一指标，而是训练总 token 数。一个 3B 模型只训 5B token，可能从头到尾 logits 都没啥异常；但一个 300M 模型训 100B token，后期照样可能炸。参数量小只是让模型“爆炸上限”更低，但训练步数一多，优化器和 loss landscape 的长期累积效应依然会让部分 head 走向极端。
我自己一个感觉：大体上，训练 token 在 20B 以内的小模型（<0.5B），裸装其实问题不大，出尖峰概率低，更多时候是学习率尾部没有调好导致的 loss 波动。但一旦你打算把数据规模冲到 50B、100B token，保险丝就变得很有必要。小模型吃那么多数据，并不等于它不会发生大模型那种退化，只是时间推迟了。
2.2 小模型出问题的表现不完全一样
大模型出问题，经常是 loss 断崖或者直接 NaN。小模型更多是：
- loss 曲线出现规律性小尖峰，幅度不算大，但每几千步一次，很烦人。
- 注意力熵下降过快，某些 head 彻底“死亡”，输出分布变得非常尖锐。
- 验证集困惑度不再下降，偶尔还会反弹。
尖峰还好查，最容易被误判的是注意力熵下降。小模型因为容量有限，本来就会有部分 head 更偏向局部信息，但如果训练到中后期，所有 head 的平均熵持续走低，或者出现个别 head 熵直接逼近 0，那大概率是 logits 尺度失控了。此时你不一定看得到巨大的损失尖峰，但模型表现已经受影响。
2.3 我的默认判断流程
我建议每个在小模型上挣扎的人先做一次排查，不要一上来就堆保险丝：
- 训练刚开始稳定后，打印 attention logits 的绝对值最大值、99 分位数，以及所有 head 的平均 attention 熵。
- 如果 logits 的绝对值最大值超过 30 到 50，就有风险，超过 100 就该处理。
- 如果某个 head 的 attention 熵已经低于 0.1，而训练还在继续，别犹豫了。
- 另外看看是不是 bf16：bf16 最大只能表示 65504，logits 一旦冲到这个附近，下一步就是 inf。
我现在的判断标准是：模型 <0.1B 且训练不到 5B token，基本不装，先观察；模型 >0.5B 且训练超过 50B token，直接上 QK-norm，softcap 视 logits 统计而定；模型在 0.2B 到 0.5B 之间，看数据规模，超过 20B token 就把 QK-norm 作为默认选项。
3. 动手给小模型装保险丝
如果你决定要装，这里给一套可以照抄的方案。
3.1 可以先看现成配置
在 HuggingFace 生态里，Gemma 2 的 config 是自带 attention softcap 的：
from transformers import Gemma2Config config = Gemma2Config( num_hidden_layers=12, hidden_size=768, intermediate_size=2048, num_attention_heads=12, num_key_value_heads=4, head_dim=64, attention_logit_softcapping=30.0, final_logit_softcapping=30.0, )但 Gemma 2 的 config 里没有 QK-norm 开关，QK-norm 通常在代码层面实现。如果你用的是 nanoGPT、GPT-2 这类自研小模型，直接改 attention 层更灵活。
3.2 QK-norm 的实现细节
QK-norm 推荐用 RMSNorm，因为它不减去均值，操作更轻，实现也就一页纸：
import torch import torch.nn as nn class RMSNormHeadDim(nn.Module): def __init__(self, head_dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(head_dim)) self.eps = eps def forward(self, x): # x: [..., head_dim] rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps) return x / rms * self.weight在 attention forward 里，reshape 成多 head 维度后再 norm：
# q, k shape: [B, heads, T, head_dim] if self.use_qk_norm: q = self.q_norm(q) # 按 head_dim 做 RMSNorm k = self.k_norm(k) q = q * self.qk_scale k = k * self.qk_scale这里的self.qk_scale我习惯初始化为 1.0，可学习也可以固定。有一个观点认为，QK-norm 后点积的标准差会变成 sqrt(head_dim)，所以 scale 设为 sqrt(head_dim) 更合理。我实测两种都行，最终看模型表现。可学习 scale 有一个额外好处：训练后期如果 scale 自动变大，说明模型自己找到了合适的注意力温度，这比手调温度参数要自然。
关键点：norm 一定在 per head 的 head_dim 上做。你如果是在hidden_size上做的，结果会有偏差，模型初期能跑但后期表现怪怪的。
3.3 softcap 叠加进 attention
在拿到点积 logits 之后，趁 softmax 之前做 cap：
# logits: [B, heads, T, T] if self.use_softcap: logits = self.beta * torch.tanh(logits / self.beta) # causal mask 在 softcap 之后再处理 logits = logits.masked_fill(~causal_mask, -1e9) probs = torch.softmax(logits, dim=-1)注意这里 mask 的处理顺序有讲究。我之前见过一个坑：把 mask 里的-inf位置先填进去，然后再做 tanh。结果tanh(-inf / beta) = -1，mask 位置变成-beta，虽然 softmax 后概率依然是 0 附近，但它会让分母带上一个额外项，在极端情况下不如直接设-1e9干净。所以我的习惯是：先 softcap，再 mask 填充一个大负数，最后 softmax。
beta 的选择，我建议起步用 50。如果你发现 logits 依然频繁触到 50，再往下调到 30。不要一开始就用特别激进的 beta，否则模型表达能力受限，训练曲线反而更差。
3.4 退火 QK-norm 的简单实现
假设你的训练步数总数是total_steps，决定在最后anneal_steps步做退火，可以这么写：
def get_qk_norm_alpha(step, total_steps, anneal_ratio=0.03): anneal_start = int(total_steps * (1.0 - anneal_ratio)) if step < anneal_start: return 1.0 progress = (step - anneal_start) / max(1, total_steps - anneal_start) return 1.0 - progress然后在 attention forward 里传入这个 alpha：
if self.use_qk_norm: q = self.q_norm(q) k = self.k_norm(k) q = q * (self.qk_scale * alpha) k = k * (self.qk_scale * alpha)这样做的好处是 alpha 降到 0 之后，qk_scale 依然存在，等于模型最终学会在一个相对自然的尺度下做注意力。你别真把 scale 也一起退火到 0，那 q、k 点积就成 0 了，注意力变成均匀分布，模型直接废了。这个坑我踩过，一度以为是退火调度写错了，最后发现是 scale 不该跟着归零。
如果你的目标是推理时不启动 QK-norm，更接近的退火方式是直接把 norm 输出和原始 q 做线性混合：
q = alpha * self.q_norm(q) + (1 - alpha) * q这个方式最终等于不使用 QK-norm，但对训练末期模型稳定性的挑战更大。我在一个小实验上对比过，效果不稳定，建议优先选择“保留 norm 结构、只退火 scale”的方案。
4. 常见问题与排查技巧实录
保险丝装上去不等于万事大吉，实际跑起来会遇到不少问题。这些是我自己记录下来的排查思路。
4.1 加了 QK-norm 之后模型反而不收敛
如果你刚开始加 QK-norm，模型出现训练初期 loss 下降变慢，甚至 perplexity 比不加更高，先检查是不是 norm 放错了位置或维度，再看是不是顺序搞错了。正确的顺序是：q、k 的线性映射之后、点积分母缩放之前做 norm。
如果顺序没问题，可能原因就是 QK-norm 改变了 q、k 的尺度，导致 attention 的温度变了，模型需要一个适应期。你可以适当调大学习率 warmup，或者给 qk_scale 一个好一点的初始值。比如 head_dim=64 时，点积原始标准差是 sqrt(64)=8，你可以把 qk_scale 初始化为 1.0，让 logits 维持在一个匹配的尺度。
4.2 softcap 把表达空间压太狠了
有一种情况特别容易迷惑人：加上 softcap 后，loss 曲线前期还挺正常，但到中后期收敛变慢，validation 一直差一口气。这时你去统计 logits 分布，发现大部分 logits 已经普遍靠近 beta，说明模型想表达更大的分数差，但被 cap 住了。
这时候别急着加 beta，先看 qk_scale。因为 QK-norm 可以把 logits 尺度压低，导致正常 logits 也在 tanh 的线性区之外，本质上是 scale 和 beta 不匹配。你想一想，如果要让 logits 在 [-15, 15] 内正常工作，beta 用 30 是够的；如果 logits 普遍到 40、50，beta 用 50 都未必有空间。所以调整顺序应该是：先保证 QK-norm 把 logits 压制在合理分布，再来调 beta。
4.3 softcap 和 mask 的顺序问题
这个前面提过，我再强调一次：
两种方式看起来差不多，实际结果有细微差别。强烈建议：logits = beta * tanh(logits / beta)之后，再masked_fill一个非常小的数，再 softmax。不要在logits本来就带-inf的情况下直接 softcap。
4.4 退火 QK-norm 后模型掉点
我做退火时第一次也遇上了。退火完成后，验证集 loss 比不退火高了一点点，虽然差得不多，但对强迫症来说挺难受。
后来排查发现，问题不是退火本身，而是退火期间学习率还处在相对高位。退火 QK-norm 相当于改变模型内部的优化目标，你需要同时把学习率降到比较低的位置，让模型在小扰动下自适应。我现在的做法是：退火开始的时候，把学习率乘一个 0.1 的系数，或者干脆让退火阶段用固定低学习率跑完。
另一个方法是把退火阶段拉长一点，别图省事只留最后几千步。我一般留总训练步数的 3%-5%，如果资源允许，5% 以上效果更好。
4.5 QK-norm 与 Flash Attention 的兼容坑
这个值得一提。Flash Attention 在内部直接算q @ k^T然后 softmax，中途不给你插手 logits 的空间。所以如果你想用 Flash Attention 并加 softcap，没法直接用官方scaled_dot_product_attention接口，只能在调用它之前把 q、k 先用 QK-norm 处理掉；但 softcap 就比较麻烦，需要自己实现 attention kernel，或者用 eager attention 先做实验。
对于小模型，序列长度通常在 2048 以下，eager attention 的额外开销完全可接受。如果序列特别长，考虑两种变通：
- 只做 QK-norm，不做 softcap，因为 QK-norm 已经把最大 root cause 解决了。
- softcap 只在 attention 层之外的其他位置用，也就是对最终输出 logits 做 cap，不去动 attention logits。
4.6 问题速查表
| 现象 | 可能原因 | 解决办法 |
|---|---|---|
| loss 后期规律性尖峰 | 某个 head 的 logits 膨胀 | 加 QK-norm，打印 logits 统计确认 |
| 加了 QK-norm 后发散 | 放错维度或顺序 | 检查 reshape 逻辑，确保在 head_dim 上 norm |
| 模型收敛但验证集差口气 | softcap 的 beta 太小 | 先调大 beta，再考虑调 qk_scale |
| 退火后验证集掉点 | 学习率没同步压低 | 退火阶段使用固定低学习率 |
| 推理去掉 QK-norm 大幅掉点 | 训练时没有做退火 | 保留 QK-norm 或做退火调度 |
| attention 熵训练中持续下降 | 模型在退化，logits 分布过尖 | QK-norm + 调整学习率，必要时 softcap |
5. 我说一点个人体会
现在我在小模型上的默认动作是：先做 500 步的体检，打印三样东西：attention logits 的最大绝对值、平均熵、以及 loss 尖峰频率。这三样都不会说谎。体检出来没事，就老老实实裸跑，别为了“跟风”给模型装上不必要的保险丝。体检出来有事，QK-norm 是第一优先，因为它是管根上的问题。
softcap 我把它当成备用保险丝：只有当 QK-norm 已经压制住大部分极端 logits，但少数位置依然偏大，或者你发现自己没办法引入 Flash Attention 的改造时，才考虑它。
至于退火 QK-norm，我更倾向于在模型量级上去以后再用。小模型很多时候训完还要做下游微调，微调本身就是一种“退火”过程，模型有机会适应没有 QK-norm 的推理环境。你如果直接部署，不一定非要退火，保留 QK-norm 部署也完全 OK。
如果你现在也在纠结自己的小模型要不要装保险丝，我的建议是：别凭感觉猜，先给你的 attention logits 装个计数器，让数据告诉你答案。等你看过一个 1B 模型在 50B token 训练中的 logits 分布变化，你对“小模型要不要装”这个问题，会有完全属于自己的判断。
