Huajishe Notes

1-BatchNorm-LayerNorm-RMSNorm

22 min read

1-BatchNorm-LayerNorm-RMSNorm

1. 归一化的直觉：为什么稳定分布能稳定训练

1.1 内部协变量偏移（Internal Covariate Shift）

训练神经网络时，每一层的输入分布随前一层参数的更新而不断变化，这种现象称为内部协变量偏移（Internal Covariate Shift, ICS）。具体来说：

  • 第 层接收到的输入是第 层输出的非线性变换
  • 第 层参数在每次更新后， 的分布发生变化
  • 第 层需要不断适应这种”飘移”的输入分布，如同在移动的靶子上射击

ICS 的结果是：第 层必须同时学习两个任务——（1）适应输入分布的漂移；（2）完成自身的任务目标。这加重了学习的负担，尤其是深层网络的前几层，漂移逐层累积，导致训练困难。

1.2 归一化的核心思想

归一化通过强制将各层激活的分布在训练过程中保持稳定，打断 ICS 的累积。具体做法是对激活值做零均值、单位方差的标准化：

其中 和 是沿某个维度计算的均值和方差。然后引入可学习的缩放 和平移 参数，恢复模型的表达能力：

1.3 归一化稳定训练的三个机制

  1. 梯度流稳定：标准化后，各层输入的尺度被控制在一个合理范围内（均值 0，方差 1），反向传播时梯度不会因为链式法则中的乘法效应而爆炸或消失
  2. 允许使用较大学习率：由于参数更新的效果不再被输入尺度的巨大差异稀释，学习率可以更大，收敛更快
  3. 隐式正则化效应：均值和方差是沿 batch（BatchNorm）或 group（GroupNorm）计算的，引入了一定程度的数据依赖噪声，这种噪声提供了类似 Dropout 的隐式正则化（但较弱）

2. BatchNorm

2.1 数学公式

设 mini-batch ，BatchNorm 对每个**通道（channel）**独立计算：

对第 个通道（神经元），计算该 batch 内该通道的均值和方差：

其中 是第 个样本第 层第 个通道的激活值。

标准化：

仿射变换（可学习）：

注意：这里的均值/方差是在 batch 维度上算的，即每个通道独立计算； 也是每个通道一个标量。

2.2 训练与推理的差异

训练时：BatchNorm 使用当前 batch 的统计量（均值 和方差 ）进行标准化。

推理时：BatchNorm 使用全局统计量（滑动平均的均值 和方差 ）：

推理时直接用 标准化，不使用 batch 统计量。

关键：若推理时切换回训练模式（如 eval() 误用），模型使用的是当前 batch 的统计量而非全局统计量，导致推理结果不稳定。这是最常见的 BatchNorm 工程错误之一。

2.3 为什么 BatchNorm 能稳定训练

设某层权重 在两个相邻 step 的输入分布分别为 和 。由于归一化，输入分布被强制拉回到标准高斯 。因此：

在前者中， 的更新直接影响 的尺度；在后者中， 更新只影响相对关系（因为 Norm 后的输入方差固定为 1），这使参数尺度的改变不会直接传导到激活方差的大幅变化。

2.4 反向传播时 BatchNorm 的梯度

设损失 对标准化输出 的梯度为 ，则对 的梯度（利用链式法则）为：

这个梯度公式的两个要点：

  • 第一项是全局梯度，与 成正比
  • 第二项是均值修正项，确保梯度在均值方向上正交于第一项，防止更新破坏归一化性质

参数 的梯度为简单的加总：

2.5 Batch Size 很小时为什么不稳定

当 batch size 很小时（如 ），batch 统计量 对真实均值/方差的估计方差增大：

时的均值估计方差是 时的 16 倍。这意味着：

  1. 每个 step 使用的均值/方差差异很大，导致同一输入在不同时刻的归一化结果不同（统计量方差大）
  2. 归一化后激活的尺度不稳定，反向传播梯度波动剧烈
  3. 在极端情况下，某些通道的方差被估计为 0 或无穷大，导致数值崩溃

经验规则：BatchNorm 在 batch size 时较为稳定； 时需要调大 momentum（）； 时建议切换到 LayerNorm 或 GroupNorm。

3. LayerNorm

3.1 数学公式

