LayerNorm不同于BatchNorm,其与batch大小无关,均值和方差 在 每个样本的特征维度 C 内计算,
适用于 变长输入(如 NLP 任务中的 Transformer)
详细的BatchNorm在之前的一篇文章进行了详细的介绍:深度学习中Batch Normalization(BN)原理、作用浅析-CSDN博客
这里主要介绍适合于Transformer架构的两个归一化操作RMSNorm 和 LayerNorm
RMSNorm 和 LayerNorm 的本质区别
RMSNorm(Root Mean Square Normalization)和 LayerNorm(Layer Normalization)都是 归一化方法,但它们的本质区别在于 是否去均值(Mean-Centering) 以及 归一化的方式。
LayerNorm公式
- μ:均值,计算的是特征
x
在d
维度上的平均值。 - σ2:方差,用于衡量数值分布的变化范围。
- γ,β:可学习的缩放和偏移参数。
- LayerNorm 不仅仅缩放数据,还会让其均值归一化到
0
,保证分布居中
RMSNorm公式
- RMSNorm 直接用
RMS(x)
归一化,而不去均值。 - RMSNorm 只调整
L2
范数的大小,不影响数据的中心位置。
代码实现
class LayerNorm(nn.Module):def __init__(self, dim, eps=1e-5):super().__init__()self.eps = epsself.weight = nn.Parameter(torch.ones(dim))self.bias = nn.Parameter(torch.zeros(dim))def forward(self, x):mean = x.mean(dim=-1, keepdim=True)std = x.var(dim=-1, keepdim=True, unbiased=False).sqrt()return self.weight * (x - mean) / (std + self.eps) + self.biasclass RMSNorm(nn.Module):def __init__(self, dim, eps=1e-5):super().__init__()self.eps = epsself.weight = nn.Parameter(torch.ones(dim))def forward(self, x):rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)return self.weight * (x / rms)
RMSNorm 的优缺点
LayerNorm 的优点
- 更稳定的梯度更新
- 由于均值归 0,梯度更新不会受到偏移影响。
- 适用于 Transformer、BERT、GPT。
- 适用于不同任务
- 既可以用于 NLP(Transformer),也可以用于 CNN。
- 适用于变长输入(如 RNN、BERT)。
- 训练和推理一致
- LayerNorm 不依赖 batch_size,在训练和推理时表现一致。
LayerNorm 的缺点
- 计算量大
- 需要 计算均值和方差,相比 RMSNorm 额外增加一次均值计算,计算量更高。
- 计算开销大,不适合大模型
- 在 大规模 Transformer(如 LLaMA) 中,LayerNorm 计算量太大,影响训练速度。
- 对 batch_size 影响较大
- 在 小 batch_size 时,LayerNorm 可能表现不稳定。
RMSNorm 的优缺点
RMSNorm 的优点
- 计算更快
- 仅计算 L2 归一化,比 LayerNorm 计算量少 约 30%。
- 适用于大模型(如 LLaMA, GPT-4)
- 在 大模型训练中,RMSNorm 比 LayerNorm 更高效。
RMSNorm 的缺点
- 不去均值,可能影响训练稳定性
- 在某些任务中,均值归 0 能稳定训练,而 RMSNorm 不能。
- 不适用于 CNN
- CNN 依赖均值信息,RMSNorm 不计算均值,可能导致训练不稳定。
总结
RMSNorm通过简化归一化过程,降低计算复杂度,提供了一种有效的归一化方法。它在保持模型性能的同时,提高了计算效率,是LayerNorm的有力替代方案。