归一化层：LayerNorm vs RMSNorm

归一化层是深度学习训练稳定性的关键组件。LayerNorm 是 Transformer 的标配，而 RMSNorm 以更低的计算成本成为大模型的新选择。

实现对比

class LayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.bias = 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 self.weight * (x - mean) / (var + self.eps).sqrt() + self.bias

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

RMSNorm 为何更受大模型青睐

RMSNorm 移除了均值计算和偏置项，计算量减少约 10-15%。在超大规模训练中，这节省的计算量非常可观。LLaMA、Mistral 等模型均采用 RMSNorm。

相关资源

• LayerNorm 论文
• RMSNorm 论文
• BatchNorm 论文