在 Transformer 结构出现之前序列建模主要依赖 RNN、LSTM 这类循环神经网络。它们虽然能按顺序处理文本但存在两个明显痛点一是长距离依赖难以建模二是无法并行计算。自注意力机制Self-Attention的提出解决了这两个问题但在实际应用中单一的注意力分布往往不够“细腻”。于是论文《Attention Is All You Need》中提出了多头注意力机制Multi-Head Attention它通过多组独立的注意力头让模型同时关注不同子空间的信息从而大幅提升表达能力。这篇文章会围绕多头注意力机制展开先回顾自注意力的计算过程再逐步拆解多头机制的原理、代码实现、维度变化和工程实践。无论你是刚接触大模型的小白还是想加深对注意力机制理解的后端开发者都建议跟着文中的代码亲自跑一遍这对后续理解 BERT、GPT、Transformer 系列模型都会有很大帮助。1. 为什么需要“多头”注意力机制1.1 单头注意力的局限性为了理解多头注意力机制我们先回顾一下单头自注意力的计算过程。自注意力的核心思想是对于输入序列中的每一个 token都计算它与其他所有 token 的关联程度然后按关联程度加权聚合所有 token 的信息。计算公式如下Attention(Q, K, V) softmax(Q * K^T / sqrt(d_k)) * V其中QQuery查询向量表示“我要找什么”。KKey键向量表示“我是什么”。VValue值向量表示“我携带的信息”。d_kK 向量的维度用于缩放点积结果防止数值过大导致 softmax 梯度消失。从数学上看单头注意力确实能捕捉 token 之间的关系。但问题在于一句话中往往同时存在多种不同的依赖关系。举个例子句子“小明从北京坐高铁到上海他非常喜欢这座城市”中“他”需要指代“小明”这是指代关系。“这座”需要关联“上海”这也是指代关系。“高铁”和“北京”“上海”之间存在语义关系。“喜欢”和“城市”之间存在主谓宾关系。如果只使用单头注意力模型只能通过一套 Q、K、V 映射学习一个注意力分布。也就是说所有 token 之间的关系都被“压缩”到了一套权重中很多细粒度信息会被平均掉。1.2 多头注意力机制的核心思路多头注意力机制的改进非常直观既然一套注意力只能关注一种关系那我们就准备多套 Q、K、V 映射让每一套关注不同类型的关系。具体来说模型会把原始的 Q、K、V 分别通过多组不同的线性变换得到多组子 Q、K、V。每组子向量独立做注意力计算得到一组输出这就是一个“头”Head。最后把所有头的输出拼接起来再经过一次线性变换得到最终的输出。这样做的好处是每个头可以关注不同的位置关系例如头 A 关注词法依存关系头 B 关注指代关系头 C 关注局部邻近关系。模型整体表达能力增强不再被单一一套注意力分布限制。不同头之间天然形成一种“集成学习”的效果类似于多个弱学习器组合成强学习器。1.3 直观理解多角度审视文本我们可以用一个生活化的例子理解多头注意力。假设你是一名侦探负责调查一起案件中的人际关系。如果你只派一个探员去调查他可能会把重点放在经济利益冲突上而忽略了情感纠纷这条线。但如果你派出多个探员分别调查利益关系、情感关系、时间线、地理轨迹最后把各组的调查结果汇总起来你对整个案件的理解就会完整得多。多头注意力机制就是这个“多探员调查团”每一个头负责从不同角度建立 token 之间的关联。2. 环境准备与基础代码结构2.1 开发环境说明在动手写代码之前我们先明确实验环境。下面是我的本地环境参考你可以根据自己的实际情况调整操作系统Ubuntu 20.04 / Windows 10 / macOS 均可Python 版本3.8 及以上深度学习框架PyTorch 1.10 或更高版本其他依赖NumPy、Matplotlib用于可视化注意力权重开发工具Jupyter Notebook 或 VS Code如果你还没有安装 PyTorch可以通过官方命令安装。下面以 CPU 版为例pip install torch torchvision如果使用 GPU 加速请前往 PyTorch 官网选择与你 CUDA 版本匹配的安装命令。2.2 项目文件结构本文的代码比较简单核心代码放在一个 Python 文件中即可。你可以按下面的结构组织项目multi_head_attention_demo/ ├── multi_head_attention.py # 多头注意力核心实现 ├── self_attention.py # 单头自注意力实现对比用 ├── transformer_encoder.py # 完整 Transformer Encoder 示例 └── README.md # 说明文档3. 逐步拆解多头注意力机制3.1 从自注意力开始首先我们用 PyTorch 实现一个最基本的自注意力模块方便后续对比。这里使用 PyTorch 的nn.Linear来生成 Q、K、V。import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, embed_dim): 单头自注意力 :param embed_dim: 输入词向量的维度 super(SelfAttention, self).__init__() self.embed_dim embed_dim self.q_linear nn.Linear(embed_dim, embed_dim) self.k_linear nn.Linear(embed_dim, embed_dim) self.v_linear nn.Linear(embed_dim, embed_dim) def forward(self, x): :param x: 输入张量形状 (batch_size, seq_len, embed_dim) :return: 注意力输出形状 (batch_size, seq_len, embed_dim) Q self.q_linear(x) # (batch_size, seq_len, embed_dim) K self.k_linear(x) V self.v_linear(x) # 计算注意力分数 d_k self.embed_dim scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) # scores 形状: (batch_size, seq_len, seq_len) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output这段代码中Q * K^T得到的形状是(batch_size, seq_len, seq_len)其中第i行第j列表示第i个 token 对第j个 token 的关注分数。除以sqrt(d_k)是为了稳定训练。3.2 多头注意力实现接下来实现多头注意力机制。核心步骤为将输入x通过线性层生成 Q、K、V。将 Q、K、V 按头数切分重塑维度。每个头独立计算注意力。合并所有头的结果。通过输出线性层。class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads): 多头注意力机制 :param embed_dim: 输入词向量维度 :param num_heads: 注意力头数量 super(MultiHeadAttention, self).__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.q_linear nn.Linear(embed_dim, embed_dim) self.k_linear nn.Linear(embed_dim, embed_dim) self.v_linear nn.Linear(embed_dim, embed_dim) self.out_linear nn.Linear(embed_dim, embed_dim) def forward(self, x): batch_size, seq_len, _ x.size() Q self.q_linear(x) # (batch_size, seq_len, embed_dim) K self.k_linear(x) V self.v_linear(x) # 将 Q、K、V 拆分为多个头 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) # 此时形状: (batch_size, num_heads, seq_len, head_dim) # 计算注意力分数 scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtypetorch.float32)) # scores 形状: (batch_size, num_heads, seq_len, seq_len) attn_weights F.softmax(scores, dim-1) context torch.matmul(attn_weights, V) # context 形状: (batch_size, num_heads, seq_len, head_dim) # 合并所有头的结果 context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim) # 恢复为 (batch_size, seq_len, embed_dim) output self.out_linear(context) return output, attn_weights下面解释几个容易困惑的维度变换过程。第一步view拆分头部Q Q.view(batch_size, seq_len, self.num_heads, self.head_dim)原始 Q 的形状是(batch_size, seq_len, embed_dim)。由于embed_dim num_heads * head_dim我们可以按最后一个维度切分成num_heads段。第二步transpose交换维度Q Q.transpose(1, 2)交换第 1 维和第 2 维后Q 的形状变为(batch_size, num_heads, seq_len, head_dim)。这样设计的目的是把“头”这一维度前移方便批量计算。第三步恢复形状context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)context经过注意力计算后形状是(batch_size, num_heads, seq_len, head_dim)。先交换第 1 维和第 2 维恢复为(batch_size, seq_len, num_heads, head_dim)然后使用.contiguous()确保内存连续再通过view将最后两个维度合并为embed_dim。这里要注意view之前必须调用.contiguous()否则 PyTorch 会报错原因是在transpose操作之后张量的内存布局可能不是连续的。3.3 多头注意力的三种常见实现方式除了上面这种手写实现实际工程中通常使用更简洁的写法。方式一使用nn.MultiheadAttentionPyTorch 官方已经封装了多头注意力层可以直接调用import torch.nn as nn mha nn.MultiheadAttention(embed_dim512, num_heads8, batch_firstTrue) # 输入形状: (batch_size, seq_len, embed_dim) query torch.randn(2, 10, 512) key torch.randn(2, 10, 512) value torch.randn(2, 10, 512) attn_output, attn_weights mha(query, key, value) print(attn_output.shape) # torch.Size([2, 10, 512])batch_firstTrue表示输入形状为(batch, seq, feature)如果不设置该参数默认输入形状为(seq, batch, feature)容易踩坑。方式二缩放点积注意力封装考虑到代码复用可以将缩放点积注意力和多头切分逻辑分开def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights这种方式将注意力计算的核心逻辑独立出来后续如果要加 mask、加 dropout都只需要改这一个函数。方式三使用F.multi_head_attention_forward这是 PyTorch 内部使用的函数参数较多但效率最高。一般不建议新手直接使用除非你在做底层优化。3.4 多头注意力的参数量计算多头注意力机制的参数量主要由线性层决定。以embed_dim 512、num_heads 8为例Q、K、V 三个线性层的参数量3 * (embed_dim * embed_dim embed_dim) 3 * (512 * 512 512) 786,432输出线性层的参数量embed_dim * embed_dim embed_dim 512 * 512 512 262,656总参数量786,432 262,656 1,049,088可以看到多头注意力机制和单头注意力的参数量几乎相同只是把一个大矩阵拆分成了多个小矩阵并行计算。这也就是为什么多头机制能在不显著增加参数量的情况下提升表达能力。3.5 为什么这里要使用层归一化LayerNorm在实际的 Transformer 结构中多头注意力层之后通常会接一个残差连接Residual Connection和层归一化Layer Normalization。层归一化的作用是对每个样本的隐藏层输出做归一化使数据分布稳定加速收敛。与 BatchNorm 不同LayerNorm 不依赖 batch 大小更适合 NLP 中变长序列的场景。在 Transformer 中层归一化的计算可以表示为LayerNorm(x MultiHeadAttention(x))其中x是输入x MultiHeadAttention(x)是残差连接。这样做的好处是即使网络层数加深梯度也能通过残差连接直接回传缓解梯度消失问题。后续代码中我会把 LayerNorm 和 Dropout 也加入组成一个完整的 Transformer Encoder 层。4. 完整实战使用 PyTorch 构建 Transformer Encoder4.1 定义位置编码Transformer 中没有循环结构因此需要显式地在输入中注入位置信息。位置编码Positional Encoding可以使用正弦余弦函数生成。import math class PositionalEncoding(nn.Module): def __init__(self, embed_dim, max_len5000): super(PositionalEncoding, self).__init__() pe torch.zeros(max_len, embed_dim) position torch.arange(0, max_len, dtypetorch.float32).unsqueeze(1) div_term torch.exp(torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) pe pe.unsqueeze(0) # 形状: (1, max_len, embed_dim) self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1), :] return x位置编码的维度与词嵌入维度相同因此可以直接相加。register_buffer的作用是将pe注册为模型的一部分它不会参与梯度更新但在model.to(device)时会随着模型一起迁移到指定设备。4.2 定义 Transformer Encoder 层下面我们组合前面实现的多头注意力、层归一化、前馈神经网络Feed-Forward NetworkFFN和残差连接构成一个完整的 Transformer Encoder 层。class TransformerEncoderLayer(nn.Module): def __init__(self, embed_dim, num_heads, ff_dim, dropout0.1): super(TransformerEncoderLayer, self).__init__() self.self_attn MultiHeadAttention(embed_dim, num_heads) self.ffn nn.Sequential( nn.Linear(embed_dim, ff_dim), nn.ReLU(), nn.Linear(ff_dim, embed_dim), ) self.norm1 nn.LayerNorm(embed_dim) self.norm2 nn.LayerNorm(embed_dim) self.dropout nn.Dropout(dropout) def forward(self, x): # 多头注意力子层 attn_output, _ self.self_attn(x) x self.norm1(x self.dropout(attn_output)) # 前馈神经网络子层 ffn_output self.ffn(x) x self.norm2(x self.dropout(ffn_output)) return x这里的前馈神经网络由两个线性层和一个 ReLU 激活函数组成。ff_dim一般是embed_dim的 4 倍左右例如embed_dim 512时ff_dim 2048。4.3 定义完整 Transformer 模型将词嵌入、位置编码和多个 Encoder 层组合起来就是一个可运行的 Transformer Encoder。class TransformerEncoder(nn.Module): def __init__(self, vocab_size, embed_dim, num_heads, ff_dim, num_layers, dropout0.1, max_len5000): super(TransformerEncoder, self).__init__() self.embedding nn.Embedding(vocab_size, embed_dim) self.positional_encoding PositionalEncoding(embed_dim, max_len) self.layers nn.ModuleList([ TransformerEncoderLayer(embed_dim, num_heads, ff_dim, dropout) for _ in range(num_layers) ]) self.dropout nn.Dropout(dropout) def forward(self, x): x self.embedding(x) x self.positional_encoding(x) x self.dropout(x) for layer in self.layers: x layer(x) return x4.4 运行与验证我们构造一组简单的输入数据看模型是否能正常运行。# 超参数设置 vocab_size 1000 # 词表大小 embed_dim 128 # 词向量维度 num_heads 8 # 注意力头数 ff_dim 512 # 前馈层维度 num_layers 3 # Encoder 层数 batch_size 2 seq_len 10 # 随机生成输入 token 序列 x torch.randint(0, vocab_size, (batch_size, seq_len)) model TransformerEncoder(vocab_size, embed_dim, num_heads, ff_dim, num_layers) output model(x) print(输入形状:, x.shape) # torch.Size([2, 10]) print(输出形状:, output.shape) # torch.Size([2, 10, 128])输出形状为(batch_size, seq_len, embed_dim)符合预期。到这里你已经实现了一个最小可运行的 Transformer Encoder其中包含多头注意力机制的全部核心逻辑。5. 可视化多头注意力权重为了更直观地理解多头注意力机制我们可以把注意力权重画出来。这里使用一个简单的中文短句作为示例。5.1 生成注意力权重我们可以修改前面MultiHeadAttention类将attn_weights返回然后绘制热力图。import matplotlib.pyplot as plt # 构造一个简单的时间序列输入 x torch.randn(1, 6, 32) # (batch_size1, seq_len6, embed_dim32) mha MultiHeadAttention(embed_dim32, num_heads4) output, attn_weights mha(x) # attn_weights 形状: (batch_size, num_heads, seq_len, seq_len) print(注意力权重形状:, attn_weights.shape) # 输出: torch.Size([1, 4, 6, 6])5.2 绘制热力图def plot_attention(attn_weights, tokens, head_idx): plt.figure(figsize(8, 6)) im plt.imshow(attn_weights[0, head_idx].detach().numpy(), cmapBlues) plt.colorbar(im) plt.xticks(range(len(tokens)), tokens, rotation45) plt.yticks(range(len(tokens)), tokens) plt.title(fAttention Head {head_idx}) plt.show() tokens [我, 爱, 自然, 语言, 处理, .] plot_attention(attn_weights, tokens, 0) plot_attention(attn_weights, tokens, 1)不同的头会呈现出不同的注意力分布。例如有的头会更关注相邻词之间的依赖关系有的头可能会关注更远距离的词与词之间的联系。这就是多头注意力机制在实际任务中的表现。6. 多头注意力机制的进阶机制6.1 因果自注意力因果自注意力Causal Self-Attention是 GPT 系列模型中的核心机制。它的特点是在预测当前 token 时不允许模型看到未来的 token。因此在计算注意力分数时需要对未来的位置设置掩码Mask。def causal_attention_mask(seq_len): mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() return mask生成的掩码矩阵中右上角的元素为True。在计算注意力分数时将所有为True的位置替换为一个非常大的负数例如-1e9这样经过 softmax 之后这些位置的权重会趋近于 0。def scaled_dot_product_attention_causal(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights在 Transformer Decoder 中每个 token 只能看到它之前的信息这是生成式模型能够自回归输出的关键。6.2 注意力 Dropout在 Transformer 中注意力权重也经常使用 Dropout。具体做法是在 softmax 之后、乘 V 之前对注意力权重矩阵做 dropout。这样可以防止模型对某些位置的依赖过强从而提升泛化能力。self.attn_dropout nn.Dropout(0.1) attn_weights F.softmax(scores, dim-1) attn_weights self.attn_dropout(attn_weights) context torch.matmul(attn_weights, V)6.3 分组查询注意力分组查询注意力Grouped Query AttentionGQA是 LLaMA 2、Mistral 等大模型中常用的优化方案。它介于多头注意力和多查询注意力之间核心思路是多个查询头共享同一组 K 和 V从而减少 KV Cache 的显存占用。经典的 MHA 每个头都有独立的 K、V而 GQA 将查询头分为若干组每组共享一套 K、V。在实际推理时KV Cache 的显存占用可以降低数倍这对于长文本生成场景尤为重要。7. 常见问题与排查思路7.1 embed_dim 必须能被 num_heads 整除吗答案是的。在传统的多头注意力实现中embed_dim必须能被num_heads整除否则无法将特征维度均匀地切分给每个头。一般情况下我们会设定head_dim embed_dim // num_heads。如果你的模型嵌入维度无法被头数整除可以考虑调整num_heads为embed_dim的因数。例如embed_dim512时可以选择 1、2、4、8、16、32、64、128、256、512 作为头数。7.2 为什么运行时报错“shape mismatch”常见的shape mismatch报错发生在多头注意力维度变换阶段。最常见的原因是替换输入张量时没有考虑 batch 维度。建议按照以下步骤排查打印Q、K、V的初始形状。检查view后是否报错如果报错说明seq_len * embed_dim与拆分后的总元素个数不一致。检查transpose是否按正确的维度交换。检查最后拼接时是否忘记调用.contiguous()。7.3 训练时 loss 不下降可能是什么原因学习率过高导致模型震荡。建议使用学习率预热warmup策略。没有对注意力分数进行缩放导致 softmax 输出接近 one-hot梯度消失。没有使用残差连接和层归一化网络过深时梯度传输受阻。掩码实现有误导致模型看到了未来信息或填充位信息。7.4 显存占用过高怎么办降低 batch size。使用梯度累积模拟更大的 batch。使用混合精度训练将部分算子切换到 FP16。对长序列使用 FlashAttention 等高效注意力实现。8. 最佳实践与工程建议8.1 使用现成的高效实现如果不是学习算法原理不建议在生产环境使用手写的多头注意力代码。PyTorch 官方提供的nn.MultiheadAttention以及 FlashAttention 库在显存占用和计算效率上都有明显优势。手写实现主要用于理解原理工程落地时建议优先考虑成熟方案。8.2 合理设置头数头数并不是越多越好。头数过多会稀释每个头可用的表达能力且增加计算量。经验法则是embed_dim越大可以使用越多的头。常见的配置有embed_dim128num_heads4embed_dim256num_heads8embed_dim512num_heads8embed_dim768num_heads12embed_dim1024num_heads168.3 正确使用 mask在实现因果自注意力时要注意 mask 的维度需要扩展到与scores相同的维度。通常需要先创建下三角矩阵然后通过unsqueeze增加 batch 和 head 维度。8.4 注意力权重可视化与调试训练过程中定期可视化注意力权重能够帮助你判断模型是否学到了合理的依赖关系。如果某个头的注意力分布几乎均匀说明这个头可能“废了”如果某些 token 对所有其他 token 的注意力都很低说明位置编码或嵌入层可能有问题。9. 从自注意力到多头的学习路线总结本文从自注意力的计算逻辑出发逐步拆解了多头注意力机制Multi-Head Attention的完整实现。核心收获可以归纳为以下几点多头注意力机制通过多组独立的 Q、K、V 映射使模型能从多个子空间捕捉序列中的不同关系。多头注意力与单头注意力的参数量几乎一致计算开销增加不大但表达能力和稳定性显著提升。多头注意力的实现关键在于维度变换拆分、转置、并行计算、拼接还原。层归一化、残差连接和 Dropout 是注意力的重要“配套措施”在实际模型中通常不可省略。因果自注意力掩码是 GPT 类模型生成能力的基础。下一步你可以继续学习Transformers 库中BertSelfAttention的源码解析。FlashAttention 如何优化显存占用。GQA 与 MQA 在大模型推理中的应用。位置编码的改进方案例如 RoPE、ALiBi。代码是最公平的试炼场。光看不练只能停留在“似懂非懂”的阶段建议你打开 Jupyter Notebook把文章里的代码亲手敲一遍在维度变换处用print打印每一步的形状你会发现自己对注意力机制的理解会迅速加深。如果遇到环境配置或代码运行问题欢迎在评论区留下你的报错信息一起讨论排查。
多头注意力机制原理与PyTorch实现深度解析
在 Transformer 结构出现之前序列建模主要依赖 RNN、LSTM 这类循环神经网络。它们虽然能按顺序处理文本但存在两个明显痛点一是长距离依赖难以建模二是无法并行计算。自注意力机制Self-Attention的提出解决了这两个问题但在实际应用中单一的注意力分布往往不够“细腻”。于是论文《Attention Is All You Need》中提出了多头注意力机制Multi-Head Attention它通过多组独立的注意力头让模型同时关注不同子空间的信息从而大幅提升表达能力。这篇文章会围绕多头注意力机制展开先回顾自注意力的计算过程再逐步拆解多头机制的原理、代码实现、维度变化和工程实践。无论你是刚接触大模型的小白还是想加深对注意力机制理解的后端开发者都建议跟着文中的代码亲自跑一遍这对后续理解 BERT、GPT、Transformer 系列模型都会有很大帮助。1. 为什么需要“多头”注意力机制1.1 单头注意力的局限性为了理解多头注意力机制我们先回顾一下单头自注意力的计算过程。自注意力的核心思想是对于输入序列中的每一个 token都计算它与其他所有 token 的关联程度然后按关联程度加权聚合所有 token 的信息。计算公式如下Attention(Q, K, V) softmax(Q * K^T / sqrt(d_k)) * V其中QQuery查询向量表示“我要找什么”。KKey键向量表示“我是什么”。VValue值向量表示“我携带的信息”。d_kK 向量的维度用于缩放点积结果防止数值过大导致 softmax 梯度消失。从数学上看单头注意力确实能捕捉 token 之间的关系。但问题在于一句话中往往同时存在多种不同的依赖关系。举个例子句子“小明从北京坐高铁到上海他非常喜欢这座城市”中“他”需要指代“小明”这是指代关系。“这座”需要关联“上海”这也是指代关系。“高铁”和“北京”“上海”之间存在语义关系。“喜欢”和“城市”之间存在主谓宾关系。如果只使用单头注意力模型只能通过一套 Q、K、V 映射学习一个注意力分布。也就是说所有 token 之间的关系都被“压缩”到了一套权重中很多细粒度信息会被平均掉。1.2 多头注意力机制的核心思路多头注意力机制的改进非常直观既然一套注意力只能关注一种关系那我们就准备多套 Q、K、V 映射让每一套关注不同类型的关系。具体来说模型会把原始的 Q、K、V 分别通过多组不同的线性变换得到多组子 Q、K、V。每组子向量独立做注意力计算得到一组输出这就是一个“头”Head。最后把所有头的输出拼接起来再经过一次线性变换得到最终的输出。这样做的好处是每个头可以关注不同的位置关系例如头 A 关注词法依存关系头 B 关注指代关系头 C 关注局部邻近关系。模型整体表达能力增强不再被单一一套注意力分布限制。不同头之间天然形成一种“集成学习”的效果类似于多个弱学习器组合成强学习器。1.3 直观理解多角度审视文本我们可以用一个生活化的例子理解多头注意力。假设你是一名侦探负责调查一起案件中的人际关系。如果你只派一个探员去调查他可能会把重点放在经济利益冲突上而忽略了情感纠纷这条线。但如果你派出多个探员分别调查利益关系、情感关系、时间线、地理轨迹最后把各组的调查结果汇总起来你对整个案件的理解就会完整得多。多头注意力机制就是这个“多探员调查团”每一个头负责从不同角度建立 token 之间的关联。2. 环境准备与基础代码结构2.1 开发环境说明在动手写代码之前我们先明确实验环境。下面是我的本地环境参考你可以根据自己的实际情况调整操作系统Ubuntu 20.04 / Windows 10 / macOS 均可Python 版本3.8 及以上深度学习框架PyTorch 1.10 或更高版本其他依赖NumPy、Matplotlib用于可视化注意力权重开发工具Jupyter Notebook 或 VS Code如果你还没有安装 PyTorch可以通过官方命令安装。下面以 CPU 版为例pip install torch torchvision如果使用 GPU 加速请前往 PyTorch 官网选择与你 CUDA 版本匹配的安装命令。2.2 项目文件结构本文的代码比较简单核心代码放在一个 Python 文件中即可。你可以按下面的结构组织项目multi_head_attention_demo/ ├── multi_head_attention.py # 多头注意力核心实现 ├── self_attention.py # 单头自注意力实现对比用 ├── transformer_encoder.py # 完整 Transformer Encoder 示例 └── README.md # 说明文档3. 逐步拆解多头注意力机制3.1 从自注意力开始首先我们用 PyTorch 实现一个最基本的自注意力模块方便后续对比。这里使用 PyTorch 的nn.Linear来生成 Q、K、V。import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, embed_dim): 单头自注意力 :param embed_dim: 输入词向量的维度 super(SelfAttention, self).__init__() self.embed_dim embed_dim self.q_linear nn.Linear(embed_dim, embed_dim) self.k_linear nn.Linear(embed_dim, embed_dim) self.v_linear nn.Linear(embed_dim, embed_dim) def forward(self, x): :param x: 输入张量形状 (batch_size, seq_len, embed_dim) :return: 注意力输出形状 (batch_size, seq_len, embed_dim) Q self.q_linear(x) # (batch_size, seq_len, embed_dim) K self.k_linear(x) V self.v_linear(x) # 计算注意力分数 d_k self.embed_dim scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) # scores 形状: (batch_size, seq_len, seq_len) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output这段代码中Q * K^T得到的形状是(batch_size, seq_len, seq_len)其中第i行第j列表示第i个 token 对第j个 token 的关注分数。除以sqrt(d_k)是为了稳定训练。3.2 多头注意力实现接下来实现多头注意力机制。核心步骤为将输入x通过线性层生成 Q、K、V。将 Q、K、V 按头数切分重塑维度。每个头独立计算注意力。合并所有头的结果。通过输出线性层。class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads): 多头注意力机制 :param embed_dim: 输入词向量维度 :param num_heads: 注意力头数量 super(MultiHeadAttention, self).__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.q_linear nn.Linear(embed_dim, embed_dim) self.k_linear nn.Linear(embed_dim, embed_dim) self.v_linear nn.Linear(embed_dim, embed_dim) self.out_linear nn.Linear(embed_dim, embed_dim) def forward(self, x): batch_size, seq_len, _ x.size() Q self.q_linear(x) # (batch_size, seq_len, embed_dim) K self.k_linear(x) V self.v_linear(x) # 将 Q、K、V 拆分为多个头 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) # 此时形状: (batch_size, num_heads, seq_len, head_dim) # 计算注意力分数 scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtypetorch.float32)) # scores 形状: (batch_size, num_heads, seq_len, seq_len) attn_weights F.softmax(scores, dim-1) context torch.matmul(attn_weights, V) # context 形状: (batch_size, num_heads, seq_len, head_dim) # 合并所有头的结果 context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim) # 恢复为 (batch_size, seq_len, embed_dim) output self.out_linear(context) return output, attn_weights下面解释几个容易困惑的维度变换过程。第一步view拆分头部Q Q.view(batch_size, seq_len, self.num_heads, self.head_dim)原始 Q 的形状是(batch_size, seq_len, embed_dim)。由于embed_dim num_heads * head_dim我们可以按最后一个维度切分成num_heads段。第二步transpose交换维度Q Q.transpose(1, 2)交换第 1 维和第 2 维后Q 的形状变为(batch_size, num_heads, seq_len, head_dim)。这样设计的目的是把“头”这一维度前移方便批量计算。第三步恢复形状context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)context经过注意力计算后形状是(batch_size, num_heads, seq_len, head_dim)。先交换第 1 维和第 2 维恢复为(batch_size, seq_len, num_heads, head_dim)然后使用.contiguous()确保内存连续再通过view将最后两个维度合并为embed_dim。这里要注意view之前必须调用.contiguous()否则 PyTorch 会报错原因是在transpose操作之后张量的内存布局可能不是连续的。3.3 多头注意力的三种常见实现方式除了上面这种手写实现实际工程中通常使用更简洁的写法。方式一使用nn.MultiheadAttentionPyTorch 官方已经封装了多头注意力层可以直接调用import torch.nn as nn mha nn.MultiheadAttention(embed_dim512, num_heads8, batch_firstTrue) # 输入形状: (batch_size, seq_len, embed_dim) query torch.randn(2, 10, 512) key torch.randn(2, 10, 512) value torch.randn(2, 10, 512) attn_output, attn_weights mha(query, key, value) print(attn_output.shape) # torch.Size([2, 10, 512])batch_firstTrue表示输入形状为(batch, seq, feature)如果不设置该参数默认输入形状为(seq, batch, feature)容易踩坑。方式二缩放点积注意力封装考虑到代码复用可以将缩放点积注意力和多头切分逻辑分开def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights这种方式将注意力计算的核心逻辑独立出来后续如果要加 mask、加 dropout都只需要改这一个函数。方式三使用F.multi_head_attention_forward这是 PyTorch 内部使用的函数参数较多但效率最高。一般不建议新手直接使用除非你在做底层优化。3.4 多头注意力的参数量计算多头注意力机制的参数量主要由线性层决定。以embed_dim 512、num_heads 8为例Q、K、V 三个线性层的参数量3 * (embed_dim * embed_dim embed_dim) 3 * (512 * 512 512) 786,432输出线性层的参数量embed_dim * embed_dim embed_dim 512 * 512 512 262,656总参数量786,432 262,656 1,049,088可以看到多头注意力机制和单头注意力的参数量几乎相同只是把一个大矩阵拆分成了多个小矩阵并行计算。这也就是为什么多头机制能在不显著增加参数量的情况下提升表达能力。3.5 为什么这里要使用层归一化LayerNorm在实际的 Transformer 结构中多头注意力层之后通常会接一个残差连接Residual Connection和层归一化Layer Normalization。层归一化的作用是对每个样本的隐藏层输出做归一化使数据分布稳定加速收敛。与 BatchNorm 不同LayerNorm 不依赖 batch 大小更适合 NLP 中变长序列的场景。在 Transformer 中层归一化的计算可以表示为LayerNorm(x MultiHeadAttention(x))其中x是输入x MultiHeadAttention(x)是残差连接。这样做的好处是即使网络层数加深梯度也能通过残差连接直接回传缓解梯度消失问题。后续代码中我会把 LayerNorm 和 Dropout 也加入组成一个完整的 Transformer Encoder 层。4. 完整实战使用 PyTorch 构建 Transformer Encoder4.1 定义位置编码Transformer 中没有循环结构因此需要显式地在输入中注入位置信息。位置编码Positional Encoding可以使用正弦余弦函数生成。import math class PositionalEncoding(nn.Module): def __init__(self, embed_dim, max_len5000): super(PositionalEncoding, self).__init__() pe torch.zeros(max_len, embed_dim) position torch.arange(0, max_len, dtypetorch.float32).unsqueeze(1) div_term torch.exp(torch.arange(0, embed_dim, 2).float() * (-math.log(10000.0) / embed_dim)) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) pe pe.unsqueeze(0) # 形状: (1, max_len, embed_dim) self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1), :] return x位置编码的维度与词嵌入维度相同因此可以直接相加。register_buffer的作用是将pe注册为模型的一部分它不会参与梯度更新但在model.to(device)时会随着模型一起迁移到指定设备。4.2 定义 Transformer Encoder 层下面我们组合前面实现的多头注意力、层归一化、前馈神经网络Feed-Forward NetworkFFN和残差连接构成一个完整的 Transformer Encoder 层。class TransformerEncoderLayer(nn.Module): def __init__(self, embed_dim, num_heads, ff_dim, dropout0.1): super(TransformerEncoderLayer, self).__init__() self.self_attn MultiHeadAttention(embed_dim, num_heads) self.ffn nn.Sequential( nn.Linear(embed_dim, ff_dim), nn.ReLU(), nn.Linear(ff_dim, embed_dim), ) self.norm1 nn.LayerNorm(embed_dim) self.norm2 nn.LayerNorm(embed_dim) self.dropout nn.Dropout(dropout) def forward(self, x): # 多头注意力子层 attn_output, _ self.self_attn(x) x self.norm1(x self.dropout(attn_output)) # 前馈神经网络子层 ffn_output self.ffn(x) x self.norm2(x self.dropout(ffn_output)) return x这里的前馈神经网络由两个线性层和一个 ReLU 激活函数组成。ff_dim一般是embed_dim的 4 倍左右例如embed_dim 512时ff_dim 2048。4.3 定义完整 Transformer 模型将词嵌入、位置编码和多个 Encoder 层组合起来就是一个可运行的 Transformer Encoder。class TransformerEncoder(nn.Module): def __init__(self, vocab_size, embed_dim, num_heads, ff_dim, num_layers, dropout0.1, max_len5000): super(TransformerEncoder, self).__init__() self.embedding nn.Embedding(vocab_size, embed_dim) self.positional_encoding PositionalEncoding(embed_dim, max_len) self.layers nn.ModuleList([ TransformerEncoderLayer(embed_dim, num_heads, ff_dim, dropout) for _ in range(num_layers) ]) self.dropout nn.Dropout(dropout) def forward(self, x): x self.embedding(x) x self.positional_encoding(x) x self.dropout(x) for layer in self.layers: x layer(x) return x4.4 运行与验证我们构造一组简单的输入数据看模型是否能正常运行。# 超参数设置 vocab_size 1000 # 词表大小 embed_dim 128 # 词向量维度 num_heads 8 # 注意力头数 ff_dim 512 # 前馈层维度 num_layers 3 # Encoder 层数 batch_size 2 seq_len 10 # 随机生成输入 token 序列 x torch.randint(0, vocab_size, (batch_size, seq_len)) model TransformerEncoder(vocab_size, embed_dim, num_heads, ff_dim, num_layers) output model(x) print(输入形状:, x.shape) # torch.Size([2, 10]) print(输出形状:, output.shape) # torch.Size([2, 10, 128])输出形状为(batch_size, seq_len, embed_dim)符合预期。到这里你已经实现了一个最小可运行的 Transformer Encoder其中包含多头注意力机制的全部核心逻辑。5. 可视化多头注意力权重为了更直观地理解多头注意力机制我们可以把注意力权重画出来。这里使用一个简单的中文短句作为示例。5.1 生成注意力权重我们可以修改前面MultiHeadAttention类将attn_weights返回然后绘制热力图。import matplotlib.pyplot as plt # 构造一个简单的时间序列输入 x torch.randn(1, 6, 32) # (batch_size1, seq_len6, embed_dim32) mha MultiHeadAttention(embed_dim32, num_heads4) output, attn_weights mha(x) # attn_weights 形状: (batch_size, num_heads, seq_len, seq_len) print(注意力权重形状:, attn_weights.shape) # 输出: torch.Size([1, 4, 6, 6])5.2 绘制热力图def plot_attention(attn_weights, tokens, head_idx): plt.figure(figsize(8, 6)) im plt.imshow(attn_weights[0, head_idx].detach().numpy(), cmapBlues) plt.colorbar(im) plt.xticks(range(len(tokens)), tokens, rotation45) plt.yticks(range(len(tokens)), tokens) plt.title(fAttention Head {head_idx}) plt.show() tokens [我, 爱, 自然, 语言, 处理, .] plot_attention(attn_weights, tokens, 0) plot_attention(attn_weights, tokens, 1)不同的头会呈现出不同的注意力分布。例如有的头会更关注相邻词之间的依赖关系有的头可能会关注更远距离的词与词之间的联系。这就是多头注意力机制在实际任务中的表现。6. 多头注意力机制的进阶机制6.1 因果自注意力因果自注意力Causal Self-Attention是 GPT 系列模型中的核心机制。它的特点是在预测当前 token 时不允许模型看到未来的 token。因此在计算注意力分数时需要对未来的位置设置掩码Mask。def causal_attention_mask(seq_len): mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() return mask生成的掩码矩阵中右上角的元素为True。在计算注意力分数时将所有为True的位置替换为一个非常大的负数例如-1e9这样经过 softmax 之后这些位置的权重会趋近于 0。def scaled_dot_product_attention_causal(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtypetorch.float32)) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights在 Transformer Decoder 中每个 token 只能看到它之前的信息这是生成式模型能够自回归输出的关键。6.2 注意力 Dropout在 Transformer 中注意力权重也经常使用 Dropout。具体做法是在 softmax 之后、乘 V 之前对注意力权重矩阵做 dropout。这样可以防止模型对某些位置的依赖过强从而提升泛化能力。self.attn_dropout nn.Dropout(0.1) attn_weights F.softmax(scores, dim-1) attn_weights self.attn_dropout(attn_weights) context torch.matmul(attn_weights, V)6.3 分组查询注意力分组查询注意力Grouped Query AttentionGQA是 LLaMA 2、Mistral 等大模型中常用的优化方案。它介于多头注意力和多查询注意力之间核心思路是多个查询头共享同一组 K 和 V从而减少 KV Cache 的显存占用。经典的 MHA 每个头都有独立的 K、V而 GQA 将查询头分为若干组每组共享一套 K、V。在实际推理时KV Cache 的显存占用可以降低数倍这对于长文本生成场景尤为重要。7. 常见问题与排查思路7.1 embed_dim 必须能被 num_heads 整除吗答案是的。在传统的多头注意力实现中embed_dim必须能被num_heads整除否则无法将特征维度均匀地切分给每个头。一般情况下我们会设定head_dim embed_dim // num_heads。如果你的模型嵌入维度无法被头数整除可以考虑调整num_heads为embed_dim的因数。例如embed_dim512时可以选择 1、2、4、8、16、32、64、128、256、512 作为头数。7.2 为什么运行时报错“shape mismatch”常见的shape mismatch报错发生在多头注意力维度变换阶段。最常见的原因是替换输入张量时没有考虑 batch 维度。建议按照以下步骤排查打印Q、K、V的初始形状。检查view后是否报错如果报错说明seq_len * embed_dim与拆分后的总元素个数不一致。检查transpose是否按正确的维度交换。检查最后拼接时是否忘记调用.contiguous()。7.3 训练时 loss 不下降可能是什么原因学习率过高导致模型震荡。建议使用学习率预热warmup策略。没有对注意力分数进行缩放导致 softmax 输出接近 one-hot梯度消失。没有使用残差连接和层归一化网络过深时梯度传输受阻。掩码实现有误导致模型看到了未来信息或填充位信息。7.4 显存占用过高怎么办降低 batch size。使用梯度累积模拟更大的 batch。使用混合精度训练将部分算子切换到 FP16。对长序列使用 FlashAttention 等高效注意力实现。8. 最佳实践与工程建议8.1 使用现成的高效实现如果不是学习算法原理不建议在生产环境使用手写的多头注意力代码。PyTorch 官方提供的nn.MultiheadAttention以及 FlashAttention 库在显存占用和计算效率上都有明显优势。手写实现主要用于理解原理工程落地时建议优先考虑成熟方案。8.2 合理设置头数头数并不是越多越好。头数过多会稀释每个头可用的表达能力且增加计算量。经验法则是embed_dim越大可以使用越多的头。常见的配置有embed_dim128num_heads4embed_dim256num_heads8embed_dim512num_heads8embed_dim768num_heads12embed_dim1024num_heads168.3 正确使用 mask在实现因果自注意力时要注意 mask 的维度需要扩展到与scores相同的维度。通常需要先创建下三角矩阵然后通过unsqueeze增加 batch 和 head 维度。8.4 注意力权重可视化与调试训练过程中定期可视化注意力权重能够帮助你判断模型是否学到了合理的依赖关系。如果某个头的注意力分布几乎均匀说明这个头可能“废了”如果某些 token 对所有其他 token 的注意力都很低说明位置编码或嵌入层可能有问题。9. 从自注意力到多头的学习路线总结本文从自注意力的计算逻辑出发逐步拆解了多头注意力机制Multi-Head Attention的完整实现。核心收获可以归纳为以下几点多头注意力机制通过多组独立的 Q、K、V 映射使模型能从多个子空间捕捉序列中的不同关系。多头注意力与单头注意力的参数量几乎一致计算开销增加不大但表达能力和稳定性显著提升。多头注意力的实现关键在于维度变换拆分、转置、并行计算、拼接还原。层归一化、残差连接和 Dropout 是注意力的重要“配套措施”在实际模型中通常不可省略。因果自注意力掩码是 GPT 类模型生成能力的基础。下一步你可以继续学习Transformers 库中BertSelfAttention的源码解析。FlashAttention 如何优化显存占用。GQA 与 MQA 在大模型推理中的应用。位置编码的改进方案例如 RoPE、ALiBi。代码是最公平的试炼场。光看不练只能停留在“似懂非懂”的阶段建议你打开 Jupyter Notebook把文章里的代码亲手敲一遍在维度变换处用print打印每一步的形状你会发现自己对注意力机制的理解会迅速加深。如果遇到环境配置或代码运行问题欢迎在评论区留下你的报错信息一起讨论排查。
![[光学原理与应用-521]：对光的错误理解与纠偏](http://pic.xiahunao.cn/yaotu/[光学原理与应用-521]：对光的错误理解与纠偏)
