我最近在练一个大概三亿参数的 decoder-only 模型，数据量不算多，满打满算几十 B token。前七万步一切正常，loss 曲线漂亮得像教科书；第七万零几十步，loss 直接从我印象里的 3.1 附近跳到了 9 以上。等我把训练挂回去、重新 dump 完 logits 分布，才发现根本不是数据或学习率的问题——QK 点积炸了。也就是从那时候起，我开始认真对待“QK 保险丝”这群技术：QK-norm、softcap、退火 QK-norm。

这篇东西我想把这三个方案放在一起讲清楚：它们各自在防什么、怎么实现的、小模型到底值不值得装。尤其最后一点，网上说法很两极：有人觉得小模型容量本来就紧张，再加 norm 等于把注意力焊死；也有人觉得小模型训练更脆，不装保险丝一个 spike 就能让你前功尽弃。我两边都踩过，下面用我自己的训练记录说话。

1. 保险丝要挡的到底是什么：attention logits 失控的完整链条

1.1 一个 300M 模型在第七万步的“意外”

那次事故的现场其实很典型：loss 在几十步内从 3.1 爬到 9.4，然后开始剧烈震荡。我第一反应是学习率出问题，或者数据 pipeline 混进了脏 batch，但检查下来都正常。真正的问题藏在 attention logits 里：我用一个采样 hook 打印了每层 QK 点积的统计值，发现最后几层的 logits max 已经冲到 60 以上，而正常训练时这个数字通常在 5 到 10 之间。

这个现象在 Transformer 训练里并不罕见。QK 点积本质上是两个向量的内积，它没有天然的上下界。如果 Q 向量和 K 向量的模长在训练中持续被放大，logits 就会被推向极端。而我们通常不直接约束模长，只会做初始化、加 scale、加 LayerNorm 这类间接控制。于是当某一步参数更新恰好让一批 Q、K 的模长同时变大，logits 就可能在极短时间内冲破阈值，形成肉眼可见的 loss spike。

1.2 logits 失控为什么会演变成梯度灾难

单看 logits 大，本身不一定致命。致命的是它后面紧接着 softmax。softmax 对输入尺度极其敏感，当 logits 从 5 涨到 50，softmax 输出会快速趋向 one-hot：最大的那个位置几乎拿到全部概率，其他位置直接归零。

一旦 softmax 变成 one-hot，注意力梯度就废了。因为被忽略的那些 token 的梯度贡献约等于零，只有被选中的那个位置在反向传播里有信号。这条路径上的参数会继续加大 logits，让下一次前向更加极端。这个正反馈循环持续几轮，loss 就会像坐过山车一样冲上去，而网络内部的权重早已被严重扰动。

这也是为什么我说保险丝防的是“失控的链条”，而不是单纯防“logits 数值大”。QK-norm、softcap 这类措施的价值，是在链条的第一环把 logits 的幅度按住，避免后面的 softmax 饱和、梯度丢失、参数震荡一环接一环地发生。

1.3 小模型为什么格外经不起这种冲击

大模型遇到 logits 尖峰时，往往表现为损失轻微抖动，因为参数量大、冗余度高，某些层被扰动后还有其他路径能兜住。小模型没有这个冗余空间。三亿参数的模型，某两层的 QK 模长一旦炸掉，几乎等于把这一路特征全部冲乱，恢复时间非常长。

小模型的训练还往往伴随更少的数据和更短的训练步数。一次严重的 spike 可能直接抹掉前面几万步积累的成果，尤其在训练后期，学习率已经降下来，模型依赖精细调整过的注意力模式，被 spike 一冲，很多局部结构就再也回不到原来的状态。

所以我的观点是：小模型不是“不需要保险丝”，而是“更需要保险丝”。但小模型也确实存在表达力被压制的副作用，这是下一节要展开的核心矛盾。

2. 三种保险丝的机制对比：QK-norm、softcap、退火 QK-norm

2.1 各自在电路的哪个位置动手

用电路来打比方：QK-norm 是在“电源端”动手，直接限制 Q、K 向量的模长；softcap 是在“负载端”动手，在 logits 进入 softmax 前加一道限幅；退火 QK-norm 则是给 QK-norm 装了一个定时开关，训练后期逐步把限制撤掉。

