RMSNorm（Root Mean Square Layer Normalization）是一种替代传统 Layer Normalization 的归一化方法，旨在通过简化计算提升效率和性能。它的作用是对输入张量进行归一化，稳定训练过程并加速收敛。
归一化：对输入张量进行归一化，使其数值分布更稳定。
简化计算：相比 Layer Normalization，RMSNorm 去除了均值计算，仅使用均方根（RMS）进行缩放，计算更高效。
加速收敛：归一化有助于缓解梯度问题，加速模型训练。
$$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \cdot \gamma $$
其中：
$\mathrm{RMS}(x) = \sqrt{\frac{1}{n}\sum_{i=1}^n x_i^2}$是输入 $x$ 的均方根。
$\gamma$是可学习的缩放参数。
import torch
import torch.nn as nn
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-8):
super(RMSNorm, self).__init__()
self.scale = nn.Parameter(torch.ones(dim)) # 可学习的缩放参数
self.eps = eps # 防止除零的小常数
def forward(self, x):
rms = torch.sqrt(torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps)
return x / rms * self.scale
# 使用示例
rms_norm = RMSNorm(dim=64)
input_tensor = torch.randn(4, 16, 64) # 输入张量
output_tensor = rms_norm(input_tensor)
print(output_tensor.shape) # 输出形状不变
Transformer 模型：RMSNorm 常用于替代 Layer Normalization，提升训练效率。
深度学习模型：适用于需要归一化的场景，如自然语言处理和计算机视觉。
RMSNorm 通过简化归一化计算，提升效率和性能，常用于 Transformer 等深度学习模型。