自注意力机制详解：从QKV原理到PyTorch手写实现与避坑指南-育师

大家都在讨论 Transformer、大模型、注意力机制，但很多朋友一上来就被 QKV 这几个字母劝退了。其实自注意力没那么玄，它干的事情一句话就能说清楚：让句子里的每个词，根据当前语境重新计算自己的表示，计算方式就是所有词两两之间先打个分，再按分数融合彼此的信息。说白了，就是让句子里的词互相“看”对方，看完一圈之后，每个词都带上了上下文的味道。

这篇文章不堆公式吓人，我会从直觉讲起，再手把手拆解 Q、K、V 到底在做什么，最后用 PyTorch 从零写一个能跑的自注意力层，顺便把我在实操中踩过的一些坑和排查心得一并分享出来。适合刚接触 NLP、想搞懂 Transformer 底层原理的同学，也适合那些能跑通代码但总觉得“差点意思”的开发者。

1. 为什么句子里的词需要“互相看”？——从 N-Gram 到注意力机制

1.1 RNN 和 N-Gram 解决不了什么问题

先来想想一个很朴素的问题：一个词的意思，是它自己决定的吗？

不是。绝大多数词的意思取决于上下文。“苹果”在“苹果很好吃”里是一种水果，在“苹果发布了新手机”里是一家公司；“打”在“打篮球”“打酱油”“打电话”里含义完全不同。你光看这个字本身，永远猜不到它到底想表达什么。

传统的 N-Gram 做法很简单：滑动窗口，看前面 N 个词。比如只看前面两个词，那“我昨天带着小狗去公园，它玩得很开心”这句话里，“它”前面的两个词是“公园”和“，”，模型根本不知道“它”指的是狗。窗口再拉大一点，N=5，能覆盖到“带着小狗”，但句子再长一点，指代信息隔了十几二十个词，N-Gram 就彻底没辙了。

RNN 的思路则是按顺序一个个处理词，用一个隐状态把“历史信息”往后传。理论上它能记住长距离依赖，但实际训练中梯度消失问题让 RNN 很难把十几步之前的信息稳稳传过来——不是不努力，是数学上就吃亏。而且 RNN 只能从左往右或从右往左扫，一个词要“看到”远处的另一个词，必须沿着时间步一步步走过去，这个路径太长了。

这时候我们需要一种更直接的方式：别管词在句子里的物理距离有多远，只要它们在语义上有关系，就应该能够直接交换信息。

1.2 注意力的本质是个“加权查询”

注意力机制最早是在机器翻译里火起来的。当时的经典套路是编码器把整个源语言句子压成一个固定长度的向量，解码器再从这个向量里解码出目标语言。这有个很明显的问题：句子太长时，一个向量塞不下所有信息，早期的信息会被“挤”丢。

所以有人提出了注意力：解码器每生成一个词，不是只看那个压缩后的全局向量，而是回头去看源语言句子的每一个词，然后决定“我现在最应该关注哪几个词”。这个过程的本质，就是带着问题去原文里做检索。

你可以把这个过程类比成在图书馆找资料：你手里拿着一张借书单（Query），书架上每本书都有标签（Key），你逐个看标签，选出和借书单最匹配的那几本，最后把书的内容（Value）记下来。注意力做的事情一模一样——Query 是当前词的意图，Key 是其他词的标签，Value 是其他词真正想表达的内容。

1.3 自注意力是“全员参与”的讨论会

理解了上面的“注意力”之后，“自注意力”只需要加一个字：“自”。它不再区分谁是源头、谁是目标，而是让句子里的每一个词同时扮演 Query、Key、Value 三个角色，大家平等地互相看一遍。

我特别喜欢用一个会议来类比：开会时每个人先陈述自己的观点（Value），同时每个人心里带着自己的问题（Query），听别人说话时快速判断对方的发言和自己的问题有没有关系（Key），有关系就认真听，没关系就一带而过。一轮讨论下来，每个人吸收了自己关心的重要信息，对议题也有了更全面的理解。自注意力就是这么一轮“全员参与的讨论会”，而且这个讨论会可以反复开很多层——每一层就是对上一轮的结果再讨论一次。

这个设计的妙处在于，它在建模长距离依赖时不再依赖“逐步传递”，而是建立了一条全连接的直通路径。“它”和“小狗”哪怕隔了五十个词，在第一层自注意力里也可以直接建立高权重连接。

