从零手写自注意力机制：PyTorch实现QKV投影、缩放点积与Mask全解析

自注意力机制Self-Attention是这几年深度学习里最值得反复琢磨的概念之一。从大语言模型到多模态架构几乎所有主流模型的核心都离不开“从自身内容里找重点”这一套机制。大多数朋友卡住的点在于公式和示意图都看得懂一旦打开真实代码就傻眼。开源Transformer代码层层封装Query、Key、Value散落在各种基类和模块里根本分不清哪一步在做什么。这篇文章我就用最“土”的方式从零手写自注意力计算把Q、K、V投影、缩放点积、Softmax、Mask这些概念全部落到具体的PyTorch张量操作上逐步展示维度变化和数值流向把书上没讲清楚的细节一次性补完。适合的读者有三类已经学过Transformer理论课、想从代码角度验证理解的人准备算法面试、想问清楚底层细节的人以及想自己动手实现却被工程框架绕晕的人。内容不涉及花哨的优化技巧专注把基础机制逐行拆透。1. Self-Attention在做什么先建立起直观认知1.1 从“找重点”到“自己找重点”处理序列数据时模型最核心的需求就是“应该关注哪些位置”。RNN的做法是沿时间步逐步传递隐状态LSTM用门控增强记忆但序列一旦拉长信息传递链路就会变深早期信息很容易衰减或被冲淡。注意力机制换了一个思路不再依赖输入顺序的递归过程而是让每个位置的向量直接去和所有其他位置计算相关度。相关度高的位置被赋予更大权重相关度低的位置权重很小。这样无论序列多长任意两个位置之间的依赖都能在一步之内建模出来。那“Self”体现在哪里Self-Attention的本质是每个位置的查询向量和同一个序列里所有位置的键向量做匹配再收集对应的值向量。注意查询、键、值全都来自同一个输入序列这一点和跨注意力Cross-Attention查询来自一个序列键值来自另一个序列有本质区别。网络上很多示意图把Q、K、V画成三个独立的矩阵新手很容易误以为它们是三个不同的输入实际上它们只是把同一个输入向量分别投影到三个不同空间罢了。理解了这个后面看代码就会顺很多——你会看到就是三个全连接层每层的输入都是同一个x没有任何神秘的地方。1.2 Q、K、V三个角色怎么理解我对QKV的理解通常借助一个检索类比你想在资料库里找一段话先要输入一个查询词Query资料库中每一条文档都有自己的索引标识Key你真正需要读到的正文内容就是Value。注意力和检索唯一的不同在于Query和Key的匹配不是二进制的“命中/未命中”而是平滑的相似度分数所有条目都参与最终结果的加权合成只是权重不同。放在Self-Attention里同一个输入单词其实在同时扮演三种角色当它在主动找别的词时它是Query当别的词要回头看它时它是Key当它被其他词“吸收”时所要贡献的信息是Value。实际计算时我们不会一个词一个词地循环而是把这些角色对应的特征向量堆叠成矩阵用矩阵乘法一次性完成两两匹配和加权汇总。1.3 为什么这套机制能成为主流Self-Attention被广泛接受的一个关键原因一步之内就能建模任意距离的依赖关系。CNN需要堆叠多层才能覆盖足够大的感受野RNN需要沿时间逐步传递隐状态而Self-Attention的注意力矩阵本身就把所有位置之间的关系直接呈现出来了。此外它的计算过程高度并行不像RNN那样一步依赖一步GPU能充分发挥矩阵运算的吞吐优势。你还能从注意力权重矩阵里直观看到模型在关注什么这给可解释性分析带来了极大便利。当然也要承认它的代价核心计算复杂度是序列长度的平方所以后来衍生出稀疏注意力、线性注意力、滑动窗口注意力等各类变体。但无论怎么变大多数方案还是要回到基础Self-Attention这个锚点上。把基础版本的代码吃透后续理解这些变体就只是加减法。2. 公式拆解每一处设计都有它的理由2.1 缩放点积注意力的完整过程Self-Attention的核心公式长这样Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V用文字拆开就是三步。第一步把Q和K做点积得到一个形状为“序列长度×序列长度”的相似度分数矩阵矩阵里每个元素表示某两个位置之间的原始匹配度。第二步将分数除以一个缩放因子并做Softmax归一化变成每一行相加等于1的权重分布。第三步把这些权重应用到V矩阵上得到加权求和后的输出。代码实现时Q、K、V还有维度约定一般Q和K的向量维度要一致V的维度可以不同。比如在多头注意力里每个头的Q/K维度是64V的维度也设为64这样实现起来整齐但从公式本身来说V的维度不强制等于Q/K。刚接触时不用太拘泥这些细节先抓住主流程。2.2 缩放因子为什么是sqrt(d_k)这是代码里最容易被当成“固定常数”略过、但面试又高频被问的一个点。当向量维度变大时两个随机向量的点积结果绝对值也会自然变大原因在于高维空间中向量越来越难以相互正交各分量相乘后累加在一起的正贡献远大于负贡献。如果点积结果方差变大进入Softmax后部分数值就会滑到饱和区Softmax的梯度在这些区域会非常小训练几乎停滞。缩放的作用就是让Softmax输入保持在一个温和的范围内不让任何一列分数提前“统治”整个概率分布。实现上我习惯用d_k ** -0.5表示缩放因子然后直接乘到分数上等价于除以sqrt(d_k)但乘法在浮点运算里能少一次中间张量生成性能上略优。代码里两种写法我都见过结果一致。2.3 Softmax为什么作用在最后一个维度这个细节看起来小但写错维度会导致模型完全学不动。注意力分数的形状通常是[batch, num_heads, seq_len, seq_len]最后两个维度分别是“查询序列”和“键序列”。Softmax必须针对“键序列”那一维做归一化也就是对每一行做归一化让给定查询位置在所有键位置上的权重和为1。所以在代码里你会看到softmax(scores, dim-1)而不是dim1或dim2。如果写错了注意力矩阵的行和列语义就颠倒模型轻则不收敛重则直接输出NaN。后面第五节我会单独列调试方法。2.4 Mask的两种基本类型Mask是自注意力里绕不开的环节原论文语境下主要分两类。第一类是padding mask。模型训练时一个batch里的句子长度几乎必然不同短句要补齐到最长长度。补齐出来的pad位置没有实际语义注意力不应该看到它们所以要把这些位置的分数设成负无穷Softmax之后权重自然变成0。第二类是look-ahead mask也叫因果掩码或future mask主要用在解码器里。它保证预测第t个token时注意力只能看到前t个位置不能看到未来的信息这是自回归生成的基础约束。两类Mask在实现层面都表现为“把对应位置填成负无穷”只不过生成的形状和用法有差异。代码里一个常见的操作是masked_fill(mask 0, float(-inf))它把mask为0的位置全部替换成负无穷Softmax处理后这些位置就彻底不参与加权。3. 从零手写Self-Attention逐行代码解析3.1 输入形状与初始参数约定为了减轻新手负担我们先从尽量简单的配置入手。假设这样一个场景batch为2一次处理两条序列序列长度为4便于打印中间张量观察d_model为16即每个token用16维向量表示输入x的形状是[2, 4, 16]含义是两条序列、每条4个token、每个token一个16维向量。这个向量可以来自Embedding层也可以是上一层Transformer Block的输出。import torch import torch.nn as nn batch_size, seq_len, d_model 2, 4, 16 torch.manual_seed(42) x torch.randn(batch_size, seq_len, d_model) print(输入x的形状:, x.shape)期望输出[2, 4, 16]。这里我按照PyTorch生态常用的batch_first约定batch放在第一个维度这是大多数现代代码的默认习惯也方便后续手动构造Mask时理解。3.2 初始化三个投影层d_k d_model d_v d_model W_q nn.Linear(d_model, d_k, biasFalse) W_k nn.Linear(d_model, d_k, biasFalse) W_v nn.Linear(d_model, d_v, biasFalse)三个线性层的输入都是d_model输出对应各自的投影维度。bias在这里我通常设成False这背后是有讲究的注意力机制中Q和K的匹配是纯点积运算引入偏置相当于给点积的每一维增加偏移量对这个相似度匹配的建模能力提升有限反而会在模型量化部署时多出一组参数要处理。很多开源实现的默认做法也是不加偏置。你可能会问三个权重矩阵为什么不能共享直觉上如果共享权重“我要找什么”和“我能提供什么”就无法分离模型表达力会明显下降。所以三个投影必须是独立参数才能让同一个位置在不同语义空间下呈现不同的角色。3.3 一次性计算Q、K、VQ W_q(x) # [2, 4, 16] K W_k(x) # [2, 4, 16] V W_v(x) # [2, 4, 16]这一步没有任何神秘的地方就是三个全连接层各自作用一遍。关键认知是虽然输入都是同一个x但因为权重矩阵不同每个token在查询空间、键空间、值空间里的表征也就不同。打印出来的维度都是[2, 4, 16]这是单头注意力的特例。工程上很多框架不会写三个独立的Linear而是写一个大Linear直接输出3倍dim再用chunk切成三份。那种写法能减少一次矩阵乘法调用、提升计算效率但可读性差一些。教学场景用三个Linear更直观。3.4 计算注意力分数scale d_k ** -0.5 scores torch.matmul(Q, K.transpose(-2, -1)) * scale print(注意力分数形状:, scores.shape)期望输出[2, 4, 4]。这里最容易绕晕的就是K.transpose(-2, -1)。-2表示倒数第二个维度也就是序列长度-1是最后一个维度也就是向量维度。转置后K变成[2, 16, 4]Q的最后一维16与K转置后的倒数第二维16对齐做矩阵乘后得到的就是每个查询位置和每个键位置之间的点积相似度。结果矩阵中第b个batch、第i行第j列的值代表序列第i个token作为查询时与第j个token作为键的匹配分数。这就是“注意力分数”这个名称的由来。关于scale的写法我习惯乘在后边而不是先除后乘因为QK^T得到的是一个完整矩阵直接乘一个标量等价于对矩阵每个元素缩放计算上更干净。数值上两种写法完全等价不用纠结。3.5 Mask怎么写以PaddingMask为例这一步非常容易写错我单独拿出来讲。假设第一个batch的真实长度是3第二个batch是4那么第一条序列的第4个token是padding第二条序列没有padding。# mask形状 [2, 4]1表示有效位置0表示padding mask torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]], dtypetorch.bool) # 扩展成 [2, 1, 1, 4]配合广播机制 mask mask.unsqueeze(1).unsqueeze(1) scores scores.masked_fill(mask 0, float(-inf))scores原本是[2, 4, 4]最后一个维度对应键位置。mask原始形状是[2, 4]4正好对应键序列长度。用两次unsqueeze在中间补上维度变成[2, 1, 1, 4]这样广播时query维度第2维会自动复制展开最后被填充的就是所有query行在对应padding列上的分数。我见过很多人在这一步纠结为什么不是mask.unsqueeze(-1)如果末尾加维度mask形状会变成[2, 4, 1]和scores的[2, 4, 4]无法直接匹配结果就是阴影面积巨大、报错或语义错误。经验之谈每次写mask都先画一下scores的shape再确定要加在第几个维度。填充-inf之后Softmax会自动把这些位置的权重算成0padding token就不会被任何位置注意。3.6 Softmax与注意力输出attn torch.softmax(scores, dim-1) print(注意力权重:, attn.shape) output torch.matmul(attn, V) print(输出形状:, output.shape)torch.softmax(scores, dim-1)在键维度上归一化每一行的权重和为1。attn形状是[2, 4, 4]。接着torch.matmul(attn, V)做的是加权求和对每个batch、每个查询位置把V的所有token向量用注意力权重加权平均。输出形状回到[2, 4, 16]和输入完全一致这是后续残差连接需要的设计。如果d_v不等于d_model这一步输出维度就会和输入不一致这时就需要在最后加一个输出投影层把维度对齐回d_model。3.7 封装成完整模块把上面的散装逻辑整理成一个可复用的PyTorch模块顺便加入Dropout和输出投影这是工程里的常规做法。class SelfAttention(nn.Module): def __init__(self, d_model, d_kNone, d_vNone, dropout0.1): super().__init__() d_k d_k if d_k else d_model d_v d_v if d_v else d_model self.scale d_k ** -0.5 self.W_q nn.Linear(d_model, d_k, biasFalse) self.W_k nn.Linear(d_model, d_k, biasFalse) self.W_v nn.Linear(d_model, d_v, biasFalse) self.W_o nn.Linear(d_v, d_model, biasFalse) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): Q self.W_q(x) K self.W_k(x) V self.W_v(x) scores torch.matmul(Q, K.transpose(-2, -1)) * self.scale if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) attn self.dropout(attn) out torch.matmul(attn, V) return self.W_o(out)几个高频疑问在这里统一回答。Dropout放在Softmax之后、加权和之前是最常见做法含义是对注意力权重做随机置零相当于正则化“哪些词被关注”避免模型过拟合到少数固定的注意力模式。注意Dropout的比例不宜过大一般取0.1到0.2太大会让注意力分布趋于均匀反而失去聚焦能力。输出投影W_o在多头上几乎必须存在因为多个头的输出拼接后维度会恢复成d_model需要再投影一次才能和残差连接对齐。即使单头我也建议保留保持和Transformer Block其他部分的结构一致。attn矩阵在训练时一般不返回因为作图和可视化才需要推理或训练时多返回一个[batch, seq, seq]的矩阵会白白增加显存占用所以函数默认只返回最终输出。3.8 用内置接口验证手写版本自己写完之后最怕的是你不知道自己写对了没有。我有一个非常实用的验证方法把自定义模块的输出和PyTorch内置的nn.MultiheadAttention输出对齐看看数值差异能不能收敛到浮点精度级别。import torch.nn.functional as F custom_attn SelfAttention(d_model16, dropout0.0) official_attn nn.MultiheadAttention(embed_dim16, num_heads1, dropout0.0, batch_firstTrue) with torch.no_grad(): official_attn.in_proj_weight.view(3, 16, 16).copy_( torch.cat([custom_attn.W_q.weight.data, custom_attn.W_k.weight.data, custom_attn.W_v.weight.data], dim0) ) official_attn.in_proj_bias.zero_() official_attn.out_proj.weight.data.copy_(custom_attn.W_o.weight.data) official_attn.out_proj.bias.zero_() custom_out custom_attn(x) official_out, _ official_attn(x, x, x, need_weightsFalse) print(最大差异:, (custom_out - official_out).abs().max().item())如果代码正确最大差异应该在1e-6以下这是浮点数精度带来的正常偏差。如果差异很大优先查三件事Mask有没有传、Softmax维度是不是-1、scale有没有漏乘。这个对照实验非常值得做一次它能彻底打消“我写的是不是官方那个注意力”的疑虑。顺带一提官方nn.MultiheadAttention从PyTorch 2.0开始内部对need_weightsFalse的路径走了更高效的融合算子数值上跟手写版本会有微小偏差但并不会影响逻辑一致性。4. 维度流向与常见工程陷阱4.1 从输入到输出的维度变化全梳理把关键张量的形状整理成一张表代入本文设定的数值batch2、seq4、d_model16、d_kd_v16。张量形状含义x[2, 4, 16]输入序列Q[2, 4, 16]查询表示K[2, 4, 16]键表示K^T[2, 16, 4]转置后的键scores[2, 4, 4]缩放后的注意力分数attn[2, 4, 4]Softmax后的权重output[2, 4, 16]注意力输出这张表在我调试代码时几乎每次都会用。建议初学者把中间张量都print出来亲眼观察一下分数范围从最初较大数值被scale拉回到温和范围再经过Softmax变成0到1之间的概率分布整个流程就会变得极其清晰。4.2 多头实现时最容易踩的维度坑多头注意力和单头唯一的本质区别是把d_model拆成num_heads份每一份独立走一遍自注意力最后再拼接起来。代码实现里涉及到四维张量变换这里踩坑率极高。常见的步骤是Q/K/V先通过一个大Linear得到[batch, seq, d_model]然后view成[batch, seq, num_heads, head_dim]再permute成[batch, num_heads, seq, head_dim]。这里view和permute的顺序极其关键。view是在内存排列上直接按块切分如果先permute再view内存里的顺序已经被打乱切出来的块就会张冠李戴产生完全错误的注意力计算。业内的习惯顺序是先view到四维再permute把num_heads提到前边之后每个头的计算就并行完成。这个顺序没有任何例外哪怕写注释也要标清楚每一步的shape。4.3 官方接口里两种Mask的区别nn.MultiheadAttention接受的Mask参数分两类名字长得像语义完全不同。attn_mask的形状是[seq_len, seq_len]或[batch * num_heads, seq_len, seq_len]它直接和注意力分数相加用于实现look-ahead mask这类结构掩码。key_padding_mask的形状是[batch, seq_len]专门标记padding位置框架内部会把它转换成对应的掩码并应用到最终分数矩阵上。坑在于很多人一上来把padding mask传给attn_mask于是形状对不上报错信息又不够直观排查半天才发现参数传错地方了。遇到这种错误第一步就是打印两个mask的shape然后对照文档确认自己用的是哪一类mask比盲目改代码快得多。4.4 数值稳定性与显存控制三点实战建议。第一如果序列长度特别长中间那个[batch, num_heads, seq_len, seq_len]的注意力矩阵会很快把显存吃光。比如seq为4096时单头单样本就要16M个浮点数多头乘上去G显存根本不够。这时候真正能帮上忙的是PyTorch内置的scaled_dot_product_attention它会自动选择FlashAttention之类的融合算子在计算过程中不显式展开完整矩阵从而大幅降低显存占用。我在自己的项目里只要序列超过512就会优先使用这个函数而不是手写展开。第二训练时使用半精度注意力分数矩阵容易出现溢出。FlashAttention的融合实现内部做了更稳健的数值处理而手写展开版在极端情况下可能会踩到数值上界。总体建议是手写版本用来学习和对齐真正训练大模型时优先框架集成的高效实现。第三torch.softmax和F.softmax底层逻辑一致但接口约定略有不同团队里统一一种写法能减少排查成本。我的习惯是用torch.softmax因为直接指定dim参数语义清楚。5. 排查技巧与常见问题速查5.1 症状一训练Loss完全不下降注意力计算最常见的坑就是缩放因子没乘或者乘错了地方。如果没有缩放点积数值很大Softmax会进入饱和区梯度接近0模型怎么训都学不动。检查第一步永远是看scale d_k ** -0.5有没有作用在分数上。另一个高频原因是Softmax维度错误。注意力矩阵的维度关系是“查询行、键列”Softmax必须在键维度上做归一化如果对查询维度做归一化行和列的语义就颠倒了。我推荐写一个单元测试用全零的Q和K构造输入点积结果全为0Softmax后应该得到均匀概率分布每个位置概率相同。一旦这个基础测试不通过马上就能暴露出维度层面的问题而不需要等训练跑很久才发现。5.2 症状二Shape报错永远找不到源头最常见的一类报错信息形如The size of tensor a (16) must match the size of tensor b (8) at non-singleton dimension 3。这类错误九成出在多头维度变换的view和permute环节。我的排错习惯很简单——在写完代码后先打印Q、K、V各路shape再逐个关键节点打印中间结果的shape谁不对一目了然。shape问题永远优先于数值问题排查因为一旦维度错了后续所有数值分析都没有意义。还有一个偷懒技巧在代码里注释每个张量的预期形状把“注释在前、代码在后”当作规范后续排错时间能省一半。5.3 症状三注意力权重看起来都差不多注意力权重全部接近于均匀分布说明模型没学到“哪个位置重要”。原因可能是训练步数太少也可能是embedding还没学好更常见的是mask写反了把有效位置全部mask掉模型眼里只剩padding自然只能学出均匀分布——因为所有可见位置的语义都一样空洞。另一个容易被忽略的是Dropout过大。如果注意力权重上的Dropout超过0.5每个位置的权重都可能被随机抹掉模型没法稳定聚焦到少数关键位置最终学出来的注意力就倾向于均匀。注意力Dropout取值控制在0.1到0.2是最稳的。5.4 症状四和内置实现输出差异巨大如果自定义实现和官方nn.MultiheadAttention输出对不上先排除浮点误差再逐项排查下面三件事。第一官方接口默认batch_firstFalse输入要求是[seq, batch, dim]如果你用的是[batch, seq, dim]且没有指定batch_firstTrue数据进入就会错位这个问题极其隐蔽。第二官方QKV投影是同一个权重矩阵拼接而成还带偏置即便你的实现不加偏置也要同步清零in_proj_bias和out_proj.bias才能做到完全对齐。第三PyTorch 2.0以后设置need_weightsFalse时内部走的是融合算子路径输出数值上与手写会有微小但可接受的偏差这属于正常现象。5.5 工程实践中的一点体会如果你只是搭实验或者跑固定框架完全没必要手写Self-Attention直接调用内置模块更省心也更高效。但如果你要做的是魔改注意力比如定制Mask策略、实现稀疏模式、调整数值绑定规则手写版本的灵活优势就体现出来了。我的经验是先完整手写一遍单头自注意力彻底搞懂每一步的维度和数值变化再去读多头和融合实现效率和深度都比直接开箱即用要强得多。最后分享一个我一直保留的调试技巧把注意力权重矩阵打出来配合特定句子观察模型关注了哪些位置。如果模型回答某个问题时注意力集中在“证据词”附近说明它在看对地方如果权重一片均匀多半不是代码问题而是训练信号出了问题。这种可视化排障方法比对着loss曲线猜原因快很多。