RMSnorm 均方根正则化

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        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):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

        rsqrt:reciprocal of the square-root

RMS:Root Mean Square Normalization

归一化后的 output 张量逐元素乘以 self.weight。这里的 self.weight 是一个形状为 (dim,) 的可学习参数，初始化为全 1。

举例说明（广播机制）

假设我们有如下输入和参数：

• 输入张量 x，形状为 (2, 3)：

  x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  
  

• 可学习参数 self.weight，形状为 (3,)：

  self.weight = torch.tensor([0.5, 1.0, 1.5])