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

通过结合熵和层间表示演化提升KV压缩的长程推理效果

前置知识

KV Cache

在Transformer推理过程中，为了加速自回归生成，系统会将计算过的键值对缓存起来，这样在生成新token时就不需要重新计算所有历史token的键值。KV cache的大小随序列长度线性增长，这是长序列推理的主要内存瓶颈之一。

本文的核心问题就是如何压缩KV cache以减少内存使用，同时保持推理性能，因此理解KV cache的工作机制是读懂本文的基础。

Attention权重

Attention机制计算当前查询向量与所有历史键向量之间的相似度，然后对值向量进行加权求和。Attention权重反映了历史token对当前token的重要性，现有的KV压缩方法主要利用这个指标来选择要保留的token。

本文的核心创新点之一就是指出单纯依赖Attention权重是不足的，需要引入额外的信息论信号，因此需要理解Attention权重的工作原理和局限性。

预测熵

信息论中的熵衡量随机变量的不确定性。在语言模型中，给定上下文后模型对下一个token的预测分布的熵反映了模型预测的不确定性程度。高熵意味着模型对下一个token的预测不确定，即当前上下文包含丰富信息。

这是本文引入的核心信号，作者发现高熵token对未来上下文有更强的远距离影响，这是与Attention权重的关键区别。

KL散度

Kullback-Leibler散度衡量两个概率分布之间的差异。在本文中，作者使用KL散度来量化从KV cache中移除某个token后，模型对未来token预测分布的变化程度，从而定义了Forward Influence指标。

这是作者用于验证和 motivating 其方法的关键指标，通过KL散度可以量化token对未来生成的影响程度。

研究动机

现有KV cache压缩方法主要依赖attention权重来估计token重要性，例如SnapKV、PyramidKV和FastKV等。这些方法从最近的观察窗口计算attention分数，保留被近期token高度关注的token。然而，这种策略存在固有局限性：它基于短期、向后看的信号，只能捕捉局部依赖关系。在长形式推理任务中，推理路径随着生成步骤动态演化，某些信息可能不会被近期上下文直接激活，但对未来的推理轨迹仍然至关重要。例如，在需要生成数千token的数学推理或代码生成任务中，关键的推理步骤或事实可能在数百个token后才能体现出其价值，但基于attention的方法可能已经将这些重要token丢弃了。

本文的目标是本文的目标是提出一种新的KV cache压缩框架，能够从前瞻性角度评估token重要性，即衡量token对未来生成步骤的贡献而非仅其对近期上下文的相关性。作者希望引入信息论信号来补充attention信号，从而在长推理轨迹中更好地保留对未来推理有用的信息，特别是在长解码场景下（如IFEval、AIME 2024和LiveCodeBench等需要生成数千token的任务），在有限的KV缓存预算下保持推理性能。

与已有工作不同的是，本文的独特切入角度是从信息论视角重新审视token重要性。作者引入了Forward Influence指标来衡量从KV cache中移除某个token后对未来预测分布的影响程度，发现高attention分数的token主要影响附近的未来上下文，而高预测不确定性（高熵）的token对远距离未来上下文有显著更强的影响。这个观察表明，现有方法过于强调局部依赖，而忽视了携带丰富语义信息的token对长程推理的重要性。本文通过结合预测熵、层间表示演化和attention分数，在序列维度和层维度两个正交方向上量化token的重要性，这是与现有方法的本质区别。

核心方法

InfoKV框架的核心思路是在KV cache压缩时不仅考虑attention信号，还引入信息论信号来综合评估token重要性。从直觉上看，高不确定性的token携带更多语言模型需要的信息，因此更可能影响未来的上下文生成。技术路线上，InfoKV首先从最终层计算每个token的预测熵来量化其信息量，然后计算每个token在早期层和最终层之间的表示演化程度（cosine距离），将两者相乘得到熵分数，最后与attention分数以加权方式组合成最终的重要性分数。对于每个Transformer层，独立计算重要性分数并保留top-ranked token，实现层级的自适应压缩。

