查看全课大纲（第 20 / 40 节）

多头注意力：并行的多种关系

小象实战讲义 · 深度学习基础：从神经网络到 Transformer

在上一节，我们学习了注意力机制的核心思想：通过查询（Query）、键（Key）、值（Value）的交互，动态地为每个位置分配权重，从而捕捉序列内部的依赖关系。但单一注意力只能学习一种模式的关系。本节我们将深入 Transformer 的多头注意力（Multi-Head Attention）机制，它通过并行运行多个注意力“头”，让模型能够同时关注来自不同表示子空间的信息，从而更强大地建模复杂关系。学完本节，你将能回答：多头注意力如何并行计算？为什么它能提升模型表达能力？在代码中如何实现和调用？

先修：缩放点积注意力的四步与 F.scaled_dot_product_attention 的用法（本章 5.1 节）、nn.Linear 的权重形状。后续：下一节讲位置编码，再往后讲残差、LayerNorm 与前馈层，最后用这些零件从零拼出一个完整的 Transformer block。

💡 核心导读

多头注意力的定义：将输入投影到多个不同的子空间，在每个子空间独立进行缩放点积注意力计算，最后将结果合并。
并行计算流程：通过线性投影、重塑、批量矩阵乘法实现高效并行，是所有头同时计算，而非循环。
形状变换全追踪：从 (batch, seq, d_model) 到 (batch, num_heads, seq, d_k)，再到合并回原始维度。
设计动机：允许模型共同关注来自不同位置的不同表示子空间的信息，类比于卷积神经网络中使用多个滤波器。
PyTorch 官方实现：使用 nn.MultiheadAttention 模块，理解其参数与默认行为。

🧭 前置知识

缩放点积注意力：公式 $\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V$ 的计算过程。
矩阵乘法形状规则：(a, b) @ (b, c) -> (a, c)，以及批量矩阵乘法的扩展。
张量重塑与转置：view()、transpose() 操作对形状的改变。
线性投影（全连接层）：nn.Linear 的作用，即 $y = xW^T + b$ 。

从单头到多头：核心思想

在《Attention Is All You Need》（Vaswani et al., 2017，arXiv:1706.03762）论文中，作者指出，与其使用一个 $d_{\text{model}}$ 维的键、值和查询执行单一的注意力函数，不如将查询、键和值用不同的、学习到的线性投影分别投影到 $d_k$ 、 $d_k$ 和 $d_v$ 维 $h$ 次（即 $h$ 个头）。然后在这些投影后的版本上并行执行注意力函数，得到 $h$ 个 $d_v$ 维的输出，最后将它们拼接起来并再次投影，得到最终值。

为什么需要多头？ 单一注意力机制在计算权重时，相当于将所有特征维度混合在一起进行相似度计算。这可能限制了模型同时关注不同类型关系的能力。例如，在一个句子中，我们可能同时需要关注语法结构（如主谓一致）和语义关联（如同义词指代）。多头注意力提供了多个独立的“表示子空间”，让每个头可以专注于学习一种特定的关系模式，最后再综合起来。

数学上，多头注意力定义为： $\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$ 其中 $\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$ 。这里，投影是参数矩阵： $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$ , $W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ , $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ ，以及输出投影 $W^O \in \mathbb{R}^{h \cdot d_v \times d_{\text{model}}}$ 。

在实现中，通常设 $d_k = d_v = d_{\text{model}} / h$ ，使得多头的总计算量与单头全维度的注意力大致相当，同时增加了并行度和表示能力。

实现详解：形状变换与并行计算

让我们通过一个完整的、可运行的代码示例，一步步拆解多头注意力的计算过程。我们将手动实现关键步骤，并与 PyTorch 官方实现进行对比。

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)  # 固定随机种子，确保可复现

# ========== 1. 定义超参数与模拟输入数据 ==========
batch_size = 2
seq_len = 5
d_model = 64        # 模型隐藏维度（也是输入输出维度）
num_heads = 8       # 注意力头的数量
d_k = d_v = d_model // num_heads  # 每个头的维度，必须能整除