LayerNorm 对单个样本的所有通道计算均值和方差，摆脱了对 batch 的依赖。

设输入 （ 为隐藏维度），LayerNorm 计算：

标准化：

仿射变换：

注意 LayerNorm 中 是与输入同维度的向量，每个样本独立归一化，不依赖 batch 统计量。

3.2 LayerNorm 在 Transformer 中的作用

Transformer 中 LayerNorm 通常出现在两个位置：

Post-norm（原始 Transformer）：

即在残差加和之后做归一化。Post-norm 使得主残差路径的信号在每层末端被标准化，这相当于在每层引入了一个”测量点”。

Pre-norm（现代 Transformer）：

即在残差分支内部先做归一化，再通过子层。Pre-norm 使得子层的输入被标准化，而不是输出被标准化。

3.3 Pre-norm vs Post-norm 的差异

维度Pre-normPost-norm
梯度流 的梯度可直接传至浅层（残差路径上的恒等映射）每层梯度都必须经过 LayerNorm，梯度流被归一化层”阻断”一部分
训练稳定性更稳定（信号方差在各子层输入处被控制）训练初期信号可能逐层放大，需配合 warmup
表示能力子层输入被标准化，输出方差由残差路径维持子层输出直接归一化，表示能力略强（经验观察）
主流程度现代大模型（GPT、BERT、LLaMA）主流原始 Transformer 使用
注意Pre-norm 在训练初期需要 warmupPost-norm 配合 warmup 难以训练深层

为什么 Pre-norm 是大模型主流

Pre-norm 在残差路径上保留了从输入到输出的恒等信息，使得深层网络的梯度流更加顺畅。这在大模型（12 层以上）和深层 Transformer 中是关键——Post-norm 在 12 层以后信号方差可能失控，需要非常精细的 warmup 才能稳定训练。

3.4 LayerNorm 在 RNN 中的使用

RNN 中 LayerNorm 通常在时间步之后、进入下一个时间步之前应用：

这使得 RNN 的隐状态方差在各时间步之间保持稳定，缓解了梯度消失/爆炸问题（类似于 LSTM 的门控机制，但 LayerNorm 是在激活上做归一化而非通过门控控制）。

4. RMSNorm

4.1 核心思想

LayerNorm 的计算包括两部分：

  1. 均值中心化（去均值）：
  2. 方差归一化（除标准差）：除以

RMSNorm（Root Mean Square Norm）提出：均值中心化对模型性能贡献有限，去掉它可以节省一半的均值计算开销，同时几乎不损失效果。

RMSNorm 的定义为：

即用**均方根（RMS）**替代标准差 ，并完全移除均值减法。仿射变换：

4.2 LayerNorm vs RMSNorm 的差异

操作LayerNormRMSNorm
均值减法有（）无
归一化因子
计算复杂度（需计算均值 + 方差 + 开方）（只需平方和 + 开方），但省去了均值计算
参数量同
效果差异在视觉任务中略优在 LLM 中基本等效果（7B 模型测试）

为什么 RMSNorm 在 LLM 中被接受

语言模型关注的是token 间关系的相对量级，而非绝对位置。均值中心化对 token 表示的相对关系影响不大，但 RMS 的归一化保留了激活的相对尺度信息。因此在 LLM 中，RMSNorm 可以替代 LayerNorm 而几乎不影响 perplexity（困惑度），同时节省约 10%~15% 的归一化计算时间。

4.3 RMSNorm 在大模型中的使用

LLaMA、Mistral 等大语言模型使用 RMSNorm 替代 LayerNorm。原因：

  • 节省归一化计算开销（在大模型的 12~80 层中累积显著）
  • 效果与 LayerNorm 基本等效（在语言建模 perplexity 上差距 < 0.5%）
  • 配合 Pre-norm 使用时，训练稳定性与 LayerNorm 无显著差异

5. 其他归一化方法

5.1 InstanceNorm

InstanceNorm 对每个样本、每个通道的每个像素（in图像）独立计算均值和方差，公式与 BatchNorm 类似，但维度是 中的 ，而非 batch：

用途：风格迁移（Style Transfer）。InstanceNorm 去除的是内容相关的全局统计量，保留了风格相关的统计量——这正是风格迁移所需要的（把”内容”的统计量去掉，保留”风格”的统计量）。

