在前面两篇文章「模型蒸馏的原理是什么？温度有什么用？起底LLM中的蒸馏技术！」和「更小、更快、更便宜，DistilBERT 模型迁移与模型蒸馏的结合！」中我们陆续介绍了大模型中频繁用到的模型蒸馏技术，包括其原理动机和从零实现。
本以为介绍完模型蒸馏以后就可以介绍了 DeepSeek 了，结果没想到刚读没两页就发现 DeepSeek 的第一代模型 DeepSeek LLM 是基于 LLaMA 1 结构训练而来，于是又开始去看 LLaMA 1。 然而看到 LLaMA 1 的网络结构部分，说是使用 RMSNorm 来替代了原有的 LayerNorm 归一化。
于是我们就开始深扒了一下，发现 LLaMA 1、 LLaMA 2、 LLaMA 3、DeepSeek LLM、DeepSeek V3 、百川大模型等等都使用了 RMSNorm 来替代了原有的 LayerNorm 。好家伙，整个一改天换面了 ，那必须要给大家介绍一下。
所以，最后又开始来看 RMSNorm 到底是个什么玩意儿了。看来，正式步入到介绍 DeepSeek 还有很长的路要走。
那一句话总结，RMSNorm 就是在 LayerNorm 基础上移除了去均值操作！
最后，作者在多种任务上对模型进行了全面评估，包括机器翻译、图像分类、图文检索以及问答任务。实验结果表明，在各种模型中，RMSNorm 的性能与 LayerNorm 相当，但在运行速度方面具有优势，提速幅度从 $7\%$ 到 $64\%$不等。这说明 RMSNorm 并不会明显牺牲模型效果，但显著加快了训练或推理过程，尤其在计算资源紧张或需要加速推理时非常有用。
关键字：归一化、批归一化、层归一化、LayerNorm、RMSNorm、LLaMA、DeepSeek
1. 标准化方法回顾#
为了让各位读者能够更加清楚地了解前因后果，在这里首先带着大家来简单回顾一下深度学习中几种常见的归一化方法，并简要叙述一下各个归一化方法所提出来的动机。
1.1 原始去均值标准化#
什么是原始去均值标准化呢？可能这个名字也不太准确，但是这里想表达的意思是后续所有的标准化方法都是在其基础之上改进而来的，所以有必要给大家交代一下。这种标准化方法主要是用于传统的机器学习中，其计算公式如下：
$$ {x}'=\frac{x-\mu }{\sqrt{Var} }\tag{1} $$其中 $\mu$ 表示每一列特征的平均值，$Var$ 表示每一列特征的方差，开根号以后便是标准差。
在式(1)中，分子部分呢便是去中心化操作，因为该维度的所有特征值都减去了均值；分子再除以分母便是归一化操作，用来消除不同维度间量纲的影响。
从图1可以看出，标准化前和标准化后都不会改变数据的分布形状以及内在的结构信息，它的主要作用便是提高模型训练的稳定性和效率。
详细原理可参见机器学习栏目「第4.2节 特征标准化教程：为什么要做归一化与标准化」
1.2 批归一化#
在深度学习中，随着网络层数的加深越是越靠近输入层的权重参数其对应的梯度也将会越小，所以这就导致靠近输出层的权重参数能够更容易的得到训练，而靠近输入层的权重参数则更新缓慢。也正是由于两端的权重参数在更新节奏上相差太大网络底部的参数更新不及时，使得顶部的参数每次都要根据底部的前向传播结果来重新适应数据的分布，从而加大了网络的训练难度。
Batch Normalization的思想便是在神经网络中添加一层归一化操作，使得网络中每一层输入的分布都尽可能地接近标准高斯分布，从而减轻这种问题。
假设现在有一个$d$维的网络层，其输出为$x=(x^{(1)},x^{(2)},...,x^{(d)})$，那么对于每一个维度都可以通过式(2)中的方法进行标准化
$$ \hat{x}^{(k)}=\frac{x^{(k)}-E[x^{(k)}]}{\sqrt{Var[x^{(k)}]}}\tag{2} $$其中，$E[x^{(k)}]$和$Var[x^{(k)}]$分别是第$k$个维度在所有样本上计算得到的期望和方差。
当然，但如果仅仅只是简单通过式(2)中的计算方法来对每个维度进行标准化，那么在某些情况下将会改变该维度原有的表示信息。例如在对Sigmoid激活函数的输入值进行标准化时，通过公式(2)标准化后的输入值可能只会趋于0附近，从而把Sigmoid变成了一个线性激活函数。为了解决这一问题，可以加入一组学习的参数$\gamma^{(k)}$和$\beta^{(k)}$来对$\hat{x}^{(k)}$进行了一次线性变换，即
$$ y^{(k)}=\gamma^{(k)}\hat{x}^{(k)}+\beta^{(k)}\tag{3} $$其中$y^{(k)}$就是最后得到的标准化结果，而$\gamma^{(k)}$和$\beta^{(k)}$也会随着网络中的权重参数一起训练，当且仅当 $\gamma^{(k)}=\sqrt{Var[x^{(k)}]}$，$\beta^{(k)}=E[x^{(k)}]$ 时，公式(3)就变成了恒等变换，也就相当于没有进行标准化（如果网络确实需要的话）。
从式(2)可以看出，Batch Normalization 的核心思想依旧是式 (1) 中的计算过程，仅仅只是加入了对其它情况的优化和处理。
详细原理可参见深度学习栏目「第6.3节 BatchNorm原理：批归一化为什么能加速训练」 内容。
1.3 层归一化#
由于批归一化的核心思想是以一个小批量数据样本为单位在对应维度上进行标准化，因此批量归一化会受到小批量样本数量的影响，同时，显而易见批归一化也不能直接用于循环神经网络。在这样的背景下，层归一化（Layer Normalization）便应运而生了。
可以看出，层归一化的出现的核心目的就是为了解决批归一化无法对变长序列进行标准化的问题。
如图2所示为文本序列批归一化序列图，其$x$轴、$y$轴和$z$轴分别表示小批量样本数量、词向量维度和样本序列长度，即图2中有5个样本，从左到右其长度分别的5、2、4、3和6，且每个词的维度为4维。如果此时对于整个小批量样本同时在$z$轴这个维度上采用批归一化进行标准化，那么在对进行第3个词的位置进行标准化时第2个样本便不会参与，进一步在对第6个词的位置进行标准化时便只有第5个词参与，而这便会影响整个归一化结果。退一步来说，即便通过这样的方式进行标准化，那么当测试样本的长度大于6时，该位置便没有对应的批归一化模型参数。
如图3所示，左右两边为分别是批归一化和层归一化的归一化维度示意图，其一共包含有2个样本和3个特征通道。从图3(a)可以看出，批归一化在进行标准化时是将所有小批量样本看成的一个整体，并逐一对每个通道进行标准化，且每个通道有各自独立的均值$\mu_i$、方差$\sigma^2_i$和权重参数$\gamma_i,\beta_i$（4个值均为标量）；对于图3(b)中的层归一化来说，它在进行标准化时是将每个样本看做一个整体，并同时对所有通道进行标准化，且每个样本均有各自独立的均值$\mu_i$和方差$\sigma^2_i$（2个值均为标量）但共享权重参数$\gamma,\beta$（2个值均为向量，维度为$p\times q\times c$，分别表示特征图的长、宽和通道数）。
因此，层归一化的核心计算逻辑依旧是式(1)当中的计算过程，只是选择了在不同维度上进行。当然，以及后来出现的组归一化也是同样的计算逻辑。
详细原理可参见深度学习栏目「第6.4节 LayerNorm原理：层归一化与 BatchNorm 的区别」 内容。
2. RMSNorm 动机#
尽管 Layer normalization (LayerNorm) 在深度学习中得到了广泛的应用，例如计算机视觉、语音识别和自然语言处理等，但由此却带来了一定量的计算开销，尤其在 RNN 等对时间延迟敏感的模型中影响更大。
2.1 LayerNorm 收敛快但更耗时#
虽然在一些简单的浅层神经网络中这部分开销可以忽略不计，但是随着网络的加深，尤其是在后续大模型的推理场景中这部分的计算开销却是巨大的。
However, the computational overhead introduced by LayerNorm makes these improvements expensive and significantly slows the underlying network, e.g. RNN in particular. Although this is negligible to small and shallow neural models with few normalization layers, this problem becomes severe when underlying networks grow larger and deeper.
也就是说，虽然 LayerNorm 使得模型在训练时更加稳定和高效，但是由此却带来了时间成本的增加。LayerNorm 可以减少训练的步数，让训练更稳定，但这些方法往往让“每一步”的计算量变大。结果是省下来的时间又被额外的计算成本吃掉了，所以从整体上看，效率提升并不明显。
the efficiency gain from faster and more stable training (in terms of number of training steps) is counter-balanced by an increased computational cost per training step, which diminishes the net efficiency, as show in Figure 1.
如图4(a)所示，在相同的训练步数下，不使用任何标准化操作的 Baseline 和 使用层归一化的 LayerNorm 达到的损失分别是 7.0 和 5.4，差距是 1.6；但是在图4(b)中，相同时间下当 Baseline 损失下降到 7.0 的时候，LayerNorm的损失才到5.9，差距只有1.1。
什么意思呢？例如 Baseline 训练300个 step，耗费了 10分钟，损失降到了7.0； LayerNorm 不考虑时间训练同样 300 个step 损失能降到 5.4（可能需要60分钟），但是如果只让其训练10分钟损失却只能降到5.9。
这说明了什么呢？
说明了：① 从“训练步数”角度看效果，LayerNorm 每一步效果好，收敛更快。② 虽然 LayerNorm 训练更稳定，但每一步更慢，所以在同样时间内，训练的步数要少于 Baseline。
作者在论证完 LayerNorm 确实增加了一定量的计算开销以后，又开始进一步论证层归一化中起主要作用的部分，目的就是为了把既耗时又不起多大作用的部分给去掉。
2.2 LayerNorm 中的核心部分#
根据式(2)可知，LayerNorm 归一化主要包含两部分计算：① 对特征进行去均值化操作；② 然后再对去均值后的特征进行缩放操作。前者是为了让模型对输入或权重中的偏移噪声不敏感，因为 LayerNorm 会自动减去均值把它“拉回”到均值为 0 的中心；后者则能确保当输入和权重被随机放大或缩小时，模型输出尽可能保持不变，例如输入被乘以一个系数（比如 10），归一化后仍然回到统一尺度，从而使输出稳定。
A well-known explanation of the success of LayerNorm is its re-centering and re-scaling invariance property. The former enables the model to be insensitive to shift noises on both inputs and weights, and the latter keeps the output representations intact when both inputs and weights are randomly scaled.
虽然这是一个学界普遍认为的什么层归一化有作用的核心原因，但是作者就是要 argue，就是不听。
没办法，奈何作者有理有据有效果。
当事实、结果摆在眼前的时候，你需要的就是找到一个合理的动机和视角把故事讲圆，这也是大家在写作论文的时候需要学习的点。
作者认为，能够使得 LayerNorm 有效的原因仅有缩放操作这一步，因为去中心化这一步并不会实际降低隐藏状态或梯度的方差。
基于这样的动机，论文作者提出了一种新的归一化方法：均方根归一化（Root Mean Square layer Normalization, RMSNorm），它使用均方根（RMS）值来对神经元输入的总和进行标准化，也就是去掉 LayerNorm 中“减去均值”这一步，只保留“缩放”这一步。
In this paper, we hypothesize that the re-scaling invariance is the reason for success of LayerNorm, rather than re-centering invariance.
当然，对于上面提到的归一化操作，我的观点是两个操作都重要，但是起更大作用的应该是缩放操作，因为必定乘以一个数远比加上一个数带来的影响更大。如果去掉作用小的部分，即提高了计算效率也能使得模型的精度损失没有那么的大，那去掉这一步也未尝不可，这其实和前面两篇文章介绍到的模型蒸馏的想法类似。
最后，作者在多种任务上对模型进行了全面评估，包括机器翻译、图像分类、图文检索以及问答任务。实验结果表明，在各种模型中，RMSNorm 的性能与 LayerNorm 相当，但在运行速度方面具有优势，提速幅度从 $7\%$ 到 $64\%$ 不等。这说明 RMSNorm 并不会明显牺牲模型效果，但显著加快了训练或推理过程，尤其在计算资源紧张或需要加速推理时非常有用。
所以，这可能也是为什么 RMSNorm 能在大模型中被广泛使用的原因。
3. RMSNorm 原理#
整体来看，RMSNorm 就是在 LayerNorm 的基础上，去掉了所有需要去均值的部分，因此计算均值这部分操作也就一并去掉了。同时，作者探索了只在部分输入上估计 RMS 的可行性，即只在前 $p\%$ 的输入上计算均方根进一步节省计算开销。所以，这也引出了一个更轻量的变种：partial RMSNorm。
我们觉得，后面这部分内容应该只是作者为了让去掉去均值化这一步更有说服力所作出的探索，算是一脉相承。因为从后续各个论文中使用 RMSNorm 时的代码，以及 PyTorch 2.4 版本后加入的 RMSNorm 模块实现来看，都没有考虑 partial RMSNorm 这样的情况。
3.1 标准化计算过程#
在介绍 RMSNorm 的计算过程之前，我们先来看一下原始标准化的计算过程。
以一个简单的前馈神经网络为例，给定输入为 $X\in \mathcal{R}^m$，经过一个网络层以后的输出为 $y\in\mathcal{R}^n$，其完整计算过程可以表示为
$$ a_i = \sum_{j=1}^mw_{ij}x_j,\;y_i=f(a_i+b_i)\tag{4} $$其中，$w_i$ 是一个权重向量，$b_i$ 是对应的偏置，$f(\cdot)$ 表示非线性变换，$a\in\mathcal{R}^n$ 表示输入 $X$ 对应的线性组合。
随着网络的加深，式(4) 中的计算过程会出现 internal covariate shift 的问题，因此便出现了批归一化（当然也包括后面的层归一化、组归一化，只是在归一化维度上有所区别，计算过程一样）。具体计算过程如下
$$ \overline{a}_i=\frac{a_i-\mu}{\sigma}g_i,\quad y_i=f(\overline{a}_i+b_i)\tag{5} $$其中 $\overline{a}_i$ 表示向量 $\overline{a}$ 中第 $i$ 个值，即 $a_i$ 标准化后的结果；$g\in\mathcal{R}^n$ 表示可学习的权重向量，也即式(3)中的 $\gamma$；$\mu$ 和 $\sigma^2$ 分别表示 $a$ 对应的均值和方差，其计算过程为
$$ \mu=\frac{1}{n}\sum_{i=1}^na_i,\quad\sigma=\sqrt{\frac{1}{n}\sum_{i=1}^n(a_i-\mu)^2}\tag{6} $$3.2 RMSNorm 计算过程#
对于 RMSNorm 来说，其标准化过程就是去掉了式(5)和式(6)中的去均值操作，具体为
$$ \overline{a}_i = \frac{a_i}{\text{RMS}(a)}g_i,\quad\text{RMS}(a)=\sqrt{\frac{1}{n}\sum_{i=1}^na^2_i}\tag{7} $$同 LayerNorm 的计算过程相比较，RMSNorm 少计算了一次均值和去均值的计算过程，而这也就大大减少了 LayerNorm 中对应部分的计算开销。
进一步，作者在此基础上又提出了更加轻量级的 partial RMSNorm 归一化方法，也就是在计算均方根的时候只考虑前 $p\%$ 的输入值，即
$$ \overline{\text{RMS}}(a) = \sqrt{\frac{1}{k}\sum_{i=1}^ka^2_i},\quad k= \lceil n\cdot p \rceil\tag{8} $$其中 $p$ 为超参数，$n$ 为向量维度，$\lceil \cdot \rceil$ 表示向上取整。
不过从后续其它地方的实现来看，partial RMSNorm 几乎都没有被用到，因此，整个式(7)的计算过程便是 RMSNorm 归一化的核心部分。
3.3 RMSNorm 从零实现#
尽管 RMSNorm 在2019年就已经提出来了，可 PyTorch 先前硬是没有支持，直到从 PyTorch 2.4 版本以后才开始支持了这一模块。当然，一个重要的原因我们猜测可能是因为 RMSNorm 提出之初大模型的发展还没有那么迅速、也没那么大，以至于都没太关注归一化层在时间上的开销。现在 RMSNorm 被各大模型广泛使用，自然也就加了进来。
如果你使用的是 PyTorch 2.4 及以后的版本，那么可以通过如下方式来使用:
1 import torch
2 import torch.nn as nn
3
4 if __name__ == '__main__':
5 batch, sentence_length, embedding_dim = 2, 4, 3
6 embedding = torch.tensor([
7 [[1, 1, 2], [2, 0, 0], [0, 0, 1.], [0, 1, 0]],
8 [[1, 0, 1], [1, 1, 2], [1, 0, 2], [0, 2, 2]]])
9 # layer_norm = nn.LayerNorm(embedding_dim)
10 # print(layer_norm(embedding))
11
12 rms_norm = nn.RMSNorm(embedding_dim)
13 print(rms_norm(embedding))可以发现，其使用方法和之前介绍的 nn.LayerNorm 一模一样。
如果你使用的是小于 PyTorch 2.4 的版本，那么首先可以通过如下代码（摘自 LLaMA 模型的实现 [2] ）来进行实现，然后进行使用。
1 class RMSNorm(torch.nn.Module):
2 def __init__(self, dim: int, eps: float = 1e-6):
3 super().__init__()
4 self.eps = eps
5 self.weight = nn.Parameter(torch.ones(dim))
6
7 def _norm(self, x):
8 a = x.pow(2).mean(-1, keepdim=True)
9 b = x * torch.rsqrt(a + self.eps)
10 return b
11
12 def forward(self, x):
13 output = self._norm(x.float()).type_as(x)
14 return output * self.weight在上述代码中，第2行中 dim 用于指定归一化部分向量的维度，eps 是平滑项防止式(7)中分母为0溢出。第5行是式(7)中对应的可学习参数 $g$。第8行则是计算 $\text{RMS}$ 的过程。第9行是计算归一化缩放后的结果，其中 torch.rsqrt() 等价于 1/torch.sqrt()。第12~14行则是依次计算最终归一化的结果。
最后，使用方式如下：
1 if __name__ == '__main__':
2 rms_norm = nn.RMSNorm(embedding_dim)
3 print(rms_norm(embedding))
4
5 rms_norm = RMSNorm(embedding_dim)
6 print(rms_norm(embedding))上述代码运行结束后的输出结果如下：
tensor([[[0.7071, 0.7071, 1.4142],
[1.7321, 0.0000, 0.0000],
[0.0000, 0.0000, 1.7321],
[0.0000, 1.7321, 0.0000]],
[[1.2247, 0.0000, 1.2247],
[0.7071, 0.7071, 1.4142],
[0.7746, 0.0000, 1.5492],
[0.0000, 1.2247, 1.2247]]], grad_fn=<MulBackward0>)
tensor([[[0.7071, 0.7071, 1.4142],
[1.7321, 0.0000, 0.0000],
[0.0000, 0.0000, 1.7320],
[0.0000, 1.7320, 0.0000]],
[[1.2247, 0.0000, 1.2247],
[0.7071, 0.7071, 1.4142],
[0.7746, 0.0000, 1.5492],
[0.0000, 1.2247, 1.2247]]], grad_fn=<MulBackward0>)4. RMSNorm 有效性分析#
在介绍完 RMSNorm 的原理和使用方法以后，我们再来简单看一下作者是如何来阐述和分析 RMSNorm 到底为什么有效的。
根据式(7)我们可以得出，RMSNorm 本质上是将原始结果 $a_i$ 缩放到了一个以 $\sqrt{n}$ 为半径的单位球面上，这保证了无论原始输入有多大，输出都将会被缩到在一个固定尺度内。通过这种方式，输出的分布在输入和权重分布缩放时保持不变，从而有助于层激活值的稳定性。
为什么呢？
由式(7)可知，$a_i$ 经RMSNorm 归一化后的结果为 $\overline{a}_i$ 。此时，我们再计算 $\overline{a}_i$ 对应的均方根
$$ \text{RMS}(\overline{a}_i)=\sqrt{\frac{1}{n}\sum_{i=1}^n\overline{a}^2_i}=\sqrt{\frac{1}{n}\sum_{i=1}^n\left(\frac{a_i}{\text{RMS(a)}}\right)^2}=\frac{1}{\text{RMS(a)}}\cdot \text{RMS(a)}=1\tag{8} $$因为 $g$ 是可学习权重参数，不影响整体性质，所以可以忽略。
根据式(8)可以知道
$$ \text{RMS}(\overline{a}_i)=\sqrt{\frac{1}{n}\sum_{i=1}^n\overline{a}^2_i}=1\tag{9} $$即
$$ \sqrt{\sum_{i=1}^n\overline{a}^2_i}=\sqrt{n}\tag{10} $$这便得出了 $a_i$ 被缩放到了一个以 $\sqrt{n}$ 为半径的单位球面上。
进一步，作者还探索发现，尽管式(7)中 $\text{RMS}(a)$ 和 $\|a\|$ 仅仅只有 $\sqrt{n}$ 倍的关系，即 $\|a\|=\sqrt{n}\cdot \text{RMS}(a)$ ，但是，如果在式(7)中将 $\text{RMS}(a)$ 替换为 $\|a\|$ 将并不会带来好的效果。作者分析认为，这是由于以输入向量的维度 $n$ 来缩放球面能让归一化在处理不同长度的向量时更具鲁棒性。
这一点其实类似于多头注意力计算的时，除以 $\sqrt{d}$ 的作用。
RMS measures the quadratic mean of inputs, which in RMSNorm forces the summed inputs into a $\sqrt{n}$-scaled unit sphere. By doing so, the output distribution remains regardless of the scaling of input and weight distributions, benefiting the stability of layer activations. Although Euclidean norm which only differs from RMS by a factor of $\sqrt{n}$ has been successfully explored, we empirically find that it does not work for layer normalization. We hypothesize that scaling the sphere with the size of the input vector is important because it makes the normalization more robust across vectors of different size.
到此，对于整个 RMSNorm 部分的内容就介绍完了，一句话总结RMSNorm 就是在 LayerNorm 基础上移除了去均值操作，以此来降低模型在训练和推理时的计算开销！
本次内容就到此结束，感谢阅读。青山不改绿水长流，我们月来客栈见！
引用#
[1] Zhang B, Sennrich R. Root mean square layer normalization[J]. Advances in Neural Information Processing Systems, 2019, 32.
[2] https://github.com/meta-llama/llama3/blob/main/llama/model.py