RMSNorm

接口功能

LayerNorm的轻量变体，仅使用均方根值归一化（无均值中心化）。

函数原型

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

    def forward(self, x: torch.Tensor):
        # 计算均方根值
        rms = x.norm(2, dim=-1, keepdim=True) * (x.size(-1) ** -0.5)
        return x / (rms + self.eps) * self.weight

参数说明