最近和几个做预训练的同行聊天话题总绕不开一个小模型稳定性的问题训练一版 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, scaleTrue): 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 万步用 cap50 的 softcaploss 曲线非常平稳后来我把 cap 改到 30曲线出现轻微震荡但模型在评测集上的困惑度反而微降。这说明 softcap 的“手感”是非常依赖任务阶段的需要在稳定性和表达力之间找平衡。3.2 和 QK-norm 组合使用时的注意点很多开源模型其实是“QK-norm softcap”一起用的比如一些采用规范化注意力的大模型。两者并不冲突QK-norm 压缩常规尺度softcap 兜底极端异常。但小模型加这套组合的时候要特别留意训练动态如果 Q、K 在 RMSNorm 之后又接了一个可学习 scale这个 scale 已经被 push 到比较大的值那么 softcap 相当于在“scale 刚大起来”的边缘做压制效果可能会互相抵消。我踩过的一个坑就是把 scale 初始化为 1.0加了 softcap30训练 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 的核心思想就是让约束强度随训练进度变化。最常用的方式是做“严格到宽松”的退火训练早期系数 λ 接近 1QK-norm 的约束完整接入训练后期系数逐渐衰减到 0让模型恢复原始的自由度。也有反向退火的方案即先不加后期加这在我试下来是“问题已经出现了才做补救”效果不如预期好因为注意力坍缩一旦发生梯度信号已经污染同修复基础损伤。4.2 动态强度调整的实现与观测具体实现不需要改模型结构只需要在 forward 里让 QK-norm 的 scale 变成一个动态系数随 step 线性或者余弦衰减def anneal_coef(current_step, total_steps, start1.0, end0.0, warmup_ratio0.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(scale0.8) softcap(cap50)适合数据极杂、训练步数超长。平衡高效型QK-norm(scale1.0) ln_anneal线性退火到 0.2适合常规规模的预训练。资源最少介入型只加softcap(cap80)适合在已有架构上做最小改动快速稳定。这三档配置都经过至少一个实际小模型项目的检验直接复用问题不大。但要提醒一点任何保险丝方案的超参数都必须配合你的训练时长、学习率、batch size 来调整不存在全局最佳的固定值。6. 实操中的坑与排查技巧6.1 常见问题的快速定位方法我把自己做训练稳定性排查时用到的判断路径整理了一下按出现频率排序现象可能原因排查方式建议处置训练中后期出现单点 loss spikeQK 点积突然变大或 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_stdQK 点积的标准差。这个值应该在 QK-norm 或 softcap 的约束下保持在一个稳定范围内如果持续攀升说明保险丝没起作用或者强度不够。qk_norm_scaleQK-norm 的可学习 scale 值。如果它持续增长说明模型在试图跨越约束边界。attention_entropy每个注意力头的平均熵。如果整体熵过高保险丝限制了表达如果某个头熵特别低可能该头已经坍缩。max_logitsQK 点积的最大值。softcap 生效时这个值应严格小于 cap如果接近 cap说明有异常 token 正在冲击上限。这四个指标配合 original loss 一起看基本能把问题定位到“是否保险丝问题”还是“优化器问题”还是“数据问题”。定位清楚之后再动手改参数效率会高很多。6.3 终极建议先加最便宜的保险如果要给一个不熟悉这个领域人建议我的结论很直接小模型加上 QK 保险丝的成本很低但收益可能比想象高。计算开销上两个 RMSNorm 对整体训练时间影响可忽略显存增加也微乎其微。调试成本才是主要成本好在有 softcap 这种几乎零参数的方案兜底。所以我的建议是当你准备训练一个 1B 以下的小模型与其赌它不会出问题不如先在最不关键的位置加一个最便宜的保险single softcap把训练跑稳再说。等模型和数据成熟了再考虑是否需要上退火 QK-norm 精细调优。小模型很容易出现“训练过程稳定评测结果不行”的问题而这类问题往往和注意力分布过度平均有关退火 QK-norm 比传统调整学习率的办法更直接。我在实际操作中还会做一件小事把 QK-norm 从一种“保险”思路理解成“初始化偏置”。它其实是在训练早期给注意力提供一个先验告诉模型“你暂时不需要用大 logits 来表达关注”慢慢再教它学会用。这个视角上的转换帮我省了不少调参时间——我不再把它当成一个防御性组件而是当成一个可设计的先验模块。这样无论在哪个训练阶段我都知道它应该起什么作用、什么时候该松手。最后再分享一个小技巧不管选哪类保险丝记得固定住随机种子对比实验否则很难判断效果是来自保险丝还是来自随机噪声。小模型训练的随机性方差比想象中大同样配置跑两次可能差 2% 到 3%如果评估指标只有 1% 的提升无法下结论。多跑几个种子统计稳定性再下判断这才是科学的方法。