5.2 GroupNorm

GroupNorm 是 BatchNorm 和 LayerNorm 的折中：将通道分成 组，每组计算均值和方差：

GroupNorm 的特点：

  • 不依赖 batch size：均值/方差在 内计算，与 batch 完全无关
  • 当 时退化为 LayerNorm；当 时退化为 InstanceNorm
  • 当 （通常 ）时，在小 batch 场景下性能与 BatchNorm 接近

小 batch 目标检测/分割场景：当 batch size 限制在 时，GroupNorm 是 BatchNorm 的最佳替代。MMDetection 和 Detectron2 默认使用 GroupNorm 作为检测器的归一化方案。

6. 归一化与残差连接的关系

6.1 残差路径上的信息流

ResNet 的残差连接为：

其中 是残差分支。若残差分支的输出方差远大于 ，则 ，残差路径的优势（梯度直接传递）被削弱——信号主要由残差分支承载。

BatchNorm 在残差分支中的位置：在标准 ResNet 中，Conv-BN-ReLU 的顺序是 BN→ReLU。但残差分支中 BN 的作用是控制 的输出方差，使得 和 的方差在同一量级，保证残差路径的有效性。

6.2 Pre-norm 中残差路径的方差控制

Pre-norm 结构：

由于 LayerNorm 在 SubLayer 之前，SubLayer 的输入方差被控制在 区间。若 SubLayer 的权重初始化使得其输出方差接近 1，则 （假设 和 SubLayer 输出独立），残差路径信号方差较为稳定。

若没有 LayerNorm 的标准化，残差路径的方差会在深层网络中逐层累积（），导致信号爆炸或消失。

7. 工程实践

7.1 训练与推理模式切换

BatchNorm 的 model.train() / model.eval()：

  • model.train()：使用当前 batch 的均值/方差；参数（）正常更新
  • model.eval()：使用全局滑动平均的均值/方差；参数不更新，BatchNorm 的统计量固定

常见错误：

  • 训练时忘记 model.train()，导致使用了 eval 模式的统计量（罕见，但会发生）
  • 更常见：推理时忘记 model.eval()，导致使用当前 batch 统计量而非全局统计量，使得推理结果随 batch 不同而波动

解决方案：始终在推理前显式调用 model.eval()；在训练前显式调用 model.train()。建议用 context manager 明确管理模式：

model.train()
for epoch in range(num_epochs):
    for batch in dataloader:
        # train step
 
model.eval()
with torch.no_grad():
    for batch in dataloader:
        # eval step

7.2 小 Batch 场景的处理

场景问题解决方案
Batch size = 1~4BatchNorm 统计量方差极大，不稳定切换到 GroupNorm（）或 LayerNorm
Batch size = 8~16BatchNorm 较稳定但逊于大 batch增大 momentum 至 ，或使用 GroupNorm
Batch size = 32+BatchNorm 正常无需特殊处理

7.3 分布式训练中的 BatchNorm

DataParallel（DP）中的 BatchNorm：DP 将 batch 分到多个 GPU，每个 GPU 计算各自的 batch 统计量，然后通过 all-reduce 同步。但同步方式导致统计量不够稳定（每个 GPU 只看到 1/N 的 batch），在大 batch 训练中反而可能不如 GroupNorm。

DistributedDataParallel（DDP）中的 BatchNorm：DDP 在每个 GPU 上独立计算 BatchNorm 统计量，通过 torch.nn.SyncBatchNorm 强制同步：先在每个 GPU 上计算本地统计量，再通过 all-reduce 求全局均值/方差，最后广播回各 GPU。SyncBatchNorm 在 DDP 训练中是必需的，否则每个 GPU 的 BatchNorm 层会使用不同的统计量，导致不一致的归一化。

经验：若 DDP 中不使用 SyncBatchNorm，训练看似正常但收敛结果会显著变差（尤其 batch size 较大时）。

7.4 混合精度训练中的坑

混合精度（AMP）训练中，BatchNorm 的统计量计算需在 FP32 中进行。原因：

  • BatchNorm 的均值/方差累积是在滑动平均中进行的，若使用 FP16，累计误差会导致 偏离真实均值
  • AMP 的 autocast 默认对 BatchNorm 层使用 FP32 推理；若手动控制 cast，可能导致数值不稳定

