深层网络的训练难题#

在深度神经网络中，每一层的输出都会成为下一层的输入。当网络只有几层时，各层激活值的尺度差异还可以接受；但随着层数加深，激活值的尺度会随层数逐层放大或收缩。如果某一层让信号平均放大一点，几十层之后这个效应就会指数累积，最终导致激活值爆炸；反过来，信号也可能逐层衰减直至消失。梯度反向传播时面临同样的问题，这就是经典的梯度爆炸与梯度消失。

早期的解释来自 Batch Normalization 原论文提出的“内部协变量偏移（Internal Covariate Shift）”：由于前面层的参数在训练中不断更新，后面每一层看到的输入分布也随之持续变化，因此每一层都必须不断适应一个“移动的输入分布”，导致训练困难。

不过，后续工作 How Does Batch Normalization Help Optimization? 通过实验表明，内部协变量偏移不足以解释 BN 的全部优化收益：即使人为向归一化后的激活注入噪声、加剧分布偏移，带归一化的网络依然训练得更快更稳。目前更被广泛接受的解释是，归一化让损失曲面变得更加平滑，参数更新方向更稳定，从而允许使用更大的学习率。至于“归一化为什么有效”这个问题本身，到今天仍没有完全统一的结论，但这并不妨碍它成为深度学习中最重要的工程手段之一。

归一化与残差流#

在 Transformer 中，问题还有一个更具体的来源：残差连接。对于输入 $\mathbf{x}$ 与某个子层 $F$，例如 Attention 或 FFN，残差连接的计算为：

这意味着每一层的输出都会在原有表示上不断累加。如果各子层输出近似不相关、方差大致恒定，那么残差流（residual stream）的方差会随层数近似线性增长。一个简单的实验就能观察到这个现象：

1
x = torch.randn(8, 128, 512)               # [N, L, C]
2
stds = [x.std().item()]
3
for _ in range(64):
4
    x = x + torch.randn(8, 128, 512) * 0.1   # 模拟各子层的输出
5
    stds.append(x.std().item())
6
# stds 大致随 sqrt(层数) 增长，64 层后尺度已明显漂移

归一化的作用，就是在每个子层读取残差流之前，把输入重新拉回一个固定的尺度，使各层看到的输入尺度保持稳定：

图中的平稳曲线表示子层读取的 Norm 输出，残差流本身仍可持续增长。这个“在残差流上进行尺度校正”的视角，是后面理解 Pre-Norm 与 Post-Norm 差异的基础。

统一形式#

BatchNorm、LayerNorm、InstanceNorm 和 GroupNorm 可以写成下面的标准化形式；稍后介绍的 RMSNorm 则省略减均值。对输入元素 $x$，先减去均值、再除以标准差，最后做一层可学习的仿射变换：

其中 $\mu$ 与 $\sigma^2$ 是统计得到的均值与方差，$\gamma$ 与 $\beta$ 是可学习参数。整个计算可以拆成三件事：

1. 去中心化：减去 $\mu$，把分布中心移到原点；
2. 缩放：除以标准差，把尺度统一为 1；
3. 仿射：乘以 $\gamma$、加上 $\beta$，提供可学习的输出尺度与偏移。它们是共享参数，不能逐输入恢复被归一化丢弃的均值与尺度；只有固定统计量时，才可以取 $\gamma = \sqrt{\sigma^2+\epsilon}$、$\beta = \mu$ 恢复恒等变换。

公式中的 $\epsilon$ 是一个很小的正数，通常为 $10^{-5}$ 或 $10^{-6}$，放在根号内防止方差为零时出现除零。另外，在 bf16/fp16 训练中，$\mu$、$\sigma^2$ 这类统计量通常会在 fp32 精度下累加，再转回原精度参与逐元素运算，以降低舍入误差。

统计的维度#

那么，各种归一化方法的区别在哪里？对于下面这四种中心化方法，一个关键区别是：

$\mu$ 和 $\sigma^2$ 沿哪些维度统计。

设输入张量为 $X \in \mathbb{R}^{N \times L \times C}$，其中 $N$ 是 batch 大小，$L$ 是序列长度，$C$ 是特征维度。几种经典方法的区别可以列成一张表：