2. 自注意力的核心细节拆解：Q、K、V 到底在做什么

2.1 三个向量是怎么来的

现在来看数学。假设输入是一串 token 的向量表示，通常记作 X，形状是 [序列长度, 词向量维度]。

所谓 Q、K、V，并不是什么从天而降的神奇变量，它们就是 X 分别乘上三个可学习的权重矩阵：

Q = X @ W_q K = X @ W_k V = X @ W_v

也就是说，每个输入向量经过三个不同的线性变换，被投影到三个不同的空间。为什么不能直接用 X 自己去做匹配？我的理解是：如果直接用原始向量既当“问题”又当“答案”，表达能力会比较受限。经过三个可学习的投影之后，模型就能学出“应该用什么标准去匹配别人”“应该用哪些特征去被别人匹配”“实际要传递什么内容”这三套不同的刻画方式，各司其职。

有人会问：Q 和 K 的维度一定等于 V 吗？不必须。经典 Transformer 里 Q、K 用的是同一个维度 d_k，V 用 d_v，两者可以不同。但为了方便和降低实现复杂度，很多实现里干脆让它们相等，这不算偷懒，是一种合理的工程选择。

要理解 Q、K、V 的关系，我再换一个更贴近生活的类比：相亲。你在聊天时向对方提问（Q），对方根据你的问题在脑海里检索自己符合的特质（K），如果匹配度高（点积得分高），那 Ta 接下来讲的实质性内容（V）你就听进去了。匹配度不高，Ta 就算说了很多，你也不太在意。

2.2 注意力分数是怎么算出来的

有了 Q、K、V，下一步是计算两个词之间的“相关程度”，也就是注意力分数。

经典公式长这样：

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

我先只看中间那个分数：一句话里第 i 个词和第 j 个词的分数，就是 Q_i 点乘 K_j。点积在几何上可以理解为一种相似度——如果两个向量方向越一致，点积越大，说明这两个词在当前的表达空间中越“有话说”。

关键是为什么要除以 sqrt(d_k)。

这要从 softmax 的性质说起。softmax 的输入如果数值太大，输出会变得非常“尖锐”，最大的那个值几乎抢走全部权重，其他值趋近于零。而 Q 和 K 里的每个元素如果是均值 0、方差 1 的随机变量，它们的点积（也就是对所有元素乘积求和）方差会变成 d_k。也就是说，d_k 越大，点积分数的绝对值越大，softmax 越容易饱和，梯度越小，训练就越难。

除以 sqrt(d_k) 之后，分数的方差被拉回到 1 左右，softmax 的输出不会太极端，梯度也就能顺畅地回流。这是一个看起来不起眼、实则影响训练稳定性的关键缩放。

2.3 加权求和：每个词的“新身份”

分数算完之后，对每一行做一次 softmax 归一化，得到的是这一行的概率分布，代表“当前这个词，应该以多大比重关注句子里的其他词”。

最后一步是把这些权重和 V 做加权求和：

输出_i = 所有词的权重_ij * V_j

这里有个容易被忽略但很重要的特征：权重是 softmax 归一化之后的，都是非负数且加起来等于 1。所以输出其实是 V_j 的一个凸组合，也就是说，每个词的新表示不会跑出所有 V 向量张成的范围，本质上是在做一次“有选择性的信息搬运”。哪个词和当前词关系越密切，它的 V 在被搬运时占的比重就越大。

我拿三个词的实际数字走一遍完整流程，感受会更直观。假设句子是“猫吃鱼”，为了方便手算，把词向量维度设为 2（实际场景往往是 64 到 1024，原理一样）。假设经过投影后，三个词的 Q、K、V 是下面这样的值：

词	Q	K	V
猫	[1.0, 0.2]	[0.8, 0.1]	[0.5, -0.2]
吃	[0.5, 0.8]	[0.3, 0.6]	[0.2, 0.7]
鱼	[0.1, 0.5]	[0.9, 0.4]	[-0.4, 0.3]

先算“猫”对三个词的注意力分数：

猫对猫：1.0 × 0.8 + 0.2 × 0.1 = 0.82
猫对吃：1.0 × 0.3 + 0.2 × 0.6 = 0.42
猫对鱼：1.0 × 0.9 + 0.2 × 0.4 = 0.98