print(f"超参数: batch={batch_size}, seq={seq_len}, d_model={d_model}, heads={num_heads}, d_k=d_v={d_k}")

# 模拟输入张量 x，形状 (batch, seq, d_model)
x = torch.randn(batch_size, seq_len, d_model)
print(f"输入 x.shape = {tuple(x.shape)}")

# ========== 2. 线性投影：生成 Q, K, V ==========
# 使用 nn.Linear 定义投影层。注意：实际中这些是模型的可学习参数。
# 这里我们创建独立的线性层来模拟 W_i^Q, W_i^K, W_i^V。
# 为了方便，我们一次性投影到所有头的总维度 (num_heads * d_k)。
W_Q = nn.Linear(d_model, num_heads * d_k, bias=False)
W_K = nn.Linear(d_model, num_heads * d_k, bias=False)
W_V = nn.Linear(d_model, num_heads * d_v, bias=False)

# 初始化权重（为了可复现性，使用 Xavier 初始化）
nn.init.xavier_uniform_(W_Q.weight)
nn.init.xavier_uniform_(W_K.weight)
nn.init.xavier_uniform_(W_V.weight)

# 执行投影: (batch, seq, d_model) -> (batch, seq, num_heads * d_k)
Q_proj = W_Q(x)  # 对应 Q * W^Q（合并了所有头）
K_proj = W_K(x)  # 对应 K * W^K
V_proj = W_V(x)  # 对应 V * W^V

print(f"投影后 Q_proj.shape = {tuple(Q_proj.shape)}")

# ========== 3. 重塑与转置：分离出"头"维度 ==========
# 目标形状: (batch, seq, num_heads, d_k) -> (batch, num_heads, seq, d_k)
# 这样可以将"头"维度放在 batch 之后，便于后续批量矩阵乘法。
new_shape_Q = (batch_size, seq_len, num_heads, d_k)
Q_heads = Q_proj.view(new_shape_Q).transpose(1, 2)  # 交换 seq 和 num_heads 维度
K_heads = K_proj.view(new_shape_Q).transpose(1, 2)
V_heads = V_proj.view(batch_size, seq_len, num_heads, d_v).transpose(1, 2)

print(f"重塑后 Q_heads.shape = {tuple(Q_heads.shape)}")  # (batch, num_heads, seq, d_k)

# ========== 4. 计算缩放点积注意力（并行所有头） ==========
# 注意力分数: (batch, num_heads, seq, d_k) @ (batch, num_heads, d_k, seq) -> (batch, num_heads, seq, seq)
attn_scores = torch.matmul(Q_heads, K_heads.transpose(-2, -1)) / (d_k ** 0.5)
attn_weights = F.softmax(attn_scores, dim=-1)  # 在最后一个维度（键序列）上做 softmax

# 加权求和: (batch, num_heads, seq, seq) @ (batch, num_heads, seq, d_v) -> (batch, num_heads, seq, d_v)
head_outputs = torch.matmul(attn_weights, V_heads)

print(f"注意力权重 attn_weights.shape = {tuple(attn_weights.shape)}")
print(f"各头输出（合并前） head_outputs.shape = {tuple(head_outputs.shape)}")

# 快速验证：注意力权重的每一行和应为 1（因为对最后一个维度做 softmax）
row_sums = attn_weights.sum(dim=-1)[0, 0, :]  # 取第一个 batch，第一个头，所有查询位置
print(f"注意力权重行和（第一个头）: {[round(v.item(), 4) for v in row_sums]}")

# ========== 5. 合并多头输出 ==========
# 将头维度移回并拼接: (batch, num_heads, seq, d_v) -> (batch, seq, num_heads, d_v) -> (batch, seq, num_heads*d_v)
concat_output = head_outputs.transpose(1, 2).contiguous().view(batch_size, seq_len, num_heads * d_v)
print(f"拼接后 concat_output.shape = {tuple(concat_output.shape)}")

