Step by Step 读懂 LLM（二）：LayerNorm、Self-Attention 与 MLP

目前为止，我们的 input string 经过了 BPE 和 embedding，得到了一个 t×d 的矩阵 X（上一篇里特征维度记作 b，从这篇开始按惯例记作 d）。其中每个横行代表一个 token 的 feature，整个矩阵代表了我们输入的这一段话。

接下来就是 Transformer 的 block，它包括 LayerNorm（近几年的模型多用 RMSNorm）、Self-Attention 和一个两层的 MLP。这几部分是这样连起来的：

注意这里其实有两个 Norm：一个在 Attention 之前，一个在 MLP 之前。这种“先 Norm 再进子层”的写法叫 pre-norm，是 GPT-2 之后几乎所有 LLM 的选择；原始 Transformer 是把 Norm 放在残差相加之后（post-norm），层数一深就很难训练^[8]。

1LayerNorm

Normalization 是 machine learning 中常见的手段，常见的主要有 Batch Normalization^[1] 和 Layer Normalization^[2]。

Batch Normalization 是针对每一个 feature，把一个 batch 里所有数据点在这个 feature 上的数值放在一起，求均值和方差，再做 normalization。在 CNN 里，feature 就是 channel，“所有数据点”是 batch 里每张图片的每一个像素位置；如果搬到 LLM 里，就是 batch 里每一段 input 的每一个 token，在词向量同一个位置上的 feature value。

BN 在过去的 CNN 模型中非常常见。人们在训练中 apply BN，用当前 batch 的均值方差做归一化，同时累计计算它的 running mean 和 running var；实际使用（推理）时，直接用训练结束时累计出来的这组 mean 和 var，不再依赖推理时的 batch。BN 之所以在 CNN 里好用，有几个前提：

• 输入尺寸相同。不同的图片通常会被 rescale 到同一尺寸，所以每个 channel 在 batch 里的每张图上都存在，统计量的含义是一致的。
• 数据之间没有关系。训练时 batch 里的每个数据都是不同的图片，打乱之后彼此独立。
• 统计量稳定。每个 channel 的均值方差来自“图片数 × 高 × 宽”个数；加上 CNN 模型相对较小，训练时 batch 可以开得比较大，均值方差就比较稳定，波动小。

这几点在 LLM 中都不成立：

• 长度不一。同一个 batch 里的句子有长有短，要用 padding 补齐。padding 的位置不是真实数据，算统计量时要专门把它们剔除；而且第 500 个位置上的 token，只在长于 500 的句子里才存在，每个位置能凑到的样本数都不一样。
• token 之间不独立，会偷看答案。同一句话里的 token 是一串有先后的序列，后面的词依赖前面的词。BN 把它们放在一起求统计量，第 i 个 token 的结果里就混进了它后面的 token。为什么这等于偷看答案？

  先看语言模型是怎么训练的：在第 i 个位置，根据前 i 个 token 猜第 i+1 个 token。为了效率，训练时一整句话是一次性并行算的：第 1 个位置猜第 2 个词、第 2 个位置猜第 3 个词……同时进行。所以每个位置要猜的“答案”，其实就是同一个矩阵里它下面那一行。后面的 Attention 会用 causal mask 保证第 i 行看不到它下面的行，模型只能老老实实地猜。

  但 BN 的均值，是把这些行加在一起算的。举个最极端的例子：batch 里只有一句话 the cat sat down。BN 之后，每个 feature 在这 4 个 token 上的平均值一定是 0（减掉的就是均值）。那么对 sat 这个位置来说，它要猜的答案 down 归一化后的向量是：

  右边三个向量，sat 通过 Attention 全都看得到（它们都在它前面，causal mask 不挡）。也就是说，模型不用理解任何语言，做一次加法就能把答案算出来：

  	feature 1	feature 2	feature 3
  the	0.23	−1.12	0.37
  cat	1.48	0.59	−0.47
  sat	−1.25	1.33	1.40
  down（答案）	−0.45	−0.80	−1.30
  −(前三行之和)	−0.45	−0.80	−1.30

  往前的位置也一样：cat 能算出“sat 和 down 加起来是多少”，虽然拆不开，也泄露了不少信息。batch 变大之后，每个 token 在均值里只占很小的一份，泄露被稀释了，但只要有一点点信号，梯度下降就会去用它：训练 loss 降得特别快，看起来效果很好。到了真正生成的时候，后面的 token 还没生成出来，统计量里根本没有答案，模型学会的这条“作弊通道”消失了，效果就会崩。

  CNN 没有这个问题，因为一张图片里没有“答案”：分类的标签不在像素里，batch 里别的图片也和这张图的标签无关，混在一起求统计量不会泄露什么。

• 统计量不稳定。文本的激活值在不同 batch 之间波动很大，比图片大得多，这是 BN 在 Transformer 里效果差的主要原因之一^[4]。而且少数 token 的激活值会大得离谱（下面第 4 点），一个这样的 token 就能把整个 batch 的统计量带偏。
• 训练和推理对不上。生成的时候，模型一次只往后算一个新 token。推理只能用训练时累计的 running mean/var，但训练时的统计量又受 batch 里句子长度、组成的影响，推理时遇到的分布很难和它一致。

