搞懂多头注意力：先投影，再分头，这一步顺序为什么不能乱

搞懂多头注意力：先投影，再分头，这一步顺序为什么不能乱？

面向LLM入门者，结合PPT数据流、张量形状、正反两种方案对比，把「分头到底是切权重还是切数据」讲透。
配套参数：$B=1,T=512,d_{\text{model}}=512,H=8,d_k=64$（Kronos-small 示例）

前言

很多同学初学Transformer多头注意力（MHA）时，很容易卡在一个问题：

多头是先把输入X切成多份，各自做QKV投影？还是先统一投影出完整QKV，再切分数据分给多个头并行计算？

这个顺序不是随便写的，是LLaMA、Qwen、GPT等主流大模型统一选择的范式：先全局线性投影，再对投影后的张量做数据拆分（分头），并行计算注意力，最后合头。
而且一个核心结论先放在前面：分头只是张量视图重塑，不拆分权重，权重始终只有一套Wq/Wk/Wv。

先回顾单头注意力的基础框图：
输入 $X$，并行 3 个独立 Linear 层得到 Q/K/V，送入同一个注意力引擎：
$$\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$
最后再过一个输出投影。
Q/K/V名字由用途决定，结构上它们只是3个并排的线性层。多头注意力，本质就是把这一套单头注意力引擎复制多份并行跑，但复制的时机，藏着最容易踩坑的知识点。

一、标准MHA完整数据流：先投影，后分头（主流方案）

完整链路

输入X [B, T, d_model]
    ↓
Wq、Wk、Wv 三套线性投影（权重只有一套，不拆分）
    ↓
Q_all、K_all、V_all [B, T, H*d_k]
    ↓
view + transpose 【数据分头，零计算，仅重塑张量视图】
    ↓
Q/K/V [B, H, T, d_k]，拆成H组独立的Qi、Ki、Vi
    ↓
H份注意力引擎并行，各自执行缩放点积注意力
    ↓
所有头结果合并（合头）
    ↓
输出投影Wo，送入FFN

分步拆解张量形状（对照参数）

1. 输入：$X\in {\mathbb{R}}^{1\times 512\times 512}$，即 $[B,T,d_{\text{model}}]$
2. 全局投影：
   $$Q_{all}=X{W}_q,K_{all}=X{W}_k,V_{all}=X{W}_v$$
   $W_q,W_k,W_v\in {\mathbb{R}}^{512\times 512}$，输出
   $Q_{all},K_{all},V_{all}\in {\mathbb{R}}^{1\times 512\times 512}$

   重点：这一步，权重没有拆分，一个Wq矩阵内部已经打包了8个头全部需要的Q子空间，前64维是头0的Q，往后依次是头1~头7。

3. 分头（view+transpose）：

   q = q_all.view(B, T, H, d_k).transpose(1,2)
   

   形状变为 $[B,H,T,d_k]=[1,8,512,64]$
   ✅ view/transpose 没有乘法、没有可训练参数，只是重新划分张量维度，只是“换个方式看这份数据”。
   做完之后，等价于8组独立的(Qi,Ki,Vi)，每个头拿到属于自己的64维子向量。
4. 并行注意力：8个注意力引擎，完全复用PPT里蓝色框的单头注意力公式，独立计算注意力分数、softmax、加权V。
5. 合头：把 8 个头的输出拼回 $[1,512,512]$（即 $[B,T,H{d}_k]$），再经过输出投影 $W_o$。

训练时梯度如何更新同一套权重

Q、K、V没有单独的Ground Truth标签，整个模型只有一个全局损失：自回归预测下一个token的交叉熵损失。
梯度从最后的loss反向回流：
每个头算出来的梯度，先汇总叠加，8个头的梯度累加之后，再去更新同一套Wq、Wk、Wv权重。
虽然权重是一套，但是梯度会自然让Wq矩阵内部不同通道片段，学习到不一样的语义匹配模式：有的头擅长局部语法，有的头擅长长距离指代。

通俗比喻：同一个镜头（Wq），一次性拍出一张完整宽幅照片（Q_all），再把照片裁成8块交给8个人检索；8个人的反馈，一起优化这同一个镜头。每一个人（头）都看过完整原图，只是检索时只用自己分到的那一块子特征。

二、反过来：先拆分输入X，再分头各自投影（非主流方案）

很多人直觉上会以为多头是这个逻辑，我们把链路写出来对比：

输入X [B, T, d_model]
    ↓
沿着特征维度切分X，得到H份子输入 X₁~X₈，每份X_i [B,T,64]
    ↓
每个头配备独立的Wq_i / Wk_i / Wv_i
    ↓
各自投影得到Qi、Ki、Vi
    ↓
独立注意力计算
    ↓
合并输出

这个方案，和标准MHA有本质区别：

1. 权重数量：H×3套独立小线性层，8个头就是24个Linear，各个头权重完全隔离，梯度互不影响；
2. 信息限制：每个头从一开始只能看见原始X的一小段64维通道，永远看不到完整token的全部特征；
3. 硬件效率：大量小矩阵GEMM，GPU Tensor Core不友好，计算碎片化，速度慢。

本质：方案B = 方案A 加了一条「块对角」约束

把两种方案写成同一个矩阵就一目了然。方案A 里头 $h$ 的输出是

$$Q_h=X{W}_q[:,64h:64(h+1)]$$

