自注意力机制（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维度是64，V的维度也设为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)，而不是dim=1或dim=2。如果写错了，注意力矩阵的行和列语义就颠倒，模型轻则不收敛，重则直接输出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, bias=False) W_k = nn.Linear(d_model, d_k, bias=False) W_v = nn.Linear(d_model, d_v, bias=False)

三个线性层的输入都是d_model，输出对应各自的投影维度。bias在这里我通常设成False，这背后是有讲究的：注意力机制中Q和K的匹配是纯点积运算，引入偏置相当于给点积的每一维增加偏移量，对这个相似度匹配的建模能力提升有限，反而会在模型量化部署时多出一组参数要处理。很多开源实现的默认做法也是不加偏置。

你可能会问，三个权重矩阵为什么不能共享？直觉上，如果共享权重，“我要找什么”和“我能提供什么”就无法分离，模型表达力会明显下降。所以三个投影必须是独立参数，才能让同一个位置在不同语义空间下呈现不同的角色。

3.3 一次性计算Q、K、V

Q = 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]], dtype=torch.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会自动把这些位置的权重算成0，padding 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_k=None, d_v=None, dropout=0.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, bias=False) self.W_k = nn.Linear(d_model, d_k, bias=False) self.W_v = nn.Linear(d_model, d_v, bias=False) self.W_o = nn.Linear(d_v, d_model, bias=False) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): 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_model=16, dropout=0.0) official_attn = nn.MultiheadAttention(embed_dim=16, num_heads=1, dropout=0.0, batch_first=True) 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], dim=0) ) 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_weights=False) print("最大差异:", (custom_out - official_out).abs().max().item())

如果代码正确，最大差异应该在1e-6以下，这是浮点数精度带来的正常偏差。如果差异很大，优先查三件事：Mask有没有传、Softmax维度是不是-1、scale有没有漏乘。这个对照实验非常值得做一次，它能彻底打消“我写的是不是官方那个注意力”的疑虑。

顺带一提，官方nn.MultiheadAttention从PyTorch 2.0开始，内部对need_weights=False的路径走了更高效的融合算子，数值上跟手写版本会有微小偏差，但并不会影响逻辑一致性。

4. 维度流向与常见工程陷阱

4.1 从输入到输出的维度变化全梳理

把关键张量的形状整理成一张表，代入本文设定的数值：batch=2、seq=4、d_model=16、d_k=d_v=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构造输入，点积结果全为0，Softmax后应该得到均匀概率分布，每个位置概率相同。一旦这个基础测试不通过，马上就能暴露出维度层面的问题，而不需要等训练跑很久才发现。

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_first=False，输入要求是[seq, batch, dim]，如果你用的是[batch, seq, dim]且没有指定batch_first=True，数据进入就会错位，这个问题极其隐蔽。第二，官方QKV投影是同一个权重矩阵拼接而成，还带偏置，即便你的实现不加偏置，也要同步清零in_proj_bias和out_proj.bias才能做到完全对齐。第三，PyTorch 2.0以后，设置need_weights=False时内部走的是融合算子路径，输出数值上与手写会有微小但可接受的偏差，这属于正常现象。

5.5 工程实践中的一点体会

如果你只是搭实验或者跑固定框架，完全没必要手写Self-Attention，直接调用内置模块更省心也更高效。但如果你要做的是魔改注意力，比如定制Mask策略、实现稀疏模式、调整数值绑定规则，手写版本的灵活优势就体现出来了。我的经验是先完整手写一遍单头自注意力，彻底搞懂每一步的维度和数值变化，再去读多头和融合实现，效率和深度都比直接开箱即用要强得多。

最后分享一个我一直保留的调试技巧：把注意力权重矩阵打出来，配合特定句子观察模型关注了哪些位置。如果模型回答某个问题时，注意力集中在“证据词”附近，说明它在看对地方；如果权重一片均匀，多半不是代码问题，而是训练信号出了问题。这种可视化排障方法，比对着loss曲线猜原因快很多。