大家在初学 Transformer 时最先接触的往往不是完整的 Encoder-Decoder 结构而是注意力机制。在注意力机制里“多头注意力Multi-Head Attention”又是一个绕不开的核心模块。可以这么说如果不理解多头注意力就很难真正读懂 Transformer 的源码更不用说去改模型、调参或者做二次开发了。这篇文章围绕“多头注意力”这个主题从单头注意力的困境讲起逐步拆解多头注意力的数学原理、代码实现、实际应用场景和工程经验。无论你是刚入门深度学习的学生还是在业务中需要使用 Transformer 系列模型的开发者这篇文章都能帮你把这一块知识补扎实。文章不依赖某个特定框架版本代码示例以 PyTorch 为主版本差异我会在涉及处做出说明。1. 为什么要有多头注意力1.1 从单头注意力说起在解释多头注意力之前我们需要先回顾一下单头注意力Scaled Dot-Product Attention是怎么工作的。假设输入是一组序列数据例如一句话中的每个词或者一张图片切分后的多个 Patch。注意力机制的核心思路是序列中的每个元素在编码时不应该只关注自己还应该关注序列中其他相关元素。具体来说每个输入元素都会被映射成三个向量Query查询向量简写 QKey键向量简写 KValue值向量简写 V注意力输出可以概括为“根据 Query 和 Key 的相关性对 Value 做加权求和”。相关性越高权重越大对应 Value 在输出中的占比越高。公式如下其中 Q、K、V 分别表示查询矩阵、键矩阵和值矩阵dk 是 Key 向量的维度。这里除以 sqrt(dk) 的目的是缩放点积结果。当维度较大时点积结果数值会变得很大导致 Softmax 输出过于尖锐梯度容易消失。缩放之后点积结果的方差保持在可控范围内模型训练更稳定。1.2 单头注意力存在的问题单头注意力虽然能捕捉序列中元素之间的相关性但它有一个明显限制每次计算时模型只能从一种角度去理解元素之间的关系。举个例子。在句子“小明喜欢打球他放学后去了操场”中“他”和“小明”之间存在指代关系同时“打球”和“操场”之间存在场景关联。这两种关系特征差异很大单头注意力在计算时会把所有相关信息混合在一起很难同时把多种关系都清晰地建模出来。换句话说单头注意力的表达空间比较有限。它更像是一个人只能从一个视角看问题容易忽略其他同样重要的信息。1.3 多头注意力到底“多”在哪里多头注意力的做法是不直接用一组 Q、K、V 完成注意力计算而是将 Q、K、V 分别通过不同的线性变换投影到多个子空间中在每个子空间独立计算注意力最后再把所有子空间的结果拼接起来经过一次线性变换输出。这样做的好处是每个头可以关注不同位置、不同语义维度的信息。不同头可以学习到互补的关系。模型整体的表达能力明显增强。可以这样理解单头注意力是一个人看问题多头注意力是一组人从不同角度同时看问题最后把大家的意见汇总。每个头有自己关注的重点比如有的头更关注相邻词关系有的头更关注远程依赖有的头更关注语法结构。2. 多头注意力的核心原理2.1 整体计算流程多头注意力的计算流程可以拆成四个步骤对输入的 Q、K、V 分别做线性变换将其映射到 h 个不同的子空间。在每个子空间中独立执行缩放点积注意力计算。将 h 个子空间的计算结果拼接起来。对拼接结果做一次线性变换得到最终输出。这里的 h 就是“头数”。通常我们会看到 h 8 或 h 12 这样的配置。2.2 数学公式拆解多头注意力的公式可以写成其中表示第 i 个头的注意力输出Q、K、V 是输入的查询、键、值矩阵、、是第 i 个头的投影权重矩阵是输出的投影权重矩阵Concat 表示拼接操作。从矩阵维度来看假设输入特征维度为 d_model头数为 h那么每个头的维度为在实际代码中多数实现不会真的创建 h 个独立的线性层而是通过一次大矩阵乘法再把结果 reshape 成多头形状来计算。这样实现更简洁计算效率也更高。2.3 为什么要用“多头”子空间角度的解释深度学习模型的隐藏层维度通常较高例如 512 维或 768 维。如果只用一次注意力计算相当于把 512 维的信息杂糅在一次相关性计算中信息之间会互相干扰。多头注意力通过将维度切分成 h 个区块每个区块在独立的子空间内计算注意力。这类似于卷积神经网络中的多个卷积核每个卷积核负责提取一种特征多个卷积核共同覆盖完整的特征空间。在实际训练中我们也能观察到不同注意力头关注的内容往往有明显差异。在文本任务中有的头关注相邻词之间的关系有的头关注句子中距离较远的指代关系在图像任务中有的头关注局部纹理有的头关注全局轮廓。这种分化是单头注意力很难做到的。2.4 权重矩阵与输出投影在标准实现中输入 Q、K、V 可能来自同一个输入序列自注意力场景也可能来自不同序列跨注意力场景。无论哪种情况多头注意力内部都会对每个头准备独立的投影矩阵。需要说明的是是每个头独立学习的一组参数也就是说头与头之间参数不共享。输出投影 W_O 则是所有头共享的负责将拼接后的多维信息融合回 d_model 维度。从参数量的角度看多头注意力的参数总量其实与单头注意力并没有本质差异。假设 d_model 512头数 h 8每个头的维度 d_k 64那么 Q、K、V 三个投影矩阵的总参数量为输出投影 W_O 的参数量同样为 512 x 512。所以多头注意力的参数量约为单头注意力的等效规模但因为中间经过了切分和拼接表达能力更强。3. 多头注意力的代码实现3.1 PyTorch 实现一个标准多头注意力下面我们使用 PyTorch 从零实现一个多头注意力模块。虽然 PyTorch 官方已经提供了 torch.nn.MultiheadAttention但从零实现能帮助我们理解内部细节。import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): 标准多头注意力实现 d_model: 输入特征维度 num_heads: 注意力头数 def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model 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 scaled_dot_product_attention(self, Q, K, V, maskNone): 缩放点积注意力 Q: [batch_size, num_heads, seq_len, d_k] K: [batch_size, num_heads, seq_len, d_k] V: [batch_size, num_heads, seq_len, d_k] mask: [batch_size, seq_len, seq_len] 或 None # 计算注意力分数 scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtypetorch.float32)) # 如果提供了 mask将 mask 中为 0 的位置替换成极小值 if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # Softmax 归一化 attention_weights F.softmax(scores, dim-1) # 对 V 加权求和 output torch.matmul(attention_weights, V) return output, attention_weights def forward(self, Q, K, V, maskNone): 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. 计算多头注意力 attn_output, attn_weights self.scaled_dot_product_attention(Q, K, V, mask) # 3. 拼接所有头 attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # 4. 输出线性变换 output self.W_o(attn_output) return output, attn_weights3.2 代码逐段解释第一步线性变换并拆成多头。Q self.W_q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)这行代码是最容易绕晕的地方。我们分解一下self.W_q(Q) 输出的形状是 [batch_size, seq_len, d_model]。view(batch_size, -1, self.num_heads, self.d_k) 将最后一维 d_model 拆成 [num_heads, d_k]。transpose(1, 2) 把维度顺序调整为 [batch_size, num_heads, seq_len, d_k]。调整之后所有头的 Q 可以一次完成矩阵乘法不需要写循环。第二步注意力计算。代码中用 torch.matmul(Q, K.transpose(-2, -1)) 计算 Q 和 K 的内积除以 sqrt(d_k) 之后做 Softmax再与 V 相乘。这里需要注意的是 mask 参数。在 Transformer 的解码器阶段我们需要避免当前位置看到未来的信息所以会传入一个上三角掩码矩阵将未来位置置为 0。代码里通过 masked_fill 将 mask 中等于 0 的位置替换成负无穷Softmax 之后这些位置的概率就趋近于 0。第三步多头结果拼接。transpose(1, 2) 将维度调整为 [batch_size, seq_len, num_heads, d_k]然后 view 合并成 [batch_size, seq_len, d_model]。第四步输出线性变换将拼接后的结果映射回 d_model 维度。3.3 实际运行示例我们可以随机生成一组输入验证模块可以正常前向传播。# 随机生成输入 batch_size 2 seq_len 10 d_model 512 num_heads 8 x torch.randn(batch_size, seq_len, d_model) # 创建多头注意力模块 mha MultiHeadAttention(d_model, num_heads) # 自注意力场景Q、K、V 都来自 x output, attn_weights mha(x, x, x) print(输出形状:, output.shape) print(注意力权重形状:, attn_weights.shape)预期输出如下输出形状: torch.Size([2, 10, 512]) 注意力权重形状: torch.Size([2, 8, 10, 10])可以看到输出的序列长度和特征维度保持不变只是每个位置的向量经过了全局信息的融合。注意力权重则是一个四维张量含义是每个 batch、每个头、每个 Query 位置对所有 Key 位置的注意力分数。3.4 使用 PyTorch 官方实现PyTorch 官方提供了现成的 torch.nn.MultiheadAttention。在日常项目中我们通常直接使用它效率更高兼容性更好。import torch import torch.nn as nn # 官方多头注意力 mha nn.MultiheadAttention(embed_dim512, num_heads8) # 输入形状: [seq_len, batch_size, embed_dim] x torch.randn(10, 2, 512) attn_output, attn_weights mha(x, x, x) print(输出形状:, attn_output.shape) print(注意力权重形状:, attn_weights.shape)注意官方接口的输入维度顺序是 [seq_len, batch_size, embed_dim]而不是常见的 [batch_size, seq_len, embed_dim]。这是很多初学者容易踩的坑。如果你习惯 batch_first 风格可以设置 batch_firstTrue。mha nn.MultiheadAttention(embed_dim512, num_heads8, batch_firstTrue) # 输入形状: [batch_size, seq_len, embed_dim] x torch.randn(2, 10, 512) attn_output, attn_weights mha(x, x, x)4. 多头注意力在实际模型中的应用4.1 文本语义中的多头注意力在 BERT、GPT 这类预训练语言模型中多头注意力是 Transformer Encoder 和 Decoder 的基础组件。以 BERT-base 为例它的配置是12 层 Transformer Block每层 12 个注意力头d_model 768每个头的维度 64很多研究者对 BERT 的注意力头做过可视化分析结果发现部分头主要负责捕捉相邻词的局部语法关系部分头可以捕捉长距离的指代关系比如代词和名词之间的关联部分头对句子的整体语义信息更敏感。这意味着多头注意力不只是“多了几个头”这么简单它实际上让模型具备了从多个粒度、多个角度理解文本的能力。4.2 图像分类中的 ViTVision TransformerViT将图像分成固定大小的 Patch例如 16x16每个 Patch 经过线性投影后变成一个 Token然后送入 Transformer。在 ViT 中多头注意力用来建模 Patch 与 Patch 之间的全局关系。相比卷积神经网络通过局部感受野逐步扩大视野ViT 的多头注意力在第一层就能让每个 Patch 直接关注到整张图像的其他 Patch。这是 Transformer 架构在图像任务上的一个重要特性。在实际图像分类任务中ViT 在数据量充足时通常能取得和 CNN 相当甚至更好的效果而且在大规模预训练后表现出很强的迁移能力。4.3 Swin Transformer 中的窗口注意力Swin Transformer 是另一类视觉 Transformer 代表模型。它并没有直接使用全局多头注意力而是将注意力限制在窗口Window内并通过移位窗口Shifted Window实现跨窗口信息交互。这种设计的出发点很实际图像分辨率高时Patch 数量很大全局注意力的计算复杂度会随序列长度平方增长显存和耗时都难以接受。Swin Transformer 通过窗口注意力将计算复杂度降为线性同时保留多尺度特征在检测、分割等密集预测任务上表现更好。从多头注意力的角度看Swin Transformer 依然使用多头机制只是每个头的注意力范围被限制在窗口内部。这说明多头注意力是一个通用模块具体怎么用、用在哪里可以根据任务灵活设计。5. 常见问题与排查思路在实际项目中使用多头注意力大家可能会遇到下面这些问题。问题现象常见原因解决思路d_model 无法被 num_heads 整除维度设置不合理修改 num_heads 或 d_model确保二者满足整除关系输出形状与输入不一致view 和 transpose 操作使用不当检查维度变换前后张量顺序使用 contiguous() 避免内存布局问题训练时注意力权重波动剧烈学习率过高或没有缩放点积检查是否除以 sqrt(d_k)适当降低学习率多头注意力可视化时某些头全为空白模型初始化或训练不充分增加训练步数检查是否使用了过大的 Dropout解码时结果异常没有传入正确的 mask检查是否需要上三角掩码mask 中未来位置应为 0 或 False显存占用过高序列长度过长使用窗口注意力、稀疏注意力或对序列进行截断5.1 头数应该怎么选头数并不是越多越好。当 d_model 固定时头数越多每个头的维度 d_k 就越小单个头的表达能力会被削减。经验上以下配置比较常见d_model 512 时常用 num_heads 8每个头维度 64d_model 768 时常用 num_heads 12每个头维度 64d_model 1024 时常用 num_heads 16每个头维度 64。可以看到很多经典模型都把单个头的维度保持在 64 附近。这可以作为一个初始参考值再根据任务效果调整。5.2 多头注意力中的 mask 到底怎么传在自回归任务中比如 GPT 这类文本生成模型每个位置只能看到它之前的信息不能看到未来。常见做法是生成一个上三角矩阵并将其作为 mask 传入注意力模块。import torch seq_len 5 # 生成上三角掩码保留左下半部分 mask torch.triu(torch.ones(seq_len, seq_len) * float(-inf), diagonal1) print(mask)输出如下tensor([[0., -inf, -inf, -inf, -inf], [0., 0., -inf, -inf, -inf], [0., 0., 0., -inf, -inf], [0., 0., 0., 0., -inf], [0., 0., 0., 0., 0.]])在计算注意力分数时将这个 mask 加到 scores 上-inf 经过 Softmax 后概率为 0从而屏蔽未来信息。5.3 实现中的一些细节坑维度顺序问题。官方 nn.MultiheadAttention 默认输入是 [seq_len, batch_size, embed_dim]与大部分模型的输入格式不一致需要在调用时特别注意。view 与 transpose 混用。transpose 操作会改变张量的内存布局后续接 view 前需要调用 contiguous()否则会报错。Dropout 的位置。注意力权重的 Dropout 通常放在 Softmax 之后而不是放在 Q、K、V 投影之前。数值稳定性。如果数据集中存在特别长的序列点积结果可能非常大建议在注意力计算时使用混合精度训练或适当增加梯度裁剪。6. 最佳实践与工程建议6.1 训练稳定性多头注意力涉及 Softmax 操作对数值范围比较敏感。在训练大型 Transformer 时我建议使用学习率预热前几千步从较小学习率逐渐上升到目标学习率配合梯度裁剪避免梯度爆炸使用 AdamW 优化器并合理设置 weight_decay如果使用混合精度训练可以显著降低显存占用加快训练速度。6.2 保存与可视化注意力权重很多情况下我们不仅需要模型预测结果还需要解释模型“看了哪里”。例如在文本分类任务中我们希望通过注意力权重看看模型关注了哪些词。使用前面的自定义模块可以方便地拿到 attention_weights。在模型推理时保存下来用热力图或辅助函数可视化。import matplotlib.pyplot as plt import numpy as np # 假设 attention_weights 形状为 [batch_size, num_heads, seq_len, seq_len] # 这里取 batch0, head0 weights attention_weights[0, 0].detach().numpy() plt.imshow(weights, cmaphot, aspectauto) plt.colorbar() plt.show()需要注意的是注意力权重并不能完全等同于“模型这样预测的原因”它只能反映模型在计算过程中的关注分布。在解释模型行为时需要结合其他分析方法一起使用。6.3 部署阶段的注意事项如果训练后的 Transformer 模型需要部署到生产环境例如使用 ONNX Runtime 或 TensorRT 加速对多头注意力的支持情况需要提前确认。部分推理引擎对 Transformer 中的多个 transpose 和 reshape 支持不够高效需要在导出时检查算子是否完整支持序列长度是动态变化的推理引擎可能需要对动态维度做额外配置如果使用的是自研多头注意力实现尽量先与官方实现做输出对比测试确保数值差异在可接受范围内。6.4 什么时候不需要多头注意力多头注意力并非所有场景的最优选择。对于特别短的序列例如只有几个 Token 的输入多头带来的表达能力增益有限反而增加参数和计算量。此时可以考虑使用单头注意力使用线性注意力替代直接使用 CNN 或 MLP。另外在一些对延迟非常敏感的线上场景中如果序列长度不大多头注意力的并行优势也体现得不明显。可以先在离线实验中对不同头数进行对比再决定最终配置。7. 代码工程中的模块化建议在真实项目里我们一般不会单独使用一个裸的多头注意力模块而是将它封装在 Transformer Block 中。下面给出一个简单的编码器层示例方便你把前面的内容组合起来。import torch import torch.nn as nn import torch.nn.functional as F class TransformerEncoderLayer(nn.Module): 单层 Transformer Encoder d_model: 输入输出维度 num_heads: 注意力头数 d_ff: 前馈网络隐藏层维度 dropout: Dropout 概率 def __init__(self, d_model, num_heads, d_ff, dropout0.1): super(TransformerEncoderLayer, self).__init__() self.mha MultiHeadAttention(d_model, num_heads) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # 多头注意力子层 残差连接 LayerNorm attn_output, _ self.mha(x, x, x, mask) x self.norm1(x self.dropout(attn_output)) # 前馈网络子层 残差连接 LayerNorm ffn_output self.ffn(x) x self.norm2(x self.dropout(ffn_output)) return x这里有两个工程细节残差连接在 LayerNorm 之前这是 Transformer 中标准的 Post-Norm 结构。也有 Pre-Norm 变体把 LayerNorm 放在子层之前训练更稳定读者可以对比了解。多头注意力和前馈网络都接 Dropout但 Dropout 加在残差支路上而不是主路径上避免过度抑制主路径信息。8. 进一步学习建议如果你希望继续深入 Transformer 和注意力机制以下几个方向值得花时间首先吃透原论文。Vaswani 等人在《Attention Is All You Need》中详细描述了多头注意力的动机、公式和实验配置虽然论文不长但信息密度很高。其次阅读优质源码。除了 PyTorch 官方实现Hugging Face Transformers 库中的代码也值得精读。它的代码组织更贴近生产环境对 BERT、GPT 等模型的实现非常完整。然后关注注意力机制的优化变体。例如 Linear Attention、Flash Attention、Sparse Attention 等。Flash Attention 在不改变数学结果的前提下大幅提升了训练和推理速度目前已经成为大模型训练的事实标准。理解多头注意力之后再学习 Flash Attention会更加顺畅。最后动手实现一个小型 Transformer。可以尝试用 Transformer 完成一个简单的中文文本分类任务或者用 ViT 完成一个小型图像分类任务。不用一开始就追求大模型重点是把数据加载、模型构建、训练、评估这条链路跑通。结语多头注意力是 Transformer 架构中最核心的组件之一。它通过将 Q、K、V 投影到多个子空间让模型能够从不同角度捕捉序列中的关系从而在表达能力上远超单头注意力。这篇文章从原理、公式、代码、应用场景到工程建议完整拆解了多头注意力的方方面面。如果你刚开始学习建议先理解公式再跟着代码实现走一遍最后再去看 Flask 或 PyTorch 官方实现。只要把多头注意力真正搞懂后续学习 BERT、GPT、ViT 甚至大模型微调都会轻松很多。如果你在实现过程中遇到过有趣的问题或者有其他关于注意力机制的疑问欢迎在评论区一起交流。