因此 LLM 更常用的是 Layer Norm：在每个 token 里，对它所有的 feature 求均值和方差做 norm。对一个 token 的向量 ：

其中 、 是可学习的缩放和偏移。RMSNorm^[3] 更简单，连均值都不减，只除以均方根：

那为什么不在同一个句子里、不同的 token 之间做 norm 呢？“跨 token”有两种做法：

• (a)对每个 feature，在整句所有 token 上求均值和方差。这相当于把 BN 用在一句话的内部。
• (b)把整句 t×d 个数放在一起，求一个均值和方差。

下图里点任意一个格子，看在不同的做法下，哪些数会和它一起算均值和方差。

只有 LayerNorm 的统计量完全来自这个 token 自己。跨 token 做 norm 行不通，有这样几个原因：

第一，训练时会偷看未来的 token。上面讲 BN 时那个“做一次加法就能算出答案”的例子，对做法 (a)、(b) 一样成立：做法 (a) 本来就是一句话内部的 BN。只要统计量跨 token，Attention 用 causal mask 小心挡住的未来信息，就会从 norm 这里漏进来。

第二，KV cache 没法做。KV cache 是推理时把前面 token 算好的中间结果（每层的 K 和 V）存起来复用，生成新 token 时就不用重算。可如果统计量跨 token，对话每延续一个 token，均值方差就变了，前面所有 token 归一化之后的值都跟着改，缓存全部作废。

第三，抹掉有用的信息。做法 (a) 会把每个特征减去它在整句上的平均值。但“整句话共同的东西”，比如话题、语言、文风，往往正是模型需要的信息。减掉之后，每个 token 只剩下“和这句话平均水平的差别”。LayerNorm 去掉的则是一个 token 向量自己的整体偏移和长度，保留的是这个向量的方向，也就是各个特征之间的相对大小，而信息主要就编码在方向里。RMSNorm 更直接：它几乎就是把向量缩放到固定长度 ，即 ，方向完全不变。

第四，一个 token 会影响其他所有 token。真实的 LLM 里，不同 token 的激活值差别很大。一个典型现象是：第一个 token（或者某些标点）在少数几个 feature 上的激活值大得离谱，可以比典型的激活值大上千倍，极端的到 10 万倍（massive activations^[5]，和 attention sink^[6] 现象有关）。在做法 (b) 下，这个 token 会把整句的方差撑得很大，其他 token 归一化后就被压得接近 0。用 LayerNorm 时，每个 token 各自缩放，这个异常 token 只影响它自己。

第五，结果依赖句子长度。同一个 token，在 10 个 token 的句子里和在 1000 个 token 的句子里，会被归一化成不同的值（图 3 里每追加一个 token，前面的值都在变，就是这个原因）。变长序列再加上 padding，又会多一层污染。

所以 LayerNorm / RMSNorm 的作用可以总结成一句话：每个 token 只管自己，把自己的向量缩放到稳定的数值范围，不依赖 batch，不依赖句子长度，也不会把别的 token 的信息混进来。token 之间交换信息这件事，全部交给下一步的 Self-Attention。

2Self-Attention

到目前为止，每个 token 的向量只知道它自己：embedding 是按 id 查表，LayerNorm 也是每个 token 各算各的。但理解一句话离不开上下文。想象一下这句话：

The tired cat crossed the street, then it slept.

要理解 it，就得知道它指的是前面的 cat。Self-Attention 做的就是这件事：让每个 token 回头看前面的 token，找到跟自己相关的，把它们的信息混进自己的向量。这是整个 Transformer 里 token 之间唯一交换信息的地方，其他部分（Norm、MLP）都是每个 token 单独算的。

具体怎么做呢？首先，每个 token 的向量分别乘上三个 d×d 的矩阵，得到三个新向量：

和上一篇讲的一样，矩阵乘法 的第 i 行就是 ，所以这一步也是每个 token 各算各的。可以这样理解这三个向量：

• query ：我在找什么。比如 it 在找“前面出现过的、能被代指的名词”。
• key ：我能被什么样的 query 找到。比如 cat 的 key 表示“我是一个动物名词”。
• value ：被找到之后，我交出去的内容。

然后用 query 和 key 做点积来打分： 越大，token i 就越想看 token j。所有的分数排在一起，正好是一个 t×t 的矩阵 。接下来三步：除以 ，盖上 causal mask，每一行做 softmax 变成权重，最后用这组权重把 value 加权平均：

下图把这个过程一步步拆开。

几个细节：

• 为什么除以 。两个 d 维随机向量的点积，方差随 d 线性增长。d 很大时分数会很大，softmax 几乎变成 one-hot，梯度也几乎为 0。除以 把分数拉回到稳定的范围。
• causal mask 就是挡住未来。把 的分数设成 ，softmax 之后权重正好是 0。图里被 mask 的格子并不是分数低：cat 其实很想看后面的 it，但训练时它不能看，因为生成 cat 的时候 it 还不存在。这和上一节 LayerNorm 不能跨 token 是同一个道理。
• 输出还是 t×d。每个 token 得到的是前面所有 value 的加权平均，最后再乘一个输出矩阵 ，加回残差。

