面向长推理的信息感知KV缓存压缩

# 面向长推理的信息感知KV缓存压缩
来源：https://arxiv.org/html/2606.26875
\\setheadertext

LUMIA Lab\\correspondingemail\\emailiconjson\.kai@sjtu\.edu\.cn (https://arxiv.org/html/2606.26875v1/mailto:[email protected])‡通讯作者\.\\setheadertitle信息感知的KV缓存压缩用于长推理

Zhuiri Xiao³Alexandra Birch²‡Zhouhan Lin¹‡ ¹LUMIA Lab上海交通大学人工智能学院 ²爱丁堡大学信息学院 ³上海交通大学

###### 摘要

大语言模型（LLM）的推理能力发展迅速，导致在预填充和解码阶段，键值（KV）缓存规模不断增大。现有的KV缓存压缩方法主要依赖注意力权重来估计标记重要性。虽然注意力能有效捕捉上下文相关性，但它忽略了与预测不确定性和标记信息量相关的互补信息论信号。本文从前瞻性视角重新审视标记重要性，并提出**前向影响**这一度量标准，用于衡量压缩后的标记如何影响未来上下文。我们的分析表明：注意力分数选中的标记主要影响附近上下文，而与高预测不确定性相关的标记对远处未来上下文的影响则显著更强。基于这一观察，我们提出**InfoKV**，一个融合信息论信号的熵感知KV缓存压缩框架。它结合了标记级预测不确定性与逐层表示演化，并在推理过程中将所得熵分数与注意力分数整合。在Llama-3.1、Llama-3.2和DeepSeek-R1上的长上下文推理基准实验表明，InfoKV在长预填充和解码场景中均持续优于现有基于注意力的KV压缩方法。¹¹我们后续将发布代码以支持可复现性。

## 1 引言

大语言模型（LLM）在长上下文理解与推理方面展现了卓越的能力\[Guo et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib6), OpenAI et al., 2026 (https://arxiv.org/html/2606.26875#bib.bib16)\]。然而，由于计算注意力时二次增长以及键值（KV）缓存内存的线性增长，其在长序列处理中的部署仍存在计算昂贵的问题\[Łańcucki et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib11), Song et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib17)\]。这一瓶颈在长形式推理任务中尤为突出，因为LLM需要处理数千个标记作为输入或输出。

为解决此问题，近年研究探索了KV缓存压缩技术，选择性保留部分历史标记。一种常见范式是基于固定观测窗口（例如最近标记）的注意力权重估计标记重要性\[Li et al., 2024 (https://arxiv.org/html/2606.26875#bib.bib13), Cai et al., 2024b (https://arxiv.org/html/2606.26875#bib.bib1), Song et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib17)\]。来自最近上下文的注意力权重较大的标记被视为重要并保留，其余标记被丢弃。此类策略在推理效率和内存减少方面显示出有希望的改进。

尽管有效，基于注意力的KV缓存压缩方法存在固有局限：它们依赖**短期、后向**信号。具体来说，重要性是从最近标记对历史标记的注意力程度推断而来，这主要捕捉局部依赖关系。然而，长形式推理可能依赖于那些未被最近上下文直接激活但对未来推理轨迹至关重要的信息。这种不匹配在长解码场景中尤为严重，因为推理路径在生成步骤中动态演化。

参见图注\(a\)短程影响（128标记范围）\.
参见图注\(b\)长程影响（14K标记范围）\.

图1：基于熵、注意力权重及其组合评分的前1%标记在100篇Arxiv-Summarization文档上的短程与长程影响对比\[Cohan et al., 2018 (https://arxiv.org/html/2606.26875#bib.bib2)\]。短程影响捕捉即时预测效应，而长程影响反映持续的长期上下文影响。在本工作中，我们证明有效的KV缓存压缩应由**前瞻性**标记效用引导，即标记对未来生成步骤的贡献程度，而不仅仅是对最近上下文的相关性。为研究这一现象，我们引入**前向影响**，它度量从KV缓存中移除标记后未来预测分布的差异。如图1 (https://arxiv.org/html/2606.26875#S1.F1)所示，注意力强调与最近上下文密切相关的标记，而熵衡量标记的信息量，高熵标记对远处未来上下文表现出显著更强且更持久的影响。

受此观察启发，我们提出一个熵感知的KV缓存压缩框架**InfoKV**，将信息论信号融入标记选择。由于熵反映了模型预测标记时的不确定性，它自然捕捉了携带更丰富语义信息的标记。为了进一步跨层表征标记重要性，我们将熵与中间层和最终层之间的表示演化相结合，这独立于序列维度。

在长预填充和长解码基准上的大量实验表明，保留信息丰富的标记显著提升了推理性能。在长预填充场景中，InfoKV在使用Llama-3.1和Llama-3.2的LongReason基准上，在不同上下文长度和缓存预算下始终优于现有的基于注意力的KV缓存压缩方法。在长解码场景中，InfoKV在使用DeepSeek-R1时进一步在IFEval、AIME 2024和LiveCodeBench上取得了显著改进，展示了其在数学推理、指令遵循和代码生成任务上的有效性。

## 2 相关工作

#### KV缓存压缩。

KV缓存压缩研究的主流方向是基于注意力模式选择性淘汰历史标记。近年方法如SnapKV\[Li et al., 2024 (https://arxiv.org/html/2606.26875#bib.bib13)\]、PyramidKV\[Cai et al., 2024b (https://arxiv.org/html/2606.26875#bib.bib1)\]和FastKV\[Jo et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib8)\]提出了启发式剪枝策略，通过注意力权重衡量标记重要性并丢弃低注意力的标记。其他工作探索了标记合并以近似全缓存的原始注意力\[Zhang et al., 2024 (https://arxiv.org/html/2606.26875#bib.bib21), Wang et al., 2024 (https://arxiv.org/html/2606.26875#bib.bib20), Wan et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib19)\]。尽管这些方法有效减少了内存使用，但它们主要依赖基于注意力的启发式方法，本质上是后向的，主要对短答案的长预填充任务有效。

#### 面向长解码的压缩。

LLM的推理能力近年来日益受到关注\[Guo et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib6), OpenAI et al., 2026 (https://arxiv.org/html/2606.26875#bib.bib16)\]。随着长推理路径的生成，解码延迟和KV缓存增长变得比预填充效率更为关键。为应对这一挑战，近年研究将KV缓存压缩从预填充阶段扩展到解码阶段。RPC\[Song et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib17)\]将SnapKV\[Li et al., 2024 (https://arxiv.org/html/2606.26875#bib.bib13)\]推广到在线解码，通过在生成过程中周期性压缩KV缓存。Expected Attention\[Devoto et al., 2025 (https://arxiv.org/html/2606.26875#bib.bib3)\]进一步估计标记对未来注意力的期望贡献。此外，FreqKV\[Kai et al., 2026 (https://arxiv.org/html/2606.26875#bib.bib10)\]提出了一种迭代频域压缩框架，支持预填充和解码压缩，实现了高效的“短训练-长测试”能力。

#### 标记重要性的信息信号。

除了基于注意力的启发式方法，近期工作探索了信息论信号，从更本质的角度表征标记重要性。与依赖上下文交互的注意力权重不同，标记携带的信息代表其原生重要性。Selective Context\[Li et al., 2023 (https://arxiv.org/html/2606.26875#bib.bib12)\]利用自信息量化标记的信息量，并剪枝LLM输入中的冗余内容。基于不确定性度量，Kai等人\[2024 (https://arxiv.org/html/2606.26875#bib.bib9)\]提出了SH2，利用预测不确定性识别信息丰富的标记并调整输出分布以提高事实性。在长形式推理背景下，SeLaR\[Fu and Luo, 2026 (https://arxiv.org/html/2606.26875#bib.bib4)\]引入了熵感知的对比正则化，通过将表示推向远离过度自信的预测来鼓励探索。本文中，我们引入信息信号以更好地反映标记对未来上下文的影响，并优化长上下文推理的KV缓存压缩。

参见图注图2：在Arxiv-Summarization的100篇文档上，不同评分策略选出的前1%标记在长生成范围内的前向影响。前2048个标记使用不同标记重要性分数进行压缩，影响度量针对未来每128个标记的块测量。组合分数在短程和长程影响之间实现了更好的平衡。

## 3 方法

### 3.1 推理过程中重新审视标记重要性

现有KV缓存压缩方法主要通过从最近观测窗口计算的注意力分数来估计标记重要性。具体而言，来自最近标记的注意力权重较大的标记被视为重要并保留在KV缓存中。尽管这类策略在维持短程依赖方面有效，但它们隐含假设：对最近上下文重要的标记对未来生成步骤仍然重要。然而，在长形式推理和扩展解码过程中，标记的相关性不断演变，具有高最近注意力分数的标记可能仅局部贡献于附近上下文，而对未来推理轨迹提供有限效用。

为了更好表征标记的长期效用，我们从信息论角度重新审视标记重要性。直觉上，与高不确定性相关的标记为语言模型携带更多信息，因此更可能影响待生成的未来上下文。如Kai等人\[2024 (https://arxiv.org/html/2606.26875#bib.bib9)\]所揭示，这些标记往往是内容词，如形容词、名词和动词变位，它们比连词、限定词和介词等功能词更具信息量。

给定一个标记序列\\{x₀, x₁, ⋯, xₙ₋₁\\}，自回归语言模型θ对下一个标记xₙ的预测概率可形式化为：

p̂(xₙ) = p_θ(xₙ ∣ x_{<n}) (1)

对于标记xₙ，我们使用预测分布的熵来衡量其不确定性：

H(xₙ) = -∑_{xₙ∈V} p̂(xₙ) log p̂(xₙ) (2)

其中V表示词汇空间。熵越高表示模型预测下一个标记时越不自信，意味着对应上下文包含更丰富的信息。

参见图注图3：每层KV缓存压缩重要性分数的计算概览。InfoKV结合了预测熵、逐层表示演化和注意力分数进行标记选择。

### 3.2 压缩KV缓存的影响估计

我们使用Llama-3.1-8B-Instruct\[Grattafiori et al., 2024 (https://arxiv.org/html/2606.26875#bib.bib5)\]进行影响估计以验证我们方法的动机。我们定义标记x_i在KV缓存中对未来上下文块\\{x_{l_c}, ⋯, x_{r_c}\\}的**前向影响**为：移除x_i后原始预测分布与预测分布之间的平均散度：

I_{l_c:r_c}(x_i) = 1/(r_c - l_c + 1) ∑_{n=l_c}^{r_c} D_KL( p_θ(x_n ∣ x_{<n}) ∥ p_θ(x_n ∣ x_{<n} ∖ {x_i}) ) (3)

其中p_θ(x_n ∣ x_{<n} ∖ {x_i})表示从KV缓存中移除x_i后对x_n的预测概率。我们使用Kullback–Leibler (KL)散度量化两个预测分布之间的差异。为简化，在估计前向影响时所有层共享相同的标记选择。

基于该度量，我们分析了由不同重要性标准（包括注意力分数、熵及其组合）选出的标记的远程贡献。具体地，我们首先按照最近观测窗口\\{x_{l_o}, ⋯, x_{r_o}\\}内的平均注意力权重对标记排序，遵循先前KV缓存压缩方法：

A_i = 1/(r_o - l_o + 1) ∑_{t=l_o}^{r_o} Attn(q_t, k_i) (4)

其中Attn(q_t, k_i)是从观测窗口中的标记到标记x_i的注意力权重，取自最后一层。

然后我们将其与基于熵的标准选出的标记进行比较。为了结合注意力权重和熵，我们使用softmax在序列维度上归一化熵的尺度，并将其加到注意力分数上：

S_i = A_i + Softmax(𝐇)_i (5)

我们使用这三种评分策略压缩文档中的前2048个标记，并估计它们在短期未来范围和长期未来范围上的前向影响。图1 (https://arxiv.org/html/2606.26875#S1.F1)表明组合分数平衡了短程和长程影响。

沿长序列的前向影响如图2 (https://arxiv.org/html/2606.26875#S2.F2)所示。它揭示了基于注意力和基于熵的重要性估计之间的明显区别。注意力分数高的标记主要影响附近未来上下文，其影响随着生成距离增加而迅速衰减。相比之下，熵高的标记对远处未来上下文的影响显著更强，表明熵能更好地捕捉与长程推理和生成相关的信息。通过结合两种信号，我们可以挑选出对短程和长程未来都有贡献的标记，从而在压缩时保留对模型长期输出最关键的信息。