# ========== 6. 最终输出投影 ==========
W_O = nn.Linear(num_heads * d_v, d_model, bias=False)
nn.init.xavier_uniform_(W_O.weight)
output = W_O(concat_output)  # (batch, seq, num_heads*d_v) -> (batch, seq, d_model)

print(f"最终输出 output.shape = {tuple(output.shape)}")
print(f"输出与输入形状一致: {output.shape == x.shape}")

超参数: batch=2, seq=5, d_model=64, heads=8, d_k=d_v=8
输入 x.shape = (2, 5, 64)
投影后 Q_proj.shape = (2, 5, 64)
重塑后 Q_heads.shape = (2, 8, 5, 8)
注意力权重 attn_weights.shape = (2, 8, 5, 5)
各头输出（合并前） head_outputs.shape = (2, 8, 5, 8)
注意力权重行和（第一个头）: [1.0, 1.0, 1.0, 1.0, 1.0]
拼接后 concat_output.shape = (2, 5, 64)
最终输出 output.shape = (2, 5, 64)
输出与输入形状一致: True

关键步骤解析：

线性投影：输入 x 分别通过三个不同的线性层，生成 Q_proj、K_proj、V_proj。注意，这里一次性投影到了 num_heads * d_k 维，包含了所有头的投影结果。
重塑与转置：通过 view() 和 transpose() 将形状从 (batch, seq, num_heads*d_k) 变为 (batch, num_heads, seq, d_k)。这里有两个容易混的细节：transpose 只改 stride 不搬内存，而 view 要求内存连续，所以反过来合并时（transpose 之后又要 view）必须先 contiguous()；把“头”维度提到序列维度之前，是为了让后续 matmul 的批量矩阵乘法一次性覆盖 batch 和 head 两维。
并行注意力计算：attn_scores 的计算使用了批量矩阵乘法，形状为 (batch, num_heads, seq, seq)。这是并行的核心：我们不是用循环依次计算 8 个头，而是一次性计算出所有头的注意力分数和权重。如果换成 PyTorch 2.x 的写法，第 4 步整段（缩放、softmax、加权求和）可以一行 F.scaled_dot_product_attention(Q_heads, K_heads, V_heads) 完成——它做的正是 5.1 节那四步，只是逐头批量执行。
合并输出：计算完每个头的输出后，通过 transpose() 和 view() 将形状变回 (batch, seq, num_heads*d_v)，然后通过最后一个线性层 W_O 投影回 d_model 维。

使用 PyTorch 官方模块

在实际开发中，我们直接使用 PyTorch 内置的 nn.MultiheadAttention 模块。它封装了上述所有步骤，并进行了高度优化。

import torch
import torch.nn as nn

torch.manual_seed(0)

# 参数与之前一致
batch_size = 2
seq_len = 5
d_model = 64
num_heads = 8

x = torch.randn(batch_size, seq_len, d_model)

# 实例化 MultiheadAttention
# 注意：embed_dim 必须是 num_heads 的整数倍
# batch_first=True 表示输入输出形状为 (batch, seq, feature)，否则为 (seq, batch, feature)
mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=num_heads, batch_first=True, dropout=0.0)

# 前向传播
# 在自注意力（self-attention）中，query, key, value 都等于输入 x
attn_output, attn_output_weights = mha(x, x, x, need_weights=True)

print("PyTorch MHA 输出形状:", tuple(attn_output.shape))
print("PyTorch MHA 注意力权重形状:", tuple(attn_output_weights.shape))
# 注意：attn_output_weights 的形状是 (batch, seq, seq)，它是所有头注意力权重的平均值。
# 若想看每个头的权重，把 average_attn_weights 设为 False（默认是 True）。

PyTorch MHA 输出形状: (2, 5, 64)
PyTorch MHA 注意力权重形状: (2, 5, 5)

nn.MultiheadAttention 模块会自动处理线性投影、重塑、注意力计算和输出投影。当 need_weights=True 时，返回的 attn_output_weights 是所有注意力头的权重平均值，形状为 (batch, seq, seq)——这是 average_attn_weights=True（默认值）的结果。想拿到 (batch, num_heads, seq, seq) 的逐头权重，调用时传 average_attn_weights=False 即可，不需要自己重写实现。

