2. 归一化层

概述

归一化层是深度学习中不可或缺的组件，用于稳定训练过程、加速收敛。本章将介绍归一化层的演进历程：

BatchNorm (CV) → LayerNorm (NLP) → RMSNorm (LLM, Qwen3.5 采用)

归一化方式	统计维度	适用场景	代表模型
BatchNorm	跨 batch 统计	计算机视觉 (CNN)	ResNet, VGG
LayerNorm	跨特征维度统计	自然语言处理 (Transformer)	BERT, GPT-2
RMSNorm	跨特征维度 (仅 RMS)	大语言模型	Qwen3.5, LLaMA

2.1 为什么需要归一化？

在深度神经网络中，每一层的输入分布会随着前层参数的更新而发生变化，这种现象称为内部协变量偏移（Internal Covariate Shift）。归一化层通过将输入标准化到稳定的分布，解决以下问题：

1. 加速收敛：归一化后的输入落在激活函数的"敏感区"，梯度更大
2. 稳定训练：减少梯度爆炸/消失的风险
3. 降低对初始化的敏感性：即使初始化不理想，归一化也能拉回合理范围

2.2 BatchNorm：批归一化

原理

Batch Normalization（Ioffe & Szegedy, 2015）是最早被广泛使用的归一化技术。核心思想：跨 batch 维度统计每个特征通道的均值和方差，然后归一化。

$${\mu }_c=\frac{1}{B\cdot H\cdot W}\sum _{b,h,w}{x}_{b,c,h,w}$$$${\sigma }_c^2=\frac{1}{B\cdot H\cdot W}\sum _{b,h,w}(x_{b,c,h,w}-{\mu }_c{)}^2$$$$\text{BatchNorm}(x)={\gamma }_c\cdot \frac{x-{\mu }_c}{\sqrt{{\sigma }_c^2+ϵ}}+{\beta }_c$$

其中：

• $B$：batch size
• $c$：特征通道索引
• ${\gamma }_c,{\beta }_c$：可学习的缩放和偏移参数

关键特点：统计量是跨 batch 计算的，同一个 batch 内的所有样本共享均值和方差。

训练 vs 推理的差异

BatchNorm 有一个重要特点：训练和推理时的行为不同。

	训练模式	推理模式
均值/方差	当前 batch 的统计量	全局移动平均（训练时累积）
用途	归一化当前 batch	归一化单个样本
是否需要 batch	✅ 是	❌ 否

为什么需要区分？

• 推理时可能只有单个样本（batch_size=1），无法计算 batch 统计量
• 解决方案：训练时用指数移动平均（EMA）累积全局统计量，推理时使用

# 训练时更新全局统计量
running_mean = (1 - momentum) * running_mean + momentum * batch_mean
running_var  = (1 - momentum) * running_var  + momentum * batch_var

代码实现

代码文件

class BatchNorm1d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.track_running_stats = track_running_stats

        # 可学习参数
        if affine:
            self.weight = nn.Parameter(torch.ones(num_features))   # γ 缩放
            self.bias = nn.Parameter(torch.zeros(num_features))    # β 偏移

        # 全局移动平均（不参与梯度更新）
        if track_running_stats:
            self.register_buffer("running_mean", torch.zeros(num_features))
            self.register_buffer("running_var", torch.ones(num_features))

    def forward(self, x):
        # x: (batch, num_features) 或 (batch, num_features, seq_len)
        if x.dim() == 2:
            reduce_dim = (0,)          # 2D: 跨 batch 统计
        elif x.dim() == 3:
            reduce_dim = (0, 2)        # 3D: 跨 batch 和 seq_len 统计

        if self.training:
            # 训练模式：用当前 batch 的统计量
            mean = x.mean(dim=reduce_dim)
            var = x.var(dim=reduce_dim, unbiased=False)
            # 更新全局移动平均
            if self.track_running_stats:
                with torch.no_grad():
                    self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
                    self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
        else:
            # 推理模式：用全局移动平均
            mean = self.running_mean
            var = self.running_var

        # 归一化
        shape = (1, -1) if x.dim() == 2 else (1, -1, 1)
        mean = mean.view(*shape)
        var = var.view(*shape)
        x_norm = (x - mean) / torch.sqrt(var + self.eps)

        # 仿射变换
        if self.affine:
            weight = self.weight.view(*shape)
            bias = self.bias.view(*shape)
            output = weight * x_norm + bias
        else:
            output = x_norm
        return output

BatchNorm 的统计维度图解

以 3D 输入 (batch=2, channels=3, seq_len=4) 为例：

输入 x:
         channel_0    channel_1    channel_2
batch_0: [a0,a1,a2,a3] [b0,b1,b2,b3] [c0,c1,c2,c3]
batch_1: [d0,d1,d2,d3] [e0,e1,e2,e3] [f0,f1,f2,f3]

BatchNorm 统计方式:
  mean_channel_0 = mean(a0,a1,a2,a3, d0,d1,d2,d3)  ← 跨 batch 和 seq_len
  mean_channel_1 = mean(b0,b1,b2,b3, e0,e1,e2,e3)
  mean_channel_2 = mean(c0,c1,c2,c3, f0,f1,f2,f3)

