为什么LLM时代都使用RMSNorm和Pre-Norm？

结论：大语言模型的预训练过程中经常会出现不稳定的问题。深度学习方法通常会采用特定的归一化策略来加强神经网络训练过程的稳定性。RMSNorm的训练效果不差且训练更快；Post-Norm很难训练，容易导致梯度消失，而Pre-Norm更稳定。

归一化分类

Batch Normalization (批归一化)

位置：传统 CNN 中卷积层后，LLM 中极少使用。
公式：
$$\mu_B = \frac{1}{B}\sum_{i=1}^B x_i, \quad \sigma_B^2 = \frac{1}{B}\sum_{i=1}^B (x_i - \mu_B)^2$$

$$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}, \quad y_i = \gamma \hat{x}_i + \beta$$

代码：

1 2	batch_norm = nn.BatchNorm1d(dim) output = batch_norm(x)

优点：

加速训练，缓解梯度消失。

缺点：

依赖批大小，小批量时不稳定。
不适用于变长序列（如 NLP 任务）。

Layer Normalization (层归一化)

位置：Transformer 的每个子层（自注意力/前馈网络）的输出后，残差连接前（Post-LN），或残差连接后（Pre-LN）。
公式：
$$\mu = \frac{1}{d}\sum_{i=1}^d x_i, \quad \sigma^2 = \frac{1}{d}\sum_{i=1}^d (x_i - \mu)^2 $$

$$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}, \quad y_i = \gamma \hat{x}_i + \beta $$
代码：

# PyTorch 实现
import torch.nn as nn
layer_norm = nn.LayerNorm(dim)
output = layer_norm(x)

优点：

对序列长度和批大小不敏感，适合变长文本。
避免 Batch Norm 对小批量的依赖问题。

缺点：

对特征维度敏感，可能丢失通道间差异性。

在LayerNorm 的基础上，微软提出的 DeepNorm 在残差连接中对之前的激活值 𝒙 按照一定比例 𝛼 进行放缩。Transformer 的层数可以被成功地扩展至 1000 层，有效提升了模型性能与训练稳定性。公式如下，Sublayer 表示Transformer 层中的前馈神经网络或自注意力模块。

$$DeepNorm(𝒙) = LayerNorm(𝛼 · 𝒙 + Sublayer(𝒙))$$

RMS Normalization (Root Mean Square Layer Normalization均方根归一化)

位置：与 Layer Norm 类似，常用于 LLM（如 LLaMA）的前馈层后。
公式：
$$\sigma = \sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2 + \epsilon} $$

$$y_i = \frac{x_i}{\sigma} \cdot \gamma $$

Transformers 代码库中LLaMA 的RMSNorm 实现代码：

class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        # 计算隐状态的均方根
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        # 将隐状态除以其均方根后重新缩放
        return self.weight * hidden_states.to(input_dtype)

优点：

计算量小（无需计算均值），训练更快。
在部分任务中表现优于 LayerNorm。

缺点：

缺少可学习的偏移参数（β），可能降低灵活性。

总结对比表

类型	计算维度	适用场景	优点	缺点
Layer Norm	特征维度（d）	NLP、Transformer	对序列长度不敏感	可能丢失通道差异性
RMS Norm	特征维度（d）	LLM（如 LLaMA）	计算高效，无均值中心化	缺少可学习偏移参数
Batch Norm	批维度（B）	CNN	加速训练，稳定梯度	依赖批大小，不适合变长数据