4. Multi-Head Attention：多头机制的设计思想，并行计算的优势

好，咱们接着聊。前面我们讲了 Scaled Dot-Product Attention，说白了就是一个 query 跟一堆 key 算相似度，然后加权求和 value。但这里有个问题——你想想看，一个注意力头，它只能关注一种关系模式。

比如在句子 "我昨天在公园里看到了一只猫" 里，"猫" 这个词可能跟 "看到" 有语法关系，跟 "公园" 有位置关系，跟 "昨天" 有时间关系。一个注意力头，它只能抓住其中一种。那怎么办？

嗯，答案就是——多头注意力（Multi-Head Attention）。

4.1 为什么需要多个头？

我个人习惯把多头注意力理解成 "多个专家并行开会"。

每个头就是一个专家，它们各自关注不同的信息维度。有的专家关注语法结构，有的专家关注语义相似度，有的专家关注位置偏移。最后把大家的意见汇总起来，得到一个更全面的理解。

我在项目中遇到过这样一个场景：用单头注意力做文本分类，效果一直上不去。后来换成 8 头注意力，同样的数据、同样的训练轮次，准确率直接提升了 3 个点。为什么？因为单头只能看到一种模式，而多头能看到多种模式。

核心思想： 多头注意力允许模型在不同的表示子空间里学习信息。每个头学习不同的投影矩阵 WQ、WK、WV，从而关注不同的特征。

4.2 多头机制的计算流程

说白了，多头注意力的计算分三步走：

  1. 线性投影 + 拆分：把 Q、K、V 分别通过 h 组不同的线性变换，得到 h 组 Qi、Ki、Vi
  2. 并行计算注意力：对每组 Qi、Ki、Vi 独立执行 Scaled Dot-Product Attention
  3. 拼接 + 再投影：把 h 个头的输出拼接起来，再经过一次线性变换得到最终结果

公式长这样：

MultiHead(Q, K, V) = Concat(head₁, head₂, ..., headₕ) · Wᴼ

其中 headᵢ = Attention(Q · Wᵢ^Q, K · Wᵢ^K, V · Wᵢ^V)

这里 h 是头的数量，Wᵢ^Q、Wᵢ^K、Wᵢ^V 是第 i 个头的投影矩阵，Wᴼ 是输出投影矩阵。

4.3 并行计算的优势

你可能会问：那为什么不直接用一个更大的单头注意力？

好问题。我刚开始学的时候也有这个疑惑。后来我意识到，并行计算才是多头注意力的灵魂。

假设我们有一个 512 维的输入，想用 8 个头。每个头处理 64 维的子空间。这 8 个头可以完全并行计算——GPU 上同时跑 8 个注意力操作，互不干扰。

如果换成单头，你只能串行地尝试不同的模式，或者用一个巨大的矩阵去拟合所有模式。前者慢，后者容易过拟合。

避坑指南： 我曾经在调参时把头的数量设得太大（比如 32 个头），结果每个头分到的维度太少（只有 16 维），注意力质量反而下降了。一般来说，头的数量 h 和每个头的维度 dₖ 要平衡，常见配置是 h=8, dₖ=64。

4.4 头的数量怎么选？

嗯，这里要注意，头的数量不是越大越好。我整理了一个经验表格：

模型规模 推荐头数 每个头维度 适用场景
小模型（d_model=256） 4 64 轻量级任务、移动端部署
中等模型（d_model=512） 8 64 通用 NLP 任务
大模型（d_model=1024） 16 64 机器翻译、文本生成
超大模型（d_model=2048） 32 64 大规模预训练

你会发现，每个头的维度 dₖ 通常固定在 64 左右。这是 Transformer 论文里的经典设置，实践证明效果不错。

4.5 代码实现要点

用 PyTorch 实现多头注意力，核心就几行：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # 一次性生成所有头的 Q、K、V 投影
        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 = self.W_Q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_K(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_V(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 2. 并行计算注意力（这里调用之前实现的 scaled_dot_product_attention）
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)
        
        # 3. 拼接 + 再投影
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        output = self.W_O(attn_output)
        
        return output

注意看，这里用 view + transpose 把矩阵拆成多头，然后一次性传入注意力函数。GPU 会并行处理所有头的计算，效率非常高。

注意事项： 拆分多头时，一定要保证 d_model 能被 num_heads 整除。否则最后一个头的维度会不对齐，导致维度错误。我刚开始写代码时就踩过这个坑，debug 了半天才发现是整除问题。

4.6 多头注意力的直观理解

你想想看，如果把单头注意力比作一个拿着放大镜的人，他只能看到局部细节。那么多头注意力就是一群拿着不同工具的人——有人拿望远镜看远处，有人拿显微镜看细节，有人拿夜视仪看暗处。他们同时工作，最后把观察结果汇总起来。

这就是多头注意力的精髓：并行地、多角度地理解输入信息。

在实际应用中，不同头确实会学到不同的模式。比如在机器翻译中，有的头关注主语-谓语关系，有的头关注形容词-名词修饰关系，有的头关注长距离依赖。这些信息被拼接在一起，让模型的理解更加全面。

好了，多头注意力就讲到这里。下一章我们会聊聊位置编码——毕竟 Transformer 没有循环结构，怎么让模型知道词的顺序呢？到时候见。

← 上一章 📖 返回目录 下一章 →