[归一化]RMSNorm

RMSNorm

输入向量$x\in {\mathbb{R}}^m$,输出向量$y\in {\mathbb{R}}^n$

线性变换：$y_i=f(a_i+b_i)$

其中：

• 非线性激活函数：$a_i=\sum _{j=1}^m{w}_{ij}{x}_j$
• $w_i$为第i个神经元的权重，$b_i$为偏置项

LayerNorm：

${\overline{a}}_i=\frac{a_i?\mu }{\sigma }{g}_i$，$y_i=f({\overline{a}}_i+b)$

其中：

• 均值$\mu =\frac{1}{n}\sum _{i=1}^n{a}_i$
• 方差$\sigma =\sqrt[]{\frac{1}{n}\sum _{i=1}^n(a_i?\mu {)}^2}$

代码实现

class LayerNorm(torch.nn.Module):
  def __int__(self, dim, eps=1e-6):
    self.eps = eps
    self.weight = nn.Parameter(dim)
  
  def forward(self, x):
    output = self._norm(x)
    return output * self.weight

RMSNorm：

${\overline{a}}_i=\frac{a_i}{RMS(a)}{g}_i$，$y_i=f({\overline{a}}_i+b)$

其中：

• $RMS(a)=\sqrt[]{\frac{1}{n}\sum _{i=1}^n{a}_i^2}$

补充：

• 不考虑re-center，效果几乎相似但效率更高
• 是LayerNorm中均值为0的特殊情况

代码实现

class RMSNorm(torch.nn.Module):
  def __init__(self, dim, eps=1e-6):
    self.eps = eps
    self.weight = nn.Parameter(torch.ones(dim))
    
  def _norm(self, x):
    return x * torch.rsqrt(x.pow(2).mean(-1, keep_dim=True) + self.eps)
  
  def forward(self, x):
    output = self._norm(x.float()).type_as(x)
    return output * self.weight