QK-norm 的实现非常直白：对每个 attention head 的 Q 向量和 K 向量分别做一次 RMSNorm，让它们在参与点积前拥有固定的模长。这样一来，QK 点积的天然尺度就被限制在 head_dim 附近，不会再跟着训练中模长的漂移而无限放大。

softcap 的做法是给注意力分数加一个非线性压限：logits = C * tanh(logits / C)。它不改变 Q、K 的模长，而是在 logits 离开点积计算后、进入 softmax 前，把绝对值大于 C 的部分平滑地压回来。这个方案保留了正常范围内的自由度，只对极端值动手。

退火 QK-norm 在思想上又进一步：它把 QK-norm 视为一个训练阶段的工具，而不是模型架构的永久组成部分。在训练前期，norm 提供稳定；在训练后期，当模型已经学会控制自身行为，就逐步把 norm 的影响退掉，让注意力恢复原本的锐利度。

2.2 三种方案的本质差异

QK-norm 的核心思路是“让 logits 的统计学分布是可控的”。RMSNorm 会让 Q、K 各自的 RMS 变成 1，所以即使输入 x 的模长跑偏再远，norm 输出也被拉回固定范围。代价是它会影响所有样本的 logits，而不是只影响极端值。也就是说，即使那些本来很正常的注意力模式，也被统一地压缩了动态范围。

softcap 的核心思路是“保留正常动态，只处理尾部”。当 logits 绝对值远小于 C 时，tanh 函数近似线性，softcap 几乎不起作用；只有当 logits 明显超过 C 的时候，它才发力把极端值压下来。这个性质让它对正常注意力分布的影响更小，但代价是需要调一个合适的 C，调太小会在正常区间引入不必要的压缩，调太大又等于没装。

退火 QK-norm 的核心思路是“按训练阶段切换策略”。它承认前期稳定性和后期表达力都是必要的，但不能同时拥有，所以用时间维度来换取空间维度。前期享受 norm 带来的稳定，后期通过退火逐步释放注意力头的自由度。

2.3 一张表看懂三种方案的取舍

这三个方案也不是互斥的。很多实践中的配置是“QK-norm 打底 + softcap 兜底 + 后期退火”，三层保险各自负责一段风险窗口。不过在小模型上，我不建议一上来就三重叠加，因为每加一层都会引入新的超参数和实现细节，出了问题反而难以定位。

2.4 从恒压电源到限流电阻的生活化类比

QK-norm 让我想到稳压电源。电网电压不稳，设备容易烧，你就在入口处加一级稳压，不管进来的电压怎么波动，输出总是固定值。代价是输出电压被统一限定，某些设备可能得不到它想要的更高电压。

softcap 更像限流电阻。它不改变正常工作电压，只在电流超过阈值的时候把它限住。设备正常运行时感觉不到它的存在，但短路时它能救命。缺点是阈值设多少需要经验，而且它不能阻止电压本身漂移，只是不让漂移造成破坏。

退火 QK-norm 则像一个训练后期允许拆掉的稳压器。电路在最容易出事故的初期需要稳压，但等整个系统稳定运行后，你希望设备能直接吃到原始电压以获得更好的性能，所以拆掉它。拆的过程不能太突然，否则设备会闪断重启，这就是退火窗口存在的意义。

3. 小模型的独特账本：head_dim、表达力损失和训练预算

3.1 head_dim 越小，QK-norm 的“限制感”越明显

小模型通常不仅参数量小，head_dim 也小。很多 300M 到 1B 的模型，head_dim 用的是 64，甚至还有 32 的。head_dim 直接决定了每个注意力头能表达多少种不同的匹配模式，也决定了 QK-norm 之后的 logits 动态范围。

QK-norm 把 Q、K 归一化到固定模长后，点积的期望标准差大约在 1 附近（乘上标准 attention scale 后）。这时候如果 head_dim 是 128，注意力头还有足够的维度去编码相对位置、句法关系、词汇匹配等不同信号；如果 head_dim 只有 32，这些信号被压缩到同一个 1 量级的 logits 空间里，头与头之间的区分度就会明显下降。

我在实践中感受到的差别是：head_dim 64 的模型加 QK-norm，下游任务分数通常只掉零点几个点；head_dim 32 的模型加 QK-norm，某些需要精确指代的评测任务能掉三到五个点。这个损失在训练曲线上不一定看得出来，因为 PPL 甚至可能更好看，但真实能力已经悄悄缩水了。

