小模型预训练注意力不稳？QK-norm与退火方案实战解析

最近在调一个小模型大概 0.4B 参数预训练跑到中段开始出问题。loss 曲线每隔一段时间就冒一个尖峰峰值高的地方几乎要 NaN训练日志里连续出现 inf 警告。那种感觉就是你知道它在某个位置已经快崩了但还没彻底崩卡在玄学边缘。拉日志一看果不其然问题出在注意力这一层。某些 head 的 QK 分数直接冲到几百softmax 后面基本就变成了 one-hot梯度传回来是一堆接近 0 的值网络更新不动然后下一批数据又来冲击这个脆弱点。最后我认认真真地试了一圈 QK-norm、softcap以及一个朋友提到的“退火 QK-norm”方案。这篇就把我在小模型上做这套“保险丝”实验的过程、取舍和踩坑记录整理出来给同样在折腾小模型训练的朋友做个参考。适合谁看正在做小模型预训练、继续训练或者长上下文扩展的人遇到过 loss 尖峰、NaN 但不确定是不是注意力炸了的人以及只是想搞明白 QK-norm 和 softcap 到底怎么选、有没有必要上退火版本的人。我会尽量把数值层面和工程层面的东西都讲透。1. 注意力 QK 部分的数值危机1.1 注意力分数为什么会长歪先把计算路径捋一遍。自注意力里对每个 head我们有score q · kᵀ / √d其中 q 是 query 向量k 是 key 向量d 是 head dimension。这个 score 会进入 softmax变成 attention 权重。问题在于q 和 k 都不是凭空出现的它们是从残差流里经过线性映射投影出来的。残差流的规模在深层网络里会缓慢积累尤其是 Pre-norm 结构的模型每一层的输入在 LayerNorm/RMSNorm 之前标准差可能远大于 1。当 q、k 向量的 L2 范数从十几涨到上百score 就变成几十甚至几百的量级。softmax 这东西指数函数在输入超过某个区间后非常敏感。score 20 时softmax 已经接近硬选择score 100 时其他位置的梯度基本就是 0这个 head 等于死掉了。更危险的是这种状态会自我强化某个 head 输出极端的注意力分布导致它接收到的梯度非常小参数几乎不更新下次还会继续输出更极端的结果直到某一刻数值溢出。我见过最夸张的情况是某个 0.4B 模型在 8k 序列长度下第 18 层的某个头q 向量范数涨到 2000 多logits max 接近 300。loss 在那个 step 直接跳了 0.3然后又花了几百步才缓回来。这种偶发尖峰比持续发散更阴险因为你会以为它只是训练波动但它实际上在持续破坏已经学好的特征。1.2 为什么小模型也会中招很多人有个直觉小模型参数量少容量小“能量”低应该不容易数值爆炸。实际上这个直觉在小模型预训练场景下是错的。小模型为了追效果往往会把学习率调得比较高、batch size 拉得比较大、序列长度也往长了推。这些操作本质上都是在增加单步更新的压力和单条序列内部的数值规模。尤其现在小模型也动辄上 8k 甚至 32k 的上下文序列变长以后注意力矩阵的尺寸变大更容易出现极端点。我做过一个对照组实验同一个 0.3B 模型batch size 从 128 涨到 512别的都不动注意力 logits 的上限从 30 左右飙到 100而且集中出现在深层。这说明小模型不是不会炸是炸的临界点比大模型更隐蔽因为它的 loss 波动幅度不大不仔细看日志根本发现不了。另一个典型场景是小模型做长上下文继续训练。这种场景下位置编码对相对位置的插值会放大某些 key 的范数注意力分布更容易变得不均匀。我在这个场景里几乎每次都建议先把 QK-norm 加上再说而不是等出问题再加。2. 三款保险丝的工作原理2.1 QK-norm把尺度拉平不直接限高QK-norm 的做法很简单在计算 score 之前先对 q 和 k 分别做一次 RMSNorm。q_norm RMSNorm(q) k_norm RMSNorm(k) score q_norm · k_normᵀ / √dRMSNorm 做的事情是让向量除以它自身的均方根再乘一个可学习的缩放参数 gamma。换句话说它抛弃了 q、k 的绝对尺度只保留方向信息然后用 gamma 重新给一个大小。这个东西和“把 logits 限制在某个上限”不是一回事。QK-norm 压缩的是 q、k 的范数波动而不是直接限制 score 的数值。打个比方它像音响的自动增益控制信号弱的时候帮你放大信号强的时候帮你压下来始终让输出落在功放的安全区间内。softmax 该尖锐还是尖锐但不会尖锐到让梯度完全消失。实操里比较关键的取舍是RMSNorm 放在 RoPE 之前还是之后。放在 RoPE 之前纯洁地控制 q、k 本体的尺度位置信息不受影响。放在 RoPE 之后RoPE 会改变向量各维度的相位但不会改变整体模长所以对范数的控制效果基本一致只是归一化时会带着位置分量一起做。我实测下来放在 RoPE 之后更稳尤其是在长上下文场景里。原因我理解是RoPE 在相对位置较大的时候会产生更大的绝对数值波动先旋转再归一化相当于把位置编码引入的尺度扰动也一并压住了。但有人会担心这样会把位置信息的一部分强度也抹掉我在小模型上试没观察到明显的位置能力损失。2.2 softcap直接给 logits 画一条软上限softcap 是另一条路线公式是logits tau * tanh(logits / tau)tanh 是双曲正切输出范围在 -1 到 1 之间。乘以 tau 之后logits 就被软性地限制在 (-tau, tau) 区间内。所谓“软”是因为它不是 clip 那种硬截断而是连续平滑地压缩。这个操作有几个性质值得知道当 logits 远小于 tau 时tanh(x / tau) 约等于 x / tau整体输出约等于 x。也就是说正常范围内的分数不会受影响。当 logits 远大于 tau 时输出趋近 tau导数趋近 0。也就是说极端分数会被“按住”但梯度也会消失模型在极端区域基本学不到东西。tau 通常取 30 到 50 之间。50 已经算比较大的因为普通训练中 logits 到 50 就已经非常极端了。softcap 像物理保险丝有人把它当急救药只在发现 logits 已经失控的时候临时加上。如果你训练全程都用它可能会遇到一个麻烦——模型的注意力始终被按在一个上限以内某些依赖“极端复制”行为的任务比如从上下文里精确找一个长段复制出来会略微掉点。这也是后来出现“退火 softcap”思路的动机训练后期把 tau 调大让模型在收敛前逐步回到自然 logits 尺度。2.3 退火 QK-norm训练后期把保险丝拆掉退火 QK-norm 是我这次实验里花了最多时间理解的东西。它本质上不是一个独立机制而是一个训练 schedule前 90% 的训练步数里保持完整的 QK-norm让模型在稳定状态下学会大部分知识最后 5% 到 10% 的步数里引入一个系数 alpha从 1 线性或余弦降到 0q alpha * RMSNorm(q) (1 - alpha) * q k alpha * RMSNorm(k) (1 - alpha) * kalpha 1 时是纯 QK-normalpha 0 时是完全不 norm 的裸 QK。退火过程结束模型就处于无 norm 的状态。为什么要这么干我理解有两个动机。第一QK-norm 是稳定器但也可能变成限制器。它把 q、k 的尺度强制拉齐模型在最有机会形成“极端注意力”的区域被摁住了。预训练后期数据里的长尾知识恰恰需要模型表现出更强的注意力选择能力如果一直被 norm 限制可能永远学不到最尖锐的模式。第二训练和推理的一致性。如果你的模型训练时带着 QK-norm但推理时要部署一个不带 norm 的版本比如为了省一点点显存或者为了兼容某些底层算子那模型在最后阶段其实是在一个失真空间里完成收敛的推理时 logits 尺度一变行为就会漂移。退火可以让模型在训练尾部就切换回真实推理环境把这种 gap 补上。说实话退火不是必需品。如果你的推理运行时时也保留 QK-norm那固定 QK-norm 就已经很好了。但如果你追求“训练后期 loss 还能往下走一点点”或者有推理去 norm 的需求退火版本值得试。3. 小模型到底要不要装3.1 先看你是否出现了“保险丝熔断”的信号不是所有小模型都需要 QK-norm。我见过很多小模型从头训到尾不加任何保险丝loss 也一直很平稳。盲目加 norm 反而可能带来轻微的能力损失和额外的计算量。所以第一步永远是判断你的模型有没有出现保险丝熔断的信号。我建议你在训练日志里重点盯这五个指标loss 曲线出现台阶式反弹或孤立尖峰尤其是尖峰过后 loss 回不到原来的下降趋势。日志里出现 inf 或 NaN哪怕只有一次。一次 NaN 之后参数可能已经损坏不恢复训练就需要回滚 checkpoint。用固定 seed 复现时相同的 token 位置反复出 NaN。这说明问题不是随机的是某个注意头结构性地失稳。某些 head 的 attention entropy 长期低于 0.05。entropy 太低意味着注意力几乎只集中在某一个 token 上这种 head 随时可能演变成死头。梯度范数出现周期性尖峰且尖峰出现在深层 Transformer block。如果以上信号一个都没有我的建议是别折腾先把训练跑完再说。如果中了一两项优先级就出来了先试固定 QK-norm再考虑 softcap最后才上退火。3.2 根据训练场景选择装法怎么选还要看你的训练场景。我整理了三个最常遇到的场景以及我的判断结果场景风险特征我的选择从零预训练小模型学习率高、数据杂、序列长最容易出现尖峰优先固定 QK-norm如果追求训练后期极致收敛尝试退火版从大模型蒸馏或继续训练小模型教师模型的注意力分布已经成型student 的 logits 尺度相对稳定先不加 norm出现尖峰时用 softcap 急救而不是改 QK长上下文增量训练位置编码外推、key 范数增长都容易导致 logits 爆炸固定 QK-norm 收益最明显建议加退火收尾第二类场景值得展开说。蒸馏小模型的时候学生的目标是模仿教师的概率分布此时学生注意力 logits 的绝对尺度在蒸馏损失里是被约束着的很少出现自主训练时的失控。给这种模型加 QK-norm等于额外改变学生的注意力表示空间容易把已经学好的分布特征弄变形。我在这类任务上吃过亏后来宁可用 softcap 兜底也不动 QK。还有一个能力层面的权衡小模型本来就要靠尖锐注意力来完成“上下文照抄”这类任务。QK-norm 把尺度拉平后某些情况下会让注意力倾向更均匀相当于削弱了模型对关键 token 的专注能力。我实测在简单的 retrieval 评测集上固定 QK-norm 会让小模型掉 1 到 2 个点但换成退火 QK-norm这个 gap 几乎就没了。这也算退火版本的一个额外收益。4. 实操参数与避坑4.1 具体插入位置和参考代码下面给一套我实际用过的参考实现方便直接抄。先是 RMSNorm 模块import torch from torch import nn class RMSNorm(nn.Module): def __init__(self, dim, eps1e-6): super().__init__() self.eps eps self.gamma nn.Parameter(torch.ones(dim)) def forward(self, x): rms x.pow(2).mean(-1, keepdimTrue).add(self.eps).sqrt() return x / rms * self.gamma然后是注意力 forward 里的插入位置。我习惯放在 RoPE 之后def forward(self, x, rotary_emb): q self.q_proj(x) k self.k_proj(x) v self.v_proj(x) # 先做 RoPE q rotary_emb(q) k rotary_emb(k) # 再做 QK-norm每 head 独立归一化 q self.q_norm(q) # q_norm 作用在最后一维 d_head k self.k_norm(k) scores torch.matmul(q, k.transpose(-2, -1)) / self.scale ...退火版本的实现稍微复杂一点关键是 alpha 要在训练循环里更新不能让它进入反向传播的梯度计算class AnnealedQKNorm(nn.Module): def __init__(self, dim, eps1e-6): super().__init__() self.norm RMSNorm(dim, epseps) self.alpha 1.0 def set_alpha(self, alpha): # alpha: 1.0 表示纯 QK-norm0.0 表示无 norm self.alpha float(alpha) def forward(self, x): if self.alpha 1.0: return self.norm(x) if self.alpha 0.0: return x return self.alpha * self.norm(x) (1.0 - self.alpha) * x训练循环里按步数算 alphaanneal_start int(total_steps * 0.92) anneal_end int(total_steps * 0.98) for step in range(total_steps): if step anneal_start: alpha max(0.0, 1.0 - (step - anneal_start) / (anneal_end - anneal_start)) for module in model.modules(): if isinstance(module, AnnealedQKNorm): module.set_alpha(alpha)注意几个细节eps 不要设太大1e-6 够用。eps 太大会让 RMSNorm 在小幅度向量上产生明显的 scale 偏差。gamma 初始用 1这是常见默认值。如果发现加了 QK-norm 之后 loss 下不去可以把 gamma 初始调到 0.8 甚至 0.5 试试实际上就是让初始注意力 logits 整体更小一点。退火窗口不要拉太长。我试过 20% 长度退火中段出现明显的 loss 反弹因为模型过早失去了保护。5% 到 10% 是合理区间。4.2 退火节奏、学习率配合与观测指标退火 QK-norm 不是单纯把 alpha 降下来就行它和训练后期的学习率策略是耦合的。我的参考配置是退火开始的同时把学习率降到最后值的 20% 左右。也就是说模型在无 norm 状态下只做小幅参数更新不至于在切换到真实 logits 尺度时发生大的震荡。退火结束后再保持当前学习率跑 2% 到 3% 的稳定训练观察 loss 是否继续下降。观测指标上别只看 loss。我每次实验都会记录所有层 q、k 向量的 L2 范数均值和 P99 分位数注意力 logits 的 max 和 mean每层每个 head 的 attention entropy 均值梯度范数的 max 值这些指标能告诉你保险丝到底在哪个部位起作用。比如如果退火后只有第 16 到 20 层的 logits 快速上升其他层都很稳那可能不需要全模型退火只对深层模块做退火就够了。局部保险丝也是一个可行的工程方案。实验日志我建议统一记到一个地方后面分析的时候按 step 对齐看比事后拼接日志高效得多。4.3 效果评估比 P99 loss而不是只比平均 loss评估加不加 QK-norm、要不要退火最直观的方式是看 loss但这里有个坑只看平均 loss 会掩盖偶发尖峰的影响。我自己的做法是额外统计“P99 loss 波动幅度”。具体来说把每个 step 的 loss 记录下来看 P99 分位数与 P50 分位数的差值或者直接记录 loss 的最大尖峰高度。这个指标才是保险丝价值的直接体现——固定 QK-norm 几乎不会降低正常区域的 loss但它能把 P99 和最大尖峰压得很实。我刚试软软的时候也踩过这个坑加 QK-norm 跑了 2000 步平均 loss 和不加几乎一样差点就把方案否了后来对比 P99 才看出差距加了 norm 的曲线 P99 波动只有不加的 40% 左右。对小模型来说少一次 NaN 回滚就省下几千步训练时间这笔账怎么算都划算。5. 常见问题与排查现象可能原因排查方向加了 QK-norm 后 loss 一直偏高norm 放早了或过强试试把 norm 移到 RoPE 之前或把 gamma 初始调低到 0.8退火过程中出现 loss 尖峰退火窗口太短或学习率没有同步降低拉长退火窗口把退火区间从 5% 调到 10%学习率先降到最后值的 20%softcap 训完出现明显平台期cap 太紧正常 logits 也被压缩把 tau 从 30 提到 50或只在 logits 超过 50 后才开始压缩分段函数某个 head 还是死掉QK-norm 不是万能可能 norm 只加了部分层检查死 head 出现在哪层只对前若干层或深层单独加 norm退火到 alpha0 后 loss 反弹模型对 norm 依赖太重不要退到 0把 alpha 下限设为 0.1保留 10% 的 norm 强度推理时如果还保留 QK-norm但训练时用了退火训练和推理不一致退火只是为推理去 norm 准备的推理保留 norm 的话直接用固定 QK-norm 即可这里说一个容易被忽略的点softcap 和 QK-norm 不是只能二选一它们是可以叠加的。QK-norm 管 q、k 的尺度softcap 管最终 logits 的上限两者作用在不同环节。如果你发现 QK-norm 加上后 logits 尖峰消失但某些大 logits 还是偶尔冒出来可以再叠一个 tau50 的 softcap 做第二道保险。我这次后期实验就是 QK-norm 退火 低强度 softcap 一起用的最后一段训练稳定得让我有点不习惯。6. 一点个人体会折腾完这一圈我最大的感受是小模型要不要装“QK 保险丝”不是按参数数量来决定的而是按训练压力来决定的。如果你的模型训练平稳、数据干净、序列不长真的没必要多此一举。QK-norm、softcap 都是防御性设计加在身上总归有一点表达力代价。但如果你已经开始频繁处理 loss 尖峰或者要往长上下文方向走那早装早安心别等到一次 NaN 把几千步训练成果毁掉再后悔。我个人的优先级排序是先固定 QK-norm把它当作第一道保险如果还有极端 logits叠一个 tau 较大的 softcap如果训练后期出现平台期、且推理环境允许去 norm再考虑退火 QK-norm。这套组合下来小模型既能稳住训练又不至于被保险丝绑住手脚。最后分享一个小技巧对比加不加 QK-norm 时记得把“P99 loss 波动幅度”和“最大尖峰高度”一起列出来。只看平均 loss你很可能得出“加了没啥用”的错误结论但看了尖峰数据你会感谢这道保险丝。