模型稳定性、z-Loss、注意力 softmax 稳定化与软截断

模型的稳定性是什么意思？

概述

模型稳定性指的是：训练过程中，损失和梯度不要出现剧烈的尖峰或爆炸，能够平稳收敛。

在 LLM 训练中，常见的不稳定现象包括：

• 损失曲线看起来正常，但梯度范数（gradient norm）突然出现巨大尖峰；
• 梯度范数持续增大，最终爆炸，训练崩溃；
• 损失突然飙升，模型无法恢复。

OLMo2 论文展示过一张图：蓝色曲线的梯度 L2 范数布满尖峰，范数值完全失控；橙色曲线则始终保持在低位，这才是理想状态。

哪些组件容易引发不稳定？

Transformer 中任何部分都可能出问题，但softmax 是公认的“问题儿童”：

• 指数运算 $e^x$ 可能产生极大值；
• 除法可能产生除零或数值溢出；
• 两个 softmax 模块：
  • 自注意力 softmax：logits $Q{K}^T/\sqrt{d_k}$ 可能过大，导致 softmax 饱和、梯度消失；
  • 输出层 softmax：logits 可能爆炸，使交叉熵损失不稳定。

下面分别讲三种稳定化技术。

z-Loss 技术——解决输出层的 softmax

概述

z-Loss 是一种正则化项，加到总损失中，目的是防止 softmax 的归一化因子 $Z$ 过大或过小。

输出层 softmax：
$$p_i=\frac{e^{z_i}}{{\sum }_j{e}^{z_j}}=\frac{e^{z_i}}{Z}$$

其中 $Z={\sum }_j{e}^{z_j}$ 是归一化因子（partition function）。

当某些 logits 非常大时，$Z$ 会爆炸，导致数值不稳定。

z-Loss 惩罚 $\log Z$ 偏离 0：
$${\mathcal{L}}_z=\lambda \cdot (\log Z{)}^2$$

总损失：
$${\mathcal{L}}_{total}={\mathcal{L}}_{CE}+\lambda \cdot (\log Z{)}^2$$

PaLM 使用 $\lambda =10^{-4}$。

数值例子

设输出层 logits 为：

z = [10, 5, 0]

计算 $Z$：
$$Z=e^{10}+e^5+e^0=22026.47+148.41+1=22175.88$$

计算 $\log Z$：
$$\log Z=\log (22175.88)\approx 10.007$$

z-Loss ($\lambda =10^{-4}$)：
$${\mathcal{L}}_z=10^{-4}\times (10.007{)}^2=10^{-4}\times 100.14=0.010014$$

对比另一组较小的 logits：

z = [2, 1, 0]

$$Z=e^2+e^1+e^0=7.389+2.718+1=11.107$$

$$\log Z=\log (11.107)\approx 2.408$$

$${\mathcal{L}}_z=10^{-4}\times (2.408{)}^2=10^{-4}\times 5.798=0.0005798$$

对比：

logits 越大，$\log Z$ 越大，z-Loss 越大。

训练时，这个惩罚项会迫使模型不要把 logits 推得过大，从而保持 softmax 数值稳定。

解决注意力层的 softmax 有什么作用？——QK 归一化

概述

标准注意力：
$$\text{logits}=\frac{Q{K}^T}{\sqrt{d_k}}$$

然后 softmax。

如果 $Q$、$K$ 的数值很大，点积会非常大，softmax 会变得极端尖锐（接近 one-hot），梯度接近 0，训练困难。

QK 归一化（QK-Norm）的做法是：在计算点积之前，先对 Q 和 K 分别做 LayerNorm/RMSNorm：
$$Q^{\prime }=\text{Norm}(Q),K^{\prime }=\text{Norm}(K)$$

$$\text{logits}=\frac{Q^{\prime }{K}^{\prime T}}{\sqrt{d_k}}$$

这样 Q、K 的数值被标准化，点积不会失控，softmax 输入保持在合理范围。

数值例子

设 $d_k=4$，某个 Query 和 Key 向量为：

Q = [10, 20, 30, 40]
K = [5, 10, 15, 20]

不做归一化：

点积：
$$Q\cdot K=10\times 5+20\times 10+30\times 15+40\times 20=50+200+450+800=1500$$

缩放：
$$1500/\sqrt{4}=1500/2=750$$

softmax 输入为 750，$e^{750}$ 直接溢出，数值不稳定。

做 QK 归一化：

对 Q 做 LayerNorm：

