@Hi_Mrinal: 这是关于KV缓存的最佳阅读，直观上非常好读 https://medium.com/@saad.ahmed1926q/kv-…

这是我在 KV 缓存方面读到的最好的一篇文章，直观易懂，非常棒。文章链接：https://medium.com/@saad.ahmed1926q/kv-cache-explained-intuitively-2b425a36dfc7… — # KV 缓存直观解释 来源：https://medium.com/@saad.ahmed1926q/kv-cache-explained-intuitively-2b425a36dfc7 Saad Ahmed Siddiqui (https://medium.com/@saad.ahmed1926q?source=post_page—byline–2b425a36dfc7—————————————) ## 目录： - 词元与嵌入 - 关于仅解码器模型的概念 - 注意力机制的直觉 - 基础语言模型中的推理 - 为什么没有 KV 缓存效率低下 - KV 缓存 - 结论 - 参考文献 ## 开始吧 在讨论语言模型时，会抛出很多术语——比如“词元”、“嵌入”、“编码器-解码器”、“注意力”等等。所以首先，我会尽量让你熟悉那些最基本、最常用的术语，了解它们会让生活轻松很多。这篇博客适合每个有兴趣理解像 ChatGPT 这样的语言模型基础构建块的人。而且，是的——甚至非机器学习背景的人也完全欢迎！我们不会直接跳进 KV 缓存，而是先看看一些预备知识。 ## 词元与嵌入 你们中的一些人可能知道，机器学习模型只能理解数字，不能理解自然语言。我们都习惯于用自然语言提示词与 LLM 交互，比如“帮我做这个”或“给我解释一下”，但我们不会直接把这些输入给模型。我们首先需要处理它，使模型能够理解。所以我们需要将提示词转换成某种模型可以处理的数值表示。 ### 词元 我们使用一种称为词元化的过程，将输入序列分割成称为词元的小块。在一个非常简单的意义上，你可以把它想象成把句子分解成单词。虽然实际上并没有那么简单，但这只是为了让你有个粗略的概念。我只是想澄清一下，免得有人来找我理论。我只是想给出一个基本概念。 按 Enter 或点击查看完整图片 词元化示例 ### 嵌入 每个词元都被分配一个数值表示。这样每个词元我们得到一个向量，称为嵌入。嵌入本质上是一个具有若干维度的向量，它捕捉了单词的含义。嵌入中的每个维度代表一些潜在特征，有助于区分不同的嵌入。让我们用一个例子来理解。 虽然在实际场景中，嵌入通常由数千个维度组成，而且我们不知道哪个维度代表哪个属性，但模型仍然可以用这些潜在特征来有效区分不同的嵌入。你可以把它想象成某个空间中的一个点，基本上如果两个单词相互关联（比如猫和狗），那么它们的嵌入会在该空间中彼此靠近；而如果单词不相关（比如马和建筑），那么它们的嵌入会相距很远。我们还会在输入嵌入中添加一种叫做位置编码的东西，这样每个单词都携带了关于它在句子中位置的一些信息。虽然我们不会深入探讨具体细节。 > 现在我知道你们肯定会想：“好了好了，我们懂了——快点进入正题吧。我们是来看语言模型的……够了！”那么事不宜迟，让我们进入真正有趣的部分。 ## 关于仅解码器模型的概念 由于这篇博客是关于KV 缓存的，所以先对解码器有一个基本理解很重要。那么我先把所有复杂的东西都去掉，只试着感受一下解码器做了什么。下面是一个非常简单的解码器示意图。需要注意的一点是，它被人为简化了——例如，我用了掩码自注意力来让事情易于理解。实际上，更常见的是多头注意力甚至像分组多查询注意力这样的东西（如果你不熟悉这些术语，不用担心细节）。另外，我选择将跳跃连接、前馈网络、归一化等组件作为黑盒处理。同样，如果你觉得这些术语吓人，也别担心——我们用不到它们。这篇博客专门关注KV 缓存，我不想偏离这个主要目标。所以我尽量只关注相关的部分。 按 Enter 或点击查看完整图片 解码器直观图示 解码器接收一些词元作为输入，并基于此尝试预测序列中下一个最可能出现的词元。它会不断重复这个过程，直到生成整个输出序列。例如，如果我们向解码器输入序列“the cat sat on the”（猫坐在……上），它很可能会预测出“mat”（垫子），前提是模型经过适当训练。现在让我们看看这个过程是如何发生的： 输入嵌入作为输入传递给解码器。但这些嵌入有一个问题。输入嵌入捕捉了它们所代表的单词的含义，但没有捕捉到该单词与序列中其他单词的关系或使用它的上下文。例如，考虑以下句子： Only Saad wants coffee.（只有萨德想要咖啡。）⇒ 这意味着只有萨德想要咖啡——其他人都不想要。 Saad only wants coffee.（萨德只想要咖啡。）⇒ 这意味着萨德只想要咖啡——他什么都不想要。 我们使用一种叫做注意力机制的东西（我们将在下一节中看到）将这些输入嵌入转化为上下文化嵌入。上下文化嵌入意味着每个词元都知道它在句子中的上下文以及与其他词元的关系。之后我们做一些“东西”，然后得到这些原始输出，称为logits。这些 logits 然后被转换为概率，模型基于这些概率生成下一个词元。所以这是对解码器非常基础、直观的观察方式。 > 现在我知道你们一定还会有疑问：“等等，刚才发生了什么？我们是怎么从输入嵌入变成上下文化嵌入的？大家一直在说的‘注意力’是什么？”那么女士们先生们——我们开始了。 ## 掩码自注意力 现在是时候看看注意力机制了。我会尽量让它直观易懂。首先，让我们理解输入是什么样子。我们知道输入序列中的每个词元都被转换为一个具有若干维度的嵌入（比如 d_model）。因此，如果输入中有 N 个词元，那么输入本质上是一个大小为 N x d_model 的矩阵。如果你不明白，别慌——我们来拆解一下。我们有 N 个嵌入，每个具有 d_model 维度。所以基本上，我们是把这 N 个具有 d_model 维度的嵌入堆叠在一起，得到一个大小为 N x d_model 的矩阵。 按 Enter 或点击查看完整图片 解码器的输入 对于每个输入嵌入，我们创建三个向量，称为查询向量、键向量和值向量。这是通过将输入嵌入与在模型训练期间学习到的矩阵 Wq、Wk 和 Wv 相乘来实现的。我们来了解一下查询和键向量的直觉。 想象序列中的所有词元都是人——第一个是 Saad，下一个是 Saad 的朋友，然后是 某个陌生人，然后是 另一个陌生人，然后是 Saad 的邻居。突然，Saad 喊道：“谁对我重要？”于是 Saad 的朋友回答：“我超级重要！”然后 Saad 的邻居说：“虽然我觉得我没那么重要，但我们确实认识——如果这算数的话。”当然，Saad 不在乎那两个陌生人。所以在这里，进行“看”的那个人，即 Saad，就是查询。被看的人——所有人，包括 Saad（这可能感觉奇怪，但在自注意力中，查询甚至也看自己）——就是键。 按 Enter 或点击查看完整图片 查询与键的直觉 在这个图中，我们可以看到： - Saad 和他的朋友关系很强。 - 死对头也与 Saad 关系很强，但方向相反。 - Saad 就是 Saad，所以是的，他与自己相关。 - NPC 们只是背景板。 - 邻居在某种程度上与 Saad 相关，但不如朋友或死对头那么强。 ### 自注意力 现在我们已经有了点直觉，让我们进入真正的数学部分。正如我们在本节开头看到的，我们将输入 X 与矩阵 Wq、Wk 和 Wv 相乘，得到查询、键和值。当 X（形状为 N x d_model）与 Wq（形状为 d_model x d_k）相乘时，我们得到矩阵 Q（形状为 N x d_k），其中每一行对应于特定词元的查询向量（具有 d_k 维度）。当 X 与 Wk（形状 d_model x d_k）相乘时，我们得到矩阵 K（形状 N x d_k），其中每一行是对应词元的键向量（具有 d_k 维度）。当 X 与 Wv（形状 d_model x d_v）相乘时，我们得到矩阵 V（形状 N x d_v），其中每一行是对应词元的值向量（具有 d_v 维度）。 按 Enter 或点击查看完整图片 计算查询、键和值 > 现在我们要看一下注意力的公式。它可能看起来令人生畏，但别担心，我们会一步一步来： 注意力公式 首先，我们看一下以下矩阵：- 所以在将 Q 与 K 的转置相乘后，我们得到一个包含原始注意力分数的矩阵。这些原始注意力分数告诉我们一个词与另一个词之间的关系有多强。让我们用一个例子来理解。考虑输入序列是“Saad is drinking coffee”（Saad 正在喝咖啡）。在下面的图中，我试图解释这个矩阵的每个元素是什么意思以及它是如何计算的。之后，我们将每个原始注意力分数除以 d_k 的平方根，这可以确保更稳定的训练。现在我们来谈谈 softmax 函数，它看起来像这样： Softmax 函数 乍一看可能吓人，但实际上很简单。它的作用是将我们的原始注意力分数转换为 0 到 1 之间的值，使得每行中的所有值加起来为 1。得到的矩阵仍然包含注意力分数，仍然告诉我们一个词与另一个词之间的关系有多强，只不过现在值在 0 到 1 之间，并且每行之和为 1。 现在我们把这个矩阵与值矩阵相乘。让我们先建立这个步骤的直觉。序列中的某些词与特定词元关系很强，具有高注意力分数。其他词可能相关但没那么强，具有较低注意力分数。最后，有些词可能完全不相关，它们的注意力分数非常接近零（比如 0.0001）。当我们执行这个矩阵乘法时，我们实际上是在根据值对当前词元的重要性（相关性）来加权值向量： - 相关词元（高分数） 对最终输出向量贡献更多。 - 不相关词元（接近零的分数） 几乎不贡献。 让我们看看“Saad”的上下文化嵌入是如何沿着其中一个维度计算的。我们假设 d_v = 2，即值向量有 2 个维度——我们称这些维度为 d1 和 d2。在上面的图中，我展示了“Saad”沿着维度 d1 的上下文化嵌入的计算过程。 所以基本上，发生的事情是：以前，“Saad”的嵌入只包含关于“Saad”的含义和它在句子中的位置的信息。但现在，我们在衡量每个其他词元的重要性——并根据这些注意力分数，一些词元强烈影响最终嵌入（如果它们与“Saad”相关），而其他词元则几乎不贡献（如果它们与“Saad”不相关）。我们最终得到的是一个上下文化嵌入，其中每个词元都知道自己相对于序列中其他词元的上下文。 ### 掩码自注意力 现在我们已经理解了自注意力，让我们看看什么是掩码自注意力。我们不希望模型在预测下一个词时通过查看未来词元来*“作弊”。所以我们应用一个掩码来阻止它关注当前词元之后的词元。例如，在序列中：“pokemon gotta catch em all”（神奇宝贝 gotta catch em all） - 单词 “pokemon” 只能看自己，因为其他所有词元都是未来词元*——所以它们被掩码了。 - 单词 “gotta” 可以看 “pokemon” 和自己。 - 单词 “catch” 可以看 “pokemon”、“gotta” 和自己。 - 以此类推。 现在，我们如何实现这个掩码呢？方法如下：在应用 softmax 之前，我们将想要隐藏的值（即未来词元对应的位置）设置为 -infinity。当应用 softmax 时，这些值会变成 0，从而有效去除在注意力计算中来自掩码（未来）词元的任何贡献。 按 Enter 或点击查看完整图片 注意力掩码 所以现在，希望你已经对注意力的工作原理以及掩码自注意力如何防止模型窥视未来有了很好的直觉。这种设置——模型只关注过去和现在的词元——被称为因果注意力，使用这种设置的模型通常被称为因果语言模型。 > 我知道这很难消化，所以现在我们转向一些更有趣且相对简单的东西。 ## 语言模型中的推理 推理只是我们使用训练好的模型对某些数据进行预测的一种花哨说法。例如，如果你问一个 LLM“法国的首都是什么？”，模型回答“巴黎”，这就是一次推理。现在，让我们看看语言模型中的推理是如何工作的。我相信你已经熟悉词元了。通常，在使用语言模型时，我们会使用一些特殊词元。特别地，有两个重要的词元需要知道： - -**句子开始** - -句子结束 词元被添加到输入序列的开头，而 用于表示生成结束，这样模型就知道何时停止产生输出。 语言模型使用仅解码器架构，我们已经大致知道它预测序列中的下一个词元。现在让我们看看这实际上是如何发生的。假设用户输入：“I am”（我） 首先，这个序列被词元化。然后，我们在开头添加一个 ``（句子开始）词元，所以输入序列变成：" I am" 接下来，每个词元被转换为嵌入。我们将这个序列传递给模型。模型计算注意力并进行更多“东西”。然后它输出称为 logits 的原始输出。我们取序列中最后一个词元的 logits，并对其应用 softmax。这给出了整个词汇表上的概率。然后模型选择概率最高的词元——假设它是“drinking”（正在喝）。这个词元现在被追加到输入序列中。现在输入变成：" I am drinking" 我们将这个更新后的序列再次输入模型。同样的过程发生。假设模型现在预测“coffee”（咖啡）。我们再次追加它，所以输入