【论文笔记】Keyformer-CSDN博客

本文链接：https://blog.csdn.net/qq_51178589/article/details/145935768

Keyformer: KV缓存通过关键标记选择实现KV缓存减少以提高生成性推理效率

Muhammad Adnan*¹ Akhil Arunkumar² Gaurav Jain²
Prashant J. Nair¹ Ilya Soloveychik² Purushotham Kamath²

¹电气与计算机工程系，不列颠哥伦比亚大学，温哥华，BC省，加拿大
²d-Matrix，圣克拉拉，加利福尼亚州，美国。
通讯作者：Muhammad Adnan adnan@ece.ubc.ca
本文为2024年第七届MLSys会议论文集，会议地点为美国加利福尼亚州圣克拉拉。版权归作者所有。

摘要
Transformer架构已成为大型语言模型（LLMs）的基础架构。在生成性语言模型中，推理过程包含两个主要阶段：提示处理和标记生成。标记生成占据了大部分计算工作量，主要涉及向量-矩阵乘法以及与键值（KV）缓存的交互。由于需要从存储系统向计算单元传输权重和KV缓存值，这一阶段受到内存带宽的限制。这种内存瓶颈在需要长文本上下文和广泛文本生成的应用中尤为明显，而这些应用对LLMs的重要性正日益增加。

本文介绍了一种名为“Keyformer”的创新性推理时方法，旨在缓解KV缓存大小带来的挑战。Keyformer利用了一个观察结果：在生成性推理中，约90%的注意力权重集中在一组特定的标记上，这些标记被称为“关键”标记。Keyformer通过一种新颖的评分函数识别这些重要标记，并仅在KV缓存中保留关键标记。这种方法在不降低模型准确性的情况下，减少了KV缓存的大小和内存带宽的使用。我们在三种基础模型（GPT-J、Cerebras-GPT和MPT）上评估了Keyformer的性能，这些模型采用了不同的位置嵌入算法。我们的评估涵盖了多种任务，重点关注涉及长上下文的总结和对话任务。我们证明，Keyformer将推理延迟降低了2.1倍，将标记生成吞吐量提高了2.4倍，同时保持了模型的准确性。

1. 引言

Transformer架构已在语言建模（Lewis et al., 2019; Brown et al., 2020; Raffel et al., 2020）、图像识别（Dosovitskiy et al., 2020）、推荐系统（Sun et al., 2019; de Souza Pereira Moreira et al., 2021; Adnan et al., 2023; Zhao et al., 2023）以及随着大型语言模型（LLMs）的出现而进行的文本生成任务中取得了特别的成功。不幸的是，LLMs的部署提出了关键的推理延迟和吞吐量问题。这主要是由于生成性推理的顺序自回归特性，尤其是在处理更大上下文的输入时。尽管取得了进展，但现代LLMs在高效处理长输入序列方面仍面临挑战，这一事实已得到最近研究的证实（Bai et al., 2023; Li et al., 2023; Chen et al., 2023; Huang et al., 2021a）。不幸的是，与更长序列相关的增加的内存和计算需求加剧了LLMs的推理延迟，并降低了吞吐量。本文提出了推理时策略，用于构建保持准确性、低延迟、高吞吐量的LLMs系统。

LLMs采用Transformer架构，并依赖于“注意力机制”来理解给定输入序列中单词之间的关系（Vaswani et al., 2017）。由于注意力机制的计算复杂度与输入序列长度呈二次方关系，因此在推理过程中会带来最大的延迟开销（Sukhbaatar et al., 2019; Dao et al., 2022; Choromanski et al., 2020）。此外，由于LLMs中标记生成的自回归特性，需要为所有先前的标记重新计算键和值向量。为了缓解这一问题，LLMs使用了一种称为键值缓存（KV缓存）的存储结构（Ott et al., 2019）。KV缓存保留了先前计算的键值对，从而消除了重新计算这些向量的高昂成本。

然而，KV缓存带来了可扩展性挑战。在标记生成过程中从片外内存访问KV缓存会引入额外的内存延迟，并受到内存带宽限制。例如，在图1(a)中所示的MPT-7B模型中，将序列长度增加16倍（从512增加到8K）会导致推理延迟增加超过50倍。此外，大约40%的总推理时间（以绿色突出显示）被KV缓存数据移动所消耗。重要的是，更大的上下文不仅增加了KV缓存的大小，还延长了其他操作所需的时间（以蓝色表示）。同样，如图1(b)所示，对于MPT-7B模型，当序列长度超过8K时，KV缓存的大小超过了模型本身的大小。因此，KV缓存大小成为了实现长序列低延迟、高吞吐量推理的障碍。

