Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning

  • 作者/机构: Yu Fu (University of California, Riverside), Zefan Cai (University of Wisconsin-Madison), Abedelkadir Asi (Microsoft), Wayne Xiong (Microsoft), Yue Dong (University of California, Riverside), Wen Xiao (Microsoft)

A1 主要贡献

本文针对大型语言模型（LLM）中常见的Key-Value（KV）缓存技术因输入长度增加而导致内存开销迅速增长的问题，提出了一种新的解决方案。先前的工作表明，并非所有token对文本生成都同等重要，并提出了层级的KV缓存压缩方法。本文认识到注意力头（attention heads）在生成过程中扮演着不同角色，因此提出了一个头级别（head-level）的KV缓存压缩方法，名为HeadKV，以及其改进版HeadKV-R2，该版本利用一种新颖的上下文推理能力评估方法进行压缩。

核心问题与研究目标：
随着LLM支持的输入长度越来越长（如GPT-4支持128K，Claude支持1M），自注意力机制导致的内存使用和延迟显著增加。KV缓存虽能提升推理速度，但其内存占用随输入长度线性增长，构成了巨大的存储和效率挑战。现有的KV缓存压缩方法大多在层级（layer-level）上进行，未能充分利用不同注意力头的重要性差异。本文的目标是设计一种更精细化的、在头级别进行KV缓存压缩的方法，从而在不牺牲性能的前提下优化内存使用。

创新点与主要贡献：
1. 提出HeadKV，一种头级别的KV缓存压缩方法：与之前在层级上分配KV缓存预算不同，HeadKV根据每个注意力头的重要性分布来分配预算。重要的头分配更多的KV缓存，不重要的头则分配更少。
2. 提出HeadKV-R2，一种集成了检索与推理能力评估的压缩方法：作者认为仅关注头的“检索”能力不足以应对需要长上下文的复杂任务（如上下文问答），这些任务同时需要检索和推理能力。因此，他们设计了一种新的重要性分数评估方法，通过构建需要模型进行推理才能找到正确答案的“大海捞针”（Needle-in-a-Haystack）测试，联合评估每个头的检索和推理能力。
3. 两步式压缩框架：该方法包含两个主要步骤（如图1所示）：
* 头级别重要性分数评估：使用特别设计的、需要上下文推理能力的“大海捞针”测试来识别对模型性能贡献大的重要头。
* 头级别KV缓存分配：在预填充（prefilling）阶段，根据第一步得到的重要性分数分布，为每个头分配相应的KV缓存预算。

  1. 优越的实验性能：在多个基准测试（如LongBench, LooGLE）和模型（Llama-3-8B-Instruct, Mistral-7B-Instruct）上的实验表明，该方法显著优于现有的强基线模型，尤其是在低资源（KV大小=64和128）设置下。值得注意的是，该方法仅保留1.5%的KV缓存，就能在上下文问答基准上达到完整KV缓存97%的性能。


图 1: 我们提出的头级别KV缓存压缩方法包括两个步骤：（1）头级别重要性分数评估（上部）：使用“大海捞针”测试识别对上下文推理能力有贡献的重要头。（2）头级别KV缓存分配（下部）：在预填充阶段，根据第一步识别的重要性分数分布为每个头分配KV缓存预算。

A3 背景知识

2.1 注意力头