此时看“猫”这个位置的输出，它会如何分配自己对句子中三个词的注意力权重？先把每个分数除以 sqrt(2)≈1.414，得到 0.58、0.30、0.69，然后 softmax 归一化。我算一下近似值：exp(0.58)=1.786，exp(0.30)=1.350，exp(0.69)=1.994，总分母 5.130，所以权重约为 0.35、0.26、0.39。

你看，即使是简化到这个地步，“猫”对“鱼”的关注度依然是最高的。而在真实模型里，经过训练之后，“猫”对“鱼”这类语义相关词往往会学到更高的权重，因为“猫”和“鱼”在类似“猫吃鱼”的语境里有强关联。最后“猫”的新向量就是 0.35 × V_猫 + 0.26 × V_吃 + 0.39 × V_鱼，相当于它吸收了“鱼”的一部分语义信息，给自己叠了一层上下文 Buff。

这就是自注意力最直观的运作方式。每个词的新表示，是由整个句子所有词的加权求和得到的，而那个权重，是它自己决定的。

3. 从零实现一个自注意力层

3.1 环境准备与最小可跑结构

这一节我们直接动手写代码。我的建议是：先跑通最简版本，再逐步加多头、加 Mask、加位置编码，每一步都搞清楚为什么要加。

环境方面，只需要 PyTorch 就行，版本 1.10 以上都可以。下面是我的完整实现。

3.2 最简版自注意力实现

先写一个单头自注意力模块。我把每一步的维度形状都写在注释里，方便你对照调试。

import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, d_model, d_k=None, d_v=None): super().__init__() self.d_model = d_model self.d_k = d_k if d_k is not None else d_model self.d_v = d_v if d_v 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_v) def forward(self, x, mask=None): # x 形状: [batch, seq_len, d_model] 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_v] # 计算注意力分数并缩放 scores = torch.matmul(Q, K.transpose(-2, -1)) # [batch, seq_len, seq_len] scores = scores / (self.d_k ** 0.5) # 如果传了 mask，把被遮蔽位置的分数设为负无穷 if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) attn_weights = F.softmax(scores, dim=-1) # [batch, seq_len, seq_len] out = torch.matmul(attn_weights, V) # [batch, seq_len, d_v] return out, attn_weights

这段代码只有十几行，但自注意力的全部核心逻辑都在里面了。我建议你把它复制到本地，用下面的测试代码跑一下：

x = torch.randn(2, 5, 16) # 2个句子，每个5个词，向量维度16 model = SelfAttention(d_model=16) out, attn = model(x) print(out.shape) # torch.Size([2, 5, 16]) print(attn.shape) # torch.Size([2, 5, 5])

跑完之后你会看到一个形状为 [2, 5, 5] 的注意力矩阵，这个矩阵第 i 行第 j 列的值就是“第 i 个词关注第 j 个词的强度”。你可以打印出来看看，每一行加起来应该都等于 1（softmax 归一化）。

3.3 多头注意力与 Mask 的正确打开方式

单头自注意力只能学一种“关注模式”，但句子中存在的依赖关系显然不止一种。拿“小明说：我昨天去了银行”这句话举例：我们需要同时捕捉“小明”和“我”的人称一致性，也需要捕捉“银行”和“去了”的语义搭配，还要捕捉“昨天”这个时间状语修饰的是哪个动词。单头注意力只能学一种匹配模式，多头则可以让不同的头各自关注不同的关系类型。

多头实现的关键代码我在下面贴出来。需要注意的是：在实际实现中，通常把 d_model 切分成 num_heads 份，每个头分到 d_k = d_model // num_heads 的维度，让每个头在自己的子空间里做自注意力。这样做还有一个隐性的好处：多个头是并行计算的，效率比串行高得多。

