☰
多头注意力机制详解:原理、代码实现与工程实践
2026/10/30 14:19:08 网站建设 项目流程

大家在初学 Transformer 时,最先接触的往往不是完整的 Encoder-Decoder 结构,而是注意力机制。在注意力机制里,“多头注意力(Multi-Head Attention)”又是一个绕不开的核心模块。可以这么说:如果不理解多头注意力,就很难真正读懂 Transformer 的源码,更不用说去改模型、调参或者做二次开发了。

这篇文章围绕“多头注意力”这个主题,从单头注意力的困境讲起,逐步拆解多头注意力的数学原理、代码实现、实际应用场景和工程经验。无论你是刚入门深度学习的学生,还是在业务中需要使用 Transformer 系列模型的开发者,这篇文章都能帮你把这一块知识补扎实。

文章不依赖某个特定框架版本,代码示例以 PyTorch 为主,版本差异我会在涉及处做出说明。

1. 为什么要有多头注意力

1.1 从单头注意力说起

在解释多头注意力之前,我们需要先回顾一下单头注意力(Scaled Dot-Product Attention)是怎么工作的。

假设输入是一组序列数据,例如一句话中的每个词,或者一张图片切分后的多个 Patch。注意力机制的核心思路是:序列中的每个元素,在编码时不应该只关注自己,还应该关注序列中其他相关元素。

具体来说,每个输入元素都会被映射成三个向量:

  • Query(查询向量,简写 Q)
  • Key(键向量,简写 K)
  • Value(值向量,简写 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 整体计算流程

多头注意力的计算流程可以拆成四个步骤:

  1. 对输入的 Q、K、V 分别做线性变换,将其映射到 h 个不同的子空间。
  2. 在每个子空间中,独立执行缩放点积注意力计算。
  3. 将 h 个子空间的计算结果拼接起来。
  4. 对拼接结果做一次线性变换,得到最终输出。

这里的 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, mask=None): """ 缩放点积注意力 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, dtype=torch.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, 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. 计算多头注意力 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_weights

3.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_dim=512, num_heads=8) # 输入形状: [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_first=True。

mha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True) # 输入形状: [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 图像分类中的 ViT

Vision Transformer(ViT)将图像分成固定大小的 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,每个头维度 64;
  • d_model = 768 时,常用 num_heads = 12,每个头维度 64;
  • d_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'), diagonal=1) 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] # 这里取 batch=0, head=0 weights = attention_weights[0, 0].detach().numpy() plt.imshow(weights, cmap='hot', aspect='auto') 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, dropout=0.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, mask=None): # 多头注意力子层 + 残差连接 + 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 甚至大模型微调,都会轻松很多。

如果你在实现过程中遇到过有趣的问题,或者有其他关于注意力机制的疑问,欢迎在评论区一起交流。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询