从零实现注意力机制：PyTorch单头与多头注意力实战

在实际深度学习项目中注意力机制是让模型学会聚焦相关信息的核心手段。它解决的问题很具体输入序列很长时模型在当前时刻应该重点关注哪些位置。传统循环网络把全部历史编码进一个固定长度的状态向量句子越长早期信息被覆盖得越严重注意力机制则提供了一条显式通路让当前查询可以动态读取所有历史位置的内容按相关性赋予权重。这正是“AI如何理解上下文”的底层答案上下文不是被压缩成一个向量而是被建模成一张随查询位置变化的权重图。这篇文章会从注意力机制的出发点讲起解释 Query、Key、Value 的计算逻辑再用 PyTorch 从零实现单头注意力和多头注意力最后补充运行验证、常见问题排查和工程实践建议。读者最好有基础的 Python 和深度学习知识但不需要提前熟悉 Transformer 的完整实现。1. 从固定向量到动态寻址注意力机制解决什么问题1.1 序列建模的瓶颈上下文信息被“压扁”在注意力机制出现之前机器翻译和文本生成任务主要依赖循环神经网络。模型读入一个句子时通常把最后一步的隐藏状态当作整个句子的语义摘要或者把每一步隐藏状态拼接起来交给解码器。问题是长句子中靠前的内容经过多步非线性变换后信息会被逐步稀释。无论前面有多少单词最终都只能塞进一个维度有限的向量里这被称为“信息瓶颈”。从工程角度看这个设计还有一个问题它没有显式的“选择”过程。解码器在生成某个词时并不知道应该回头读源句子的哪个部分。注意力机制正是针对这个缺陷提出的不需要把上下文强行压成固定向量而是允许模型在每一步生成时重新“翻看”源序列并根据当前意图找到最相关的源位置。1.2 注意力的直观含义查询、键、值注意力机制有三个核心概念Query、Key、Value。用检索场景类比会更直观Query 是当前的问题。Key 是候选信息的标签。Value 是候选信息本身。模型先计算 Query 和每个 Key 的匹配程度得到一个分数再把分数转成概率最后用概率对 Value 做加权求和。匹配程度越高对应的 Value 在输出中占的比重越大。放到文本场景里当前正在翻译的词是 Query源句子中每个单词的表示是 Key 和 Value。模型计算当前词与源句每个词的相关度相关度高的词会更多地影响当前词的生成。整个过程是动态的不同生成位置得到的注意力权重不一样。1.3 典型应用场景注意力机制并不只在自然语言处理里出现。表格里的场景都适用场景Query 含义Key/Value 含义输出作用机器翻译当前解码状态源语言各位置状态生成当前目标词文本摘要已生成摘要状态原文各词表示生成下一摘要词图像描述当前单词状态图像各区域特征生成描述当前区域推荐系统用户近期行为表示候选物品特征预测点击概率语音识别当前声学状态语音帧特征输出文字或音素理解了这些场景就能理解为什么注意力机制被称为“通用组件”它不依赖具体模态只要数据能组织成序列或集合就能用注意力计算相关性。1.4 从注意力到自注意力早期的注意力通常发生在编码器和解码器之间叫跨注意力。后来研究者发现即使只把序列本身映射成 Query、Key、Value也能建立位置之间的关系。这就是自注意力机制序列中的每个位置都与序列中的其他位置计算相关性。自注意力让模型可以直接建立任意两个位置之间的依赖而不像循环网络那样必须经过中间步骤。这也是 Transformer 能处理长距离依赖的关键原因。下文实现的重点就放在自注意力上。2. 核心计算流程注意力分数、Softmax 与加权求和2.1 完整的注意力计算公式缩放点积注意力的计算分为三步用矩阵乘法计算 Query 与所有 Key 的点积得到注意力分数。将分数除以缩放因子再经过 Softmax 转成概率分布。用概率分布对 Value 加权求和得到输出。公式如下Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V其中d_k是 Key 的维度。Q 的每一行对应一个查询位置K 的每一行对应一个键位置QK^T的(i, j)元素就是第 i 个查询与第 j 个键之间的相关度。这个公式是 Transformer 系列模型中最常用的一种注意力实现也叫缩放点积注意力。它和早期加性注意力的区别在于点积运算可以写成矩阵乘法GPU 上执行效率高实现也简单。2.2 为什么需要缩放因子 sqrt(d_k)如果 d_k 很大点积结果会非常大。比如两个维度为 64 的向量每个分量的均值是 0、方差是 1那么点积的方差大约是 64标准差是 8。当数值进入 Softmax 时较大的输入会让输出分布变得非常尖锐靠近 1 的位置几乎饱和梯度变得极小模型难以训练。除以sqrt(d_k)可以把点积的方差拉回到 1 附近让 Softmax 输入保持在合理范围。这是缩放因子存在的根本原因。缩放方式输入分布训练效果不缩放方差接近 d_k数值偏大Softmax 饱和梯度小除以 sqrt(d_k)方差接近 1Softmax 分布平滑梯度稳定2.3 Mask 的作用填充掩码与因果掩码实际训练时不能直接对所有位置计算注意力需要根据任务引入掩码。填充掩码一个批次里句子长度不一短句末尾要补填充符填充位置不应该参与注意力计算。因果掩码自回归生成任务中当前位置只能看到自己和之前的位置不能看到未来信息因此需要把上三角位置屏蔽。实现时通常把被屏蔽位置的分数设为一个很小的负数例如-1e9让 Softmax 之后的权重接近 0。import torch seq_len 4 # 1 表示保留0 表示屏蔽 key_padding_mask torch.tensor([[1, 1, 0, 1]], dtypetorch.float32) mask key_padding_mask.unsqueeze(1) # 变成 [1, 1, 4] scores torch.randn(1, 4, 4) scores scores.masked_fill(mask 0, -1e9)2.4 常见注意力变体对比变体计算方式复杂度特点加性注意力通过前馈网络计算相关性较高早期模型常用表达能力更强点积注意力直接内积低实现简单矩阵运算高效缩放点积注意力点积后除以 sqrt(d_k)低Transformer 默认方案多头注意力多组 QKV 并行线性倍增能关注不同子空间自注意力序列自身生成 QKVO(n^2)建立任意位置依赖工程上做选型时通常优先用缩放点积注意力加多头结构只有特殊场景才考虑稀疏近似或线性注意力。3. 用 PyTorch 从零实现单头注意力3.1 环境准备学习环境不需要 GPU。准备 Python 3.8 以上、PyTorch 1.10 以上即可。CPU 上跑几个小张量验证逻辑完全够用。python -m venv attention_env source attention_env/bin/activate pip install torch安装完成后可以用一行命令确认版本python -c import torch; print(torch.__version__)如果原始环境里已经安装过 PyTorch只需要确认版本不要太旧因为masked_fill、torch.matmul这些 API 在旧版本中行为基本一致。3.2 最小实现缩放点积注意力下面的类接收 Q、K、V 和可选 mask输出注意力计算结果和权重矩阵。import math import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductAttention(nn.Module): def __init__(self, dropout0.0): super().__init__() self.dropout nn.Dropout(dropout) def forward(self, query, key, value, maskNone): d_k query.size(-1) scores torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) output torch.matmul(attn_weights, value) return output, attn_weights这里有一个关键约定mask 0表示屏蔽即值为 0 的位置被替换成-1e9得到接近 0 的注意力权重。不同框架的 mask 语义可能相反迁移代码时一定要先确认。3.3 测试单头注意力构造一个小张量验证输出形状和权重分布torch.manual_seed(42) query torch.randn(2, 4, 8) key torch.randn(2, 4, 8) value torch.randn(2, 4, 8) attention ScaledDotProductAttention(dropout0.0) output, attn_weights attention(query, key, value) print(output.shape) print(attn_weights.shape) print(attn_weights.sum(dim-1))预期输出output.shape为[2, 4, 8]与输入维度一致。attn_weights.shape为[2, 4, 4]其中第二个 4 是 key 的数量。attn_weights.sum(dim-1)的每一行都接近 1因为 Softmax 会把每一行归一化成概率分布。这一步跑通后说明注意力的基本计算链路没有问题。4. 自注意力与多头注意力模型如何建立上下文联系4.1 自注意力的特殊之处自注意力的输入只有一个序列xQ、K、V 都由同一个x经过线性变换得到。这样做的意义是每个位置都可以基于自身内容生成查询去“询问”其他位置的内容从而提炼上下文。定义输入形状为[batch_size, seq_len, d_model]其中seq_len是序列长度d_model是每个位置的向量维度。自注意力计算完成后输出形状仍然保持[batch_size, seq_len, d_model]。这个“形状不变”的特性让注意力层可以自由堆叠和替换其他模块。4.2 多头注意力多组子空间单头注意力只能学习一种相关性度量。实际任务里一个词可能同时关心语法搭配、语义相似、指代关系等多方面信息单头很难覆盖。多头注意力的做法是把 Q、K、V 拆成多组每组在低维子空间里独立计算注意力最后拼接起来。以下实现展示了拆分、计算、合并的完整过程class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0 self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads 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) self.dropout nn.Dropout(dropout) def forward(self, query, key, value, maskNone): batch_size query.size(0) Q self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k) K self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k) V self.w_v(value).view(batch_size, -1, self.n_heads, self.d_k) Q Q.transpose(1, 2) K K.transpose(1, 2) V V.transpose(1, 2) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn F.softmax(scores, dim-1) attn self.dropout(attn) context torch.matmul(attn, V) context context.transpose(1, 2).contiguous() context context.view(batch_size, -1, self.d_model) output self.w_o(context) return output, attn要点如下d_model // n_heads是每个头的维度必须整除否则无法正确拆分组装。view把最后一维拆成[n_heads, d_k]。transpose(1, 2)后形状变成[batch_size, n_heads, seq_len, d_k]方便对每个头独立计算。reshape前先contiguous()否则view会报错。4.3 位置编码为什么不能省自注意力对位置不敏感。交换输入序列中两个位置如果嵌入向量相同计算结果会完全相同。这样模型无法区分“猫追老鼠”和“老鼠追猫”。因此 Transformer 类模型会在输入嵌入上叠加位置编码给每个位置一个可区分的偏置。常见方案有正弦位置编码和可学习位置嵌入。处理超长文本时位置编码的外推能力直接影响模型对长上下文的适应性。4.4 上下文长度与计算复杂度自注意力的复杂度是 O(n^2 d)其中 n 是序列长度d 是向量维度。序列长度翻倍计算量约翻四倍。这是标准注意力在超长文本场景下面临的最大挑战也是后面要讨论各种高效注意力变体的原因。序列长度注意力分数矩阵大小计算量趋势512512 x 512基线10241024 x 1024约 4 倍20482048 x 2048约 16 倍81928192 x 8192约 256 倍正因如此上下文长度增加时不仅要关注显存还要关注训练时间和推理延迟。5. 运行验证观察注意力权重与输出5.1 用一个具体序列验证多头注意力构造一个形状为[2, 4, 8]的随机序列d_model8n_heads2torch.manual_seed(42) x torch.randn(2, 4, 8) mha MultiHeadAttention(d_model8, n_heads2, dropout0.1) output, attn_weights mha(x, x, x) print(output shape:, output.shape) print(attn shape:, attn_weights.shape)输出结果类似output shape: torch.Size([2, 4, 8]) attn shape: torch.Size([2, 2, 4, 4])attn_weights的第二维是 2代表两个注意力头。每个头内部都有一个[4, 4]的矩阵(i, j)表示第 i 个位置对第 j 个位置的注意力权重。5.2 如何解读注意力权重打印第一个样本第一个头的权重print(attn_weights[0, 0].round(decimals3))输出是一个 4 行 4 列的矩阵每一行和为 1。某一行中权重集中在对角线附近说明当前位置主要依赖自身权重分散到其他行说明模型在结合上下文。在真实模型里解读注意力权重需要谨慎。权重高并不意味着“因果解释”一定成立但它能反映模型计算时分配的计算重心。工程上常用注意力权重做可解释性分析、异常检测和剪枝依据。5.3 验证清单跑通代码后建议按以下清单逐项检查输出形状是否与输入一致。注意力权重最后一维是否满足 Softmax 归一化即每行之和接近 1。传入 mask 后被屏蔽位置的权重是否为 0。训练模式和推理模式下 Dropout 行为是否不同。多头拆分后拼接的结果是否与单头输出形状一致。注意只确认程序不报错还不够要验证输入输出语义是否正确。比如 Mask 传反了程序照样能跑但模型会学到错误依赖。6. 常见问题排查6.1 维度不匹配问题现象常见原因检查方式解决建议matmul报维度错误Q 和 K 的最后一维不一致打印query.shape、key.shape统一 d_model 或检查拆分逻辑view报错transpose后直接使用view检查是否调用contiguous()先contiguous()再view多头拼接后形状不对拆分维度与合并维度不一致逐步打印每步形状确认d_model % n_heads 0排查顺序先确认输入最外层是不是[batch, seq_len, d_model]再确认经过线性层后维度是否改变最后检查多头拆分时view的维度顺序。6.2 训练出现 NaN常见原因有四个学习率过大导致梯度爆炸。Softmax 输入包含inf常见于 mask 填充值设得过大或过小。fp16 混合精度下较大 score 超过半精度表示范围。数据里存在 NaN 或无穷值。检查方式torch.isnan(scores).any() torch.isinf(scores).any()处理建议把 mask 填充值从-1e9调整为-torch.inf效果更稳定。降低学习率加入梯度裁剪。混合精度训练时对 attention score 做数值稳定处理例如减去每行最大值后再做 Softmax。数据预处理时过滤异常值。6.3 Mask 没生效现象是填充位置仍然有较大注意力权重。通常原因是 mask 与 score 的形状不匹配或 mask 的 true/false 语义反了。检查步骤打印mask.shape和scores.shape确认能否广播。打印mask 0的分布确认哪些位置被标记为屏蔽。检查是masked_fill(mask 0, -1e9)还是masked_fill(mask, -1e9)两种写法的输入含义完全不同。6.4 注意力分布过于平滑如果所有位置权重都接近均匀分布模型很难学到有效上下文。可能原因训练不充分线性层尚未收敛。缺少位置编码序列顺序信息丢失导致部分位置的注意力没有区分度。学习率过高或过低模型没有稳定下降。dropout设置过大训练阶段把注意力权重打散。定位方式打印训练早期和训练后期的注意力矩阵观察分布是否逐渐集中。如果始终均匀优先检查位置编码和学习率。6.5 上下文长度超过训练范围直接把训练时 512 长度的模型用到 4096 长度常见表现是分数异常、注意力发散或性能明显下降。原因一般是绝对位置编码没有外推能力。解决思路使用支持长度外推的位置编码例如旋转位置编码或对位置编码做插值。训练时加入更长的序列样本。把超长文本切块或做摘要压缩后再进入模型。换用稀疏注意力或窗口注意力控制超长输入的计算成本。7. 工程实践与长文本优化7.1 学习环境与生产环境的差别学习环境里验证的是“能不能跑通”。生产环境要额外关注稳定性、显存、延迟和可维护性。关注点学习环境生产环境数据量小批量随机数据真实业务数据精度fp32混合精度或量化上下文长度固定较短长短不一需做截断或分段Mask可能忽略必须处理 padding 和因果掩码监控不关注需要记录 loss、梯度范数、注意力分布回滚不需要需要版本化模型和配置生产环境还要考虑推理时的 KV Cache。标准注意力在每次生成时要重新计算所有历史位置的 Key 和 Value浪费严重。使用 KV Cache 后历史结果可以缓存复用只计算当前新增位置。7.2 超参数选择建议下表给出常用参考值具体以实验为准参数常见范围影响d_model256 到 1024越大表达能力越强计算量越大n_heads8 到 16越多越细粒度但超过 d_model 整除限制会报错dropout0.1 到 0.3过大容易欠拟合过小容易过拟合学习率1e-4 到 3e-4过大梯度不稳过小收敛慢序列长度512 到 4096越长计算量增长越快这里的区间只是典型起点。实际项目要通过小规模实验确定不要直接照搬。7.3 长上下文优化方向当序列长度超过 2048 时标准注意力的计算成本往往不可接受。常见优化方向Flash Attention利用分块计算和在线 Softmax 减少显存读取是当前训练和推理的主流选择。窗口注意力每个位置只关注附近固定范围降低复杂度到 O(n)。稀疏注意力按预定义模式选择部分位置计算注意力。线性注意力把 Softmax 换成核函数近似把复杂度降到线性。分组查询注意力多个查询头共享一组 Key 和 Value减少 KV Cache 占用。这些方案有各自的精度和实现成本选型时要考虑硬件环境、模型规模和业务对精度的要求。8. 扩展方向从标准注意力到高效注意力8.1 Flash Attention标准注意力在计算过程中会把完整的[seq_len, seq_len]分数矩阵写入显存再进行 Softmax。序列越长显存占用越大。Flash Attention 通过分块和在线归一化避免实例化完整分数矩阵同时减少显存读写。使用细节上推荐优先使用官方实现的优化注意力 API而不是手写分块逻辑。例如 PyTorch 提供的高效注意力算子底层已经做了大量优化。8.2 线性注意力与稀疏注意力线性注意力把exp(QK^T)近似为phi(Q) phi(K)^T从而利用矩阵结合律先计算phi(K)^T V把复杂度从 O(n^2) 降到 O(n)。缺点是表达能力可能有损失长序列任务需要验证效果。稀疏注意力则通过预定义模式减少参与计算的位置适合文档、图像等局部结构明显的任务。窗口内计算完整注意力窗口外不计算直观且容易实现。8.3 GQA 与 KV Cache推理场景中自回归生成每一步都要产出新 token并读取全部历史 token 的 Key 和 Value。如果模型参数量很大KV Cache 的显存开销会非常高。分组查询注意力GQA让一组 Query 共享一份 Key 和 Value显著减少缓存量代价是轻微的表达能力损失。从工程角度KV Cache 的淘汰策略、精度压缩和异步管理是把大模型部署到高并发服务时必须解决的问题。8.4 学习路径建议如果刚开始学习注意力机制建议按这个顺序练习手写单头缩放点积注意力理解 QKV 和 Softmax。手写多头注意力掌握拆分与拼接。加入 padding mask 和 causal mask验证屏蔽效果。在真实小数据集上训练一个简单分类器或翻译模型观察注意力可视化。阅读 Transformer 原文中的实现细节重点看位置编码和残差结构。再切换到 Flash Attention 和长文本优化从原理层面理解为什么标准注意力在长序列上昂贵。注意力机制是理解大模型上下文的入口但它本身不复杂。掌握了 QKV 的拆解、权重的计算、mask 的控制和复杂度的来源后面再学习 Transformer 结构、参数高效微调、长文本推理都会顺畅很多。最值得花时间的地方不是背诵公式而是亲手实现一遍并把每个张量的形状变化写清楚。