图解大模型注意力计算过程以及KV Cache原理Decoder-only 架构的大模型，其任务是根据历史信息逐步预测下一

一、大模型生成过程

首先，我们看看大模型生成过程中每一步是怎么进行的。由于现有大模型基本都是Decoder-only架构，可以按照以下流程进行理解：

大模型生成过程

<Begin>是一个起始符，用于标记句子的开头。当模型的输入只有<Begin>的时候，模型输出概率最大的词I，表示在模型看来，以I作为句子的实际开头是合理的；当模型的输入变成<Begin>和I的时候，模型预测出下一个词大概率是have。依此类推，整个生成过程就是把模型预测出的词拼接到输入的句子中去，组成新的输入句子后，再让模型预测这个新输入句子的下一个词汇是什么。

到这里，我们会发现一个问题，随着输入句子的长度越来越长，Decoder模块需要计算的内容越来越多，预测新的词的速度应该是越来越慢才对，为什么我们实际体验中的大模型输出都很流畅，甚至有的平均输出达到几百token/s？

二、Decoder模块

回顾一下Decoder模块的结构： GPT1 Decoder模块

紫色部分的归一化和黄色部分的前向传播属于标准的神经网络层都有的结构，这里不做讨论。Transformers架构中最复杂也最容易搞混的是其中的注意力部分，即图中的红色区域。原始Decoder模块共由12层注意力层进行堆叠，前一层的输出作为下一层的输入，因此，我们可以简化成只有一层注意力层进行分析。

每次输入一个新词，都会由橙色区域把该词变成一个固定维度的向量化表示，这个向量化表示包含了语义信息和位置信息，比如把I变成如下向量化表示：

"I" → [0.2, -0.5, 1.3, ..., 0.7]

接下来，该向量会与3个矩阵（WQ、WK、WV）分别相乘，得到新的3个向量，分别叫做q、k、v。这3个矩阵是在训练过程中不断调整的参数，是学习目标之一。至于为什么是这3个矩阵，这里先不管，只要记得这个计算过程就行了。注意力计算公式如下：

$Attention(Q,K,V)=softmax( \frac{QK^T}{\sqrt{d_k}})V$

还有，根据前面描述，Decoder-only架构的大模型生成过程就是不断把模型预测出的词拼接到输入的句子中去，组成新的输入句子后，再让模型预测这个新输入句子的下一个词汇是什么。因此，Decoder-only定义了一种称为掩码注意力的机制，即每个词只能看到当前词与所有的历史词，计算注意力的时候也只能利用这些信息计算，不能用未来词计算注意力。掩码注意力机制

三、注意力计算过程

对于某个单词的向量，与WQ相乘后，得到向量表示q: 在这里插入图片描述

对于有多个输入的，比如现在句子的输入长度是2个词，则可以用矩阵来表示: 在这里插入图片描述

同样的，与WQ相乘类似，我们可以把x和WK和WV分别相乘，同样可以得到k和v向量，这里省略这部分图。

在这里插入图片描述 ${q_1} {k_1}^T$ 的结果是一个标量，而标量的Softmax结果为1，所以最终该向量在注意力的作用下，结果依旧还是 $v_1$ 。这非常合理，因为目前的输入只有一个词。

在这里插入图片描述

到了第二个词的时候，产生的结果跟前面不一样了，不再是一个标量或者一维向量，这时候的QK相乘的结果是一个矩阵。等式右侧的结果中，第二行是新计算出来的，因此第二行是没问题的。问题在于第一行，多了一个 ${q_1} {k_2}^T$ ，这与前面的结果不符。也与前面规定的计算每一步时只能看到前面的信息这句定理相违背。

这时候可以在原QK矩阵相乘的基础上再加上一个掩码矩阵，这个掩码矩阵是一个下三角矩阵，对角线及以下元素为0，对角线以上元素为-∞。经过叠加掩码矩阵后，未来位置的值变为-∞ ，而经Softmax转换后结果变为0，从而实现未来位置的注意力权重为0。

在这里插入图片描述

如果只是为了隐藏后面位置的信息，可以与一个下三角为1的矩阵进行点乘运算，这样也能保留历史和当前位置的注意力信息。但是由于 ${q} {k}^T$ 这个向量中每个元素的值有正有负，甚至也有可能为0，这样在进行Softmax的时候就无法与未来信息区分开。因此这一步的掩码矩阵主要考虑的是Softmax的数学特性。