class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads=8): super().__init__() assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除" self.num_heads = num_heads self.d_model = d_model self.d_k = d_model // num_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.out_proj = nn.Linear(d_model, d_model) def forward(self, x, mask=None): # x: [batch, seq_len, d_model] batch, seq_len, _ = x.shape Q = self.W_q(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) K = self.W_k(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) V = self.W_v(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 此时形状都是 [batch, num_heads, seq_len, d_k] scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5) # [batch, num_heads, seq_len, seq_len] if mask is not None: # mask 需要扩到和 scores 一样的维度 # mask 原始形状: [batch, 1, 1, seq_len] 或 [batch, 1, seq_len, seq_len] scores = scores.masked_fill(mask == 0, float('-inf')) attn_weights = F.softmax(scores, dim=-1) # 加权求和 out = torch.matmul(attn_weights, V) # [batch, num_heads, seq_len, d_k] # 把头合并回去 out = out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model) out = self.out_proj(out) # 最后再接一层输出投影 return out, attn_weights

在实际 NLP 任务里，Mask 是必须严肃对待的东西，有两种最典型：

第一种是padding mask。一个 batch 里句子长度不一样，我们会把短的句子用特殊 token 补齐到一样长。这些填充位置不携带任何语义信息，要是允许它们参与注意力计算，等于往上下文里掺了一堆噪音。所以要把这些位置的注意力分数全部设为负无穷，softmax 之后权重就会变成 0。

第二种是causal mask，也就是因果掩码，主要用于自回归式的语言模型。它保证在预测第 t 个词的时候，模型只能看到前 t-1 个词，不能偷看未来的词。实现方法是构造一个上三角矩阵，把上三角部分（包括对角线右上方）全部遮掉。这样训练和推理的时候行为一致，否则模型会在训练时通过“偷看答案”把损失降到很低，一推理就露馅。

关于 Mask 的坑，后面第 4 节我会专门展开讲。

3.4 位置信息去哪了

自注意力对位置的感知是天然缺失的。你可以把整个序列理解成一个“袋子”——把词的顺序打乱，注意力分数的计算结果完全不变，因为 Q 和 K 的点积只看内容关系，不看相对位置。

这显然不行。语言是强烈依赖顺序的：“我打你”和“你打我”，词完全一样，意思完全相反，但自注意力会把这两个句子当成同一个表示。所以必须显式地把位置信息塞进去。

常见做法是在输入向量上加一个位置编码向量，经典实现是 sin 和 cos 交替：

