从 GPT-2 到 LLaMA · EP01｜RMSNorm：简化归一化

学完 GPT-2 之后，第一个要换的零件：LayerNorm → RMSNorm。一页纸就能手算完，适合作为"进阶篇「从 GPT-2 到 LLaMA」"的开篇。

一、为什么先看归一化

GPT-2 和今天主流的大模型，主体结构其实是同一个东西：一层层 Transformer Block 堆叠。

变化集中在少数几个零件上。这个系列要做的事，就是一个零件一个零件地看：

GPT-2 里原来是什么
现在换成了什么
为什么换

第一个零件选归一化。原因很简单：它最小，改动最少，是上手"零件替换"最好的练习。

你在 GPT-2 里用的是 LayerNorm（回顾推理篇 EP08 LayerNorm 那一期）。LLaMA 这一系列的模型，把它换成了 RMSNorm。

二、先回顾：LayerNorm 做了什么

对一个 token 的向量 $x$ （长度为 $d$ ），LayerNorm 做三步：

\mu = \frac{1}{d}\sum_{i=1}^{d} x_i \qquad \sigma^2 = \frac{1}{d}\sum_{i=1}^{d} (x_i - \mu)^2

\text{LayerNorm}(x)_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma_i + \beta_i

拆开看，它同时做了两件事：

平移（re-centering）：减去均值 $\mu$ ，让向量以 0 为中心
缩放（re-scaling）：除以标准差，让向量的"大小"统一

再用 $\gamma$ 和 $\beta$ 做一次可学习的缩放和平移。

三、问题：两件事都有必要吗？

RMSNorm 这篇论文（Zhang & Sennrich, 2019）提出的假设是：

LayerNorm 起作用的主要是缩放，平移不是必须的。

所以直接把平移去掉：不算均值，也不加 $\beta$ 。

需要说明的是，这是一个经验性的发现，论文是通过实验验证"去掉平移后效果相近、计算更省"，并不是严格的数学证明。它之所以被广泛采用，是因为后来很多大模型实际用起来效果没有变差。

四、RMSNorm 的公式

只保留缩放，而且缩放的基准从"标准差"换成"均方根（Root Mean Square）"：

\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}

\text{RMSNorm}(x)_i = \frac{x_i}{\text{RMS}(x)} \cdot \gamma_i

对比一下：

	LayerNorm	RMSNorm
要算的统计量	均值 + 方差	只有均方
是否减均值	是	否
可学习参数	$\gamma$ 、 $\beta$ （共 $2d$ 个）	只有 $\gamma$ （共 $d$ 个）

形状标注（以 猫吃鱼 为例，batch=1，3 个 token，d_model=4）：

x            : (1, 3, 4)
x²           : (1, 3, 4)
mean(x²)     : (1, 3, 1)   ← 沿最后一维求平均，keepdim
RMS(x)       : (1, 3, 1)
x / RMS(x)   : (1, 3, 4)   ← 广播
γ            : (4,)
输出          : (1, 3, 4)

五、纸上手算：猫吃鱼

假设三个 token 经过前面层之后，向量如下（d_model=4， $\gamma$ 取初始值全 1， $\epsilon$ 忽略不计）：

猫 = [ 1,  2,  3,  4]
吃 = [ 2,  0, -2,  4]
鱼 = [-1,  1, -1,  1]

注：这里为了好算，用的是示意数字，不是上一系列里猫吃鱼的真实向量。如果想和之前保持一致，可以把你原来的向量代进去，算法完全一样。

1. 猫 = [1, 2, 3, 4]

LayerNorm

均值 $\mu = (1+2+3+4)/4 = 2.5$
偏差 $= [-1.5, -0.5, 0.5, 1.5]$
方差 $\sigma^2 = (2.25 + 0.25 + 0.25 + 2.25)/4 = 1.25$ ，标准差 $\approx 1.118$
结果 $\approx [-1.342,\ -0.447,\ 0.447,\ 1.342]$

RMSNorm

均方 $= (1 + 4 + 9 + 16)/4 = 7.5$
RMS $= \sqrt{7.5} \approx 2.739$
结果 $\approx [0.365,\ 0.730,\ 1.095,\ 1.461]$

2. 吃 = [2, 0, -2, 4]

LayerNorm

均值 $\mu = 1$ ，偏差 $= [1, -1, -3, 3]$
方差 $\sigma^2 = (1+1+9+9)/4 = 5$ ，标准差 $\approx 2.236$
结果 $\approx [0.447,\ -0.447,\ -1.342,\ 1.342]$

RMSNorm

均方 $= (4 + 0 + 4 + 16)/4 = 6$
RMS $= \sqrt{6} \approx 2.449$
结果 $\approx [0.816,\ 0,\ -0.816,\ 1.633]$

注意看：原向量里的 0，在 LayerNorm 之后变成了 -0.447（被平移了），而 RMSNorm 之后依然是 0。

