注意力机制这几个字但凡碰过Transformer的人都不可能绕开。它既是整个架构的灵魂也是绝大多数模型面试必问的点更是很多人从看懂架构图到真正手写实现之间最大的一道坎。这篇内容打算聊透注意力机制的第一阶段把数学公式一步步落地成可运行的代码把Q、K、V的来源和语义边界搞清楚把缩放点积的前因后果说明白。适合正在啃Transformer源码的读者、准备模型相关面试的人、以及所有对为什么这么设计有一种执念的人。先说一下这次计划的边界。我只聊单头注意力不聊多头、不做完整Transformer块。原因很简单多头注意力是单头的拆分和并行Encoder Decoder的完整结构也只是在注意力外围套壳不把单头搞扎实后面全是空中楼阁。另外我会尽量用先建直觉、再讲公式、最后上代码的顺序推进每一步都告诉你这里的设计动机是什么踩过什么样的坑。1. 注意力机制到底在解决什么问题1.1 一个关于检索的直觉认知先抛开所有公式用生活场景建立直觉。假设你走进一座图书馆想找一本关于深度学习的数学基础的书但你不知道具体书名。你的做法是先拿这张需求描述去和书架上的每一本书比对每本书的标题、目录、简介都和你的需求有相关性有的高度吻合有的完全无关。最后你大概率会取走那本相关性最高的书同时顺手翻一翻相关度比较高的几本。这就是注意力的核心思想你不是把整个图书馆白都背下来而是带着一个查询目标主动去查找和它最匹配的信息再按匹配程度决定吸收多少内容。在早年的序列模型里模型的记忆方式非常粗暴。一个编码器把整个输入句子压缩成最后一个隐状态向量解码器只能从这个固定长度的向量里获取全部信息句子一长前面的内容基本被压缩没了。后来出现了Bahdanau注意力解码时每生成一个词就回头去看编码器所有位置的隐状态按相似度做加权相当于每次生成时都能重新翻阅原句。这套机制已经非常接近人类翻译的查词习惯效果也明显提升。但到这里仍然是编码器产生记忆、解码器去查的两段式结构。Transformer往前再走了一大步让每个位置自己去关注同一层里的其他位置也就是Self-Attention。句子中每个单词都同时扮演提问者和被查阅者两个角色直接两两交互不再借助循环结构的顺序传递。这样带来的结果就是任意两个位置之间的依赖距离缩短为一步长距离依赖问题被彻底换了一种解法而且所有位置可以并行计算效率和表现都上了一个台阶。1.2 Self-Attention和传统结构的本质差别RNN的序列建模受限于时间步的递归第t个词要看到第1个词需要经过t-1次状态传递路径长、信息衰减、无法利用GPU并行。CNN虽然可以通过大卷积核扩大感受野但感受野是有限的加深层数才能覆盖更远的距离代价也不小。Self-Attention的做法是输入L个token直接把所有token两两计算相似度每个token都能在一层之内看到整个序列的全部内容。这种全连接式的关系建模能力让它对长距离依赖的处理变得极其自然。当然代价也很直接——计算量是序列长度的平方L个token要算L²次两两交互。这也是为什么后来有一堆优化工作都在尝试降低这个平方复杂度但那是另一篇文章的故事了这里先不展开。另外还有一个容易被忽略的点Self-Attention本身是集合操作它对顺序不敏感。把句子里的词打乱顺序两两之间的相似度关系不会变这显然不符合语言的特点。所以Transformer在输入里加了位置编码先给每个token注入位置信息再做注意力这才能让模型分得清我爱你和你爱我的区别。如果你在复现代码时发现结果怪怪的先检查一下是不是忘了位置编码这种低级错误我见过不止一次。2. 数学原理从输入到注意力输出分四步走2.1 QKV到底是怎么来的为什么偏偏是三个矩阵进入公式之前要把Q、K、V三个字母的身份对齐。Q是查询QueryK是键KeyV是值Value。你可以继续用图书馆的例子Q是你脑子里的需求描述K是每本书的标题和索引号V是书本本身的内容。先拿Q去和所有K比对算出每本书和需求的相关度再用这些相关度作为权重把V加权加起来得到最终检索结果。如果只有一个输入序列X那Q、K、V都从X变化而来。具体做法是拿X分别乘三个参数矩阵W_Q、W_K、W_V线性变换出三个新矩阵。有人会问既然是同一个X为什么不直接用X自己算相似度非要先变换一遍这里的关键是语义角色的分离。如果不做变换X和X做点积只是单纯在原始特征空间里算cosine相似度但我提问和我被提问在语言里是两种不同性质的任务。一个词作为查询时它想知道的是我要找什么相关信息作为键时它要回答的是我身上有哪些可被匹配的特征。这两种语义用同一份向量表达信息容量不够而且模型没有机会针对两种角色分别做特征变换。引入三个可学习的参数矩阵本质上就是让模型在训练过程中学会什么样的查询表示最适合去搜索、什么样的键表示最适合被搜索、什么样的值表示最适合参与加权输出。三个角色的特征空间被拉开模型表达能力自然更强。用最直白的话总结Q、K、V是把同一段信息投影到三种不同功能视角下的产物。这里的W_Q、W_K、W_V全部是可学习参数维度都是(d_model, d_k)d_k通常取d_model除以头数。这样每个头都把原来的高维空间切分成一个子空间可以关注不同类型的模式。2.2 点积、缩放、Softmax每一步在做什么拿到Q、K、V之后第一步是计算Q和K的点积。Q的形状是(L, d_k)K的形状也是(L, d_k)把K转置成(d_k, L)做矩阵乘法得到(L, L)的相似度矩阵。这里的第i行第j列就是第i个token作为查询和第j个token作为键的相似度得分。点积的直觉是两个向量方向越一致、模长越大内积越大说明越匹配。但这里有一个陷阱。如果d_k比较大比如64或128两个随机向量点积的方差也会变大导致相似度得分分布变得很极端有的分非常高有的分非常低。把这个极端分布扔进softmax会得到什么结果高分会把概率几乎全部吸走低分直接趋近于0整个概率分布变成一个几乎one-hot的形态。此时对应位置的梯度会非常小模型几乎学不动。为了解决这个问题原论文在点积之后除以sqrt(d_k)把得分的方差拉回一个稳定区间。为什么要用sqrt(d_k)而不是d_k这里有个简单的推导。假设Q和K中的每个元素都是均值为0、方差为1的随机变量那么两个长度为d_k的向量做点积结果的期望是0方差是d_k。除以sqrt(d_k)之后方差就恢复成1。这个操作纯粹是为了数值稳定让softmax的输入分布保持在一个合理的尺度内避免梯度消失。我见过很多人把这个缩放当成可有可无的细节实际去掉它训练时非常容易出现loss震荡甚至不收敛的情况尤其是模型变深以后。缩放之后的相似度矩阵进入softmax按行做归一化每行加起来等于1。softmax的指数函数会把负分压小、正分放大让模型能表达重点关注哪些位置、忽略哪些位置的决策。最后把归一化后的概率分布作为权重和V矩阵加权求和每一行的输出都是所有位置V向量的加权平均权重就是该行对应的注意力分布。这一步相当于带着从整个序列中检索出来的信息结果继续往下传递。2.3 一个公式看懂全局把上面的逻辑压缩成一个公式Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这个公式里阻断的顺序是先算相似度再缩放稳定再概率化再加权。每一步单独拎出来都不复杂合在一起就成了Transformer的基石。建议阅读源码的时候把每个中间结果的shape都在纸上画一遍例如输入是(L, d_model)经过W_Q得到(L, d_k)的Q和K的(L, d_k)转置相乘得到(L, L)的相似度矩阵再和(L, d_k)的V相乘输出又回到(L, d_k)。整个过程中你只需要盯住哪两个维度的矩阵相乘、结果的shape又是什么盯住一次就再也不会乱了。3. 代码实现从零手写单头注意力3.1 准备工作与一个容易复现的小样例先安装好PyTorch建议用1.13以上版本。为了验证结果我会用一个小尺寸样例batch_size2序列长度4d_model8。这个规模你可以非常清楚地观察每一步张量维度的变化比直接上手大模型要容易得多。初始化的代码如下加上一时注释方便后面逐行对照import torch import torch.nn as nn import torch.nn.functional as F import math torch.manual_seed(42) batch_size 2 seq_len 4 d_model 8 d_k 8 # 单头所以d_k等于d_model x torch.randn(batch_size, seq_len, d_model) print(输入x形状:, x.shape) # (2, 4, 8)这里我把输入形状设计成PyTorch标准的(batch, seq_len, d_model)格式后面所有计算都基于这个布局展开。如果哪一步的shape突然不匹配了多半是维度顺序理解错了别慌用print(张量.shape)一行一行查很快能定位问题。3.2 QKV投影与缩放点积注意力的完整实现下面直接给出单头注意力的完整PyTorch实现。我先定义三个线性层做QKV投影再实现缩放点积的attention函数全部写在一个模块里便于直接复用。class SingleHeadAttention(nn.Module): def __init__(self, d_model, d_k): super().__init__() self.d_k d_k 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_k, biasFalse) def forward(self, x, maskNone, return_attnFalse): Q self.W_Q(x) # (batch, seq, d_k) K self.W_K(x) # (batch, seq, d_k) V self.W_V(x) # (batch, seq, d_k) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # scores形状: (batch, seq, seq) 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) # output形状: (batch, seq, d_k) if return_attn: return output, attn_weights return output逐行解释几个关键点第一W_Q、W_K、W_V都是nn.Linear默认不带bias这是因为原始论文里就没有在投影层加偏置加不加对结果影响很小但去掉可以减少一点参数量。第二K.transpose(-2, -1)是把每个batch内部的序列维度转置注意是倒数两个维度交换不要写成K.transpose(1, 2)除非你确定输入没有batch维。第三masked_fill(mask 0, -1e9)这一步很关键它的作用是把mask中为0的位置填充成非常小的数这样经过softmax时几乎就是0等价于不关注这些位置。为什么要填-1e9而不是0因为softmax里e的0次方等于1会被当成正常值参与归一化填一个极小的负数e的负大数次方趋近于0才能真正屏蔽掉。这是初学最容易踩的坑之一。第四F.softmax在最后一个维度上做归一化也就是对每个query在所有key上的分数归一化方向不能搞错。我还特意加了return_attn参数用来输出注意力权重矩阵调试时特别有用。你可以把注意力权重打印出来看看模型是不是真的关注到了合理的位置这一步对检查和理解内部行为意义很大。3.3 用一个具体算例验证输出维度写代码最怕的就是维度错乱。这里我建议你把上面模块实例化然后运行一下观察每一步的shape变化model SingleHeadAttention(d_model, d_k) output, attn_weights model(x, return_attnTrue) print(输出形状:, output.shape) # (2, 4, 8) print(注意力权重形状:, attn_weights.shape) # (2, 4, 4)attn_weights是(2, 4, 4)的分布矩阵第i行的含义是第i个token对整个序列4个token的注意力权重分布每一行加起来应该等于1。你可以直接验证一下print(attn_weights[0].sum(dim-1)) # 每行之和都应该是1如果输出不是全1说明softmax的dim用错了或者没生效。这个简单的验证能帮你省掉大量排查维度的时间。我在调试自己的实现时几乎每次都会先打印注意力权重的行和确认归一化没有问题再继续往下看。3.4 和PyTorch内置的MultiheadAttention对齐验证手写实现最怕的就是跟官方实现不一致一旦出错整个模型效果都会跑偏。好在我们可以拿PyTorch内置的nn.MultiheadAttention来交叉验证。这里有一个容易让人困惑的地方官方的nn.MultiheadAttention在输入单头时行为和我们上面的实现是否完全一致严格说计算逻辑一致但官方模块和其他接口约定不同需要做一些对齐操作。验证的思路是构造相同的Q、K、V分别用手写模块和官方模块计算一遍对比输出是否一致。由于两个模块的初始权重不同直接对比输出肯定不一样。我先把手写模块的Q、K、V投影权重设成官方模块的对应权重再做推理就能验证计算逻辑是否一致。official nn.MultiheadAttention(embed_dimd_model, num_heads1, biasFalse, batch_firstTrue) # 把手写模块的投影权重同步为官方权重 with torch.no_grad(): model.W_Q.weight.copy_(official.in_proj_weight[0:d_model]) model.W_K.weight.copy_(official.in_proj_weight[d_model:2*d_model]) model.W_V.weight.copy_(official.in_proj_weight[2*d_model:3*d_model]) o_out, o_weights official(x, x, x, average_attn_weightsTrue) my_out, my_weights model(x, return_attnTrue) print(官方输出:, o_out) print(手写输出:, my_out) print(最大误差:, (o_out - my_out).abs().max().item())期待的结果是最大误差在1e-6量级基本可以认为两个实现数学上等价。如果你用这个验证方式发现差异很大先检查一下official里batch_firstTrue是否正确再确认in_proj_weight的切片位置有没有拿对。官方实现里当query、key、value都指向同一个输入x时会用三个连续的投影矩阵一次性把Q、K、V算出来分别存储在一段连续的权重大向量里。所以需要把这个大向量按顺序切片赋予到我的三个独立Linear层中。这种对齐验证是我每次手写注意力实现都会做的一步测试。它不能证明你的模块能训练好但能证明你至少在数学实现层面没有出大差错排查范围瞬间缩小很多。4. 常见问题与排查技巧实录4.1 Mask到底应该怎么填填在哪里注意力计算中有时候你需要屏蔽掉某些位置比如文本padding部分、翻译解码时的未来信息。理论上mask有两种常见处理方式一种是把对应相似度分数直接设为负无穷另一种是设为一个绝对值很大的负数比如-1e9。推荐文章的做法是用-1e9原因前面提过softmax对非屏蔽位置的归一化不受影响而屏蔽位置经过指数计算后几乎为0。如果你直接填0那softmax中e的0次方等于1被屏蔽的位置仍然占有了概率质量等于没屏蔽。实际操作中还要注意一个细节mask矩阵的shape。常见的是(seq_len, seq_len)的布尔矩阵或者(batch, seq_len)的padding mask。在矩阵乘法中scores的形状是(2, 4, 4)如果你用的是(batch, seq_len)的形状就需要先unsqueeze扩展维度变成(batch, 1, seq_len)让mask能正确广播到最后一个维度上。我见过很多人在这块一头雾水其实只要记住一个规则mask要广播到scores的倒数第二维也就是query所在的维度上所以要么是(seq_len, seq_len)要么是(batch, 1, seq_len)。4.2 维度顺序和transpose的经典坑PyTorch里张量的shape顺序是(batch, seq_len, d_model)但在做矩阵乘的时候K需要转置最后一个维度和倒数第二个维度。如果直接用K.transpose(1, 2)在只有三维张量时才等价于transpose(-2, -1)。但一旦你后续加了多头、或者把batch和seq换成了不确定的动态维度这种硬编码写法就特别容易出问题。我这个教训很具体有段时间为了省事把所有transpose都写成了transpose(1, 2)后来换了一个带头维的数据结构形状从(batch, head, seq, d_k)变成了四维直接报错。改成统一使用transpose(-2, -1)之后这个问题再也没有出现过。建议你也养成用负数索引写transpose的习惯这是PyTorch里最稳妥的写法。另一个维度坑是在输出层。单头注意力输入时d_modeld_k输出仍是d_k形状和输入一致看起来没问题。但进入多头时每个头输出的d_k会小于d_model需要先拼接再线性变换回d_model。很多初学多头的代码跑不起来多半是卡在怎么把多个头的输出合回去这一步。这个我会在下一篇展开但你现在就要有这个意识。4.3 训练不收敛或NaN多半是数值稳定性问题如果你去掉缩放因子sqrt(d_k)训练时会有很高的概率出现loss震荡。原因我在前面推导过dk越大点积方差越大softmax输出越接近one-hot反向传播时这部分梯度会非常小。这种问题很难从loss曲线的形态直接判断因为不是完全发散的NaN就是loss降不下去。我自己的排查经验是先固定随机种子构造一个极小的训练集把单头注意力单独拎出来跑一遍观察loss能不能稳定下降。如果去掉缩放能跑但掉得慢加上缩放却掉得更快那基本就是缩放因子用错或者位置搞错。另外softmax本身也有数值不稳定性当scores里出现非常大的正数时指数函数会直接上溢出到inf这时候softmax输出全是NaN。稳健的写法是在softmax里先对每行减掉最大值再算指数也就是所谓的max-trick。PyTorch内置的F.softmax已经做了这个处理你可以放心用但如果你是自己实现的softmax一定记得写这一步。最后提醒一点如果你在注意力权重的计算之前dropout应该放在softmax之后、和V相乘之前。这是Transformer原论文的标准做法目的是对注意力分布加噪声防止过拟合。放错位置不仅起不到正则效果还可能破坏softmax的归一化性质导致权重之和不再等于1。5. 一点个人经验和后续预告代码写到能跑通、能对齐官方实现大概只是理解注意力的第一步。我自己的感觉是真正吃透这个机制靠的不是看公式而是反复用手写实现去替换官方模块然后把每个变量名的shape和含义都写出来对着算一遍。最开始可能很慢但一旦过了这道坎后面再看多头、再看Transformer Encoder会觉得整个框架都是顺下来的。还有一个贴心的小建议如果你在调试过程中发现输出很怪先别怀疑注意力公式先从输入数据、mask、维度入手排查。我统计过自己写代码时遇到的bug百分之八十都出在shape和mask广播上公式本身反而是最不容易写错的部分。这篇只展开了单头注意力注意力机制的完整应用还需要覆盖多头拆分的逻辑、mask变体的细节、以及注意力在Encoder和Decoder之间是如何交互的。下一篇我会直接基于这里实现的SingleHeadAttention扩展成完整的MultiHeadAttention并且用翻译任务的实例验证它的实际训练效果。如果你自己把文中这段代码跑通了那下一阶段的内容对你来说会非常顺。
注意力机制这几个字但凡碰过Transformer的人都不可能绕开。它既是整个架构的灵魂也是绝大多数模型面试必问的点更是很多人从看懂架构图到真正手写实现之间最大的一道坎。这篇内容打算聊透注意力机制的第一阶段把数学公式一步步落地成可运行的代码把Q、K、V的来源和语义边界搞清楚把缩放点积的前因后果说明白。适合正在啃Transformer源码的读者、准备模型相关面试的人、以及所有对为什么这么设计有一种执念的人。先说一下这次计划的边界。我只聊单头注意力不聊多头、不做完整Transformer块。原因很简单多头注意力是单头的拆分和并行Encoder Decoder的完整结构也只是在注意力外围套壳不把单头搞扎实后面全是空中楼阁。另外我会尽量用先建直觉、再讲公式、最后上代码的顺序推进每一步都告诉你这里的设计动机是什么踩过什么样的坑。1. 注意力机制到底在解决什么问题1.1 一个关于检索的直觉认知先抛开所有公式用生活场景建立直觉。假设你走进一座图书馆想找一本关于深度学习的数学基础的书但你不知道具体书名。你的做法是先拿这张需求描述去和书架上的每一本书比对每本书的标题、目录、简介都和你的需求有相关性有的高度吻合有的完全无关。最后你大概率会取走那本相关性最高的书同时顺手翻一翻相关度比较高的几本。这就是注意力的核心思想你不是把整个图书馆白都背下来而是带着一个查询目标主动去查找和它最匹配的信息再按匹配程度决定吸收多少内容。在早年的序列模型里模型的记忆方式非常粗暴。一个编码器把整个输入句子压缩成最后一个隐状态向量解码器只能从这个固定长度的向量里获取全部信息句子一长前面的内容基本被压缩没了。后来出现了Bahdanau注意力解码时每生成一个词就回头去看编码器所有位置的隐状态按相似度做加权相当于每次生成时都能重新翻阅原句。这套机制已经非常接近人类翻译的查词习惯效果也明显提升。但到这里仍然是编码器产生记忆、解码器去查的两段式结构。Transformer往前再走了一大步让每个位置自己去关注同一层里的其他位置也就是Self-Attention。句子中每个单词都同时扮演提问者和被查阅者两个角色直接两两交互不再借助循环结构的顺序传递。这样带来的结果就是任意两个位置之间的依赖距离缩短为一步长距离依赖问题被彻底换了一种解法而且所有位置可以并行计算效率和表现都上了一个台阶。1.2 Self-Attention和传统结构的本质差别RNN的序列建模受限于时间步的递归第t个词要看到第1个词需要经过t-1次状态传递路径长、信息衰减、无法利用GPU并行。CNN虽然可以通过大卷积核扩大感受野但感受野是有限的加深层数才能覆盖更远的距离代价也不小。Self-Attention的做法是输入L个token直接把所有token两两计算相似度每个token都能在一层之内看到整个序列的全部内容。这种全连接式的关系建模能力让它对长距离依赖的处理变得极其自然。当然代价也很直接——计算量是序列长度的平方L个token要算L²次两两交互。这也是为什么后来有一堆优化工作都在尝试降低这个平方复杂度但那是另一篇文章的故事了这里先不展开。另外还有一个容易被忽略的点Self-Attention本身是集合操作它对顺序不敏感。把句子里的词打乱顺序两两之间的相似度关系不会变这显然不符合语言的特点。所以Transformer在输入里加了位置编码先给每个token注入位置信息再做注意力这才能让模型分得清我爱你和你爱我的区别。如果你在复现代码时发现结果怪怪的先检查一下是不是忘了位置编码这种低级错误我见过不止一次。2. 数学原理从输入到注意力输出分四步走2.1 QKV到底是怎么来的为什么偏偏是三个矩阵进入公式之前要把Q、K、V三个字母的身份对齐。Q是查询QueryK是键KeyV是值Value。你可以继续用图书馆的例子Q是你脑子里的需求描述K是每本书的标题和索引号V是书本本身的内容。先拿Q去和所有K比对算出每本书和需求的相关度再用这些相关度作为权重把V加权加起来得到最终检索结果。如果只有一个输入序列X那Q、K、V都从X变化而来。具体做法是拿X分别乘三个参数矩阵W_Q、W_K、W_V线性变换出三个新矩阵。有人会问既然是同一个X为什么不直接用X自己算相似度非要先变换一遍这里的关键是语义角色的分离。如果不做变换X和X做点积只是单纯在原始特征空间里算cosine相似度但我提问和我被提问在语言里是两种不同性质的任务。一个词作为查询时它想知道的是我要找什么相关信息作为键时它要回答的是我身上有哪些可被匹配的特征。这两种语义用同一份向量表达信息容量不够而且模型没有机会针对两种角色分别做特征变换。引入三个可学习的参数矩阵本质上就是让模型在训练过程中学会什么样的查询表示最适合去搜索、什么样的键表示最适合被搜索、什么样的值表示最适合参与加权输出。三个角色的特征空间被拉开模型表达能力自然更强。用最直白的话总结Q、K、V是把同一段信息投影到三种不同功能视角下的产物。这里的W_Q、W_K、W_V全部是可学习参数维度都是(d_model, d_k)d_k通常取d_model除以头数。这样每个头都把原来的高维空间切分成一个子空间可以关注不同类型的模式。2.2 点积、缩放、Softmax每一步在做什么拿到Q、K、V之后第一步是计算Q和K的点积。Q的形状是(L, d_k)K的形状也是(L, d_k)把K转置成(d_k, L)做矩阵乘法得到(L, L)的相似度矩阵。这里的第i行第j列就是第i个token作为查询和第j个token作为键的相似度得分。点积的直觉是两个向量方向越一致、模长越大内积越大说明越匹配。但这里有一个陷阱。如果d_k比较大比如64或128两个随机向量点积的方差也会变大导致相似度得分分布变得很极端有的分非常高有的分非常低。把这个极端分布扔进softmax会得到什么结果高分会把概率几乎全部吸走低分直接趋近于0整个概率分布变成一个几乎one-hot的形态。此时对应位置的梯度会非常小模型几乎学不动。为了解决这个问题原论文在点积之后除以sqrt(d_k)把得分的方差拉回一个稳定区间。为什么要用sqrt(d_k)而不是d_k这里有个简单的推导。假设Q和K中的每个元素都是均值为0、方差为1的随机变量那么两个长度为d_k的向量做点积结果的期望是0方差是d_k。除以sqrt(d_k)之后方差就恢复成1。这个操作纯粹是为了数值稳定让softmax的输入分布保持在一个合理的尺度内避免梯度消失。我见过很多人把这个缩放当成可有可无的细节实际去掉它训练时非常容易出现loss震荡甚至不收敛的情况尤其是模型变深以后。缩放之后的相似度矩阵进入softmax按行做归一化每行加起来等于1。softmax的指数函数会把负分压小、正分放大让模型能表达重点关注哪些位置、忽略哪些位置的决策。最后把归一化后的概率分布作为权重和V矩阵加权求和每一行的输出都是所有位置V向量的加权平均权重就是该行对应的注意力分布。这一步相当于带着从整个序列中检索出来的信息结果继续往下传递。2.3 一个公式看懂全局把上面的逻辑压缩成一个公式Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这个公式里阻断的顺序是先算相似度再缩放稳定再概率化再加权。每一步单独拎出来都不复杂合在一起就成了Transformer的基石。建议阅读源码的时候把每个中间结果的shape都在纸上画一遍例如输入是(L, d_model)经过W_Q得到(L, d_k)的Q和K的(L, d_k)转置相乘得到(L, L)的相似度矩阵再和(L, d_k)的V相乘输出又回到(L, d_k)。整个过程中你只需要盯住哪两个维度的矩阵相乘、结果的shape又是什么盯住一次就再也不会乱了。3. 代码实现从零手写单头注意力3.1 准备工作与一个容易复现的小样例先安装好PyTorch建议用1.13以上版本。为了验证结果我会用一个小尺寸样例batch_size2序列长度4d_model8。这个规模你可以非常清楚地观察每一步张量维度的变化比直接上手大模型要容易得多。初始化的代码如下加上一时注释方便后面逐行对照import torch import torch.nn as nn import torch.nn.functional as F import math torch.manual_seed(42) batch_size 2 seq_len 4 d_model 8 d_k 8 # 单头所以d_k等于d_model x torch.randn(batch_size, seq_len, d_model) print(输入x形状:, x.shape) # (2, 4, 8)这里我把输入形状设计成PyTorch标准的(batch, seq_len, d_model)格式后面所有计算都基于这个布局展开。如果哪一步的shape突然不匹配了多半是维度顺序理解错了别慌用print(张量.shape)一行一行查很快能定位问题。3.2 QKV投影与缩放点积注意力的完整实现下面直接给出单头注意力的完整PyTorch实现。我先定义三个线性层做QKV投影再实现缩放点积的attention函数全部写在一个模块里便于直接复用。class SingleHeadAttention(nn.Module): def __init__(self, d_model, d_k): super().__init__() self.d_k d_k 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_k, biasFalse) def forward(self, x, maskNone, return_attnFalse): Q self.W_Q(x) # (batch, seq, d_k) K self.W_K(x) # (batch, seq, d_k) V self.W_V(x) # (batch, seq, d_k) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # scores形状: (batch, seq, seq) 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) # output形状: (batch, seq, d_k) if return_attn: return output, attn_weights return output逐行解释几个关键点第一W_Q、W_K、W_V都是nn.Linear默认不带bias这是因为原始论文里就没有在投影层加偏置加不加对结果影响很小但去掉可以减少一点参数量。第二K.transpose(-2, -1)是把每个batch内部的序列维度转置注意是倒数两个维度交换不要写成K.transpose(1, 2)除非你确定输入没有batch维。第三masked_fill(mask 0, -1e9)这一步很关键它的作用是把mask中为0的位置填充成非常小的数这样经过softmax时几乎就是0等价于不关注这些位置。为什么要填-1e9而不是0因为softmax里e的0次方等于1会被当成正常值参与归一化填一个极小的负数e的负大数次方趋近于0才能真正屏蔽掉。这是初学最容易踩的坑之一。第四F.softmax在最后一个维度上做归一化也就是对每个query在所有key上的分数归一化方向不能搞错。我还特意加了return_attn参数用来输出注意力权重矩阵调试时特别有用。你可以把注意力权重打印出来看看模型是不是真的关注到了合理的位置这一步对检查和理解内部行为意义很大。3.3 用一个具体算例验证输出维度写代码最怕的就是维度错乱。这里我建议你把上面模块实例化然后运行一下观察每一步的shape变化model SingleHeadAttention(d_model, d_k) output, attn_weights model(x, return_attnTrue) print(输出形状:, output.shape) # (2, 4, 8) print(注意力权重形状:, attn_weights.shape) # (2, 4, 4)attn_weights是(2, 4, 4)的分布矩阵第i行的含义是第i个token对整个序列4个token的注意力权重分布每一行加起来应该等于1。你可以直接验证一下print(attn_weights[0].sum(dim-1)) # 每行之和都应该是1如果输出不是全1说明softmax的dim用错了或者没生效。这个简单的验证能帮你省掉大量排查维度的时间。我在调试自己的实现时几乎每次都会先打印注意力权重的行和确认归一化没有问题再继续往下看。3.4 和PyTorch内置的MultiheadAttention对齐验证手写实现最怕的就是跟官方实现不一致一旦出错整个模型效果都会跑偏。好在我们可以拿PyTorch内置的nn.MultiheadAttention来交叉验证。这里有一个容易让人困惑的地方官方的nn.MultiheadAttention在输入单头时行为和我们上面的实现是否完全一致严格说计算逻辑一致但官方模块和其他接口约定不同需要做一些对齐操作。验证的思路是构造相同的Q、K、V分别用手写模块和官方模块计算一遍对比输出是否一致。由于两个模块的初始权重不同直接对比输出肯定不一样。我先把手写模块的Q、K、V投影权重设成官方模块的对应权重再做推理就能验证计算逻辑是否一致。official nn.MultiheadAttention(embed_dimd_model, num_heads1, biasFalse, batch_firstTrue) # 把手写模块的投影权重同步为官方权重 with torch.no_grad(): model.W_Q.weight.copy_(official.in_proj_weight[0:d_model]) model.W_K.weight.copy_(official.in_proj_weight[d_model:2*d_model]) model.W_V.weight.copy_(official.in_proj_weight[2*d_model:3*d_model]) o_out, o_weights official(x, x, x, average_attn_weightsTrue) my_out, my_weights model(x, return_attnTrue) print(官方输出:, o_out) print(手写输出:, my_out) print(最大误差:, (o_out - my_out).abs().max().item())期待的结果是最大误差在1e-6量级基本可以认为两个实现数学上等价。如果你用这个验证方式发现差异很大先检查一下official里batch_firstTrue是否正确再确认in_proj_weight的切片位置有没有拿对。官方实现里当query、key、value都指向同一个输入x时会用三个连续的投影矩阵一次性把Q、K、V算出来分别存储在一段连续的权重大向量里。所以需要把这个大向量按顺序切片赋予到我的三个独立Linear层中。这种对齐验证是我每次手写注意力实现都会做的一步测试。它不能证明你的模块能训练好但能证明你至少在数学实现层面没有出大差错排查范围瞬间缩小很多。4. 常见问题与排查技巧实录4.1 Mask到底应该怎么填填在哪里注意力计算中有时候你需要屏蔽掉某些位置比如文本padding部分、翻译解码时的未来信息。理论上mask有两种常见处理方式一种是把对应相似度分数直接设为负无穷另一种是设为一个绝对值很大的负数比如-1e9。推荐文章的做法是用-1e9原因前面提过softmax对非屏蔽位置的归一化不受影响而屏蔽位置经过指数计算后几乎为0。如果你直接填0那softmax中e的0次方等于1被屏蔽的位置仍然占有了概率质量等于没屏蔽。实际操作中还要注意一个细节mask矩阵的shape。常见的是(seq_len, seq_len)的布尔矩阵或者(batch, seq_len)的padding mask。在矩阵乘法中scores的形状是(2, 4, 4)如果你用的是(batch, seq_len)的形状就需要先unsqueeze扩展维度变成(batch, 1, seq_len)让mask能正确广播到最后一个维度上。我见过很多人在这块一头雾水其实只要记住一个规则mask要广播到scores的倒数第二维也就是query所在的维度上所以要么是(seq_len, seq_len)要么是(batch, 1, seq_len)。4.2 维度顺序和transpose的经典坑PyTorch里张量的shape顺序是(batch, seq_len, d_model)但在做矩阵乘的时候K需要转置最后一个维度和倒数第二个维度。如果直接用K.transpose(1, 2)在只有三维张量时才等价于transpose(-2, -1)。但一旦你后续加了多头、或者把batch和seq换成了不确定的动态维度这种硬编码写法就特别容易出问题。我这个教训很具体有段时间为了省事把所有transpose都写成了transpose(1, 2)后来换了一个带头维的数据结构形状从(batch, head, seq, d_k)变成了四维直接报错。改成统一使用transpose(-2, -1)之后这个问题再也没有出现过。建议你也养成用负数索引写transpose的习惯这是PyTorch里最稳妥的写法。另一个维度坑是在输出层。单头注意力输入时d_modeld_k输出仍是d_k形状和输入一致看起来没问题。但进入多头时每个头输出的d_k会小于d_model需要先拼接再线性变换回d_model。很多初学多头的代码跑不起来多半是卡在怎么把多个头的输出合回去这一步。这个我会在下一篇展开但你现在就要有这个意识。4.3 训练不收敛或NaN多半是数值稳定性问题如果你去掉缩放因子sqrt(d_k)训练时会有很高的概率出现loss震荡。原因我在前面推导过dk越大点积方差越大softmax输出越接近one-hot反向传播时这部分梯度会非常小。这种问题很难从loss曲线的形态直接判断因为不是完全发散的NaN就是loss降不下去。我自己的排查经验是先固定随机种子构造一个极小的训练集把单头注意力单独拎出来跑一遍观察loss能不能稳定下降。如果去掉缩放能跑但掉得慢加上缩放却掉得更快那基本就是缩放因子用错或者位置搞错。另外softmax本身也有数值不稳定性当scores里出现非常大的正数时指数函数会直接上溢出到inf这时候softmax输出全是NaN。稳健的写法是在softmax里先对每行减掉最大值再算指数也就是所谓的max-trick。PyTorch内置的F.softmax已经做了这个处理你可以放心用但如果你是自己实现的softmax一定记得写这一步。最后提醒一点如果你在注意力权重的计算之前dropout应该放在softmax之后、和V相乘之前。这是Transformer原论文的标准做法目的是对注意力分布加噪声防止过拟合。放错位置不仅起不到正则效果还可能破坏softmax的归一化性质导致权重之和不再等于1。5. 一点个人经验和后续预告代码写到能跑通、能对齐官方实现大概只是理解注意力的第一步。我自己的感觉是真正吃透这个机制靠的不是看公式而是反复用手写实现去替换官方模块然后把每个变量名的shape和含义都写出来对着算一遍。最开始可能很慢但一旦过了这道坎后面再看多头、再看Transformer Encoder会觉得整个框架都是顺下来的。还有一个贴心的小建议如果你在调试过程中发现输出很怪先别怀疑注意力公式先从输入数据、mask、维度入手排查。我统计过自己写代码时遇到的bug百分之八十都出在shape和mask广播上公式本身反而是最不容易写错的部分。这篇只展开了单头注意力注意力机制的完整应用还需要覆盖多头拆分的逻辑、mask变体的细节、以及注意力在Encoder和Decoder之间是如何交互的。下一篇我会直接基于这里实现的SingleHeadAttention扩展成完整的MultiHeadAttention并且用翻译任务的实例验证它的实际训练效果。如果你自己把文中这段代码跑通了那下一阶段的内容对你来说会非常顺。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
