Transformer 解剖：从 Attention 到推理系统

第 3 章 Multi-Head Attention：分头的代数与几何意义

作者杨艺韬 · 6,439 字 · 发布于 2026-05-09 · 更新于 2026-09-25

第 2 章我们把 Self-Attention 的内核拆开看了。但实际工程里你看到的不会是单头 Attention，永远是 Multi-Head Attention——多个头并行计算。Llama 70B 有 64 个头，GPT-3 有 96 个头，DeepSeek-V3 有 128 个头。

为什么要多头？把同一件事并行做多遍有什么好处？这一章我们要彻底讲清楚两件事：第一，Multi-Head 在数学和几何上到底在做什么——它不是「把单头跑 N 遍取平均」；第二，从原始 MHA 到今天主流的 GQA / MQA，这条演化路线背后的工程压力是什么。

读完这章你会理解：为什么 Llama 3 的 70B 模型只有 8 个 KV 头但有 64 个 Q 头、为什么 GPT-3 用 MHA 而 Llama 3 改用了 GQA、为什么 PaLM 用 MQA 而 Mistral 7B 用 GQA。这些选择背后都是同一道工程账。

3.1 单头不够：语言里同时存在多种关系

我们先停下来想一个问题：自然语言里，一个 token 和上下文之间的关系只有一种吗？

回到第 2 章那个例子：

The animal didn't cross the street because it was too tired.

"it" 这个 token 至少有这几条不同的关系链：

指代关系：it → animal（共指）
语法关系：it 是 was 的主语
语义关系：it / tired 描述同一个实体的状态
位置关系：it 处在主从句的从句开头
修辞关系：整个 because-从句解释为什么 didn't cross

如果你只用一个 Self-Attention 头来处理，所有这些关系都得挤在同一个注意力分布里。结果就是：要么模型只学到最强的那种关系（比如指代），其他的都被压制；要么模型试图同时学所有关系，导致每种都学得不够好。

更深的问题是几何上的。一个头的注意力权重 $\alpha_{ij}$ 是由一组 $W_Q, W_K$ 投影出来的 $q_i, k_j$ 的点积决定的。 $W_Q, W_K$ 一旦选定，整个注意力分布就被限制在它们投影出来的那个 $d_k$ 维子空间里。一组投影很难同时照顾好多类相关性——想用同一组投影矩阵同时编码「语法」和「语义」这两种很不一样的语言学结构，注意力分布只能互相妥协。

Multi-Head 的解决方案非常直接：让模型有多组 $W_Q, W_K, W_V$ ，每一组独立学一种关系，然后把它们的输出拼起来。每个头活在自己那个子空间里、专心学一类关系，所有头加起来就覆盖了语言里多样的相关性结构。

flowchart LR
  X[输入 x_i] --> H1["Head 1<br/>学指代关系"]
  X --> H2["Head 2<br/>学语法关系"]
  X --> H3["Head 3<br/>学语义关系"]
  X --> HD[...]
  X --> HN["Head h<br/>学修辞关系"]
  H1 --> CONCAT[拼接所有头的输出]
  H2 --> CONCAT
  H3 --> CONCAT
  HD --> CONCAT
  HN --> CONCAT
  CONCAT --> WO["× W_O 投影"] --> OUT[最终输出]

注意「学指代 / 学语法 / 学语义」是为了讲解直觉给的标签。模型不会真的有这些标签——它们是模型从数据中自发涌现出来的分工。后面 3.5 节会用 BERT 上的实证分析和我们自己训的 mini-GPT 的注意力图让你看到这种分工长什么样。

3.2 Multi-Head 的代数定义

把直觉翻译成数学。给定输入 $X \in \mathbb{R}^{T \times d_{\text{model}}}$ 、头数 $h$ 、每个头的维度 $d_k = d_v = d_{\text{model}} / h$ 。

第一步：每个头有自己的一组投影矩阵 $W_Q^{(i)}, W_K^{(i)}, W_V^{(i)} \in \mathbb{R}^{d_{\text{model}} \times d_k}$ ，其中 $i = 1, \dots, h$ 。

第二步：每个头独立做一次完整的 Self-Attention：

