引子

在之前的学习过程中，总是听到有关KV Cache的话题，关于它有多么大，如何成为大模型推理的存储容量和带宽瓶颈，等等。 曾经的我对大模型推理^[1]的原理和细节并无太大兴趣，也没有进行深入了解。 对KV Cache这个词，我望文生义的理解为一堆“Key-Value”的键值对，类似编程语言中“map”或是“dict”那样的数据结构。 这个暑假，由于后续的一些安排，我开始试图学习有关大模型推理原理的知识。我发现KV Cache的含义并不是我曾臆想的那样。因此，作为对我近期所学知识的梳理与总结，我打算在这篇文章中梳理一般的Transformer模型推理的计算过程，并简单说说KV Cache，以及其他的一些LLM领域的关键词，在这个计算过程中所处的位置。

不得不说，这个领域已经有无数篇文章了。我获取LLM相关知识的方式正是阅读这些文章以及观看一些科普视频。这其中当然还有LLM本身的协助。它们的寓教于乐能力显然强过我的笔力。因此，即便这篇文章是以“从头开始认识”的视角写成，我也很难建议任何人使用我这篇文章作为学习资料。写这篇文章的目的仅仅是方便我整理自己的粗浅思绪罢了。如果读者有任何想点评的，或是发现了什么疏漏和错误，非常欢迎你向我指出。

输出一段话

或许还是从输入输出开始说起吧。从日常使用中也能体验到，大语言模型的输出是逐字（准确的说是逐token）产生的。大语言模型的输出是一个迭代的过程。首先，它以用户提供的所有文本作为输入，经过一堆计算，输出它的第一个词（token）。随后，原始的输入，加上大语言模型刚刚输出的那第一个词，组合起来，再一次作为输入送入模型，经过计算，可以输出第二个词。如此往复，则可以输出完整的段落和文章。

1
2
3
4

第一次：[用户输入] -(模型计算)-> [第一个词]
第二次：[用户输入+第一个词] -(模型计算)-> [第二个词]
第三次：[用户输入+第一个词+第二个词] -(模型计算)-> [第三个词]
......

Q/A:

• 迭代输出，每次送入同一个模型？
  是的。输出第二个词，所采用的模型结构、参数都和输出第一个词时完全一样，只是输入在原始基础上有追加。当然，这只是数学上的相同，并不是具体实现上的约束。

为了扫清理解的阻碍，有必要阐释token究竟是什么。Token是模型处理文本所采用的最小单位，在多数时候对应一个词、词的一部分、或者标点。 或许可以想象模型有一本token的“词典”。 GPT3模型有五万个^[2]预设的token，也就是词典中有五万多个词。现在的先进模型，这个值大概在几万到几十万不等。模型只认识这些预设好的“词”，因此针对用户的文本输入，有必要进行“分词”（学名token化），将输入一个一个对应到已有的词典中。

Q/A:

• 哪些词会被视为token？
  词典有多大，具体有哪些词，在模型训练之前就确定了。词典的大小是模型规模设计的一部分。而具体使用哪些词填充词典，则是设计者使用算法，根据大量文本的统计分布总结的。

• 如果用户输入的词比较生僻，没有收录进模型的“词典”怎么办？
  在常用词之外，还会有所有字节码对应的token作为备选。因此请放心，不会存在输入无法token化的情况。

• 多国语言的相似词汇会如何处理？会合并为一个token吗？
  只要字符不同（比如汉字和英文），不论含义是否相似，都将分配不同token^[3]。多国语言抢占有限词典，这是词典制定时面临的现实。在统计语料的过程中，占比大的那种语言，更有可能分到更多的token。对于token少的那个语言，模型性能可能不佳，但就像第二问说的，总能有对应token。

• 如果token词典出现问题，后续计算会不会受严重影响？
  会的。一个典型例子是gpt4o时代的一个乌龙。4o采用的词表中“给主人留下些什么吧”9字对应一个token，导致模型根本没能理解这个token的含义，对这句话的翻译五花八门。

输出一个词

明确token的概念后，让我们将目光聚焦于输出大段文本这个多轮次过程的其中的一个轮次，即输出一个token的过程。

一个轮次的流程大概是这样：

• 首先，当然要对输入文本分词，按照词表拆分为一串token
• Embedding：将每个token变为一个高维度的向量表示。原输入的一堆token在这一步变为了一堆向量，它们可拼成一个矩阵，姑且记作，其中的每一列^[4]代表一个token向量；
• Attention：是transformer最具代表性的步骤呢。具体细节先按下不表，其效果是更新了 的值；
• 多层感知机（MLP） / 全连接层: 更像传统前馈神经网络。其效果是更新了 的值；
• 重复：重复执行“Attention-MLP”的组合， 被逐渐修改。重复次数也即模型的“层数”，每层参数并不相同；
• Unembedding：完成所有层计算后，理想情况下，此时 的最后一列^[4]，将包含模型对原始文本的下一个词的预测信息。取 的最后一列，将它翻译成最合适的token（依据token的概率分布选择），然后输出token对应的文本。

