搞清楚 BatchNorm、LayerNorm、RMSNorm 到底在干嘛

训练神经网络的时候，Normalization 几乎是标配。但 BatchNorm、LayerNorm、RMSNorm 这些东西，我之前一直没彻底搞明白它们的区别。公式都差不多，都是减均值除标准差，那到底差在哪？

这篇把它们放一起对比一下，顺便理一理各自适用的场景。

先从最基础的问题开始：为什么要 Normalize？

神经网络训练时有个烦人的问题：每一层的输入分布会随着前面层参数的更新而变化。这个现象叫 Internal Covariate Shift（ICS）。

直觉上理解：你在训练第 5 层的时候，第 4 层的输出分布变了，那第 5 层学到的东西可能就不太对了。就像你在练习投篮，但篮筐一直在移动。

但这个解释其实有点问题。2018 年 MIT 的一篇论文 How Does Batch Normalization Help Optimization? 做了实验，发现 BatchNorm 的效果和 ICS 关系不大。他们甚至故意往 BatchNorm 之后注入噪声来增加 ICS，结果模型训练得依然很好。

那 Normalization 到底在帮什么忙？

更靠谱的解释是：Normalization 让 loss landscape 变得更平滑了。

没有 Normalization 的时候，loss 曲面可能很崎岖，梯度方向变化剧烈，得用很小的学习率才能稳定训练。加了 Normalization 之后，loss 曲面变平滑了，梯度方向更一致，可以用更大的学习率，收敛更快。

从数学上看，Normalization 实际上是在约束每一层输出的尺度。如果没有这个约束，权重的微小变化可能导致输出的巨大变化，梯度就会爆炸。Normalization 相当于给每一层加了一个”稳压器”。

BatchNorm：开创者，但有硬伤

BatchNorm 是 2015 年 Ioffe 和 Szegedy 提出的，算是 Normalization 家族的开山之作。

公式推导

假设一个 mini-batch 包含 $m$ 个样本 ${x_1, x_2, …, x_m}$，对于某一个特征维度（或者 CNN 里的某个通道），BatchNorm 的计算过程是：

第一步：计算 batch 内的均值 $\mu_B = \frac{1}{m} \sum_{i=1}^{m} x_i$

第二步：计算 batch 内的方差 $\sigma_B^2 = \frac{1}{m} \sum_{i=1}^{m} (x_i - \mu_B)^2$

第三步：归一化 $\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$

这里 $\epsilon$ 是一个很小的数（比如 $10^{-5}$），防止除零。

第四步：缩放和偏移 $y_i = \gamma \hat{x}_i + \beta$

等等，刚归一化完又加 $\gamma$ 和 $\beta$，这不是白干了吗？

不是的。归一化是强制把分布限制在均值 0、方差 1，但这可能不是最优的分布。$\gamma$ 和 $\beta$ 是可学习的参数，让网络自己决定需要什么样的分布。极端情况下，如果网络学到 $\gamma = \sigma_B$，$\beta = \mu_B$，那就等于没有归一化。

训练和推理的不一致

BatchNorm 有个麻烦的地方：训练时用的是当前 batch 的统计量，但推理时 batch size 可能是 1，没法算统计量。

解决方案是用 移动平均。训练过程中维护全局的均值和方差：

\[\mu_{running} = (1 - \alpha) \cdot \mu_{running} + \alpha \cdot \mu_B\] \[\sigma^2_{running} = (1 - \alpha) \cdot \sigma^2_{running} + \alpha \cdot \sigma^2_B\]

推理时就用这个 $\mu_{running}$ 和 $\sigma^2_{running}$。

这带来一个问题：训练和推理的行为不一致。如果你忘了调用 model.eval()，或者训练时的 batch 分布和推理时差异很大，就可能出问题。我之前就踩过这个坑，调了半天才发现是 BatchNorm 的锅。

对 batch size 的依赖

BatchNorm 对 batch size 很敏感。batch size 太小的话，$\mu_B$ 和 $\sigma^2_B$ 的估计就不准，方差可能很大。

一般来说，batch size 至少要 16 以上，BatchNorm 才能稳定工作。但有些任务（比如目标检测、分割）由于显存限制，batch size 经常只有 2 或 4，这时候 BatchNorm 就不太行了。

后来有人提出了 GroupNorm 来解决这个问题，但那是另一个故事了。

BatchNorm 的梯度

这部分稍微硬核一点，但理解梯度对于理解 BatchNorm 的行为很重要。

设 $\hat{x}_i = \frac{x_i - \mu}{\sigma}$，我们来推导 $\frac{\partial L}{\partial x_i}$。

由于 $\mu$ 和 $\sigma$ 都依赖于所有的 $x_i$，所以求导会有点复杂。设 $\frac{\partial L}{\partial \hat{x}_i} = g_i$，那么：

\[\frac{\partial L}{\partial x_i} = \frac{1}{\sigma} \left( g_i - \frac{1}{m}\sum_j g_j - \frac{\hat{x}_i}{m}\sum_j g_j \hat{x}_j \right)\]