\text{head}_i = \text{Attention}(X W_Q^{(i)}, X W_K^{(i)}, X W_V^{(i)})

每个头的输出形状是 $(T, d_v)$ 。

第三步：把 $h$ 个头的输出沿特征维拼起来（不是相加！这是常见误解），得到形状 $(T, h \cdot d_v) = (T, d_{\text{model}})$ ：

\text{Concat}(\text{head}_1, \dots, \text{head}_h) \in \mathbb{R}^{T \times d_{\text{model}}}

第四步：再过一次输出投影 $W_O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$ ：

\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W_O

合起来就是这一个公式：

\text{MultiHead}(X) = \text{Concat}_{i=1}^{h}\left[\text{softmax}\left(\frac{(X W_Q^{(i)})(X W_K^{(i)})^T}{\sqrt{d_k}}\right) X W_V^{(i)}\right] W_O

看着复杂，但每一步都是第 2 章已经讲过的运算，只是被并行复制了 $h$ 份。

Concat 是「沿最后一维拼接」——Head1 (T, d_v) 和 Head2 (T, d_v) 拼起来变成 (T, 2·d_v)。这个操作在 PyTorch 里是 torch.cat([head1, head2, ...], dim=-1)。

3.3 关键的工程实现：一次大矩乘

上面的描述里我们说 " $h$ 个头独立做"，听起来像要跑 $h$ 次 for 循环——这在 GPU 上显然效率不高。实际工程上，所有头都被打包到一次大矩阵乘法里：

把所有头的投影矩阵沿特征维拼起来：

W_Q^{\text{big}} = [W_Q^{(1)} | W_Q^{(2)} | \dots | W_Q^{(h)}] \in \mathbb{R}^{d_{\text{model}} \times h \cdot d_k}

由于 $h \cdot d_k = d_{\text{model}}$ ，所以 $W_Q^{\text{big}} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$ 。 $W_K^{\text{big}}, W_V^{\text{big}}$ 同理。

一次大矩乘把所有头的 Q/K/V 算出来：

Q^{\text{big}} = X W_Q^{\text{big}} \in \mathbb{R}^{T \times d_{\text{model}}}

接下来用一个 reshape，把 $d_{\text{model}}$ 这一维拆成 $(h, d_k)$ ：

Q^{\text{big}} \in \mathbb{R}^{T \times d_{\text{model}}} \xrightarrow{\text{reshape}} Q \in \mathbb{R}^{h \times T \times d_k}

PyTorch 里这一步是 Q.view(T, h, d_k).transpose(0, 1)。然后所有头的 attention 就可以用一个带 batch 维的矩阵乘法统一算：

# Q, K, V 形状都是 (h, T, d_k)
S = Q @ K.transpose(-2, -1)        # (h, T, T)
S = S / math.sqrt(d_k)
A = torch.softmax(S, dim=-1)        # (h, T, T)
out = A @ V                         # (h, T, d_k)
out = out.transpose(0, 1).contiguous().view(T, d_model)
out = out @ W_O                     # (T, d_model)

整个过程没有 for 循环。所有 $h$ 个头的全部计算被压到 6 次矩阵乘法（QKV 投影 3 次、 $QK^T$ 一次、 $AV$ 一次、输出投影一次，其中 $QK^T$ 和 $AV$ 这两次是带 batch 维的）加 1 次 softmax 里——这是 Multi-Head 在工程上能高效跑的根本原因。

flowchart LR
  X["X<br/>(T, d_model)"] --> WQB["× W_Q_big<br/>(d_model, d_model)"]
  WQB --> QB["Q_big<br/>(T, d_model)"] --> RS["reshape<br/>(h, T, d_k)"]
  RS --> Q["Q<br/>(h, T, d_k)"]
  X --> KV["类似得到 K, V<br/>(h, T, d_k)"]
  Q --> ATT["batched<br/>scaled dot product<br/>attention"]
  KV --> ATT
  ATT --> CAT["transpose + reshape<br/>(T, d_model)"]
  CAT --> WO["× W_O"] --> OUT["MultiHead 输出<br/>(T, d_model)"]

3.4 参数量与计算量

来算一下 Multi-Head Attention 的参数量。每个头有 $W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}$ ，每个是 $d_{\text{model}} \times d_k$ ；总共 $h$ 个头：

