RMSNorm：Root Mean Square Layer Normalization

主要贡献：提出归一化方法 RMSNorm，并通过实验证明归一化方法 re-scale 具有重要作用，re-center 收效甚微，基于此，放弃 re-center 步骤，仅保留 re-scale，降低了计算复杂度，提升了训练效率
论文：https://arxiv.org/abs/1910.07467

研究背景

• 提升模型的训练效率是一个长期存在的挑战，深度学习的多年发展中，优化器的改进、损失函数的优化、数据预处理步骤、模型参数初始化、训练权重归一化等，众多的优化方法目的只有一个——提升训练速度，加速模型收敛；
• BatchNorm 在序列问题上存在局限性，序列长度的不同导致 BatchNorm 计算得到的均值和方差不具有代表性，难以胜任序列任务的归一化方式；
• LayerNorm 能够较好的处理序列任务的归一化，作者对 LayerNorm 的两个核心机制 re-center、re-scale 进行研究，假设 re-scale 是 LayerNorm 最主要的机制，并通过实验验证猜想提出 RMSNorm，具有更少的计算量，更快的训练效率，可匹敌的收敛性能；
• 进一步提出 $p-\text{RMSNorm}$，基于独立同分布的假设，使用 $p\%$ 的参数估计整体的缩放系数，理论上进一步减少了计算量，但实验发现训练效率并未显著提升；

研究内容

$\text{LayerNorm}$

假设一个标准的前馈神经网络（线性层），输入向量 $x\in {\mathbb{R}}^m$，输出向量 $y\in {\mathbb{R}}^n$，神经网络的工作原理可由如下公式表示
$$a_i=\sum _{j=1}^m{w}_{ij}{x}_j,y_i=f\left(a_i+b_i\right)$$其中 $x_j$ 是输入向量第 $j$ 个值，$w_{ij}$ (矩阵 $n\times m$)是第 $i$ 个神经元对 $x_j$ 的权重，$f(\cdot )$ 表示非线性激活（逐元相乘），$b_i$ 代表第 $i$ 个神经元的偏置 (bias)

神经网络模型基于独立同分布的假设，如果不显式的干预中间层的分布，神经网路模型会遭遇内部斜变量偏移的问题（每层分布的变化随着网络层的加深其效应会被放大，导致不同层的分布不仅相同，违背了基于独立同分布的假设，并且模型难以训练和收敛【一种分布其实就是一种学习到的推理特征】），如下公式描述了 LayerNorm 均值和方差的计算方式
$${\overline{a}}_i=\frac{a_i-\mu }{\sigma }{g}_i,y_i=f\left({\overline{a}}_i+b_i\right)$$其中，${\bar{a}}_i$ 是向量 $\overline{a}\in {\mathbb{R}}^n$ 的第 $i$ 个值，$g\in {\mathbb{R}}^n$ 表示增益参数初始设置为1，$\mu$ 和 ${\sigma }^2$ 则是计算得到的均值和方差

$\text{RMSNorm}$

RMSNorm 公式如下，仅对数据分布做缩放变化：
$${\overline{a}}_i=\frac{a_i}{RMS(a)}{g}_i,\text{其中 }RMS(a)=\sqrt{\frac{1}{n}\sum _{i=1}^n{a}_i^2}$$

不变性分析

不变性衡量的是归一化后的模型输出是否会随着输入和权重矩阵的变化而产生显著改变，不同的归一化方法展现出不同的不变性特征，这对模型的稳健性有着重要的影响。考虑一个 RMSNorm 的一般形式：
$$y=f\left(\frac{Wx}{RMS(a)}\odot g+b\right)$$其中 $\odot$ 表示逐元相乘，假设权重矩阵缩放 $\delta$ 倍，即 $W^{\prime }=\delta W$，可做如下推导：
$$y^{\prime }=f\left(\frac{W^{\prime }x}{RMS\left(a^{\prime }\right)}\odot g+b\right)=f\left(\frac{\delta Wx}{RMS\left(\delta a\right)}\odot g+b\right)=f\left(\frac{\delta Wx}{\delta RMS\left(a\right)}\odot g+b\right)=y$$

由于 $RMS(\delta a)=\sqrt{\frac{1}{n}{\sum }_{i=1}^n(\delta a_i{)}^2}=\sqrt{\frac{{\delta }^2}{n}{\sum }_{i=1}^n(a_i{)}^2}=\delta \sqrt{\frac{1}{n}{\sum }_{i=1}^n(\delta a_i{)}^2}=\delta RMS(a)$

因此可以发现，当权重矩阵缩放时，输出具有不变性，同理将输入缩放 $\delta$ 倍，输出依然具有不变性

因为 $a=Wx$，因此无论是权重还是输入进行缩放，$a$ 都会同步缩放，因此输出结果依然具有不变性

