PatchKV: Weight-Space Compensation of KV Cache

📄 arXiv: 2609.39329v1 📥 PDF

作者: Chanryeol Lee, Chanhyuk Lee, Yeonwoo Choi, Donggyun Kim, Seunghoon Hong

分类: cs.LG

发布日期: 2026-09-30

备注: NeurIPS 2026. Code available at: https://github.com/cusasak/PatchKV


💡 一句话要点

提出PatchKV以解决长上下文推理中的KV缓存瓶颈问题

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 长上下文推理 KV缓存 压缩方法 岭回归 自然语言处理 模型权重 问答系统 性能提升

📋 核心要点

  1. 长上下文推理面临KV缓存线性增长带来的内存瓶颈,现有压缩方法在激进压缩预算下性能显著下降。
  2. PatchKV通过将上下文信息存储在模型权重中,提出了一种无训练的补偿机制,以提高KV缓存的有效性。
  3. 在多种基准测试中,PatchKV在长上下文问答和数学任务上均表现出显著的性能提升,改善了缓存压缩方法的效果。

📝 摘要(中文)

长上下文推理在大型语言模型(LLMs)中受到键值(KV)缓存线性增长内存的限制。现有的压缩方法通过令牌驱逐或近似来减少缓存,但在激进的压缩预算下性能急剧下降。我们提出了PatchKV,这是一种无训练框架,通过将部分上下文存储在模型权重中来补偿KV缓存压缩方法。PatchKV将现成的压缩KV缓存与上下文特定的权重补丁配对,该补丁在上下文加载时计算一次,并为后续查询提供服务。权重补丁通过岭回归以封闭形式导出,通过对齐在完整缓存和压缩缓存下的上下文派生参考查询令牌的块状激活来实现。一旦合并到模型中,补丁在单上下文、多查询设置中保持前向图和每查询推理成本不变。在长上下文问答(SCBench,最多170K个令牌,SQuAD,NIAH)和数学(GSM8K)基准测试中,PatchKV始终改善缓存压缩方法,表明在激进预算下补偿它们的另一种方向。

🔬 方法详解

问题定义:论文要解决的问题是长上下文推理中KV缓存的线性增长导致的内存瓶颈。现有的压缩方法在面对激进的压缩预算时,性能急剧下降,无法有效支持长上下文的推理需求。

核心思路:PatchKV的核心思路是通过将部分上下文信息存储在模型的权重中,来补偿KV缓存的压缩效果。这种设计使得在不增加推理成本的情况下,能够有效利用压缩后的KV缓存。

技术框架:PatchKV的整体架构包括两个主要模块:压缩的KV缓存和上下文特定的权重补丁。权重补丁在上下文加载时计算一次,并在后续查询中使用。

关键创新:PatchKV的主要创新在于通过岭回归以封闭形式导出权重补丁,从而对齐完整缓存和压缩缓存下的块状激活。这一方法与现有的令牌驱逐或近似方法本质上不同,提供了一种新的补偿机制。

关键设计:在设计中,PatchKV使用了岭回归作为关键的数学工具,以确保权重补丁的有效性。此外,补丁的合并过程保持了前向图和每查询推理成本不变,确保了模型的高效性。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

在长上下文问答(SCBench,最多170K个令牌)和数学(GSM8K)基准测试中,PatchKV显著改善了缓存压缩方法的性能,展示了在激进压缩预算下的有效性。具体实验结果表明,PatchKV在多个模型架构上均实现了性能提升,验证了其创新性和实用性。

🎯 应用场景

该研究的潜在应用领域包括长文本处理、问答系统和其他需要处理大规模上下文的自然语言处理任务。通过提高KV缓存的有效性,PatchKV能够在资源受限的环境中实现更高效的推理,具有重要的实际价值和未来影响。

📄 摘要(原文)

Long-context inference with Large Language Models (LLMs) is bottlenecked by the linearly growing memory of the key-value (KV) cache. Existing compression methods reduce the cache through token eviction or approximation, but degrade sharply at aggressive compression budgets. We propose PatchKV, a training-free framework that compensates KV cache compression methods by carrying part of the context in the model's weights. PatchKV pairs an off-the-shelf compressed KV cache with a context-specific weight patch, which is computed once at context-loading time and served for downstream queries for the context. The weight patch is derived in closed form via ridge regression, by aligning the block-wise activations of context-derived reference query tokens under the full cache and the compressed cache. Once merged into the model, the patch leaves the forward graph and per-query inference cost unchanged in the single-context, multi-query setting. Across long-context QA (SCBench with up to 170K tokens, SQuAD, NIAH) and math (GSM8K) benchmarks on three model architectures, PatchKV consistently improves cache compression methods, suggesting an alternative direction to compensate them at aggressive budgets.