画成图会更直观；同样一个 $(N, L, C)$ 张量，四种方法覆盖的区域各不相同：

可以看到，它们的公式完全一致，只是统计量的“覆盖范围”不同。接下来我们分别看看这些方法在 LLM 场景下的取舍。

BatchNorm 的局限性#

将 BatchNorm 搬到 LLM 上，会遇到几个绕不开的问题：

1. 统计量被 padding 污染。语言模型的序列长度不一，batch 内需要 padding 对齐，而 padding 位置的值会进入 $\mu_c$、$\sigma_c^2$ 的统计，使统计量随 batch 的 padding 比例漂移；
2. 小 batch 下统计噪声大。统计量依赖整个 batch，batch 越小噪声越大，而大模型训练的 per-device batch 往往很小；
3. 训练与推理不一致。推理依赖 running mean/var，但语言模型面对的输入长度、领域分布差异极大，固定的 running 统计量很难代表推理时的真实分布；
4. 逐 token 解码难以沿用训练时的统计方式。单请求每步只有一个新位置，直接重算 batch 统计量会很不稳定；标准 BN 在推理时可以使用 running 统计量，并非无法执行；
5. 跨卡同步有额外开销。普通 BN 可以使用各卡局部统计量；若使用 SyncBatchNorm 聚合跨卡统计量，则需额外通信。

代码实现#

先实现一个 BatchNorm 教学版本，对 [N, L, C] 输入沿 (N, L) 统计，并复现 running mean/var 的更新逻辑：

1
class BatchNorm(nn.Module):
2
    def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1) -> None:
3
        super().__init__()
4
        self.eps = eps
5
        self.momentum = momentum
6
        self.weight = nn.Parameter(torch.ones(num_features))
7
        self.bias = nn.Parameter(torch.zeros(num_features))
8
        self.register_buffer("running_mean", torch.zeros(num_features))
9
        self.register_buffer("running_var", torch.ones(num_features))
10

11
    def forward(self, x: torch.Tensor) -> torch.Tensor:
12
        # x: [N, L, C]
13
        if self.training:
14
            mean = x.mean(dim=(0, 1))                    # [C]
15
            var = x.var(dim=(0, 1), unbiased=False)      # [C]
16
            with torch.no_grad():
17
                self.running_mean.mul_(1 - self.momentum).add_(self.momentum * mean)
18
                # PyTorch 惯例：running_var 用无偏估计
19
                self.running_var.mul_(1 - self.momentum).add_(
20
                    self.momentum * x.var(dim=(0, 1), unbiased=True)
21
                )
22
        else:
23
            mean, var = self.running_mean, self.running_var
24
        return (x - mean) / torch.sqrt(var + self.eps) * self.weight + self.bias

注意两处方差的口径不一样。归一化当前 batch 用有偏方差，分母是 $NL$，这样归一化后的方差恰好约为 1；如果错用无偏估计，分母变成 $NL - 1$，归一化后方差就只剩 $\frac{NL - 1}{NL}$。更新 running_var 时则反过来用无偏方差：它是推理阶段对总体方差的长期估计，样本量小时有偏估计会系统性偏低，这也是 PyTorch BatchNorm 的做法。另外 $NL = 1$ 时无偏估计会除零，得到 NaN。

而 padding 污染的问题可以直接观察到：给序列补零之后，统计量会明显偏移。

1
x = torch.randn(4, 16, 8) + 2      # [N, L, C]
2
x_padded = torch.cat([x, torch.zeros(4, 16, 8)], dim=1)  # [N, 2L, C]，补长一倍
3

4
x.mean(dim=(0, 1))          # ≈ 2
5
x_padded.mean(dim=(0, 1))   # ≈ 1，均值被拉走，方差也从 ≈1 变成 ≈1.5

batch 内容一变，统计量就变。归根结底，BatchNorm 的统计依赖于 batch 与序列维度，而语言模型的 batch 构成和序列长度都在不断变化；Transformer 需要一种只依赖单个 token 自身特征的归一化方法，这正是 LayerNorm。

性质#

