搞懂Transformer的人很多,但真能不看资料、把自注意力的来龙去脉讲明白的人,说实话不多。大部分人停留在"Q、K、V三个矩阵,算一下相似度,加权求和"这个层面,再往下问一句"为什么点积结果要除以√d_k",就卡住了。这篇文章我想把这层窗户纸捅破,从自注意力机制在整个Transformer中的位置、QKV矩阵的由来、注意力分数的计算细节、多头注意力的设计动机,一路聊到位置编码和工程实战里那些容易忽略的坑。适合已经跑过Transformer代码、想真正理解底层原理的读者,也适合准备面试、需要系统性梳理这部分知识的人。
1. 先搞清楚自注意力在整个Transformer里的位置
1.1 那张经典结构图的右侧其实是另一个故事
很多人在初学Transformer时,盯着那张经典结构图看半天,左边的Encoder堆着Masked Multi-Head Attention和Feed Forward,右边的Decoder又有一层Cross Attention,很容易把注意力机制的来龙去脉搞混。我第一次看的时候也是一脸懵,后来才意识到一个关键点:Encoder里的自注意力是"自己跟自己玩",Decoder里的Cross Attention是"跟编码器的输出玩",Masked Self-Attention则是"自己跟自己玩但戴着面具",三种注意力虽然核心运算一模一样,但输入来源和作用对象有本质区别。
理解自注意力机制,我建议先从Encoder的Self-Attention入手。它最纯粹,也最容易建立直觉。它的任务一句话概括就是:让序列里的每一个token,都能根据其他所有token的信息,来更新自己的表示。这个"更新"不是简单拼接,也不是取平均,而是学到"我应该重点关注谁"。
这个思路本身就颠覆了RNN的串行理念。RNN处理序列靠一个隐藏状态逐步传递,想获取远处信息就得一步步传过来,路径很长。自注意力则是一步到位,任意两个位置的token直接计算相关性,不管它们在序列里隔多远,信息传递路径长度始终是1。这也是Transformer能处理长距离依赖的根本原因。
1.2 自注意力要解决的是"一句话里的词怎么互相打招呼"
用一个生活化的类比来建立直觉。想象你在一个热闹的会议室里,你的发言内容(一个词的含义),不应该只是你自己的发言稿决定的,还应该受到会议室里其他人反应的影响。某个人点头赞同,你的发言里关于那部分内容的分量就该加重;某个人摇头,那部分就该适当地被忽略。自注意力做的就是这个事情——序列里每个词,都"张望"一圈序列里所有其他词,判断哪个词跟自己关系大,哪个词关系小,然后按照亲疏程度把信息汇聚到自己的表示里。
这个"亲疏程度"不是预先定死的,而是从数据里学出来的。两个词之间该有多亲近,取决于它们的"意图"和"特征"是否匹配。怎么判断匹配?把每个词变成两个不同维度的向量——一个代表"我想找什么"(Query),一个代表"我能提供什么"(Key),然后算Query和Key的匹配度。匹配度高的两个词,信息融合得就多。
这个设计后来被很多研究者展开过,衍生出了各种各样的注意力变体,但万变不离其宗,核心还是那套query匹配key、value带信息汇聚的逻辑。搞懂这一层的设计意图,后面看Flash Attention也好、稀疏注意力也好,都不会被绕晕。
2. Q、K、V矩阵:输入是怎么拆成三个视角的
2.1 从词向量到Query、Key、Value的投影过程
现在进入最关键的实操部分。假设输入是一句话,经过Embedding之后得到了一组向量,每个token对应一个d_model维的向量,比如d_model=512。这个向量带着token的语义信息,但还不够。在自注意力里,每个token需要同时扮演三个角色:它要找别人(Query)、它要被别人找到(Key)、它被找到之后要提供的内容(Value)。
怎么从同一个输入向量得到三个角色?靠三个可学习的权重矩阵W_Q、W_K、W_V。输入向量X分别跟这三个矩阵做矩阵乘法,得到Q、K、V。矩阵的尺寸是d_model分别映射到d_k、d_k、d_v维,通常d_k = d_v = d_model / num_heads,比如8个头的话就是64维。
这三个权重矩阵不是手工设计的,而是训练过程中学出来的。这也意味着,每个token的Q向量和K向量,是在"如何表达自己的搜索意图"和"如何被他人搜索到"这两个维度上被重新编码过的。用生活类比就是:同一张简历,你分别做了三个版本,一个写满"我想找什么样的机会"(Query),一个写满"我的关键标签"(Key),一个写满"我的完整技能包"(Value)。别人拿他们的查询来匹配你的标签,匹配上了就把你的技能包取走一部分。
很多资料到这里就一笔带过了,但有一个细节值得注意:在Transformer的原始实现里,Q、K、V虽然是三个矩阵,但它们是从同一个输入X投影出来的,也就是说,对于Encoder的自注意力而言,X是同一个。到了Decoder的Cross Attention,Q来自解码端上一层输出,K和V来自编码器的输出,这就是"交叉"的含义。理解了这层,你就理解了为什么代码里Self-Attention和Cross-Attention的API长得很像,只是传入的输入不同而已。
2.2 为什么这三块要分开投影,而不是直接用原始向量
这里其实藏着一个更深的设计动机。原始向量X是Embedding空间中的表示,它融合了词义、语法、位置等多重信息。如果直接用原始向量互相点积算相似度,结果会受到各种混合信息的干扰。Q、K、V的独立投影给了模型三个自由度:它可以把注意力计算所需的"搜索意图"和"被搜索特征"单独抽出来,让匹配更纯粹;同时Value的独立投影让"信息内容"和"匹配机制"解耦,想调整内容时不必牵动匹配方式。
你可以试着自己做个小实验。拿一小段真实词向量不做任何投影直接做注意力,再把投影也去掉换成单位矩阵,模型在表达能力上的差距会非常明显。投影层本质上就是在给模型"腾挪空间",让后续的多头机制有地方发挥。这也是为什么后面很多改进模型仍然保留QKV投影,只是在投影之后加了各种稀疏化或低秩约束,而不是干脆删掉这一层。
动手推一遍维度关系会更清楚。假设输入的Embedding矩阵是[batch_size, seq_len, d_model],W_Q是[d_model, d_k],那么Q就是[batch_size, seq_len, d_k]。Q和K的点积结果维度是[batch_size, seq_len, seq_len],正好是一个注意力矩阵,第i行第j列表示第i个token对第j个token的关注权重。这一步做完,后续的归一化和加权求和就水到渠成了。
3. 注意力分数的计算细节与缩放因子的数学直觉
3.1 Q和K的点积到底在算什么
注意力分数的计算公式长这样:
score = Q · K^T / √d_k
Q的每一行是当前token的查询向量,K^T的每一列是另一个token的键向量,点积结果衡量两者的内积相似度。在高维空间里,两个向量的方向越接近,点积越大;方向越垂直,点积越小。这个相似度体现了"当前token的信息需求"与"其他token的信息标签"之间的匹配程度。
点积相似度有两个性质值得留意。第一,它没有归一化约束,理论上取值范围是全体实数,大小取决于向量模长和方向。第二,它对向量模长敏感,如果某个token的向量模长整体偏大,它作为Key的时候更容易被高亮,这种"模长干扰"在训练初期尤其明显。
训练初期权重矩阵还没有学好,Q、K向量的模长分布很不稳定,如果直接用未缩放的注意力分数进softmax,很容易出现某些位置的分数特别大,softmax输出接近one-hot,梯度几乎为零,模型学不动。缩放因子√d_k就是针对这个问题的第一道防线。
3.2 为什么是√d_k而不是d_k或者其他值
很多人知道要除以√d_k,但不知道这个系数怎么来的。答案其实藏在向量维度和方差的关系里。假设Q和K的每个分量都是均值为0、方差为1的独立随机变量,那么两者的点积结果,也就是d_k个独立乘积的和,它的方差近似等于d_k。这是概率论里的基本结论:相互独立随机变量乘积的期望是各自期望的乘积,E[xy] = E[x]E[y] + Cov(x,y),在均值为0且独立时协方差为0,方差逐项累加。
于是点积结果的方差是d_k,标准差是√d_k。如果不做任何处理,随着d_k变大,点积结果的标准差也在变大,数据分布变得越来越"尖锐",直接进softmax就会落入梯度饱和区。除以√d_k,相当于把标准差重新压回1左右,让点积结果的分布保持在一个相对平缓的区间内,softmax输入的数值范围可控,梯度流动更顺畅。
那为什么不用d_k做分母?如果除以d_k,相当于把方差压到了1/d_k,对低维情况还好,但对64维甚至更大的维度来说,分数会被压得过小,所有位置之间的差异被抹平,softmax输出趋于均匀分布,注意力就失去了选择性。√d_k是最平衡的选择:它不改变方差量级,只是把方差标准化到1附近。
3.3 softmax归一化与加权求和的前向计算
注意力分数矩阵经过缩放之后,对每一行做softmax,得到的就是权重分布。这个分布的每一行之和为1,代表了当前token应该分配给序列里其他各个token的关注比例。说完softmax,下一步就是拿着这组权重去加权求和Value向量。
实现上,这步计算是注意力权重矩阵和V矩阵相乘:Attention = softmax(Q · K^T / √d_k) · V。输出向量的每一个分量都是"关注度加权的V向量混合结果"。在这个过程中,Value向量里的信息被按照关注度比例提取出来,融入了当前token的位置。经过这一层之后,每个token的新表示里都携带了整个序列中所有与它相关的上下文信息。
这里有一个新手容易犯的概念错误:以为加权求和之后就直接输出了,实际上自注意力层的输出后面还跟着一个线性投影和残差连接。原始结构的顺序是,注意力输出经过一个输出投影矩阵W_O,然后跟输入X做残差相加,再做LayerNorm。残差的作用是让模型在优化过程中保留原始信息,防止深层网络退化。LN则负责把数值分布稳定下来,方便后续前馈网络处理。
4. 多头注意力:为什么要分成多组头并行计算
4.1 单头注意力的局限与多头的工作方式
如果一个自注意力层只有一个注意力头,那么它在计算attention时,只能产生一组注意力权重。这一组权重只能表达"一个维度上的亲疏关系"。但语言里的关系是极其复杂的,词和词之间可能同时存在语法依赖、语义关联、指代关系、修辞呼应等多种不同的关系模式。单头注意力就像一个只能从一个角度观察问题的人,很难同时捕捉这些不同维度的联系。
多头注意力做的事情很简单:不把整个d_model投影成一组Q、K、V,而是拆成h组,每组用独立的W_Q^i、W_K^i、W_V^i投影到较低的维度(d_k = d_model / h),每组独立计算注意力,得到h个不同的注意力加权结果,最后拼接起来,再过一层线性投影。这就是"多头"的直观含义。
实际操作中过程略有不同,一般是先把输入X复制h份或者一次性用大矩阵投影,然后用reshape和transpose把它们切分成h组,分别做bmm(批量矩阵乘法),最后再concat。这个接口设计在现代深度学习框架里已经很成熟了,但理解它是怎么拼起来的依然重要。
4.2 头与头之间捕捉的信息差异
头与头之间学了什么,学术界做过不少可视化分析。结论很有意思:有的头主要关注相邻词之间的关系,有点像语法上的近邻依赖;有的头会关注远距离的核心名词,像是在做指代消解;还有一些头专门关注特殊标点或者句法边界。这些头跑的就是同一个Transformer,只是不同头在训练中自发分化出了不同的分工。
这种分化不是人工指定的,而是数据里自然涌现的。训练目标并不要求"你应该当语法头"或者"你应该当指代头",但多头结构和梯度下降的组合,让模型的表达能力在高维空间中自适应地占据了不同子空间。这也是为什么多头个数、每个头的维度这些超参数对最终效果有直接影响,实际上它们决定了模型能同时探索多少种"关系视角"。
从工程角度说,多头带来的另一个好处是稳定的梯度传播路径。注意力机制本质上是一个带权重的求平均操作,单头情况下某个token的输出对另一个token的梯度路径过于单一,一旦这个权重接近0,相关梯度就断了。多头给每条信息传递路径提供了一个"冗余备份",即使某个头学到的注意力权重很小,其他头仍然能维持必要的梯度联系。这个性质对深层Transformer的训练稳定性帮助很大。
4.3 实际模型里是怎么设置头数和维度
最常见的配置是d_model=512、num_heads=8、d_k=d_v=64。为什么8是个常用选择?不是算出来的,而是实验调出来的经验值,需要兼顾表达能力和计算开销。头数太少,关系视角不够;头数太多,每个头的维度太小,单个头能承载的信息量不足,表达能力反而下降。有研究表明,某些任务中头数超过某个阈值后,部分头的注意力模式高度相似,实际上是在重复劳动,白白浪费计算量。
如果你在魔改模型,可以观察一下注意力头的退化现象。训练一个多头模型,把各头在验证集上的注意力矩阵拿出来做相关性分析,你会发现有不少头学出来的模式相当接近。这时候你可以尝试做"头剪枝"——砍掉一部分冗余头,模型效果往往不会明显下降,但推理速度能提升不少。这个技巧在一些模型压缩的论文里经常用到,实际操作起来也不复杂。
还必须提一下:多头注意力在代码实现时,如果只是简单地把Q、K、V拆开分别做bmm,那其实每个头是独立的。但很多框架的矩阵乘法优化依赖于张量内存布局的连续性,所以实际代码里往往是先合在一起计算,再用reshape和permute来拆分,最后concat之后再投影。写的时候注意维度的顺序,别把batch和head两个维度搞混了,这个问题在手动实现Transformer时很容易踩坑。
5. 自注意力相对RNN/CNN的根本优势
5.1 长距离依赖的路径长度问题
RNN处理长距离依赖的能力有限,这个结论不是直觉,而是有理论支撑的。在RNN里,信息从序列的第1个位置传到第10个位置,中间要经过9个循环单元。每一层传递都会经过非线性变换和权重的矩阵乘法,梯度在反向传播时要沿着时间步逐层回传,信号要么指数级衰减,要么指数级爆炸。即便用上LSTM/GRU的门控机制,也只能缓解衰减,没法根治远距离梯度传递困难。
CNN的情况类似。对于标准的卷积神经网络,感受野取决于卷积核大小和层数。想让远距离位置之间建立联系,要么加大卷积核,要么加深网络,两种做法都伴随参数量的快速增长或信息瓶颈。Transformer就不一样了,自注意力在任何两个位置之间直接建立连接,跳跃所有中间层,路径长度恒定为1。从这个角度说,它理论上天然具备处理长距离依赖的能力。
这也是为什么Transformer在NLP任务上优势巨大,后来也被扩展到语音、图像等序列建模领域。只要数据可以被组织成序列,自注意力就能用同样一套机制去捕捉长程联系,这种通用性本身也是它流行起来的原因之一。
5.2 并行性与复杂度权衡
但万能药是不存在的。RNN是串行结构,第t步必须等第t-1步算完,GPU并行化收益有限。CNN在并行性上已经好很多,但感受野受限。Transformer的自注意力在计算上有一个让人又爱又恨的特性:任意两个位置之间都直接相关,所以整个序列可以一次性并行算完,训练时大大缩短了时间。
代价就是O(n²)的空间和时间的复杂度。n是序列长度,注意力矩阵是n×n的正方形矩阵。n=512的时候还好说,n=2048的时候就不太舒服了,等到处理长文档或高分辨率图像特征,序列长度轻松破万,O(n²)直接变成不可承受之重。这就是后面一系列稀疏注意力、局部注意力、线性注意力、Flash Attention等优化方案的动机源头,核心问题都是把注意力从完全的O(n²)里拖出来,尽量逼近线性复杂度。
理解这个权衡很重要,因为"Transformer适合长序列吗"这个问题没有绝对答案,而是看你对复杂度的容忍程度和对精度损失的接受程度。
6. 位置编码:给没有顺序感的自注意力注入秩序
6.1 自注意力为什么必须知道位置
自注意力机制本身对序列顺序是浑然不知的。因为注意力计算是集合层面的交互,把输入序列的顺序打乱,Q、K、V的点积结果完全不受影响,最后的输出对应的只是token的集合关系,而不是它们的顺序关系。你如果不给Transformer提供位置信息,它看到的句子和打乱的词袋之间没有实质区别,这就严重影响了语言建模的能力。
解决方案是显式地把位置信息编码成向量,加在输入token的Embedding上。Transformer原始论文用的是三角函数位置编码,每个位置pos对应一个d_model维的向量,偶数维度用sin(pos/10000^(2i/d_model)),奇数维度用cos(pos/10000^(2i/d_model))。这个公式看起来有点吓人,但它有一个非常漂亮的性质:任意两个位置的编码向量的点积,只依赖于它们之间的距离差,一个位置向量可以用其他位置向量的线性组合来表示,这让模型更容易学到相对位置的模式。
6.2 三角函数编码为什么能编码相对位置
为什么可以用sin和cos的组合表示位置?核心在于三角函数的坐标旋转特性。如果我们把位置编码看成是将token位置投影到高维的"坐标轴"上,每个维度对应一个不同频率的正余弦波,那么位置p的信息就可以和高频波、低频波建立联系。不同的频率组合起来,足以让模型区分不同位置,也能泛化到训练时没见过的更长序列。
后来很多研究用相对位置编码替换了绝对位置编码,最简单的做法是直接在注意力分数里加入一个可学习的相对位置偏置,这样模型对"相隔多少距离"更敏感,而不是对"具体在哪"敏感。这个改动在许多任务上都有明显提升。不过对于理解底层原理来说,先掌握绝对位置编码的直觉就够了,毕竟所有变体都在解决同一个问题:把顺序信息以某种方式注入到模型中。
还有一个容易忽略的细节:位置编码是加在输入上还是拼在输入上?原始实现是加在输入Embedding上。拼接会扩大输入维度和参数量,加法实现简洁且保留了原始语义信息,同时让位置信息和词嵌入信息在训练中自然融合。实际训练中,词向量和位置向量加起来之后的数据分布,对后续的QKV投影会产生直接影响,所以也有人尝试过把位置编码放在注意力计算内部而不是输入上,效果各有优劣,这也是一个可以自己动手做对比实验的方向。
7. 工程实战中值得注意的细节与优化
7.1 显存占用与序列长度的平方复杂度
实战中跑Transformer,尤其是想训练长序列时,最先撞上的总是显存。注意力分数矩阵[batch_size, num_heads, seq_len, seq_len]占的空间随着seq_len平方增长。以batch_size=8、num_heads=8为例,seq_len=1024时,这个矩阵大约需要8×8×1024×1024×4个字节,也就是256MB浮点显存。seq_len到2048,直接翻四倍到1GB,这还只是一个层的中间变量,多层叠加起来相当可观。
Flash Attention的思路就是把这个矩阵切块,在SRAM级别的内存里做分块softmax计算,避免把它完整写回显存,从而把显存占用降了一个数量级。如果你自己写注意力层,也尽量用算子融合的方式减少中间张量物化,能在很大程度上缓解显存压力。
另一个工程技巧是attention_mask的设置。处理Padding时,在注意力分数矩阵的对应位置加上一个极大的负值(通常-1e9)而不是0,这样softmax之后这些位置的概率趋近于0,既屏蔽了无效位置的信息,也避免了数值问题。如果你用的是PyTorch,直接传mask给缩放点积注意力算子,它内部会帮你处理,但如果自己手写,这一步千万不能漏,漏了nans就来了。
7.2 训练稳定性与初始化
自注意力层对初始化和学习率的敏感度比传统RNN高。模型越深,这个问题越明显。我实际调过的模型里,有时候只是把初始化的标准差从0.02调到0.01,训练Loss曲线就能从爆炸边缘拉回来。原因在于Q、K矩阵的初始权重大小直接影响注意力分数矩阵的数值范围,如果初始值偏大,softmax很容易饱和,反向传播梯度极小,模块等于没用。
学习率方面,原始论文用了warmup加线性衰减的策略。前几步用很小的学习率让模型适应初始状态,之后再逐渐加到峰值,一方面是因为Adam优化器在初始步里二阶动量估计不够准确,另一方面也是防止注意力分布在前几步就被推入极端。这个策略现在依然是训练Transformer类模型的标配做法。
Dropout的位置也值得注意。原始实现是在注意力权重矩阵计算完后,对权重做了一次dropout,也就是随机丢弃部分注意力连接,然后再和Value加权求和。这块dropout的作用是防止模型对特定位置的注意力过于依赖,提升泛化能力。如果输入的token本身有embedding dropout,输出层又有最后的linear dropout,多个位置叠加起来,训练和推理行为会跟单点dropout差异很大,推理时要记得关闭所有dropout,否则结果会不一致。
7.3 手动实现自注意力时的常见错误
写自注意力代码踩过的坑,基本都是维度错误和mask使用不当。最常见的包括:多头拆分后忘记把batch和head两维合并回去,导致后续投影reshape出错;注意力分数维度在转置Q和K的时候搞反;mask的维度是[batch, seq_len]还是[batch, seq_len, seq_len]搞混。我建议新手在实现时每一步都用print检查shape,尤其是Q@K^T这个操作前后,多花两分钟确认,能省很多debug时间。
还有个容易被忽视的细节是对Q、K、V做层归一化与否的选择。标准的Transformer不在QKV投影前做LayerNorm,但在一些变体中,有人把LayerNorm放到投影之前,发现训练更稳定,这本质上是在控制Q、K向量的模长变化范围,跟缩放因子的作用有些类似。你可以自己做实验对比,看看在特定任务上有没有区别。
7.4 长序列优化的几条实用路径
如果你想处理超过模型原本训练长度的序列,几条常用路径是:位置编码插值、滑动窗口注意力、全局token加局部窗口的组合。位置编码插值比较好理解,线性插值或NTK-aware插值都是让模型在更长序列上依然能感知相对位置的方案。滑动窗口注意力则是设定一个窗口大小,每个token只参加窗口内的注意力计算,把复杂度从O(n²)降到O(n×w)。实际操作中,滑窗配合全局token可以兼顾局部细粒度关系和全局粗粒度的上下文,这也是很多长文本模型的基本架构。
实测下来,对大多数中等规模的任务,先把序列长度压到模型能舒服处理的范围内,比急着上各种稀疏注意力方案更见效。很多任务不一定真的需要几万token的全局注意力,往往把输入切块、保留有效内容,比增加一个复杂的注意力机制提升更明显。这也是工程里最朴素的真理:先分析数据和任务,再决定用什么模型和机制。