RMSNorm（Root Mean Square Layer Normalization）

RMSNorm 是一种类似于 Layer Normalization 的归一化技术，但它使用的是均方根（Root Mean Square, RMS）而非标准差。RMSNorm 在实际应用中表现出更好的稳定性和计算效率。

• 均值：描述数据的中心位置。
  $$\mu =\frac{1}{n}\sum _{i=1}^n{x}_i$$
• 方差：描述数据的离散程度，单位是原数据单位的平方。
  $${\sigma }^2=\frac{1}{n}\sum _{i=1}^n(x_i-\mu {)}^2$$
• 标准差：描述数据的离散程度，单位与原数据相同。
  $$\sigma =\sqrt{\frac{1}{n}\sum _{i=1}^n(x_i-\mu {)}^2}$$
• 均方根：描述数据的总体水平，常用于信号处理等领域。
  $$\text{RMS}=\sqrt{\frac{1}{n}\sum _{i=1}^n{x}_i^2}$$

数学公式

假设输入张量 $\mathbf{x}$ 的形状为 $(N,D)$，其中 $N$ 是批量大小，$D$是特征维度。RMSNorm 的计算步骤如下：
$$\begin{array}{rl}& 1.求均方根:\text{RMS}(\mathbf{x})=\sqrt{\frac{1}{D}\sum _{i=1}^D{x}_i^2}\\ & 2.归一化输入:\mathbf{y}=\frac{\mathbf{x}}{\text{RMS}(\mathbf{x})+ϵ}\\ & 3.缩放:\mathbf{z}=\gamma \cdot \mathbf{y},\gamma 是一个可学习的参数.\end{array}$$

代码实现

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

    def forward(self, x):
        variance = torch.var(x, dim=-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * x

总结

1. RMSNorm 的数学公式：
   • 计算均方根（RMS）。
   • 归一化输入。
   • 缩放归一化后的输出。

RMSNorm 相对于传统的 Layer Normalization 更加稳定，尤其是在长序列处理中表现更优。

LayerNorm和BatchNorm 参考