PE(pos, 2i) = sin(pos / 10000^(2i / d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

用周期函数的原因在于，不同频率的 sin/cos 组合能产生一种“编码指纹”，让模型通过比较位置编码之间的差值，间接感知到相对位置。另一个更现代的做法是让位置向量变成可学习的参数，直接当成 embedding 训练。两种方案各有优势：sinusoidal 版本长度可外推，可学习版本在长文本场景可能要专门处理外推问题。

值得注意的是，位置编码是加在输入 Embedding 上的而不是拼接到向量后面，这样既改进了输入表达，又没有让向量维度变大一倍，节省了参数和显存。

4. 实操中的坑与排查心得

4.1 不除以 sqrt(d_k) 会怎样

我刚开始自己实现注意力时，觉得除以 sqrt(d_k) 这种小细节没啥大不了，直接省略。结果训练一个中小规模的翻译模型，loss 降得非常慢，甚至有时候直接发散。

原因是缩放因子对梯度健康度的影响超出直觉。刚才已经分析过：不缩放时，点积的方差随 d_k 增大而增大，softmax 在分数大时梯度近乎为零，参数几乎学不动。你以为模型在训练，实际上它在原地徘徊。

这里给一个非常直观的验证方法。你把 d_k 设成 128，然后在没有任何缩放的情况下，随机初始化模型跑一个 batch，打印一下注意力权重的分布。你会发现权重很快变成 one-hot 形状，即某一列接近 1，其他列接近 0。这不是模型学得好，而是 softmax 饱和了。除以 sqrt(d_k) 之后再看，权重分布就稀疏但不极端，这才留出了学习空间。

4.2 Mask 写错的几种典型表现

Mask 是新手最容易写错的地方，我自己就犯过不止一次。

我总结一下最常见的三种错误表现，你如果遇到类似现象可以快速对号入座：

训练 loss 极低但推理质量很差：概率最高的原因是因果 mask 写错了，模型看到了未来的词。因为训练时模型“作弊”，所以 loss 下降得很顺利，一点也不像训练有问题，骗了你整个训练周期。排查方法也很简单：把推理时随机改一两个预测词，观察后续生成是否出现一致性崩坏。
注意力权重里有整行为 NaN：通常是某个 attention 行被 mask 成了全负无穷，softmax 对全负无穷的输入求出来的结果是 NaN。这种情况一般出现在 mask 的形状不对，比如把 padding mask 错误地覆盖到了所有序列的某个位置上。检查 mask 生成逻辑里是否有维度层面的错位。
每个 batch 的注意力矩阵几乎不变：大概率是 mask 被广播错了空间维度，或干脆没有 mask，模型就把所有位置都当作有效信息，学到的是同质的注意力。你可以把 mask 打印出来看一眼，刚开始做实验时，打印 shape 和几行数值，能省掉很多隐性调试时间。

Mask 在多头实现里的形状正确写法是：对于 padding mask，通常保持 [batch, 1, 1, seq_len]，这样在广播到 [batch, num_heads, seq_len, seq_len] 时不会有歧义。而 causal mask 是 [batch, 1, seq_len, seq_len]，因为每个句子自己的上三角位置都不同，不能再简单压缩成一个向量。用 masked_fill 之前，确认你的 mask 里 1 表示“保留”、0 表示“屏蔽”，不要和常见的 bool mask（True 表示屏蔽）搞混，这两套语义很容易打架。

4.3 训练不稳定的排查清单

自注意力层训练不稳定，背后往往不是单一原因，而是一串因素叠加。我按优先级整理了一份排查清单：

学习率：Transformer 系列对学习率极其敏感。常见做法是先用 Adam + warmup，把学习率从 0 线性升至预设峰值。峰值太大表现为 loss 前期暴增，峰值太小表现为收敛缓慢。如果 loss 曲线像心电图一样上下抖，先减学习率，再查归一化层。
初始化：有时不是模型的问题，而是 Linear 权重的初始范围不合适。PyTorch 默认的初始化方式一般没问题，但如果你自己实现了带缩放的结构，记得手动检查各层输出的方差量级。输出方差太大或太小，都会导致训练不稳。
LayerNorm 的位置：是 Pre-Norm（先归一化再进入注意力）还是 Post-Norm（先注意力再归一化），对训练稳定性有很大影响。对于深层 Transformer，Pre-Norm 往往更稳定，因为残差路径里少了一个有风险的变换。想要快速让模型训起来，优先考虑 Pre-Norm。
梯度裁剪：即使其他都正常，长序列 + 深模型仍然容易偶发梯度爆炸，加个梯度裁剪（比如设置 max_norm=1.0）能有效兜底。

4.4 注意力可视化：验证模型到底学到了什么

写代码跑通是最低要求，你肯定还希望能确认模型真的学到了有意义的东西。我的习惯是把注意力矩阵画成热力图，一眼就能看出关系模式。

越低的层级往往学到的是相邻词的局部关系——这很好理解，语法上的相邻依赖本来就多。越高的层级越容易出现长距离的、语义性的关联，比如代词和它指代的名词之间会出现一条明显的亮带。如果你观察到所有头在高层都还是纯局部注意力，那模型大概率没学到深层语义结构，这时候就该检查数据量、训练步数或者模型容量了。

做可视化用的就是每个 head 的 attn_weights，形状是 [batch, num_heads, seq_len, seq_len]，用 matplotlib 的 imshow 画出来即可。一个小技巧是：不要只画一个头的，画多个头放一起对比，方便看出多头之间是不是真的分化了。如果八个头画出来几乎一样，那说明多头并没有各司其职，模型容量存在浪费。

4.5 什么时候不值得用自注意力

写这篇文章绝不是劝你什么任务都上自注意力。自注意力的计算复杂度是 O(n²) 的，n 是序列长度。序列长度 1000 时，计算量是 100 万数量级的点积；序列长度 10000 时，直接到 1 亿。这不是线性增长，是二次方的爆炸。

如果任务本身是短文本分类、文本匹配，几百个词的序列长度，用标准的自注意力完全没问题；但如果你的输入是万词级别的长文档或者超长时序数据，就应该考虑更快的手段——比如滑动窗口注意力、局部注意力、稀疏注意力，或者直接先用 CNN / LSTM 提取局部特征再做全局交互。有些场景里，一个轻量级 CNN 一个卷积核扫过去，速度比自注意力快几个数量级，效果也不差。工具是死的，人和场景才是活的，不要为了炫技而牺牲工程合理性。