每个 channel 有独立的 mean 和 var，共 3 组统计量

BatchNorm 在 NLP 中的问题

BatchNorm 在 CV（CNN）中非常成功，但在 NLP（Transformer）中存在严重问题：

问题	原因
变长序列	NLP 序列长度不一，padding 位置的统计量无意义
小 batch	LLM 训练 batch size 小，统计量不稳定
序列依赖	同一序列内 token 有依赖关系，跨 batch 统计不合理
推理不一致	推理时只有单个样本，行为与训练时不同

这就是为什么 Transformer 使用 LayerNorm 而非 BatchNorm。

2.3 LayerNorm：层归一化

原理

Layer Normalization（Ba et al., 2016）对每个样本的最后一维（特征维度）进行归一化：

$$\mu =\frac{1}{D}\sum _{i=1}^D{x}_i$$$${\sigma }^2=\frac{1}{D}\sum _{i=1}^D(x_i-\mu {)}^2$$$$\text{LayerNorm}(x)=\gamma \odot \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$

其中：

• $D$：特征维度（如 hidden_size）
• $\gamma$：可学习的缩放参数（初始化为 1）
• $\beta$：可学习的偏移参数（初始化为 0）
• $ϵ$：防止除零的小常数（如 ${10}^{-5}$）

关键特点：LayerNorm 同时计算了均值和方差。

代码实现

代码文件

import torch
import torch.nn as nn


class LayerNormalization(nn.Module):
    """
    层归一化（Layer Normalization）
    对每个样本的指定维度做归一化，减少内部协变量偏移，加速模型训练
    """

    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
        super().__init__()
        self.normalized_shape = normalized_shape
        self.eps = eps
        self.elementwise_affine = elementwise_affine

        if elementwise_affine:
            self.gamma = nn.Parameter(torch.ones(normalized_shape))   # 缩放
            self.beta = nn.Parameter(torch.zeros(normalized_shape))   # 偏移
        else:
            self.register_parameter('gamma', None)
            self.register_parameter('beta', None)

    def forward(self, x: torch.Tensor):
        # x: (batch_size, ..., normalized_shape)
        # 1. 计算均值
        mu = x.mean(dim=-1, keepdim=True)
        # 2. 计算方差（有偏方差，除以 N 而非 N-1）
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        # 3. 归一化
        x_norm = (x - mu) / torch.sqrt(var + self.eps)
        # 4. 仿射变换
        if self.elementwise_affine:
            output = self.gamma * x_norm + self.beta
        else:
            output = x_norm
        return output

计算过程图解

输入 x: (batch=2, seq_len=3, hidden=4)

样本1-token1: [1.0, 2.0, 3.0, 4.0]
  ├── μ = (1+2+3+4)/4 = 2.5
  ├── σ² = ((1-2.5)² + (2-2.5)² + (3-2.5)² + (4-2.5)²)/4 = 1.25
  ├── 归一化: [(1-2.5)/√1.25, (2-2.5)/√1.25, (3-2.5)/√1.25, (4-2.5)/√1.25]
  │         = [-1.34, -0.45, 0.45, 1.34]
  └── × γ + β (γ=1, β=0 时结果不变)

2.4 RMSNorm：均方根归一化（Qwen3.5 采用）

原理

RMSNorm（Root Mean Square Normalization）是 LayerNorm 的简化版本，由 Zhang & Sennrich 在 2019 年提出。核心改进：去掉均值中心化，只用均方根（RMS）进行归一化。

$$\text{RMS}(x)=\sqrt{\frac{1}{D}\sum _{i=1}^D{x}_i^2+ϵ}$$$$\text{RMSNorm}(x)=\gamma \odot \frac{x}{\text{RMS}(x)}$$

与 LayerNorm 的关键区别：

• LayerNorm：先减均值，再除标准差 → $\frac{x-\mu }{\sigma }$
• RMSNorm：直接除以 RMS → $\frac{x}{\text{RMS}(x)}$
• RMSNorm 没有偏移参数 $\beta$，也没有均值中心化

为什么 RMSNorm 更高效？

1. 省去均值计算：不需要计算 $\mu =\frac{1}{D}\sum x_i$
2. 省去减均值操作：不需要 $x-\mu$
3. 无偏移参数：少一组可学习参数

在大规模模型中，归一化层会被调用数十亿次，这些简化带来的计算节省是显著的。

代码实现

basic 版本

代码文件

class RMSNormalization(nn.Module):
    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
        super().__init__()
        self.normalized_shape = normalized_shape
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if elementwise_affine:
            self.gamma = nn.Parameter(torch.ones(normalized_shape))
        else:
            self.register_parameter("gamma", None)

    def forward(self, x):
        # 计算 RMS
        rms = torch.sqrt((x ** 2).mean(dim=-1, keepdim=True) + self.eps)
        # 归一化
        x_norm = x / rms
        # 仿射变换
        if self.elementwise_affine:
            output = self.gamma * x_norm
        else:
            output = x_norm
        return output