先前的研究已经探索了在处理更长序列时缓解注意力机制的内存和计算需求的方法（Zaheer et al., 2020; Kitaev, 2020; Wang et al., 2020; Beltagy et al., 2020）。尽管系统级优化方法如FlexGen（Sheng et al., 2023）、Flash Attention（Dao et al., 2022）、Paged Attention（Kwon et al., 2023）和多维分区（Pope et al., 2023）旨在提高生成性AI的可扩展性，但它们通常忽略了KV缓存大小这一根本挑战。多查询（Shazeer, 2019）和分组查询注意力（Ainslie et al., 2023）等技术通过减少特定注意力头对KV缓存的写入来减少KV缓存大小，但这些方法通常需要资源密集型的模型重新训练或微调。这在各种加速器已经在现场部署的情况下变得复杂。因此，迫切需要推理时KV缓存减少技术。这更加具有挑战性，因为任何提出的技术都必须符合模型准确性的严格约束。例如，MLPerf（Reddi et al., 2020）要求对LLMs应用的任何优化都必须保持模型准确性在基线的99%到99.9%之间。

为了解决这些问题，我们引入了Keyformer¹，这是一种在推理时动态减少KV缓存大小的新方法。Keyformer通过智能地丢弃不必要的标记而不损失准确性来实现这一点。Keyformer的关键见解如图2所示，我们还在其中与现有的最先进的推理优化方法进行了比较。图2(a)展示了传统的“全注意力”（Brown et al., 2020）机制，其中每个新生成的标记都会关注序列中所有前面的标记。图2(b)描述了“窗口注意力”（Child et al., 2019），它维护一个固定大小的最近标记滑动窗口，从而减少了KV缓存的大小。然而，这种方法限制了模型从过去捕获全面语义信息的能力，导致文本生成质量下降和准确性降低。图2©展示了一个变体，称为“扩张窗口注意力”，其准确性限制与窗口注意力相似。

为了应对这一挑战，Keyformer利用了一个洞见：某些标记比其他标记更重要。具体而言，它观察到约90%的注意力权重集中在一小部分被称为关键标记的标记上。这些标记对于LLMs理解上下文至关重要，但可能会超出窗口注意力的滑动窗口范围。Keyformer引入了一种混合注意力方法，如图2(d)所示，该方法在生成下一个标记时结合了最近的标记和前面的关键标记。我们的实验表明，与H2O（Zhang et al., 2023）等最先进的方法相比，Keyformer表现出显著的改进。这是因为与仅根据注意力分数识别“重量级选手”的H2O不同，Keyformer在识别关键标记时考虑了丢弃标记的重要性。

我们在多种模型（包括GPT-J（Wang & Komatsuzaki, 2021）、Cerebras-GPT（Dey et al., 2023）和MPT（Team et al., 2023））上评估了Keyformer，涵盖了总结和对话等多种长序列任务。即使KV缓存减少了50%，Keyformer也能保持准确性，同时将推理延迟降低了2.1倍，并将标记生成吞吐量提高了2.4倍。

¹ https://github.com/d-matrix-ai/keyformer-llm

2. 背景与动机

2.1 大型语言模型中的推理过程

在语言建模中，任务是基于前面的标记 $x_1, x_2, \ldots, x_n$ 来估计下一个标记的概率。对于生成性大型语言模型（LLMs），推理过程分为两个阶段：

提示处理阶段：此阶段对输入上下文进行因果处理，使模型能够为上下文中的所有标记生成键和值。这些键值对随后被存储在KV缓存中。
标记生成阶段：此阶段依次且自回归地生成文本。每个标记通过模型的所有层生成，值得注意的是，下一个标记的生成依赖于之前生成的标记及其顺序。

为了提高推理效率，避免在所有层中重复且复杂地计算键（K）和值（V）张量，这些张量被缓存起来，这被称为KV缓存。KV缓存会在每次标记生成步骤中依次填充（Strati et al., 2024），直到文本生成过程完成。

2.2 通过利用稀疏性减少KV缓存大小

为了应对KV缓存不断扩大的挑战，我们考虑一个由n个标记组成的序列 $S_n$ 及其针对单个注意力头和层的KV缓存内容。在全注意力机制中，KV缓存组件包含 $n$ 个键和值，这些键和值与 $S_n$ 成比例增长。为了缓解这一问题，我们可以将KV缓存的大小缩小到适合较短序列的大小，记为 $S_k$ 。这涉及到使用较少的标记，从 $n$ 个标记减少到 $k$ 个标记，其中 $S_k$ 是 $S_n$ 的一个子集，且 $k < n$ 。这种减少可以通过利用LLMs注意力机制中的固有稀疏性来实现。

