来自未来的回望：基于反因果惊奇度的键值缓存管理

# 来自未来的回响：基于反因果惊异度的键值缓存管理

来源：https://arxiv.org/html/2607.27600

Stephen Gould¹,²，Anton van den Hengel¹,³  
¹Metacognition AI  
²澳大利亚国立大学  
³阿德莱德大学  
\{steve,anton\}@metacognitionai\.com  

###### 摘要

近年来，通过压缩和驱逐策略进行键值（KV）缓存管理已成为一个重要研究方向。大型语言模型（LLM）及其多模态变体在输出生成过程中的计算需求，可以通过缓存后续缩放点积注意力操作所需的先前键值计算结果来部分缓解。然而，这又带来了另一个问题：KV缓存的大小随上下文长度线性增长，当提示或生成输出较长时，会迅速耗尽所有可用的GPU内存。KV缓存管理会周期性从缓存中剔除条目，从而在降低内存占用同时，尽力保留足够的信息以保证准确生成。其副产品是更快的推理速度。我们提出了一种简单而有效的KV驱逐方案，其动机源于一个洞见：能够由更近期的token很好地预测出的过往token是冗余的，它们关联的键值可以从缓存中移除。为了对条目进行驱逐评分，我们按照原始顺序在token上运行模型，重用KV缓存中已存储的键值表示，并应用反因果注意力掩码，使每个位置仅关注其未来上下文。该方案处于分布内（in-distribution），与实际缓存内容直接相关，且无需额外训练。为进一步降低成本，我们还提出了一种快速的单层近似方法，将反因果传递限制在最后一个Transformer层，从而在每次刷新周期中获得显著加速，同时仅产生微小的精度损失。我们在各种开源LLM和基准数据集上评估了我们的策略，结果显示与其它最先进方法相比，我们的方法具有竞争力或更优的表现。参考代码可在 https://github.com/metacognitionai/counter_causal 获取。

## 1 引言