3.2 被牺牲的注意力锐利度在哪些任务上看得见

注意力锐利度指的是模型能否让某一两个 token 获得极高权重。文档检索、多跳问答、代码补全里精确的变量引用，都依赖这种锐利选择能力。QK-norm 把 logits 的整体尺度压下来后，softmax 的分布会变得更平，注意力分数之间的差距变小，锐利度自然下降。

小模型本身在锐利度上就没有优势，因为参数量限制了它精确建模长程依赖的能力。再加一层模长限制，等于把一份本来就不丰厚的预算继续压低。我做过一个对照实验（规模不大，仅代表我自己的观察）：同一个 300M 模型，训练数据完全相同，有 QK-norm 的版本 PPL 略好，但在一个简单的文档检索任务上，准确率比无 QK-norm 版本低了约 4 个点。

这让我意识到一个很重要的问题：PPL 是“平均视角”，注意力锐利度是“极端视角”。保险丝把极端情况稳住的同时，也可能把极端能力一并抹平。小模型用户如果特别在意检索类能力，必须认真权衡这笔账。

3.3 数据预算与算力约束改变了性价比

但这笔账的另一头是训练稳定性。小模型常见的训练场景是：数据量不大、算力有限、只打算跑一次训练。在这种预算下，一次中等规模的 loss spike 可能就让你损失几十上百个小时的算力。QK-norm 把 spike 的概率压到很低，这种确定性带来的收益，比那些“理论上存在但很难复现”的表达力损失更值钱。

我的个人阈值判断是：如果你的模型小于 1B、训练数据在 100B token 以内、head_dim 不超过 64、下游任务以通用生成和对话为主，那么装 QK-norm 的数学期望是正的。它牺牲的那点锐利度，往往可以在退火阶段拿回来。如果你是在做搜索引擎、做密集检索、做需要精确指代的任务，或者你有很多次重训机会，那可以优先考虑不装，或者只装一个宽松的 softcap。

4. 实战：给 300M 模型加上 QK-norm 与 softcap 的最小改动

4.1 完整参考实现：QK-norm 的插入位置

我用一段接近 PyTorch 风格的伪代码来说明。关键点是：norm 必须作用在 head_dim 维度上，而不是整个 hidden_dim 上。很多初学者把 RMSNorm 直接接到 q_proj 输出后面，发现效果不对，就是维度搞错了。

import math import torch import torch.nn as nn class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): # x: [..., head_dim] x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return x * self.weight class Attention(nn.Module): def __init__(self, hidden_dim, num_heads, head_dim): super().__init__() self.num_heads = num_heads self.head_dim = head_dim self.q_proj = nn.Linear(hidden_dim, num_heads * head_dim, bias=False) self.k_proj = nn.Linear(hidden_dim, num_heads * head_dim, bias=False) self.v_proj = nn.Linear(hidden_dim, num_heads * head_dim, bias=False) self.o_proj = nn.Linear(num_heads * head_dim, hidden_dim, bias=False) self.q_norm = RMSNorm(head_dim) self.k_norm = RMSNorm(head_dim) def forward(self, x): bsz, seq_len, _ = x.shape q = self.q_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim) k = self.k_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim) v = self.v_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim) # QK-norm 放在这里：对最后一维 head_dim 做归一化 q = self.q_norm(q) k = self.k_norm(k) # 如果用了 RoPE，建议在 norm 之后做旋转，顺序更稳定 # q, k = apply_rope(q, k) q = q.transpose(1, 2) # [bsz, heads, seq_len, head_dim] k = k.transpose(1, 2) v = v.transpose(1, 2) attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim) attn = torch.softmax(attn, dim=-1) out = torch.matmul(attn, v) out = out.transpose(1, 2).contiguous().view(bsz, seq_len, -1) return self.o_proj(out)

这里有几个细节值得强调。第一，q_proj 和 k_proj 的 bias 最好设为 False，因为 norm 对偏置比较敏感，加了 bias 反而会让归一化后的向量带上一个固定的偏移。第二，RMSNorm 的 weight 初始化为 1 就好，不需要额外去调 scale。第三，如果你用的是 GQA 或 MQA，K 的 head 数比 Q 少，k_norm 要按 K 的实际 head 数来建，不能直接复制 Q 的配置。

4.2 softcap 的配置细节：避免把注意力压成“一锅粥”