• 均值：$(10+20+30+40)/4=25$
• 标准差：$\sqrt{((10-25{)}^2+(20-25{)}^2+(30-25{)}^2+(40-25{)}^2)/4}=\sqrt{(225+25+25+225)/4}=\sqrt{125}\approx 11.18$
• 归一化：
  $$Q^{\prime }=[(10-25)/11.18,(20-25)/11.18,(30-25)/11.18,(40-25)/11.18]$$
  $$=[-1.342,-0.447,0.447,1.342]$$

对 K 做 LayerNorm：

• 均值：$(5+10+15+20)/4=12.5$
• 标准差：$\sqrt{((5-12.5{)}^2+(10-12.5{)}^2+(15-12.5{)}^2+(20-12.5{)}^2)/4}=\sqrt{(56.25+6.25+6.25+56.25)/4}=\sqrt{31.25}\approx 5.590$
• 归一化：
  $$K^{\prime }=[(5-12.5)/5.590,(10-12.5)/5.590,(15-12.5)/5.590,(20-12.5)/5.590]$$
  $$=[-1.342,-0.447,0.447,1.342]$$

点积：
$$Q^{\prime }\cdot K^{\prime }=(-1.342)(-1.342)+(-0.447)(-0.447)+(0.447)(0.447)+(1.342)(1.342)$$
$$=1.801+0.200+0.200+1.801=4.002$$

缩放：
$$4.002/2=2.001$$

softmax 输入为 2.001，完全在合理范围内。

对比：

作用

• 控制 softmax 输入的数值范围，防止饱和和梯度消失；
• 允许使用更激进的学习率，优化器发挥更大作用；
• Gemma2、DCLM、OLMo2、Chameleon 等模型采用；
• NVIDIA 团队实验：QK 归一化能提升效果，而软截断反而使困惑度变差（11.19 → 更差）。

对输入 softmax 的 logits 进行软截断

概述

软截断（Soft-clipping / Logit-capping）在 softmax 之前，用 tanh 把 logits 压缩到有界范围：
$${\text{logits}}_{\text{clipped}}=\text{cap}\cdot \tanh \left(\frac{\text{logits}}{\text{cap}}\right)$$

其中 cap 是阈值，Gemma2/OLMo2 使用 cap = 30.0。

当 logits 远大于 cap 时，tanh 饱和到 1，输出接近 cap；当 logits 很小时，tanh 近似线性，几乎不改变。

数值例子

设 cap = 30。

logits = 100：
$$100/30=3.333$$
$$\tanh (3.333)\approx 0.9974$$
$$30\times 0.9974=29.92$$

原本 100 被压缩到 29.92。

logits = 10：
$$10/30=0.333$$
$$\tanh (0.333)\approx 0.3215$$
$$30\times 0.3215=9.645$$

原本 10 被压缩到 9.645，变化很小。

logits = 0：
$$\tanh (0)=0$$
$$30\times 0=0$$

不变。

logits = -100：
$$-100/30=-3.333$$
$$\tanh (-3.333)\approx -0.9974$$
$$30\times (-0.9974)=-29.92$$

被压缩到 -29.92。

对比表：

作用

• 防止极端 logits 导致 softmax 饱和；
• 保持梯度有效；
• Gemma2、OLMo2 采用；
• 但 NVIDIA 团队实验发现：基线困惑度 11.19，软截断反而使效果变差；QK 归一化却能提升效果。

为什么软截断可能变差？ 因为 tanh 在饱和区梯度接近 0，反而削弱了模型学习极端重要 token 的能力；

而 QK 归一化是从源头控制数值，保留了线性区的梯度。

三种技术的对比

总结

模型稳定性：训练中损失和梯度不出现尖峰、爆炸，能平稳收敛。softmax 是最容易出问题的组件。

z-Loss：在输出层 softmax 上，惩罚归一化因子 $Z$ 过大，${\mathcal{L}}_z=\lambda (\log Z{)}^2$。数值上，logits 从 [2,1,0] 增大到 [10,5,0]，z-Loss 从 0.00058 增到 0.010，迫使模型保持 logits 温和。

QK 归一化：在注意力 softmax 前对 Q、K 做 LayerNorm，把点积从 1500 压到 4.002，避免 softmax 饱和，允许更大学习率。

软截断：用 cap · tanh(logits/cap) 把 logits 压到 $(-cap,cap)$，100 变 29.92。能防止极端值，但可能削弱极端重要 token 的学习，NVIDIA 实验显示效果可能变差。

一句话：稳定性技术都是围绕 softmax 做文章——z-Loss 管输出层归一化因子，QK 归一化管注意力输入范围，软截断管 logits 极端值；三者各有适用场景，QK 归一化目前证据最积极。