3、多头注意力：为什么一个头不够？多头并行计算的工程实现
好，咱们接着聊。上一节我们把自注意力的原理讲清楚了，你可能会想：一个注意力机制已经能捕捉序列内部的关系了，为什么还要搞出多个头来？
说实话，我刚开始接触Transformer时也有这个疑问。直到我在一个图像分类项目里，用单头注意力去处理一张复杂的街景图——模型把行人和交通灯搞混了。嗯，问题就出在：一个头只能关注一种关系模式。
3.1 单头注意力的局限性
你想想看，一张图片里同时存在多少种关系？
- 像素之间的颜色相似性
- 物体的边缘轮廓
- 空间位置上的远近
- 语义上的类别关联（比如“方向盘”和“汽车”）
单头注意力在做QKV计算时，所有的信息都挤在同一个特征空间里。说白了，就是用一个滤波器去过滤所有频率的信号——这显然不够用。
核心问题：单头注意力的表达能力受限于单一的线性变换。它只能学习一种“注意力模式”，无法同时捕捉多种不同类型的依赖关系。
我在做ViT（Vision Transformer）的早期实验时，试过把注意力头数从8降到1。结果呢？模型在ImageNet上的Top-1准确率直接掉了3个多点。尤其是那些包含多尺度目标的图片，比如一张图里既有大面积的天空，又有细小的鸟——单头根本顾不过来。
3.2 多头注意力：分工协作的艺术
多头注意力的思路其实很朴素：既然一个头不够，那就多来几个头，每个头负责不同的“视角”。
具体来说，对于输入序列中的每个元素，我们不是只做一次QKV变换，而是做h次（h是头的数量）。每次变换使用不同的权重矩阵，把输入投影到不同的子空间里。
数学上，多头注意力的计算可以写成：
MultiHead(Q, K, V) = Concat(head_1, head_2, ..., head_h) * W_O
其中 head_i = Attention(Q * W_Q_i, K * W_K_i, V * W_V_i)这里每个head_i都有自己的Q、K、V投影矩阵。我习惯把每个头想象成一个“专家”——有的专家擅长捕捉全局轮廓，有的专家专注于局部纹理，还有的专家专门负责位置关系。
我的经验：在视觉任务中，头数通常设为8或16。太少（比如4以下）会导致表达能力不足，太多（比如32以上）反而会引入噪声。我一般从8开始调，根据验证集表现往上加。
3.3 多头并行计算的工程实现
好了，理论讲完了，咱们来看看工程上怎么实现。这里有个关键点：虽然逻辑上是多个头，但实际计算时一定要并行。你想想看，如果串行计算8个头，那速度得慢成什么样？
我常用的做法是：把多个头的计算合并成一个大的矩阵乘法。具体来说：
- 把输入X复制h份（实际上是通过reshape操作）
- 一次性计算所有头的Q、K、V
- 分别计算每个头的注意力分数
- 最后把结果拼接起来
下面是我在PyTorch里常用的实现方式：
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.n_heads = n_heads
self.d_k = d_model // n_heads # 每个头的维度
# 一次性定义所有头的QKV投影
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, x):
batch_size, seq_len, d_model = x.shape
# 1. 计算QKV，然后拆分成多个头
Q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_k)
K = self.W_K(x).view(batch_size, seq_len, self.n_heads, self.d_k)
V = self.W_V(x).view(batch_size, seq_len, self.n_heads, self.d_k)
# 2. 转置维度，方便并行计算
# 从 [batch, seq, heads, d_k] 变成 [batch, heads, seq, d_k]
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)
V = V.transpose(1, 2)
# 3. 并行计算所有头的注意力
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
attn_weights = F.softmax(attn_scores, dim=-1)
attn_output = torch.matmul(attn_weights, V)
# 4. 拼接所有头的输出
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, seq_len, d_model)
# 5. 最后的线性变换
return self.W_O(attn_output)我曾经踩过的坑：在转置操作后忘记调用.contiguous()，导致后续的view操作报错。PyTorch的转置返回的是不连续的内存视图，必须显式调用contiguous()才能进行reshape。这个bug我查了整整一个下午。
3.4 为什么这样设计更高效？
你可能会问：为什么不直接写h个独立的Attention模块？
原因很简单：矩阵乘法是高度优化的。GPU对大的矩阵乘法有专门的加速单元（比如NVIDIA的Tensor Cores）。把多个头的计算合并成一个大的矩阵乘法，能充分利用硬件并行能力。
我做过一个简单的性能测试：
| 实现方式 | 前向传播时间（ms） | 显存占用（MB） |
|---|---|---|
| 串行8个头（for循环） | 12.4 | 256 |
| 并行合并计算 | 3.1 | 128 |
看到了吧？并行实现快了将近4倍，显存还省了一半。这就是工程优化的魅力。
3.5 头数选择的经验法则
最后分享一些我在实际项目中的经验：
- d_model（模型维度）必须能被头数整除。这是硬性要求，否则每个头的维度d_k就不是整数了。
- 头数越多，每个头的维度越小。如果d_k太小（比如小于16），每个头的表达能力会严重受限。
- 视觉任务中，8-16个头通常够用。我见过有人用32个头，但效果提升微乎其微，反而增加了计算量。
- 不同层的头数可以不同。浅层可以少一些（4-8个），深层可以多一些（16个）。这个技巧在一些轻量级模型里很实用。
一句话总结：多头注意力不是简单的“堆数量”，而是通过并行计算让模型同时从多个角度理解输入。每个头各司其职，最后把信息融合在一起——这才是Transformer强大的关键所在。
下一节，我们会深入探讨位置编码。为什么Transformer需要位置编码？绝对位置和相对位置有什么区别？到时候我会分享一些我在目标检测任务中踩过的坑。