softcap 的代码更简单，通常在注意力分数计算完之后、softmax 之前插入：

attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim) softcap = 50.0 attn = softcap * torch.tanh(attn / softcap) attn = torch.softmax(attn, dim=-1)

但配置 C 值需要先了解你的 logits 分布。我的做法是在训练日志里记录 logits 的 max、p99、mean、std，跑上两三千步，看正常区间在哪。如果正常 logits 大多在 5 以内、极少超过 15，那 C 设 30 或者 50 都行，基本不干扰正常区间。如果正常 logits 就已经到 30，那 C 至少要设 80，否则 tanh 的非线性区会在日常计算里就生效，把注意力分布压得太平。

softcap 的陷阱是：它只能限幅，不能恢复梯度。当 logits 绝对值远超 C 时，tanh 的导数趋近于零，这部分梯度会被截断。所以它适合兜底，不适合作为唯一的稳定性手段。真正频繁出现 logits 尖峰时，我更推荐用 QK-norm 从源头控制。

4.3 训练监控：统计 qk logits 而不是只看 loss

装保险丝之后，必须配套监控。loss 是结果指标，logits 分布是原因指标。你在 loss 上看到 spike 时，往往已经晚了。我习惯在训练循环里每隔固定步数采样一次：

def log_qk_stats(step, attn): with torch.no_grad(): max_logit = attn.max().item() mean_logit = attn.mean().item() std_logit = attn.std().item() p99_logit = torch.quantile(attn.flatten(), 0.99).item() print(f"step={step} max={max_logit:.2f} p99={p99_logit:.2f} " f"mean={mean_logit:.2f} std={std_logit:.2f}")

如果 p99 持续上升，比如从 8 涨到 12 再到 18，别等它炸。这时候就要考虑调整初始化、降低学习率，或者提前把 QK-norm 的介入权重拉高。小模型最忌讳的就是“loss 还好就不管 logits”，因为 loss 的平滑很容易掩盖内部分布的缓慢漂移。

5. 退火 QK-norm 的实操：什么时候退、怎么退、退完发生了什么

5.1 为什么要退：norm 会锁住注意力的“锐化”窗口

前面说过，QK-norm 会把注意力分布的动态范围限制在一个较窄的区间里。这个限制在训练前期是保护，在训练后期却可能变成天花板。模型在后期需要精细地区分“哪个 token 重要、哪个 token 不重要”，需要注意力分数之间的差距足够大，而 norm 不允许这个差距无限拉大。

退火 QK-norm 的做法，就是在训练后期把 norm 的影响逐步撤掉，让模型重新获得调整注意力锐度的空间。它和 warmup 刚好相反：warmup 是让优化器从平滑过渡到激进，退火是让模型从受限过渡到自由。一个好的退火计划，应该让模型在几乎没有感知的情况下完成这个切换。

5.2 三种退火实现与我的推荐方案

第一种是插值法，也是我最常用的。定义退火系数 lambda，它从 1 平滑衰减到 0：

def anneal_factor(step, start_step, end_step): if step < start_step: return 1.0 if step >= end_step: return 0.0 t = (step - start_step) / (end_step - start_step) # cos 从 1 到 0，比线性更平滑 return 0.5 * (1.0 + math.cos(math.pi * t)) # 在 attention forward 中 lam = anneal_factor(step, start_step, end_step) q = lam * self.q_norm(q) + (1 - lam) * q k = lam * self.k_norm(k) + (1 - lam) * k

第二种是直接切换法：在某个步数直接把 norm 从计算图里拿掉。这个方案简单，但极容易导致 loss spike，因为 QK logits 会在一步之间从 norm 约束内跳到无约束状态。除非你的学习率已经降到非常低，否则我不推荐。

第三种是空间法：先只退 K 侧 norm，保留 Q 侧 norm 一段时间；或者反过来。这样可以让模型分阶段适应，但实现起来要多写不少状态逻辑。我的经验是插值法已经足够平滑，不需要再去搞分侧退火，增加不必要的复杂度。

5.3 退火窗口、学习率与常见翻车点

退火窗口的长度，我一般取总训练步数的 5% 到 15%。总步数一万步时，窗口至少给一千步；总步数十万步时，窗口五千到一万步比较稳妥。起点通常放在总步数的 85% 到 95% 之间。太早退，模型还没稳定，norm 的保护价值没有充分实现；太晚退，剩余步数不够模型充分适应无 norm 的动力学。

