RMSNorm：均方根层标准化

在深度神经网络中，LayerNorm用于帮助稳定化训练且提升模型的拟合能力。这是因为LayerNorm对输入和权重矩阵具有re-centering与re-scaling不变的特性。然而，随着网络加深，尤其是RNN，因计算量越来越大导致LayerNorm带来性能的成本越来越高。由此，在LayerNorm的re-centering不变属于不必要的假设下，RMSNorm基于均方根进行标准化，且使模型拥有re-scaling不变的特性和学习率的隐式自适应能力。

LayerNorm

给定输入向量$\mathbf{x}\in\mathbb{R}^m$，全连接网络通过线性变换与非线性激活函数把它映射为输出向量$\mathbf{y}\in\mathbb{R}^n$。
$$
\begin{aligned}
a_i=\sum_{j=1}^mw_{ij}x_j,\quad y_i=f(a_i+b_i)
\end{aligned}\tag{1}
$$
神经网络往往受到内部协方差偏移的影响，可以理解为每层的输入随着之前层参数的改变为改变，从而导致模型训练的不稳定。为了减少偏移的影响，LayerNorm利用均值与方差对输入和进行标准化
$$
\begin{aligned}
\bar{a}_i=\frac{a_i-\mu}{\sigma}g_i,\quad y_i=f(\bar{a}_i+b_i)
\end{aligned}\tag{2}
$$
式(2)中$g\in\mathbb{R}^n$为增益参数，用于缩放标准输入和，起始设置为1。均值与方差的计算方式为
$$
\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}\tag{3}
$$
对于LayerNorm有效的解释，它拥有re-centering和re-scaling不变的特性。其中，前者使模型对输入与权重参数因噪音的迁移不敏感，而后者可在输入与权重参数被随机的缩放时仍保证输出表示完整。

RMSNorm

作者们以re-scaling不变性为LayerNorm有效的原因，而不是re-centering不变性，作为假设，提出了RMSNorm
$$
\begin{aligned}
\bar{a}_i=\frac{a_i}{RMS(\mathbf{a})}g_i,\quad where\quad RMS(\mathbf{a})=\sqrt{\frac{1}{n}\sum_{i=1}^na_i^2}
\end{aligned}\tag{4}
$$
由于不需要计算均值与方差，只需要计算均方根，因此RMSNorm减少了计算量，提高了LayerNorm的效率。实验结果表明，RMSNorm拥有与LayerNorm相似的性能，或甚至更有效。

不变性分析

不变性测量的是标准化后模型输出的改变是否与输入和权重矩阵高度一致。考虑RMSNorm的通用形式
$$
\begin{aligned}
\mathbf{y}=f(\frac{\mathbf{W}x}{RMS(\mathbf{a})}\odot\mathbf{g}+\mathbf{b})
\end{aligned}\tag{5}
$$
式(5)中$\odot$为元素级别的相乘。

RMSNorm对输入和权重矩阵的缩放不变，这是因为RMS拥有线性特性
$$
\begin{aligned}
RMS(\alpha\mathbf{x})=\alpha RMS(\mathbf{x})
\end{aligned}\tag{6}
$$
若权重参数矩阵$\mathbf{W}$被因子$\delta$缩放，也就是${\mathbf{W}}'=\delta\mathbf{W}$，那么最后一层的输出不改变
$$
\begin{aligned}
{y}'=f(\frac{\mathbf{W}'x}{RMS({a}')}\odot g+b)=f(\frac{\delta\mathbf{W}x}{\delta RMS(a)}\odot g+b)=y
\end{aligned}\tag{7}
$$
对输入数据的缩放，RMSNorm拥有相同的特性。

梯度分析

RMSNorm是通过随机梯度下降训练的方式增强神经网络，即梯度的鲁棒性对模型参数更新和模型的拟合很重要。Santurkar等人认为标准化方法的有效来源于优化地形的增强的平滑性，而不是每层输入增强的稳定性。因此，也研究了模型的梯度。

给定损失函数$\mathcal{L}$，式(4)的反向传播利用关于参数$\mathbf{g},\mathbf{b}$的梯度进行反向传播
$$
\begin{aligned}
\frac{\partial\mathcal{L}}{\partial\mathbf{b}}=\frac{\partial\mathcal{L}}{\partial\mathbf{v}},\quad \frac{\partial\mathcal{L}}{\partial\mathbf{g}}=\frac{\partial \mathcal{L}}{\partial\mathbf{v}}\odot\frac{\mathbf{W}\mathbf{x}}{RMS(\mathbf{a})}
\end{aligned}\tag{8}
$$
式(8)中$\mathbf{v}=\frac{{\mathbf{W}}'x}{RMS({a}')}\odot g+b$。梯度$\frac{\partial\mathcal{L}}{\partial b}$与$\frac{\partial\mathcal{L}}{\partial g}$均对输入$\mathbf{x}$和权重参数矩阵$\mathbf{W}$的缩放拥有不变性。

参数矩阵$\mathbf{W}$的梯度
$$
\begin{aligned}
\frac{\partial\mathcal{L}}{\partial\mathbf{W}}=\sum_{i=1}^n[\mathbf{x}^T\otimes (diag(g\odot\frac{\partial\mathcal{L}}{\partial\mathbf{v}})\times\mathbf{R})]_i,\quad where\quad\mathbf{R}=\frac{1}{RMS(\mathbf{a})}(\mathbf{I}-\frac{(\mathbf{W}\mathbf{x})(\mathbf{W}\mathbf{x})^T}{nRMS(\mathbf{a})^2})
\end{aligned}\tag{9}
$$
式(9)中$diag(\cdot)$表示输入的对角矩阵，$\otimes$为Kronecker积。

对于$\mathbf{R}$，通过分析证明其与输入和权重矩阵的缩放呈现负相关。若进行缩放输入$\mathbf{x}({\mathbf{x}'=\delta\mathbf{x}})$或权重矩阵$({\mathbf{W}}'=\delta\mathbf{W})$，那么
$$
\begin{aligned}
{\mathbf{R}}'=\frac{1}{\delta RMS(\mathbf{a})}(\mathbf{I}-\frac{(\delta\mathbf{W}\mathbf{x})(\delta\mathbf{W}\mathbf{x})^T}{n\delta^2 RMS(\mathbf{a})^w})=\frac{1}{\delta}\mathbf{R}
\end{aligned}\tag{10}
$$
由此，可得权重矩阵的梯度对输入的缩放拥有不变性，而与权重矩阵的缩放呈现负相关。权重矩阵的梯度对输入缩放的不敏感性确保了平滑性且提升了学习的稳定性。另一方面，负相关性可视为隐式学习率的自适应器，动态的控制梯度的范数，用于避免权重矩阵有较大的范数且提升模型的拟合能力。

$p$RMSNorm

RMSNorm的re-scaling不变特性是由RMS的线性特性造成的。若每层的神经元拥有独立同分布的结构，那么均方根RMS可被这些神经元的子集表示。作者们也提出了部分RMSNorm，被称为pRMSNorm。给定无标准化输入$\mathbf{a}$，pRMSNorm利用其前$p%$的元素估计RMS统计：$\overline{RMS}(a)=\sqrt{\frac{1}{k}\sum_{i=1}^k a_i^2}$，其中$k=\lceil n\cdot p\rceil$表示用于估计RMS的元素数量。其中，实验表明$\overline{RMS}$过小会造成梯度爆炸，这是因为其为RMS的有偏估计。因此，作者推荐$p=6.25\%$