Skip to content

均方根归一化(RMSNorm)简介与推导

简介

均方根归一化(Root Mean Square Normalization,简称 RMSNorm) 是一种在深度学习模型中用于规范化神经网络层输入的方法。与层归一化(LayerNorm)和批归一化(BatchNorm)类似,RMSNorm 的目标是通过规范化输入数据来加速训练过程、提高模型稳定性,并提升模型的泛化能力。不同于 BatchNorm 依赖于批次维度,RMSNorm 不依赖于批次大小,因而在小批次训练或序列模型中表现尤为出色。此外,RMSNorm 通过简化归一化过程,减少了计算复杂度。

工作原理

RMSNorm 的核心思想是基于均方根(Root Mean Square,RMS)来规范化输入向量。具体来说,它通过计算输入向量的均方根值,接着使用该值对输入进行缩放,从而实现归一化。与 LayerNorm 类似,RMSNorm 也是在单个样本内部进行归一化,但它省略了均值的计算,简化了归一化过程。

归一化步骤

对于一个输入向量 \(\mathbf{x} = (x_1, x_2, \dots, x_H)\),RMSNorm 的计算过程如下:

  1. 计算均方根(RMS)值 (\(\text{RMS}\)): $$ \text{RMS} = \sqrt{ \frac{1}{H} \sum_{i=1}^{H} x_i^2 } $$

  2. 归一化: $$ \hat{x}_i = \frac{x_i}{\text{RMS} + \epsilon} $$ 其中,\(\epsilon\) 是一个极小的常数,防止分母为零。

  3. 缩放和平移: $$ y_i = \gamma \cdot \hat{x}_i + \beta $$ 这里,\(\gamma\)\(\beta\) 是可学习的参数,用于调整归一化后的输出。

最终公式

将上述步骤综合起来,RMSNorm 的数学表达式为:

\[ \text{RMSNorm}(\mathbf{x}) = \gamma \cdot \frac{\mathbf{x}}{\sqrt{ \frac{1}{H} \sum_{i=1}^{H} x_i^2 } + \epsilon} + \beta \]

常用符号

符号 含义
\(\mathbf{x}\) 输入向量
\(H\) 特征维度
\(\text{RMS}\) 均方根值
\(\gamma\) 缩放参数
\(\beta\) 平移参数
\(\epsilon\) 稳定常数

一个具体的例子

假设我们有一个 \(2 \times 4\) 的矩阵作为输入,即一个批次包含 2 个样本,每个样本的特征维度是 4:

\[ X = \begin{bmatrix} 1.0 & 2.0 & 3.0 & 4.0 \\ -1.0 & -2.0 & -3.0 & -4.0 \end{bmatrix} \]

为了简化计算,我们先不考虑 \(\epsilon\)(设为 \(0\)),并且假设缩放参数 \(\gamma\) 为全 \(1\) 向量,平移参数 \(\beta\) 为全 \(0\) 向量。

步骤 1: 计算每个样本的均方根值 (\(\text{RMS}\))

  • 对于第一个样本 \([1.0, 2.0, 3.0, 4.0]\): $$ \text{RMS}_1 = \sqrt{ \frac{1^2 + 2^2 + 3^2 + 4^2}{4} } = \sqrt{ \frac{1 + 4 + 9 + 16}{4} } = \sqrt{ \frac{30}{4} } = \sqrt{7.5} \approx 2.7386 $$

  • 对于第二个样本 \([-1.0, -2.0, -3.0, -4.0]\): $$ \text{RMS}_2 = \sqrt{ \frac{(-1)^2 + (-2)^2 + (-3)^2 + (-4)^2}{4} } = \sqrt{ \frac{1 + 4 + 9 + 16}{4} } = \sqrt{7.5} \approx 2.7386 $$

步骤 2: 进行归一化

  • 对于第一个样本: $$ \hat{x}{11} = \frac{1.0}{2.7386} \approx 0.3651 \ \hat{x} \approx 0.7303 \ \hat{x}} = \frac{2.0}{2.7386{13} = \frac{3.0}{2.7386} \approx 1.0954 \ \hat{x} \approx 1.4606 $$} = \frac{4.0}{2.7386

  • 对于第二个样本: $$ \hat{x}{21} = \frac{-1.0}{2.7386} \approx -0.3651 \ \hat{x} \approx -0.7303 \ \hat{x}} = \frac{-2.0}{2.7386{23} = \frac{-3.0}{2.7386} \approx -1.0954 \ \hat{x} \approx -1.4606 $$} = \frac{-4.0}{2.7386

步骤 3: 应用缩放和平移

这里 \(\gamma\) 为全 \(1\)\(\beta\) 为全 \(0\),所以结果不变。

因此,最终的 RMSNorm 输出矩阵 \(Y\) 为:

\[ Y \approx \begin{bmatrix} 0.3651 & 0.7303 & 1.0954 & 1.4606 \\ -0.3651 & -0.7303 & -1.0954 & -1.4606 \end{bmatrix} \]

RMSNorm 与其他归一化方法的比较

特性 BatchNorm LayerNorm RMSNorm
归一化维度 批次维度和特征维度 特征维度 特征维度
依赖批次大小 需要较大的批次大小 不依赖批次大小 不依赖批次大小
计算复杂度 相对较高 中等 较低
适用场景 图像分类等需要大批次的任务 序列模型如 RNN、Transformer 序列模型及小批次训练
参数数量 \(\gamma\)\(\beta\) 按特征维度 \(\gamma\)\(\beta\) 按特征维度 \(\gamma\)\(\beta\) 按特征维度

优点与缺点

优点

  1. 计算效率高:由于不需要计算均值,RMSNorm 在计算上比 LayerNorm 更为高效,尤其是在高维度数据中。
  2. 稳定性好:通过均方根值进行归一化,RMSNorm 能有效减小输入向量的尺度变化,提升模型训练的稳定性。
  3. 适用性广:不依赖于批次大小,使其在小批次训练和序列模型中表现良好。

缺点

  1. 信息损失:省略了均值的计算,可能会丢失部分关于输入分布的信息,影响模型的表达能力。
  2. 性能提升有限:在某些任务中,RMSNorm 相较于 LayerNorm 或 BatchNorm 的性能提升可能不明显。