深度学习的归一化层

一、常用 Norm 介绍

LayerNorm

LayerNorm 是一种神经网络中的归一化技术，用于稳定训练过程（缓解梯度消失/爆炸），加速收敛。它的核心思想是 对单个样本的某一层所有神经元的输出进行归一化

① 计算方式

在 Transformer 或 GPT 这类模型中，输入张量的形状通常是 [batch_size, seq_len, hidden_dim]，而 LayerNorm 会沿着 hidden_dim 维度进行归一化。

对于形状为 [batch_size, seq_len, hidden_dim] 的输入，LayerNorm 的归一化是 对最后一个维度（hidden_dim）独立进行的，即：

  • 每个 token 的特征向量（长度为 hidden_dim）单独归一化。
  • 均值和方差的计算：对每个样本的每个时间步（即 [hidden_dim] 维度）计算均值和方差。

具体计算步骤如下：

给定输入张量 x（形状 [batch_size, seq_len, hidden_dim]）：

  1. 计算均值和方差：
    • 对每个样本的每个时间步（即 hidden_dim 维度）计算：
      [
      \mu = \frac{1}{d} \sum_{i=1}^d x_i, \quad \sigma^2 = \frac{1}{d} \sum_{i=1}^d (x_i - \mu)^2
      ]
      （其中 ( d = \text{hidden_dim} )）
  2. 归一化：
    [
    \hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}
    ]
  3. 缩放和平移：
    [
    y = \gamma \hat{x} + \beta
    ]
    （gamma 和 beta 是可学习的参数，形状为 [hidden_dim]）

② 具体例子

假设：

  • batch_size = 2
  • seq_len = 3
  • hidden_dim = 4
  • 输入数据：
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    import torch
    x = torch.tensor([
        # 样本 1 (3 tokens, 每个 token 4 维)
        [[1.0, 2.0, 3.0, 4.0],
         [2.0, 3.0, 4.0, 5.0],
         [3.0, 4.0, 5.0, 6.0]],
        
        # 样本 2 (3 tokens, 每个 token 4 维)
        [[-1.0, -2.0, -3.0, -4.0],
         [-2.0, -3.0, -4.0, -5.0],
         [-3.0, -4.0, -5.0, -6.0]]
    ])  # 形状 [2, 3, 4]

I. 手动计算 LayerNorm

以第一个样本的第一个 token [1.0, 2.0, 3.0, 4.0] 为例：

  1. 计算均值：
    [
    \mu = \frac{1 + 2 + 3 + 4}{4} = 2.5
    ]
  2. 计算方差：
    [
    \sigma^2 = \frac{(1-2.5)^2 + (2-2.5)^2 + (3-2.5)^2 + (4-2.5)^2}{4} = 1.25
    ]
  3. 归一化：
    [
    \hat{x} = \frac{[1.0, 2.0, 3.0, 4.0] - 2.5}{\sqrt{1.25 + 1e-5}} \approx [-1.3416, -0.4472, 0.4472, 1.3416]
    ]
  4. 缩放和平移（假设 gamma=1, beta=0）：
    [
    y = 1.0 \times \hat{x} + 0.0 = \hat{x}
    ]

II. PyTorch 验证

1
2
3
layernorm = torch.nn.LayerNorm(4)  # 对 hidden_dim 归一化
output = layernorm(x)
print(output[0, 0])  # 第一个样本的第一个 token

输出：

1
tensor([-1.3416, -0.4472,  0.4472,  1.3416], grad_fn=<SelectBackward>)

与手动计算一致！

③ PyTorch实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, hidden_dim, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(hidden_dim))  # 缩放参数
        self.beta = nn.Parameter(torch.zeros(hidden_dim))   # 平移参数
        self.eps = eps

    def forward(self, x):
        mu = x.mean(-1, keepdim=True)                     # 计算均值
        var = x.var(-1, keepdim=True, unbiased=False)     # 计算方差
        x_hat = (x - mu) / torch.sqrt(var + self.eps)     # 归一化
        return self.gamma * x_hat + self.beta             # 缩放和平移

RMSNorm

近年来，许多大型语言模型（如 LLaMA、GPT-NeoX、PaLM）开始采用 RMSNorm（Root Mean Square Layer Normalization）替代传统的 LayerNorm，主要原因包括 计算效率更高、训练稳定性更好、对初始化敏感度更低