顺带澄清一个常被误解的点：多头并不比单头更“贵”。总投影维度固定为 $d_{\text{model}}$ ，每个头的 $d_k = d_{\text{model}}/h$ ，所以 $h$ 个头的 Q/K/V 投影参数量与单头完全相同；注意力矩阵的规模看起来是 $h \times (\text{seq}\times\text{seq})$ ，但每个头的 $d_k$ 缩小了 $h$ 倍，总浮点运算量仍在同一量级。多头换来的是表示能力：不同头可以各自关注位置邻近、句法依赖、指代关系等不同模式，而不是逼着一个头同时兼顾所有关系。

还有一个实现细节：nn.MultiheadAttention 的 in_proj_weight 形状是 $(3\,d_{\text{model}}, d_{\text{model}})$ ——Q、K、V 三组投影权重被拼成一个矩阵，一次矩阵乘就同时算出三者。这解释了为什么在参数列表里找不到单独的 q_proj_weight；自己手写实现时拆成三个 nn.Linear 更直观，两者数学等价。

为什么除以 $\sqrt{d_k}$ 在多头上依然重要？

在单头注意力中，我们解释了缩放因子 $\frac{1}{\sqrt{d_k}}$ 的作用：防止点积结果过大导致 softmax 梯度消失。在多头注意力中，每个头的 $d_k$ 通常比原始的 $d_{\text{model}}$ 小得多（例如 $d_{\text{model}}=512$ , $h=8$ , $d_k=64$ ）。但即便如此，缩放仍然至关重要。

因为每个头的查询和键是独立投影的，点积的方差仍然与 $d_k$ 成正比（假设投影权重初始化得当）。如果不缩放，对于较大的 $d_k$ （即使只有 64），点积的绝对值也可能很大，将 softmax 推入梯度极小的饱和区，影响训练稳定性。因此，在每个头的注意力计算中，都必须进行缩放。

📝 动手练一练

题目：假设我们将 num_heads 从 8 减少到 4，同时保持 d_model=64 不变。请手动计算新的 d_k 和 d_v，并修改上面第一个代码块中的相关参数，重新运行。观察 Q_heads 和 attn_weights 的形状发生了什么变化？最终输出的形状是否保持不变？

👉 点击查看参考答案

思路：d_k = d_v = d_model // num_heads。减少头数，每个头的维度会增加。注意力计算的核心步骤不变，但中间张量在“头”维度上的大小会改变。

修改与运行：

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

batch_size, seq_len, d_model = 2, 5, 64
num_heads = 4                      # 从 8 改为 4
d_k = d_v = d_model // num_heads   # 每个头变成 16 维

x = torch.randn(batch_size, seq_len, d_model)

W_Q = nn.Linear(d_model, num_heads * d_k, bias=False)
W_K = nn.Linear(d_model, num_heads * d_k, bias=False)
W_V = nn.Linear(d_model, num_heads * d_v, bias=False)
for w in (W_Q, W_K, W_V):
    nn.init.xavier_uniform_(w.weight)

# 投影 + 拆头：(batch, seq, heads*d_k) -> (batch, heads, seq, d_k)
Q_heads = W_Q(x).view(batch_size, seq_len, num_heads, d_k).transpose(1, 2)
K_heads = W_K(x).view(batch_size, seq_len, num_heads, d_k).transpose(1, 2)
V_heads = W_V(x).view(batch_size, seq_len, num_heads, d_v).transpose(1, 2)
print(f"Q_heads.shape = {tuple(Q_heads.shape)}")

# 完整的注意力权重（所有头一起算）
attn_weights = F.softmax(
    torch.matmul(Q_heads, K_heads.transpose(-2, -1)) / (d_k ** 0.5), dim=-1
)
print(f"attn_weights.shape = {tuple(attn_weights.shape)}")
print(f"每个头内部仍是 (seq, seq) = {tuple(attn_weights[0, 0].shape)}")

