★ FEATURED ARTICLE 大家都在讨论 Transformer、大模型、注意力机制但很多朋友一上来就被 QKV 这几个字母劝退了。其实自注意力没那么玄它干的事情一句话就能说清楚让句子里的每个词根据当前语境重新计算自己的表示计算方式就是所有词两两之间先打个分再按分数融合彼此的信息。说白了就是让句子里的词互相“看”对方看完一圈之后每个词都带上了上下文的味道。这篇文章不堆公式吓人我会从直觉讲起再手把手拆解 Q、K、V 到底在做什么最后用 PyTorch 从零写一个能跑的自注意力层顺便把我在实操中踩过的一些坑和排查心得一并分享出来。适合刚接触 NLP、想搞懂 Transformer 底层原理的同学也适合那些能跑通代码但总觉得“差点意思”的开发者。1. 为什么句子里的词需要“互相看”——从 N-Gram 到注意力机制1.1 RNN 和 N-Gram 解决不了什么问题先来想想一个很朴素的问题一个词的意思是它自己决定的吗不是。绝大多数词的意思取决于上下文。“苹果”在“苹果很好吃”里是一种水果在“苹果发布了新手机”里是一家公司“打”在“打篮球”“打酱油”“打电话”里含义完全不同。你光看这个字本身永远猜不到它到底想表达什么。传统的 N-Gram 做法很简单滑动窗口看前面 N 个词。比如只看前面两个词那“我昨天带着小狗去公园它玩得很开心”这句话里“它”前面的两个词是“公园”和“”模型根本不知道“它”指的是狗。窗口再拉大一点N5能覆盖到“带着小狗”但句子再长一点指代信息隔了十几二十个词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_kV 用 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 是下面这样的值词QKV猫[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.786exp(0.30)1.350exp(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_kNone, d_vNone): 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, maskNone): # 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_model16) 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 个词的强度”。你可以打印出来看看每一行加起来应该都等于 1softmax 归一化。3.3 多头注意力与 Mask 的正确打开方式单头自注意力只能学一种“关注模式”但句子中存在的依赖关系显然不止一种。拿“小明说我昨天去了银行”这句话举例我们需要同时捕捉“小明”和“我”的人称一致性也需要捕捉“银行”和“去了”的语义搭配还要捕捉“昨天”这个时间状语修饰的是哪个动词。单头注意力只能学一种匹配模式多头则可以让不同的头各自关注不同的关系类型。多头实现的关键代码我在下面贴出来。需要注意的是在实际实现中通常把 d_model 切分成 num_heads 份每个头分到 d_k d_model // num_heads 的维度让每个头在自己的子空间里做自注意力。这样做还有一个隐性的好处多个头是并行计算的效率比串行高得多。class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads8): 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, maskNone): # 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, 2i1) 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 maskTrue 表示屏蔽搞混这两套语义很容易打架。4.3 训练不稳定的排查清单自注意力层训练不稳定背后往往不是单一原因而是一串因素叠加。我按优先级整理了一份排查清单学习率Transformer 系列对学习率极其敏感。常见做法是先用 Adam warmup把学习率从 0 线性升至预设峰值。峰值太大表现为 loss 前期暴增峰值太小表现为收敛缓慢。如果 loss 曲线像心电图一样上下抖先减学习率再查归一化层。初始化有时不是模型的问题而是 Linear 权重的初始范围不合适。PyTorch 默认的初始化方式一般没问题但如果你自己实现了带缩放的结构记得手动检查各层输出的方差量级。输出方差太大或太小都会导致训练不稳。LayerNorm 的位置是 Pre-Norm先归一化再进入注意力还是 Post-Norm先注意力再归一化对训练稳定性有很大影响。对于深层 TransformerPre-Norm 往往更稳定因为残差路径里少了一个有风险的变换。想要快速让模型训起来优先考虑 Pre-Norm。梯度裁剪即使其他都正常长序列 深模型仍然容易偶发梯度爆炸加个梯度裁剪比如设置 max_norm1.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 一个卷积核扫过去速度比自注意力快几个数量级效果也不差。工具是死的人和场景才是活的不要为了炫技而牺牲工程合理性。最后再分享一点我的实际体会自注意力这个东西你光看公式会觉得就那么几行但真正想把它用扎实必须亲手写一遍、亲手调一遍、亲手画一画注意力热力图。我建议你有空的时候把 GPT-2 的源码或者某些开源实现里的注意力模块拿出来一行一行对照着自己写的版本找差异这一步能帮你补上很多论文里没细讲的工程细节收益很大。如果你刚开始接触这块先别急着一次实现 12 层。把一个单层自注意力彻底吃透搞明白 QKV 和 mask 的维度流转后面看再大的模型都是一层层的重复堆叠底子就落在今天讲的内容上。
阅读完成 · 觉得有帮助？