这个式子有什么含义？

$g_i$：来自上游的梯度
$-\frac{1}{m}\sum_j g_j$：减去梯度的均值（类似于中心化）
$-\frac{\hat{x}_i}{m}\sum_j g_j \hat{x}_j$：减去和方向相关的分量

后两项可以理解为一种”去相关”操作，让梯度分布更均匀，这也是为什么 BatchNorm 能让训练更稳定。

LayerNorm：Transformer 的标配

BatchNorm 在 CNN 上效果很好，但到了 RNN 和 Transformer 这边就不太行了。原因很简单：序列任务里，不同样本的长度可能不一样，在 batch 维度上取统计量不太合理。

2016 年 Ba 等人提出了 LayerNorm，思路是：不跨样本，只在单个样本的特征维度上做归一化。

公式

对于一个样本的特征向量 $\mathbf{x} = (x_1, x_2, …, x_D)$：

\[\mu = \frac{1}{D} \sum_{i=1}^{D} x_i\] \[\sigma^2 = \frac{1}{D} \sum_{i=1}^{D} (x_i - \mu)^2\] \[\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}\] \[y_i = \gamma_i \hat{x}_i + \beta_i\]

注意这里 $\gamma$ 和 $\beta$ 是对每个特征位置单独设置的，形状是 $(D,)$。

和 BatchNorm 的对比

用一张图来说明会更清楚。假设输入形状是 $(B, T, D)$：

输入张量: (B, T, D)
         B = batch size
         T = 序列长度
         D = 特征维度（hidden size）

BatchNorm: 对每个特征位置 d ∈ [1,D]
           计算所有 (b,t) 位置的均值和方差
           → 归一化维度: (B, T)
           → 每个 d 共享一组统计量

LayerNorm: 对每个 (b,t) 位置
           计算所有 d ∈ [1,D] 的均值和方差
           → 归一化维度: (D,)
           → 每个 (b,t) 位置独立计算

关键区别：

BatchNorm 让同一个特征在不同样本之间对齐
LayerNorm 让同一个样本的不同特征之间对齐

为什么 Transformer 用 LayerNorm？

几个原因：

不依赖 batch：每个样本独立计算，batch size 是 1 也能正常工作
处理变长序列：不同样本长度不一样也没问题
训练推理一致：不需要维护移动平均，train() 和 eval() 行为相同

但 LayerNorm 也有缺点。它假设同一层的不同特征之间有可比性，这在某些情况下可能不成立。比如，如果特征的第 1 维表示”年龄”，第 2 维表示”收入”，把它们一起归一化就有点奇怪。

不过在 Transformer 里，hidden state 的每一维没有明确的语义，所以这个问题不大。

LayerNorm 的位置：Pre-Norm vs Post-Norm

原始 Transformer（Vaswani 2017）用的是 Post-Norm：

# Post-Norm
x = x + Attention(x)
x = LayerNorm(x)
x = x + FFN(x)
x = LayerNorm(x)

后来发现 Pre-Norm 训练更稳定，尤其是在深层网络里：

# Pre-Norm
x = x + Attention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

为什么 Pre-Norm 更稳定？

直觉上，Pre-Norm 相当于给残差路径加了一个”开关”。在训练初期，Attention 和 FFN 的输出可能很不稳定，但由于有 LayerNorm 在前面”压住”，输出的尺度不会太离谱。而 Post-Norm 是先加了残差再归一化，如果 Attention 输出炸了，加上残差之后再归一化可能已经来不及了。

从梯度角度看，Pre-Norm 的梯度可以直接通过残差连接流回去，不会被 LayerNorm 截断。而 Post-Norm 的梯度必须经过 LayerNorm，可能会被”调制”。

现在大部分 LLM（GPT、LLaMA、Qwen 等）都用 Pre-Norm。

RMSNorm：去掉均值，更快更简单

2019 年 Zhang 和 Sennrich 提出了 RMSNorm（Root Mean Square Normalization），核心思想是：LayerNorm 里减均值这一步，真的有必要吗？

公式

RMSNorm 只用 RMS（均方根）来归一化：

\[\text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{D} \sum_{i=1}^{D} x_i^2}\] \[\hat{x}_i = \frac{x_i}{\text{RMS}(\mathbf{x})}\] \[y\_i = \gamma_i \hat{x}_i\]

注意这里没有 $\beta$（偏移项）。原论文里也没加，因为实验发现加不加效果差不多。

和 LayerNorm 的对比

把两者拆开看：

LayerNorm 做了两件事：

减均值（中心化）：$x - \mu$
除标准差（缩放）：$\frac{1}{\sigma}$

RMSNorm 只做一件事：

除 RMS（缩放）：$\frac{1}{\text{RMS}}$

数学上，LayerNorm 和 RMSNorm 的关系是：

\[\text{RMS}^2 = \sigma^2 + \mu^2\]

这是因为：

