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 Norm

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 里句子长度、组成的影响，推理时遇到的分布很难和它一致。

同一个句子里的 token 之间做 norm

既然跨句子做 norm 不行，那只在同一个句子里、不同的 token 之间做 norm 呢？这样不依赖 batch，也碰不到别的句子。具体有两种做法：

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

下图里点任意一个格子，看在不同的做法下，哪些数会和它一起算均值和方差。图里也放上了 LayerNorm 作为对照：它只在每个 token 自己内部做归一化，最后一小节细讲。

可以看到，(a) 和 (b) 虽然只在一句话内部，统计量里还是混进了排在后面的 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 只剩下“和这句话平均水平的差别”。

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

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

Layer Norm：每个 token 只管自己

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

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

和做法 (a) 减掉“整句话的平均”不同，LayerNorm 去掉的只是一个 token 向量自己的整体偏移和长度，保留的是这个向量的方向，也就是各个特征之间的相对大小，而信息主要就编码在方向里。RMSNorm 更直接：它几乎就是把向量缩放到固定长度 ，即 ，方向完全不变。

回头看前面那些问题：依赖 batch 和 padding、偷看未来、KV cache 失效、一个异常 token 拖累整句、结果依赖句子长度，全都来自“统计量跨了句子或者跨了 token”。LayerNorm 的统计量只来自这个 token 自己，这些问题就都没有了。所以 LayerNorm / RMSNorm 的作用可以总结成一句话：每个 token 只管自己，把自己的向量缩放到稳定的数值范围，不依赖 batch，不依赖句子长度，也不会把别的 token 的信息混进来。token 之间交换信息这件事，全部交给下一步的 Self-Attention。

2Self-Attention

讲 Self-Attention 之前，首先要讲清楚什么是 attention。

在 NLP 领域还在用 RNN 的时候，模型是从左到右一个词一个词地读的。每读到一个词 ，就把它和前面压缩好的 hidden state 合在一起，更新出新的 ：

每个位置拿来做预测的，就是这一个固定长度的 ，它由前面所有的词压缩而成。这样哪怕前面有 1000 个词，也会被压成同一个固定长度的 vector 来做预测；而且每读一个新词 都要被更新一次，越早的词被冲淡得越厉害，越近的词影响力越大。对前文信息的压缩非常严重。

于是后来人们开始想：是不是除了压缩完的 ，也可以让模型直接回头看看前文每个位置的内容，再进行处理？这就是最开始的 attention 机制^[14]。它最早用在机器翻译上：一个 RNN（encoder）先读完原句，另一个 RNN（decoder）再一个词一个词地生成译文，生成每个译文词的时候，都回头看一遍原句的每一个位置。具体分三步：

1. 每个位置的内容。首先，每个词都有一个自己输出的 hidden state ：它从这个词本身的词向量出发，结合前文的 hidden state，经过这一步的 RNN 单元处理之后得到。可以把它看作这个位置的结果向量，也是之后每一步要回头看的“前文内容”。（原论文的 encoder 从左往右、从右往左各读一遍，所以 里有这个词两边的上下文。）
2. 打分。接下来需要决定一个权重，也就是给每个词多少 attention。当时的做法是一个小 MLP：输入是当前这一步的向量（decoder 上一步的 hidden state ）和原句某个词的向量 ，各乘一个矩阵后相加，过一个 tanh，再和一个向量 点乘，变成一个一维的数值： 这个过程是当前这一步和原句每个词都做一遍，所以每个词都得到一个“分数”。
3. softmax 加权。紧接着对所有这些分数做 softmax，得到一组加起来等于 1 的权重；然后用这组权重把每个词的向量加起来，得到一个加权向量 ，交给 decoder 去生成这一步的词：

举个例子，把 the cat on the mat 翻译成“垫子上的猫”。中文和英文的语序正好是反的：生成“垫子”的时候，权重应该集中在 mat 上；到了生成“猫”，又要跳回前面的 cat。

把每一步的权重排成一张表，就是原论文里那张有名的“对齐图”：哪个译文词对应哪个原文词，模型自己学了出来，没有人教过它。而且这是“软”的对齐，用 softmax 加权而不是硬选一个位置，所以整个过程可导，可以和翻译模型一起端到端训练。

这个算法后来被延伸成了今天 QKV 版本的 self-attention，逻辑其实还是一样的：

1. 前文词的 value：作为加权的基础值；
2. attention score：作为加权的权重。

还是一样的 pattern，只不过计算方法做了升级。主要升级了三处。