核心创新点在于引入信息论信号（特别是预测熵）来补充attention信号，并证明了这两种信号的互补性。现有方法只关注token对近期上下文的相关性，而InfoKV通过测量预测不确定性来识别携带丰富语义信息的token。作者通过Forward Influence分析发现，基于attention选择的token主要影响短距离未来上下文（128 token范围内），而基于熵选择的token对长距离未来上下文（14K token范围内）有更强且更持久的影响。另一个关键创新是结合层间表示演化，通过计算早期层和最终层之间的cosine距离来量化token在不同层之间的语义收敛程度。表示变化大的token意味着包含更多未解决的语义信息，因此更有价值。

方法步骤详情

InfoKV的方法分为四个主要步骤。第一步，对于序列中的每个token，从最终层计算其预测熵，其中是模型在给定历史上下文下对下一个token的预测概率分布。为了提供更稳定的不确定性估计，作者采用Top-k限制熵，只使用预测概率最高的k个token来计算熵，避免低概率token的噪声干扰。第二步，对于每个层l，计算token在早期层l和最终层L之间的表示演化距离，其中和分别是token在层l和层L的隐藏表示。第三步，将表示距离与熵结合得到熵分数，为了确保最终层的熵分数不为零，给添加偏置。第四步，对于每个层l，将attention分数和熵分数以加权方式组合，其中控制attention和熵的平衡，最后保留每个层中分数最高的token。

技术新颖性

InfoKV的技术新颖性体现在三个层面。首先，这是首次将Forward Influence分析引入KV cache压缩领域，通过KL散度量化移除token对未来预测分布的影响，从信息论角度揭示了现有方法的局限性。其次，提出了在序列维度和层维度两个正交方向上量化token重要性的框架，序列维度通过预测熵捕捉信息量，层维度通过表示演化捕捉语义收敛程度，这种多维度的信号融合是全新的设计。最后，在实验方法上，作者不仅评估了传统的长prefilling任务（LongReason），还引入了长解码任务（IFEval、AIME 2024、LiveCodeBench），并使用推理模型（DeepSeek-R1）进行评估，这些任务需要生成数千token，更能体现前瞻性token选择的价值，这是KV cache压缩文献中相对较少探索的场景。

**Figure 2**: Forward influence of top-1% tokens selected by different scoring strategies over long generation horizons on 100 documents from Arxiv-Summarization

**Figure 3**: The overview of how to compute the importance score for KV cache compression in each layer

实验结果

在LongReason长上下文推理基准上，InfoKV在多个配置下实现了SOTA或极具竞争力的性能。对于Llama-3.1-8B-Instruct，在40%保留率下，InfoKV在16k/32k/64k上下文长度上的平均准确率（带CoT）分别为52.53%，而SnapKV为51.09%，PyramidKV为50.76%，Expected Attention为51.13%。在20%保留率下，InfoKV的平均准确率为49.79%，仍然优于SnapKV的49.37%和PyramidKV的48.57%。对于Llama-3.2-3B-Instruct，InfoKV同样展现出优势，在40%保留率下平均准确率为43.95%，而SnapKV为43.03%，PyramidKV为42.86%。更重要的是，InfoKV的优势随着序列长度增加而更加明显，这验证了熵信号在长程推理中的价值。在长解码任务上，InfoKV在DeepSeek-R1-Distill-Qwen-7B和R1-Distill-Llama-8B上优于RPC基线。特别值得注意的是，在IFEval任务上，R1-Distill-Llama-8B使用InfoKV在25%和12.5%保留率下甚至超越了完整缓存，这表明长推理轨迹中存在大量冗余，适当压缩反而能让模型更专注于信息量大的上下文。消融实验显示，Top-256限制熵在alpha=0.9和tau=1时达到最佳性能，证实了attention和熵的互补性。

**Table 1**: Accuracy (%) comparison on LongReason across different cache rates and prefill lengths (16k – 64k)

**Table 2**: Comparison of uniform budget and adaptive budget on IFEval

**Figure 4**: Performance on the three categories of long decoding benchmarks

**Figure 5**: Performance of InfoKV with different tau and top-k

**Figure 6**: Performance of InfoKV with top-256 restricted entropy and full-vocabulary entropy across a range of alpha

查看结构化数据