尽管在Transformer训练过程中存在巨大的计算需求，但其注意力机制中存在固有的稀疏性。然而，稀疏性的程度可能因特定的下游任务而异。图3a展示了不同模型在使用CNN/DailyMail数据集进行总结任务时的注意力稀疏性水平差异。这种差异体现在模型的不同层面，包括整个模型、各个层以及模型的不同部分。

2.3 通过使用关键标记提高性能

在图3b中，累积分布函数（CDF）描述了注意力分数与总上下文的比例之间的关系。值得注意的是，在文本生成过程中，一小部分标记获得了最多的注意力权重。这突显了特定关键标记的重要性，以及它们在理解上下文和促进文本生成中的关键作用。然而，动态确定哪些标记是关键标记，尤其是在输入序列中包含在推理过程中未知或未见过的标记时，是一个相当大的挑战。

2.3.1 利用评分函数识别关键标记

我们为每个标记引入一个评分函数 $f_\theta$ ，以识别 $n$ 个标记中的 $k$ 个关键标记。在多头注意力机制中，注意力分数决定了单个标记与其他所有标记之间的连接程度，如公式（1）所示：

$\text{注意力分数} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) \quad (1)$

自然的选择是利用注意力分数作为评分函数，记为 $f_\theta(\text{acc}; \text{attn})$ 。这种方法在先前的最先进工作中很常见，例如H2O（Zhang et al., 2023）。它将那些在提示和标记生成阶段持续获得较高注意力分数的标记识别为最重要或关键的标记。

我们可以根据它们的累积注意力分数选择并保留这些 $k$ 个标记，从而创建我们所称的“关键注意力”。然而，仅依赖于这些 $k$ 个标记进行注意力操作无法提供必要的准确性，导致性能不佳。这一点在图3c中有所体现。

在这一比较中，“窗口注意力”和“关键注意力”在性能上均劣于全注意力，即使在将窗口和关键标记参数减少一半（ $n /2$ ）时也是如此。尽管减少窗口和关键标记的大小对于最小化KV缓存的大小至关重要，但这也导致了准确性的显著下降。这种下降主要是由于关键标记注意力中最近上下文的丢失以及窗口注意力中重要上下文的丢失。基于这一观察，我们提出了一种混合方法，将选定的关键标记与最近标记结合起来，以减少KV缓存的大小，同时保留准确性。

2.3.2 问题：分数分布不均匀

图4展示了全注意力的注意力分数分布（ $f_\theta$ ），如公式（2）所述。当KV缓存减少时，分数较低的标记被丢弃。这改变了评分函数，如公式（3）所示，因为项 $\sum_{m=n-k}^{n} e^{x_m}$ 变为零。

$f_\theta(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}, \quad i = 1, 2, \ldots, n \quad (2)$

$f_\theta(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{k} e^{x_j} + \sum_{m=n-k}^{n} e^{x_m}} \quad (3)$

这种标记的移除破坏了评分函数的分布。这是因为丢弃的标记的注意力权重不均匀地分配给了减少的KV缓存中的标记。这种不均匀分布是由于固有的softmax函数的性质引起的。

图4通过比较全注意力的评分函数分布和减少KV缓存后的分布来说明这一现象。当评分分布不均匀时，模型可能无法关注序列中最具相关的标记。因此，这可能导致上下文信息的丢失、准确性降低以及潜在的文本生成质量下降。

2.3.3 动机：衰减评分函数

一种直接的方法是通过衰减因子来衰减评分函数，以抵消由于丢弃标记而导致的多余注意力分数。假设 $\alpha$ 是衰减因子，它将评分函数修改为 $\bar{f}_\theta = \alpha f_\theta$ 。理想情况下，我们希望将评分函数衰减到一个与 $\sum_{m=n-k}^{n} e^{x_m}$ 相当的因子，其中 $n - k$ 表示被丢弃的标记。

图3b显示，即使KV缓存大小减少了50%，关键标记的平均累积注意力分数仍然保持在较高水平，大约在90%到95%之间。我们在Cerebras-GPT-6.7B（Dey et al., 2023）模型上进行了范围值扫描，以探索不同衰减因子（ $\alpha$ ）对整体模型质量的影响。这一分析是在KV缓存大小设置为50%且最近比率（表示最近生成标记的百分比）为20%的情况下进行的。

然而，如图5所示，即使应用了衰减因子，也无法达到全注意力模型的相同质量。这种差异是由于减少的KV缓存中剩余标记的评分分布发生了显著变化。这些发现强调了仅依赖于基于累积注意力分数的评分函数 $f_\theta(\text{acc}; \text{attn})$ 来识别关键标记的不足。因此，解决评分函数中丢弃标记的影响对于实现更高的模型准确性或满足MLPerf等基准的准确性要求至关重要。

