如果说Transformer是一台精心设计的机器那注意力机制就是它的动力核心。我在刚接触这块内容时最大的困惑不是公式本身有多难而是论文里几行矩阵运算和实际能跑的代码之间始终隔着一层“知道原理但写不出来”的窗户纸。这篇内容会从Q查询、K键、V值三个向量开始把缩放点积注意力的数学原理逐步拆开再用PyTorch代码从头实现一遍所有维度变化都会标清楚。适合两类人一是刚入门的同学需要一个从公式到实现的完整闭环二是已经用过现成模型但想搞清楚内部到底发生了什么的人。1. 先搞清楚注意力机制在Transformer里到底解决了什么问题1.1 在注意力出现之前序列建模为什么让人头疼很多讲注意力的文章一上来就怼公式但我觉得先理解“为什么需要它”更重要。在Transformer之前处理序列数据的主流工具是RNN或者LSTM。它们的特点是逐个时间步处理信息先读第一个词更新一个隐藏状态再读第二个词继续更新……这带来一个无法回避的问题——信息必须被“压缩”到一个固定长度的向量里。如果一句话有20个词目标是把整句话的语义塞进这个向量如果一句话有200个词还是塞进同一个向量。信息太多的时候后面的位置往往会“记不住”前面的关键内容尤其当关键内容和当前位置隔得很远时中间经过的每一个时间步都相当于一次有损压缩早期信息会被慢慢稀释掉。我记自己第一次训练RNN做文本分类时误差基本上全出在长句子上。句子短还行一旦超过四十个词模型就明显开始“装失忆”。当时还以为是调参不到位后来才发现这是串行建模的结构性瓶颈。另一个问题是并行性。RNN必须等前一个时间步算完才能算下一步GPU的并行能力在这里几乎发挥不出来。Transformer的核心贡献之一就是用注意力机制把“序列建模”变成了“矩阵运算”所有位置可以同时计算这才让大规模训练有了可行性。1.2 把注意力理解成一个“软检索”过程注意力机制最直观的类比是查字典或者更准确地说是带权重的检索。假设你要从一堆文档里找资料你手上有一个问题每个文档有一串标签文档本身是内容。你会先拿自己的问题去匹配标签找到最相关的文档然后提取它的内容。注意力机制做的是一模一样的事情QQuery就是你的问题KKey就是每个文档的标签VValue就是文档的实际内容。但有一个重要区别它不是“只挑最匹配的那一个文档”而是所有文档都拿一点相关度高的多拿一点相关度低的少拿一点。这个“拿多少”的比例就是softmax算出来的注意力权重。这个设计对并行计算特别友好。因为检索过程变成了Query矩阵和Key矩阵的批量点积所有查询、所有键可以一次性算完。为什么Transformer能在GPU上高效跑起来核心原因就在这一步。2. 从Query、Key、Value开始注意力公式的每一处设计意图2.1 Q、K、V三个向量从哪里来在实际模型中Q、K、V并不是凭空出现的它们是把输入向量分别乘以三个不同的权重矩阵得到的。假设输入序列为x形状是[batch_size, seq_len, d_model]其中d_model是每个token的向量维度。为了计算注意力需要定义三个投影矩阵W_q、W_k、W_v。于是Q x W_qK x W_kV x W_v这里的d_k是Q和K的维度d_v是V的维度。在Transformer原文里d_k和d_v通常相等但数学上并不强制。即使不相等也可以正确计算因为真正参与score计算的是Q和KV只负责提供最终的“内容”所以d_v独立于d_k是合法的。我见过不少初学者把d_k和d_v混为一谈其实它们的分工完全不同。为了方便推导下面的实现统一采用d_k d_v这也是绝大多数开源模型默认的做法。符号含义典型维度单头实际作用Q查询向量d_k负责发起“匹配请求”K键向量d_k负责响应查询请求V值向量d_v负责提供实际内容seq_len序列长度任意句子包含的token数量batch_size批大小任意一次性处理的句子数2.2 点积为什么能衡量两个向量的相似度Q和K做内积数学上叫点积几何意义是两个向量在方向上有多“一致”。对于两个维度相同的向量q和k点积是q · k |q| * |k| * cos(θ)其中θ是两个向量的夹角。夹角越小cos(θ)越接近1点积越大夹角接近90度时点积接近0方向相反时点积为负。关键在于如果两个向量方向相似说明它们的语义模式接近值得分配更多注意力如果方向差异很大说明关联弱注意力权重就低。这正是“匹配标签”的数学化表达。但注意点积的大小不仅取决于夹角还取决于向量长度。如果两个向量长度都很大即使夹角一般点积也可能很大。这就会导致一个问题点积值的绝对大小被人为拉高。后面讨论缩放因子时会重点展开。2.3 softmax为什么必须放在缩放之后计算出所有位置的匹配分数之后需要把它们转成“加起来等于1的权重”这一步用softmax完成softmax(z)_i exp(z_i) / Σ_j exp(z_j)softmax的特点是指数放大差异如果某个分数比另一个大一点经过指数之后会被放大得更加明显这让模型可以“有主见”地选择关注谁。但这里有讲究softmax的输入端如果数值过大输出会迅速饱和。举个例子对于向量[0, 3]softmax结果大约是[0.047, 0.953]已经有点极端了。如果输入变成[0, 30]softmax结果几乎就是[0, 1]另一个位置的概率小到可以忽略。这本身不一定是坏事——模型确实希望注意力集中。但问题是当softmax输出接近one-hot时它在那个位置上的梯度会非常小网络几乎学不到有用的信号。训练初期尤其危险参数刚开始随机初始化匹配分数本身没有意义如果这时候注意力直接“锁死”在某一个位置上后续很难再纠正过来。所以在softmax之前必须把分数控制在一个合理范围内。这就是下一节要讨论的缩放因子的由来。3. 除以√d_k不是玄学是保梯度的一次关键操作3.1 一组数值实验直接告诉你为什么需要缩放先看缩放因子长什么样。标准的注意力公式是Attention(Q, K, V) softmax(Q K^T / √d_k) V问题来了为什么一定要除以√d_k答案是控制方差。我们来推一下。假设q_i和k_i都是均值为0、方差为1的独立随机变量这是对向量元素的常见假设实际经过归一化后的向量也大致如此。点积的计算是score q_1*k_1 q_2*k_2 ... q_{d_k}*k_{d_k}每一项q_i*k_i的均值是0方差是1因为q_i和k_i独立乘积的方差等于各自方差之积即1。现在把d_k项加起来总的方差变成d_k。换句话说随着向量维度d_k增大点积值的标准差也在增大而且增长倍率是√d_k。可以做一个直观的实验d_k点积分数的标准差约QK点积典型取值范围约11[-3, 3]164[-12, 12]648[-24, 24]51222.6[-68, 68]如果直接把范围约[-24, 24]的分数喂给softmax会发生什么exp(24)大约是2.6e10而exp(0)是1。这意味着几乎所有的概率都堆在了最大值位置上其他位置直接“归零”注意力矩阵退化成近似one-hot矩阵。那层“近似one-hot”后面的梯度非常小。在反向传播时softmax那些被压到接近0的位置梯度几乎传不过去模型权重更新会极其缓慢甚至卡住不动。解决方法是把点积除以√d_k。这样做之后无论d_k多大score的方差都被拉回到1附近softmax输入保持在一个梯度友好的区间。score_scaled (q · k) / √d_k这个√d_k不是一个可以随便省略的超参数——它就是对方差的一个标准化和BatchNorm的均值方差归一化在思路上是一回事。3.2 除了缩放还有没有其他替代方案有人可能会问既然点积受向量长度影响那直接用余弦相似度不就行了吗余弦相似度会先把向量归一化天然让每个score限制在[-1, 1]之间看起来更优雅。实际工程中确实有模型这样试过但有代价。余弦相似度放弃了向量的模长信息而模长本身可能携带语义重要性。某些token的特征比较“强烈”或“微弱”这种强度在点积里会被保留下来在余弦相似度里会被抹平。Transformer最终选择了缩放点积因为它既控制了方差又完整保留了信息。我在调模型时做过对比在同样任务上把√d_k换成固定常数比如“除10”效果都略差。这个缩放因子的数值是有推导依据的不建议随意改动。4. 从公式到代码用PyTorch手写一个单头注意力4.1 先写一个最朴素的版本确保跑得通import torch import torch.nn.functional as F torch.manual_seed(0) batch_size 2 seq_len 6 d_model 16 d_k 16 # 随机输入 x torch.randn(batch_size, 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) # 计算Q、K、V Q torch.matmul(x, W_q) # [batch, seq_len, d_k] K torch.matmul(x, W_k) # [batch, seq_len, d_k] V torch.matmul(x, W_v) # [batch, seq_len, d_v] # 计算注意力分数并缩放 scores torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5) # [batch, seq_len, seq_len] # softmax归一化 attn_weights F.softmax(scores, dim-1) # [batch, seq_len, seq_len] # 加权求和 out torch.matmul(attn_weights, V) # [batch, seq_len, d_v] print(Q shape:, Q.shape) print(K shape:, K.shape) print(scores shape:, scores.shape) print(attn_weights shape:, attn_weights.shape) print(out shape:, out.shape)输出Q shape: torch.Size([2, 6, 16]) K shape: torch.Size([2, 6, 16]) scores shape: torch.Size([2, 6, 6]) attn_weights shape: torch.Size([2, 6, 6]) out shape: torch.Size([2, 6, 16])这个版本里需要注意的关键点计算Q K^T时K必须把最后两个维度转置。K.transpose(-2, -1)把形状从[batch, seq_len, d_k]变成[batch, d_k, seq_len]这样才能和Q做矩阵乘法得到[batch, seq_len, seq_len]的分数矩阵。很多人写代码时卡在这一步报错永远是“维度不匹配”。记住了Q K^T转置的是K不是Q。4.2 用nn.Linear实现可学习投影手工矩阵乘法和真正模型之间还有差距真正的模型需要让W_q、W_k、W_v变成可训练参数并且通常在输入维度上进行线性变换。用nn.Linear可以一步到位import torch.nn as nn class SingleHeadAttention(nn.Module): def __init__(self, d_model, d_k, d_v): super().__init__() self.d_k d_k self.d_v d_v self.W_q nn.Linear(d_model, d_k) self.W_k nn.Linear(d_model, d_k) self.W_v nn.Linear(d_model, d_v) def forward(self, x, maskNone): 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_v] scores torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5) scores scores.masked_fill(mask 0, float(-inf)) if mask is not None else scores attn_weights F.softmax(scores, dim-1) out torch.matmul(attn_weights, V) return out, attn_weights这里有一个容易忽略的细节nn.Linear的权重矩阵在内部存储为[out_features, in_features]而PyTorch在计算时会自动对输入做矩阵乘法原理上等同于输入右乘权重矩阵的转置。也就是说nn.Linear做的是x W.T b我们不需要手动填转置但需要知道它内部确实有这一步。关于biasQKV投影通常是保留bias的。虽然严格来说注意力计算里Q、K被除以√d_k之后bias会造成一定的偏移但实践中影响不大而且保留bias可以让模型有更多拟合空间。省不省bias更多是工程习惯问题。4.3 补上mask训练时最容易被忽略的关键细节实际训练中我们经常需要屏蔽某些位置的注意力。最常见的是两种padding mask句子有填充符虽然长度为6但只有前4个是真实token后2个是填充不能让真实token去关注填充位置因果maskcausal mask在自回归生成模型中当前位置不能看到未来的token只能看自己和过去的位置。实现方式是在softmax之前把需要屏蔽的位置设置为负无穷。为什么这样做而不是对softmax结果置零因为softmax的分母会包含所有位置的指数值如果先把注意力权重算出来再置零分母已经被污染了剩下位置的概率之和就不再是1。只有先把分数设为负无穷再经过softmax那些位置的指数自然为0被softmax自动排除才能保证有效位置的权重之和仍然为1。# padding mask key_padding_mask torch.zeros(batch_size, seq_len) key_padding_mask[0, 4:] 1 # 第0个样本后2位是padding key_padding_mask[1, 5:] 1 # 第1个样本后1位是padding # reshape为可广播的形状[batch, 1, seq_len] mask key_padding_mask.unsqueeze(1) # 因果mask causal_mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool()在forward里通常这样处理if key_padding_mask is not None: scores scores.masked_fill(key_padding_mask.unsqueeze(1) 1, float(-inf)) if causal_mask is not None: scores scores.masked_fill(causal_mask, float(-inf))有一个细节需要注意padding mask应该作用在“被查询的Key”一侧也就是要对score[i, j]中的j位置做屏蔽mask的形状对应Key的长度。而Query侧如果本身是padding可以不管——因为在后续计算loss的时候通常会对这些位置做mask它们的输出不会影响梯度。如果只在Attention模块内部处理只屏蔽Key就已经足够。4.4 用一个小例子完整验证mask行为我们用之前定义过的SingleHeadAttention构造一个含有padding的batch观察输出的注意力权重x torch.randn(batch_size, seq_len, d_model) attn SingleHeadAttention(d_model16, d_k16, d_v16) key_padding_mask torch.zeros(batch_size, seq_len) key_padding_mask[0, 4:] 1 key_padding_mask[1, 5:] 1 out, attn_weights attn(x, key_padding_maskkey_padding_mask) print(attn_weights[0, 0]) # 第0个样本第0个token对所有位置的注意力你会看到第4和第5个位置的注意力权重趋近于0而前4个位置加上第0个token自身位置的概率之和约等于1。这就是mask起效的直观验证。我在自己调试模型时为了防止mask写错位置习惯在每次训练前跑一个极小规模的“冒烟测试”直接打印一个batch的attn_weights检查padding位置的权重是否为0。很多训练loss不下降的奇怪问题最后都查到了mask写反了上面。5. 多头注意力并不是把单头跑h遍维度压缩背后的计算账本5.1 为什么是“拆分”而不是“复制”多个独立注意力标准Transformer里不是只做一个注意力头而是把原来的d_model拆成h个低维子空间每个头独立计算注意力最后拼接回来。常见的配置是d_model512h8那么每个头的维度就是64。为什么不是各自用512维然后得到8个512维的结果再拼起来因为那样计算量会变成8倍模型会大到无法训练。“拆分”的妙处在于总计算量和单头512维几乎一致。8个头的QKV投影每个都是512 - 64加起来正好是一个512 - 512投影的计算量。8个头的注意力分数矩阵也都是[batch, 64, 64]假设seq_len64并行的计算量和单头[batch, 512, 512]的分数矩阵在复杂度上是同一量级的。方案QKV投影计算量注意力分数矩阵大小单头 512 维1 × (512×512)[seq, 512, 512]8头 × 64维8 × (512×64) 同总量8 × [seq, 64, 64]所以多头并不是“多花8倍钱买冗余”而是在同样总参数量下让模型拥有8组不同的投影矩阵。每个头从自己的角度去判断“谁和谁相关”有的头可能聚焦局部相邻token有的头可能关注跨长距离的依赖。5.2 拼接后的输出投影层到底在做什么class MultiHeadAttention(nn.Module): def __init__(self, d_model, h, d_k): super().__init__() self.h h self.d_k d_k self.d_model d_model self.W_q nn.Linear(d_model, h * d_k) self.W_k nn.Linear(d_model, h * d_k) self.W_v nn.Linear(d_model, h * d_k) self.W_o nn.Linear(h * d_k, d_model) def forward(self, x, maskNone): batch, seq_len, _ x.shape Q self.W_q(x).view(batch, seq_len, self.h, self.d_k).transpose(1, 2) K self.W_k(x).view(batch, seq_len, self.h, self.d_k).transpose(1, 2) V self.W_v(x).view(batch, seq_len, self.h, self.d_k).transpose(1, 2) scores torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5) scores scores.masked_fill(mask 0, float(-inf)) if mask is not None else scores attn_weights F.softmax(scores, dim-1) attn_out torch.matmul(attn_weights, V) # [batch, h, seq_len, d_k] attn_out attn_out.transpose(1, 2).contiguous().view(batch, seq_len, self.h * self.d_k) out self.W_o(attn_out) # [batch, seq_len, d_model] return out, attn_weights每个头算完之后把[batch, h, seq_len, d_k]重新拆回[batch, seq_len, h * d_k]再经过一个输出投影W_o变回d_model维度。这个投影的作用是把多个头的子空间信息混合回统一的表示空间。这里还有一个容易踩的坑从[batch, h, seq_len, d_k]转回[batch, seq_len, h*d_k]时必须先把第1和第2维转置再做view。如果你不先转置直接view会把不同头的维度顺序搞乱。PyTorch的张量内存连续性问题也会出来捣乱——所以还要记得加.contiguous()。5.3 一个常见的误解多头真的会自动分化吗大多数人读到“多头注意力可以让不同头关注不同方面”时会默认这是一个自动形成的性质。实测下来这句话只能算是“设计意图”而不是“必然结果”。没有任何loss函数明确约束每个头必须学得不一样。模型只需要最终的预测能对不关心头之间是否有差异。所以经常出现的情况是8个头里面有4个学到几乎一模一样的模式剩下几个也高度相似。这就是所谓“头同质化”现象。我在调自己的模型时会额外写一个调试脚本把每一层的attn_weights存下来统计每个头注意力分布的平均熵然后画出来看。如果发现大部分头的熵值都差不多、注意力模式高度重合就该考虑是否真的需要这么多头或者是否应该调整初始化方式。6. 从论文到工程实际调试中经常踩的坑6.1 维度错误为什么有些人用matmul成功有些人用bmm报错注意力计算里最常见的错误就是矩阵相乘维度写错。PyTorch里torch.matmul和torch.bmm都能做矩阵乘法但行为略有不同。torch.bmm要求两个输入都是3D张量它只会把最后一维和倒数第二维做乘法batch维保持不变。比如Q是[batch, seq, d_k]K^T是[batch, d_k, seq]bmm没问题。但很多人会把transpose写反# 错误示范 scores torch.bmm(K, Q.transpose(-2, -1)) # 顺序反了得到的形状完全不对另一种情况是输入不小心多了一维或少了一维bmm直接报错而matmul会自动做广播结果可能在意外处成功但语义是错的。这一点很迷惑人。我的建议很简单在关键维度变换之后立刻加一行print或断言不要靠肉眼判断。维度问题占了注意力实现80%的调试时间提前断言能省下一大片头发。assert scores.shape (batch_size, seq_len, seq_len), fscores shape mismatch: {scores.shape}6.2 训练初期loss不下降先检查缩放和初始化如果代码整体能跑但loss死活不降优先级最高的怀疑对象有两个一个是没做√d_k缩放另一个是初始化参数范围不合适。没做缩放的情况下score的方差很大softmax输出接近one-hot梯度消失模型几乎学不进去。这种情况在训练初期尤其明显loss可能从第一步开始就不动。初始化方面如果W_q、W_k的参数初始值过大即使做了缩放也会让logits分布偏移。实践里常见的做法是让Q和K投影使用更小的初始化标准差比如1/√d_k的量级。这个细节在训练深层Transformer时效果差别很明显。如果发现训练曲线非常平可以先用一段随机数据做“过拟合测试”。模型理论上能逐渐记住这几十条随机数据loss应该下降并趋向0。如果连随机数据都学不动那问题大概率出在Attention实现本身。6.3 attention dropout的位置为什么加在权重上而不是logits上在训练过程中通常还会对注意力权重做dropout防止co-adaptation。容易被忽略的是dropout应该加在softmax之后也就是attn_weights上而不是加在scores上。原因很简单dropout是随机置零如果把dropout加在scores上等于随机把某些分数变成0然后再进softmax分母会变化概率分布会被整体扭曲。而加在softmax之后的attn_weights上随机置零某些位置的权重恰好是“随机忽略某些连接”的效果语义更干净。实现就一行attn_weights F.dropout(attn_weights, p0.1, trainingself.training)注意推理时trainingFalsedropout会被自动关闭模型行为变成确定性。提示给attention加dropout时建议把p控制得很小。0.1是比较常见的取值p超过0.2时我发现注意力模式容易被扰动收敛会明显变慢。写在最后的一点经验我踩过很多次坑以后才确认一件事注意力机制里真正重要的不是记住公式而是理解每一行代码改动对应数学上的什么变化。比如缩放因子表面看只是除一个数实际上是在控制softmax输入分布的方差mask加在softmax前还是后直接决定概率分布是否归一化多头拆分的维度设计决定了一个模型的参数量和计算效率。如果你正在调试自己的注意力实现我建议先跑通单头版本验证维度和mask行为都正确再加多头和dropout。一步吃透再走下一步比一口气写完一堆代码再回头debug要快得多。这篇是第一阶段的上半部分下一阶段会继续把多头注意力的完整工程细节、位置编码的数学动机和不同mask组合的使用场景展开。