Attention机制这几年几乎成了深度学习入门绕不开的一道坎尤其是Self-Attention从BERT到GPT再到各类多模态模型到处都有它的身影。公式倒不算复杂先算Q和K的点积得到相似度分数缩放之后softmax变成权重再对V做加权求和。可很多人看完公式一打开代码还是懵Q是哪儿来的转置怎么转mask填什么值为什么点积还要除个根号d_k这个系列上一篇把Attention的整体框架梳理了一遍这篇我们专门啃Self-Attention的代码。我会直接从代码入手先用十几行写一个最简实现把流程跑通再封装成一个能放进训练循环的模块然后重点拆解缩放因子和mask这两个最容易被忽视又最容易出错的细节最后聊几个我在实际调试中踩过的坑。适合刚看完公式、正准备动手写代码以及写了很多次维度还是报错的同学。1. 动手之前先对齐公式Self-Attention每一步在做什么1.1 从输入到Q、K、V为什么不能直接用X自己算标准公式是Attention(Q,K,V) softmax(QK^T / √d_k) V。输入是X形状通常是(batch, seq_len, d_model)。其中batch是样本数seq_len是序列长度d_model是每个token的向量维度。很多初学者第一反应是既然叫Self-Attention让X自己和自己算不就行了为什么还要先经过三个不同的线性投影如果直接用X算点积X和自身的相似度天然最高这个“匹配分数”缺少可学习的余地而且所有token都带着同样的原始语义在匹配信息交互太粗糙。引入W_Q、W_K、W_V这组可学习参数后模型就能学出三套不同的表示一个负责表达“我想找什么”查询一个负责表达“我有什么可以被匹配”键一个负责表达“匹配上之后我实际提供什么”值。用个生活化点的类比Q是你脑子里拟好的搜索关键词K是待检索文档的标题索引V是文档正文。关键词和标题匹配上了就去取对应的正文内容。至于怎么把句子变成关键词、标题、正文这就是三个线性层在训练中自己学出来的。所以Self-Attention的第一步代码一定是三个投影而不是直接拿X去点乘。1.2 四步计算流程数据是怎么流动的不看公式单纯看数据流Self-Attention就是四步投影X经过三个线性层得到Q、K、V形状都是(batch, seq_len, d_k)。单头情况下d_k通常等于d_model多头时d_k d_model / num_heads。计算相似度矩阵Q和K的转置做矩阵乘法得到(batch, seq_len, seq_len)的分数矩阵。第i行第j列表示序列第i个token对第j个token的匹配程度。缩放与归一化除以√d_k然后按最后一个维度做softmax每一行变成概率分布概率含义是“第i个token在聚合信息时应该分多少注意力给第j个token”。加权聚合用这个概率矩阵去乘V每个token的输出变成所有token的V按注意力权重加权求和的结果形状仍然是(batch, seq_len, d_k)。看这个流程就知道为什么Self-Attention可以做残差连接输入是(batch, seq_len, d_model)输出是(batch, seq_len, d_k)只要保证d_k和d_model一致输入输出形状就完全对齐。不同张量的维度变化可以整理成一张表张量形状含义x(batch, seq_len, d_model)输入序列表示Q(batch, seq_len, d_k)查询向量表达“要找什么”K(batch, seq_len, d_k)键向量表达“能被什么匹配”V(batch, seq_len, d_k)值向量表达“匹配后提供什么”scores(batch, seq_len, seq_len)注意力分数未归一化attn(batch, seq_len, seq_len)注意力权重每行和为1out(batch, seq_len, d_k)加权聚合后的输出维度追踪是写代码时最容易出错的地方我建议新手把这张表打印出来贴在显示器旁边真的比记公式管用。2. 最简实现先用十几行代码把原理跑通我建议大家第一版不要写类不要管dropout不要管mask就用一个函数把核心流程写出来。因为这样报错少逻辑也容易检查。工程上后面会封装但先把核心运算写对。2.1 十行版本的自注意力函数这里先用手动矩阵乘法不引入nn.Linear目的是让新手看见“投影本质上就是矩阵乘法”。w_q、w_k、w_v是外部传入的权重矩阵形状分别是(d_model, d_k)。import torch import torch.nn.functional as F import math def self_attention(x, w_q, w_k, w_v): # x: (batch, seq_len, d_model) Q x w_q # (batch, seq_len, d_k) K x w_k # (batch, seq_len, d_k) V x w_v # (batch, seq_len, d_k) d_k Q.shape[-1] scores Q K.transpose(-2, -1) / math.sqrt(d_k) # (batch, seq_len, seq_len) attn F.softmax(scores, dim-1) # 按最后一个维度归一化 out attn V # (batch, seq_len, d_k) return out, attn关于转置K的形状是(batch, seq_len, d_k)要和Q做(batch, seq_len, d_k) (batch, d_k, seq_len)的乘法需要把K后面两维交换所以用K.transpose(-2, -1)。这里用运算符即可PyTorch对batch矩阵乘法支持得很好不需要手动写循环。为什么转置用(-2, -1)而不是(1, 2)transpose的参数是维度索引倒数第二维和倒数第一维也就是交换右边两位。这种写法在函数里更稳健不管batch维度是几维永远只交换矩阵维度。2.2 验证一下输出形状和权重分布写完第一件事是验证我习惯先用小随机数跑一遍只看两个东西形状对不对attention权重行和是否为1。torch.manual_seed(42) batch, seq_len, d_model, d_k 2, 5, 8, 8 x torch.randn(batch, seq_len, d_model) w_q torch.randn(d_model, d_k) w_k torch.randn(d_model, d_k) w_v torch.randn(d_model, d_k) out, attn self_attention(x, w_q, w_k, w_v) print(out.shape) # torch.Size([2, 5, 8]) print(attn.shape) # torch.Size([2, 5, 5]) print(attn.sum(dim-1)) # 每行都应该等于1attn.sum(dim-1)的结果每一行都应该等于1因为softmax按dim-1归一化了。如果这里不对多半是softmax的维度写错了。我记得自己第一次写的时候把dim写成了1当时输入是三维张量输出行和不等于1排查了挺久才发现维度索引的问题。2.3 手动推一次2×2的注意力分数到底长什么样为了彻底理解拿一个极简例子手动算一遍。假设d_k1序列长度2两个token的query分别是q10.5、q2-0.3key分别是k10.8、k2-0.1。scores矩阵就是score_11 0.5 * 0.8 0.4score_12 0.5 * (-0.1) -0.05score_21 -0.3 * 0.8 -0.24score_22 -0.3 * (-0.1) 0.03d_k1缩放后数值不变softmax按行归一化。以第一行为例exp(0.4)1.49exp(-0.05)0.95softmax后第一个token对自身权重约0.61对第二个token约0.39。这个矩阵的含义是第一个token在生成输出时61%的信息来自自己39%来自第二个token。加权聚合时输出就是0.61v1 0.39v2。这一步做完你会发现Self-Attention的本质其实就是一种数据相关的加权平均而权重完全由Q、K的匹配结果决定。理解了这个小例子后面所有代码里的张量形状都不难看懂。3. 工程封装把它做成一个能训练的Module最简函数只适合理解原理真要用在模型里必须包成一个nn.Module。模块需要自己管理参数、支持mask、可插拔到任意网络层。3.1 设计要点线性层、dropout、输出投影封装之前先处理几个设计选择用nn.Linear而不是手动矩阵乘法这样参数会自动注册到模型里优化器能正确更新。Q、K、V的投影维度。单头实现里d_k可以等于d_model。如果项目里有降维需求也可以单独指定一个更小的d_k相当于把token表示压缩到低维子空间再计算注意力。输出投影out_proj。Transformer block里通常会在注意力后接一个线性层把输出映射回d_model方便和残差连接相加。dropout放在哪里应该放在softmax得到的attention权重上而不是Q/K/V上更不是输出上。原因是注意力权重才是模型需要“随机失活”的地方防止某个位置被过度依赖。这个细节很多新手会放错位置。3.2 完整代码带mask支持的SelfAttention模块import torch import torch.nn as nn import torch.nn.functional as F import math class SelfAttention(nn.Module): def __init__(self, d_model, d_kNone, dropout0.1): super().__init__() self.d_model d_model self.d_k d_k if d_k is not None else d_model self.w_q nn.Linear(d_model, self.d_k) self.w_k nn.Linear(d_model, self.d_k) self.w_v nn.Linear(d_model, self.d_k) self.out_proj nn.Linear(self.d_k, d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # x: (batch, seq_len, d_model) # mask: (batch, seq_len)1表示保留0表示遮蔽 batch, seq_len, _ x.shape Q self.w_q(x) # (batch, seq_len, d_k) K self.w_k(x) # (batch, seq_len, d_k) V self.w_v(x) # (batch, seq_len, d_k) scores Q K.transpose(-2, -1) / math.sqrt(self.d_k) # (batch, seq_len, seq_len) if mask is not None: scores scores.masked_fill(mask.unsqueeze(1) 0, float(-inf)) attn F.softmax(scores, dim-1) attn self.dropout(attn) out attn V # (batch, seq_len, d_k) output self.out_proj(out) # (batch, seq_len, d_model) return output, attn这个版本可以直接丢进任意一个Transformer block。mask参数设计成“保留位为1、遮蔽位为0”的布尔掩码比较直观也能和我后面要讲的两种mask自由组合。注意masked_fill里我特意加了mask.unsqueeze(1)这一步是把(batch, seq_len)的mask变成(batch, 1, seq_len)。scores是(batch, seq_len, seq_len)只有把mask的第二维扩成1它才会自动广播到所有query位置。如果不加这步mask的形状会和scores不匹配轻则报错重则在batch等于seq_len时侥幸通过得到完全错误的结果。3.3 和PyTorch内置实现对比PyTorch官方有nn.MultiheadAttention但它的接口和国家都不一定和手写实现一致。它需要(q, k, v)三个独立输入输入输出维度也不是按d_model统一管理默认还有embedding维度包装。对初学者来说直接用它就像坐过山车出了问题根本不知道是哪里错了。PyTorch 2.0之后提供了F.scaled_dot_product_attention它内部有flash attention和memory-efficient后端速度更快但接口更“裸”一些直接输入q、k、v不管理权重投影。如果拿我自己写的SelfAttention和官方接口对比可以这样理解我的手写版是一个“权重管理 自定义mask 数值控制”的壳子内部计算逻辑和公式完全对齐生产环境里你可以把内部核心换成F.scaled_dot_product_attention保留外面的权重投影和mask逻辑既有了性能又保留了对流程的控制权。4. 缩放因子和Mask两个最容易出问题的细节4.1 为什么公式里要除以根号d_k这是初学者最经典的疑问也是实际调参时真正有影响的细节。如果q和k的每个分量都是独立随机变量均值0、方差1那么点积q·k等于所有分量乘积之和。多个分量乘积的累加会导致方差也累加所以在d_k维的情况下点积结果的方差大约也是d_k。也就是说维度越高点积数值越分散。而softmax本质上是exp的归一化输入值稍微拉大一点exp的差异就会被急剧放大softmax输出的分布会变得非常“尖锐”——接近one-hot。一旦某个位置的概率近1其他位置趋近0反向传播时这些位置的梯度几乎消失模型很难学下去。除以√d_k的作用是把方差拉回到1附近。我用一个小实验验证过当d_k64时随机初始化的模型在没缩放情况下attention分数标准差经常在8左右softmax之后某个位置的概率经常直接到0.99以上加上缩放之后分布温和很多训练稳定不少。所以在写法上除以math.sqrt(self.d_k)这行一定不能省。场景不加缩放加缩放scores标准差≈√d_k偏高≈1softmax输出容易接近one-hot分布更平滑梯度偏小训练慢正常深层模型稳定性差稳定4.2 Padding Mask为什么填0不行必须填-inf实际训练时一个batch里序列长度通常不同我们会把短序列padding到同一个长度。padding位置没有任何信息不应该参与attention计算。最直觉的想法可能是“把padding位置的分数设成0”这是错的。softmax计算的是exp(score)的比值0经过exp之后是1它照样会占到概率权重。假设某一行有三个位置分数分别是5、2、0softmax之后padding位置的分量大概是exp(0)/(exp(5)exp(2)exp(0)) ≈ 0.001看起来很小但注意力权重求和时它仍然贡献了非零信息。而且随着序列变长所有padding位置的权重累计起来就不是小数目了。正确的做法是把不需要的位置设成一个非常小的数比如-inf这样exp(-inf)0归一化之后对应位置的权重就是0。这也是正本部分实现的原理if mask is not None: mask mask.unsqueeze(1) # (batch, 1, seq_len) scores scores.masked_fill(mask 0, float(-inf))如何从token id序列构造padding mask呢常规做法是这样的def build_padding_mask(token_ids, pad_token_id0): # token_ids: (batch, seq_len) return (token_ids ! pad_token_id).long()得到的是1表示有效、0表示无效的mask。把这个mask传进SelfAttention的forward即可。4.3 Causal MaskDecoder为什么只能看左边在Decoder的生成场景里当前位置不能看到未来的token。比如翻译任务中生成第三个词时输入只允许有前两个词和当前词不能看到后面的词。这个用下三角矩阵实现seq_len 5 causal_mask torch.tril(torch.ones(seq_len, seq_len)).bool()形状是(seq_len, seq_len)第i行第j列表示query位置i是否允许看到key位置j允许条件是j i。打印出来长这样[[ True, False, False, False, False], [ True, True, False, False, False], [ True, True, True, False, False], [ True, True, True, True, False], [ True, True, True, True, True]]把False的位置填成-infsoftmax之后权重就是0。在实际Decoder里通常需要同时加padding mask和causal mask。组合方式也很简单两个mask按位与# padding_mask: (batch, 1, seq_len) # causal_mask: (seq_len, seq_len) 或扩展到 (1, seq_len, seq_len) combined_mask padding_mask causal_mask scores scores.masked_fill(combined_mask 0, float(-inf))不过两者的shape需要广播对齐padding_mask是(batch, 1, seq_len)causal_mask是(seq_len, seq_len)对齐后变成(batch, seq_len, seq_len)。这个细节在工程里非常容易踩坑我见过不少人在Encoder里忘了扩展维度在Decoder里忘了和padding mask合并最后attention矩阵形状完全乱了。4.4 数值稳定性-inf在混合精度下会变成NaNfloat(-inf)在fp32下没问题但混合精度训练时如果用fp16inf的表示范围有限有些硬件会把-inf转换成NaN。所以工程上并不总填-inf有时填一个足够小的负数比如-1e9。fp16能表示的数值绝对值上限大约65504-1e9在fp16里可能被溢出成NaN所以更稳妥的选择是-10000。标准建议是纯fp32训练用-inf没问题开启AMP混合精度时优先用-10000这类较小值或者对mask后的分数做显式的isnan防护。另外还有一个隐藏的坑如果某一行全部被mask掉那整行都是-infsoftmax计算时会得到NaN。所以mask设计时一定要保证每一行至少有一个True。我自己调试时就遇到过这种情况检查和修改后就稳定了。5. 常见问题与排查实录5.1 矩阵维度报错transpose、dim、view全搞清楚最常见的报错是matmul shape mismatch特征就是“sizes of tensors must match”。核心是盯死三个shapeQ是(batch, seq_len, d_k)K转置后是(batch, d_k, seq_len)V是(batch, seq_len, d_k)。很多人在把d_model改成d_k时忘了改后面两个线性层的输出维度导致Q、K、V形状不一致。还有一类问题是view和reshape混用。view要求内存连续如果之前做过transpose直接view会报错。需要先contiguous()再view或者干脆用reshape它会自动处理内存重排。建议新手全程使用reshape省心。5.2 attention分布不对劲太尖或太平训练初期可以打印attn的分布来观察。用attn.max(dim-1)看每一行的最大概率如果最大概率几乎等于1说明注意力太尖锐模型每个token几乎只盯一个位置泛化差如果每行都接近uniform说明模型还没学到有效的匹配关系可能缩放因子没加对或者投影参数初始化不合适。我处理过一个真实casemask写反了保留位写0、遮蔽位写1导致所有有效token被遮蔽padding位置反而被保留模型训练快速退化。打印attention权重后发现每行都是均匀分布查代码才发现mask逻辑反了改回来后恢复正常。5.3 怎么确认自己写的实现是对的最直接的办法构造一个d_model1、seq_len较小的输入手动算一遍分数、权重和输出和代码输出对比误差应该接近0。这也是单元测试的核心思想。再分享一个我常用的链路验证方法让w_q w_k w_v I单位矩阵此时投影不改变向量Self-Attention退化为X自己和自己做点积softmax加权平均。如果这个特殊case的形状和数值都连续至少说明核心链路没有基本错误。不过这只是链路验证真正正确性还是需要和公式推导对齐。def check_forward(): torch.manual_seed(0) model SelfAttention(d_model8, d_k8) x torch.randn(2, 3, 8) out, attn model(x) assert out.shape (2, 3, 8) assert attn.shape (2, 3, 3) assert torch.allclose(attn.sum(-1), torch.ones(2, 3), atol1e-5) print(forward check passed) check_forward()注意assert attention行和为1是判断softmax维度是否正确最廉价的方式。任何一次实现改动后跑一遍能排除掉大部分低级错误。5.4 模型训练出现NaNNaN的排查顺序输入数据本身是否有NaNattention scores里是否有NaN尤其是mask导致的全-inf行以及softmax在fp16下是否触发了数值溢出。用torch.isnan(scores).any()逐层打印哪个张量出现NaN就从哪个环节排查。另一个和Attention无关但经常误伤的原因是学习率过大导致梯度爆炸。这时检查模型参数的梯度范数如果确实爆炸先把学习率降下来再说。我见过不少人为了NaN问题在Attention代码上反复折腾最后发现是学习率太高。6. 从单头到多头下一步怎么改6.1 为什么要多头单头注意力只能学一种匹配模式也就是一组Q、K、V投影。多头相当于同时学多组不同的匹配模式每个头可以关注不同的位置、不同的子空间关系。比如在文本里一个头可能关注语法依赖另一个头可能关注相邻词关系综合起来表达更丰富。从代码上讲多头就是在单头基础上做两件事把d_model切分成num_heads份每份独立计算注意力最后拼接起来。6.2 最小改造切分、计算、拼接假设d_model8head数4每个头d_k2。投影输出shape从(batch, seq_len, 8)变成(batch, seq_len, 4, 2)然后调换顺序变成(batch, 4, seq_len, 2)。这样每个头就独立成为batch维上的一“份”可以复用单头的矩阵操作。算完之后再把(batch, 4, seq_len, 2)转回(batch, seq_len, 4, 2)reshape成(batch, seq_len, 8)过一遍out_proj。def split_heads(tensor, num_heads): # tensor: (batch, seq_len, d_model) batch, seq_len, d_model tensor.shape d_k d_model // num_heads # (batch, seq_len, num_heads, d_k) - (batch, num_heads, seq_len, d_k) return tensor.view(batch, seq_len, num_heads, d_k).transpose(1, 2) def concat_heads(tensor, num_heads): # tensor: (batch, num_heads, seq_len, d_k) batch, num_heads, seq_len, d_k tensor.shape d_model num_heads * d_k # (batch, seq_len, num_heads, d_k) - (batch, seq_len, d_model) return tensor.transpose(1, 2).reshape(batch, seq_len, d_model)通俗地打个比方单头注意力像一个人只戴一副眼镜看问题多头是同时戴好几副不同角度的眼镜每副眼镜过滤出不同的特征最后综合判断。多头注意力的完整实现和计算细节我会放在系列下一篇里继续拆这篇把单头吃透就足够应付大部分面试和项目场景了。最后说点我的个人体会。Self-Attention学起来最笨但最有效的方法就是亲手推一遍shape和数值。别看公式里只有三个字母真正动手的时候Q、K、V三个投影、mask的广播维度、缩放因子这些细节全都会浮出来。这个系列后面还会继续拆多头注意力和完整的Transformer编码器到时候你会发现底层全都是这边写过的几行代码。先把手头的单头实现反复跑明白后面的路会顺很多。
