KV-Kaizen: Learning Context-Adaptive Cache Compression Choices
作者: Joao Monteiro, Louis Béthune, Anastasiia Filippova, Sonia Laguna, David Grangier, Marco Cuturi
分类: cs.AI
发布日期: 2026-09-29
💡 一句话要点
提出KV-Kaizen以解决长文本处理中的缓存压缩问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长文本处理 缓存压缩 大语言模型 自适应选择 内存优化
📋 核心要点
- 现有方法通过丢弃不相关的标记来压缩缓存，但可能导致后续决策不利。
- 本文提出KV-Kaizen，通过学习选择器实现上下文自适应的缓存配置，避免驱逐标记。
- 实验结果显示，KV-Kaizen在长上下文任务中显著减少缓存大小，同时保持模型准确性。
📝 摘要（中文）
随着使用大语言模型（LLM）处理文本的上下文大小增加，键值（KV）缓存的大小可能超过原始模型权重分配的内存。这会对LLM的吞吐量产生负面影响，因为解码过程受限于内存，且解码成本随着缓存大小的增加而增加。近期的研究通过丢弃最不相关的标记来缓解这一瓶颈。然而，驱逐操作可能会在后续阶段造成不利影响。本文提出KV-Kaizen方法，通过学习选择器，根据上下文生成每层的缓存配置，从而实现缓存压缩，而无需驱逐标记。实验表明，该方法在指令跟随和推理任务中达到了准确性与缓存大小的帕累托前沿，且在长上下文任务中表现优于驱逐方法，显著减少解码时间缓存的大小，同时保持准确性。
🔬 方法详解
问题定义：本文旨在解决在处理长文本时，KV缓存大小超出内存限制的问题。现有方法通过丢弃标记来压缩缓存，但可能导致后续决策的不利影响。
核心思路：论文提出KV-Kaizen方法，利用学习选择器根据上下文动态生成每层的缓存配置，从而实现缓存压缩而不需驱逐标记。这种方法通过共享缓存、降低缓存精度和截断低秩表示来实现。
技术框架：整体架构包括一个选择器模块，该模块在推理前运行一次，根据上下文生成每层的缓存配置。选择器的设计考虑了深度、精度和秩三个维度，以适应不同的上下文需求。
关键创新：最重要的技术创新在于选择器的自适应性和局部组合能力，能够在不牺牲准确性的前提下实现大幅度的内存节省。这与现有方法的全局均匀处理形成鲜明对比。
关键设计：选择器的参数设置和损失函数经过精心设计，以确保在不同上下文下的有效性。具体细节包括缓存共享策略、精度控制和低秩表示的截断方法。选择器的运行效率也经过优化，以适应实时推理需求。
🖼️ 关键图片
📊 实验亮点
实验结果表明，KV-Kaizen在长上下文任务中实现了32倍的解码时间缓存缩小，同时保持了模型的准确性。与无学习和后处理基线相比，选择器在准确性与缓存大小的权衡中达到了帕累托前沿，4倍的缓存大小减少未导致7B参数模型的准确性下降。
🎯 应用场景
KV-Kaizen方法在自然语言处理、对话系统和长文本生成等领域具有广泛的应用潜力。通过有效压缩缓存，该方法能够在资源有限的环境中提升大语言模型的性能，降低计算成本，进而推动更大规模模型的预训练和应用。
📄 摘要（原文）
As the context size of text processed with an LLM grows, the size of KV caches can outstrip the memory allocated for the original model weights. This impacts LLM throughput negatively, since decoding is memory-bound and decode cost grows with cache size. Recent work alleviates this bottleneck by discarding the least relevant tokens. Eviction introduces a tension, since a one-off decision to discard content may prove detrimental later. Instead, we focus on alternative choices that can lead to cache compression without evicting tokens. We achieve this by learning a selector that is able to produce, based on context, a per-layer cache configuration towards an overall compression budget. The selector operates along three axes: sharing one cache across layers (depth), caching at fewer bits (precision), or truncating the low-rank latent cache representations (rank). We call the resulting method KV-Kaizen, for the many small per-layer choices it compounds. We observe that these interventions taken independently and uniformly over all layers limit achievable compression because they degrade accuracy. Crucially, composing them locally and adaptively to the context can instead preserve accuracy while achieving large memory savings. At inference, the selector runs once, before pre-fill. In evaluations on instruction following and reasoning tasks, our selectors reach the Pareto frontier of accuracy against cache size, against learning-free and post-hoc baselines. On long-context tasks, KV-Kaizen improves on eviction and can be composed with it, reaching a 32x smaller decode-time cache on a 14B model while preserving accuracy. A 4x cache size reduction incurs no accuracy degradation from 7B parameters up, and a compressed model is more accurate than a smaller uncompressed one with the same cache size. Together, these findings support pre-training large models and compressing them only afterwards.