LayerNorm 有几个值得注意的数学性质：

• 平移不变性：$\text{LN}(\mathbf{x} + c\mathbf{1}) = \text{LN}(\mathbf{x})$，所有分量同时加上常数，归一化结果不变；
• 缩放不变性：忽略 $\epsilon$ 且方差非零时，对 $\alpha > 0$ 有 $\text{LN}(\alpha\mathbf{x}) = \text{LN}(\mathbf{x})$；固定的正 $\epsilon$ 会破坏严格等式，仅当其相对方差很小时近似成立。

这意味着 LayerNorm 丢失了输入的整体平移信息，并在常见数值范围内弱化整体尺度差异。可学习的 $\gamma$、$\beta$ 可以调整输出，却不能恢复每个输入原有的均值与尺度。从反向传播的角度看，由于均值与方差同样参与计算，LN 的梯度中会出现“减去均值方向的梯度分量”这一项，这使得它同时约束了梯度的尺度，是其稳定训练的原因之一。

代价方面，LayerNorm 需要先统计均值和方差，再做逐元素归一化，属于典型的访存受限（memory-bound）算子：计算量不大，但要完整读写几遍激活值。在 LLM 中归一化层极多，这部分开销累积起来也不可忽视；这也是后来 RMSNorm 出现的一个动机。

代码实现#

1
class LayerNorm(nn.Module):
2
    def __init__(self, d_model: int, eps: float = 1e-5) -> None:
3
        super().__init__()
4
        self.eps = eps
5
        self.weight = nn.Parameter(torch.ones(d_model))
6
        self.bias = nn.Parameter(torch.zeros(d_model))
7

8
    def forward(self, x: torch.Tensor) -> torch.Tensor:
9
        # x: [N, L, C]
10
        mean = x.mean(dim=-1, keepdim=True)                       # [N, L, 1]
11
        var = x.var(dim=-1, unbiased=False, keepdim=True)         # [N, L, 1]
12
        return (x - mean) / torch.sqrt(var + self.eps) * self.weight + self.bias

与 nn.LayerNorm 数值对拍，使用默认参数与零初始化 bias，输出应在浮点误差内一致：

1
x = torch.randn(4, 16, 768)  # [N, L, C]
2
assert torch.allclose(LayerNorm(768)(x), nn.LayerNorm(768)(x), atol=1e-6)

工程细节#

实现层面有几个值得注意的点：

• 统计量在 fp32 下计算。bf16 激活先 upcast 到 fp32 求均方根，再转回原精度做逐元素乘法，避免累加误差；
• 用 rsqrt 表达倒数平方根。这有利于内核优化；具体速度与误差可以通过基准测试比较；
• fused kernel。Liger Kernel、FlashInfer 等提供将统计、缩放与乘权重融合的实现；nn.RMSNorm 是否走融合内核取决于 PyTorch 版本、设备和编译路径；
• Gemma 的小差别。Gemma 系列的 RMSNorm 权重形式是 $(1 + w)$ 而不是 $w$，参数初始化为 0 而不是 1。数学上等价，但读代码时容易困惑，值得留意。

代码实现#

1
class RMSNorm(nn.Module):
2
    def __init__(self, d_model: int, eps: float = 1e-6) -> None:
3
        super().__init__()
4
        self.eps = eps
5
        self.weight = nn.Parameter(torch.ones(d_model))
6

7
    def forward(self, x: torch.Tensor) -> torch.Tensor:
8
        # x: [N, L, C]
9
        dtype = x.dtype
10
        x = x.float()                                        # 统计量用 fp32
11
        x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)  # [N,L,C] × [N,L,1]
12
        return self.weight * x.to(dtype)                     # 转回原精度再乘权重，沿用 LLaMA 顺序

与 nn.RMSNorm 对拍时注意把 eps 显式对齐，两者的默认值不同：

1
x = torch.randn(4, 16, 768)  # [N, L, C]
2
assert torch.allclose(
3
    RMSNorm(768, eps=1e-6)(x),
4
    nn.RMSNorm(768, eps=1e-6)(x),
5
    atol=1e-6,
6
)

