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])