人工智能

归一化

Properties1
aliasesNormalization, 批归一化, BatchNorm, 层归一化, LayerNorm, RMSNorm

6 min read

直觉理解

深层网络中，每一层的输入都是前一层的输出。参数一更新，后面各层看到的输入分布就跟着变，数值可能越来越大或越来越小，训练变得不稳定，对学习率和初始化也格外敏感。

BatchNorm 和 LayerNorm 先在指定维度上做中心化与缩放，再应用可学习的仿射变换；RMSNorm 则不减均值。它们改善数值尺度和优化条件，但不能保证仿射变换后的输出仍是零均值、单位方差，也不保证每层输入始终落在固定范围内。

几种归一化的区别，在于在哪些维度上统计均值和方差。

定义与公式

对一组数值 ，BatchNorm 和 LayerNorm 的标准化及仿射变换形式是

其中 、 是这组数的均值和方差， 是防止除零的小常数，（缩放）和 （偏移）是可学习参数。

批归一化（BatchNorm）

对每个特征（通道），在批次及相应空间维度上统计均值和方差。例如形状为 的图像通常按每个通道在 上统计。由 Ioffe 和 Szegedy 在 2015 年提出。

  • 训练时用当前批次的统计量，同时维护一个滑动平均；推理时改用滑动平均，所以训练和推理的行为不同。
  • 效果依赖批次大小，批次太小时统计量噪声很大。
  • 对变长序列不友好：不同样本长度不同，填充位置会干扰统计。所以它在计算机视觉里很常用，在 Transformer 里基本不用。

层归一化（LayerNorm）

对每个样本内部指定的特征维度统计均值和方差，与其他样本无关。Transformer 通常对每个 token 的最后一个隐藏维度归一化；不必把整个样本的所有 token 合在一起统计。由 Ba 等人在 2016 年提出。

  • 训练和推理行为一致，不受批次大小影响。
  • 对序列模型天然适用：每个 token 的向量各自归一化。
  • 是 Transformer 的标准配置。

RMSNorm

层归一化的简化版，去掉了减均值这一步，只用均方根（root mean square）做缩放，也没有偏移参数：

由 Zhang 和 Sennrich 在 2019 年提出。计算更少，效果与层归一化相当，LLaMA 等大语言模型普遍采用。

归一化放在哪里：Post-LN 与 Pre-LN

在 Transformer 里，归一化与残差连接的相对位置很关键：

  • Post-LN（原始 Transformer）：，先加残差再归一化。
  • Pre-LN（GPT-2 起的主流做法）：，先归一化再进子层，残差通路保持「干净」。

相关论文在其分析和实验设置中发现，Pre-LN 的初始化梯度更稳定，某些设置可以免去学习率预热。不能据此认为所有 Pre-LN 大模型都不需要 warmup；实际配方仍常使用预热，且也存在其他归一化布局。详见 Transformer。

代码示例

import torch
import torch.nn as nn
 
class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))
 
    def forward(self, x):
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x / rms * self.weight
 
x = torch.randn(2, 4, 8)                  # (batch, 序列长度, 特征维度)
 
# 手写层归一化，与 nn.LayerNorm 对照
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True, unbiased=False)
manual = (x - mean) / torch.sqrt(var + 1e-5)
print(torch.allclose(manual, nn.LayerNorm(8, elementwise_affine=False)(x), atol=1e-6))  # True
 
print(RMSNorm(8)(x).shape)                # torch.Size([2, 4, 8])
 
# 批归一化在 batch 维上统计，输入形状是 (batch, 通道)
print(nn.BatchNorm1d(8)(torch.randn(16, 8)).shape)   # torch.Size([16, 8])

在 AI 中的作用

归一化是让深层网络能训练起来的关键部件之一，与残差连接配合使用。视觉模型（如 ResNet）以批归一化为主；Transformer 与大语言模型使用层归一化或 RMSNorm，并采用 Pre-LN 结构。在大模型训练中，归一化的位置和数值精度还与训练稳定性（损失尖峰）密切相关。

参考资料

相关

Graph View

Backlinks