KV Cache预算降至1.5%！他们用进化算法把大模型内存占用砍下来了|焦点精选

机器之心报道

编辑：张倩

图源：https://x.com/rohanpaul_ai/status/1966820906916815156

(资料图片仅供参考)

键值缓存（KV cache）已经成为大模型快速运行的核心技术，它就像一个「记忆库」，能够保存之前计算过的结果并重复使用，这样就不用每次都重新计算同样的内容。

但是，这个记忆库有个问题：输入的文本越长，需要的存储空间就越大，而且模型处理长文本时会变得非常慢。

为了应对这些挑战，现有的 KV cache 压缩方法主要依赖基于规则的启发式方法。当前的方法可以归类为三种范式：

• 跨所有层的固定位置信息保留
• 基于注意力权重的均匀层级分配淘汰机制
• 具有预定义深度衰减的金字塔策略

虽然这些方法在降低内存占用方面有效，但它们未能考虑两个关键问题：

• transformer 层在信息处理中的不同功能角色
• 缓存与任务性能之间的动态关系

仅依赖基于规则的 KV cache 预算分层分配，可能导致任务相关信息无法被最优地保留。

针对这些限制，来自中国科学院大学、中国科学院自动化研究所的 Bohan Yu 和苏黎世联邦理工学院的 Yekun Chai 受到（Chai 等，2022）的启发，采用进化算法直接基于任务性能搜索最优的 KV cache 分配

• 论文标题：EvolKV: Evolutionary KV Cache Compression for LLM Inference
• 论文链接：https://arxiv.org/pdf/2509.08315

他们引入了EvolKV，这是一个进化框架，能够自适应地在 transformer 层之间分配 KV cache 预算，如图 1 所示。它将每层 KV cache 预算制定为优化变量，将其分为组，并采用进化算法迭代搜索能够直接最大化下游任务适应度得分的组别配置。通过将任务驱动优化与层特定缓存剪枝相结合，EvolKV 实现了与不同层的不同贡献相匹配的细粒度、性能感知分配。

与刚性启发式方法相比，EvolKV 为以下游任务目标为导向的逐层 KV cache 预算分配提供了一个灵活而有效的机制。首先，它将层 / 组级缓存预算制定为可学习参数，其中，作者将层分组为优化单元以实现高效搜索。然后，它们使用黑盒进化优化方法直接最大化下游任务的性能。

通过这种方式，他们的方法能够实现任务感知的细粒度缓存分配，自动适应每个组或层的功能贡献。具体而言，它能够适应多样化的评估标准，如准确率和 F1 分数，并在没有预定义假设的情况下发现非均匀分布（即偏离启发式固定长度或金字塔模式的模式）。

作者在 Mistral 7B-Instruct 和 Llama-3-8B-Instruct 上进行了全面实验，在四个不同的基准测试（十一项任务）上评估 EvolKV，涵盖长上下文检索、长上下文推理和数学任务。结果表明，任务优化的 KV cache 分配产生了一致的改进：

• 在 Needle-in-a-Haystack 基准测试中，EvolKV 比最佳基线提高了多达 13%
• 在 RULER 基准测试中，EvolKV 比最强基线提升了多达 3.6%
• 在 LongBench 评估中，它在广泛的目标 KV cache 预算范围内（从 128 到 2048）始终优于所有基线方法，并且在仅使用完整模型 1.5% 的 KV cache 预算的情况下，其性能显著超过了完整模型
• 对于 GSM8K，在 128 KV cache 预算下，EvolKV 比最强基线在准确率上提高了多达 7 个百分点，保持了完整模型性能的 95.7%，而最强基线在 512 KV cache 预算下仅保持了 84.5%。

EvolKV 详解

EvolKV 是一个动态的、任务驱动的进化框架，通过利用下游任务的性能反馈来自适应地为每一层分配 KV cache 预算。图 2a 展示了 EvolKV 与其他方法之间预算分配的比较。

进化压缩的优化目标

KV Cache 预算的分组

进化压缩的迭代过程

本文中的 KV cache 预算优化以分组方式进行，如算法 1 所示，从底层到顶层依次进行。在优化每个组时，先前优化组的 KV cache 预算被固定为各自的最优方案，而其余组保持其初始值。如果候选方案 S_g 获得比当前最佳方案更高的适应度分数 r，则相应地更新当前组的 KV cache 预算。这个过程反复进行，直到所有组都被优化。

KV 缓存预算补全

为了确保评估的公平性，作者对总大小偏离目标的 KV cache 预算优化结果进行补全。具体而言，他们首先计算实际总 KV cache 预算与目标总预算 T = c・L 之间的差异，记为∆_cache = T − A。然后根据各层在 A 中的原始占比，将这个差异按比例重新分配到各层。补全后的 KV cache 预算为 B = {b_1, b_2, . . . , b_L}，其中，i ∈ {1, 2, . . . , L}。

实验结果

在 LongBench 上的结果

表 1 报告了在 Mistral-7B-Instruct 上使用 16 个 LongBench 子数据集的评估结果，所有训练样本已被移除。在所有评估的 KV cache 预算中，EvolKV 始终获得最高的平均性能，优于所有基于规则的基线方法。此外，在包括 MultiFieldQA-en、2WikiMultihopQA、MuSiQue、TriviaQA 和 PassageRetrieval-en 在内的几个子数据集上，EvolKV 不仅与未压缩的完整模型保持竞争力，甚至在某些 KV cache 预算下超越了完整模型。

表 2 展示了 Llama-3-8B-Instruct 上的类似结果，同样排除了训练样本。EvolKV 在所有 KV cache 预算下都表现出了优异的性能。值得注意的是，在缓存预算为 128 时，EvolKV 在 TREC 子集上比最强基线高出 7.69 个百分点，突出了其对多样化下游任务的强适应性。

在 GSM8K 上的结果

图 3b 展示了 EvolKV 为 Llama-3-8B-Instruct 优化的 KV cache 预算分配。表 3 报告了 Llama-3-8B-Instruct 和 Mistral-7B-Instruct 两个模型对应的测试集准确率。在所有配置中，EvolKV 在两个模型上都始终优于基线方法。具体而言，在 Llama-3-8B-Instruct 上，它相比最强竞争对手取得了显著改进，在 KV 缓存预算为 128、256 和 512 时，准确率分别至少提升了 7.28、2.05 和 7.58 个百分点。值得注意的是，EvolKV 使用减少的缓存预算（c = 512）就达到了完整模型性能的 95.7%，显著优于所有基线方法，其中最佳基线结果仅达到 84.5%。

在 NIAH 和 RULER 上的结果

作者在 NIAH 上评估了 EvolKV 以及所有基线模型的长上下文检索能力。图 7 展示了这部分评估的结果：与基线方法相比，EvolKV 在 Llama3-8B-Instruct 上取得了超过 4 个百分点的改进，在 Mistral-7B-Instruct 上取得了超过 13 个百分点的显著提升。这些结果表明，EvolKV 有效地探索并利用了模型在长上下文检索中的潜在层级 KV cache 分配。

作者在 RULER 基准测试上进一步评估了 NIAH 中优化的 KV 缓存分配。如表 4 所示，EvolKV 在平均得分上始终优于所有基线方法，在 Mistral-7B-Instruct 上提升了多达 0.99 分，在 Llama-3-8B-Instruct 上提升了 3.6 分。这些结果进一步证明了 EvolKV 强大的泛化能力、长上下文检索和推理能力，因为优化的 KV 预算可以有效地迁移到其他基准评估中，这表明 EvolKV 揭示了潜在的层级分配策略。

更多细节请参见原论文。