llama 2 改进之 RMSNorm

RMSNorm
在这里插入图片描述

论文：https://openreview.net/pdf?id=SygkZ3MTJE
Github：https://github.com/bzhangGo/rmsnorm?tab=readme-ov-file
在这里插入图片描述
论文假设LayerNorm中的重新居中不变性是可有可无的，并提出了均方根层归一化(RMSNorm)。RMSNorm根据均方根(RMS)将一层神经元的总和输入正则化，得到模型重新缩放不变性特性和隐式学习率适应能力

LayerNorm 公式

深度学习当中，没有线性激活函数的预测公式

$\begin{aligned}a_i=\sum_{j=1}^mw_{ij}x_j,\quad y_i=f\left(a_i+b_i\right),\end{aligned}$

通过激活函数后，其中，随着前一层的更新，层的输入分布会发生变化。这可能会对参数梯度的稳定性产生负面影响，延迟模型收敛。为了减少这种转变，LayerNorm 对求和的输入进行归一化，以固定它们的均值和方差，如下所示：

$\begin{aligned}\bar{a}_i=\frac{a_i-\mu}{\sigma}g_i,\quad y_i=f\left(\bar{a}_i+b_i\right),\end{aligned}$

其中 $\bar{a}_i$ 是向量 $\bar{a}\in\mathbb{R}^n$ 的第 $i$ 个值，作为 $\alpha_i$ 的归一化替代值用于层激活。 $\mathbf{g}\in\mathbb{R}^n$ 是增益参数，用于重新调整标准化求和输入的大小，一开始设置为 1。 $\mu$ 和 $\sigma^2$ 分别是根据原始求和输入估计的均值和方差统计量。

$\begin{aligned}\mu=\frac{1}{n}\sum_{i=1}^na_i,\quad\sigma=\sqrt{\frac{1}{n}\sum_{i=1}^n(a_i-\mu)^2}.\end{aligned}$

在本文中，假设重新缩放不变性是LayerNorm成功的原因，而不是重新定中心不变性。我们提出了RMSNorm，它只关注重新缩放不变性，并简单地根据均方根（RMS）统计对求和输入进行正则化：
$\begin{aligned}\bar{a}_i=\frac{a_i}{\text{RMS}(\mathbf{a})}g_i,\quad\text{where RMS}(\mathbf{a})=\sqrt{\frac{1}{n}\sum_{i=1}^na_i^2}.\end{aligned}$

python实现

# root mean square layer normalization
def rln(x, s):_eps = 1e-5output = x / tensor.sqrt((x * x).mean(1)[:,None] + _eps)output = s[None, :] * outputreturn output# layer normalization
def ln(x, b, s):_eps = 1e-5output = (x - x.mean(1)[:,None]) / tensor.sqrt((x.var(1)[:,None] + _eps))output = s[None, :] * output + b[None,:]return output

使用pytorch来写RMSNorm的函数

import torch
import torch.nn as nnclass RMSNorm(nn.Module):def __init__(self, d, p=-1., eps=1e-8, bias=False):"""Root Mean Square Layer Normalization:param d: model size:param p: partial RMSNorm, valid value [0, 1], default -1.0 (disabled):param eps:  epsilon value, default 1e-8:param bias: whether use bias term for RMSNorm, disabled bydefault because RMSNorm doesn't enforce re-centering invariance."""super(RMSNorm, self).__init__()self.eps = epsself.d = dself.p = pself.bias = biasself.scale = nn.Parameter(torch.ones(d))self.register_parameter("scale", self.scale)if self.bias:self.offset = nn.Parameter(torch.zeros(d))self.register_parameter("offset", self.offset)def forward(self, x):if self.p < 0. or self.p > 1.:norm_x = x.norm(2, dim=-1, keepdim=True)d_x = self.delse:partial_size = int(self.d * self.p)partial_x, _ = torch.split(x, [partial_size, self.d - partial_size], dim=-1)norm_x = partial_x.norm(2, dim=-1, keepdim=True)d_x = partial_sizerms_x = norm_x * d_x ** (-1. / 2)x_normed = x / (rms_x + self.eps)if self.bias:return self.scale * x_normed + self.offsetreturn self.scale * x_normed