docs/core/norm.md

core/norm.py · RMSNorm

代码

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return (self.weight * self.norm(x.float())).type_as(x)

公式

$$\mathrm{R}\mathrm{M}\mathrm{S}(x)=\sqrt{\frac{1}{d}\sum _{i=1}^d{x}_i^2+\epsilon },y=\frac{x}{\mathrm{R}\mathrm{M}\mathrm{S}(x)}\odot \gamma \backslash mathrm\{RMS\}(x)=\backslash sqrt\{\backslash frac\{1\}\{d\}\backslash sum\_\{i=1\}^\{d\}x\_i^2+\backslash varepsilon\},\backslash qquad\; y=\backslash frac\{x\}\{\backslash mathrm\{RMS\}(x)\}\backslash odot\; \backslash gamma$$

要点（面试）

• 相比 LayerNorm 省去均值中心化，只做缩放，参数更少、更稳定。
• 在 float32 里算归一化再 .type_as(x) 还原，避免低精度下数值爆炸。
• gamma（即 weight）初始化为 1，等价于先不做缩放。
• LLaMA / MiniMind 全程用 RMSNorm，不使用 bias。