\text{params}_{QKV} = h \cdot 3 \cdot (d_{\text{model}} \cdot d_k) = 3 \cdot d_{\text{model}}^2

（用了 $h \cdot d_k = d_{\text{model}}$ ）

加上输出投影 $W_O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$ ：

\text{params}_{\text{MHA}} = 3 d_{\text{model}}^2 + d_{\text{model}}^2 = 4 d_{\text{model}}^2

注意：Multi-Head 的总参数量等于「等价单头但 $d_k = d_{\text{model}}$ 」的参数量——分头并没有让参数变多。这是一个非常重要的观察：MHA 不是用更多参数换更强表达力，而是用同样的参数实现了结构上更丰富的表达（多个独立子空间）。

计算量呢？一次 Self-Attention 的计算量主要是两块：

QKV 投影： $O(T \cdot d_{\text{model}}^2)$
Attention 矩阵（QK^T、softmax、AV）： $O(T^2 \cdot d_{\text{model}})$

对 Multi-Head 来说总数和单头一致（投影的参数量一致；attention 矩阵那一项是 $h \cdot T^2 \cdot d_k = T^2 \cdot d_{\text{model}}$ ），分头只是把计算重新组织。但当 $T$ 很大时，第 2 项 $O(T^2 \cdot d_{\text{model}})$ 是主导——这是 Transformer 在长上下文下的瓶颈，第 13 章会展开。

3.5 不同头到底学到了什么

直觉上我们说「不同头学不同关系」，是真的吗？这个问题在 BERT 出来之后被研究得很透彻。

2019 年 Clark et al. 的论文 What Does BERT Look At? An Analysis of BERT's Attention 系统分析了 BERT-base（12 层 × 12 头 = 144 个头）的注意力分布，发现了几类典型的「专家头」：

位置头：有些头几乎只关注前一个 token 或后一个 token（局部位置关系）。
句法头：有些头集中关注当前 token 的句法依存对象（如动词的宾语、名词的限定词、介词的宾语）。
共指头：有个头（第 5 层第 4 头）会把代词和它指代的名词连起来（"it" 与 "animal"），共指消解的准确率明显高于字符串匹配基线。
分隔符头：相当多的头把大量注意力放在特殊 token 上——浅层偏向 [CLS]，第 6–10 层有一半以上的注意力落在 [SEP] 上，深层则偏向句号和逗号。论文的解释是：这不是模型在「重视分隔符」，而是某个头的功能用不上时，把注意力停到 [SEP] 上当作「no-op」。softmax「每行必须加起来等于 1」，注意力不能凭空消失，只能找个位置倾倒——后来的研究把这类位置称为「注意力 sink」。

flowchart TB
  subgraph BERT-base 注意力分工 据 Clark 等 2019 简化
    L1["浅层 - 关注 CLS 的头多，看前后相邻 token 的位置头也多在较早的层"]
    L4["中层 - 句法依存头（如 8-10 宾语、8-11 限定词），5-4 共指头"]
    L8["第 6-10 层 - 一半以上注意力落在 SEP 上"]
    L12["深层 - 注意力偏向句号与逗号"]
  end
  L1 --> L4 --> L8 --> L12

这个分工不是模型被显式训练出来的——它是从大规模语言建模这个单一目标里自发涌现的。这是 Multi-Head 设计真正的妙处：你不需要告诉模型「这个头学语法、那个头学指代」，模型自己会按照「让 loss 最低」的原则自动分工。

下面这张图是我们用第 9 章那个 1.7M 参数的 mini-GPT（4 层 × 4 头）在唐诗上训完后，把《静夜思》整首灌进去得到的最后一层 4 个头的 attention 模式（由本专栏配套代码里的画图脚本生成）。可以看到各个头已经形成了不同的模式：

mini-GPT 最后一层 4 个头的 attention 模式

（图里「举」「头」「乡」显示为 <unk>：训练语料取自 chinese-poetry 仓库的《全唐诗》，是繁体字，简体的这几个字不在词表里。）