3. Keyformer：直觉与设计

Keyformer通过利用解码器层中的固有稀疏性，使用最近标记的混合方法来识别关键标记，并通过应用正则化来调整丢弃标记对评分函数的影响，从而在推理时动态减少KV缓存的大小。

3.1 概率分布正则化

我们在提示处理阶段战略性地从上下文中移除 $n - k$ 个标记，以在生成过程中保持一个包含 $k$ 个标记的恒定KV缓存大小，防止不必要的内存扩展。此后，Keyformer采用概率分布正则化技术。通过向未归一化的logits中引入额外的分布（ $\zeta$ ），使模型在推理时能够适应未知上下文，从而保持鲁棒性和适应性。

Keyformer向由 $QK^T$ 得到的未归一化logits中添加噪声，如公式（4）所示：

$y_i = x_i + \zeta_i, \quad \text{其中} \quad x_i = \frac{Q[i, :]K[:, i]^T}{\sqrt{d_k}} \quad (4)$

这里， $y_i$ 是调整后的logits， $x_i$ 是未归一化的logits， $\zeta_i$ 是为正则化而添加的分布。

3.2 正则化分布的选择

添加到未归一化logits中的正则化分布对关键标记的识别和模型质量有显著影响。因此，我们试图利用LLMs的语义学来选择合适的分布。

3.2.1 直觉：偏向初始标记

先前的研究（例如流式LLMs [Xiao et al., 2023] 和H2O模型 [Zhang et al., 2023]）表明存在对初始标记的偏向。这种偏向源于累积注意力分数倾向于偏向初始标记，因为它们在解码迭代过程中具有累积效应。我们提出使用一种偏斜分布来利用这种偏向，并有效地模拟最大值（关键标记）的分布。这种分布偏向初始标记，同时保持不对称的轮廓，增强从最近上下文窗口中提取的标记的表示。

3.2.2 Gumbel分布的引入

我们的分布选择受到Gumbel分布的启发。Gumbel分布特别适合于我们的关键标记识别任务，因为它描述了一组样本中最大值的分布，并且偏向初始标记。这使其成为模拟长序列中关键标记的理想选择。

Gumbel分布的概率密度函数（pdf）如下：

$f_{\text{Gumbel}}(\zeta_i) = e^{-\zeta_i - e^{-\zeta_i}} \quad (5)$

调整后的logits的pdf为：

$f_{\text{Gumbel}}(y_i) = e^{-(y_i - x_i) - e^{-(y_i - x_i)}} \quad (6)$

此外，Gumbel分布在统计理论中具有重要意义，它体现了Gumbel极限定理，即常见的概率分布（如正态、指数、均匀等）会收敛到Gumbel分布。

3.2.3 Keyformer评分函数的均匀性

理论上，选择一个在归一化后促进均匀性的正则化分布有助于识别关键标记。为了量化归一化后概率分布的均匀性，我们使用熵（ $-\sum p_i \log(p_i)$ ）作为衡量标准。我们的分析表明，基于Gumbel的logit调整有助于实现更均匀的分布，从而证明了其作为关键标记识别正则化技术的有效性。

3.3 Keyformer评分函数

我们提出了一种新的Keyformer评分函数 $f_\theta(\text{Keyformer})$ ，以解决基于累积注意力的评分函数（ $f_\theta(\text{acc}; \text{attn})$ ）的局限性。新的评分函数将Gumbel噪声分布整合到未归一化的logits中，但未能考虑丢弃标记对底层概率分布的影响。为了解决这一问题，我们引入了一个温度参数（ $\tau$ ），如公式（9）所示：

$f_\theta(x_i) = \frac{e^{(x_i + \zeta_i)/\tau}}{\sum_{j=1}^{k} e^{(x_j + \zeta_j)/\tau}}, \quad i = 1, 2, \ldots, k \quad (9)$

这种概率评分函数类似于Gumbel Softmax，它为离散随机变量提供了连续的松弛。这与我们识别过去标记的一个子集 $S_k \subset S_n$ 的主要目标一致，这些标记能够传达与原始完整标记集相同的语义信息。

3.3.1 温度参数（ $\tau$ ）的重要性

温度参数（ $\tau$ ）在调节概率分布的平滑度方面起着关键作用。当 $\tau$ 的值较高（ $\tau \to \infty$ ）时，概率趋于均匀，为所有标记分配相等的分数。相反，当 $\tau$ 的值较低（ $\tau \to 0$ ）时，概率分布更加尖锐，根据它们的未归一化logits优先考虑特定标记。这个参数在丢弃标记时至关重要，因为它们不能在不重新计算其键和值的情况下重新引入。