$\begin{bmatrix} q_1k_1^t & q_1k_2^t & q_1k_3^t\\ q_2k_1^t &q_2k_2^t & q_2k_3^t\\ q_3k_1^t & q_3k_2^t & q_3k_3^t \end{bmatrix} \cdot \begin{bmatrix} 1 & 0 & 0\\ 1 &1 & 0\\ 1 & 1 &1 \end{bmatrix} = \begin{bmatrix} q_1k_1^t & 0 & 0\\ q_2k_1^t &q_2k_2^t & 0\\ q_3k_1^t & q_3k_2^t & q_3k_3^t \end{bmatrix}$

总之，在叠加了掩码矩阵之后，注意力权重矩阵也同样变成下三角矩阵，对角线以上的元素全为0，表示每个词只能看到当前词以及之前所有词的信息，当前词后面的信息是完全看不见的。在这里插入图片描述

在输入长度为2个词的时候，输出的结果发生了变化，输出矩阵的第二行不再等于v2，而是变成了v2'。从本质上来说，是因为k1和v1参与了计算；而从宏观上看，可以认为当前词最终的向量表示为叠加了不同权重系数的历史上所有词向量表示的结果之和，在不同词向量表示的权重系数实际上就是注意力权重。

当输入长度有三个词的时候，同样需要生成q3、k3和v3。为了表述方便，下图用浅色部分表示新出现的变量，之前存在过的向量用深色表示。

在这里插入图片描述

经过前面两轮，这时候我们发现，在所求结果中，只有v3'是需要新计算的，而v1和v2'完全可以复用之前的计算结果，直接把前2个向量拼接到v3'这个向量之上即可。

接下来继续简化这个过程，逐步展示出需要新计算的部分：在这里插入图片描述

到这里我们就明白了什么是KV Cache了，即把每一步计算得到的K和V矩阵保存起来，以供下一步计算使用，典型的以空间换时间的优化思路。

同时到这里也能回答那个经典的问题了：为什么只有KV Cache而没有Q Cache？ 这是因为当前想要的计算结果只与当前的q、k、v和历史上的K、V相关，历史上的Q根本不参与计算。

看到这里，细心的读者可能会有疑问：不对，尽管过去出现的K和V能进行缓存，避免了与WK和WV的重复矩阵运算。但是输入序列长度的增加，新输入的q还是必须要与K和V进行计算的，随着输入长度变长，当前词需要计算与历史词的注意力也越来越多，生成每个新 Token 的计算量理论上会线性增长，但实际上生成速度并不会显著变慢。这实际上不属于KV Cache的内容了，KV Cache主要解决的是重复计算的问题，而对于大矩阵计算和推理过程的优化，则涉及到其他一些算法，如Flash Attention和vLLM中的PagedAttention等。这些算法一般要求预先分配出比较大的显存，同时将显存进行分块处理，将多步计算过程合并为一次计算，从而提升推理速度。

四、一些思考

经过上述分析过程，QK两个矩阵相乘结果本质上等于一个权重，而用于表示最终结果的是V，即值向量矩阵。理论上也可以把QK两个矩阵相乘变为一个权重矩阵，甚至可以变成多个矩阵相乘都行，只要满足输出权重矩阵的维度即可。之所以不用一个权重矩阵是因为单一的权重矩阵会缺少互信息，Q矩阵表示当前词需要关注哪些信息，而K矩阵表示其他词如何关注当前词，通过这2个矩阵丰富了模型的表征能力。而多个矩阵相乘虽然可以带来多个维度的表征，但实际上带来巨大计算量的同时可能收效甚微，因此一般只用2个矩阵来表示注意力。

此外，Decoder-only 架构的大模型，其任务是根据历史信息逐步预测下一个词。若使用双向注意力，在训练时模型能看到完整序列，但推理时还是只能看到历史部分，这会导致训练和推理不一致，从而使得模型性能严重下降。这种方式在训练时无疑极大增加了模型的训练难度和收敛时间，但模型的上限同样很大，如果做好的话潜力不可限量。当大部分人都在用Encoder-only和Encoder-Decoder进行刷榜的时候，能坚持做Decoder-only确实非常有远见，不得不佩服OpenAI的前首席科学家Ilya Sutskever。而如今GPT4.5和Grok3的现状同样验证了他提出的预训练即将终结的观点。在大模型架构不会显著改进的情况下，推理时计算应该是现阶段大模型唯一的发展路径了。