用的是完整 512 维输入，乘 $W_q$ 的第 $h$ 个 64 列块（形状 $512\times 64$）。
方案B 里头 $h$ 的输出是

$$Q_h=X[:,64h:64(h+1)]W_{q,h}$$

只用到 $X$ 的第 $h$ 个 64 维切片（形状 $64\times 64$）。这两者要等价，方案B 拼回的
$512\times 512$ 权重矩阵必须长这样：

$$W_q^{\prime }=\left[\begin{array}{cccc}{W}_{q,0}& 0& \cdots & 0\\ 0& W_{q,1}& \cdots & 0\\ ⋮& ⋮& \ddots & ⋮\\ 0& 0& \cdots & W_{q,7}\end{array}\right]$$

也就是块对角矩阵——所有跨头通道都被强行置零。所以「信息被截断」只是表象，
真正的损失是：方案B 把方案A 的假设空间限制成块对角的一个子集，头之间无法共享
输入通道。方案A 的 $W_q$ 里，任意一个头的输出都可以用到全部 512 维输入，这正是它
表达能力更强的来源。

通俗比喻：先把原图直接撕碎成8块碎片，每一块碎片单独配一个镜头，每个镜头只看碎片，看不到整张图。信息被物理隔断，表达能力天然受限。

补充：还有一种变体：X不切，但是每个头单独一套Wq_i。每个头可以看见完整X，但是参数量爆炸，性价比极低，工业界几乎不用。

三、两种方案核心对比表

项目	方案A：先投影，后分头（标准MHA，LLaMA/Qwen/GPT）	方案B：先切X，分头再投影（非主流）
权重结构	Wq/Wk/Wv 各1套大矩阵，权重共享	H×3套独立小权重矩阵，互相隔离
每个头的信息范围	每个头都来自完整X投影，能看到token全部原始信息	每个头只能拿到X的一部分通道，信息被截断
梯度更新	多个头梯度累加，联合优化同一套权重	每个头只更新自己的权重，互不干扰
QKV 参数量	$3\times 512\times 512=786,432$	$8\times 3\times 64\times 64=98,304$
加输出投影 $W_o$ 后总计	$1,048,576$	$360,448$
GPU计算效率	单个大GEMM，硬件加速友好	大量小GEMM，碎片化，速度慢
表达能力	强，底层共享基础特征，多头协同学习不同关联模式	弱，每个头视野被锁死，很难捕捉复杂依赖

⚠️ 一个容易迷惑的点：方案B参数量看着更小，但是能力上限更低，根源就是信息被提前切开。这也是为什么主流LLM全部采用「先投影，后分头」。

四、Mermaid流程图直观对照

五、延伸：GQA分组查询注意力会不会改变这个顺序？

不会。
现在 LLaMA2/3、Qwen 大量使用的 GQA（分组查询注意力），依然遵守：先全局投影，再分头。
只是 K/V 的输出通道更少——K/V 头数 $H_{kv}<H$，多个 Q 头共享同一份 K/V 张量，
用来减少 KV Cache 显存占用，投影、分头的时序不变。

唯一需要改的是分头时的形状：Q 用 view(B, T, H, d_k)，K/V 用 view(B, T, H_kv, d_k)，
然后把 K/V 在头维度上重复 $H/H_{kv}$ 次（或直接在 Attention 里广播），再参与计算。
可见标准 MHA 只是 $H_{kv}=H$ 的特例。

六、一句话总结

✅ 标准多头注意力的分头，切的是投影之后的数据张量，不是切权重，也不是提前切输入X；
✅ 顺序不能颠倒的核心原因：保证每一个注意力头，都能基于完整的token特征去学习不同的关联模式，同时兼顾GPU大矩阵计算的硬件优势；
✅ QKV三个投影矩阵本身没有单独的监督标签，依靠全局预测损失，自动分化成查询、键、值三种不同的子空间；
✅ 「先切 X 再分头」并不是完全等价于「先投影再分头」——前者要求权重矩阵是块对角的，
假设空间被大幅缩小，这才是它表达能力弱、且无人使用的根本原因。

七、极简 PyTorch 示意代码（标准 MHA）

import torch
import torch.nn as nn

B, T, d_model, H, d_k = 1, 512, 512, 8, 64
w_q = nn.Linear(d_model, H*d_k, bias=False)
w_k = nn.Linear(d_model, H*d_k, bias=False)
w_v = nn.Linear(d_model, H*d_k, bias=False)
w_o = nn.Linear(H*d_k, d_model, bias=False)

x = torch.randn(B, T, d_model)
# 1. 全局投影，一套权重
q_all = w_q(x)
k_all = w_k(x)
v_all = w_v(x)
# 2. 分头：仅重塑数据视图，无权重
q = q_all.view(B, T, H, d_k).transpose(1,2)
k = k_all.view(B, T, H, d_k).transpose(1,2)
v = v_all.view(B, T, H, d_k).transpose(1,2)
# 3. 缩放点积注意力
attn_score = q @ k.transpose(-1,-2) / (d_k**0.5)
attn_weight = attn_score.softmax(dim=-1)
attn_out = attn_weight @ v
# 4. 合头
attn_out = attn_out.transpose(1,2).contiguous().view(B, T, H*d_k)
out = w_o(attn_out)
print(out.shape)   # torch.Size([1, 512, 512])