Head 1、2、4 都以看近邻为主：每个字的注意力集中在自己往前 1～3 个位置上，只是偏移量各不相同。Head 3 除了紧挨着自己的那条对角线，还有一条较淡的平行带，大约落在往前 6～7 个 token 处（比如「是」看 <bos>、「地」看「床」、「霜」看「明」），跨过了一整句五言——像是在拿上一句相近位置的字做参照。不过这只是一个 1.71M 小模型在一条样本上的图，偏移也不严格等于句长 5，不宜解读成学会了对仗。

即便这么小的模型，不同的头也各自形成了不同的注意力模式，这和 BERT 实验里「不同头学不同关系」的观察方向一致——只不过我们的模型比 BERT-base（约 1.1 亿参数）小了近两个数量级。

但有一个发现让后续的工作产生了深远影响：很多头其实是冗余的。Voita et al. 在 Analyzing Multi-Head Self-Attention（2019）里对翻译模型做了剪头实验，Michel et al. 在 Are Sixteen Heads Really Better than One?（2019）里对 BERT 和翻译模型都做了，结论一致：相当一部分头可以被剪掉而几乎不影响下游任务表现——不少层甚至可以只保留 1 个头。

这个观察也给另一条工程优化路线提供了旁证：既然很多头是冗余的，能不能在保留 query 多头的同时，让 key 和 value 共享？——这就是 MQA / GQA 的思路。不过 MQA 论文本身的直接动机不是头冗余，而是下一节要讲的增量解码时反复读取 K、V 的显存带宽开销。

3.6 推理瓶颈推动的演化：MHA → MQA → GQA

进入大模型时代后，Multi-Head 的代价开始显形。问题不在训练阶段，而在推理阶段的 KV Cache。

第 15 章会详细讲 KV Cache，这里先给一个直觉：自回归生成时，每生成一个新 token，模型不需要重新算前面所有 token 的 K 和 V——它们可以缓存下来反复用。这个缓存的大小，对一个有 $h$ 个头、序列长度 $T$ 、每个头维度 $d_k$ 、 $L$ 层的模型来说是：

\text{KV Cache size} = 2 \cdot L \cdot h \cdot T \cdot d_k \cdot \text{bytes per element}

「2」是 K 和 V 各一份。我们用 Llama-3-70B 算一笔账（先假设它没有用 GQA、保持 MHA）： $L = 80$ 、 $h = 64$ 、 $d_k = 128$ ，FP16 就是 2 bytes。在 $T = 4096$ 时：

2 \cdot 80 \cdot 64 \cdot 4096 \cdot 128 \cdot 2 \approx 10 \text{ GiB}

光是一条序列的 KV Cache 就要 10 GiB——而 Llama-3-70B 模型权重本身在 FP16 下约 140 GB。如果你想 batch=8 同时跑 8 个用户的请求，KV Cache 一项就要 80 GiB——已经把整张 H100 的显存吃完。这是「如果继续用 MHA」的反事实——下面我们会看到正是这个压力推动了 GQA / MLA 的演化。

这就是 KV Cache 的工程压力：它正比于 head 数 $h$ 。如果能降低有效的 KV head 数，KV Cache 立刻线性下降。

Multi-Query Attention (MQA)

最激进的优化是 2019 年 Noam Shazeer 提出的 MQA：所有 query 头共享同一组 key 和 value：

\text{head}_i = \text{Attention}(X W_Q^{(i)}, X W_K^{\text{shared}}, X W_V^{\text{shared}})

也就是说， $Q$ 还是 $h$ 个头，但 $K$ 和 $V$ 只有 1 个头。

KV Cache 立刻从 $2 L h T d_k$ 降到 $2 L T d_k$ ——降低了 $h$ 倍。Llama-3-70B 在上面 MHA 假设下的 10 GiB 直接砍到约 160 MiB。

但 MQA 有代价：因为所有 query 都用同一份 key、value，模型表达能力降低。GQA 论文（Ainslie et al., 2023）在 T5 上的对照实验显示 MQA 相比 MHA 有可见的质量下降；另一边 PaLM 论文（Chowdhery et al., 2022）则报告 MQA 对模型质量和训练速度的影响是中性的，同时显著降低自回归解码成本。两者设置不同（GQA 论文是把训好的 MHA 检查点改成 MQA 再少量续训，PaLM 是从头按 MQA 训练），结论不宜直接互相外推。

