前阵子有个做 1B 级 decoder-only 模型的朋友来找我说训练跑到第 4000 步左右 loss 曲线开始隔三差五冒尖每次 spike 都要掉接近 0.5得花上千步才爬回来。他看了一圈开源仓库发现 QK-norm 和 softcap 在主流大模型里几乎是标配于是拿着这个问题来问我小模型要不要也装上这根“QK 保险丝”这个问题我很熟。最近手头正好同时跑了一条 70M、一条 300M、一条 1B 的预训练实验中间试过 QK-norm、attention logits softcap以及训练末期把 QK-norm 退火anneal到近乎关掉的方案。这几种手段在社区里经常被放进同一个句子里讨论但它们防的是截然不同的风险。这篇文章我把这几根保险丝拆开讲清楚最后给出一套可以照抄的分档决策和默认配置。标题里说的“小模型”我按预训练领域常说的 100M ~ 1B 参数来理解训练数据在 50B ~ 300B token 这个量级。至于那些 7B 起步的大模型装 QK-norm 已经是行业默认小模型往往让人犹豫的主要原因是加上去多一套算子、多几个初始化要调而且很多人直觉上觉得“参数小、深度浅应该不会像大模型那样炸”。这个直觉一半对一半不对。下面我从注意力 logits 为什么失控说起。1. 先看 QK-norm 和 softcap 在防什么注意力 logits 失控的完整链路要回答“小模型要不要装”得先搞明白 QK-norm 和 softcap 到底在防什么。它们防的是同一个东西的两个阶段注意力 logits 的数值失控。Transformer 多头注意力里每一层都在做这样一个点积logits Q K^T * scale其中 scale 通常是 1 / sqrt(d_head)d_head 是每个注意力头负责的维度。比较理想的情况是Q、K 投影之后的每个维度方差保持在 1 附近这样一来 logits 的方差大概也是 1 的量级序列长度 4k 时一批 logits 的最大值也就是 3~5。这个范围对 softmax 来说非常健康梯度既不会消失也不会爆炸。但训练到中后段事情会变。Q、K 投影权重是参与训练的优化器为了降低当前 loss完全有动力把某些 head 的 Q/K 向量尺度推大从而让某几个 token 的注意力分数特别高形成“锐利”的注意力分布。这种锐利分布在部分任务上确实能压 loss但它有很强的正反馈logits 越大softmax 越接近 one-hot对应路径的梯度占绝对主导其他 token 几乎拿不到梯度。模型就会越来越依赖那几个特殊 token更新信号越来越偏参数越走越歪直到某一步 logits 冲到 30、50loss 直接拉出一个 spike。这中间还有一层隐蔽的数学问题d_head 越大越容易失控。点积本身是 d_head 维的求和如果 Q/K 每维方差被放大到 2~3在训练中很容易发生logits 的典型范围就会从 ±3 放大到接近 ±10。如果再叠上长序列里位置编码的相互作用某些 token pair 出现系统性大值就一只脚踩进了饱和区。小模型虽然 d_head 往往不大64 居多但仍然要面对这个机制只是概率和幅度不同。1.1 为什么说 QK-norm 是“拆掉失控的发动机”QK-norm 的做法很直白在 Q、K 投影之后、做点积之前各自过一个 RMSNorm有的实现用 LayerNorm把向量的 RMS 重新拉回固定值。这样无论训练中 Q/K 内部怎么“膨胀”最终参与点积的向量长度都被重新归一化了。关键的数学约束是柯西不等式|q · k| ≤ ||q|| · ||k||归一化之后两个范数都被压到固定尺度单个 logits 想无限往高冲就失去了基础。我习惯把 QK-norm 叫“拆发动机”它没有直接限制 logits 多高而是把能推动 logits 变高的“动力源”——也就是 Q/K 向量长度的自由度——拿掉了。方向由向量方向表达大小由固定缩放统一控制。需要补充的是norm 之后的缩放要仔细选。RMSNorm 会把 Q/K 每一维的 RMS 变成 1等价于向量的 L2 范数被固定为 sqrt(d_head)。这时候点积之后继续乘 1/sqrt(d_head)logits 的典型标准差仍然在 1 附近注意力分布和标准初始状态差不多这是最不容易出错的配置。我见过不少第一次加 QK-norm 的同学把 scale 漏掉或者改成了别的值结果注意力要么变成均匀分布要么重新变得过于尖锐。Mistral、Gemma、DeepSeek 这几个系里都能看到 QK-norm 的身影实现细节差异很大有的带可学习 scale有的不带但核心逻辑都是同一件事让 Q/K 的模长不再参与梯度博弈。1.2 softcap 是其次的闸门硬截断和 tanh 软截断softcap 直接作用在 logits 上跟 QK-norm 是两条互补路径。最常见的实现是 tanh softcaplogits C * tanh(logits / C)C 是超参比如 30、50、100。logits 的绝对值被压到 C 以内即使某一步真的爆炸模型也不会立刻 loss NaN而是变成“平滑压缩”后的结果。这相当于给注意力分布加了一层保险闸门。也有实现直接用 clamp(logits, -C, C)简单且没有额外算子。但 clamp 的问题是超过边界的 logits 在 softmax 里会“贴边”堆积多个 token 共享同一个上限梯度语义变差tanh 版本因为连续光滑logits 很大时梯度趋近 0更像“我不想管这个死活”反而不容易过拟合出新的错误结构。大模型社区里近年更流行 tanh softcapLLaMA 3 的 attention 就用了这个思路。我自己的实验也倾向于 tanh softcapC 从 30 起步再根据 logits 直方图调整。1.3 两个保险丝一起装会怎样QK-norm 加 softcap 同时开等于从“输入端”和“输出端”各设一道防线。一个大致的分工是QK-norm 保证训练过程中 logits 的尺度不会整体漂移softcap 负责兜住个别极端 token pair。一起开通常不会冲突但要注意两个坑。一是别把 softcap 的 C 设得太小比如 20 以下。QK-norm 已经限制了大量 logitsC 太小会压制正常的对比度让注意力变软最终 loss 反而变差。二是一旦开了 QK-norm想让注意力更锐利应该调 norm 路径上的额外缩放而不是无限压 C。固定 standard scale然后靠更大的 C 去放行效果最可预期。2. 小模型的 attention 为什么也会冒尖三大实际风险来源大模型必须装 QK-norm多半是因为层深、头多、序列长失稳概率随规模指数上升。小模型是不是就可以不装了我在 70M / 300M / 1B 三条线上跑下来结论是小模型通常不会像 7B 那样系统性高频地炸但依然有三种场景非常容易冒尖而且一旦冒尖恢复代价并不会因为模型小就变便宜。2.1 层数少不代表不会炸特殊 token 和局部参数崩溃小模型深度浅层间叠加效应弱所以 QK-norm 的收益确实比深模型小。但注意力 logits 失控并不只来自“多层放大”。最典型的是特殊 token 竞争。训练语料里 BOS/EOS、数字、罕见 Unicode、重复出现的模板词这些 token 天然会和上下文里的常规 token 形成强烈的注意力竞争。小模型因为容量有限更容易把“注意力资源”押在一个固定 token 上比如每句话都去看 BOS 或者去看上一个 token。某个 batch 里如果这类 token 大量出现logits 就会在一层内被局部放大不需要多层配合也能 spike。我还遇到过一种“局部参数崩溃”的形态某个注意力头在某个训练阶段演化成几乎只输出一个固定方向K 投影的某几维方差持续膨胀导致该头对所有 query 都给出近似相同的极值 logits。这种崩溃在大模型上表现为某个 head 失效在小模型上则直接表现为 batch loss 的孤立尖峰。QK-norm 对这类问题很有效因为它直接把 K 向量的长度压住了崩溃头失去“靠长度得分”的空间。2.2 bf16 低精度训练是隐形的帮凶第二个风险源来自低精度训练这块经常被低估。很多小模型实验为了提高吞吐已经切到 bf16 混合精度q/k 投影后的张量也是 bf16 存储。bf16 的指数范围够大不容易溢出但尾数只有 7~8 bit表示一个大数比如 20.0时相邻可表示数之间的间隔接近 0.1。Q 向量和 K 向量本身动态范围越大点积结果在 bf16 下被量化掉的误差就越显著。QK-norm 在这里有一个隐藏优势经过归一化后Q/K 向量的每个元素都被压到比较小的动态范围里bf16 表达中的相对量化误差反而更均匀。也就是说即使不出现剧烈 spikeQK-norm 也能让低精度下的注意力分数更干净一些。我在这几个实验里统计过打开 QK-norm 后bf16 全精度对比的 logits 相对误差大致能下降一个数量级虽然具体数值和硬件 kernel 实现有关但趋势是稳定的。当然“下降一个数量级”这说法要看基数。如果模型规模小、d_head64、bf16 点积本来误差就不大QK-norm 带来的精度提升可能对最终 loss 毫无感知。但反过来如果你的 loss 曲线有那种“高频小毛刺”切换上下文长度或数据分布后冒出一连串小 spike那十有八九和低精度表达有关上 QK-norm 比调学习率更对症。2.3 从 loss spike 日志反推“是哪种炸法”排查 loss spike 时我最反感“先调学习率再说”的做法。更有效的路径是把 spike 附近的梯度、logits 统计一起拿出来看。我自己常用的探针很小就是在 attention forward 里每 N 步采样几批数据记录这批数据在所有层的 logits max、logits 均值、q/k 的 L2 范数 p90。怎么读这些指标呢logits max 平时 8~15spike 时冲到 30且 q/k L2 范数同步上涨大概率是 Q/K 向量长度失控首选 QK-norm。logits max 平时只有 4~6spike 时跳到 20 左右但 q/k 范数没怎么变通常是特定 token/位置组合的偶发尖峰softcap 就够用。logits 和 q/k 范数都正常但 loss spike 延迟半拍到下一步更可能是某个 batch 本身数据异常或梯度累积问题不要动 QK-norm先检查 loss 重算逻辑和数据清洗。这个区分逻辑其实也呼应了标题——QK 保险丝不是万能药。很多小模型实验真正的问题根本不在 QK而在数据处理和 LR schedule贸然加 QK-norm 只会增加变量最后调了一周才发现源头是 data sampler。3. 退火 QK-norm把保险丝改成可拔掉的手闸如果你决定给小模型上 QK-norm接下来马上会碰到另一个问题训练结束后这个 norm 留不留如果留着推理时每层都要多被执行一两次 RMSNorm虽然单个 norm 不贵但在小模型部署到边缘设备时算子少一个是一个。更关键的是有些小模型实验的目的是产出“无 QK-norm 依赖”的基线模型方便后续跟其他不带 norm 的架构公平对比。于是就有了退火 QK-normQK-norm annealing的做法。3.1 为什么要把好好的 QK-norm 退火到不存在直接说动机有三类。第一类是推理端省算子。QK-norm 本身计算量占比很小在 7B 以上模型里完全可忽略但在 300M 的移动端模型里每层多两个 RMSNorm 会让 latency 增加几个百分点积少成多。如果最终产品想裁掉它最好的办法不是训练完硬删而是在训练末期让模型适应没有 norm 的分布。第二类是为了让 logits 分布与部署环境对齐。很多推理框架或量化工具会假设注意力 logits 是某种尺度分布例如做 KV cache 量化时会期望 logits 的绝对值在一个相对固定的 range 里。带 QK-norm 的 logits 分布往往比纯 raw 点积更“瘦”一旦部署端按裸模型校准可能出偏差。退火到最后让模型学会在 raw 尺度上保持稳定量化时更省事。第三类是放松 QK-norm 对注意力对比度的限制。QK-norm 本质上引入了“无长度信息”的约束这对某些需要极端锐利的任务未必最优。在最后阶段退到低 alpha让模型有机会用 Q/K 向量长度重新编码部分信息实测中偶尔能补回一点点下游分数。3.2 一个可落地的 alpha 退火实现我推荐一种简单的向量级插值方案而不是在 logits 层面直接混。原因是向量层面的干预语义更清晰训练早期 alpha1完全走 norm 路径训练末期 alpha 线性从 1 降到 0模型逐步切回 raw q/k 路径。# attention forward 内部伪代码 # alpha 由训练 loop 传入1.0 表示完全 QK-norm0.0 表示无 QK-norm q_raw self.q_proj(hidden_states) k_raw self.k_proj(hidden_states) if self.qk_norm_alpha 0: q_n self.q_norm(q_raw) # RMSNorm无 affine k_n self.k_norm(k_raw) q q_n * self.alpha q_raw * (1 - self.alpha) k k_n * self.alpha k_raw * (1 - self.alpha) else: q q_raw k k_raw if self.use_rope: q, k apply_rope(q, k, positions) scale self.head_dim ** -0.5 # 标准缩放退火期间保持不变 logits torch.matmul(q, k.transpose(-2, -1)) * scale if self.softcap_scale is not None: logits self.softcap_scale * torch.tanh(logits / self.softcap_scale)这段代码里有两个细节容易被忽略。第一个是退火时不要开可学习的 qk scale。如果 scale 是可学习的优化器会在 alpha 减小时偷偷把 scale 调大来对冲等到 alpha0 时 scale 可能已经长成一个危险值最后几步容易爆新的 spike。退火期间务必把 scale 固定住。第二个细节是这个插值方案会让中间态的 q/k 向量既有 norm 路径的成分也有 raw 路径的成分logits 分布会有一个过渡窗口模型需要一点时间重新适应所以退火窗口要留足。如果你想保留 norm 到最后一刻又想省推理算子还有一个偷懒方案训练最后 N 步直接把 alpha 置 0然后观察 loss 反弹。有效但反弹幅度一般比线性退火大最终 checkpoint 可能要多跑一小段纯 raw 路径才能稳住。我实测下来的经验是线性退火更平滑但会额外占用几千步最终 loss 和完全保留 norm 的 checkpoint 差不太多不同规模、不同数据上差异很大大体在 0.01~0.03 这个量级。3.3 退火窗口多长才不明显反弹这个问题没有标准答案我给出我在小模型上的观察范围供参考。退火窗口占训练总步数的 3% 以内alpha 从 1 到 0 会伴随明显的 loss 反弹反弹后需要再训一小段才能回到之前水平。5%~10%通常能看到一个先涨后平的曲线反弹幅度在 0.02~0.05 量级大部分模型可以接受。超过 10%有时反而没有必要因为小模型本来总步数就不多浪费太多步数在退火上不划算除非你同时在做序列长度切换。如果训练 schedule 里本来就有“最后 10% 线性降到最小 LR”的阶段我一般把 QK-norm 退火和 LR 退火重叠在一起开两个变化同时进行模型反而更容易平滑过渡。不过要小心如果 LR 已经降到接近 0而 alpha 还在快速下降模型几乎没有能力修复切换带来的分布变化所以最好让 alpha 的下降起点比 LR 的终点早一点留出几百步用微跌后的 LR 稳定。还有一类更高级的做法是“退火 softcap”即 softcap_scale 在训练后期从 30 线性增大到 1000等效于把 cap 放宽到近似关闭。原理和退火 QK-norm 一样把保护措施慢慢撤掉让模型在没有保护的情况下仍能维持。两者可以联动但我建议不要同时退火 QK-norm 和 softcap变量太多一旦 loss 反弹不容易定位是哪个保险丝的问题。4. 小模型到底要不要装规模分档与默认配置现在回到标题的核心问题小模型要不要装 QK 保险丝。我的主张是不要把它当成“标配”或“不配”而是按模型规模、训练数据、部署目标三件事来分档。下面这张表是我现在做预训练时用的默认参考。参数规模d_head 典型值默认 QK-norm默认 softcap说明 100M32~64不装可选 C30先查 logits 统计数据质量优先级最高100M ~ 300M64推荐开最后可退火C50低成本高保障尤其要切长上下文时300M ~ 1B64~128装C50~100默认装训练末期按部署需求退火1B ~ 3B64~128装C100和主流开源配置看齐别在这上面省这个表不是拍脑袋。我自己的观察是100M 以下的模型即使出现 logits 尖峰往往也是数据里某个 token 群体的偶发问题加 QK-norm 属于用更多参数换稳定但没有控制住真正的根因。100M 以上开始训练总 token 数普遍在 50B~200B 这个区间训练中途经常需要切换上下文长度或混合数据集QK-norm 的兜底价值明显变大。4.1 装的位置与超参选择全装还是只装上半层一个常见问题是QK-norm 要每个注意力层都装吗还是只装前几层、后几层在小模型上我的建议是要么全装要么全不装不要只装一半。原因不是技术上有害而是“只装一半”会让超参排查变得极其痛苦如果有的层有 norm、有的层没有logits 的整体尺度在各层之间不一致一旦出现 spike你很难判断是哪一层加了 norm 才让它稳定的也很难判断是哪一层没加才让它炸的。实在想省算力我试过另一种替代方案只对 K 做 normQ 不做。K 是“被检索的键”在长上下文中更容易出极端值只 norm K 已经能挡住不少尖峰且算子只有原来的 1/2。但这只适合“轻度保护”场景。如果要开退火还是建议 Q/K 都 norm否则 alpha 降到 0 时两边路径的语义不对称退火会更不稳定。超参方面最简单的组合是RMSNorm不带 affine 参数。省掉 scale/bias 两个可学习量也更容易避免训练初期 scale 漂移。norm 之后保持标准缩放1 / sqrt(d_head)不要乱改 scale。softcap C 先设 50观察一段时间如果 logits 直方图的 99.9 分位数长期低于 15可以调大到 100 或者关掉。如果开了 RoPEnorm 放在 RoPE 前后都可以因为 RoPE 是正交变换不改变向量的范数。我习惯先 RoPE 再 norm原因是很多开源实现这样排方便对照。4.2 我当前的实际基线配置顺手贴一条我现在跑 300M 模型的默认 config都是实测调整过的d_model768d_head64num_heads12layers12每个 attention 前向里q_proj / k_proj 输出后接 RMSNorm无 affine随后继续乘标准缩放1 / sqrt(64)也就是 0.125logits 层做 tanh softcapC50优化器 AdamW初始 LR 3e-3没有 QK-norm 的对照实验我把 LR 降到 2.5e-3因为没 norm 时 LR 稍高一点就会出现隐约的小毛刺训练最后 8% 步数把 QK-norm alpha 从 1 线性降到 0softcap 保持不动推理阶段alpha 设为 0也就是删掉 q/k 的 norm保留 softcap。这套配置在当前几个 300M~1B 的实验上跑得很稳loss 没有出现超过 0.3 的孤立 spikes。作为对照我把 QK-norm 关掉、其他不变重跑了一次大部分时候曲线几乎重合但在上下文长度从 512 切到 2048 的那几百步以及某个包含大量重复广告文本的 batch 出现时loss 会高出 0.3~0.5。也就是说对小模型而言 QK-norm 多数时候是“不那么必需的保险丝”但确实管住了切换期和脏数据期这两类最大的痛点。5. 建议的排查顺序先用日志定位再决定上哪根保险丝最后这段给实操路径。如果你正在跑一个小模型第一条建议永远不是“马上加 QK-norm”而是把 attention 的 logits 和 q/k 范数统计埋进训练日志。绝大多数失稳问题靠统计数据比靠感觉快得多。5.1 给 attention logits 加统计探针在 attention forward 里面加一段代码不要每步都算训练时每 50 步取一个 batch 采样即可if trainer.global_step % 50 0 and random.random() 0.1: with torch.no_grad(): # logits 是 [batch, heads, seq, seq]采样一部分即可 sample logits.flatten().float() stats { attn_logits_max: sample.max().item(), attn_logits_p99: torch.quantile(sample, 0.99).item(), attn_logits_mean: sample.mean().item(), q_l2_p90: torch.norm(q.float(), dim-1).quantile(0.90).item(), k_l2_p90: torch.norm(k.float(), dim-1).quantile(0.90).item(), } log_metrics(stats)这个探针的开销几乎可以忽略。注意要在spike 发生前就开始记录否则事后没有对比依据。拿到这些指标后按我在 2.3 里说的三类区分方式判断。我自己用的阈值大致是logits p99 平时小于 10max 平时小于 20如果 max 超过 30 或者 p99 超过 15就进入值得干预的区间了。5.2 分级处方从观察、softcap 到 QK-norm按严重程度给一个分级处方方便快速照做。第一档logits 几乎健康loss spike 频率极低每 5000 步一次以下。不做任何架构改动先检查 data sampler 和 loss 重算。很多人纠结的 QK-norm 其实没有用武之地。第二档logits max 偶尔超过 30但 q/k 范数稳定。加 tanh softcapC50跑 500 步看变化。如果尖峰次数下降但仍有残留把 C 调到 30或者配合 grad clip 1.0。这一步基本不影响最终 loss成本极低。第三档q/k 范数明显膨胀、logits 高频冒尖。上 QK-normscale 保持1 / sqrt(d_head)训练初期多给 1000 步观察不要因为前面 loss 比对照高 0.05 就立刻回滚norm 需要一点时间让模型习惯新的 logits 分布。第四档已经上了 QK-norm 还炸。这时候不要继续加保险丝回头查初始化、LR、batch size 和 cosine schedule大概率是别的问题。曾见过有人把 softcap、QK-norm、额外 grad clip 全叠上也压不住 spike最后发现是某个 head 的 bias 初始化为 10导致 logits 一开始就偏瘫。5.3 最后分享两个小教训一个是退火别开太晚。我做过一次 1B 模型总步数 20 万步把 QK-norm 退火排到最后 1 万步才开始。alpha 刚开始下降的那几百步没事我以为很稳结果第 2 万步直接出现 0.7 的 loss 反弹后面用 3000 步才磨平。后来把退火窗口提前到总步数的 10%虽然中途也有一点小起伏但最终 checkpoint 明显更干净。另一个教训是关于 softcap 的超参。softcap 开太大就形同虚设开太小会影响模型在少数高置信场景下的表达能力。我最早的实验把 C 设成 5想着“越紧越安全”结果注意力几乎失去了锐利度几个需要长程定位的任务指标掉了不少。后来把 C 调到 30、50 才正常。所以如果你拿不准 C先设 50 是相对安全的起步值真正值得调的是 logits 尺度而不是过度依赖 C。还有一个经常被忽略的前提如果你最终部署时既不想留 QK-norm也不想留 softcap退火只能帮你“软着陆”但不能消除所有风险。部署前最好用一小段验证集跑一下 logits 分布对比确认移除后的分布没有进入饱和区。我在一个 300M 模型上做过一次完整移除logits max 平均上升了 4~6虽然没炸但量化位宽需要放宽一档否则精度损失明显。这些都是成本提前想清楚能省很多返工。在我现在的实践中300M 以上的预训练默认装 QK-norm 和 softcap 两道保险训练末期按部署目标决定要不要退火300M 以下优先做日志探针等看到证据再上保险丝。这个思路不一定适合所有人的资源和场景但它至少帮你把“小模型要不要装 QK 保险丝”这个问题从玄学变成了可以靠数据回答的问题。
前阵子有个做 1B 级 decoder-only 模型的朋友来找我说训练跑到第 4000 步左右 loss 曲线开始隔三差五冒尖每次 spike 都要掉接近 0.5得花上千步才爬回来。他看了一圈开源仓库发现 QK-norm 和 softcap 在主流大模型里几乎是标配于是拿着这个问题来问我小模型要不要也装上这根“QK 保险丝”这个问题我很熟。最近手头正好同时跑了一条 70M、一条 300M、一条 1B 的预训练实验中间试过 QK-norm、attention logits softcap以及训练末期把 QK-norm 退火anneal到近乎关掉的方案。这几种手段在社区里经常被放进同一个句子里讨论但它们防的是截然不同的风险。这篇文章我把这几根保险丝拆开讲清楚最后给出一套可以照抄的分档决策和默认配置。标题里说的“小模型”我按预训练领域常说的 100M ~ 1B 参数来理解训练数据在 50B ~ 300B token 这个量级。至于那些 7B 起步的大模型装 QK-norm 已经是行业默认小模型往往让人犹豫的主要原因是加上去多一套算子、多几个初始化要调而且很多人直觉上觉得“参数小、深度浅应该不会像大模型那样炸”。这个直觉一半对一半不对。下面我从注意力 logits 为什么失控说起。1. 先看 QK-norm 和 softcap 在防什么注意力 logits 失控的完整链路要回答“小模型要不要装”得先搞明白 QK-norm 和 softcap 到底在防什么。它们防的是同一个东西的两个阶段注意力 logits 的数值失控。Transformer 多头注意力里每一层都在做这样一个点积logits Q K^T * scale其中 scale 通常是 1 / sqrt(d_head)d_head 是每个注意力头负责的维度。比较理想的情况是Q、K 投影之后的每个维度方差保持在 1 附近这样一来 logits 的方差大概也是 1 的量级序列长度 4k 时一批 logits 的最大值也就是 3~5。这个范围对 softmax 来说非常健康梯度既不会消失也不会爆炸。但训练到中后段事情会变。Q、K 投影权重是参与训练的优化器为了降低当前 loss完全有动力把某些 head 的 Q/K 向量尺度推大从而让某几个 token 的注意力分数特别高形成“锐利”的注意力分布。这种锐利分布在部分任务上确实能压 loss但它有很强的正反馈logits 越大softmax 越接近 one-hot对应路径的梯度占绝对主导其他 token 几乎拿不到梯度。模型就会越来越依赖那几个特殊 token更新信号越来越偏参数越走越歪直到某一步 logits 冲到 30、50loss 直接拉出一个 spike。这中间还有一层隐蔽的数学问题d_head 越大越容易失控。点积本身是 d_head 维的求和如果 Q/K 每维方差被放大到 2~3在训练中很容易发生logits 的典型范围就会从 ±3 放大到接近 ±10。如果再叠上长序列里位置编码的相互作用某些 token pair 出现系统性大值就一只脚踩进了饱和区。小模型虽然 d_head 往往不大64 居多但仍然要面对这个机制只是概率和幅度不同。1.1 为什么说 QK-norm 是“拆掉失控的发动机”QK-norm 的做法很直白在 Q、K 投影之后、做点积之前各自过一个 RMSNorm有的实现用 LayerNorm把向量的 RMS 重新拉回固定值。这样无论训练中 Q/K 内部怎么“膨胀”最终参与点积的向量长度都被重新归一化了。关键的数学约束是柯西不等式|q · k| ≤ ||q|| · ||k||归一化之后两个范数都被压到固定尺度单个 logits 想无限往高冲就失去了基础。我习惯把 QK-norm 叫“拆发动机”它没有直接限制 logits 多高而是把能推动 logits 变高的“动力源”——也就是 Q/K 向量长度的自由度——拿掉了。方向由向量方向表达大小由固定缩放统一控制。需要补充的是norm 之后的缩放要仔细选。RMSNorm 会把 Q/K 每一维的 RMS 变成 1等价于向量的 L2 范数被固定为 sqrt(d_head)。这时候点积之后继续乘 1/sqrt(d_head)logits 的典型标准差仍然在 1 附近注意力分布和标准初始状态差不多这是最不容易出错的配置。我见过不少第一次加 QK-norm 的同学把 scale 漏掉或者改成了别的值结果注意力要么变成均匀分布要么重新变得过于尖锐。Mistral、Gemma、DeepSeek 这几个系里都能看到 QK-norm 的身影实现细节差异很大有的带可学习 scale有的不带但核心逻辑都是同一件事让 Q/K 的模长不再参与梯度博弈。1.2 softcap 是其次的闸门硬截断和 tanh 软截断softcap 直接作用在 logits 上跟 QK-norm 是两条互补路径。最常见的实现是 tanh softcaplogits C * tanh(logits / C)C 是超参比如 30、50、100。logits 的绝对值被压到 C 以内即使某一步真的爆炸模型也不会立刻 loss NaN而是变成“平滑压缩”后的结果。这相当于给注意力分布加了一层保险闸门。也有实现直接用 clamp(logits, -C, C)简单且没有额外算子。但 clamp 的问题是超过边界的 logits 在 softmax 里会“贴边”堆积多个 token 共享同一个上限梯度语义变差tanh 版本因为连续光滑logits 很大时梯度趋近 0更像“我不想管这个死活”反而不容易过拟合出新的错误结构。大模型社区里近年更流行 tanh softcapLLaMA 3 的 attention 就用了这个思路。我自己的实验也倾向于 tanh softcapC 从 30 起步再根据 logits 直方图调整。1.3 两个保险丝一起装会怎样QK-norm 加 softcap 同时开等于从“输入端”和“输出端”各设一道防线。一个大致的分工是QK-norm 保证训练过程中 logits 的尺度不会整体漂移softcap 负责兜住个别极端 token pair。一起开通常不会冲突但要注意两个坑。一是别把 softcap 的 C 设得太小比如 20 以下。QK-norm 已经限制了大量 logitsC 太小会压制正常的对比度让注意力变软最终 loss 反而变差。二是一旦开了 QK-norm想让注意力更锐利应该调 norm 路径上的额外缩放而不是无限压 C。固定 standard scale然后靠更大的 C 去放行效果最可预期。2. 小模型的 attention 为什么也会冒尖三大实际风险来源大模型必须装 QK-norm多半是因为层深、头多、序列长失稳概率随规模指数上升。小模型是不是就可以不装了我在 70M / 300M / 1B 三条线上跑下来结论是小模型通常不会像 7B 那样系统性高频地炸但依然有三种场景非常容易冒尖而且一旦冒尖恢复代价并不会因为模型小就变便宜。2.1 层数少不代表不会炸特殊 token 和局部参数崩溃小模型深度浅层间叠加效应弱所以 QK-norm 的收益确实比深模型小。但注意力 logits 失控并不只来自“多层放大”。最典型的是特殊 token 竞争。训练语料里 BOS/EOS、数字、罕见 Unicode、重复出现的模板词这些 token 天然会和上下文里的常规 token 形成强烈的注意力竞争。小模型因为容量有限更容易把“注意力资源”押在一个固定 token 上比如每句话都去看 BOS 或者去看上一个 token。某个 batch 里如果这类 token 大量出现logits 就会在一层内被局部放大不需要多层配合也能 spike。我还遇到过一种“局部参数崩溃”的形态某个注意力头在某个训练阶段演化成几乎只输出一个固定方向K 投影的某几维方差持续膨胀导致该头对所有 query 都给出近似相同的极值 logits。这种崩溃在大模型上表现为某个 head 失效在小模型上则直接表现为 batch loss 的孤立尖峰。QK-norm 对这类问题很有效因为它直接把 K 向量的长度压住了崩溃头失去“靠长度得分”的空间。2.2 bf16 低精度训练是隐形的帮凶第二个风险源来自低精度训练这块经常被低估。很多小模型实验为了提高吞吐已经切到 bf16 混合精度q/k 投影后的张量也是 bf16 存储。bf16 的指数范围够大不容易溢出但尾数只有 7~8 bit表示一个大数比如 20.0时相邻可表示数之间的间隔接近 0.1。Q 向量和 K 向量本身动态范围越大点积结果在 bf16 下被量化掉的误差就越显著。QK-norm 在这里有一个隐藏优势经过归一化后Q/K 向量的每个元素都被压到比较小的动态范围里bf16 表达中的相对量化误差反而更均匀。也就是说即使不出现剧烈 spikeQK-norm 也能让低精度下的注意力分数更干净一些。我在这几个实验里统计过打开 QK-norm 后bf16 全精度对比的 logits 相对误差大致能下降一个数量级虽然具体数值和硬件 kernel 实现有关但趋势是稳定的。当然“下降一个数量级”这说法要看基数。如果模型规模小、d_head64、bf16 点积本来误差就不大QK-norm 带来的精度提升可能对最终 loss 毫无感知。但反过来如果你的 loss 曲线有那种“高频小毛刺”切换上下文长度或数据分布后冒出一连串小 spike那十有八九和低精度表达有关上 QK-norm 比调学习率更对症。2.3 从 loss spike 日志反推“是哪种炸法”排查 loss spike 时我最反感“先调学习率再说”的做法。更有效的路径是把 spike 附近的梯度、logits 统计一起拿出来看。我自己常用的探针很小就是在 attention forward 里每 N 步采样几批数据记录这批数据在所有层的 logits max、logits 均值、q/k 的 L2 范数 p90。怎么读这些指标呢logits max 平时 8~15spike 时冲到 30且 q/k L2 范数同步上涨大概率是 Q/K 向量长度失控首选 QK-norm。logits max 平时只有 4~6spike 时跳到 20 左右但 q/k 范数没怎么变通常是特定 token/位置组合的偶发尖峰softcap 就够用。logits 和 q/k 范数都正常但 loss spike 延迟半拍到下一步更可能是某个 batch 本身数据异常或梯度累积问题不要动 QK-norm先检查 loss 重算逻辑和数据清洗。这个区分逻辑其实也呼应了标题——QK 保险丝不是万能药。很多小模型实验真正的问题根本不在 QK而在数据处理和 LR schedule贸然加 QK-norm 只会增加变量最后调了一周才发现源头是 data sampler。3. 退火 QK-norm把保险丝改成可拔掉的手闸如果你决定给小模型上 QK-norm接下来马上会碰到另一个问题训练结束后这个 norm 留不留如果留着推理时每层都要多被执行一两次 RMSNorm虽然单个 norm 不贵但在小模型部署到边缘设备时算子少一个是一个。更关键的是有些小模型实验的目的是产出“无 QK-norm 依赖”的基线模型方便后续跟其他不带 norm 的架构公平对比。于是就有了退火 QK-normQK-norm annealing的做法。3.1 为什么要把好好的 QK-norm 退火到不存在直接说动机有三类。第一类是推理端省算子。QK-norm 本身计算量占比很小在 7B 以上模型里完全可忽略但在 300M 的移动端模型里每层多两个 RMSNorm 会让 latency 增加几个百分点积少成多。如果最终产品想裁掉它最好的办法不是训练完硬删而是在训练末期让模型适应没有 norm 的分布。第二类是为了让 logits 分布与部署环境对齐。很多推理框架或量化工具会假设注意力 logits 是某种尺度分布例如做 KV cache 量化时会期望 logits 的绝对值在一个相对固定的 range 里。带 QK-norm 的 logits 分布往往比纯 raw 点积更“瘦”一旦部署端按裸模型校准可能出偏差。退火到最后让模型学会在 raw 尺度上保持稳定量化时更省事。第三类是放松 QK-norm 对注意力对比度的限制。QK-norm 本质上引入了“无长度信息”的约束这对某些需要极端锐利的任务未必最优。在最后阶段退到低 alpha让模型有机会用 Q/K 向量长度重新编码部分信息实测中偶尔能补回一点点下游分数。3.2 一个可落地的 alpha 退火实现我推荐一种简单的向量级插值方案而不是在 logits 层面直接混。原因是向量层面的干预语义更清晰训练早期 alpha1完全走 norm 路径训练末期 alpha 线性从 1 降到 0模型逐步切回 raw q/k 路径。# attention forward 内部伪代码 # alpha 由训练 loop 传入1.0 表示完全 QK-norm0.0 表示无 QK-norm q_raw self.q_proj(hidden_states) k_raw self.k_proj(hidden_states) if self.qk_norm_alpha 0: q_n self.q_norm(q_raw) # RMSNorm无 affine k_n self.k_norm(k_raw) q q_n * self.alpha q_raw * (1 - self.alpha) k k_n * self.alpha k_raw * (1 - self.alpha) else: q q_raw k k_raw if self.use_rope: q, k apply_rope(q, k, positions) scale self.head_dim ** -0.5 # 标准缩放退火期间保持不变 logits torch.matmul(q, k.transpose(-2, -1)) * scale if self.softcap_scale is not None: logits self.softcap_scale * torch.tanh(logits / self.softcap_scale)这段代码里有两个细节容易被忽略。第一个是退火时不要开可学习的 qk scale。如果 scale 是可学习的优化器会在 alpha 减小时偷偷把 scale 调大来对冲等到 alpha0 时 scale 可能已经长成一个危险值最后几步容易爆新的 spike。退火期间务必把 scale 固定住。第二个细节是这个插值方案会让中间态的 q/k 向量既有 norm 路径的成分也有 raw 路径的成分logits 分布会有一个过渡窗口模型需要一点时间重新适应所以退火窗口要留足。如果你想保留 norm 到最后一刻又想省推理算子还有一个偷懒方案训练最后 N 步直接把 alpha 置 0然后观察 loss 反弹。有效但反弹幅度一般比线性退火大最终 checkpoint 可能要多跑一小段纯 raw 路径才能稳住。我实测下来的经验是线性退火更平滑但会额外占用几千步最终 loss 和完全保留 norm 的 checkpoint 差不太多不同规模、不同数据上差异很大大体在 0.01~0.03 这个量级。3.3 退火窗口多长才不明显反弹这个问题没有标准答案我给出我在小模型上的观察范围供参考。退火窗口占训练总步数的 3% 以内alpha 从 1 到 0 会伴随明显的 loss 反弹反弹后需要再训一小段才能回到之前水平。5%~10%通常能看到一个先涨后平的曲线反弹幅度在 0.02~0.05 量级大部分模型可以接受。超过 10%有时反而没有必要因为小模型本来总步数就不多浪费太多步数在退火上不划算除非你同时在做序列长度切换。如果训练 schedule 里本来就有“最后 10% 线性降到最小 LR”的阶段我一般把 QK-norm 退火和 LR 退火重叠在一起开两个变化同时进行模型反而更容易平滑过渡。不过要小心如果 LR 已经降到接近 0而 alpha 还在快速下降模型几乎没有能力修复切换带来的分布变化所以最好让 alpha 的下降起点比 LR 的终点早一点留出几百步用微跌后的 LR 稳定。还有一类更高级的做法是“退火 softcap”即 softcap_scale 在训练后期从 30 线性增大到 1000等效于把 cap 放宽到近似关闭。原理和退火 QK-norm 一样把保护措施慢慢撤掉让模型在没有保护的情况下仍能维持。两者可以联动但我建议不要同时退火 QK-norm 和 softcap变量太多一旦 loss 反弹不容易定位是哪个保险丝的问题。4. 小模型到底要不要装规模分档与默认配置现在回到标题的核心问题小模型要不要装 QK 保险丝。我的主张是不要把它当成“标配”或“不配”而是按模型规模、训练数据、部署目标三件事来分档。下面这张表是我现在做预训练时用的默认参考。参数规模d_head 典型值默认 QK-norm默认 softcap说明 100M32~64不装可选 C30先查 logits 统计数据质量优先级最高100M ~ 300M64推荐开最后可退火C50低成本高保障尤其要切长上下文时300M ~ 1B64~128装C50~100默认装训练末期按部署需求退火1B ~ 3B64~128装C100和主流开源配置看齐别在这上面省这个表不是拍脑袋。我自己的观察是100M 以下的模型即使出现 logits 尖峰往往也是数据里某个 token 群体的偶发问题加 QK-norm 属于用更多参数换稳定但没有控制住真正的根因。100M 以上开始训练总 token 数普遍在 50B~200B 这个区间训练中途经常需要切换上下文长度或混合数据集QK-norm 的兜底价值明显变大。4.1 装的位置与超参选择全装还是只装上半层一个常见问题是QK-norm 要每个注意力层都装吗还是只装前几层、后几层在小模型上我的建议是要么全装要么全不装不要只装一半。原因不是技术上有害而是“只装一半”会让超参排查变得极其痛苦如果有的层有 norm、有的层没有logits 的整体尺度在各层之间不一致一旦出现 spike你很难判断是哪一层加了 norm 才让它稳定的也很难判断是哪一层没加才让它炸的。实在想省算力我试过另一种替代方案只对 K 做 normQ 不做。K 是“被检索的键”在长上下文中更容易出极端值只 norm K 已经能挡住不少尖峰且算子只有原来的 1/2。但这只适合“轻度保护”场景。如果要开退火还是建议 Q/K 都 norm否则 alpha 降到 0 时两边路径的语义不对称退火会更不稳定。超参方面最简单的组合是RMSNorm不带 affine 参数。省掉 scale/bias 两个可学习量也更容易避免训练初期 scale 漂移。norm 之后保持标准缩放1 / sqrt(d_head)不要乱改 scale。softcap C 先设 50观察一段时间如果 logits 直方图的 99.9 分位数长期低于 15可以调大到 100 或者关掉。如果开了 RoPEnorm 放在 RoPE 前后都可以因为 RoPE 是正交变换不改变向量的范数。我习惯先 RoPE 再 norm原因是很多开源实现这样排方便对照。4.2 我当前的实际基线配置顺手贴一条我现在跑 300M 模型的默认 config都是实测调整过的d_model768d_head64num_heads12layers12每个 attention 前向里q_proj / k_proj 输出后接 RMSNorm无 affine随后继续乘标准缩放1 / sqrt(64)也就是 0.125logits 层做 tanh softcapC50优化器 AdamW初始 LR 3e-3没有 QK-norm 的对照实验我把 LR 降到 2.5e-3因为没 norm 时 LR 稍高一点就会出现隐约的小毛刺训练最后 8% 步数把 QK-norm alpha 从 1 线性降到 0softcap 保持不动推理阶段alpha 设为 0也就是删掉 q/k 的 norm保留 softcap。这套配置在当前几个 300M~1B 的实验上跑得很稳loss 没有出现超过 0.3 的孤立 spikes。作为对照我把 QK-norm 关掉、其他不变重跑了一次大部分时候曲线几乎重合但在上下文长度从 512 切到 2048 的那几百步以及某个包含大量重复广告文本的 batch 出现时loss 会高出 0.3~0.5。也就是说对小模型而言 QK-norm 多数时候是“不那么必需的保险丝”但确实管住了切换期和脏数据期这两类最大的痛点。5. 建议的排查顺序先用日志定位再决定上哪根保险丝最后这段给实操路径。如果你正在跑一个小模型第一条建议永远不是“马上加 QK-norm”而是把 attention 的 logits 和 q/k 范数统计埋进训练日志。绝大多数失稳问题靠统计数据比靠感觉快得多。5.1 给 attention logits 加统计探针在 attention forward 里面加一段代码不要每步都算训练时每 50 步取一个 batch 采样即可if trainer.global_step % 50 0 and random.random() 0.1: with torch.no_grad(): # logits 是 [batch, heads, seq, seq]采样一部分即可 sample logits.flatten().float() stats { attn_logits_max: sample.max().item(), attn_logits_p99: torch.quantile(sample, 0.99).item(), attn_logits_mean: sample.mean().item(), q_l2_p90: torch.norm(q.float(), dim-1).quantile(0.90).item(), k_l2_p90: torch.norm(k.float(), dim-1).quantile(0.90).item(), } log_metrics(stats)这个探针的开销几乎可以忽略。注意要在spike 发生前就开始记录否则事后没有对比依据。拿到这些指标后按我在 2.3 里说的三类区分方式判断。我自己用的阈值大致是logits p99 平时小于 10max 平时小于 20如果 max 超过 30 或者 p99 超过 15就进入值得干预的区间了。5.2 分级处方从观察、softcap 到 QK-norm按严重程度给一个分级处方方便快速照做。第一档logits 几乎健康loss spike 频率极低每 5000 步一次以下。不做任何架构改动先检查 data sampler 和 loss 重算。很多人纠结的 QK-norm 其实没有用武之地。第二档logits max 偶尔超过 30但 q/k 范数稳定。加 tanh softcapC50跑 500 步看变化。如果尖峰次数下降但仍有残留把 C 调到 30或者配合 grad clip 1.0。这一步基本不影响最终 loss成本极低。第三档q/k 范数明显膨胀、logits 高频冒尖。上 QK-normscale 保持1 / sqrt(d_head)训练初期多给 1000 步观察不要因为前面 loss 比对照高 0.05 就立刻回滚norm 需要一点时间让模型习惯新的 logits 分布。第四档已经上了 QK-norm 还炸。这时候不要继续加保险丝回头查初始化、LR、batch size 和 cosine schedule大概率是别的问题。曾见过有人把 softcap、QK-norm、额外 grad clip 全叠上也压不住 spike最后发现是某个 head 的 bias 初始化为 10导致 logits 一开始就偏瘫。5.3 最后分享两个小教训一个是退火别开太晚。我做过一次 1B 模型总步数 20 万步把 QK-norm 退火排到最后 1 万步才开始。alpha 刚开始下降的那几百步没事我以为很稳结果第 2 万步直接出现 0.7 的 loss 反弹后面用 3000 步才磨平。后来把退火窗口提前到总步数的 10%虽然中途也有一点小起伏但最终 checkpoint 明显更干净。另一个教训是关于 softcap 的超参。softcap 开太大就形同虚设开太小会影响模型在少数高置信场景下的表达能力。我最早的实验把 C 设成 5想着“越紧越安全”结果注意力几乎失去了锐利度几个需要长程定位的任务指标掉了不少。后来把 C 调到 30、50 才正常。所以如果你拿不准 C先设 50 是相对安全的起步值真正值得调的是 logits 尺度而不是过度依赖 C。还有一个经常被忽略的前提如果你最终部署时既不想留 QK-norm也不想留 softcap退火只能帮你“软着陆”但不能消除所有风险。部署前最好用一小段验证集跑一下 logits 分布对比确认移除后的分布没有进入饱和区。我在一个 300M 模型上做过一次完整移除logits max 平均上升了 4~6虽然没炸但量化位宽需要放宽一档否则精度损失明显。这些都是成本提前想清楚能省很多返工。在我现在的实践中300M 以上的预训练默认装 QK-norm 和 softcap 两道保险训练末期按部署目标决定要不要退火300M 以下优先做日志探针等看到证据再上保险丝。这个思路不一定适合所有人的资源和场景但它至少帮你把“小模型要不要装 QK 保险丝”这个问题从玄学变成了可以靠数据回答的问题。