多头注意力的功能分析。多头注意力是Transformer架构【【索引31，Attention is all you need，2017，Advances in Neural Information Processing Systems】](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf)的核心组件，已有大量工作分析了单个注意力头的作用【【索引32 ，Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned，2019，Association for Computational Linguistics】](https://aclanthology.org/P19-1580)、【【索引24 ，In-context learning and induction heads，2022】】、
【【索引15，Cutting off the head ends the conflict: A mechanism for interpreting and mitigating knowledge conflicts in language models，2024】】、【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】、【【索引43，Attention heads of large language models: A survey，2024】】，其目标通常是剪枝冗余的头【【索引29，Layer-wise pruning of transformer attention heads for efficient language modeling，22021】】、【【索引18，A fast post-training pruning framework for transformers，2022】】或增强可解释性【【索引24，In-context learning and induction heads，2022】】、【【索引15，Cutting off the head ends the conflict: A mechanism for interpreting and mitigating knowledge conflicts in language models，2024】】、【【索引43，Attention heads of large language models: A survey，2024】】。例如，Voita等人【【索引32，Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned，2019，Association for Computational Linguistics】](https://aclanthology.org/P19-1580)发现，在机器翻译任务中只有一小部分头起关键作用，它们通常负责管理位置信息、句法分析或关注罕见词。类似地，Olsson等人【【索引24 ，In-context learning and induction heads，2022】】识别出“归纳头”（induction heads），它们实现了一种复制粘贴机制，使模型能够复制先前遇到过的序列。此外，Zheng等人【【索引43，Attention heads of large language models: A survey，2024】】全面回顾了近期表征注意力头不同功能的研究，将其分为四类：知识回忆、上下文识别、潜在推理和表达准备。一项密切相关的工作由Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】展开，他们发现了在知识获取中起关键作用的“检索头”（retrieval heads），并使用“大海捞针”测试进行验证。这些关于头级别功能的见解为我们提出的头级别KV缓存压缩方法奠定了基础，该方法旨在在压缩过程中同时保留检索和推理能力。

2.2 KV缓存压缩

提升LLM计算效率的研究。提升大型语言模型（LLM）的计算效率，尤其是针对极长输入的情况，吸引了大量研究兴趣【【索引27，Fast transformer decoding: One write-head is all you need，2019】】、【【索引5，Accelerating large language model decoding with speculative sampling，2023】】、【【索引6，FlashAttention-2: Faster attention with better parallelism and work partitioning，2024】】、【【索引44，A survey on efficient inference for large language models，2024】】、【【索引40，Chunkattention: Efficient self-attention with prefix-aware kv cache and two-phase partition，2024】】，其中包括缓存多头注意力中的键（key）和值（value）向量以提高生成效率的方法【【索引9，Model tells you what to discard: Adaptive KV cache compression for LLMs，2024】】、【【索引35，Efficient streaming language models with attention sinks，2024】】、【【索引42，H2o: Heavyhitter oracle for efficient generative inference of large language models，2023】】、【【索引23，Snapkv: Llm knows what you are looking for before generation，2024】】、【【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024】】。全量KV缓存方法的一个挑战是内存使用量随输入长度线性增长，导致了显著的内存管理难题。为了减少内存使用并提高推理效率，研究者们提出了多种KV缓存压缩技术。例如，Xiao等人【【索引35，Efficient streaming language models with attention sinks，2024】】通过StreamingLLM解决了“注意力沉没”（attention sink）问题，只保留前k个token的KV缓存。Zhang等人【【索引42，H2o: Heavyhitter oracle for efficient generative inference of large language models，2023】】应用Heavy Hitter Oracle策略来选择关键的缓存条目，而Li等人【【索引23，Snapkv: Llm knows what you are looking for before generation，2024】】则利用最后α个token的注意力分数来识别相关条目。Cai等人【【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024】】引入了PyramidKV，基于注意力矩阵的模式为较高层分配较小的缓存预算。

现有压缩方法的局限性与本文方法的差异。尽管这些方法提升了性能和效率，但它们大多依赖于固定的层级分配策略，这可能无法为下游任务完全优化KV缓存的分配。Feng等人【【索引8，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference，2024】】提出了使用注意力分数进行动态的头级别分配，但仍依赖于层级预算。类似地，Tang等人【【索引30，Razorattention: Efficient kv cache compression through retrieval heads，2024】】利用检索头的分布进行头级别分配，但对关键的头保留了完整的KV缓存，限制了真正的头级别压缩。相比之下，我们的方法仅根据头级别的重要性分数来分配KV缓存，不受层级约束，从而实现了更有效的压缩，并且在不增加解码延迟或内存使用的情况下优于基线方法。

A2 方法细节

本节介绍了我们提出的头级别KV缓存压缩方法，该方法由三个关键部分组成：（1）识别重要头并计算头级别的重要性分数分布（3.1），（2）利用这些分布在各个头之间高效地分配KV缓存预算（3.2），以及（3）确定在每个头内部保留哪些键（Key）和值（Value）向量（3.3）。

3.1 头级别重要性分数评估

通过重要性分布进行预算分配。在头级别实现精确的预算分配，需要识别出对于给定任务哪些头最重要，哪些最不重要。通过利用这种重要性分布，我们可以为更关键的头分配更大的KV缓存预算，为次要的头分配较小的预算。为了实现这一点，我们受Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】的启发，提出了一种新颖的重要性分数评估方法，使我们能够有效地评估每个注意力头的重要性，以实现最佳的KV缓存分配。


图 2: 头识别示例对比：左图为Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】用于识别检索头分布的“大海捞针”测试示例，右图为我们提出的用于识别检索-推理头分布的“大海捞针”测试示例。

检索头（Retrieval Heads）的识别方法。Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】使用“大海捞针”测试1和自定义的检索示例如图2所示，来评估每个头的重要性分数。在这些示例中，一个无法用模型参数知识回答的问题q，与一个答案k（“针”）配对，并将这个答案插入到“干草堆”c的不同位置pi。模型被要求从组合输入(k, c)中检索出确切的答案k。

检索头的重要性分数计算。在每个解码步骤t，如果(1)注意力分数最高的token arg max(ah)与生成的token匹配，并且(2)该token是插入答案k的一部分，那么头h就会获得一部分重要性分数。每个头h的最终重要性分数计算如下：

$$\begin{aligned} S_h = \sum_{t=1}^{N} \mathcal{N}^t, \text{ where } \mathcal{N}^t = \begin{cases} \frac{1}{N} & \text{if } \arg max(\boldsymbol{a}_h^t) \in \boldsymbol{k} \\ 0 & \text{otherwise} \end{cases} \end{aligned}$$

其中N是插入答案k的长度，$a_h^t$是头h在t-sh解码步对组合输入的注意力分数。通过使用插入位置pi的多种设置，他们获得了头级别的检索头分布。

直接使用检索头分布存在的问题。直接使用此分布存在两个问题：（1）它侧重于检索-粘贴机制，缺乏对复杂问题所需的上下文和推理技能的考量；（2）该分布对于有效的头级别KV缓存分配来说过于稀疏，由于严格的精确匹配要求，近70%的头获得的重要性分数为零【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】。

检索-推理（R2）头的提出。为了解决这些问题，我们提出了一种新的重要性分数评估方法，该方法同时考虑了头的检索和推理能力，从而能够更准确地评估其重要性。

R2示例的构建方法。首先，我们通过在Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】的检索示例中加入明确的上下文推理步骤来构建检索-推理示例，如图2的检索-推理示例部分所示。我们进一步将插入的“针”修改为三个部分：k = (r, c1, c2)，其中r是推理步骤，c1和c2是针对细化问题q的不同答案。模型必须利用r进行推理，以检索并生成正确答案c2，同时避免选择错误答案c1。

R2重要性分数的 refined 评估方法。其次，我们改进了评估方法，重点关注整个正确答案c2（图1中的Correct Copy），因为所有token都与问题q相关。这种方法与上下文问答（Contextual QA）任务的要求一致，该任务需要同时具备检索和推理能力。通过考虑完整的正确答案，每个头h的重要性分数不再仅仅依赖于注意力分数最高的token。头h的重要性分数计算如下：

$$\begin{aligned} S_{h}=\sum_{t=1}^{N} \sum_{i=1}^{N} \mathcal{M}_{i}^{t}, \quad \text { where } \mathcal{M}_{i}^{t}= \begin{cases}\frac{a_{i}}{N} & \text { if top-} i\left(\boldsymbol{a}_{h}^{t}\right) \in \boldsymbol{c}^{2} \\ 0 & \text { otherwise }\end{cases} \end{aligned}$$

其中，$a_i \in a_h^t$ 表示来自头h的第i高的注意力分数，而$top-i(a_h^t)$是第t个解码步骤中得分第i高的token。我们通过考虑整个正确答案并增加每个头评估的token数量来计算重要性分数（如图1中的重要性分数评估）。直观上，对正确答案k有更高注意力的头应该获得更高的重要性分数。我们进一步使用注意力权重来细化分数，从而产生更准确的分布，如公式2所示。

3.2 头级别KV缓存分配

基于重要性分数的预算分配。在评估出每个头的重要性分数后，我们便可以识别出关键的头并相应地分配KV缓存预算。本节将解释如何将这些分布融入到头级别的KV缓存分配中。

预备知识：多头注意力机制。在多头注意力机制中，对于第l层的每个头h，嵌入的输入$X = \{x1, x2, . . . , xn\} \in R^{n \times d}$会通过查询矩阵$W_Q^l$、键矩阵$W_K^l$和值矩阵$W_V^l \in R^{d \times d_h}$映射到不同的子空间：

$$\boldsymbol{Q}_{h}^{l}=\boldsymbol{X} \boldsymbol{W}_{Q}^{l} \in \mathbb{R}^{n \times d_{h}} ; \quad \boldsymbol{K}_{h}^{l}=\boldsymbol{X} \boldsymbol{W}_{K}^{l} \in \mathbb{R}^{n \times d_{h}} ; \quad \boldsymbol{V}_{h}^{l}=\boldsymbol{X} \boldsymbol{W}_{V}^{l} \in \mathbb{R}^{n \times d_{h}} .$$

KV缓存压缩的目标。为了优化内存并提升效率，KV缓存压缩方法【【索引35，Efficient streaming language models with attention sinks，2024】】、【【索引23，Snapkv: Llm knows what you are looking for before generation，2024】】、【【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024】】被用来丢弃不重要的KV缓存条目，同时保持性能。对于每个头h，压缩后的KV缓存被减少到$K_{lh} \in R^{s \times d_h}$和$V_{lh} \in R^{s \times d_h}$，其中$s \ll n$，从而显著提高了计算效率。

头级别分配与现有方法的对比。先前在预填充阶段进行KV缓存压缩的工作【【索引35，Efficient streaming language models with attention sinks，2024】】、【【索引23，Snapkv: Llm knows what you are looking for before generation，2024】】、【【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024】】、【【索引8，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference，2024】】仅限于层级分配，即每层使用统一或动态的预算，但对层内的所有头一视同仁。尽管Feng等人【【索引8，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference，2024】】引入了头级别的信息，但他们的方法仍依赖于层级分配作为先决条件。

预算分配策略。基于头级别的重要性分布，我们提出了一种全面的KV缓存分配策略。每个头h最初被分配一个固定的KV缓存大小b，并关联一个重要性分数Sh。为了实现动态分配，我们通过从每个头中提取一部分预算来创建一个共享的预算池B，剩余部分作为基本预算。这个过程在图1的头级别分配部分有图示。然后，预算池B根据各头的重要性分数Sh按比例分配。重要性分数分布S经过L1归一化，以确保所有Sh的总和为1。最终的头级别KV缓存分配如下：

$$B = \frac{b}{\beta} \times L \times H; \quad \boldsymbol{b}_h = (b - \frac{b}{\beta}) + S_h \times B$$

其中b是每个头的初始固定预算，β是一个控制动态预算池大小的超参数，L和H分别是当前LLM的层数和头数。

最终保留的KV缓存组成。在形成动态预算池B之前，最后的α个指令token会被保留下来，以指导选择过程，具体细节在3.3节中阐述。每个头保留的KV缓存包括：（1）基本预算($b - b\beta$)，（2）与其重要性分数成正比的动态预算($S_h \times B$)，以及（3）最后的α个指令token。

3.3 KV缓存选择

基于注意力的选择策略。在使用上述算法确定要保留的KV缓存条目数量后，我们应用了先前工作【【索引23，Snapkv: Llm knows what you are looking for before generation，2024】】、【【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024】】、【【索引8，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference，2024】】中的基于注意力的选择策略，以保留最相关的条目。具体来说，最后的α个指令token（局部窗口）用于指导每个头的KV缓存选择。从这些局部窗口到其余token的注意力分数通过一个池化层进行聚合，得分较高的token被认为更重要，并被保留在缓存中。

A4 实验环境

  • 模型:

    • Llama-3-8B-Instruct【【索引7，The llama 3 herd of models，2024】】
    • Mistral-7B-Instruct【【索引12，Mistral 7b，2023】】

  • 数据集:

    • LongBench【【索引3，LongBench: A bilingual, multitask benchmark for long context understanding，2024】】: 用于评估上下文推理能力，使用了其中的"Single-Doc QA"和"Multi-Doc QA"类别的数据集。
    • LooGLE【【索引20，Loogle: Can long-context language models understand long contexts?，2024】】: 专注于长依赖问答任务，使用了其中四个与QA相关的任务。
    • 具体数据集细节见附录表5。

  • 基线方法:

    • SnapKV【【索引23，Snapkv: Llm knows what you are looking for before generation，2024】】: 使用最后α个token作为局部窗口指导KV缓存选择。
    • PyramidKV【【索引4，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024】】: 遵循金字塔注意力模式，为底层分配更多KV缓存。
    • Ada-KV【【索引8，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference，2024】】: 在每层内根据头的注意力集中度动态分配预算。本文使用其性能更优的变体Ada-SnapKV作为基线。

  • 软件与实现细节:

    • 本文提出的方法（HeadKV-R和HeadKV-R2）使用SnapKV作为预算分配后的选择策略。
    • 所有方法中，局部窗口大小α统一设置为8。
    • 超参数β（控制共享预算池大小）从集合 {1.005, 1.01, 1.1, 1.2, 1.5, 2, 5, 10} 中选择，并报告最佳性能。

  • 硬件配置:

    • 实验使用了Mistral-7B-Instruct模型，最大序列长度为32K，并默认启用FlashAttention【【索引6，FlashAttention-2: Faster attention with better parallelism and work partitioning，2024】】来进行内存和延迟评估。具体的GPU型号和数量未明确提及。

A4 实验结果

主要结果

实验内容: 在LongBench和LooGLE基准上，比较了本文提出的HeadKV-R（基于检索头）和HeadKV-R2（基于检索-推理头）与SnapKV、PyramidKV、Ada-SnapKV以及FullKV（完整KV缓存）的性能。评估了不同KV缓存大小（64, 128, 256, 512, 1024）下的表现。

实验结果:

  • 总体性能: 如表1所示，本文的头级别KV缓存压缩方法（HeadKV-R和HeadKV-R2）在几乎所有任务上都一致地优于强基线方法，尤其是在资源受限的设置下（KV大小为64和128）。
  • HeadKV-R2的优势: HeadKV-R2（使用检索-推理头分布）的性能显著优于HeadKV-R（使用标准检索头分布），这凸显了将推理能力纳入重要性评估的价值。
  • 超越FullKV: 在Llama-3-8B-Instruct模型上，当KV大小为1024时，HeadKV-R2的平均分（32.95）甚至超过了FullKV的平均分（32.90），证明了该方法在选择性保留关键信息方面的有效性。
  • 不同KV大小下的性能: 如图3所示，随着KV缓存大小的增加，所有压缩方法的性能都得到提升。但本文提出的方法在所有压缩级别上都保持领先，尤其是在极小的KV大小（如64，仅占总token的0.7%）下，优势更为明显。

分析结论: 头级别分配策略和检索-推理（R2）分布是性能提升的关键。头级别分配比层级别分配更精细，能根据每个头的具体重要性动态调整预算，从而更有效地保留关键信息。而R2分布相比仅关注检索的分布，能更好地捕捉对复杂上下文问答任务至关重要的头，从而带来显著的性能增益。


表 1: Llama-3-8B-Instruct和Mistral-7B-Instruct在LongBench和LooGLE基准上的性能对比。我们的头级别KV缓存压缩方法优于所有基线，尤其是在低资源设置（KV大小=128）下。在Llama-3-8B-Instruct上（KV大小=1024），它甚至超过了FullKV的结果（32.90 vs 32.95），凸显了将上下文推理纳入头选择的好处。


图 3: 不同KV缓存大小（64, 128, 256, 512, 1024）下的结果，显示了在LongBench基准的六个数据集上的平均准确率，平均输入长度为8,683个token。值得注意的是，64的KV缓存大小仅保留了总token的0.7%。

检索-推理头的重要性

实验内容: 进行了一项消融研究，以评估本文提出的两项改进（1. 引入检索-推理示例；2. 改进重要性分数估计算法）对性能的影响。比较了三种设置：
1. HeadKV-R: 使用标准的检索头分布。
2. HeadKV-ER (Enhanced-Retrieval): 使用原始的检索示例，但采用本文改进的重要性分数估计算法（关注整个答案而非单个token）。
3. HeadKV-R2: 同时使用检索-推理示例和改进的估计算法。

实验结果:
- 分数估计算法的效果: 如表2所示，HeadKV-ER的性能优于HeadKV-R，表明关注整个答案的估计算法是有效的。
- 检索-推理示例的效果: HeadKV-R2的性能一致优于其他两种方法，证明了引入需要推理的示例对于识别关键头至关重要。
- 分布的可视化: 如图4所示，标准的检索头分布非常稀疏（许多头得分为0），这限制了其区分头重要性的能力。而增强后的检索头分布和检索-推理头分布则密集得多，能为KV缓存分配提供更有效的信息。

分析结论: 本文提出的两项改进都是有效的。改进的重要性分数估计算法通过考虑整个答案，使得分数分布更密集、信息更丰富。而引入检索-推理示例则能更好地捕捉模型在复杂任务中所需的综合能力，从而实现更优的性能。


图 4: Llama-3-8B-Instruct结果的头可视化。检索头分布过于稀疏，无法有效区分头的重要性，而我们的检索-推理头分布更密集，能更好地进行区分。Mistral-7B-Instruct的结果见附录图7。

表 2: LongBench基准上的消融研究结果。HeadKV-R利用标准检索头分布。HeadKV-ER使用Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】的检索示例，但采用我们提出的重要性分数评估方法。HeadKV-R2同时利用我们提出的重要性分数评估方法和检索-推理示例。

长上下文检索与推理能力

实验内容:
1. 大海捞针（Needle-in-a-Haystack）测试: 评估不同KV缓存压缩方法在长上下文中的信息检索能力。所有方法均设置KV大小为128。
2. 草堆中的推理（Reasoning-in-a-Haystack）测试: 评估模型在长上下文中进行推理的能力。该测试将多个需要推理的“针”（来自bAbI基准）插入“草堆”中，模型必须找到并综合这些信息才能得出正确答案。

实验结果:
- 大海捞针测试: 如图5所示，在Llama-3-8B-Instruct上，本文的头级别压缩方法（HeadKV-R和HeadKV-R2）在长上下文检索任务上显著优于所有基线方法，能够更有效地保留重要信息。
- 草堆中的推理测试: 如表3所示，本文的方法在长上下文推理任务上也显著优于基线。HeadKV-R2凭借其对检索-推理能力的侧重，取得了最佳的平均准确率。有趣的是，即使是HeadKV-R也优于其他基线，这可能是因为检索头与推理头存在重叠，且bAbI任务的答案本身也包含在“针”中，因此强化检索能力本身也有帮助。

分析结论: 本文提出的头级别压缩方法，特别是HeadKV-R2，不仅能提升标准基准上的性能，还能有效保持模型在长上下文场景下的核心能力，包括精确的信息检索和复杂的上下文推理。


图 5: Llama-3-8B-Instruct在KV缓存=128下的“大海捞针”测试结果。我们基于SnapKV构建的头级别KV缓存方法显著优于所有强基线。此外，我们的检索-推理头分布保持并提升了长上下文检索能力。Mistral-7B-Instruct的结果见附录图8，与Llama-3-8B-Instruct的结果一致。

表 3: KV缓存=128下的“草堆中的推理”测试结果。最终结果是每个长度下QA1-QA5任务的平均值。与“大海捞针”测试不同，此测试将多个“针”插入“草堆”，要求模型通过推理来提取正确答案。

内存与延迟

实验内容: 使用Mistral-7B-Instruct模型（最大序列长度32K）评估本文方法与基线方法的计算效率，包括解码延迟和峰值内存使用。

实验结果:
- 解码延迟: 如图6左侧所示，本文提出的方法（HeadKV）与其他KV缓存压缩方法（SnapKV, PyramidKV, Ada-SnapKV）的解码延迟几乎相同。所有压缩方法的延迟都远低于FullKV。这表明本文方法的预填充阶段开销（在生成长度=1时体现）可以忽略不计。
- 峰值内存使用: 如图6右侧所示，本文方法的峰值内存使用量与其他压缩基线相当，并且都显著低于FullKV。

分析结论: 本文提出的头级别KV缓存压缩方法在实现优越性能的同时，保持了与其他先进压缩方法相当的计算效率。它能够在不增加额外解码延迟或内存开销的情况下，提供更好的模型准确率。


图 6: 解码延迟和峰值内存使用结果。我们提出的方法与其他KV缓存压缩基线保持了相当的计算效率。

A5 结论

本文介绍了HeadKV-R2，一种创新的头级别KV缓存压缩方法，它包含两个核心部分。首先，我们提出了一种根据注意力头的重要性分数来分配KV缓存预算的策略。其次，我们开发了一种新颖的方法来评估这些分数，通过构建同时需要检索和推理能力的示例来实现。通过为更重要的头分配更大的KV缓存预算，为次要的头分配较小的预算，我们的方法比层级别的KV缓存压缩方法能更有效地保留重要的KV缓存。我们在多个基准、模型和长上下文能力测试中对HeadKV进行了全面评估，综合结果表明，我们的方法在保持计算效率的同时，取得了更优越的性能。

对于未来的工作，进一步探索各种类型的注意力头可能会提供有价值的见解，例如那些参与上下文学习【【索引24，In-context learning and induction heads，2022】】和真实性【【索引21，Inference-time intervention: Eliciting truthful answers from a language model，2023】】的头。例如，更加重视真实性头可能有助于缓解幻觉等问题，从而提高模型输出的事实准确性。此外，开发一种通用的、任务特定的分数评估算法也值得研究。一个潜在的方向是利用特定任务的梯度来更有效地分配KV缓存预算，从而实现定制化的压缩，以增强模型在不同应用中的性能。

A6 附录

A 详细结果

详细实验数据。在表4中，我们展示了主论文中图3的详细结果。


表 4: Llama-3-8B-Instruct和Mistral-7B-Instruct在不同KV缓存大小（64, 128, 256, 512, 1024）下的详细结果。

B 数据集详情

数据集信息。表5提供了我们实验中使用的数据集的详细信息。为了评估不同KV缓存压缩方法的检索和上下文推理能力，我们使用了来自LongBench【【索引3，LongBench: A bilingual, multitask benchmark for long context understanding，2024】】的六个问答数据集和来自LooGLE【【索引20，Loogle: Can long-context language models understand long contexts?，2024】】的四个额外问答数据集作为基准。对于LooGLE数据集，我们随机选择了100个样本来构建本研究中使用的数据集。

表 5: 数据集详情。

C LONGBENCH结果

其他LongBench任务的评估。我们在表6中列出了在LongBench【【索引3，LongBench: A bilingual, multitask benchmark for long context understanding，2024】】其余数据集上的评估结果。我们的头级别KV缓存压缩方法在这些任务上也优于其他基线，尤其是在Llama-3-8B-Instruct和Mistral-7B-Instruct的摘要和代码数据集上表现突出。

表 6: LongBench上的结果（KV缓存=128）。


图 7: Mistral-7B-Instruct上的头可视化结果。我们增强后的分布也能覆盖重要的检索头。


图 8: Mistral-7B-Instruct上的“大海捞针”测试结果。结果与Llama-3-8B-Instruct上的一致，我们的检索-推理头分布（HeadKV-R2）仍然优于其他强基线。

D “草堆中的推理”测试详情

任务示例。如图12所示，我们为“草堆中的推理”测试2中使用的每个任务提供了示例。我们使用了Kuratov等人【【索引17，Babilong: Testing the limits of llms with long context reasoning-in-a-haystack，2024】】整理的数据集来进行该测试。具体来说，他们将图12中列出的“输入”作为“针”，并用句点将其分割成不同的句子。然后将这些句子插入到“草堆”中的不同位置，以构成最终的“草堆中的推理”测试样本。

详细结果分析。如图13所示，我们提供了不同上下文长度和任务的详细结果。我们提出的头级别KV缓存压缩方法在所有方法中取得了最佳的平均分。值得注意的是，在Llama-3-8B-Instruct模型上，我们的方法在QA3任务上的表现超过了Full KV缓存，而根据Kuratov等人【【索引17，Babilong: Testing the limits of llms with long context reasoning-in-a-haystack，2024】】的说法，QA3是五个任务中最具挑战性的。这表明我们的检索-推理头分布能更好地捕捉语义相关和重要的内容，从而生成正确的答案。

E 构建检索-推理示例

示例构建的目标。为了指导头级别的KV缓存压缩，我们需要获得每个头的重要性分数。为此，我们手动构建了特定的示例，以确保模型在“大海捞针”实验中依赖于注意力头而不是内部知识来回答问题。因此，我们基于Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】的检索示例，通过向示例中引入不同的推理路径来构建检索-推理示例，以强调上下文推理能力。一个构建的检索-推理示例如图2右侧所示。除此之外，我们还反转了问题，总共创建了两个检索-推理示例。

实验设置。遵循Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】中概述的设置，在“大海捞针”实验中，我们使用模型的最大训练长度作为最大“草堆”长度，并均匀选择5个不同的长度值作为实际的“草堆”长度。对于每个“草堆”长度，问题被插入到10个不同的深度，从当前“草堆”长度的开始到结束均匀分布。总共，我们为每个模型生成100个示例来收集检索-推理头分布。

示例设计的考量。增加推理内容r和错误答案c1的目的是为了给推理过程增加复杂性和上下文，我们相信这可以根据头是否支持准确的推理模式来区分不同的头。我们设想它们的作用如下：


图 9: 构建的检索-推理示例。它们被用于进行“大海捞针”实验以获得检索-推理头分布。

与上下文推理要求的对齐。基于对上下文推理数据集的深入分析，我们知道相应问题的答案仍然会出现在输入中，但会伴有各种干扰项。因此，模型仍然依赖于检索-粘贴机制来获得真实答案。原始的检索头评估方法没有考虑到这种现象，但我们通过增加逻辑并模拟错误答案来解决这个问题，以获得更准确的分布。

引入多样化的推理路径。通过结合推理内容和错误答案，我们模拟了两种不同的潜在推理路径。错误答案充当干扰项，我们希望能找到那些即使在错误答案与正确答案结构几乎相同的情况下，仍能专注于正确答案的头。我们期望在评估方程中，将正确答案c2视为基准真值，这与原始的检索头评估方法类似。

专注于正确的推理路径。构建检索-推理示例的目的是获得每个头的重要性分数，然后用它来指导头级别的KV缓存预算分配。因此，我们的目标是识别重要的头，而不是那些关注错误答案的头。通过强调重要的头，那些关注错误答案的头会因为共享同一个全局预算池而被自然地忽略。

与标准检索头评估的对齐。我们遵循了Wu等人【【索引34，Retrieval head mechanistically explains long-context factuality，2024】】的设置，并基于“大海捞针”实验来确定检索-推理头的分布。由于“大海捞针”实验只输出正确答案，我们选择专注于正确答案，以确保在公式2中每个头能达到的最大重要性分数为1。增加额外的逻辑会破坏这一特性，可能会影响最终的分布。

F 伪代码

实现核心步骤。代码清单1展示了基于获得的重要性分数分布进行头级别KV缓存预算分配所需的核心步骤。

  • obtain_head_budget函数: 该函数根据先前获得的重要性分数分布动态确定每个头的预算，如公式4所示。由于重要性分数分布是静态的，不会随输入改变，该函数只需在模型初始化时执行一次。
  • head_kv函数: 该函数为不同的头分配预算。首先，我们使用SnapKV作为选择策略来指导具体的选择过程。在获得预算并确定选择策略后，我们根据预算执行选择并获得最终保留的结果。
 2 def obtain_head_budget(...):
 3     # Load saved importance score distribution
 4     with open(data_path, 'r') as file:
 5         head_list = json.loads(file.readline())
 6     # accumulate the importance score and convert to tensor
 7     head_score_list = [np.mean(l[1]) for l in head_list.items()]
 8     head_score_list = torch.tensor(head_score_list / sum(head_score_list))
 9     # Obtain the importance score distribution
10     total_attention = head_score_list.reshape(num_hidden_layers, num_attention_heads)
11     # Construct shared budget pool and define the minimum KV cache budget
12     total_pool_capacity = (base_capacity // beta) * num_hidden_layers * num_attention_heads
13     min_num = (base_capacity - base_capacity // self.beta)
14     # Head-level Allocation based on the importance score
15     head_capacity = torch.round(total_attention * total_pool_capacity + min_num).int()
16 
17 def head_kv(...):
18     # calculate attn_score from the observation windows to input. (SnapKV)
19     attn_score = calcul_attn_sore(self, key_states, query_states)
20     # Sorted based on attn_score and obtain indices.
21     _,indices = attn_score.sort(dim=-1,descending=True)
22     # Obtain cached index for each head.
23     for head_idx in range(num_heads):
24         cache_index = indices[head_idx][...,:head_capacity[self.layer_idx][head_idx]]
25         # Expand the indices to match the head dimension for gathering. (Same as SnapKV)
26         cache_index = cache_index.view(1, 1, -1, 1).expand(-1, -1, -1, head_dim)
27         # Gather the compressed past key and value states based on the selected indices. (Same as SnapKV)
28         top_Kcache = origin_heads_key_states[head_idx].gather(dim=2,index=cache_index)
29         top_Vcache = origin_heads_value_states[head_idx].gather(dim=2,index=cache_index)
30         # Merge with obervation windows
31         selected_k = torch.cat([top_Kcache,origin_heads_key_states[head_idx][:, :, -self. window_size:, :]],dim=2)
32         selected_v = torch.cat([top_Vcache,origin_heads_value_states[head_idx][:, :, -self. window_size:, :]],dim=2)
33 
34     # Combine together
35     heads_key_states.append(selected_k.view(-1, head_dim))
36     heads_value_states.append(selected_v.view(-1, head_dim))
37     # Merge together
38     heads_key_states = torch.cat(heads_key_states, dim=0)
39     heads_value_states = torch.cat(heads_value_states, dim=0)
40 
41     return heads_key_states,heads_value_states


列表 1: HeadKV的伪PyTorch风格实现。
图 10: 不同β值的结果，显示了在LongBench基准的六个数据集上的平均准确率。

G 超参数分析

唯一的超参数β。我们提出的方法引入的唯一超参数是β，它定义了共享全局预算池B的大小。其他超参数，如指令token数量α，与PyramidKV代码库3中的设置保持一致。我们也确保了所有其他超参数在基线和我们提出的方法中都是一致的。对于超参数β，如我们在4.1节中所说，它是从1.005, 1.01, 1.1, 1.2, 1.5, 2, 5, 10中选择的，我们报告的是基于平均分数的最佳性能，而不是为每个数据集选择一个β。在KV缓存=128的情况下，Llama-3-8B-Instruct和Mistral-7B-Instruct上的详细结果如图11所示。

β值的影响。对于超参数β，根据公式4，一个较小的值代表一个较大的共享预算池B，这意味着KV缓存分配更依赖于重要性分数分布。结果显示，Head-R2在β较小时表现更好，表明我们的检索-推理头分布在指导KV缓存预算分配方面更有效。值得注意的是，在我们为Llama-3-8B-Instruct和Mistral-7B-Instruct设置的不同β值下，Head-R和Head-R2的表现都优于当前的SOTA方法Ada-SnapKV。这证明了我们提出的头级别KV缓存压缩方法的性能稳定性。

H 讨论

与上下文压缩方法的比较。除了KV缓存压缩方法，另一种缓解预填充阶段长上下文场景并增强模型处理长文本能力的方法是上下文压缩【【索引13，LLMLingua: Compressing prompts for accelerated inference of large language models，2023】】、【【索引14，LongLLMLingua: Accelerating and enhancing LLMs in long context scenarios via prompt compression，2024】】、【【索引10，In-context autoencoder for context compression in a large language model，2024】】、【【索引25，Dodo: Dynamic contextual compression for decoder-only LMs，2024】】。在上下文压缩领域，也有不同的方向探索如何压缩提示。例如，Jiang等人【【索引13，LLMLingua: Compressing prompts for accelerated inference of large language models，2023】】、【【索引14，LongLLMLingua: Accelerating and enhancing LLMs in long context scenarios via prompt compression，2024】】提出定义一个额外的模块来压缩原始提示，同时确保压缩前后提示的意义和连贯性保持不变。压缩后的文本将作为后续LLM的输入。另一方面，Ge等人【【索引10，In-context autoencoder for context compression in a large language model，2024】】和Qin等人【【索引25，Dodo: Dynamic contextual compression for decoder-only LMs，2024】】采用了不同的压缩策略，将原始提示压缩到内存槽而不是文本中。Ge等人【【索引10，In-context autoencoder for context compression in a large language model，2024】】提出了ICAE方法，它使用一个额外训练的上下文自编码器将输入压缩成固定长度的内存槽，然后将其用作LLM的输入。他们对提出的上下文自编码器进行了预训练和微调，以使其具备能力和泛化性。Qin等人【【索引25，Dodo: Dynamic contextual compression for decoder-only LMs，2024】】进一步利用模型的内部隐藏状态作为压缩后保留的内容，并使用一个额外的评分器进行相应的选择。

KV缓存压缩的特点。KV缓存压缩的目标是KV缓存本身，而上下文压缩的目标是原始输入或从输入中派生的相应表示。从上下文压缩的角度来看，当前的KV缓存压缩方法可以被视为使用模型自身的知识来压缩输入上下文，而不依赖于可训练的模块进行选择。例如，我们的方法所基于的SnapKV【【索引23，Snapkv: Llm knows what you are looking for before generation，2024】】方法，使用最后α个token作为观察窗口，并根据这些token的注意力选择性地保留KV缓存。与LLMLingua【【索引13，LLMLingua: Compressing prompts for accelerated inference of large language models，2023】】和ICAE【【索引10，In-context autoencoder for context compression in a large language model，2024】】等上下文压缩方法相比，当前的KV缓存压缩方法更简单，因为它们不需要定义或训练额外的组件。它们也更容易实现更高的计算效率，因为这些KV缓存压缩方法避免了依赖外部组件来获取压缩输入，并且不需要重新计算键和值矩阵。

I 局限性

并行执行的开销。尽管头级别的KV缓存压缩方法可以与其他基线方法保持相当的计算效率，但由于需要额外的工作来动态管理每个头，它们可能会引入额外的开销，尤其是在并行执行中。与头级别操作相关的同步和计算成本可能会在多GPU环境中降低解码速度，从而导致性能和效率之间的权衡。因此，如何进一步优化头级别KV缓存压缩以减轻并行开销并保持实际部署的优势，仍然是一个重要的挑战。


图 11: 不同β值的结果。我们利用scikit-learn从六个数据集中各提取15%作为验证集（HeadKV-validation），其余数据用作测试集（HeadKV-test）。黑线表示根据验证集确定的最佳β值。
图 12: “草堆中的推理”测试中使用的每个任务的示例。


图 13: “草堆中的推理”测试在五个数据集上的详细结果。