3. 鱼 = [-1, 1, -1, 1]

LayerNorm

均值 $\mu = 0$ ，方差 $= 1$ ，标准差 $= 1$
结果 $= [-1, 1, -1, 1]$

RMSNorm

均方 $= 1$ ，RMS $= 1$
结果 $= [-1, 1, -1, 1]$

两者完全一样。

汇总

token	原向量	LayerNorm	RMSNorm
猫	[1, 2, 3, 4]	[-1.34, -0.45, 0.45, 1.34]	[0.37, 0.73, 1.10, 1.46]
吃	[2, 0, -2, 4]	[0.45, -0.45, -1.34, 1.34]	[0.82, 0, -0.82, 1.63]
鱼	[-1, 1, -1, 1]	[-1, 1, -1, 1]	[-1, 1, -1, 1]

从这三行能读出两个结论：

当向量的均值本来就接近 0 时，LayerNorm 和 RMSNorm 几乎没有区别（鱼）。因为减掉的均值本来就是 0。
均值不为 0 时，LayerNorm 会整体平移，RMSNorm 保留原来的相对结构（猫、吃）。

这也是为什么"去掉平移"有可能行得通：很多时候平移这一步贡献有限。

六、一点几何直观

两种归一化的输出，长度（范数）其实是一样的。

以刚才的结果验算：三个 token 的输出，各分量平方和都是 4，也就是 $d$ 。

用几何语言说：

RMSNorm：把向量缩放到半径为 $\sqrt{d}$ 的球面上。
LayerNorm：先把向量投影到"各分量之和为 0"的超平面上（这就是减均值），再缩放到同一个球面。

所以 LayerNorm 比 RMSNorm 多了一步投影。RMSNorm 的假设就是：这一步投影，可以不要。

七、类比：调音量

可以把一个 token 向量想象成一段声音：

RMSNorm：只调音量，把响度统一到同一档，声音的波形不变。
LayerNorm：除了调音量，还把声音整体的"直流偏置"（整体偏高或偏低的部分）也去掉。

多数时候，直流偏置对"听起来是什么"影响不大，所以只调音量也够用，而且更省事。

八、PyTorch 实现

import torch
import torch.nn as nn


class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))  # γ，没有 β

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        mean_sq = x.pow(2).mean(dim=-1, keepdim=True)    # (batch, seq_len, 1)
        rms = torch.sqrt(mean_sq + self.eps)             # (batch, seq_len, 1)
        return x / rms * self.weight                     # (batch, seq_len, d_model)

用刚才手算的例子验证：

x = torch.tensor([[[ 1.,  2.,  3., 4.],
                   [ 2.,  0., -2., 4.],
                   [-1.,  1., -1., 1.]]])   # (1, 3, 4)

norm = RMSNorm(4)
print(norm(x))

输出应该和上面手算的结果一致（保留几位小数后）：

[[[ 0.3651,  0.7303,  1.0954,  1.4606],
  [ 0.8165,  0.0000, -0.8165,  1.6330],
  [-1.0000,  1.0000, -1.0000,  1.0000]]]

替换到 GPT-2 里：把 Block 中所有的 nn.LayerNorm(d_model) 换成 RMSNorm(d_model)，包括最后输出层前的那一个归一化。其他代码不用动。

两个小提示：

较新版本的 PyTorch 已经内置了 nn.RMSNorm，具体从哪个版本开始，以你本地版本的文档为准。自己手写一遍，主要是为了看清楚它在算什么。
参数量几乎没有变化。以 GPT-2 small 为例（d=768，12 层，每层 2 个归一化，再加最后 1 个，共 25 个），每个少 768 个 $\beta$ ，一共只少约 1.9 万个参数，相对 1 亿多参数可以忽略。省的主要是计算（少算一次均值，少一次减法），而不是参数。

九、小结

	GPT-2	现代大模型（如 LLaMA 系列）
归一化方式	LayerNorm	RMSNorm
做的事情	平移 + 缩放	只缩放
可学习参数	γ、β	只有 γ
依据	—	经验发现：平移不是关键，去掉后效果相近、计算更省

一句话总结：RMSNorm 就是 LayerNorm 去掉"减均值"和 β，只保留缩放。

十、延伸阅读

Zhang & Sennrich, Root Mean Square Layer Normalization（2019）：RMSNorm 原论文，关于"平移不重要"的假设和实验结果，以原文为准。
Touvron 等，LLaMA（2023）：把 RMSNorm、SwiGLU、RoPE 合到一起用的代表。

下期预告

下一期是 EP02 · SwiGLU：升级前馈网络。归一化之后，我们把 GPT-2 的 GELU FFN 换成门控前馈网络，用“猫吃鱼”手算两路投影如何协作。RoPE 将在 EP03 介绍。

本系列适合已经学完 GPT-2 结构的读者。如果文中公式有任何疑问，欢迎在评论区留言。