Transformer 实现机制 · 第 10 课
多头注意力 —— 一个头一种眼光
目录上这课的问题挂了快一个月了:一套 Q/K/V 打出的一种分数,照顾得过来吗?第 8 课埋过一句伏笔——"吃行的最大配比给了它自己……多头注意力里不同头还会各有所爱"。今天收口。先看单头的窘境:那套注意力整句话只算出一张配比表,猫行看鱼就是一个分数点。可一句话里的关系五花八门:代词找指代对象、动词找宾语、形容词找它修饰的名词、语气词找情感落点……一套眼镜一种眼光,戴十副眼镜的人看得更全——这就是多头注意力的全部动机。
一、朴素方案,与它的账单
最小改动谁都会:想要几种眼光,多装几套 Q/K/V 就是了——h 套相互独立的 W_q/W_k/W_v,每套完整 d_model 维,各算各的注意力,输出拼一起。「猫吃鱼」想要"指代 + 语法"两种眼光?装两套。
先别急着写代码,算笔账(以 GPT-2 为例:d_model = 768,h = 12 头):
| 做法 | 参数量 | 点评 |
|---|---|---|
| 朴素:h 套完整眼镜 | 12 × 3 × 768² ≈ 2123 万 | 一层注意力顶过去九层,账单爆炸 |
| 原论文:切分 + 融合(下节) | 3 × 768² + 768² ≈ 236 万 | 与"单头 + 输出投影"完全相同 |
朴素方案还有一个隐性浪费:12 套眼镜全在同一个 768 维全空间里各自为政,学出来的东西大概率高度重叠——花了 9 倍的钱,买了 12 份雷同的目光。问题不是"要不要多套眼镜",而是"怎么便宜地多套"。
二、原论文的迭代:切分,不是重复
《Attention Is All You Need》的答案漂亮在换了个问法:与其给每人发一副新望远镜,不如把一台望远镜切成几段,各段各看各的。具体三步:
- 切:把 d_model 维切成 h 份,每头拿 d_k = d_model ÷ h 维(GPT-2:768 ÷ 12 = 64)。q、k、v 全部照切。
- 各算各的:每头在自己那个低维子空间里,独立走第 8 课的四步(打分 → 缩放 → softmax → 加权 V;第 9 课的掩码每头都要戴)——h 张小配比表,互不干扰。
- 拼 + 融合:h 份输出拼回 d_model 维(concat),再过一个可训练矩阵 WO 融合——让各头"开会碰头",决定谁的话听几成。
三、手算:把「猫吃鱼」切成两头
老朋友第 8 课的 q/k/v(4 维小整数)直接上手术台:h = 2,每头拿 d_k = 4 ÷ 2 = 2 维——头 1 拿维 0、1,头 2 拿维 2、3,一个数都不改:
| 词 | q(切前 → 切后) | k | v |
|---|---|---|---|
| 猫 | [1,1,0,0] → 头1 [1,1] · 头2 [0,0] | [1,0,1,0] → [1,0] · [1,0] | [1,0,1,0] → [1,0] · [1,0] |
| 吃 | [0,2,0,2] → [0,2] · [0,2] | [0,1,0,1] → [0,1] · [0,1] | [0,1,0,1] → [0,1] · [0,1] |
| 鱼 | [0,0,1,1] → [0,0] · [1,1] | [1,1,0,0] → [1,1] · [0,0] | [2,0,0,2] → [2,0] · [0,2] |
手算:头 1 独立四步(猫行;每头都戴第 9 课掩码,句首只能看自己)
打分(d_k = 2,÷√2):q_猫·k_猫 = 1,q_猫·k_吃 = 1,q_猫·k_鱼 = 2 → 缩放 [0.71, 0.71, 1.41]。
softmax → 猫行配比 [0.25, 0.25, 0.50]——头 1 盯上了鱼(50%)。
加权 v:0.25×[1,0] + 0.25×[0,1] + 0.50×[2,0] = [1.26, 0.25]。(吃行 [0.11, 0.45, 0.45]、鱼行 [0.33, 0.33, 0.33],下面代码对答案。)
softmax → 猫行配比 [0.25, 0.25, 0.50]——头 1 盯上了鱼(50%)。
加权 v:0.25×[1,0] + 0.25×[0,1] + 0.50×[2,0] = [1.26, 0.25]。(吃行 [0.11, 0.45, 0.45]、鱼行 [0.33, 0.33, 0.33],下面代码对答案。)
手算:头 2 独立四步(猫行)
打分:q_猫(头2)= [0,0],与谁点积都是 0 → 缩放 [0, 0, 0] → softmax → [0.33, 0.33, 0.33]——均匀巡逻。
加权 v:0.33×[1,0] + 0.33×[0,1] + 0.33×[0,2] = [0.33, 1.00]。(吃行 [0.16, 0.67, 0.16]、鱼行 [0.40, 0.40, 0.20]。)
加权 v:0.33×[1,0] + 0.33×[0,1] + 0.33×[0,2] = [0.33, 1.00]。(吃行 [0.16, 0.67, 0.16]、鱼行 [0.40, 0.40, 0.20]。)
停一秒,看这个对照:同一只猫,头 1 盯鱼(50%),头 2 雨露均沾(33%×3)——两套眼镜,两种眼光,没有谁对谁错(真实模型里哪种眼光有用,训练说了算)。这就是第 8 课那句伏笔的完整兑现:吃行 self-0.67 的"自恋",只是单头的一种口味;多头把所有口味分给不同的头去尝。
手算:Concat 与 W^O 融合
Concat(纯拼列,无参数):猫的 concat = 头1 [1.26, 0.25] ⊕ 头2 [0.33, 1.00] = [1.26, 0.25, 0.33, 1.00](前 2 格头 1 的,后 2 格头 2 的)。
W^O(可训练,d_model × d_model):把 4 格跨头混合成最终表示——好比两位调查员交报告,主管决定各听几成。它也是多头不涨参数量的关键一笔(见 §五)。
W^O(可训练,d_model × d_model):把 4 格跨头混合成最终表示——好比两位调查员交报告,主管决定各听几成。它也是多头不涨参数量的关键一笔(见 §五)。
四、动画:切分流水线走一遍
动画 · 多头切分流水线:完整 q → 切 → 各头独立算 → concat → W^O
8 步播放/单步;权重是手工挑的示意——「盯鱼头」「巡逻头」的分工要靠训练才长得出来,随机初始化的多头没有真语义(下面代码里两头的差异同样只是"随机噪声不同"而已)。
五、公式:论文原式,逐符号解码
公式 + 逐符号讲解(多头注意力——每个符号什么意思、数字从哪来)
$$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1,\dots,\text{head}_h)\,W^O, \qquad \text{head}_i = \text{Attention}(QW_i^Q,\; KW_i^K,\; VW_i^V)$$
$h$(头数):超参数,GPT-2 取 12——"装几套眼镜"。$i$ = 1…h(头编号):第 i 套眼镜,不是位置编号(位置在论文里用别的字母,别混)。
$W_i^Q, W_i^K, W_i^V$:第 i 头自己的三副小眼镜,每副形状 $(d_{model} \times d_k)$,$d_k = d_{model} \div h$(GPT-2:768×64)。它们是切出来的:整个大矩阵 $W^Q$$(d_{model} \times d_{model})$ 按列均分给 h 个头,真实实现一次矩阵乘完再切(下面代码演示)。
$\text{head}_i$:第 i 头独立走完四步的输出,形状 $L \times d_k$——猫行那格就是手算的 [1.26, 0.25]。
$\text{Concat}$:把 h 份 $L \times d_k$ 沿列拼回 $L \times d_{model}$——纯重排,没有参数。
$W^O$(O = Output,输出融合矩阵):形状 $d_{model} \times d_{model}$,把拼好的表示跨头混合——"会议纪要整合器"。
参数量账:眼镜 $h \times 3 \times d_{model} \times d_k = 3 \times d_{model}^2$(切分不改总量),加 $W^O$ 一份 $d_{model}^2$,共 $4 d_{model}^2$——与"单头 + 输出投影"分毫不差。多头买到的多种眼光,边际成本为零。
下标防混淆:$W^O$ 的 O 是 Output,不是数字 0;$W_i$ 的 i 是头号;$d_k$ 在这里 = $d_{model} \div h$(每头的维度),不再是第 8 课单头语境里"= d_model"的 d_k——同一个符号,语境不同,认符号要看它当前代表什么。
$W_i^Q, W_i^K, W_i^V$:第 i 头自己的三副小眼镜,每副形状 $(d_{model} \times d_k)$,$d_k = d_{model} \div h$(GPT-2:768×64)。它们是切出来的:整个大矩阵 $W^Q$$(d_{model} \times d_{model})$ 按列均分给 h 个头,真实实现一次矩阵乘完再切(下面代码演示)。
$\text{head}_i$:第 i 头独立走完四步的输出,形状 $L \times d_k$——猫行那格就是手算的 [1.26, 0.25]。
$\text{Concat}$:把 h 份 $L \times d_k$ 沿列拼回 $L \times d_{model}$——纯重排,没有参数。
$W^O$(O = Output,输出融合矩阵):形状 $d_{model} \times d_{model}$,把拼好的表示跨头混合——"会议纪要整合器"。
参数量账:眼镜 $h \times 3 \times d_{model} \times d_k = 3 \times d_{model}^2$(切分不改总量),加 $W^O$ 一份 $d_{model}^2$,共 $4 d_{model}^2$——与"单头 + 输出投影"分毫不差。多头买到的多种眼光,边际成本为零。
下标防混淆:$W^O$ 的 O 是 Output,不是数字 0;$W_i$ 的 i 是头号;$d_k$ 在这里 = $d_{model} \div h$(每头的维度),不再是第 8 课单头语境里"= d_model"的 d_k——同一个符号,语境不同,认符号要看它当前代表什么。
六、代码:切分账本 + GPT 同款多头类
# ============================================================
# 【演示目标】① 用数字验证"多头是切分不是重复计算":朴素 h 套
# 完整眼镜 vs 原论文切分,参数量差多少;② 形状体检:X 过 h=2 头
# 的完整形状链,每步形状对上课文手算。
# 【思路】
# 1. d_model=768(GPT-2)、h=12:
# 朴素 = 12 套 3×(768×768);切分 = 12 套 3×(768×64) + W^O(768×768)
# 2. 形状链:X(3 词×8 维) → 切 2 头 → 每头 Q_i (3,4) → 配比 (3,3)
# → 头上下文 (3,4) → concat (3,8) → W^O(8,8) → 输出 (3,8)
# 读法: 两种做法参数量差 9 倍;切分总量与「单头+W^O」分毫不差——
# 多头 = 免费的多眼光。形状链与「猫吃鱼」手算逐格对上。
# ============================================================
import numpy as np
d_model, h = 768, 12
naive = h * 3 * d_model * d_model # h 套完整眼镜
paper = h * 3 * d_model * (d_model // h) + d_model * d_model # 切分 + W^O
print("朴素 h 套完整眼镜:", format(naive, ","), "参数")
print("原论文 切分+WO :", format(paper, ","), "参数 ← 与单头+输出投影相等:", paper == 4 * d_model * d_model)
print("单头+WO 应为 :", format(4 * d_model * d_model, ","), "\n")
# ---- 形状体检(玩具:d_model=8, h=2, d_k=4)----
rng = np.random.RandomState(10)
X = rng.randn(3, 8) # 「猫吃鱼」,3 词 × 8 维
d_k = 8 // 2 # 每头维度 = d_model ÷ h = 4
Wq, Wk, Wv = (rng.randn(8, 8) * 0.5 for _ in range(3)) # 三个大矩阵(真实实现就这么存)
heads = []
for i in range(2):
Q_i = X @ Wq[:, i*4:(i+1)*4] # 按列切给第 i 头:Q_i (3,4)
K_i = X @ Wk[:, i*4:(i+1)*4]
V_i = X @ Wv[:, i*4:(i+1)*4]
S = Q_i @ K_i.T / np.sqrt(d_k) # 每头独立打分
S[np.triu(np.ones((3, 3), dtype=bool), 1)] = -1e9 # 第 9 课掩码,每头都要戴
e = np.exp(S - S.max(axis=1, keepdims=True))
A_i = e / e.sum(axis=1, keepdims=True) # 每头独立 softmax → (3,3)
heads.append(A_i @ V_i) # 每头独立加权 V → (3,4)
print("头%d: Q_i%s A_i%s 输出%s" % (i+1, Q_i.shape, A_i.shape, heads[-1].shape))
Z = np.concatenate(heads, axis=1) # concat → (3,8)
Wo = rng.randn(8, 8) # 融合矩阵
print("concat:", Z.shape, "→ @ W^O(8,8) →", (Z @ Wo).shape, " ← 进出同形状")
# ============================================================
# 【演示目标】把多头注意力焊成类(GPT 同款骨架,含第 9 课掩码):
# 大矩阵乘一次 → 按头切 → 每头四步 → concat → W^O;并看两个头
# 在同一句话上的配比差别(各有所爱——虽然此时还只是随机差异)。
# 【思路】
# A. 类:Wq/Wk/Wv 存成完整 (d_model × d_model) 大矩阵(真实做法),
# 前向时 reshape 成 (h, d_model, d_k)——一次乘完再切,省 h 次小乘法
# B. 体检:X(4 词×8 维), h=2 → 打印每头的配比表、行和、输出形状
# C. 观察:头1/头2 配比逐格不同;但随机权重下的差异没有语义,
# 训练后才会分化出真正的「指代头」「语法头」
# 读法: 每头配比表上三角全 0(第 9 课掩码在每头都工作);两头分布
# 肉眼可见不同;输出形状与输入相同——多头也不改形状,只换看法。
# ============================================================
import numpy as np
rng = np.random.RandomState(11)
def softmax(x, axis=-1):
e = np.exp(x - x.max(axis=axis, keepdims=True))
return e / e.sum(axis=axis, keepdims=True)
class MultiHeadAttention:
"""多头因果注意力:GPT 每一层的同款骨架(Vizuara 1.16 的 numpy 版)"""
def __init__(self, d_model, h):
assert d_model % h == 0, "d_model 必须能被头数整除(切分的前提!)"
self.h, self.d_model = h, d_model
self.d_k = d_model // h # 每头维度 = d_model ÷ h
self.Wq = rng.randn(d_model, d_model) * 0.4 # 大矩阵(内含全部头的眼镜)
self.Wk = rng.randn(d_model, d_model) * 0.4
self.Wv = rng.randn(d_model, d_model) * 0.4
self.Wo = rng.randn(d_model, d_model) * 0.4 # 输出融合
def forward(self, X):
L, d = X.shape
# 一次大乘法 + 切头:reshape 成 (h, L, d_k)——比循环 h 次小乘法快
Q = (X @ self.Wq).reshape(L, self.h, self.d_k).transpose(1, 0, 2) # (h, L, d_k)
K = (X @ self.Wk).reshape(L, self.h, self.d_k).transpose(1, 0, 2)
V = (X @ self.Wv).reshape(L, self.h, self.d_k).transpose(1, 0, 2)
S = Q @ K.transpose(0, 2, 1) / np.sqrt(self.d_k) # 每头独立打分 → (h, L, L)
mask = np.triu(np.ones((L, L), dtype=bool), 1) # 第 9 课:每头都要戴掩码
S[:, mask] = -1e9
A = softmax(S, axis=-1) # 每头独立 softmax → (h, L, L)
heads = A @ V # 每头独立加权 V → (h, L, d_k)
Z = heads.transpose(1, 0, 2).reshape(L, d) # concat → (L, d_model)
return Z @ self.Wo, A # W^O 融合 + 配比表(供观察)
X = rng.randn(4, 8) # 4 词 × 8 维
mha = MultiHeadAttention(d_model=8, h=2)
np.set_printoptions(precision=2, suppress=True)
Z, A = mha.forward(X)
print("头 1 配比表(4×4,上三角=未来=0):")
print(A[0])
print("头 2 配比表:")
print(A[1])
print("两头同格对比(第 3 行):", A[0][3].round(2), "vs", A[1][3].round(2), "← 各有所爱(随机版)")
total = sum(v.size for v in mha.__dict__.values() if hasattr(v, "size"))
print("参数量账:Wq/Wk/Wv/Wo 各", mha.Wq.size, "个 → 共", total, "= 4×d_model² =", 4 * 64, "✓")
print("输出 Z:", Z.shape, "—— 进 (4,8) 出 (4,8),形状不变量依旧")
本课胜利
① 单头的窘境:一套眼镜一种眼光、一张配比表,句中多种关系只能被平均;② 朴素 h 套完整眼镜 = 2123 万参数,账单爆炸;③ 原论文的迭代:切分不是重复——d_k = d_model ÷ h,各头低维子空间独立四步(每头戴掩码),concat 回来过 W^O 融合;④ 「猫吃鱼」h=2 手算:头 1 猫行盯鱼 50%、头 2 均匀 33%——同一只猫,两种眼光(第 8 课伏笔兑现);⑤ 公式 MultiHead = Concat(head₁…head_h)W^O 逐符号全认识;⑥ 参数量账:切分 + W^O = 单头 + W^O = 4d²,多种眼光边际成本为零。你刚亲手切过的这套流水线,就是论文里的 Multi-Head Attention——GPT-2 的 12 副眼镜,你现在能一副副数出来了。下一课:block 里还剩两个没拆的零件——网络深了为什么不稳?谁来做稳定器?
检索练习
本课主读材料
去读(约 12 分钟)
Vizuara《The Transformers》1.16 节:多头注意力的完整推演——重点对图:切分/多头并行计算/concat 的流程图,与本课的 SVG 和动画一一对应。
图解补充:The Illustrated Transformer 的多头一节(多组彩色 Q/K/V 箭头,和本课配色同款思路)。
延伸(可选):Are Sixteen Heads Really Better than One?——训练好的模型里不少头可以剪掉,「各头分工」比想象中更有弹性。
我是你的老师,别客气
"头数怎么选""h 不整除 d_model 行不行""W^O 能不能省""为什么切 k、v 也要切 q"——这类问题都在正道上,贴回来问。下一课拆 block 最后两个零件:残差与 LayerNorm,回答"网络深了为什么不稳"。