退火期间的学习率也应该配合。如果训练计划里本来就有 cosine decay，退火窗口大概率落在学习率已经比较低的区间，这是天然的配合。但如果你用的是常数学习率或者分段常数，那退火开始时要主动把 lr 降到峰值的 0.1 到 0.3 倍，否则更容易翻车。

我实际踩过的坑有两个。第一个是 lambda 下降太快。有一次我把一万步的退火窗口设成了两千步，结果退火到一半 loss 开始出现低频振荡，每次振荡都伴随一轮 logits 尖峰。改成 cos 计划、窗口拉长到五千步之后，问题立刻消失。第二个是退火结束后马上做了一次 short warmup 再训练，导致已有模式被破坏。退火结束后模型应该进入一种“静默适应”状态，学习率保持很低，不要给它新的剧烈扰动。

退火完成后怎么验证？我会看三件事：loss 曲线在退火区间有没有异常抬升；qk logits 的 max 和 p99 是否逐步脱离 norm 约束并稳定在一个新的区间；下游任务的 benchmark 是否比不退火的版本更好或者至少不差。如果退火后 loss 短暂上升但很快回落，通常说明模型正在重新校准注意力尺度，可以接受；如果 loss 持续不降，就要回退窗口长度或者停止退火。

6. 要不要装的决策清单：别盲信保险丝

6.1 按任务类型与训练阶段选择的决策框架

我在实际项目中总结了一套比较实用的判断顺序，核心是先看模型规模和数据预算，再看下游任务对注意力锐利度的敏感程度。

这套建议不是说没装保险丝就一定出事，而是说在不确定性和预算约束下，哪些选择更稳。小模型用户最常见的错误是两个极端：要么完全不信保险丝，等 spike 出现后才痛苦地重训；要么把 norm 焊死在架构里，从头到尾不退火，最后发现模型检索能力差得莫名其妙。

6.2 不装保险丝的例外场景

有一种情况我特别想强调：如果你的模型要在推理阶段做极长上下文的扩展，QK-norm 需要谨慎。因为长上下文推理时，新增 token 的 K 向量可能和训练时看到的有分布差异，norm 会把这种差异强行拉到和训练一致的尺度，短上下文看是保护，长上下文看可能反而限制了模型对新位置的适应能力。

另外，如果你在模型上还叠加了很重的量化方案，低 bit 量化对 Q、K 向量的模长很敏感。QK-norm 让 Q/K 的尺度变得很集中，对量化是好事；但退火之后你又把尺度放开了，这会让量化误差重新变得不可控。所以退火前要评估量化方案能不能接受无 norm 的 logits 分布，不能为了训练期性能牺牲部署期稳定性。

6.3 我的个人取舍：一次性训练与研究性调试分开对待

现在再回到“小模型要不要装”这个问题，我的答案已经很清楚：取决于你还有没有第二次机会。

如果你只有一次训练预算，跑完就上线，那我强烈建议装。无论是 QK-norm、softcap 还是退火组合，它们的核心价值不是“让分数更高”，而是“让训练过程可控”。一次训练里少一次 spike，比在评测集上高零点几个点重要得多。

如果你有研究性质的调试空间，可以反复训练、能接受失败，那我会建议先不装，只加监控。等 logits 分布真的出现问题时，再装上保险丝做对照。很多时候你会惊讶地发现，某些模型在特定数据配比下可以一直稳定跑下去，根本不需要 norm。盲装反而会让你的实验多一个变量，干扰你对模型本身的判断。

我还想分享一个“半自动保险丝”的思路：在训练代码里写一个监控器，当 qk logits 的 p99 连续 N 步超过某个阈值时，自动把 QK-norm 临时插入计算图，等分布稳定后再逐步移除。这个方案适合超长训练任务，等于给模型装了一个可自动触发的智能保险丝。代价是实现复杂，但在资源有限的项目里非常值得投入。

保险丝的本质不是永远限制电流，而是让电路安稳度过最容易出事故的阶段。训练前期接上它，训练后期有计划地拆掉它，比永远焊死在电路里聪明得多。这也是我从三亿参数的教训里最深的体会：任何训练技巧都要服务于“这一阶段的真实风险”，而不是服务于某种教条。