每日基础课

把一层网络输出的向量想成一排音量忽大忽小的麦克风。数值过大或过小，后面的层就很难稳定学习。RMSNorm 先计算整组数的均方根，再用它缩放每个元素，让整体尺度稳定。

公式是 $$y_i = g_i \frac{x_i}{\sqrt{\frac{1}{d}\sum_{j=1}^{d}x_j^2+\epsilon}}$$。

• $x_i$：输入向量第 $i$ 个数；$d$：最后一维的长度。
• 平方后求平均：只关心整体能量，不让正负号抵消。
• 开根号：把量纲拉回原来尺度；$\epsilon$：防止除以 0。
• $g_i$：可学习缩放参数，让模型自己决定每个维度最终应有多大。

实际使用时通常对每个 token 的 hidden dimension 单独归一化。输入形状若是 `[batch, time, hidden]`，RMSNorm 就沿 `hidden` 维计算。LLaMA 等模型常把它放在注意力层或前馈层之前，也就是 Pre-Norm 结构。

先沿最后一维计算均方根，再相除并乘可学习权重。

python
import torch
from torch import nn

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

    def forward(self, x):
        rms = torch.sqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (x.float() / rms * self.weight.float()).to(x.dtype)

torch.manual_seed(0)
x = torch.randn(2, 3, 8)
mine = MyRMSNorm(8)
official = nn.RMSNorm(8, eps=1e-6)
official.weight.data.copy_(mine.weight.data)

y1, y2 = mine(x), official(x)
print('shape:', y1.shape)
print('max error:', (y1 - y2).abs().max().item())
assert torch.allclose(y1, y2, atol=1e-5)