LayerNorm vs RMSNorm 技术对比
1. 核心概念
LayerNorm (层归一化)
- 思想:对单个样本的所有特征维度进行归一化
- 目标:使每个样本的特征分布 μ = 0 \mu=0 μ=0, σ = 1 \sigma=1 σ=1
- 特点:同时调整均值和方差
RMSNorm (均方根归一化)
- 思想:基于均方根的简化归一化方法
- 目标:仅通过缩放调整特征幅度
- 特点:不调整均值,仅缩放方差
2. 数学表达
LayerNorm公式
μ = 1 d ∑ i = 1 d x i σ = 1 d ∑ i = 1 d ( x i − μ ) 2 x ^ i = x i − μ σ + ϵ y i = γ i x ^ i + β i \begin{aligned} \mu &= \frac{1}{d}\sum_{i=1}^d x_i \\ \sigma &= \sqrt{\frac{1}{d}\sum_{i=1}^d (x_i - \mu)^2} \\ \hat{x}_i &= \frac{x_i - \mu}{\sigma + \epsilon} \\ y_i &= \gamma_i \hat{x}_i + \beta_i \end{aligned} μσx^iyi=d1i=1∑dxi=d1i=1∑d(xi−μ)2=σ+ϵxi−μ=γix^i+βi
RMSNorm公式
σ = 1 d ∑ i = 1 d x i 2 x ^ i = x i σ + ϵ y i = γ i x ^ i \begin{aligned} \sigma &= \sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2} \\ \hat{x}_i &= \frac{x_i}{\sigma + \epsilon} \\ y_i &= \gamma_i \hat{x}_i \end{aligned} σx^iyi=d1i=1∑dxi2=σ+ϵxi=γix^i
3. 关键差异对比
特性 | LayerNorm | RMSNorm |
---|---|---|
均值调整 | ✓ \checkmark ✓ | × \times × |
方差调整 | 标准差归一化 | 均方根归一化 |
计算复杂度 | 较高 | 较低 |
可学习参数 | γ , β \gamma,\beta γ,β | 通常仅 γ \gamma γ |
适用场景 | 对分布敏感的任务 | 高效计算场景 |
4. 优缺点分析
LayerNorm
- 优点:
- 分布调整更全面
- 适合精细控制的任务
- 缺点:
- 计算开销较大
RMSNorm
- 优点:
- 计算效率高(省去均值计算)
- 适合大规模模型
- 缺点:
- 忽略均值可能影响效果
5. 实现代码
PyTorch实现
# LayerNorm
import torch.nn as nn
layernorm = nn.LayerNorm(feature_dim)# RMSNorm
class RMSNorm(nn.Module):def __init__(self, dim):super().__init__()self.scale = nn.Parameter(torch.ones(dim))def forward(self, x):rms = x.pow(2).mean(-1, keepdim=True).sqrt() + 1e-8return x / rms * self.scale
备注
有问题随时交流~~