在公式（10）中，我们展示了 $\tau$ 在每个解码迭代 $t$ 中的动态性质。为此，我们定义了一个从 $\tau_{\text{init}}$ 到 $\tau_{\text{end}}$ 的 $\tau$ 范围。在每个解码步骤中，我们通过 $\Delta \tau$ 增加 $\tau$ ，其值由 $\tau$ 的范围和正在生成的文本长度 $T$ 决定：

$\tau = \tau_{\text{init}} + t\Delta \tau, \quad \Delta \tau = \frac{\tau_{\text{end}} - \tau_{\text{init}}}{T} \quad (10)$

这种策略基于这样一个前提：随着更多标记被丢弃，我们需要一个更均匀或更随机的概率分布。通过实证分析，我们发现将 $\tau_{\text{init}}$ 设置为1， $\tau_{\text{end}}$ 设置为2可以产生最佳结果。这一决定与我们在提示阶段保持非随机评分函数的目标一致，此时所有标记都可用。当 $\tau$ 设置为1时，Gumbel Softmax方法几乎等同于标准Softmax。随着我们进入解码迭代并丢弃更多标记以保持恒定的KV缓存大小，我们通过 $\Delta \tau$ 系统地增加评分函数 $f_\theta$ 中的随机性。

3.3.2 利用评分函数的累积

评分函数的累积对于根据其在解码步骤中的一致行为识别关键标记至关重要。如果没有累积，标记的丢弃将仅依赖于当前标记与前面标记的相关性。尽管当前标记的相关性在识别关键标记中很重要，但它们的行为应在大多数生成的标记中保持一致。为了根据这种一致行为识别关键标记，我们在提示和标记生成阶段累积评分函数（ $f_\theta$ ），如图6所示。

3.4 Keyformer算法

图6展示了Keyformer的概述。我们强调了其在基于稀疏化的标记丢弃、最近标记和关键标记的混合使用，以及基于Gumbel Softmax的关键标记识别评分函数方面的关键功能。在提示处理阶段，Keyformer计算提示长度 $S_n$ 中所有 $n$ 个标记的键和值，以预测第一个标记。鉴于KV缓存预算，Keyformer保留一个包含 $w$ 个最近标记的最近窗口，同时从 $n - w$ 个标记窗口中丢弃 $n - k$ 个标记，从而识别出 $k - w$ 个标记。根据Keyformer评分函数，从 $n - w$ 个标记窗口中选择排名前 $k - w$ 的标记。关键标记（ $k - w$ ）和最近标记（ $w$ ）的组合构成了减少的KV缓存。由于在提示处理阶段没有丢弃标记，Keyformer使用温度参数 $\tau_{\text{init}} = 1$ 来近似Softmax概率分布。这在解码步骤1中有所体现。在标记生成阶段，Keyformer使用减少的KV缓存进行操作。第一个生成的标记仅关注KV缓存中的 $k$ 个标记，如解码步骤2所示。最近窗口 $w$ 向右移动一个标记，而评分函数 $f_\theta$ 与前一个解码步骤的评分函数累积。在标记生成阶段的每个解码步骤中，从大小为 $k + 1 - w$ 的窗口中识别出 $k - w$ 个关键标记。由于我们在“关键标记”窗口中添加和移除标记，我们可以在保持恒定KV缓存大小的同时提高准确性。此外，随着更多标记被丢弃，温度参数 $\tau$ 通过 $\Delta \tau$ 增加，以调整评分函数的概率分布。Keyformer的详细算法如算法1所示。

4. 评估

我们在三种重要的模型家族中对Keyformer进行了评估：GPT-J（Wang & Komatsuzaki, 2021）、Cerebras-GPT（Dey et al., 2023）和MPT（Team et al., 2023）。这些模型分别采用了不同的位置编码技术。GPT-J使用了RoPE（Su et al., 2022），Cerebras-GPT使用了可学习的位置嵌入，而MPT使用了ALiBi（Press et al., 2021）。通过包含具有不同位置编码方法的模型，我们确保了我们的发现具有代表性和普适性。所有评估均使用固定的beam size为4。

4.1 实验设置

我们在两个代表性文本生成任务上进行了评估：使用CNN/DailyMail（See et al., 2017）和GovReport（Huang et al., 2021b）数据集的总结任务，以及使用SODA数据集（Kim et al., 2022）的对话任务。GPT-J模型针对总结任务进行了微调，而Cerebras-GPT和MPT是预训练模型。我们使用了MPT-chat版本的MPT模型进行对话任务，该模型针对对话生成进行了微调。所有模型均使用2k的序列长度进行预训练。

