layer norm和 rms norm 对比

Zero_to_zero12342025-06-05 17:46

Layer norm

python 复制代码

# Layer Norm 公式
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True)
output = (x - mean) / sqrt(var + eps) * gamma + beta

特点：

减去均值（去中心化）
除以标准差（标准化）
包含可学习参数 gamma 和 beta
计算复杂度相对较高

RMS Norm（Root Mean Square归一化）：

python 复制代码

# RMS Norm 公式
rms = sqrt(mean(x²))
output = x / rms * gamma

特点：

不减去均值（保持中心）
只除以RMS值
只有一个可学习参数 gamma
计算更简单高效

对比

代码对比

python 复制代码

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(dim))
        self.beta = nn.Parameter(torch.zeros(dim))
        self.eps = eps
    
    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        var = x.var(-1, keepdim=True, unbiased=False)
        return (x - mean) / torch.sqrt(var + self.eps) * self.gamma + self.beta

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(dim))
        self.eps = eps
    
    def forward(self, x):
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
        return x / rms * self.gamma