多頭注意力和位置編碼

上一課寫出了一個注意力頭。要用它搭一個 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 層的模型沒有位置嵌入，為什麼也能分辨兩句的順序？

因果掩碼讓每個位置只能看到自己和前面的字，第一層裡不同位置看到的字數不同，輸出也就和順序有關了。第二層再讀取這些輸出時，就間接得到了位置資訊。有研究發現，不加位置編碼的因果語言模型也能學到位置資訊，但實際的模型仍然會明確地加入位置編碼。