为了处理长文档总结，我们使用了MPT-storywriter版本的MPT模型，该模型针对虚构故事写作进行了微调。该模型支持65k的上下文长度，并能生成长达84k标记的内容。此外，我们还评估了来自lmeval-harness框架（Gao et al., 2021）的四个任务：PIQA（Bisk et al., 2020）、Winogrande（Sakaguchi et al., 2021）、OpenBookQA（Mihaylov et al., 2018）和COPA（Roemmele et al., 2011）。这些任务涉及自回归语言模型的少量样本评估，使用NVIDIA A100（80GB）GPU进行执行。

4.2 基线方法

为了评估Keyformer的准确性，我们将其与**全注意力（Full Attention）进行了比较。全注意力作为我们的基准，代表了准确性的黄金标准。我们的目标是实现与全注意力相当的准确性，目标范围为99%到99.9%。这一目标与行业基准测试机构（如MLPerf）设定的高质量标准一致（Reddi et al., 2020）。此外，我们还与窗口注意力（Window Attention）**和最近的H2O模型（Zhang et al., 2023）进行了比较，将KV缓存大小从提示长度的20%调整到90%。

4.3 准确性结果

为了评估KV缓存减少对文本生成质量的影响，我们依赖于ROUGE分数（Lin, 2004），这是一个广泛使用的评估指标，用于衡量流畅性和连贯性。ROUGE通过计算生成文本与参考文本之间的n-gram重叠来提供标准化的文本质量衡量。根据MLPerf，ROUGE分数（包括ROUGE-1、ROUGE-2和ROUGE-L）应达到其原始值的99%到99.9%，以满足总结任务的要求。因此，即使在减少KV缓存的情况下，我们的模型也应保持所需的ROUGE分数。图7展示了Keyformer与其他方法（全注意力、窗口注意力和H2O）在不同KV缓存大小下的准确性比较。图中重点关注了ROUGE-2分数，该分数衡量了二元组的重叠。ROUGE-1和ROUGE-L的趋势在附录A.5中进行了详细说明。

结果突显了先前上下文对模型性能的重要性。例如，窗口注意力仅依赖于最近的标记，导致显著的准确性损失。因此，识别关键标记对于实现所需的模型准确性至关重要。在各种KV缓存预算下，Keyformer始终优于最先进的H2O。这表明Keyformer识别的关键标记比H2O识别的“重量级选手”更重要。例如，Keyformer仅使用70%的KV缓存就达到了目标ROUGE分数，而H2O即使在更大的KV缓存预算下也未能达到这一目标。此外，Keyformer在某些情况下甚至超过了全注意力的基线准确性，最高可达1.73%（总结任务中GPT-J-6B为0.9%，Cerebras-GPT-6.7B为1.73%）。这证明了Keyformer评分函数中引入的Gumbel噪声的正则化效应及其对关键标记识别的积极影响。

4.3.1 长上下文总结

我们还评估了Keyformer在处理长上下文时减少KV缓存的有效性，同时保持准确性。这一评估使用了预训练的MPT-7B-storywriter模型，其上下文长度为65k。我们使用了政府报告（GovReport）数据集，该数据集包含由政府研究机构撰写的报告，具有更长的总结和文档。该数据集需要对上下文有深入的理解，以便提取关键信息以进行总结。图8展示了Keyformer、H2O和全注意力之间的准确性比较。值得注意的是，即使KV缓存大小仅为50%，Keyformer也能保持所需的99%准确性阈值，而H2O在同一KV缓存大小下的准确性显著较低。

4.4 性能结果

为了评估Keyformer在减少KV缓存后的性能优势，我们考虑了两个关键的推理指标：推理延迟和目标模型的标记生成吞吐量。我们的Keyformer实现无缝集成到Huggingface（Wolf et al., 2019）模型卡片中，确保易于采用。我们在模型和KV缓存超出GPU HBM内存容量的情况下禁用了CPU卸载，以确保评估的一致性。我们使用MPT-7B-storywriter模型生成了与每个提示等量的标记，并测试了不同的提示和生成长度组合。

图9展示了在保持准确性的条件下Keyformer的推理加速比，表1显示了与基于全注意力的方法相比，标记生成吞吐量的提升。在KV缓存减少50%的情况下，Keyformer显著降低了推理延迟，实现了2.1倍的加速比，并且在相同的batch size下处理了两倍的batch size，从而将标记生成吞吐量提高了2倍（相同batch size）和2.4倍（更大的batch size）。

4.4.1 性能提升分析

为了理解Keyformer带来的性能提升，我们主要考虑两个因素：

减少的KV缓存：较小的KV缓存显著减少了从GPU HBM片外内存的数据传输。
优化的点积计算：KV缓存中的标记数量从 $n$ 减少到 $k$ ，从而优化了多头注意力块中的点积计算（ $QK^T$ ）V。