第一，从“译文看原文”变成“一句话看自己”。Self-attention 里没有两句话，每个 token 回头看的是同一句话里它前面的 token（包括它自己）。为什么需要这一步？到目前为止，每个 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 拆成 Q、K、V 三个向量。还是按上面的逻辑来想：value 是一定需要的，它是每个位置交出去的内容；然后我们需要一组给所有 value 的权重，这个权重取决于当前词本身，以及其他每个词的特征。所以每个 token 的向量分别乘上三个 d×d 的矩阵，得到三个新向量：

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

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

那为什么要三个？K 和 V 都长在被看的那个词身上，能不能用同一个？或者干脆 QKV 都用同一个？

• K 和 V 干的是两件事。K 决定“我能不能被找到”，V 决定“被找到之后交出什么”。就像去图书馆，按索引卡上的书名、分类找书，这是 K；找到之后读的是书的正文，这是 V。cat 能被 it 找到，靠的是“名词、动物、单数”这些特征，但 it 真正要拿走的，是“猫”本身的意思，以及它前面被说成 tired。如果 K = V，哪些特征让匹配分数高，哪些特征就会被搬过去，模型没法“按 A 去找、拿走 B”。其实最早的 attention 里，key 和 value 就是同一个 ，只不过打分时乘了 ，decoder 拿到 之后又乘了自己的矩阵，已经部分分开了；Transformer 给 V 也单独一个投影，把这件事做彻底。
• Q 和 K 分开，关系才能是单向的。如果 Q 和 K 用同一个矩阵，i 对 j 的分数就永远等于 j 对 i 的分数。但语言里的关系大多是单向的：it 要找 cat，cat 并不需要找 it。而且一个向量和自己的点积通常最大，每个 token 都会倾向于盯着自己。Reformer 为了省计算共用过 Q 和 K，就得专门加一条规则，禁止 token 注意自己^[16]。
• 都不投影、直接用 x 也不行。那样分数就是两个词的 embedding 有多像，注意力只会落在和自己相似的词上。但 it 和 cat 的 embedding 一点都不像。我们要的是“相关”，不是“相似”。三个投影矩阵，就是让模型自己学“什么叫匹配”、“匹配之后搬什么”。

第三，打分从 MLP 换成点积。 越大，token i 就越想看 token j。点积打分在 2015 年就有人用了^[15]，它最大的好处是所有 (i, j) 的分数可以一次写成一个矩阵乘法 ，正好是一个 t×t 的矩阵，GPU 算起来极快；而 MLP 打分每一对都要单独过一次 tanh。点积的数值会随维度变大，所以要先除以 。接下来和最早的 attention 一样：盖上 causal mask，每一行做 softmax 变成权重，再用这组权重把 value 加权平均：

和最早的 attention 对照一下：

先只看一个词。下图里，it 向前文的每个词（包括它自己）连一条线，线上标出每一步的计算：先用 q·k 打分，再除以 、做 softmax 变成权重，最后用权重把每个词的 v 加起来。

每个词都在做同样的事。把所有词的这一行分数叠在一起，就是一个 t×t 的矩阵 ：第 i 行就是第 i 个词对前文的打分，右上角是“未来”，被 causal mask 挡住。下图用矩阵的视角把整个过程再走一遍。

几个细节：

• 为什么除以 。两个 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 维，乘上 ：

多头到底改变了什么？下图跳过打分，直接看最后一步“用权重把 v 加起来”。把 cat 和 tired 交出的 v 写成一个矩阵：每一行是一个词的 v，每一列是一个维度。为了好看清，假设前两维存的是“是谁”，后两维存的是“什么状态”。格子越亮，说明它最后进入输出的越多。

单头的输出是 v 的整行加权平均：权重是一行一个数，cat 那一行和 tired 那一行都只能整行按同一个比例进来。所以输出只能落在右图那条斜线上：多拿一点 cat，就必然少拿一点 tired。it 两样都想要，最好也只能各拿一半。

多头把列切开，每一段维度有自己的一组权重。头 1 管前两维，可以把权重全给 cat；头 2 管后两维，可以全给 tired。输出就变成“前一段完整地来自 cat，后一段完整地来自 tired”，右上角那颗星终于能到达了。

所以多头并不增加计算量（乘法次数完全一样），它改变的是输出能长成什么样：从“几个 v 的整行平均”，变成“每一段维度可以各自从不同的词那里拿”。反过来，如果两个头想找的是同一个词，或者 it 本来只需要一种信息，切不切开就没什么区别。

推广到 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 的向量和它对得上时，对应的神经元被点亮（图 9 里中间层亮起来的那几个）；第二层 的对应那一行，就是这个神经元被点亮时要写回 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.
14. Bahdanau, Cho, Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015.
15. Luong, Pham, Manning. Effective Approaches to Attention-based Neural Machine Translation. EMNLP 2015.
16. Kitaev, Kaiser, Levskaya. Reformer: The Efficient Transformer. ICLR 2020.