多头注意力

上面这个版本有一个限制：每个 token 只有一次 softmax，也就是只有一组权重。但一个 token 往往同时需要好几种信息。还是看 it：它既需要知道“我指的是谁”（cat），也许还需要知道“它当时是什么状态”（tired）。一组权重只能在这两者之间折中，结果是哪个都看不清楚。

多头（multi-head）注意力^[7]的做法是：把 d 维切成 H 段，每段 维，叫一个“头”。每个头用自己那一段的 q、k 打分，各做各的 softmax，再用自己的权重去混合自己那一段的 v，最后把 H 段拼回 d 维，乘上 ：

下图是一个 4 维、两个头的例子。拖动 it 的 query，比较单头和双头。

可以看到，单头和双头用的是同一组 q、k、v，乘法的次数也完全一样，多头并不增加计算量。唯一的区别是 softmax 做了几次：单头把 4 维的分数加在一起做一次 softmax，两个头想找的东西被摊成一个分布；双头分开做，每个头各自变尖，it 的前两维从 cat 那里拿，后两维从 tired 那里拿。

推广到 32 个头，可以这样理解：每个 value 被切成 32 段，每段用它所属那个头的权重单独加权求和，再拼回去。最后的 把这 32 段重新混合起来，所以“第几段来自第几个头”只是中间状态，不会一直保留到输出。

现在的模型还会让几个 query 头共用同一组 K、V（GQA^[12]，例如 Llama 3 是 32 个 query 头共用 8 组 K、V），这样 KV cache 能小好几倍。这个等讲 decode 的时候再细说。

3MLP

Attention 之后，每个 token 的向量已经混进了上下文的信息。接下来的 MLP 又回到了每个 token 单独计算：同一个两层的全连接网络，作用在每一行上。

第一层把 d 维升到 4d 维，经过激活函数 ，第二层再降回 d 维，加回残差。如果说 Attention 是 token 之间“交流”，MLP 就是每个 token 拿着交流来的信息，自己“想一想”。

这里的激活函数，原始 Transformer 用的是 ReLU，GPT-2 换成了更平滑的 GELU。现在的 Llama 等模型用的是 SwiGLU^[10]：第一层拆成两个矩阵，一个过 SiLU 激活，一个不过，两者逐元素相乘，起到“门控”的作用：

多了一个矩阵，所以中间层通常从 4d 缩到大约 ，让参数量和原来差不多（Llama 2 7B 里 d = 4096，中间层是 11008）。结构上它仍然是“升维、激活、降维”的两层网络。

MLP 看起来比 Attention 简单得多，但模型的大部分参数都在这里。一层 Attention 有 四个 d×d 矩阵，共 ；一层标准 MLP 是 加 ，共 ，是 Attention 的两倍。

MLP 占了将近三分之二的参数，所以有一种很流行的理解：MLP 是模型存放“知识”的地方^[13]。第一层 的每一列像一个“模式探测器”，当 token 的向量和它对得上时，对应的神经元被点亮（图 7 里中间层亮起来的那几个）；第二层 的对应那一行，就是这个神经元被点亮时要写回 token 向量的内容。这和 Attention 的 key/value 很像，只不过 key 和 value 不是来自上下文，而是训练时学到、固定在权重里的。

也正因为 MLP 参数最多、又是每个 token 单独计算，MoE（Mixture of Experts）模型就是从这里下手：把一个 MLP 换成很多个“专家” MLP，每个 token 只选其中几个来算，参数量大大增加，但每个 token 的计算量不变。

到这里，一个 Transformer block 就走完了：Norm → Attention → 加回残差 → Norm → MLP → 加回残差。输入是 t×d，输出还是 t×d，每一行仍然对应一个 token，只是已经混进了上下文，又经过了 MLP 的加工。真实的模型把这个 block 堆叠几十层（Llama 2 7B 是 32 层）。

参考

1. Ioffe, Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. ICML 2015.
2. Ba, Kiros, Hinton. Layer Normalization. 2016.
3. Zhang, Sennrich. Root Mean Square Layer Normalization. NeurIPS 2019.
4. Shen, Yao, Gholami, Mahoney, Keutzer. PowerNorm: Rethinking Batch Normalization in Transformers. ICML 2020.
5. Sun, Chen, Kolter, Liu. Massive Activations in Large Language Models. COLM 2024.
6. Xiao, Tian, Chen, Han, Lewis. Efficient Streaming Language Models with Attention Sinks. ICLR 2024.
7. Vaswani et al. Attention Is All You Need. NeurIPS 2017.
8. Xiong et al. On Layer Normalization in the Transformer Architecture. ICML 2020.
9. Elhage et al. A Mathematical Framework for Transformer Circuits. 2021.
10. Shazeer. GLU Variants Improve Transformer. 2020.
11. Touvron et al. Llama 2: Open Foundation and Fine-Tuned Chat Models. 2023.
12. Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023.
13. Geva, Schuster, Berant, Levy. Transformer Feed-Forward Layers Are Key-Value Memories. EMNLP 2021.