PyTorch 的 AMP 自动处理：PyTorch 的自动混合精度（torch.cuda.amp.autocast）会自动将 BatchNorm 的 forward 保持在 FP32；但 backward 时梯度仍可能是 FP16，需显式使用 GradScaler。

7.5 BatchNorm 统计量与 EMA 的使用问题

问题 1：使用预训练模型的统计量

若加载预训练的 BatchNorm 权重（如从 ResNet 预训练模型迁移到新数据集），滑动平均的统计量（）可能不适用于新数据集的输入分布。解决方案：

  • 在新数据集上重新运行统计量计算（model.eval(); model.apply(reset_bn_stats)）
  • 或在新数据集上 fine-tune 时，不 frozen 住 BatchNorm 层（使其更新统计量）

问题 2：EMA 过时

若训练中断后从 checkpoint 恢复，但 EMA 没有正确加载，则推理时使用的统计量来自恢复点而非全局最优点的估计。解决方案：始终检查 model.eval() 时加载的 running_mean 和 running_var 是否与训练末期的值一致。

8. 各种归一化方法的对比表

方法归一化维度Batch 依赖适用场景小 batch 稳定性计算成本
BatchNorm每个通道，独立计算 batch 维度是CNN、视觉分类检测差（batch size < 16 时差）中等（需计算均值方差）
LayerNorm每个样本的所有通道否Transformer、RNN、NLP优（完全独立于 batch）中等（需计算均值方差）
RMSNorm每个样本的所有通道否LLM（替代 LayerNorm）优低（省去均值计算）
GroupNorm每个样本，通道分组否目标检测、小 batch 分割良（ 时接近 BatchNorm）中等
InstanceNorm每个样本、每个通道、每个空间位置否风格迁移、图像生成优低（空间维度计算）
SyncBatchNorm与 BatchNorm 相同，但全局同步是（全局）DDP 分布式训练同 BatchNorm高（all-reduce 通信）

各方法适用架构速查

架构推荐归一化位置
CNN（ResNet、VGG）BatchNorm（训练）+ SyncBatchNorm（DDP）Conv-BN-ReLU 顺序
Vision Transformer（ViT）LayerNorm 或 RMSNormPre-norm 或 Post-norm
NLP Transformer（BERT、GPT）LayerNorm（Pre-norm）或 RMSNorm（Pre-norm）Pre-norm（主流）
RNN / LSTMLayerNorm（时间步级别）在循环连接之后
目标检测（YOLO、FCOS）GroupNorm（batch size 小时）替换 BatchNorm
风格迁移（StyleGAN）InstanceNorm每层独立
大语言模型（7B+ LLaMA、Mistral）RMSNormPre-norm + RMSNorm（无均值中心化）

9. 实战选择建议

9.1 决策树

架构类型
├── CNN（图像分类、检测、分割）
│   ├── Batch size ≥ 16 → BatchNorm（标准配置）
│   └── Batch size < 16 → GroupNorm（$G=32$）或 LayerNorm
├── Transformer（ViT、NLP）
│   ├── 大模型（> 1B）→ RMSNorm + Pre-norm（效率+稳定性）
│   └── 小模型 / 标准 → LayerNorm + Pre-norm
├── RNN / LSTM
│   └── LayerNorm（在循环单元内，LayerNorm 的位置影响梯度流）
└── 风格迁移 / 图像生成
    └── InstanceNorm（去除内容统计量）

9.2 工程经验总结

  1. BatchNorm 最怕小 batch：当 batch size 降至 8 以下时，GroupNorm 是最直接有效的替代，无需改变网络结构
  2. Transformer 用 LayerNorm 是标准：BERT/GPT 时期 LayerNorm 是主流；LLaMA/Mistral 换成 RMSNorm 是工程优化，不影响性能
  3. Pre-norm vs Post-norm 中：深层 Transformer（> 6 层）强烈建议 Pre-norm；6 层以下差异不明显
  4. 分布式训练必须用 SyncBatchNorm：DDP 中普通 BatchNorm 会导致不同 GPU 使用不同统计量
  5. BatchNorm 的 running stats 不可忽视：加载预训练模型时，检查 running_mean / running_var 是否与训练终止时一致

Graph View