写这篇文章的起因是我在调试一个生成模型的时候盯着某层的注意力图看了很久越看越觉得不对劲有些头和语义明显不对应有的头几乎学到同一个模式。旁边一位同学问我是不是QKV矩阵初始化有问题。我下意识点头但回头一想真要我把QKV从头到尾讲清楚我发现自己其实也是半桶水。于是就有了这篇总结把Transformer里始终绕不开的QKV矩阵从头捋一遍它们到底从哪来各自在算什么为什么少了其中任何一个注意力机制就转不起来。适合刚接触Transformer的初学者也适合那些已经能跑通模型、但对注意力内部结构还模糊的同学。我不会堆公式推导尽量用大白话再用实际调试经验来对照。1. 先搞明白一件事QKV到底在干嘛1.1 一个词典类比Q是找K是被找V是内容我第一次接触注意力机制的时候看到文档里写Query、Key、Value脑子里直接跳出数据库查询。后来发现这个类比其实非常准只是很少有人把它延伸到算法细节里。你可以想象你有一本巨大的词典词典里每个词条都有一张索引卡卡上写着这个词大概在讲什么主题这就是Key词条正文则是真正的内容这就是Value。现在你想查一个东西脑子里会先形成一个我要找什么的模糊描述这就是Query。注意力的过程就是拿你的Query去和词典里每一张索引卡比对相似度谁跟你的需求最匹配就把对应的正文内容提取出来加权融合成输出。Key负责被匹配Value负责提供内容Query负责发起匹配。这个三角关系在Transformer的每一层中都在重复发生。从这个角度看QKV不是三个独立的神秘矩阵而是一套检索-打分-取值的工作流。你可以把Q理解为当前token在问谁和我的语义相关K理解成其他token在回答我这边包含哪些语义特征可以让你来匹配V则是如果你觉得我相关就把我的内容给你用。1.2 从RNN到Transformer为什么要引入QKV既然QKV是检索思路那就有个问题为什么以前的循环神经网络不这么干非得等Transformer出来才普及RNN的核心问题是信息要沿着时间步一步步往后传。你在第1个词看到的信息传到第50个词的时候理论上还在脑子里但实际上早就被中间步骤稀释得差不多了。长距离依赖靠的是隐状态这个窄窄的管道信息要依次经过每个词才能传过去一旦链路太长梯度消失就找上门了。注意力机制换个思路不再按顺序传而是让当前词回头看所有历史词直接问一句你们谁跟我当前的需求相关这样一来距离就不再是问题。最早序列到序列模型里注意力是在编码器和解码器之间做的解码器每一步都能直接看到编码器所有位置的信息。但当输入变长如果还是每次从头到尾扫一遍计算量会非常吓人。Transformer的贡献是把注意力用到每个位置上让序列任意两个位置都能直接建立联系并且用矩阵运算一次性算完所有位置之间的匹配关系。这里就离不开QKV的设计它让每个token同时生成三种角色然后通过一个矩阵乘法并行计算所有token两两之间的相似度。从运算的角度讲QKV本质上是对输入做三次不同的线性变换让同一个token在不同角色下分别表征。没有这套拆分你就只能拿原始向量去算相似度那相当于让同一个人既当裁判又当运动员表达力和灵活性都会差很多。2. QKV的来历它们不是天上掉下来的2.1 词向量到QKV一次线性变换不少人以为QKV是模型输入自带的东西其实不然。输入只有一个向量——通常叫token embedding可能是随机初始化的词向量也可能是预训练好的词嵌入。QKV全部是后续变换出来的。具体来说每个token的向量x会分别乘上三个权重矩阵得到对应的q、k、vq x · Wqk x · Wkv x · Wv这里Wq、Wk、Wv就是QKV矩阵。它们都是可训练参数形状是输入维度 × 注意力维度。输入维度是token向量的大小比如768或者1024注意力维度则是每个头的head_dim。关键点在于这个变换是逐token独立做的没有序列之间的信息交互。也就是说QKV只负责把当前token自己的特征映射成三种不同的表示真正的token间信息交换发生在后面的注意力分数计算里。我见过不少人写自定义Transformer时在这里犯过一个经典错误直接把输入向量重复三份当作QKV用不乘任何权重。那样模型完全丧失学习能力因为Q和K永远一模一样注意力分数矩阵必然是常数。三个矩阵存在的意义就是让网络自己决定哪些维度用来做匹配、哪些维度用来做内容输出这个分工是学出来的。2.2 W矩阵是怎么被训练出来的那Wq、Wk、Wv是怎么学到合适数值的答案是梯度下降。训练过程里模型先随机初始化这些矩阵然后用数据不断迭代。你可以把这些矩阵看成多个旋钮组合每个矩阵里的每个数值都是一个旋钮模型通过损失函数反馈决定往哪个方向拧旋钮让最终输出更接近期望结果。比如某个Wk矩阵里的某个参数变大一点可能就会让某两个token的相似度变高导致这个token更多地被提取出来。如果训练数据告诉你这么做是好的梯度就会把这个参数往大调如果不好就反向调。这里有一个细节值得提初始化方式很关键。Transformer里默认的初始化会对权重做方差缩放把数值控制在一定范围内。如果初始化太大QK^T算出来的分数可能一开始就爆炸softmax直接饱和成one-hot梯度几乎传不动如果初始化太小所有分数接近注意力退化成平均池化学得又慢。我之前在一个小模型上手动改了初始化方差发现同样的数据、同样的步数收敛速度能差出好几倍可见这一步不是无关紧要的。另外有人会问QKV矩阵在训练过程中有没有什么规律我的观察是训练早期它们变化很快放出来的注意力图基本都是乱的和语义几乎没什么对应。训练到中期部分头开始出现关注下一个词关注标点关注句首这类模式到后期头之间的分化越来越明显这其实对应了多头机制的含义。2.3 训练与推理时QKV的表现差异这里想专门讲一个容易被忽略的点QKV在训练和推理阶段的表现不完全一样。训练阶段所有位置的token可以并行计算QKV整个序列一次性算完。由于所有位置都可见Q和K可以任意交互注意力分数矩阵是完整的方阵。推理阶段就不同了尤其是在自回归生成时每走一步只能看到当前位置和之前的位置。如果不做处理新token的Q会去匹配所有位置包括未来的token那信息泄漏就来了。解决办法是加因果掩码在计算注意力分数时把未来位置的分数强行设成负无穷softmax之后这些位置权重归零。掩码加在QK^T的结果上不影响QKV本身的生成。还有一个和推理强相关的概念叫KV Cache。生成第t个token时前面所有token的K和V其实已经算过了结果一直没变。为了避免每步重复计算推理框架会把前序的K、V存起来每一步只算新token自己的q、k、v然后用新的q去匹配缓存里的所有k再对缓存里的所有v做加权。这个优化就是KV Cache它直接利用了K和V在推理过程中保持不变的特性。理解这一点对优化生成速度很有帮助。一个直观感受是序列长度越长KV Cache占的内存越大因为缓存大小是层数 × 头数 × 序列长度 × 维度这个量级。很多模型在长上下文场景下的内存瓶颈往往不在参数本身而在这些缓存出来的K和V。3. QKV的作用它们怎么配合完成注意力3.1 注意力分数的完整计算Q和K求相似度有了Q、K、V之后下一步就是算注意力分数。Transformer用的是点积注意力思想很直接两个向量的点积越大方向越接近就认为越相似。对序列里任意两个位置i和j我们先拿位置i的q向量和位置j的k向量做内积得到分数s_ij qi · kj。这个分数就表示位置i对位置j的关注程度。把所有位置都算一遍就得到一个完整的分数矩阵。矩阵运算是这样的如果Q的形状是[seq_len, d_k]K的形状也是[seq_len, d_k]那么Q K.T就是[seq_len, seq_len]的分数矩阵。每一行是某个token对所有token的分数分布每一列是所有token对某个token的分数分布。有了分数矩阵接下里要做两步先除以缩放因子再做softmax归一化。最后拿着归一化后的权重去加权V得到输出。我把这个计算过程整理成一个简单伪代码方便你对照官方实现scores Q K.T # 所有token两两点积 scores scores / sqrt(d_k) # 缩放防止数值过大 weights softmax(scores, dim-1) # 按行归一化每行和为1 output weights V # 加权整合被选中内容这个运算在多个注意力头之间是并行做的。输入维度768头数12那么每个头分到64维的Q、K、V各自独立算注意力最后再拼回去。3.2 为什么除以根号dk很多人卡在为什么要除以根号dk这个问题上。解释其实不复杂它和数值稳定性有关。点积得分的大小会随维度增加而变大。想象两个随机向量维度越高它们的内积的期望和方差都会变大具体来说方差约等于dk。如果dk是64甚至128点积结果动辄几百上千一旦进入softmax指数函数会把这些差距拉得极大结果就是一个远古时期的问题大多数权重几乎归零梯度在小数点后淹没。除以根号dk的本质作用是把点积分数的方差压回1左右。这样softmax的输入整体落在合理区间梯度更容易流动。你可以把它理解成给分数做了一次归一化让模型在维度变化时保持行为稳定。我在调参时曾经把缩放因子手动去掉训练直接炸了loss飙升注意力图瞬间全变成极端分布。恢复缩放以后模型才稳定下来。从那之后我对这个除法的态度是别动它除非你清楚自己在做什么。补充一个细节缩放因子里用的是dk也就是单个头的维度不是整个embedding维度。多头机制下每个头的点积规模取决于head_dim所以缩放要和head_dim匹配。很多开源实现里代码写的是头维度开根号你找不到维度的逻辑谜底就在这。3.3 V值如何加权聚合softmax的意义权重算好之后注意力输出其实就是对所有value做加权求和output_i sum_j weight_ij * v_j。这里V向量携带了token的实际语义内容而权重决定了我到底要多少来自你的内容。如果把Q和K比作匹配器V就是载体。匹配器决定信息的流向载体决定流过去的具体内容。两者是协作关系而非竞争关系这也是为什么QKV三个矩阵缺一不可。softmax在这里有两个作用一是把分数变成非负且和为1的概率分布二是引入竞争机制。同一个query的注意力权重会此消彼长一个位置被强调其他位置份额就会减少。这让模型能做出相对尖锐的选择而不是和稀泥。但softmax也有个副作用它无法输出零权重只是把极小的权重压到近乎为零。所以注意力图谱上很多看起来完全不关注的位置其实只是权重极小不是真正的零。这个特性在实际调试时要注意不然你会误以为模型在某个位置有硬性屏蔽实际上只是概率分配的结果。4. 多头注意力把QKV拆成多组模型才够灵活4.1 多头拆分的原理单头注意力的问题是一个匹配规则用到底太单一。比如找下一个词和找句子的主语是两种完全不同的检索需求单靠一组QKV很难同时兼顾。多头注意力就是把QKV各切成若干份每一份独立走一遍注意力计算。每个头相当于一个并行的检索器有的头负责局部相邻词有的头负责长距离依赖有的头负责标点或语法结构。最后把所有头的结果拼接起来再过一层线性投影得到最终输出。从参数角度看多头并没有增加太多额外参数因为每个头的QKV都是从同一个embedding里分出来的总维度和单头时基本一致。它的意义在于让模型用不同的子空间做不同匹配这种解耦提升了表达效率。我训练模型时专门观察过头部分化。开始的时候各头输出几乎一样看起来非常冗余训练充分后头之间开始分化。有的头注意力集中在下几个token有的头集中到句首有的头能稳定关注到语义对照词。这就是多头带来的分工合作。4.2 MHA、MQA、GQAQKV在变体里的不同形态如果QKV这么有用那是否每个头都要独立一套后来的工作确实在共享这件事上做过不少文章。MHAMulti-Head Attention标准多头每个头有自己独立的Wq、Wk、Wv灵活度最高参数量和计算量也最大。MQAMulti-Query Attention所有头共享同一组K和V只让各自Q独立。这样KV Cache直接降成一份推理显存占用大幅下降代价是表达能力略有折损。早期的一些大规模推理场景用过。GQAGrouped-Query Attention折中方案把多个头分成若干组组内共享K和V组间独立。既保留了部分多样性又显著降低了KV Cache开销。现在很多新模型在长上下文场景下都采用这个思路。选哪个结构本质上是在表达能力和推理效率之间做取舍。如果你在跑长文本生成KV Cache是最明显的瓶颈之一那GQA几乎是必选项如果任务比较短、又特别看重效果标准MHA更稳。这个选择可以在模型结构层面直接改不需要动训练数据。我自己实际测试过一个对比同一批数据MHA版本比GQA版本在评估指标上高大约百分之几但推理速度慢了快一倍显存占用也高一大截。在短序列场景这点差距无所谓一旦序列拉到几千甚至上万GQA的优势就非常明显。5. 实际踩坑心得QKV在真实模型里长什么样5.1 查看已训练模型的QKV理论讲了很多不如直接动手看。我调试的时候最常用的方式是取某个中间层的Q、K、V分析它们的范数和余弦相似度。一个常见的观测手法是输入几个句子把某一层的Q、K、V全部取出来计算它们的L2范数。如果某个维度特别大或特别小往往说明对应头输出分布有问题。我见过一种情况某层模型有个头V的范数比其他头大几千倍结果注意力输出的幅度被这个头带偏层归一化前数值直接飞起。定位到之后发现是初始化的时候出了点小问题重训之后恢复正常。还可以做法对Q和K做余弦相似度观察残差连接前和残差连接后的变化幅度。如果相似度过高说明这个头的匹配模式太单一对语义的区分不够如果相似度过低可能是因为分布已经偏移到训练数据之外也可能是头与头之间严重重叠导致资源浪费。5.2 常见误区澄清关于QKV有几个误区我见得太多了值得专门拿出来说。第一个误区是Q和K应该完全不一样。实际并非如此。很多训练充分的模型里同一个头对应的Q和K在部分维度上仍然有较高相关性这不是bug。好的分工不是让Q和K完全不相似而是让它们的差异是有用信息的编码结果。第二个误区是多头越多越好。增加头数确实能提升表达力但也会带来优化难度和计算成本。头数超过一定阈值之后收益会递减甚至出现各头学习到相似模式导致冗余。某些论文里甚至观察到一些头在训练后期几乎可以剪掉而不影响性能。第三个误区是QKV是对称的。Q和K在注意力计算中虽然都参与了相似度计算但角色完全不对等Q来自当前tokenK来自被检索token。调换它们的位置会让模型关注方向完全反转行为天差地别。所以一定要记住Q是需求方K是供给方的索引。第四个误区是注意力权重高就是模型在重点判断这个位置很重要。注意力高只说明这个位置的信息被大量聚合到当前输出中但不代表它一定是对分类或生成最关键的token。有些高频词比如的“了”也会获得很高注意力因为它们确实频繁地充当语法连接点但这不等于它们是语义核心。5.3 一些小技巧KV Cache与因果掩码最后分享几个实用技巧。先讲KV Cache的落地细节。推理框架里KV Cache一般按层、按头存储每一层的K和V是分开的两个数组。生成新token时只需计算当前token的q然后把它和缓存里所有K拼起来做矩阵乘法V同理。为了让这一步高效许多实现还会把缓存预分配为最大长度避免写的时候频繁扩容。因果掩码的实现也值得留意。常见做法是构造一个上三角矩阵把右上角部分填成负无穷加到分数矩阵上。注意这里要加在缩放之后、softmax之前。有些框架提供了更快的融合算子直接把掩码融进注意力计算里减少一次显存读写。序列特别长的时候这一步优化收益非常明显。还有一个调试经验如果发现模型输出重复、卡在某个词循环可以考虑检查注意力权重的熵。正常训练的模型注意力权重分布通常不会极度集中在一个位置如果某个位置权重接近1、其他位置几乎为零说明softmax温度出了问题或者模型已经退化到只看上一个词的懒人策略了。这时可以试试调整推理超参数比如降低温度或者检查是否有层把注意力分数数值规模拉得异常。关于QKV的最后一个体会是不需要害怕这些矩阵它们只是可训练的线性变换。搞懂它们的关键不是背公式而是理解注意力机制在做一件什么事让每个token主动去检索其他token的信息然后决定要多少内容。Wq、Wk、Wv就是那套检索规则训练过程不过是让这套规则逐步逼近任务需要的匹配模式。你只要多看几层实际输出多跑几次可视化脚本很快就能建立直观感觉。等到看注意力图不再是一堆无意义的色块能看出某个头在盯什么、某个头在找什么的时候你就真正把Transformer里的QKV彻底搞懂了。