至于“RMSNorm 到底快多少”，可以直接写个基准对比：同样 [T, d] 的输入，分别跑 nn.LayerNorm 与 nn.RMSNorm 各千次取平均。由于两者都是访存受限算子，实际收益主要来自少维护一组 $\beta$ 参数、少一次均值规约，具体加速比例值得用实测数字说话。GPU 基准还应先预热，并同步设备或使用设备计时器。

Post-Norm#

原始 Transformer 采用的是 Post-Norm，即先残差相加、再归一化：

Post-Norm 的每一层输出都被强制拉回固定尺度，表达能力较强；但代价是残差路径上每一层都要穿过一次归一化。On Layer Normalization in the Transformer Architecture 的分析表明，Post-Norm 在训练初期靠近输出层的梯度会异常放大，因此深层时通常需要配合 warmup 才能稳定收敛。

Pre-Norm#

现代 LLM 几乎全部使用 Pre-Norm，即先归一化、再进入子层、最后残差相加：

这样残差流本身就是一条没有归一化打断的“高速公路”，梯度可以直达每一层，训练非常稳定，对 warmup 也不敏感。代价则来自前面提到的方差累加：残差流的尺度随层数不断增长，深层子层的输出相对残差流越来越小，相当于深层模块的贡献被稀释。此外还有一个容易遗漏的细节：因为残差流上不再有归一化，所以模型最后通常还需要一个 final norm，在 lm_head 之前把输出拉回固定尺度。

两种接线方式的对比一目了然：

折中方案#

Post-Norm 与 Pre-Norm 各有代价，因此也出现了一些折中设计：

• DeepNorm：DeepNet 在 Post-Norm 的残差通路上给 $\mathbf{x}$ 乘一个常数 $\alpha$，即 $\mathbf{x} \leftarrow \text{LN}(\alpha\mathbf{x} + F(\mathbf{x}))$，并配套缩小初始化幅度，让 Post-Norm 结构也能稳定堆叠到上千层；
• Sandwich Norm：在子层前后各放一个归一化，同时约束输入与输出幅度，CogView 曾采用；Gemma 2 中每个子层前后各一个 RMSNorm 的设计也是同样的思路。

代码实现#

两种接线方式的差别只在 forward 里一行代码的顺序：

1
class PostNormBlock(nn.Module):
2
    def __init__(self, d_model: int, sublayer: nn.Module) -> None:
3
        super().__init__()
4
        self.sublayer = sublayer
5
        self.norm = nn.RMSNorm(d_model)
6

7
    def forward(self, x: torch.Tensor) -> torch.Tensor:
8
        # x: [N, L, C]
9
        return self.norm(x + self.sublayer(x))   # 残差之后归一化
10

11

12
class PreNormBlock(nn.Module):
13
    def __init__(self, d_model: int, sublayer: nn.Module) -> None:
14
        super().__init__()
15
        self.sublayer = sublayer
16
        self.norm = nn.RMSNorm(d_model)
17

18
    def forward(self, x: torch.Tensor) -> torch.Tensor:
19
        # x: [N, L, C]
20
        return x + self.sublayer(self.norm(x))   # 归一化后进子层

如果想直观感受两者的差异，可以堆一个几十层的 mini Transformer，固定学习率、去掉 warmup 分别训练：观察 Post-Norm 更容易出现的 loss 震荡，以及 Pre-Norm 更平稳的梯度；这个实验比单看公式更直观。

代码实现#

在 attention 中插入 qk norm 只需在投影之后、RoPE 之前加两行：

1
self.q_norm = RMSNorm(head_dim)
2
self.k_norm = RMSNorm(head_dim)
3

4
def forward(self, x):
5
    # x: [B, L, d_model]
6
    q = self.q_norm(self.q_proj(x).view(B, L, H, D))  # [B,L,H·D] -> [B,L,H,D]
7
    k = self.k_norm(self.k_proj(x).view(B, L, H, D))  # [B,L,H,D]
8
    q = apply_rope(q, cos, sin)   # norm 在 RoPE 之前
9
    k = apply_rope(k, cos, sin)   # apply_rope 见 EP.3
10
    ...