如果仅对输出中的部分或权重的部分进行缩放，输出结果则不具有不变性，因为 RMS 的线性性质被破坏

简单举例，$A=[1,2,3]$, 现在对部分元素缩放得到 $A^{\prime }=[1,2\ast 1.5,3]$，$2.16=RMS(A)\ne RMS(A^{\prime })=2.52$

梯度分析

模型梯度的稳健性对于参数更新和模型收敛至关重要（论文 How Does Batch Normalization Help Optimization? 指出归一化方法之所以成功，并非源于为层输入增加了稳定性，而是因为优化态势的平滑度有所提高），这里将对 RMSNorm 进行梯度分析。给定一个损失函数 $\mathcal{L}$，根据链式法则，求 $\mathcal{L}$ 对 $g$、$b$ 的导数如下：
$$\frac{\partial \mathcal{L}}{\partial b}=\frac{\partial \mathcal{L}}{\partial v}\odot \frac{\partial v}{\partial b}=\frac{\partial \mathcal{L}}{\partial v},\frac{\partial \mathcal{L}}{\partial g}=\frac{\partial \mathcal{L}}{\partial v}\odot \frac{Wx}{RMS(a)}$$其中 $v=\frac{Wx}{RMS(a)}\odot g+b$

根据 $v=\frac{Wx}{RMS(a)}\odot g+b$ 可得 $\frac{\partial v}{\partial b}=1$，所以 $\frac{\partial \mathcal{L}}{\partial b}=\frac{\partial \mathcal{L}}{\partial v}$，此外这里的 $\mathcal{L}$ 并非一个真正的损失函数，而是 RMSNorm 的输出结果，在反向传播时，对 RMSNorm 中的变量求导数

此外我们将分析 RMSNorm 对权重 $W$ 的梯度
$$\frac{\partial \mathcal{L}}{\partial W}=\sum _{i=1}^n{\left[x^T\otimes \left(diag\left(g\odot \frac{\partial \mathcal{L}}{\partial v}\right)\times R\right)\right]}_i,\text{其中 }R=\frac{1}{RMS(a)}\left(I-\frac{(Wx)(Wx{)}^T}{nRMS(a{)}^2}\right)$$

这里进行示例推导，推导过程以标量形式示例，矩阵导数类似：

• 根据链式法则 $\frac{\partial \mathcal{L}}{\partial W}=\frac{\partial \mathcal{L}}{\partial v}\frac{\partial v}{\partial W},v=\frac{Wx}{RMS(a)}\odot g+b$
• $\frac{\partial v}{\partial W}=\frac{\partial (\frac{Wx}{RMS(a)}\odot g+b)}{\partial W}=g\odot \frac{\partial (\frac{Wx}{RMS(a)})}{\partial W}$，这里 $b,g$ 相对于 $W$ 为常数
• $\frac{\partial (\frac{Wx}{RMS(a)})}{\partial W}=x\frac{RMS(a)-W\frac{\partial RMS(a)}{\partial W}}{RM{S}^2(a)}$
• $\frac{\partial RMS(a)}{\partial W}=\frac{\partial RMS(a)}{\partial a}\frac{\partial a}{\partial W}$，由于 $a=Wx$，所以 $\frac{\partial a}{\partial W}=x$，假设$s={\sum }_{i=1}^n{a}_i^2$，则 $\frac{\partial RMS(a)}{\partial a}=\frac{\partial RMS(a)}{\partial s}\frac{\partial s}{\partial a}$，又 $RMS(a)=\sqrt{\frac{s}{n}}$，所以 $\frac{\partial RMS(a)}{\partial s}=\frac{1}{2\sqrt{n}}\frac{1}{\sqrt{s}}=\frac{1}{2n\cdot RMS(a)}$，又 $\frac{\partial s}{\partial a}=2a$，所以 $\frac{\partial RMS(a)}{\partial W}=\frac{1}{2n\cdot RMS(a)}\cdot x\cdot 2a=\frac{a\cdot x}{n\cdot RMS(a)}=\frac{Wx\cdot x}{n\cdot RMS(a)}$
• 经过上述步骤，根据链式法则进行组合逐步往上带入可得 $\frac{\partial (\frac{Wx}{RMS(a)})}{\partial W}=x\frac{RMS(a)-W\frac{\partial RMS(a)}{\partial W}}{RM{S}^2(a)}=\frac{\partial (\frac{Wx}{RMS(a)})}{\partial W}=x\frac{RMS(a)-W\frac{Wx\cdot x}{n\cdot RMS(a)}}{RM{S}^2(a)}=x\cdot \frac{1}{RMS(a)}(1-\frac{(Wx)(Wx)}{nRM{S}^2(a)})$
• 因此 $\frac{\partial \mathcal{L}}{\partial W}=\frac{\partial \mathcal{L}}{\partial v}\cdot g\cdot x\cdot \frac{1}{RMS(a)}(1-\frac{(Wx)(Wx)}{nRM{S}^2(a)})=x\cdot g\cdot \frac{\partial \mathcal{L}}{\partial v}\cdot \frac{1}{RMS(a)}(1-\frac{(Wx)(Wx)}{nRM{S}^2(a)})$
• 最后考虑到矩阵形式，整理结果 $\frac{\partial \mathcal{L}}{\partial W}={\sum }_{i=1}^n{\left[x^T\otimes \left(diag\left(g\odot \frac{\partial \mathcal{L}}{\partial v}\right)\times R\right)\right]}_i$，其中 $R=\frac{1}{RMS(a)}\left(I-\frac{(Wx)(Wx{)}^T}{nRMS(a{)}^2}\right)$

