4. 多头注意力机制：结构与并行计算

多头注意力机制，说白了就是把注意力“复制”好几份，让模型从不同角度去看问题。你想想看，单头注意力就像只用一只眼睛看世界，虽然也能看清，但总有些细节会遗漏。多头就不一样了，它相当于同时用多只眼睛，每只关注一个子空间，最后再把信息融合起来。

我刚开始做Transformer模型时，对多头这块的理解其实有点偏差。当时觉得不就是多算几次注意力然后拼起来嘛，能有多大区别？直到有一次做文本分类任务，单头注意力怎么调都学不到某些关键短语的依赖关系。换成8个头之后，效果立竿见影。嗯，从那以后我再也不敢小看多头这个设计了。

4.1 为什么需要多头？

单头注意力有个天然局限：它只能在一个子空间里计算 query、key、value 的相似度。不同的语义关系，比如语法关系、指代关系、主题关系，往往分布在不同的特征子空间里。单头注意力很难同时兼顾这些不同的关系。

多头注意力的核心思路，就是通过多组独立的线性变换，把输入映射到不同的子空间。每个头独立执行注意力计算，最后拼接起来再投影回原始维度。这样模型就能在不同的表示子空间里捕捉多样化的依赖信息。

核心优势

  • 每个头关注不同的特征子空间
  • 联合学习不同类型的注意力模式
  • 不增加序列长度的计算复杂度，仅增加线性投影的计算量

我记得之前有篇论文专门分析过，不同头其实学到了不同的东西：有的头主要关注语法依赖，比如动词和宾语的关系；有的头则更关注位置近邻，比如相邻词之间的交互。这就是多头机制的妙处。

4.2 多头注意力的数学结构

多头注意力可以看作是对单头注意力的一种“批处理”扩展。我们假设头数为 h，每个头的维度为 d_k = d_model / h。整个计算过程分为三步：

  1. 线性投影与分割：对 Q、K、V 分别做 h 组不同的线性变换，得到 h 组低维的 query、key、value
  2. 并行注意力计算：对每组低维数据独立执行缩放点积注意力
  3. 拼接与输出投影：将 h 个头的输出拼接起来，再做一次线性变换得到最终输出

公式其实不复杂，我直接用数学形式写出来：

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) * W^O

其中：
head_i = Attention(Q * W_i^Q, K * W_i^K, V * W_i^V)

投影矩阵：
W_i^Q ∈ R^(d_model x d_k)
W_i^K ∈ R^(d_model x d_k)
W_i^V ∈ R^(d_model x d_v)
W^O   ∈ R^(h * d_v x d_model)

这里有个细节需要注意：每个头的 d_k 和 d_v 通常等于 d_model / h。这样一来，多头注意力的总计算量跟单头（维度为 d_model）基本持平。

避坑指南

我曾经在实现时把每个头的维度直接设成 d_model，结果计算量暴增了 h 倍，显存直接爆了。记住，每个头的维度是 d_model/h，不是 d_model。

4.3 并行计算的实现技巧

多头注意力的“并行”主要体现在两点：头之间的并行和批次内的并行。实际实现时，我们通常不会真的用循环去算每个头，而是通过张量变形一次性搞定。

我个人习惯用 PyTorch 实现时，把维度变换做在 batch 维度后面。具体步骤是这样的：

  1. 输入 Q、K、V 的形状是 [batch_size, seq_len, d_model]
  2. 通过线性投影得到 Q'、K'、V'，形状不变
  3. 将 Q'、K'、V' reshape 成 [batch_size, seq_len, h, d_k]
  4. 把 h 维度放到 batch 维度后面，变成 [batch_size, h, seq_len, d_k]
  5. 对 [batch_size * h, seq_len, d_k] 统一做注意力计算
  6. 算完后恢复形状，拼接，再投影

这种实现方式，本质上是把 h 个头当成 batch 维度来处理。GPU 特别喜欢这种大矩阵乘法，比循环快得多。

代码实现大概是这样的：

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, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)

        # 线性投影 + 变形分头
        Q = self.w_q(Q).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        K = self.w_k(K).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        V = self.w_v(V).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)

        # 批量注意力计算
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)

        # 拼接 + 输出投影
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.h * self.d_k)

        return self.w_o(attn_output)

注意：transpose 之后一定要做 .contiguous()，否则 view 会报错。我当年被这个坑过，debug 了整整一个下午。

4.4 头数的选择与影响

头数 h 是个超参数，常见的选择有 8、12、16。头数太多或太少都有问题：

头数 单头维度 特点 常见场景
4 d_model/4 每个头维度较大，但多样性不足 小模型、轻量级部署
8 d_model/8 平衡好，最常用的配置 BERT base、GPT-2
12 d_model/12 多样性高，但单头维度较小 BERT large、GPT-3
16 d_model/16 每个头维度很小，可能信息丢失 超大模型

你想想看，如果 d_model=768，头数设为 16，每个头只有 48 维。这么低的维度，能学到什么丰富的模式吗？我个人的经验是，头数最好不要超过 d_model/64，否则每个头的表达能力会严重受限。

实用建议

如果你刚开始做实验，直接上 8 个头基本不会错。等模型稳定了，可以用注意力图分析工具看看每个头到底在关注什么，再决定要不要调头数。

4.5 多头注意力核心流程

下面这张 SVG 图把多头注意力的完整流程画出来了，你可以对照着理解整个过程：

多头注意力机制核心流程 输入 Q/K/V 线性投影 头 1 分割 头 2 分割 ... 头 h 分割 并行注意力1 并行注意力2 并行注意力h 拼接 Concat 输出投影 W^O 关键说明 每个头独立计算注意力 不共享参数 拼接后通过 W^O 融合 总计算量与单头持平

这张图展示的是标准的多头注意力流程。从左到右看：输入先做线性投影，然后分割成 h 个头。每个头独立做注意力计算，互不干扰。最后把 h 个头的输出拼起来，再做一次线性投影得到最终结果。

箭头部分我用了不同颜色区分不同头的路径，但核心思想就是一个：分而治之，并行计算，最后融合。

4.6 实现中的常见问题

最后我总结几个实现多头注意力时容易踩的坑，都是我亲身经历过的：

维度陷阱

  • d_model 必须能被 h 整除，否则 d_k 不是整数，会报错
  • transpose 之后一定不要忘了 .contiguous()
  • mask 的形状要适配 [batch_size, h, seq_len, seq_len] 这个维度

调试小技巧

我曾经在第一次写多头注意力时，死活不收敛。后来我打印出每个头的注意力权重矩阵，发现好几个头学出来的模式一模一样。原来是我初始化有问题，线性投影的参数太相似了。用不同的随机种子初始化每组的 W^Q、W^K、W^V 就能解决。

嗯，多头注意力机制的核心内容大概就这些。你只要能理解并行计算和子空间分割这两个点，后面看什么变体都会很轻松。

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