多头注意力和位置编码

上一课写出了一个注意力头。要用它搭一个 GPT，还差两样东西：多头，以及位置信息。

python heads_positions.py

这一课的实验用的是 gpt.py 里的模型，但不需要训练，看的是结构本身的性质。

多头：同时看好几样东西

一个字要从前面收集的信息往往不止一种。写"白日依山尽"的"尽"时，模型可能需要同时知道：前面是什么景物（白日、山），这是第几个字（五言诗到第五个字该押韵或断句），上一句用了什么词。

一个注意力头只有一套 q 和 k，只能按一种标准打分，得出一套权重。多头注意力的办法是：把向量切成几段，每段各自做一次注意力，最后再拼回来。每个头有自己的 q、k、v，可以学会关注不同的东西。

== 2. 多头注意力：32 维的向量切成 4 个头，每个头 8 维
  q 的形状 (1, 5, 32)
  拆成多个头之后 (1, 4, 5, 8)：(句子数, 头数, 字数, 每个头的维度)
  注意力分数 (1, 4, 5, 5)：每个头都有自己的一张 5×5 的表
  一个注意力层的参数：qkv 3168 个，proj 1056 个，头数多少都不影响这个数

在 gpt.py 里，实现只是多了几次改变形状：

q, k, v = self.qkv(x).split(C, dim=2)
# 拆成多个头：(B, T, C) -> (B, 头数, T, 每个头的维度)
q, k, v = (t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) for t in (q, k, v))

self.qkv 是一个线性层，一次算出 q、k、v 三个向量，再用 split 分开。然后把 32 维的向量看成 4 段、每段 8 维，并把"头"这个维度挪到前面。之后的打分、掩码、softmax、加权求和，和上一课一模一样，只是 PyTorch 的矩阵乘法会自动对每个头分别计算，一次算完 4 个头。

算完之后，把 4 个头的结果拼回 32 维，再经过一个线性层 proj 混合一下，让各个头收集到的信息能互相结合：

out = out.transpose(1, 2).contiguous().view(B, T, C)  # 各个头拼回去
return self.proj(out)

注意最后一行输出：参数的数量和头数无关。4 个头每头 8 维，和 1 个头 32 维，用的是同样大小的矩阵，只是切分的方式不同。多头不增加参数，也几乎不增加计算量，却让模型能同时学几种不同的关注方式。

我们的 GPT 用 4 个头，每个头 32 维（向量总共 128 维）。大模型通常有几十个头，每个头 64 或 128 维。

注意力分不清顺序

回头看上一课的注意力公式：打分只看 q 和 k 的内容，加权求和也只看 v 的内容，没有任何地方用到"这个字在第几个位置"。

这意味着，对注意力来说，"白日依山尽"和"山依日白尽"可能没有区别。

实验：让模型读这两句，它们最后一个字都是"尽"，前面四个字换了顺序。比较模型在最后一个位置的输出（也就是对下一个字的打分）差多少：

model.pos_emb.weight.zero_()  # 把位置嵌入清零，等于没有位置信息
a = model(torch.tensor([tok.encode("白日依山尽")]))[0][0, -1]
b = model(torch.tensor([tok.encode("山依日白尽")]))[0][0, -1]

== 1. '白日依山尽' 和 '山依日白尽'，最后一个位置的输出差多少
  1 层：没有位置嵌入 3.7e-08，有位置嵌入 2.8e-03
  2 层：没有位置嵌入 1.9e-02，有位置嵌入 1.5e-02

先看第一行，只有 1 层的模型。没有位置信息时，两句的输出差别是 3.7e-08，也就是只剩浮点数的舍入误差，完全一样。对最后一个位置来说，它面前是同样的五个字，只是顺序不同。注意力对它们打分、加权求和，结果和顺序无关。

这对语言模型是致命的。"白日依山尽"是诗，"山依日白尽"不通；"我打你"和"你打我"意思相反。模型必须知道顺序。

意料之外的第二行

第二行是 2 层的模型。按上面的道理，没有位置嵌入时两句的输出也该一样。可实验结果是差了 1.9e-02，和有位置嵌入时的 1.5e-02 差不多大。

我写这个实验时也没料到。原因出在因果掩码上：

