3、多头注意力机制：为什么需要多头？如何并行计算？
好，咱们接着聊。上一节我们把自注意力的计算流程捋了一遍，你可能会想：一个注意力头够用吗？
说实话，不够。我在早期做文本分类项目时，就吃过这个亏。单头注意力就像一个人看问题，视角太单一了。它只能关注到一种关系模式，比如词与词之间的语法依赖，但语义、位置、情感这些信息就顾不上了。
为什么需要多头？
说白了，就是让模型从多个角度去理解输入。你想想看，一句话里包含的信息是多维的：
- 有的头关注语法结构（主语-谓语-宾语）
- 有的头关注语义相似性（同义词、上下位词）
- 有的头关注位置关系（距离远近、顺序）
- 有的头关注情感色彩（正面/负面词汇）
单头注意力只能学到一种加权方式。多头注意力呢？它把查询、键、值分别投影到多个子空间，每个子空间独立做注意力计算。最后再把结果拼起来。
核心公式：
MultiHead(Q, K, V) = Concat(head₁, head₂, ..., headₕ) · Wᴼ
其中 headᵢ = Attention(Q · Wᵢᑫ, K · Wᵢᴷ, V · Wᵢⱽ)
每个头都有自己的投影矩阵 Wᵢᑫ、Wᵢᴷ、Wᵢⱽ。这些矩阵把输入映射到不同的子空间。我习惯把每个头想象成一个「专家」，各自负责不同的分析维度。
举个例子，假设输入是「我喜欢吃苹果」：
- 头1可能关注「我」和「喜欢」的主谓关系
- 头2可能关注「苹果」和「吃」的动宾关系
- 头3可能关注「喜欢」和「苹果」的情感关联
你看，同一个句子，不同头看到的东西完全不一样。最后拼接起来，模型就拥有了「多视角」的理解能力。
我的经验：头数不是越多越好。我在一个翻译任务中试过32个头，结果训练慢、过拟合严重。后来降到8个头，效果反而更好。一般8-16个头是比较稳妥的选择。
如何并行计算？
嗯，这里有个关键问题：多个头怎么算？一个一个算吗？那太慢了。
实际上，多头注意力是完全并行的。怎么做？把多个头的投影矩阵拼成一个大矩阵，一次矩阵乘法搞定所有头的投影。
具体来说：
- 把 Wᵢᑫ、Wᵢᴷ、Wᵢⱽ 分别拼成 Wᑫ、Wᴷ、Wⱽ（形状为 d_model × (h · d_k)）
- 一次矩阵乘法得到所有头的 Q、K、V
- 然后对每个头独立计算注意力
- 最后拼接结果
看代码更直观：
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, h):
super().__init__()
self.h = h
self.d_k = d_model // h
# 一次性定义所有头的投影矩阵
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def forward(self, Q, K, V, mask=None):
batch_size = Q.size(0)
# 1. 线性投影，得到所有头的Q、K、V
Q = self.W_q(Q) # (batch, seq_len, d_model)
K = self.W_k(K)
V = self.W_v(V)
# 2. 拆分成多个头
# (batch, seq_len, h, d_k) -> (batch, h, seq_len, d_k)
Q = Q.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
K = K.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
V = V.view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
# 3. 计算注意力分数（所有头同时计算）
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn_weights = torch.softmax(scores, dim=-1)
# 4. 加权求和
context = torch.matmul(attn_weights, V)
# 5. 拼接所有头的结果
context = context.transpose(1, 2).contiguous().view(
batch_size, -1, self.h * self.d_k
)
# 6. 最后的线性变换
output = self.W_o(context)
return output注意看第2步：我们用 view 和 transpose 把张量从 (batch, seq_len, d_model) 变成 (batch, h, seq_len, d_k)。这样每个头的数据就独立了，但计算时还是一起算的。
避坑指南：我曾经在实现时忘记做 contiguous() 操作，导致 view 报错。transpose 后张量内存不连续，必须调用 contiguous() 才能 reshape。这个小坑花了我半天时间排查。
头数怎么选？
这个没有标准答案，但有几个经验法则：
| 模型规模 | 推荐头数 | d_k（每个头的维度） |
|---|---|---|
| 小模型（d_model=256） | 4-8 | 32-64 |
| 中等模型（d_model=512） | 8-16 | 32-64 |
| 大模型（d_model=1024） | 16-32 | 32-64 |
我个人习惯保持 d_k 在 32-64 之间。如果 d_k 太小，每个头能表达的信息就有限；太大又浪费计算资源。
还有一个细节：d_model 必须能被 h 整除。否则最后一个头维度不一样，拼接时会出问题。
总结一下
多头注意力说白了就是「分而治之」：
- 把输入投影到多个子空间
- 每个子空间独立做注意力
- 最后拼接融合
并行计算的关键在于矩阵运算。把所有头的投影合并成一次大矩阵乘法，再通过 reshape 和 transpose 来分离各个头。这样既利用了 GPU 的并行能力，又实现了多视角分析。
下一节我们会聊位置编码。嗯，没有位置编码的 Transformer 就像没有坐标的地图——词与词之间不知道谁先谁后，这问题可不小。