将输入 $x$ 和权重矩阵 $W$ 进行缩放，讨论 RMSNorm 的梯度不变性，我们将 $\delta x$ 或者 $\delta W$ 代入 $R$ 得到$$R^{\prime }=\frac{1}{\delta RMS(a)}\left(I-\frac{(\delta Wx)(\delta Wx{)}^T}{n{\delta }^{2}RMS(a{)}^2}\right)=\frac{1}{\delta }R$$
所以

• 将 $\delta x$ 代入 $\frac{\partial \mathcal{L}}{\partial W}={\sum }_{i=1}^n{\left[\delta x^T\otimes \left(diag\left(g\odot \frac{\partial \mathcal{L}}{\partial v}\right)\times \frac{1}{\delta }R\right)\right]}_i=\frac{\partial \mathcal{L}}{\partial W}$，表明 RMSNorm 的梯度对于输入 $x$ 具有缩放不变性，降低梯度对输入缩放的敏感度可确保其平滑性，并提高学习的稳定性
• 同理将 $\delta x$ 代入 $\frac{\partial \mathcal{L}}{\partial W}$ 得到 $\frac{1}{\delta }\frac{\partial \mathcal{L}}{\partial W}$，表明 RMSNorm 的梯度对于权重矩阵 $W$ 具有反比性质，这种负相关关系起到了一种隐性学习率调节器的作用，能够动态地控制梯度的范数，从而避免出现大范数的权重矩阵，并促进模型的收敛【权重过大，梯度会更小，权重过小，梯度会过大】

$p\text{RMSNorm}$

用部分数据来描述整体的分布，即将部分值用于计算 $RMS(\cdot )$，$\overline{RMS}(a)=\sqrt{\frac{1}{k}{\sum }_{i=1}^k{a}_i^2}$，其中 $k=\lceil n\cdot \%p\rceil$ 表示用于均方根估计的元素数量，$\lceil \cdot \rceil$ 表示向上取整，不难证明 $p\text{RMSNorm}$ 与 $\text{RMSNorm}$ 具有相同的性质

作者指出，当 $m$ 较小时，梯度容易出现爆炸现象，我感觉这里像是笔误，这里的 $m$ 可能是输入 $x$ 的维度，也可能是 $p$ 写成了 $m$，此外，作者通过实验发现，当采用 $p=6.25\%$ 时能够成功实现令人满意的收敛

实验结果

作者做了非常充足的实验，包含在不同的框架（PyTorch、Tensorflow、Theano ）、不同的模型架构（RNN 变体、卷积和自注意力模型），不同的激活函数（sigmoid、tanh、softmax）、不同的归一化方式（无归一化、LayerNorm）等，作者实验太多了，这里就不贴图了，直接看下方结论部分

结论

• 实验证实了作者的猜想，在 LayerNorm 中，相比于 re-centering，re-scaling 起主要作用；
• $\text{RMSNorm}$ 没有 re-centering，计算量更少，训练速度更快，并且精度不低于$\text{LayerNorm}$，根据不同的网络架构、训练框架、硬件，提升性能在 $7\%-64\%$；
• $p\text{RMSNorm}$ 计算量低于 $\text{RMSNorm}$，在各项试验中，精度也低于 $\text{RMSNorm}$，但是训练速度并未显著提升，甚至可能更慢，作者认为可能是底层张量的切片算法拉低了训练效率；
• 序列任务中，不引入归一化、有概率导致梯度爆炸，训练失败，序列任务中精度 $\text{RMSNorm}>\text{LayerNorm}>\text{BatchNorm}$；
• 图像任务精度 $\text{BatchNorm}>\text{RMSNorm}>\text{LayerNorm}$；
• RMSNorm 对输入具有缩放不变性，降低梯度对输入缩放的敏感度可确保其平滑性，并提高学习的稳定性；
• RMSNorm 的梯度对于权重矩阵具有反比性质，权重过大，梯度更小，权重过小，梯度更大，能够动态地控制梯度的范数，从而避免出现大范数的权重矩阵；