miniqwen 版本（Qwen3.5 风格，零初始化）

代码文件

class RMSNorm(nn.Module):
    """
    Qwen3.5 RMSNorm 个人实现

    Qwen3.5 的 RMSNorm 特点:
    - weight 初始化为 0（不是 1）
    - 公式: output = (1 + weight) * norm(x)
      这样初始时相当于恒等变换
    """

    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.zeros(dim))  # 注意：零初始化！

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        output = self._norm(x.float())               # 归一化时用 float32
        output = output * (1.0 + self.weight.float()) # (1 + weight) * norm(x)
        return output.type_as(x)                       # 转回原始 dtype

Qwen3.5 的零初始化技巧

这是 Qwen3.5 相比标准 RMSNorm 最重要的改进：

标准 RMSNorm：$\text{output}=\gamma \odot \frac{x}{\text{RMS}(x)}$，其中 $\gamma$ 初始化为 1

Qwen3.5 RMSNorm：$\text{output}=(1+w)\odot \frac{x}{\text{RMS}(x)}$，其中 $w$ 初始化为 0

为什么这样做？

个人认为有利于训练稳定：在训练初期，如果梯度较大，gamma变化幅度较大，可能会破坏初始的信号流。而Qwen3.5 的做法更类似于残差连接，带来了更好的训练稳定。

对比：

初始化方式	初始时的输出	效果
标准 $\gamma =1$	$\text{norm}(x)$	可能偏离原始信号
Qwen3.5 $w=0$	$\text{norm}(x)$	初始时更接近恒等变换

数值稳定性

Qwen3.5 的 RMSNorm 还有一个细节：归一化时使用 float32 计算。

output = self._norm(x.float())               # 归一化时用 float32
output = output * (1.0 + self.weight.float()) # 乘法也用 float32
return output.type_as(x)                      # 转回原始 dtype (如 bfloat16)

这是因为：

• 混合精度训练时，模型参数可能是 bfloat16 或 float16
• 平方和、均值等操作在低精度下容易溢出或损失精度
• 先在 float32 下计算，再转回原始 dtype

2.5 LayerNorm vs RMSNorm 对比

特性	LayerNorm	RMSNorm
归一化方式	$\frac{x-\mu }{\sigma }$	$\frac{x}{\text{RMS}(x)}$
需要均值计算	✅ 是	❌ 否
需要方差计算	✅ 是	✅ 是（RMS²）
可学习参数	$\gamma$ (缩放) + $\beta$ (偏移)	只有 $\gamma$ (缩放)
计算开销	较高	较低
代表模型	BERT, GPT-2, T5	Qwen3.5, LLaMA, Gemma

2.6 归一化层在 Transformer 中的位置

在 Transformer 中，归一化层有两种放置方式：

Post-Norm（原始 Transformer）

x → Attention → Add(x, ·) → LayerNorm → FFN → Add(·, ·) → LayerNorm

Pre-Norm（现代 Transformer，Qwen3.5选择）

x → LayerNorm → Attention → Add(x, ·) → LayerNorm → FFN → Add(·, ·)

Pre-Norm 的优势：

• 残差连接直接相加原始输入，梯度流更顺畅
• 训练更稳定，不需要 warmup
• 允许使用更大的学习率

Qwen3.5 使用 Pre-Norm 风格。

class DecoderLayer(nn.Module):
    def __init__(self, config, layer_idx):
        # ...
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, hidden_states, ...):
        # ...
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)  # 先归一化
        hidden_states = self.self_attn(hidden_states, ...)   # 再做注意力
        hidden_states = residual + hidden_states              # 残差连接
        # ...

2.7 总结

三种归一化对比

归一化方式	统计量	公式	可学习参数	适用场景	代表模型
BatchNorm	跨 batch 均值和方差	$\frac{x-{\mu }_{batch}}{{\sigma }_{batch}}$	$\gamma ,\beta$	CV (CNN)	ResNet, VGG
LayerNorm	跨特征 均值和方差	$\frac{x-\mu }{\sigma }$	$\gamma ,\beta$	NLP (Transformer)	BERT, GPT-2
RMSNorm	跨特征 RMS	$\frac{x}{\text{RMS}(x)}$	$\gamma$	LLM	Qwen3.5, LLaMA

为什么 LLM 偏爱 RMSNorm？

原因	说明
计算效率	去掉均值中心化，减少约 25% 计算
无偏移参数	少一组参数，节省显存
训练稳定性	零初始化 + (1+w) 公式，初始时接近恒等变换
数值稳定	float32 精度计算，避免低精度溢出

Qwen3.5 的选择：

• 使用 RMSNorm，$ϵ={10}^{-6}$
• weight 零初始化 + $(1+w)$ 公式 → 初始时接近恒等变换
• float32 精度计算，保证数值稳定性
• Pre-Norm

---

上一章：1. MLP | 下一章：3. 位置编码