这只是流程的大致说明，接下来马上进行逐个的详细描述，因此这里没有Q/A哦。

Embedding & Unembedding

输入文本变为token序列后，第一步就是将每个token变为高维向量。在随后的计算中，这些高维向量的值会逐渐更新，直到完成所有层数的计算，取出更新后的，原文本末尾的token对应的向量，将它反向变回token，则完成了一个token的输出。这第一步就叫embedding，最后一步叫unembedding。

token变成的向量究竟有多少维？ 这取决于模型的规模设计。以GPT3-175B为例，它的token向量维数是12288.

每个token会变成什么向量？ 这是模型参数的一部分。这部分参数可以记作一个 矩阵^[5]，矩阵的每一列代表一个token应该变成什么向量。在这个定义下，矩阵的行数（下文采用 表示）是token向量维数，列数（下文采用 表示）是模型使用的这个token“词典”中“收录”的词的个数（记作 ）： 在Embedding过程中，根据每个token在“词典”中的编号 ，即可在这个矩阵中找到它所对应的向量。在这个意义上，这个 在运算上并不像是矩阵，倒更像是一个“查找表”。

这些参数有什么含义？在人工智能领域，询问参数的含义似乎永远无法得到明确的答复。这些参数是在训练过程中使用优化算法逐渐确定下来的，是计算的结果，并不具备什么人为赋予的含义。在推理时它们是固定已知量，不随输入文本改变。如果硬要深究它们的含义，只能说，这些高维向量的取值和方向，某种意义上代表了token的初步的语言含义。

值得补充的是，因为语言有严格的语序，单纯查表映射无法体现词的先后顺序。因此 Embedding 还有一个关键步骤叫位置编码（Positional Encoding），它会将 token 在文本中所处位置的信息“注入”到向量中。

经由embedding步骤，原始文本（已被转化为token的序列）变为了多个向量。这些向量按照文本顺序拼接起来，成为一个矩阵 。这个矩阵就是模型接下来要面对的输入。仍然采用列向量表示，则：

描述完初始的Embedding步骤，现在描述末尾的Unembedding步骤。 经由多层Attention和前馈网络处理， 发生了诸多变化。 将最终的 的最后一列 反向变换回token的步骤是unembedding。这是一步矩阵向量乘：

输出的 向量学名logits^[6]，每个分量代表词典中每个token的得分，越大代表模型越倾向于输出那个token。得分经 函数处理为标准的概率分布（都为正且和为1），再根据概率分布选择恰当token，作为模型的输出^[7]。

是unembedding步骤的参数矩阵。 可以作为独立的参数交由模型自由选取和优化，也可以和 绑定起来。比如，有一种做法是直接将 的转置作为 。那么计算的效果就是：

这样天然具备不错的效果。因为点乘可以提取两个向量的方向信息，如果 和某个 token 对应的向量 接近，那么点乘自然就更大。此外，高维空间有个特点：两个没啥关系的向量通常会“相当正交”（即点乘很接近0）。因此，这种点乘会比较有效的把最相关的token显露出来。转置法有效的省去了整个 的参数量，据说^[8]被不少模型所采用。

似乎Embedding和Unembedding两步能说的就这些。

也许读者还会有一个困惑：如果最终输出只在乎 的最后一列，那计算 的其它列究竟有什么用？这个问题可以从两个方面回答。

一方面，从训练的角度，生成的 的其余列也可以用于计算损失，从而优化模型参数，提高训练效率。

另一方面，从推理的角度，在模型推理过程中， 的列与列之间将频繁的“交换信息”，因此即便最终不需要 的剩余列，它们也需要参与计算并不断更新。而涉及“交换信息”的最重要的步骤，就是接下来将描述的“Attention”步骤。

Attention

Attention 是 Transformer 网络结构中负责“交换信息”的核心模块。它能够根据上下文动态计算不同位置之间的信息关联，使每个 token 都能够融合来自其它 token 的信息，也是 Transformer 区别于传统序列模型的重要特征之一。

