ReFreeKV: 迈向无阈值KV缓存压缩

摘要

为降低大语言模型推理过程中的内存消耗，已有少量方法被提出用于KV缓存剪枝。尽管这些技术能在许多数据集上实现无损内存压缩，但它们通常依赖于一个尚未被充分强调的条件：即需预先设定一个输入/领域特定的KV缓存预算阈值，才能达到最优性能。然而，这种输入敏感的设计在实际场景中可能受到显著限制，因为开放领域的输入涵盖不同领域、长度和难度级别，阈值选择缺乏明确边界。因此，这种对输入敏感阈值的依赖可能成为根本性局限，导致任意输入的性能大幅下降。在本工作中，我们提出一项新目标，旨在解除稳健KV压缩中的阈值约束，倡导采用"无阈值"方法，在保持完整缓存性能的同时自适应调整预算分配。随后，我们提出一种创新方法ReFreeKV，作为该目标的首个实例化方案。在涵盖不同上下文长度、任务类型和模型规模的13个数据集上进行的广泛实验，验证了其有效性与高效性。我们的代码已公开于https://github.com/Patrick-Ni/ReFreeKV。

English

To reduce memory consumption during LLM inference, a handful of methods have been proposed for KV cache pruning. While these techniques can accomplish lossless memory reduction on many datasets, they often hinge on an under-emphasized condition: an input/domain-specific threshold for KV cache budget needs to be pre-determined to achieve the optimal performance. However, such input-sensitive design may be considerably limited in real-world scenarios, as open-domain inputs span diverse domains, lengths and difficulty levels, without clear boundaries for threshold selection. As a result, the dependence of such input-sensitive threshold can be a fundamental limitation that causes large degradation on arbitrary inputs. In this work, we propose a new objective that lifts the threshold constraints for robust KV compression, advocating for "threshold-free" methods that adaptively adjust budget allocation while preserving full-cache performance. We then propose a novel method, ReFreeKV, serving as the first instantiation of this objective. Extensive experiments across 13 datasets with diverse context lengths, task types, and model sizes demonstrate its efficacy and efficiency. Our code is publicly released at https://github.com/Patrick-Ni/ReFreeKV.