llama 2 改进之 RMSNorm

news2025/11/1 4:47:54

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-5
    output = x / tensor.sqrt((x * x).mean(1)[:,None] + _eps)
    output = s[None, :] * output
    return output

# layer normalization
def ln(x, b, s):
    _eps = 1e-5
    output = (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 nn


class 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 by
            default because RMSNorm doesn't enforce re-centering invariance.
        """
        super(RMSNorm, self).__init__()

        self.eps = eps
        self.d = d
        self.p = p
        self.bias = bias

        self.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.d
        else:
            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_size

        rms_x = norm_x * d_x ** (-1. / 2)
        x_normed = x / (rms_x + self.eps)

        if self.bias:
            return self.scale * x_normed + self.offset

        return self.scale * x_normed