如何让模型理解上下文之间的相关性，一直是人工智能领域的重要课题。传统卷积网络可以通过局部感受野实现特征融合，但信息交互范围受到卷积核大小的限制；而循环神经网络虽然能够处理序列信息，却难以充分利用并行计算能力。Transformer 所采用的 Attention 机制，使单层计算中任意位置之间都能够直接建立联系，在捕获长距离依赖关系的同时，还具备高度并行化的计算能力，因此成为现代大规模语言模型的重要基础。

进入数学之前恐怕还是先明确符号。 Transformer是多层“Attention 前馈网络”依次处理。 假设这个网络有L层，那么我们可以把Embedding完成的Token向量拼成的矩阵记作 ， 完成一次“Attention-前馈网络” 之后变为 ，以此类推，完成全部L层，即将进行Unembedding的矩阵变为 。我们这里讨论的Attention，是每层的前半步骤，姑且使用 吧。虽然有些蠢，但还蛮形象的：

完整的注意力运算是“多头”机制，是“单头”的多路并行和加总。因此先详述“单头”机制。

单头注意力机制

第l层Single-headed Attention运算（也即 变换到 的计算）是^[9]： 各个关键项的维度（略去标识所处层数的括号）是：

注意到这里出现了一个新的维度量 。 是注意力计算所使用的维数，k代表key。一般来说 ，因此将输入 映射为 的过程是一个降维的过程。在作为案例的GPT3-175B中，，相比 来说是妥妥的降维。

下面分步骤讲述。沿袭列向量线性代数的优良传统，我们从右往左看，先从 、 矩阵说起。将 和 两个矩阵各自作用于 ，形成 和 矩阵。 和 分别代表Query和Key。上文已经提到，这是一个降维的过程。 根据3B1B在视频中的类比， 矩阵类似于每一个 （ 的每一列） 提出了一个“问题”， 矩阵类似于每一个 给出了一个“答案”。 从降维的角度考虑，生成 和 的过程，相当于从原先高维的向量 中 提取了某个方向的、偏具体的特征 ，并且在这个特征下判断token和token之间的相关性。

判断相关性的方法是做点乘，也就是 ， 这个矩阵的(i,j)元是 。

可以按列理解这个矩阵的含义。矩阵的第j列，代表了所有 生成的key对 的query 的相关性如何。 在接下来的步骤中，所有token将依据相关性，以不同比例调整 ，从而实现将其它相关token的“含义”“赋予”第j个token。

对 按列进行softmax处理，使得每列变为合法的概率分布。随后，通过 生成，V代表Value。每一列 代表 “如果要向其它向量 ‘赋予’自身 的信息，需要对 施加怎样的修正或偏移量” 。以softmax得到的分布作为权重加权所有的，就是表达式中 与经softmax之后的 的矩阵乘。再和原始的 相加，一个注意力计算就完成了对 的值的更新，由 迈入 。

从这一串文字描述可以看出，K和V并非“由key找value”的关系。从产生来源角度，K、Q、V都是由输入矩阵X产生，关系有 、 、 , 并无 。 硬要说的话，k与v的关系是“ 和 的点乘（归一化后）作为权重用来加总 ”这种间接的关系。 因此，不论怎么看，KV Cache都不可能是“由K指向V”的散列表，这纯粹是臆想了。

在进入多头注意力机制之前，有必要补充关于 “masking” 步骤的内容。上述的这种注意力机制被称为 “自注意力” 机制。其特点在于，产生key的 和产生query的 是同一个 （即，一段文本内部token彼此交换信息）。据说，一些多模态场景可能存在两串不同内容的序列互相交换信息的情况，但那不在我们的讨论范围内。对于自注意力机制，通常的做法是对点乘结果矩阵 进行 “掩膜”（Masking） ，把 时 项改写为 ，经softmax后变为0。这种做法相当于 强制要求文本序列中靠后的 不允许改变靠前的 。

每个词只能理解前文，不能理解后文，这不会影响模型性能吗？事实上，模型总可以将前文的关键信息传递给后文，让靠后的 承载更多的信息。YouTube一条评论说，对大模型的分析发现，模型倾向于将句子的信息嵌入到句子末尾的标点符号里。因此这种做法是不会影响模型性能的。

同时，masking带来至少两个好处：一方面，在训练过程中，它避免了模型“偷看”下文从而“作弊”，使最终输出 每一列都能计算损失，训练模型，提高训练效率；另一方面，它直接使得KV Caching成为可能，不过这就是后话了。恐怕得等本文把“第一轮次”，输出第一个token的过程梳理完，进入“第二轮次”的时候，才能提到Caching的内容，可能要拖到下一篇文章了。

这大概就是“单头”注意力机制的主要内容。

多头注意力机制

描述完“单头”，“多头”机制就比较容易描述了。

wip

前馈网络