① 计算方式

RMSNorm 是 LayerNorm 的简化版本，由 Zhang & Sennrich (2019) 提出，其核心改进是 移除均值中心化（Mean Subtraction），仅对输入进行缩放（Scale）。

具体计算步骤：

  1. 计算均方根（RMS）：
    [
    \text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}
    ]
    （其中 ( d ) 是 hidden_dim，( \epsilon ) 是小常数防止除零）
  2. 归一化：
    [
    \hat{x}_i = \frac{x_i}{\text{RMS}(\mathbf{x})}
    ]
  3. 缩放（可学习参数）：
    [
    y_i = \gamma_i \hat{x}_i
    ]
    （注意：RMSNorm 没有偏置项 ( \beta )）

② 具体例子

假设我们有一个 单个样本（忽略 batch_size 和 seq_len），其 hidden_dim=4，输入向量为：
[
\mathbf{x} = [1.0, 2.0, 3.0, 4.0]
]

  1. 计算均方根（RMS）
    [
    \text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{4} (1.0^2 + 2.0^2 + 3.0^2 + 4.0^2)} = \sqrt{\frac{1 + 4 + 9 + 16}{4}} = \sqrt{7.5} \approx 2.7386
    ]
    （假设 ( \epsilon = 1e-5 ) 忽略不计）

  2. 归一化
    [
    \hat{x}_i = \frac{x_i}{\text{RMS}(\mathbf{x})} \implies \hat{\mathbf{x}} \approx \left[ \frac{1.0}{2.7386}, \frac{2.0}{2.7386}, \frac{3.0}{2.7386}, \frac{4.0}{2.7386} \right] \approx [0.3651, 0.7303, 1.0954, 1.4606]
    ]

  3. 缩放（假设可学习参数 ( \gamma = [1, 1, 1, 1] )）
    [
    \mathbf{y} = \gamma \cdot \hat{\mathbf{x}} \approx [0.3651, 0.7303, 1.0954, 1.4606]
    ]

③ PyTorch实现

1
2
3
4
5
6
7
8
9
10
class RMSNorm(nn.Module):
    def __init__(self, hidden_dim, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(hidden_dim))  # 仅缩放参数
        self.eps = eps

    def forward(self, x):
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps  # 计算 RMS
        x_hat = x / rms                                  # 归一化
        return self.gamma * x_hat                        # 仅缩放

二、Norm 对比

RMSNorm 对比 LayerNorm

关键差异

操作 RMSNorm LayerNorm
均值中心化 ❌ 不减去均值 ✅ 减去均值
分母 均方根（RMS） 标准差（含均值中心化）
可学习参数 仅缩放（gamma） 缩放 + 平移（gamma, beta）

为什么 RMSNorm 更适合 LLM？

  1. 计算效率更高
    • LayerNorm 需要计算均值（Mean）和方差（Variance），涉及两次逐元素操作（减均值、除标准差）。
    • RMSNorm 仅需计算均方根（RMS），减少一次计算（无需减均值），**提速约 10%~20%**（对大规模模型显著）。
  2. 训练稳定性更好
    • LayerNorm 的均值中心化可能导致某些情况下梯度不稳定（尤其是初始化阶段或极端值出现时）。
    • RMSNorm 仅缩放输入，对初始化敏感度更低，实验显示在深层网络中更稳定。
  3. 性能无显著损失
    • 论文实验表明，RMSNorm 在语言模型任务中与 LayerNorm 表现相当，甚至在某些场景下更优（如长序列建模）。
    • 例如，LLaMA-1/2、GPT-NeoX 等模型均采用 RMSNorm，未观察到性能下降。

关键区别总结

特性 LayerNorm RMSNorm
归一化方式 减均值 + 除标准差 仅除均方根（RMS）
可学习参数 gamma 和 beta（缩放+平移） 仅 gamma（缩放）
计算复杂度 较高（需计算均值和方差） 较低（仅计算 RMS）
训练稳定性 对初始化敏感 更稳定
适用场景 传统 Transformer / 需要严格零均值的任务（如某些语音或图像模型） 大规模 LLM（如 LLaMA、GPT-NeoX）