Grouped-Query Attention (GQA)

Ainslie et al. 在 2023 年提出的折中方案：把 query 头分组，组内共享一组 key 和 value。

设 $h_q$ 个 query 头， $h_{kv}$ 个 KV 头（且 $h_q$ 是 $h_{kv}$ 的整数倍）。每 $h_q / h_{kv}$ 个 query 头共享一组 K、V：

\text{group}_g\text{ shared K, V; queries } q_{(g-1) \cdot (h_q/h_{kv})+1}, \dots, q_{g\cdot(h_q/h_{kv})}

最直白的实现并不需要一套专门的「分组 attention」内核：把缓存里的 $h_{kv}$ 份 K、V 沿头维复制 $n_{rep} = h_q / h_{kv}$ 份，再走原来的 MHA 路径——transformers-4815a0a/src/transformers/models/llama/modeling_llama.py:179 的 repeat_kv 就是干这件事的（先 expand 出一个视图，再 reshape 成 $h_q$ 个头——后一步会真的把数据拷贝一份出来）。也有不复制的路径：transformers 的 SDPA 集成在条件允许时直接给 scaled_dot_product_attention 传 enable_gqa=True（transformers-4815a0a/src/transformers/integrations/sdpa_attention.py:98-102），条件不满足才退回 repeat_kv。所以要记清楚 GQA 省的是什么：省的是KV Cache 的体积和读它的带宽，不是 attention 计算量—— $QK^\top$ 和 $AV$ 仍然按 $h_q$ 个头算。配置侧则是 transformers-4815a0a/src/transformers/models/llama/configuration_llama.py:89：num_key_value_heads 不填时直接等于 num_attention_heads，也就是退化回 MHA。

GQA 是 MHA 和 MQA 的连续插值：

$h_{kv} = h_q$ → MHA（每个 query 头独享 K、V）
$h_{kv} = 1$ → MQA（所有 query 头共享 K、V）
中间值 → GQA

flowchart TB
  subgraph mha ["MHA: h=8"]
    Q1[Q1-K1-V1]
    Q2[Q2-K2-V2]
    Q3[Q3-K3-V3]
    Q4[Q4-K4-V4]
    Q5[Q5-K5-V5]
    Q6[Q6-K6-V6]
    Q7[Q7-K7-V7]
    Q8[Q8-K8-V8]
  end
  subgraph gqa ["GQA: hq=8, hkv=4"]
    GQ1["Q1, Q2 共享 K1-V1"]
    GQ2["Q3, Q4 共享 K2-V2"]
    GQ3["Q5, Q6 共享 K3-V3"]
    GQ4["Q7, Q8 共享 K4-V4"]
  end
  subgraph mqa ["MQA: hq=8, hkv=1"]
    MQ["Q1, Q2, ..., Q8 全部共享 K-V"]
  end

GQA 在 Llama-2 之后成为主流。Llama-3-70B 用的就是 $h_q = 64, h_{kv} = 8$ 的 GQA——KV Cache 比 MHA 降到原来的 1/8（4K 上下文下从 10 GiB 降到 1.25 GiB），同时质量损失很小（GQA 论文在 T5-XXL 上的实验里，GQA-8 的平均分 47.1，MHA 为 47.2）。

主流模型的注意力变种汇总：

模型	h_q	h_kv	类型
GPT-3 175B	96	96	MHA
Llama-1 7B	32	32	MHA
Llama-2 70B	64	8	GQA
Llama-3 70B	64	8	GQA
Mistral 7B	32	8	GQA
PaLM 540B	48	1	MQA
GPT-4	未公开	未公开	未公开（OpenAI 未披露架构细节）
DeepSeek-V2/V3	128	—（KV 压成一份低秩隐向量）	MLA（更激进的 KV 压缩，第 15 章会讲）

可以看到 2023 年之后开源模型大多从 MHA 切到了 GQA（DeepSeek 走的是 MLA），主要驱动就是推理时的 KV Cache 压力。MQA 太激进，GQA 是「质量损失小、显存收益大」的甜点。

3.7 几何上理解 Multi-Head：把空间切成子空间