值得注意的是，在以内存为瓶颈的LLMs中，性能提升主要来自于减少KV缓存的数据传输，而不是矩阵乘法。然而，Keyformer的关键标记识别过程引入了一些额外开销，主要是由于Gumbel Softmax。图10展示了Keyformer在50% KV缓存减少情况下的性能提升，包括Gumbel Softmax的额外开销。结果表明，Keyformer将KV缓存数据传输减少了2.9倍，并在4k序列长度的注意力模块的点积计算中提高了1.3倍的计算效率。

4.5 少样本评估

我们使用lm-eval-harness框架中的四个任务和预训练模型进行了少样本实验，以评估Keyformer在推理时不同样本数量下的性能。表2展示了0-shot和5-shot的结果，表明Keyformer在所有任务和样本设置中始终优于先前的方法。即使KV缓存大小减少了50%，它也能实现接近全注意力基线的准确性。

4.6 消融研究

4.6.1 共享与每层评分函数

评分函数 $f_\theta$ 定义了上下文中的关键标记。在具有堆叠解码器层的生成性LLMs中，评分函数 $f_\theta$ 可以在所有层之间共享（共享），或者为每一层单独分配（每层）。在每层方法中， $f_\theta(\text{每层})$ 为每个解码器层分配一个独立的评分函数，并在每个解码阶段进行累积。相反， $f_\theta(\text{共享})$ 使用一个全局评分函数用于所有解码器层，并在解码器层和解码阶段进行累积。

表3展示了在保持原始位置信息和KV缓存大小不变的情况下，每层和共享评分函数之间的准确性比较。值得注意的是，使用每层评分函数的准确性更高。这与直觉一致，即Transformer在不同层之间学习层次化的文本表示，较低层捕获局部语法和语义特征，而较高层捕获更抽象和复杂的模式（Geng et al., 2023）。因此，为每一层配备独立的评分函数有助于识别特定于该层的关键标记。

4.6.2 新旧位置信息的比较

我们研究了减少KV缓存大小对Keyformer中键的位置信息的影响。我们探索了两种方法：Keyformer（原始位置）和Keyformer（新位置）。在Keyformer（原始位置）中，位置信息反映了文本中标记的原始位置。相反，Keyformer（新位置）使用基于KV缓存中标记新排列的位置。

表3展示了在保持一致的KV缓存大小和评分函数的情况下，不同方法的准确性比较。值得注意的是，当使用原始位置信息时，Keyformer在准确性方面表现出色。然而，引入新的位置信息会导致准确性略有下降。尽管如此，即使使用新的位置信息，Keyformer仍然优于最先进的H2O。

4.6.3 不同分布的比较

我们进行了消融研究，以评估不同的logit调整分布对模型准确性或关键标记识别的影响。我们评估了三种正则化策略，并将其与基于Gumbel的logit调整进行了比较。这些策略包括：不进行logit调整、常数logit调整和基于高斯分布的logit调整。

不进行logit调整：为了研究省略正则化对未归一化logits的影响，我们进行了不进行logit调整的实验，其中 $y_i = x_i + \zeta_i$ ，类似于H2O（Zhang et al., 2023）中的方法。
常数logit调整：为了研究对所有未归一化logits进行常数正则化的影响，我们进行了常数logit调整实验，设置 $y_i = x_i + c$ ，其中 $c$ 是添加到每个未归一化logit中的常数。
高斯logit调整：我们还使用了对称的高斯分布进行logit调整。高斯概率密度函数（pdf）如下：

$f_{\text{高斯}}(\zeta_i) = \frac{1}{\sqrt{2\pi\sigma^2}} \exp\left(-\frac{(\zeta_i - \mu)^2}{2\sigma^2}\right) \quad (11)$

$f_{\text{高斯}}(y_i) = \frac{1}{\sqrt{2\pi\sigma^2}} \exp\left(-\frac{(y_i - (x_i + \mu))^2}{2\sigma^2}\right) \quad (12)$

表4提供了基于不同logit调整方法的实证评估结果。我们使用了标准的Gumbel pdf（ $\mu = 0.5772$ 和 $\sigma = 1.2825$ ），并将其与具有相同均值和方差的高斯pdf进行了比较。常数logit调整使用了 $c = 0.5772$ 。结果表明，Gumbel分布（以其对初始标记的偏斜性而闻名）是一种有效的关键标记识别正则化机制。

4.6.4 最近窗口与关键标记窗口比例

我们进行了敏感性研究，以研究最近标记比例 $w$ 对KV缓存大小的影响。这导致了关键标记数量（ $k - w$ ）的变化。附录A.4中的结果表明，当最近标记比例 $w$ 落在20%到30%的范围内时，模型表现更好。这一观察结果与我们的假设一致，即最近和关键标记对于LLMs推理至关重要。