head_outputs = torch.matmul(attn_weights, V_heads)
concat = head_outputs.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
print(f"合并回 (batch, seq, d_model) = {tuple(concat.shape)}")

Q_heads.shape = (2, 4, 5, 16)
attn_weights.shape = (2, 4, 5, 5)
每个头内部仍是 (seq, seq) = (5, 5)
合并回 (batch, seq, d_model) = (2, 5, 64)

结论：头数从 8 减到 4，d_k 从 8 变成 16，Q_heads 从 (2, 8, 5, 8) 变成 (2, 4, 5, 16)，完整的 attn_weights 也从 (2, 8, 5, 5) 变成 (2, 4, 5, 5)——头这一维变小了，每个头内部的 (5, 5) 注意力矩阵并没有变。合并后的形状始终是 (batch, seq, d_model)，这正是多头设计的关键：改变内部表示方式，不改变输入输出的接口维度。

常见错误

错误理解或写法	现象或报错	原因	正确做法
`d_model` 不能被 `num_heads` 整除	`RuntimeError: embed_dim must be divisible by num_heads`	PyTorch 的 `nn.MultiheadAttention` 要求每个头的维度是整数。	设计模型时确保 `embed_dim % num_heads == 0`。例如 `d_model=512`, `num_heads=8`（可整除），`num_heads=7`（不可整除）。
忘记在注意力分数上除以 $\sqrt{d_k}$	模型可能训练不稳定，收敛慢或效果差。	点积值过大，softmax 梯度极小，阻碍学习。	在计算 `attn_scores = Q @ K.T` 后，务必执行 `attn_scores = attn_scores / (d_k ** 0.5)`。
重塑张量时形状计算错误	`RuntimeError: shape '[2, 5, 8, 8]' is invalid for input of size ...`	`view()` 要求新形状的元素总数必须与原张量一致。计算错误会导致总数对不上。	仔细计算：`batch * seq * num_heads * d_k` 必须等于输入张量的 `numel()`。使用 `x.view(batch, seq, num_heads, -1)` 让 PyTorch 自动推断最后一维。
混淆 `batch_first` 参数	输入形状与模块预期不匹配，导致维度错误。	`nn.MultiheadAttention` 默认输入形状为 `(seq, batch, dim)`。	如果数据是 `(batch, seq, dim)`，创建模块时务必设置 `batch_first=True`。
误以为 `attn_output_weights` 是每个头的权重	得到的权重形状是 `(batch, seq, seq)`，而不是 `(batch, num_heads, seq, seq)`。	`average_attn_weights` 默认为 `True`，模块把所有头平均成一矩阵返回。	需要逐头权重时传 `average_attn_weights=False`；只做整体可视化时用默认的平均权重即可。

本章小结

多头注意力的核心：通过多组独立的线性投影，将输入映射到多个子空间，并行计算注意力，最后合并结果，以捕捉不同类型的依赖关系。
并行计算流程：线性投影 → 重塑分离头维度 → 批量矩阵乘法计算缩放点积注意力 → 合并头 → 输出投影。关键形状变换为 (batch, seq, d_model) ↔ (batch, num_heads, seq, d_k)。
缩放因子的必要性：即使每个头的 $d_k$ 较小，除以 $\sqrt{d_k}$ 仍是稳定训练、防止梯度消失的关键步骤。
PyTorch 官方模块：使用 nn.MultiheadAttention(embed_dim, num_heads, batch_first=...)，它封装了所有步骤并进行了优化。注意其默认输入格式和返回的注意力权重是平均值。
设计约束：embed_dim 必须能被 num_heads 整除，以确保每个头维度一致。

📋 行动清单

运行本节第一个完整代码块，在打印输出的每个步骤暂停，对照代码理解每一步张量形状的变化。
将 num_heads 改为 1（单头），再改为 d_model（即 d_k=1），重新运行并观察形状变化，思考极端情况下的意义。
查阅 PyTorch 官方文档中 nn.MultiheadAttention 的参数说明，特别是 add_bias_kv、add_zero_attn 等选项的作用（本节未展开）。