深层网络的训练难题#

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

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

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

归一化与残差流#

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

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

1
x = torch.randn(8, 128, 512)
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 层后尺度已明显漂移

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

这个“在残差流上进行尺度校正”的视角，是后面理解 Pre-Norm 与 Post-Norm 差异的基础。

统一形式#

尽管归一化方法种类繁多，但它们都可以写成同一个形式。对输入元素 $x$，先减去均值、再除以标准差，最后做一层可学习的仿射变换：

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

1. 去中心化：减去 $\mu$，把分布中心移到原点；
2. 缩放：除以标准差，把尺度统一为 1；
3. 仿射：乘以 $\gamma$、加上 $\beta$，把被“过度约束”的分布恢复出表达能力——如果恒等变换是最优解，网络可以学到 $\gamma = \sigma$、$\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. 自回归解码没有 batch 统计可言。逐 token 生成时，每一步只有一个新 token，“沿 batch 统计”这个前提本身就不成立；
5. 分布式训练需要跨卡同步统计量，带来额外的通信开销。

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

代码实现#

先实现一个标准的 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))
15
            var = x.var(dim=(0, 1), unbiased=False)
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

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

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

4
x.mean(dim=(0, 1))          # ≈ 0
5
x_padded.mean(dim=(0, 1))   # 被 padding 拉向 0，方差也被压小一半

batch 内容一变，统计量就变——这正是 BatchNorm 不适合语言模型的根本原因。

性质#

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

• 平移不变性：$\text{LN}(\mathbf{x} + c\mathbf{1}) = \text{LN}(\mathbf{x})$，所有分量同时加上常数，归一化结果不变；
• 缩放不变性：$\text{LN}(\alpha\mathbf{x}) = \text{LN}(\mathbf{x})$（$\alpha > 0$），输入整体缩放不改变归一化结果。

这也意味着 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
        mean = x.mean(dim=-1, keepdim=True)
10
        var = x.var(dim=-1, unbiased=False, keepdim=True)
11
        return (x - mean) / torch.sqrt(var + self.eps) * self.weight + self.bias

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

1
x = torch.randn(4, 16, 768)
2
assert torch.allclose(LayerNorm(768)(x), nn.LayerNorm(768)(x), atol=1e-6)

工程细节#

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

• 统计量在 fp32 下计算。bf16 激活先 upcast 到 fp32 求均方根，再转回原精度做逐元素乘法，避免累加误差；
• 用 rsqrt 而不是 sqrt + 除法。rsqrt 在硬件上通常是单条指令，速度更快、数值上也更稳；
• fused kernel。PyTorch 的 nn.RMSNorm，以及 Liger Kernel、FlashInfer 等推理框架中的实现，都是把“统计 + 缩放 + 乘权重”融合成一个 kernel，减少访存次数；
• 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
        dtype = x.dtype
9
        x = x.float()                                        # 统计量用 fp32
10
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
11
        return (x / rms * self.weight).to(dtype)

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

1
x = torch.randn(4, 16, 768)
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$ 参数、少一次均值规约，量级通常在百分之几到十几，远不如 fused kernel 与 fp32 统计带来的差异显著——这一点值得用实测数字说话，而不是想当然。

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
        return self.norm(x + self.sublayer(x))   # 残差之后归一化
9

10

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

17
    def forward(self, x: torch.Tensor) -> torch.Tensor:
18
        return x + self.sublayer(self.norm(x))   # 归一化后进子层

如果想直观感受两者的差异，可以堆一个几十层的 mini Transformer，固定学习率、去掉 warmup 分别训练：Post-Norm 的 loss 通常在几十个 step 内就开始震荡甚至发散，而 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
    q = self.q_norm(self.q_proj(x).view(B, L, H, D))
6
    k = self.k_norm(self.k_proj(x).view(B, L, H, D))
7
    q = apply_rope(q, cos, sin)   # norm 在 RoPE 之前
8
    k = apply_rope(k, cos, sin)   # apply_rope 见 EP.3
9
    ...