4.6.5 与注意力池的比较

最近的研究引入了StreamingLLM（Xiao et al., 2023），它引入了“注意力池”的概念。StreamingLLM使语言模型（LLMs）能够在没有微调的情况下处理无限序列长度。这是通过保留前四个标记（称为“注意力池”）和一个最近标记的移动窗口来实现的，其中 $w = k - 4$ 。为了比较StreamingLLM和Keyformer，我们保持了60%的KV缓存大小。表3展示了准确性比较，结果表明StreamingLLM在仅依赖前四个标记作为注意力池和剩余标记来自最近窗口的情况下，在总结任务中表现不佳（附录A.7）。

5. 相关工作

5.1 注意力加速

以往的研究主要集中在提高基于Transformer（Vaswani et al., 2017）模型的推理速度。PoWER-BERT（Goyal et al., 2020）利用词向量消除技术，通过利用编码器模型中的冗余来加速推理。Linformer（Wang et al., 2020）尝试将注意力机制从二次复杂度降低到线性复杂度。Reformer（Kitaev, 2020）通过局部敏感哈希（LSH）减少注意力复杂度。Linear Transformers（Katharopoulos et al., 2020）通过存储累积状态而不是保留每个表示来优化计算。FLAT（Kao et al., 2023）提出了优化的数据流方法，而其他研究（Wang et al., 2022）则通过重叠通信与依赖计算来提高注意力执行效率。相比之下，Keyformer旨在通过减少标记数量来减少KV缓存，从而加速注意力计算。

5.2 稀疏注意力

另一条研究路线是通过稀疏化注意力机制来减少注意力块的计算和内存需求。BigBird（Zaheer et al., 2020）结合随机、窗口化和全局注意力，在稀疏化注意力块的同时保持Transformer的准确性。LongFormer（Beltagy et al., 2020）也利用窗口化注意力和基于任务的局部注意力来实现稀疏注意力。Spatten（Wang et al., 2021）在头部和标记级别引入稀疏性。然而，这些方法需要专门的架构来利用稀疏性。此外，这些工作并未涉及推理优化。

5.3 KV缓存减少

El-Attention（Yan et al., 2021）通过修改多头注意力模块来减少KV缓存大小，利用键和值在增量解码中的稳定性进行跨层重用。相比之下，H2O（Zhang et al., 2023）通过识别“重量级选手”并将它们保留在KV缓存中来减少其大小，但忽略了在移除先前标记后注意力分数分布的变化，从而导致准确性权衡。其他方法（Liu et al., 2023; Anagnostidis et al., 2023）在粗粒度和细粒度级别引入稀疏性，目标是在推理时消除特定的头部和标记。然而，这些方法需要针对特定任务的预测器和预训练模型的微调。另一种方法（Mu et al., 2023）通过将提示压缩为摘要标记来减少KV缓存。Landmark Attention（Mohtashami & Jaggi, 2023）通过在词汇表中引入额外的地标标记来表示标记块，需要计算密集型的重训练或微调以整合摘要或地标标记。

6. 未来工作

最近的技术，如多查询注意力（Multi-Query Attention, MQA）（Shazeer, 2019）和分组查询注意力（Group-Query Attention, GQA）（Ainslie et al., 2023），旨在通过减少注意力头的数量来训练基础模型。然而，这些模型通常在特定任务的微调之后才会被使用。尽管对Keyformer与这些模型的详细评估将推迟到未来工作，但值得注意的是，Keyformer仍然可以应用于基于MQA或GQA的模型之上。这是因为Keyformer无论头的数量多少，都会丢弃冗余标记。此外，我们计划通过将Keyformer集成到LLMs的注意力块中，用基于Keyformer的softmax替换标准softmax，从而在训练期间引入稀疏性，以解决Transformer的二次计算和内存复杂性问题。这一方向旨在在不牺牲准确性的前提下，提高对更长上下文的可扩展性。

7. 结论

大型语言模型（LLMs）的进步推动了对更长上下文和广泛文本生成的需求，这些模型在数百万标记的序列上进行训练。然而，这种趋势对系统内存带宽造成压力，导致执行成本增加。在更长的上下文中，KV缓存的大小——主要负责内存带宽消耗和推理延迟——超过了模型参数的大小。为了解决这一问题，我们提出了Keyformer，它通过丢弃跨头部、层和束的关键标记来有效减少KV缓存的大小，最多可达50%，而无需牺牲准确性。Keyformer可以在推理时应用于LLMs，无需微调，同时还能提高延迟和标记生成吞吐量。