我们在 3.1 节用「不同头学不同关系」给了直觉，3.5 节用 BERT 实验给了证据，3.6 节讲了工程演化。这一小节我们再换一个角度——几何视角——给 Multi-Head 一个更精确的解读。

把 $W_Q^{(i)} \in \mathbb{R}^{d_{\text{model}} \times d_k}$ 看作一个线性投影：把 $d_{\text{model}}$ 维空间投到一个 $d_k$ 维子空间。每个头有自己的 $W_Q^{(i)}, W_K^{(i)}$ 投影对，把输入投到一个专属的 $d_k$ 维子空间，在这个子空间里计算注意力。

如果所有头的投影矩阵随机初始化，这些子空间在 $d_{\text{model}}$ 维空间里大致是相互正交（或低相关）的——也就是说，每个头看的是输入信息的一个不同的「视角」。在子空间 1 看上去高度相关的两个 token，在子空间 2 可能完全没关系。

模型训练的过程，就是让每个 $W_Q^{(i)}, W_K^{(i)}$ 旋转到一个「让某种特定相关性凸显出来」的方向。最终：

头 1 的子空间编码了「相邻位置的关系」
头 2 的子空间编码了「主谓一致」
头 3 的子空间编码了「指代」
……

这不是模型「学到了多套独立的 attention」，而是同一个 d_model 维空间被同时切分成多个互补的视角，每个视角看到的是同一份输入的不同侧面。

这个几何观点能解释为什么 Multi-Head 的总参数量不变但表达力增强：参数量取决于「自由度」，而 $h$ 个头的 $W_Q^{(i)} \in \mathbb{R}^{d_{\text{model}} \times d_k}$ 拼起来等价于一个 $d_{\text{model}} \times d_{\text{model}}$ 的大矩阵——参数自由度和单头的 $d_k = d_{\text{model}}$ 一样。但 Multi-Head 通过对最后这个矩阵施加分块结构（每 $d_k$ 列对应一个独立的 attention 子空间），约束了模型一定要把表达力切成 $h$ 个独立部分，强行避免了所有注意力都坍缩到同一种关系上。

约束就是先验。Multi-Head 的先验是「语言里有多种关系，必须分别建模」。这个先验和数据本身吻合得很好，所以 MHA 在实际效果上反而比单头大 $d_k = d_{\text{model}}$ 的 attention 要好（原论文 Table 3 (A) 行：在计算量不变的前提下，单头比最佳设置低 0.9 BLEU）。

3.8 如何选 head 数

现在你看模型卡片里 num_attention_heads 这个数字，应该不再是个抽象选项。来对照几条工程经验法则：

法则一： $d_{\text{model}}$ 必须能被 $h$ 整除。否则没法均匀分头。常见约定是 $d_k = d_v = 64$ （Vaswani 原论文这么用）或 128（Llama 系列）， $h$ 由 $d_{\text{model}} / d_k$ 决定。

法则二：每个头的维度 $d_k$ 不能太小。2.5 节的 $\sqrt{d_k}$ 管的是分数的尺度，管不了表达能力：每个头的分数矩阵 $QK^T$ 秩不超过 $d_k$ ， $d_k$ 太小（比如 8）时单个头能表示的注意力模式就很受限——Bhojanapalli 等人（2020，arXiv:2002.07028）把这称为多头注意力的「低秩瓶颈」。经验上 $d_k \ge 32$ 可以当作下限，多数主流模型是 64 或 128。

法则三：头数太多不一定更好。Multi-Head 的「先验」在 $h$ 太大时会变成「过度切分」——每个子空间太小，反而让模型学不到有意义的表达。所以头数应当随 $d_{\text{model}}$ 一起涨：GPT-3 175B 的 96 头、Llama-1 7B 的 32 头，每个头都是 128 维。

法则四：推理压力大时用 GQA。如果要在线服务大量用户、KV Cache 是瓶颈， $h_q / h_{kv}$ 取 4 到 8 是被工业验证过的甜点区间（Llama-3 70B 是 8，Llama-3 8B 与 Mistral 7B 是 4）。

3.9 Multi-Head 的训练动力学