\[\text{RMS}^2 = \frac{1}{D}\sum x\_i^2 = \frac{1}{D}\sum (x\_i - \mu + \mu)^2 = \sigma^2 + \mu^2\]

如果 $\mu \approx 0$，那么 $\text{RMS} \approx \sigma$，两者就差不多了。

为什么去掉均值也可以？

这是个好问题。我的理解是：

Pre-Norm 架构下，均值的影响不大。残差连接本身会累积偏移，LayerNorm 减均值只是暂时”清零”了一下，但过了残差连接偏移又回来了。
可学习的 $\gamma$ 可以补偿。如果网络真的需要某种偏移，$\gamma$ 可以学到。
实验说话。论文里做了大量实验，发现 RMSNorm 和 LayerNorm 效果几乎一样。

计算效率

RMSNorm 省掉了：

计算均值 $\mu$
减均值 $x - \mu$
偏移项 $\beta$ 的加法和存储

这些节省对于 LLM 这种规模的模型来说很可观。论文报告说大概能减少 7%-15% 的 Normalization 计算时间。

考虑到 Transformer 里 Norm 层的数量（每个 layer 有两个），这个优化是值得的。

实现对比

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(dim))
        self.beta = nn.Parameter(torch.zeros(dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * x_norm + self.beta


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        x_norm = x / rms
        return self.gamma * x_norm

RMSNorm 的代码更短，计算更少。

其他变体：GroupNorm、InstanceNorm

既然都整理了，顺便把其他几种 Norm 也说一下。

InstanceNorm

InstanceNorm 最早用在风格迁移任务里。它对每个样本的每个通道单独做归一化。

对于形状 $(B, C, H, W)$ 的输入，InstanceNorm 在 $(H, W)$ 维度上计算统计量：

\[\mu\_{b,c} = \frac{1}{HW} \sum\_{h,w} x\_{b,c,h,w}\]

可以理解为：每张图的每个通道，单独归一化。

用途：风格迁移、图像生成。在这些任务里，每张图的风格信息很重要，不希望被 batch 里其他图片”污染”。

GroupNorm

GroupNorm 是 2018 年 He 等人提出的，介于 LayerNorm 和 InstanceNorm 之间。

它把 $C$ 个通道分成 $G$ 个组，每组内做归一化：

\[\mu_{b,g} = \frac{1}{(C/G) \cdot HW} \sum_{c \in \text{group } g} \sum_{h,w} x_{b,c,h,w}\]

当 $G = C$ 时，GroupNorm 退化成 InstanceNorm。当 $G = 1$ 时，GroupNorm 接近 LayerNorm。

用途：小 batch size 的 CNN 任务（目标检测、分割等）。在 batch size = 1 或 2 的情况下，GroupNorm 效果比 BatchNorm 好很多。

总结对比

方法	归一化维度	依赖 batch	主要用途
BatchNorm	$(B, H, W)$ 对每个 $C$	是	CNN（batch size >= 16）
LayerNorm	$(D)$ 对每个 $(B, T)$	否	Transformer
RMSNorm	$(D)$ 对每个 $(B, T)$	否	LLM
InstanceNorm	$(H, W)$ 对每个 $(B, C)$	否	风格迁移
GroupNorm	$(C/G, H, W)$ 对每个 $(B, G)$	否	小 batch CNN

深入理解：Normalization 的几何意义

这部分稍微抽象一点，但有助于理解为什么这些 Norm 能 work。

把 Normalization 看成几何变换：

减均值：把数据中心移到原点（平移）
除标准差：把数据缩放到单位球面（缩放）

LayerNorm 是把每个样本的特征向量投影到单位球面上（严格来说是经过缩放后的球面）。

RMSNorm 更直接：就是把向量除以它的（加权）模长。

从这个角度看，Normalization 是在限制每一层输出的”能量”或”幅度”，防止信号在层与层之间传播时不断放大。

实践建议

最后给一些工程上的建议：

CNN（batch size >= 16）：用 BatchNorm，效果经过充分验证
CNN（batch size < 16）：考虑 GroupNorm 或者同步 BatchNorm
Transformer（小模型）：LayerNorm 和 RMSNorm 都可以
LLM（大模型）：优先 RMSNorm，省计算
训练不稳定：检查 Norm 的位置（Pre-Norm 通常更稳定）
推理行为异常：检查是否正确设置了 eval() 模式（BatchNorm 的坑）

关于 $\epsilon$ 的选择：一般 $10^{-5}$ 或 $10^{-6}$ 就够了。如果用 float16/bfloat16 训练，可以稍微调大一点（$10^{-4}$）防止数值问题。

写到这里，这些 Norm 方法的脉络应该比较清楚了：

BatchNorm 解决了深层 CNN 的训练问题，但依赖 batch → LayerNorm 去掉 batch 依赖，适配了 Transformer → RMSNorm 进一步简化，去掉均值计算，大模型更友好。

选择哪种 Norm，主要看你的模型结构和计算约束。没有绝对的好坏，只有适不适合。