多头注意力机制是 Transformer 的核心模块，也是很多深度学习初学者从 RNN 进入 Transformer 架构时最需要啃下来的硬骨头。它要解决的实际问题很明确：单组注意力权重只能刻画一种位置关系，模型没有办法同时捕捉词与词之间多种粒度的关联；多头注意力通过并行多组映射，让不同注意力头各自关注不同的表示子空间，然后把结果拼回去。这篇文章适合刚学完注意力机制基础、准备手撕 Transformer 源码，或者在 PyTorch 里做注意力模块修改和部署的读者。下面按“原理拆解—公式理解—代码实现—验证排查—选型建议”的顺序展开。
1. 先搞清楚多头注意力到底在解决什么问题
1.1 QKV 三件套：查询、键、值分别承担什么角色
不管单头还是多头，注意力机制的核心都围绕着三个向量：Query（查询）、Key（键）、Value（值）。很多人第一次看源码时被这三个字母绕晕，其实可以换成一个很朴素的类比。
Query 像你手里拿着的搜索条件，Key 像每个候选内容贴的标签，Value 是候选内容本身。系统做的事情是拿 Query 去和每个 Key 做相似度匹配，相似度越高的候选，它的 Value 在最终输出里占的权重越大。放到文本里看，一个 token 的 Query 会去和序列里所有 token 的 Key 比对，得到的权重再作用于所有 token 的 Value，最后加权求和得到这个 token 的注意力输出。
在 Transformer 的自注意力场景里，Q、K、V 全部来自同一个输入序列。这意味着每个位置都能看到整个序列的全局信息，而不是像 RNN 那样只能沿着时间步依次传递状态。自注意力的优势在于直接建模任意两个位置之间的关系，距离远近不再是问题。缺点也很明显，计算复杂度是序列长度的平方，后面会单独说。
1.2 单头注意力的限制在哪里
单头注意力只有一组 W_Q、W_K、W_V 投影矩阵，也就是说所有特征都挤在同一套空间里做相似度计算。这样做的问题在于：一句话里同时存在语法关系、语义关联、指代关系、局部搭配等多种信息，一套投影很难同时把这些不同维度的关系都分开刻画。
举一个很直观的例子。在“小明把苹果递给小红，她觉得很高兴”这句话里，“她”到底指谁，需要结合上下文判断；同时“递给”和“苹果”“小红”之间还有动作、对象的语义搭配。单一注意力分布只能输出一组权重，模型被迫把所有关系压缩到同一个概率分布里。这样不是不能用，但表达效率会明显下降，尤其在长序列、复杂任务上。
单头注意力的计算效率也不是最优的。硬件上我们通常希望矩阵运算的维度足够整齐、并行度足够高。一个独立的大注意力头虽然也能算，但不会有多个小头并行计算来得灵活。多头机制天然把向量分成多份，每个头独立走一遍注意力的计算流程，既增加了模型容量，又保持了计算的并行性。
1.3 多头并行的真实意义
多头注意力的做法是：把 Q、K、V 分别通过线性投影映射到多个低维子空间，每个子空间独立计算注意力，最后把所有头的输出拼接起来，再做一次线性投影。
这里容易有个误解：多头并行不是为了增加参数量，而是为了让模型拥有多套“观察角度”。每个头有自己独立的投影矩阵，学习阶段会自动分化出不同的关注模式。有的头可能更关注相邻词，有的头更关注长距离指代，有的头可能倾向于关注句法关系。你不一定能在每个任务里都明确看出某个头“负责什么”，但整体上多头带来的表达能力提升是稳定的。
多头并行还有一个工程上的好处：每个头的维度只有总维度除以头数，单头里的矩阵乘法规模变小，很多运算可以并行执行。在 GPU 上，只要维度切分合理，多头注意力比一个超大单头注意力更容易吃满算力。这也是为什么现在主流 Transformer 基本都用多头配置。
2. 公式拆解：从缩放点积注意力到多头输出拼接
2.1 缩放点积注意力为什么除以 sqrt(d_k)
多头注意力内部的核心计算单元是缩放点积注意力，公式写出来是：
Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ VQ 和 K 的最后一维 d_k 是每个头的维度。Q 和 K 做矩阵乘法之后，得到的点积数值会随着 d_k 增大而变大。原因不复杂：如果 Q 和 K 里的每个元素都来自均值 0、方差 1 的分布，那么 d_k 个独立元素相乘累加后，点积的方差大约是 d_k。维度越大，方差越大，点积值的分布就越分散。
如果直接把这样的大数值送进 softmax，会出现一个经典问题：softmax 在输入绝对值很大时，梯度会变得非常小，模型训练起来很吃力，甚至出现梯度消失。除以 sqrt(d_k) 是为了把点积的方差重新拉回到 1 附近，让 softmax 的输入落在一个梯度相对敏感的区域。这个缩放是 Transformer 论文里非常关键的一个设计点，手写实现时不能省。
2.2 多头切分、并行计算、拼接和输出投影
多头注意力的完整流程可以拆成五步：
- 输入 X 分别经过 W_Q、W_K、W_V 三个投影得到 Q、K、V。
- 把 Q、K、V 从最后一维切分成 num_heads 份，每份是一个头。
- 每个头独立计算缩放点积注意力。
- 把所有头的输出在特征维度上拼接起来。
- 拼接结果经过输出投影 W_O 得到最终输出。
用维度来描述会更清晰。假设输入形状是 (batch_size, seq_len, embed_dim)，embed_dim 是 512，num_heads 是 8，那么每个头的 head_dim 就是 512 / 8 = 64。Q、K、V 都从 (batch, seq, 512) 被 reshape 成 (batch, seq, 8, 64)，再转置成 (batch, 8, seq, 64)。这样每个头就在自己的 64 维子空间里做注意力，互不干扰。
最终输出形状还是 (batch, seq, 512)，不会改变序列长度和特征维度。这一点对搭建深层 Transformer 很重要，因为多头注意力后面通常接残差连接和 LayerNorm，输入输出形状必须保持一致。
2.3 头数、维度、dropout 等参数怎么理解
多头注意力涉及的参数不多，但每个都要理解到位。
embed_dim 是输入特征的维度，也是每个 Q/K/V 投影后的总维度。这个值通常和模型的主维度一致，BERT base 是 768，Transformer base 是 512。
num_heads 是注意力头数。核心约束是 embed_dim 必须能整除 num_heads。Transformer base 用 8 个头，BERT base 用 12 个头。头数越多，每个头能处理的子空间维度越小，单个头的表达能力会变弱；头数太少，又失去了多角度建模的意义。
head_dim 是每个头的计算维度，等于 embed_dim / num_heads。缩放因子 sqrt(head_dim) 用的是这个值，而不是 embed_dim。很多刚手写注意力的人在这里容易写错，直接用 embed_dim 去缩放，结果训练时数值表现异常。
dropout 在注意力里通常加在两个地方。一是注意力权重上，也就是 softmax 之后随机丢弃部分权重；二是残差结构里，这不是多头注意力本身的部分，但实际工程里默认会加上。注意力权重做 dropout 的目的是防止模型过度依赖某些固定位置关系，对泛化有帮助。
3. PyTorch 手写一个可运行的多头注意力模块
3.1 环境准备和代码组织
建议环境是 Python 3.9 以上，PyTorch 2.0 以上，CUDA 版本按自己的显卡驱动来选择。如果没有 GPU，CPU 上跑小样例完全没问题，只是训练大模型时不现实。
代码组织上，我一般会把多头注意力单独放在一个attention.py文件里，方便后续复用。不急着和完整 Transformer 混在一起写，先把注意力模块验证好，再接 Feed Forward、残差和 LayerNorm。
动手前先确认几件事：
- PyTorch 已安装，能正常
import torch。 - 如果是用 GPU，先执行
torch.cuda.is_available()确认环境。 - 输入数据先设计成小形状，比如 batch_size 为 2，序列长度为 8，embed_dim 为 64。
- 不要一上来就模拟 512 长度的大句子，先保证模块能跑通。
3.2 完整实现代码
下面是一个适合学习和二次修改的多头注意力实现。我没有用nn.MultiheadAttention，而是手动写 Q/K/V 投影和注意力计算，这样整个流程透明，后续加 mask、加相对位置编码都更容易。
import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads, dropout=0.0): super().__init__() assert embed_dim % num_heads == 0, \ "embed_dim 必须能被 num_heads 整除" self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.dropout = dropout self.w_q = nn.Linear(embed_dim, embed_dim) self.w_k = nn.Linear(embed_dim, embed_dim) self.w_v = nn.Linear(embed_dim, embed_dim) self.w_o = nn.Linear(embed_dim, embed_dim) def forward(self, x, mask=None): batch_size, seq_len, embed_dim = x.shape q = self.w_q(x) k = self.w_k(x) v = self.w_v(x) # 切分多头并转置为 (batch, num_heads, seq_len, head_dim) q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 缩放点积注意力 scale = self.head_dim ** 0.5 attn_scores = torch.matmul(q, k.transpose(-2, -1)) / scale if mask is not None: attn_scores = attn_scores.masked_fill(mask == 0, float("-inf")) attn_probs = F.softmax(attn_scores, dim=-1) attn_probs = F.dropout(attn_probs, p=self.dropout, training=self.training) out = torch.matmul(attn_probs, v) # 恢复为 (batch, seq_len, embed_dim) out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim) out = self.w_o(out) return out这个实现还有几个可以优化的地方。比如 Q、K、V 三个投影可以合并成一个nn.Linear(embed_dim, embed_dim * 3)再做 split，推理时能减少 kernel 启动次数。但学习阶段分开写更容易理解，也方便单独调试每个投影。
3.3 用小样例验证输出形状和计算正确性
写完之后不要急着集成到 Transformer 里，先跑一个最小样例。
torch.manual_seed(42) batch_size = 2 seq_len = 8 embed_dim = 64 num_heads = 8 x = torch.randn(batch_size, seq_len, embed_dim) mha = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=0.1) out = mha(x) print(out.shape) # 期望 torch.Size([2, 8, 64])正常情况下输出形状应该和输入完全一致。这里最容易犯的错误是忘了contiguous()，直接view会报错；或者输出维度写错，把多头维度漏掉了。
验证计算正确性可以分两步。第一步只检查形状，第二步检查梯度。把输出对输入求梯度，确认没有出现全零梯度：
loss = out.mean() loss.backward() print(x.grad is not None) print(x.grad.shape)如果梯度是正常的，说明前向和反向通道基本没有问题。之后再用同样的输入和参数，和nn.MultiheadAttention的输出对比，数值会在一定误差范围内接近。具体差异来自初始化方式，不用强求完全一致。
4. 运行时的验证标准与性能边界
4.1 输出张量、梯度回传、数值稳定性怎么检查
一个多头注意力模块是否正常，不只是看能不能跑通，还要看几个判断标准。
输出张量的形状必须和输入保持一致，这是最基础的验收项。如果输入是 (batch, seq, embed_dim)，输出也必须长一样，否则后面接残差连接时会直接报形状错误。
梯度回传要检查每个线性层的权重梯度都不为空。可以用一个简单的循环遍历模型的参数：
for name, param in mha.named_parameters(): print(name, param.grad is not None)正常的模型每个可学习参数都应该有梯度。如果某个参数的梯度一直是 None，说明前向计算里它没有参与，通常是代码里的引用写错了。
数值稳定性方面，建议在训练早期观察注意力权重的分布。如果大量权重集中在少数几个位置，说明 softmax 的输入分布有问题，要检查缩放因子是否用对了。如果训练过程中 loss 出现 NaN，优先怀疑的问题是 mask 把 -inf 带进了数值计算，或者学习率设置过大。
4.2 显存和速度：序列长度是主要瓶颈
多头注意力的显存占用是序列长度的平方级。原因在于注意力分数矩阵的形状是 (batch, num_heads, seq_len, seq_len)。序列长度翻倍，这个矩阵的内存占用变成原来的四倍。相比之下，batch 大小增大带来的显存增长是线性的，反而更好控制。
实际经验是：如果 512 长度的序列在 8GB 显存上跑不动，优先降低序列长度，而不是减小 batch。序列长度从 512 降到 256，注意力矩阵的显存直接降到四分之一；而 batch 从 8 降到 4，只是整体减半。
速度优化上，可以从几个方向入手。如果只用 PyTorch 原生算子，可以尝试把 Q/K/V 投影合并，减少Linear的调用次数。如果序列很长，可以考虑 PyTorch 2.0 的scaled_dot_product_attention，它在底层会自动选择 flash attention 或者 memory-efficient attention，对长序列的加速非常明显。自己手写注意力主要是为了学习原理，真正跑大模型时建议还是用官方优化过的实现。
4.3 和 nn.MultiheadAttention 对比，什么时候用官方实现
PyTorch 官方提供的nn.MultiheadAttention参数接口更复杂，但内部逻辑完整，支持 mask、key padding mask、需要返回注意力权重等场景。它的接口签名是(query, key, value)，并不是只接受一个 x，这一点和手写的自注意力版本不一样，很多人在第一次使用时容易搞混。
我的建议是：学习阶段用手写实现，做研究验证用官方实现。手写实现的好处是任何一步都能打印中间结果，方便调 mask 和 debug。官方实现的好处是性能优化做得更好，底层算子经过充分测试，批量训练时不容易因为自己代码的问题导致不稳定。
如果用的是 PyTorch 2.1 以上的版本，可以直接试F.scaled_dot_product_attention替代手写 attention。它支持is_causal=True开启因果掩码，写 decoder 的时候很方便，代码量能省很多。
4.4 精度格式：fp32、fp16、bf16、tf32 的实际影响
这个点很多人会忽略，但真正训练或部署时会直接撞上。多头注意力里涉及大量的矩阵乘法，正好是混合精度影响最明显的部分。
fp32 是默认精度，数值稳定，但显存占用大、计算速度慢。fp16 能把显存减半，在支持半精度加速的 GPU 上明显变快，但数值范围小，训练时容易溢出。bf16 的指数范围和 fp32 一样，不容易溢出，但尾数精度低，更适合训练阶段的大规模矩阵运算。tf32 是 NVIDIA 显卡在 Ampere 架构之后提供的一种格式，它不改变显存占用，但能让 fp32 的矩阵乘法在硬件上加速，精度损失很小。
实际选型思路很简单：训练阶段先全程 fp32 保证正确，稳定后再开混合精度；内存不够时先试 bf16，而不是直接上 fp16；推理部署看重速度的话再用 fp16。不管用什么格式，注意力里的 softmax 在很多实现里仍然用 fp32 计算，这是有意为之，避免精度损失在指数运算里放大。
5. 常见报错和排查顺序
5.1 维度不匹配和头部切分错误
手写多头注意力最常见的报错就是维度不匹配。典型报错信息是shape '...' is invalid for input of size ...。
这类问题先按一个固定顺序排查：先确认输入 x 是三维张量(batch, seq, embed_dim)，再确认 embed_dim 能被 num_heads 整除，最后看 view 和 transpose 的目标形状计算是否一致。
很多人在这里把 batch 和 seq 的顺序搞混。view(batch_size, seq_len, self.num_heads, self.head_dim)之后transpose(1, 2)，得到(batch, num_heads, seq, head_dim)。这个顺序是后续矩阵乘法的前提，不能错。如果发现 attention score 的维度变成了(num_heads, batch, seq, seq)，多半是 transpose 的维度写错了。
调试时建议打印每个中间变量的 shape，不要靠猜。一个很小的注意力模块，中间张量也就 5 到 6 个，全部打印出来就能定位问题。
5.2 mask 形状错误和 -inf 处理
mask 是一个容易踩坑的地方。Transformer decoder 里的 causal mask 形状应该是(seq_len, seq_len)，key padding mask 形状通常是(batch_size, seq_len)。传到 attention 函数之前，要把 mask 扩展成(batch, num_heads, seq, seq)或者(batch, 1, seq, seq)，否则 broadcast 规则会让人很难受。
mask 的值也值得注意。如果代码里用的是masked_fill(mask == 0, float("-inf"))，那么有效位置应该是 1，填充位置是 0。这个约定本身没有问题，但很多人会不小心传成 bool 类型的 mask，导致mask == 0的结果和预期相反。
还有一个隐蔽的问题：如果把 -inf 填进 softmax，但后面又对 attention score 做 dropout，-inf 位置上可能出现 NaN。如果 loss 出现 NaN，先把 mask 相关的 -inf 去掉看看是否稳定，再排查精度问题。
5.3 训练时 loss 不降、梯度消失或数值不稳定
如果 forward 没问题但训练效果很差，不要急着改模型结构，先从三个方向检查。
第一，检查缩放因子。这是最容易被忽视的地方。如果忘了除以 sqrt(head_dim)，attention score 的数值会偏大，softmax 输出接近 one-hot，梯度很容易消失。
第二，检查初始化。PyTorch 的nn.Linear默认初始化通常没问题，但如果自己手动初始化了 Q/K/V，要确认标准差是 0.02 这个量级，而不是 1.0 这种大范围。投影矩阵过大会让 attention 一开始就饱和。
第三，检查 dropout 位置。注意力权重的 dropout 应该只在训练时生效。如果测试时忘记把模型切到eval()模式，dropout 会让输出随机变化，推理结果不稳定。
5.4 长序列 OOM 的处理
长序列下显存溢出是必然会发生的问题，尤其当输入长度超过 1024 时，注意力矩阵的占用量会很夸张。处理顺序是：先减 batch，再减序列长度，最后再考虑换注意力实现。
如果业务上确实需要处理长序列，可以考虑几种方案。第一种是用F.scaled_dot_product_attention，让 PyTorch 自动选择 flash attention。第二种是把注意力改成稀疏形式，比如局部窗口注意力。第三种是换用 Linear Attention 这类近似方案，但实现复杂度会明显上升。新手阶段不建议一上来就换复杂注意力变体，先把标准多头注意力跑稳最重要。
6. 实用边界与选型建议
6.1 多头不是越多越好
很多人在改模型时有一种惯性：头数不够就加头，加到头模型效果就变好。这个思路不一定对。
多头注意力的参数总量主要来自四个投影矩阵，头数变化并不会明显改变参数量。但头数增加意味着每个头的维度变小。比如 embed_dim 是 128，8 个头每个头只有 16 维，再往上加到 16 个头，每个头只有 8 维，过小的维度会让每个头能表达的信息非常有限，效果反而下降。
实际经验是：embed_dim 为 256 到 512 时，用 8 个头通常是比较稳的选择；embed_dim 为 768 以上时，可以考虑 12 个头。头数最好不要超过 embed_dim 能够支持的上限，更不要随意设置成不能被 embed_dim 整除的值。
如果发现多头注意力的效果不理想，优先调的是投影维度、dropout 和层数，而不是单纯加头。
6.2 自注意力、交叉注意力和掩码注意力的使用场景
多了解几种注意力变体，代码复用起来会更顺手。自注意力是 Q/K/V 都来自同一个序列，Transformer encoder 里用的就是这种。交叉注意力是 Q 来自 decoder 侧，K/V 来自 encoder 侧，比如翻译任务里 decoder 需要去 encoder 的输出里找信息。掩码注意力是在自注意力基础上加 causal mask，保证当前位置只能看到过去和当前，看不到未来。
在视觉模型里，Vision Transformer 把图像切成 patch 后直接做自注意力；Swin Transformer 在窗口内部做自注意力，再通过窗口移动建立跨窗口联系。这类结构本质上还是多头注意力，只是输入组织方式不同。如果你已经掌握了多头注意力的代码，理解这些变体只需要了解输入怎么切分、mask 怎么加，不涉及全新概念。
6.3 从学习 Demo 到生产部署的转换路线
如果你只是学原理，手写实现完全够用。但要做训练或部署，建议分步升级。
第一步，把手写 Q/K/V 换成官方nn.MultiheadAttention或者F.scaled_dot_product_attention，性能会好很多。第二步，补上 key padding mask 和 causal mask 的处理逻辑，确保 batch 内不同长度样本能正常计算。第三步，考虑精度格式和混用，在显存受限时打开混合精度训练。第四步，如果要把模型部署到服务里，注意把模型切到 eval 模式，去掉 dropout，固定随机种子，验证多次推理结果一致。
踩过几次坑之后我发现，很多注意力模块的问题不是原理不懂，而是输入形状、mask 约定和精度格式这些前置条件没有处理好。先把单头跑明白，再上多头；先把小序列跑通，再上长序列；先把 fp32 跑稳定，再折腾混合精度。这个顺序基本不会出错。
模型：原理、Python实现与数学建模实战)