大型语言模型（LLM）部署和普及中的一个关键创新是引入了键值（KV）缓存。通过提供一种机制来存储Transformer架构中缩放点积注意力所需先前计算出的键和值（Vaswani 等人，2017 (https://arxiv.org/html/2607.27600#bib.bib148)），KV缓存显著降低了计算量，使LLM能够在推理期间扩展到长上下文窗口。然而，这些计算上的节省是有代价的。KV缓存占用相当大的内存空间，随着上下文窗口增长到数千个token以上，存储该缓存所需的GPU内存变得难以承受。因此，近年来在寻找减少KV缓存内存占用的策略方面已有大量研究工作。

一条研究路线侧重于在尽量保持生成输出质量的同时，从KV缓存中驱逐条目。一个显而易见的方法是以滑动窗口方式，在KV缓存达到某个预设最大尺寸后移除最旧的条目（Beltagy 等人，2020 (https://arxiv.org/html/2607.27600#bib.bib56)）。另一种方法是在将输入提示提供给模型之前对其进行压缩，从而在模型开始生成输出之前减少上下文窗口中的token数量（Li 等人，2023 (https://arxiv.org/html/2607.27600#bib.bib78)；Fei 等人，2024 (https://arxiv.org/html/2607.27600#bib.bib79)）。显然，这仅在模型随后生成的文本较短时有效，例如回答多项选择题或需要简短回答的提示，或者仅适用于如思维链（CoT）推理流程中的阶段性应用（Xia 等人，2025 (https://arxiv.org/html/2607.27600#bib.bib71)）。

更复杂的方法会动态尝试识别缓存中最具影响力的条目并保留它们，同时移除无关条目（Zhang 等人，2023 (https://arxiv.org/html/2607.27600#bib.bib69)；Oren 等人，2024 (https://arxiv.org/html/2607.27600#bib.bib579)；Cetin 等人，2025 (https://arxiv.org/html/2607.27600#bib.bib68)）。这些方法通常通过分析输出生成期间计算的注意力矩阵，并认为获得高注意力分数的键具有影响力，应予以保留（连同其对应的值）。这些键关联的token在文献中被称为“重击者”（heavy hitters）（Zhang 等人，2023 (https://arxiv.org/html/2607.27600#bib.bib69)）。然而，这些基于注意力的分数容易受到*自我强化偏差*的影响：保留在缓存中的token会持续获得来自后续解码步骤的注意力，从而累积更高的分数，使它们在未来的驱逐周期中更可能被保留。相反，一个事实关键但低频、仅被提及一次的token可能吸引的注意力过少，以至于无法在驱逐中幸存，尽管它可能对回答后续查询至关重要。我们在第4节 (https://arxiv.org/html/2607.27600#S4) 中通过实验证明了这种失效模式，其中基于注意力的方法在某些任务/模型/缓存大小组合下精度下降更快。

在本文中，我们提出了一种新颖的驱逐策略，其基础观察是：如果我们能够根据一个过去token之后的token来预测该token，那么该过去token所提供的信息已经以某种方式嵌入到了后续上下文中。为捕捉这一思想，我们定义了一个称为*反因果惊异度*的量，它根据关联token可预测性的好坏来衡量保留键值对的收益。惊异度低的token对后续生成过程几乎没有进一步用处，可以从KV缓存中移除。相反，惊异度高的token无法由当前位于其后的token预测。因此，它们可能包含对LLM尚未生成的输出至关重要的独特信息，应予以保留。该策略还有一个额外好处，即更倾向于保留较新的token，这有助于LLM类似滑动窗口方法但以更原则化的方式关注最新上下文。重要的是，我们的方法不需要任何训练，并且可以被任何使用KV缓存进行自回归token生成的模型使用。

我们用于评分是否驱逐KV缓存条目的机制，是按原始token顺序执行一次前向传递，重用生成过程中已计算的缓存键和值，并应用反因果（上三角）注意力掩码，使每个位置仅关注缓存中的未来token。每次刷新的成本为 O(Ln²)，其中 L 是Transformer层数，n 是当前缓存大小，但可以同时对一批token执行。我们还提出了一种快速的单层近似方法，将反因果传递限制在最后一个Transformer层，使每次刷新成本降至 O(n²)，并实现7到9倍的加速（在Qwen2.5-7B上，512-token缓存大小的每次刷新低至7.9毫秒）。我们在多个基准数据集以及两个不同规模的开源权重大型语言模型家族上进行了实验。这些基准涵盖具有长而详细提示的任务，或需要长文生成的任务。这使我们能够在LLM推理的预填充阶段和解码（生成）阶段都测试我们的方法。结果表明，我们的方法与现有方法相比具有竞争力或表现更优。

## 2 背景与相关工作

Transformer架构是现代大型语言模型（LLM）及其多模态扩展的基础，通过缩放点积注意力实现任意长度序列建模（Vaswani 等人，2017 (https://arxiv.org/html/2607.27600#bib.bib148)）。我们实验中使用的当代模型，如Qwen2\.5和LLaMA 3\.1，执行自回归解码，其中先前计算的键值（KV）对被缓存以避免冗余计算（Yang 等人，2025 (https://arxiv.org/html/2607.27600#bib.bib51)；Grattafiori 等人，2024 (https://arxiv.org/html/2607.27600#bib.bib52)）。虽然KV缓存对高效推理至关重要，但它随序列长度（和注意力层数）线性增长，造成内存和带宽瓶颈，尤其是在长上下文生成时。除了增长问题，最近系统层面的工作集中于改进KV缓存的存储和访问方式。例如，分页或虚拟化KV缓存设计（如vLLM风格系统所推广的（Kwon 等人，2023 (https://arxiv.org/html/2607.27600#bib.bib58)））将缓存视为一种虚拟内存，实现高效批处理、前缀共享以及跨请求复用。在算法层面，FlashAttention（Dao 等人，2022 (https://arxiv.org/html/2607.27600#bib.bib57)）等技术在计算注意力时高效地在GPU上的内存单元之间移动数据，以减少内存读写。这些方法在不改变缓存语义内容或注意力计算精确性的情况下提高了吞吐量和内存利用率。它们与修改或修剪存储表示的方法互补，因此与本文工作互补。

### 持久与增强内存
一些工作通过将KV缓存与持久性或外部内存扩展，以纳入当前上下文窗口之外的信息。Sukhbaatar 等人（2020 (https://arxiv.org/html/2607.27600#bib.bib73)）引入了跨序列持久化的学习记忆槽，有效增强了模型的工作记忆。最近，Eyuboglu 等人（2025 (https://arxiv.org/html/2607.27600#bib.bib67)）提出了“cartridges”（墨盒）：由专门语料库派生的预训练KV缓存，可在推理时加载，减少对长上下文演示的需求。Xing 等人（2026 (https://arxiv.org/html/2607.27600#bib.bib72)）分析了预填充阶段的KV缓存动态，并自适应地在“快速”和“慢速”推理模式之间切换，仅对困难输入选择性增加计算。Behrouz 等人（2024 (https://arxiv.org/html/2607.27600#bib.bib74)）的补充工作探索了使用基于惊异度的信号优先处理重要信息的长期神经存储，尽管该操作是前瞻性的，并未严格定位为KV缓存管理。

### 输入上下文压缩
另一类方法在处理前或处理过程中减少有效输入长度。Li 等人（2023 (https://arxiv.org/html/2607.27600#bib.bib78)）使用自信息度量识别输入序列中的冗余以过滤词汇单元，而Fei 等人（2024 (https://arxiv.org/html/2607.27600#bib.bib79)）将输入分段为主题连贯的块，并使用辅助预训练模型对每个段进行摘要。这些方法主要作为预处理步骤运行，因此不直接解决自回归解码期间KV缓存的增长问题。其他方向探索在线压缩，例如在生成过程中增量摘要早期上下文，用紧凑的语义表示来换取精确的token级保真度（Ge 等人，2024 (https://arxiv.org/html/2607.27600#bib.bib54)；Li 等人，2024 (https://arxiv.org/html/2607.27600#bib.bib55)）。

### 动态上下文驱逐
与压缩不同，上下文驱逐方法在生成过程中显式管理KV缓存，使其非常适合长文解码。Zhang 等人（2023 (https://arxiv.org/html/2607.27600#bib.bib69)）提出了H2O框架，该框架识别对注意力分数影响不成比例的重击token。他们建议在缓存中维持重击token与近期token之间的平衡。Oren 等人（2024 (https://arxiv.org/html/2607.27600#bib.bib579)）提出TOVA，该方法使用最后一个Transformer层中最远处理token的注意力权重来对缓存token进行驱逐评分。与H2O不同，TOVA将驱逐分数与累积的生成注意力解耦，但和H2O一样，它依赖于前向传递的注意力，并且容易受到自我强化偏差的影响，这在预填充阶段应用时可能是一个问题，因为中间上下文token为多跳任务提供了较差的驱逐信号。Cai 等人（2024 (https://arxiv.org/html/2607.27600#bib.bib580)）提出PyramidKV，该方法是基于注意力模式在后期层中越来越集中的观察，跨Transformer层非均匀分配KV缓存预算（早期层槽位较少，后期层较多）。这与我们的评分标准是正交的，原则上可以在未来工作中与反因果驱逐相结合。Kim 等人（2025 (https://arxiv.org/html/2607.27600#bib.bib581)）提出KVzip，通过选择那些移除后对上下文重建损害最小的条目来压缩KV缓存；该方法是查询无关的，需要一次辅助重建传递。与我们的方法类似，它执行一次额外的推理步骤，但使用因果重建目标。Feng 等人（2024 (https://arxiv.org/html/2607.27600#bib.bib582)）提出Ada-KV，该方法基于注意力熵动态分配每个头和每个层的缓存预算，使预算适应每个输入的内容，而不是使用固定的全局预算。这种逐层自适应分配与驱逐评分互补。Cetin 等人（2025 (https://arxiv.org/html/2607.27600#bib.bib68)）引入了神经注意力记忆模型（NAMM），该模型使用从注意力矩阵派生的特征（包括允许token关注未来位置的反向注意力）来学习token驱逐策略。他们的方法训练一个轻量级模型对token进行驱逐评分，而我们的方法则直接估计预测效用，不需要额外训练。其他相关方法近似或稀疏化注意力（Child 等人，2019 (https://arxiv.org/html/2607.27600#bib.bib60)；Jiang 等人，2024 (https://arxiv.org/html/2607.27600#bib.bib59)），而非显式驱逐token，例如通过低秩近似（Song 等人，2024 (https://arxiv.org/html/2607.27600#bib.bib61)）或在完整注意力计算之前进行选择性剪枝。这些方法模糊了架构修改与运行时缓存管理之间的界限。

### 低位表示与量化
与token级策略正交，一些工作通过浮点表示压缩来减少KV缓存内存。Hooper 等人（2024 (https://arxiv.org/html/2607.27600#bib.bib76)）和Zandieh 等人（2026 (https://arxiv.org/html/2607.27600#bib.bib70)）探索了KV张量的低位量化，在试图保持模型质量的同时显著减少内存占用。这些方法与驱逐和压缩互补，因为它们减少了每个token的内存成本，而不是存储的token数量。

简而言之，现有的KV缓存管理方法可分为四大类：（i）通过持久性或外部机制增强内存，（ii）压缩或摘要输入上下文，（iii）通过驱逐或稀疏化动态管理缓存，以及（iv）通过量化降低表示成本。诸如FlashAttention之类的算法级优化进一步提高了效率，同时不改变模型行为。我们的工作最直接地建立在上下文驱逐的基础上，侧重于一种原则性的、无需训练的方法，用于在预填充和自回归解码期间估计并保留缓存中最具预测性的元素。