• 第 1 层里，每个位置只能看到自己和前面的字。"白日依山尽"的第 2 个位置看到的是"白日"，"山依日白尽"的第 2 个位置看到的是"山依"。所以第 1 层在第 2、3、4 个位置的输出，两句是不一样的。
• 第 2 层里，最后一个位置去看前面各个位置第 1 层的输出，而这些输出已经不一样了。

也就是说，因果掩码本身就泄露了顺序：一个位置能看到几个字，间接告诉了它自己在第几个位置。只要有两层以上，模型就能从中得到一些位置信息。

这不是我们代码的偶然现象。2022 年有一篇论文专门研究了这件事（Haviv 等，《Transformer Language Models without Positional Encodings Still Learn Positional Information》），发现不加任何位置编码的因果语言模型，也能学到位置信息，效果只比加了位置编码的差一点。

不过，靠掩码"猜"出来的位置信息是间接的、模糊的。实际的模型都会明确地加入位置信息。

位置嵌入

最简单的办法，GPT-2 用的就是它：给每个位置也准备一个向量，和字的向量相加。

self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
...
pos = torch.arange(start, start + T, device=idx.device)
x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))

pos_emb 是一张 128 行的表，第 0 行是第 0 个位置的向量，第 1 行是第 1 个位置的，以此类推。这些向量和其他参数一样，通过训练学到。

相加之后，同一个字在不同位置的向量就不一样了。"白"在第 1 个位置和第 4 个位置，得到的是两个不同的向量，注意力就能区分顺序。上面的实验里，1 层的模型加上位置嵌入后，两句的输出差别从 3.7e-08 变成了 2.8e-03：模型能分辨这两句了（这个模型还没训练，所以差别不大，也还没有意义）。

这种做法有一个限制：表只有 128 行，模型最多只能处理 128 个位置。训练时没见过的位置，它没有对应的向量。

旋转位置编码

截至 2026 年 9 月，主流的开源大模型（Llama、通义千问、DeepSeek 等）大多不用 GPT-2 那种位置嵌入，而是用旋转位置编码（RoPE，出自 2021 年的论文 RoFormer）。

它的想法是：不把位置加到字的向量上，而是在计算注意力分数之前，把 q 和 k 按照各自的位置"旋转"一个角度。位置越靠后，转得越多。这样一来，两个位置的 q 和 k 做点积时，结果只和它们之间相隔多远有关，而和它们各自的绝对位置无关。

"月"在"明"后面一个字，不管这两个字出现在一首诗的开头还是结尾，它们之间的关系都是一样的。相对位置比绝对位置更符合语言的规律，而且更容易推广到比训练时更长的文字上。

这门课的 GPT 保持用最简单的位置嵌入，因为它的效果在我们的实验里完全够用，而且好理解。练习 3 可以动手试试 RoPE。

练习

1. 在 heads_positions.py 里再加一组：3 层的模型，没有位置嵌入时，两句的输出差多少？
2. 把 gpt.py 训练用的模型从 4 个头改成 1 个头、8 个头（n_embd 保持 128），用下一课的 train.py 各训练 1000 步，比较验证损失。
3. 挑战：参考 RoFormer 论文或者一个开源模型的代码，给 CausalSelfAttention 加上旋转位置编码，并去掉 pos_emb，看看训练的效果。

自测

1. 多头注意力是怎么做的？它会增加参数吗？

把每个字的向量切成几段，每段各自做一次注意力（各有自己的 q、k、v 和权重），再把结果拼回来，经过一个线性层混合。参数的数量和只用一个头时一样，只是矩阵的切分方式不同，好处是每个头可以学会关注不同的东西。

2. 为什么注意力需要位置信息？

注意力的打分和加权求和只看向量的内容，不看位置。对只有一层的模型，把前面的字换个顺序，最后一个位置的输出完全不变。而语言的意思依赖顺序，所以要把位置信息加进去。

3. 实验里，2 层的模型没有位置嵌入，为什么也能分辨两句的顺序？

因果掩码让每个位置只能看到自己和前面的字，第一层里不同位置看到的字数不同，输出也就和顺序有关了。第二层再读取这些输出时，就间接得到了位置信息。有研究发现，不加位置编码的因果语言模型也能学到位置信息，但实际的模型仍然会明确地加入位置编码。