任务	指标	本文	基线	提升
LongReason长上下文推理	准确率（平均，带CoT）	52.53%（Llama-3.1-8B，40%保留率）	51.09%（SnapKV），50.76%（PyramidKV），51.13%（Expected Attention）	相比SnapKV提升1.44个百分点，相比PyramidKV提升1.77个百分点
IFEval指令遵循	性能分数（DeepSeek-R1-Distill-Llama-8B）	60.07%（25%保留率），63.77%（12.5%保留率）	完整缓存（约60%左右）	25%保留率下超越完整缓存，12.5%保留率下显著超越
AIME 2024数学推理	性能分数（DeepSeek-R1-Distill-Qwen-7B）	在25%和12.5%保留率下保持竞争力	RPC基线	在各种保留率下优于RPC
LiveCodeBench代码生成	性能分数（DeepSeek-R1-Distill-Llama-8B）	在25%和12.5%保留率下保持竞争力	RPC基线	在各种保留率下优于RPC

局限与改进

作者在论文中坦诚了几个局限性。首先，虽然熵比基于attention的指标展现出更强的前向影响，但它仍然是未来效用的间接近似而非显式优化目标，这意味着可能存在更直接的方法来优化token选择。其次，作者观察到的自适应层间预算分配策略对某些模型（如R1-Distill-Llama-8B）有改进，但对其他模型（如R1-Distill-Qwen-7B）会导致性能下降，说明不同架构在各层的信息分布可能存在差异，需要更鲁棒和架构感知的分配策略。此外，论文中的实验主要在开源模型（Llama系列和DeepSeek-R1）上进行，未在闭源模型上验证效果的泛化性。从方法本身来看，InfoKV需要计算额外的熵分数和表示演化距离，这会引入一定的计算开销，虽然论文未详细量化这部分开销，但在实际部署中可能成为考虑因素。

独立分析的弱点

从独立分析的角度来看，InfoKV存在几个潜在弱点。首先，方法依赖最终层的预测熵，但在长推理过程中，模型的理解会随着生成的进行而演化，早期计算的熵分数可能不能准确反映token在后续推理阶段的实际价值，这是一个随时间变化的动态选择问题，而InfoKV采用静态策略。改进方向可以考虑引入动态更新机制，根据解码过程中的上下文变化重新评估已生成token的重要性。其次，Top-k限制熵中的k值需要针对不同任务和模型进行调优，论文中固定k=256可能在某些场景下不是最优的，可以考虑自适应选择k的方法。第三，层间表示演化距离的计算需要额外的内存来存储早期层的隐藏状态，这在极端内存受限的场景下可能成为瓶颈，可以通过只采样部分层或使用代理指标来减少开销。最后，方法假设信息量大的token总是有用的，但在某些推理任务中，确定性高的token（如已确认的事实或中间结论）可能同样重要，当前的熵偏向可能会过度丢弃这类token。

未来方向

基于本文的成果，有几个有前景的未来研究方向。首先，作者提出的自适应层间预算分配虽然结果不一致，但这个思路值得进一步探索，可以研究不同架构的层间信息分布模式，设计更鲁棒的自适应策略，或者结合强化学习来动态学习预算分配策略。其次，Forward Influence作为验证指标，本身可以作为优化目标，可以设计直接最大化Forward Influence的token选择算法，而非间接地使用熵作为代理。第三，可以探索其他信息论信号，如互信息、条件熵或VIB原则，这些可能提供更丰富的token重要性刻画。第四，当前方法主要关注token级别的压缩，未来可以探索更细粒度的压缩策略，如在特征维度或注意力头维度进行选择性压缩，这可能在保持性能的同时实现更高的压缩比。最后，可以研究InfoKV在多模态模型中的应用，如在视觉-语言模型中，图像和文本token可能需要不同的重要性评估策略。

复现评估

作者在论文末尾声明代码将稍后发布，这意味着代码尚未开源，这给复现带来了一定困难。从实验设置来看，作者使用了标准的数据集和开源模型，这些都可以通过公开渠道获取。论文提供了详细的超参数设置，并在附录中包含了额外的实现细节。从算力需求来看，长上下文推理任务（特别是64k长度）和长解码任务（最多32768 token输出）需要相当的GPU内存，估计需要80GB或更大的显存。整体来说，一旦代码开源，复现难度应该是中等，主要挑战在于硬件要求和长序列任务的调试复杂性。建议作者在代码发布时提供预训练模型、完整的实验脚本和环境配置文件，以降低复现门槛。

面向长推理的信息感知KV缓存压缩 Information-Aware KV Cache Compression for Long Reasoning