LLM 常见归一化方法解析
在大语言模型（LLM）的架构中，归一化（Normalization）技术对于模型的训练稳定性、收敛速度以及最终性能起着至关重要的作用。Transformer 及其变体广泛采用了不同的归一化策略。本文将详细解析 LayerNorm、RMSNorm 和 DeepNorm 的原理、代码实现及适用场景，并对比 PreLN 与 PostLN 在 Transformer 中的位置差异。
1. Layer Norm 的计算公式与实现
Layer Normalization (LayerNorm) 是 Transformer 架构中最基础的归一化方法之一。它对单个样本的特征维度进行归一化，使其均值为 0，方差为 1，然后通过可学习的缩放参数 γ 和平移参数 β 进行调整。
数学原理
给定输入特征 $x$，LayerNorm 的计算过程如下：
$$\mu = \frac{1}{N} \sum_{i=1}^{N} x_i$$ $$\sigma^2 = \frac{1}{N} \sum_{i=1}^{N} (x_i - \mu)^2$$ $$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$$ $$y_i = \gamma \cdot \hat{x}_i + \beta$$
其中：
- $\mu$ 为 $x$ 的均值。
- $\sigma$ 为 $x$ 的标准差。
- $\gamma$ 和 $\beta$ 是可训练的模型参数，分别控制新分布的方差和均值。
- $\epsilon$ 是一个极小值（如 $1e-6$），添加到方差上以避免分母为 0。
PyTorch 代码实现
import torch
import torch.nn as nn
def layer_norm(feature):
# 获取特征形状
size = feature.shape
# 初始化可学习参数 gamma 和 beta
alpha = nn.Parameter(torch.ones(size[-1]))
beta = nn.Parameter(torch.ones(size[-1]))
input_dtype = feature.dtype
# 转换为 float32 以稳定计算
feature = feature.to(torch.float32)
# 计算均值和标准差
mean = feature.mean(-1, keepdim=True)
std = feature.std(-1, keepdim=True, unbiased=False)
# 归一化并应用缩放平移
normalized = (feature - mean) / (std + 1e-6)
output = alpha * normalized + beta
return output.to(input_dtype)
2. RMS Norm 的计算公式与实现
RMSNorm (Root Mean Square Layer Normalization) 是 LayerNorm 的一种简化变体，主要应用于大语言模型（如 LLaMA 系列）。它去除了均值减去的步骤，仅保留缩放部分。
数学原理
RMSNorm 的核心思想是假设数据的均值已经接近 0，或者减去均值带来的收益不如直接对 RMS 进行归一化显著。其公式如下：