为了完整，我们再讲一下 Multi-Head 在训练过程中的「动力学」——也就是它在反向传播下怎么演化。这一小节稍微高阶，跳过不影响后续阅读。

每个头独立做注意力，反向传播时也独立。 $W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}$ 的梯度只沿第 $i$ 个头这条路径回传——在同一层内，头与头之间的计算路径是分离的，它们在输出投影 $W_O$ （它把所有头的输出合到一起做最后的变换）以及之后的残差流、后续各层处才汇合。

这意味着两件事：

第一，初始化非常重要。如果两个头的初值相同或非常接近，它们在训练中会沿着相似的梯度方向演化，最终学到几乎相同的 attention 模式——「头崩塌」（head collapse）。这就是为什么各家实现都让三套投影矩阵独立随机初始化（PyTorch 的 nn.MultiheadAttention 在 _reset_parameters 里对 in-proj 权重整体做 Xavier 均匀初始化），各头的初值互不相关，起点自然拉得开。

第二，稀疏注意力下的训练不稳定。如果某个头的 attention 分布在初始化阶段就接近饱和（softmax 输出几乎是 one-hot），它的梯度会很小——头几乎不再更新，很难离开初始的那个模式。这也是 $\sqrt{d_k}$ 缩放（第 2.5 节）必须存在的原因之一。

GPT-2 / GPT-3 时代曾流行过一种做法：把残差分支输出投影的初始化 std 按层数缩小（GPT-2 乘 $1/\sqrt{N}$ ， $N$ 为残差层数），配合把 LayerNorm 挪到子层输入端（Pre-LN），来稳定训练。Llama 沿用 Pre-LN，并把 LayerNorm 换成了 RMSNorm（第 5 章会讲）。

3.10 为什么 Cross-Attention 也用 Multi-Head

Multi-Head 的设计在 Self-Attention 里讲得最多，但它在 Cross-Attention 里同样适用——而且同样重要。

Cross-Attention 是 Encoder-Decoder 架构里 Decoder 的一种 attention 变体： $Q$ 来自 Decoder， $K, V$ 来自 Encoder：

\text{CrossAttention}(D, E) = \text{softmax}\left(\frac{(D W_Q)(E W_K)^T}{\sqrt{d_k}}\right) E W_V

机器翻译里 Cross-Attention 是 Decoder 「回头看英文原文」的机制（参考第 1 章 Bahdanau Attention）。它同样可以做成 Multi-Head——多个头并行「看」原文，每个头关注一种对齐关系（词级对齐、短语级对齐、句法对齐等）。

不过到了 Decoder-only 模型时代（GPT 系列），Encoder 和 Cross-Attention 都被消解了——所有事都用 Self-Attention 加上因果掩码完成。第 6 章会讲清楚这次架构演化的代价和收益。

本章小结

单头 Self-Attention 很难同时兼顾多种关系——它的注意力分布被一组 $W_Q, W_K$ 投影出的那个 $d_k$ 维子空间限制住了。
Multi-Head 通过把空间切成 $h$ 个互补子空间，让每个头独立学一类相关性（语法、共指、语义、位置等），最后用 $W_O$ 把它们合在一起。
MHA 的总参数量等于等价单头——分头不是加参数，而是给同样的参数量加结构约束。
工程实现上所有头被打包到一次大矩阵乘法，没有 for 循环，GPU 上完全并行。
不同头有自发分工：BERT 实验里能识别出位置头、句法头、共指头、sink 头等专家。但很多头是冗余的——这也为 MQA / GQA 共享 K、V 提供了旁证。
MHA → MQA → GQA 是被推理时 KV Cache 压力推动的演化。今天开源主流模型大多用 GQA（DeepSeek 用 MLA），典型配置是 $h_q$ 为 $h_{kv}$ 的 4–8 倍。
几何视角：Multi-Head 是把 $d_{\text{model}}$ 空间切成 $h$ 个子空间，每个子空间从输入的某个独立视角看相关性。

下一章我们解决一个第 2 章里挖下的坑：Self-Attention 不感知位置。(animal, it) 和 (it, animal) 在 Self-Attention 看来是一样的。语言里位置非常重要，必须从外部注入。第 4 章会讲位置编码——从原始 sinusoidal 到 RoPE 的完整演化路径。