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

AIGC动态13小时前发布 almosthuman2014

45 0 0

文章摘要

键值缓存（KV cache）是大模型快速运行的核心技术，但存在输入文本越长、所需存储空间越大且处理长文本变慢的问题。现有 KV cache 压缩方法主要依赖基于规则的启发式方法，未考虑 transformer 层在信息处理中的不同功能角色以及缓存与任务性能之间的动态关系。

为解决这些问题，研究人员引入了 EvolKV，这是一个进化框架，能自适应地在 transformer 层之间分配 KV cache 预算。它将每层 KV cache 预算制定为优化变量，分组后采用进化算法迭代搜索，实现与不同层的不同贡献相匹配的细粒度、性能感知分配。与刚性启发式方法相比，EvolKV 为下游任务目标导向的逐层 KV cache 预算分配提供了灵活有效的机制。

EvolKV 详解如下：以进化算法的工作原理为基础，将下游任务的性能反馈当作适应度分数，指导每一层的 KV cache 压缩。其优化目标是找到既能最大化任务性能，又能接近目标平均 KV cache 预算的最优方案。为提高优化效率，研究人员引入组大小参数，将 KV cache 预算分组。KV cache 预算优化以分组方式从底层到顶层依次进行，若候选方案适应度分数更高，则更新当前组的 KV cache 预算。最后，对总大小偏离目标的 KV cache 预算优化结果进行补全。

实验结果显示，EvolKV 在多个基准测试中表现出色。在 LongBench 上，EvolKV 在 Mistral – 7B – Instruct 和 Llama – 3 – 8B – Instruct 模型的所有评估 KV cache 预算中，始终获得最高平均性能，优于所有基于规则的基线方法，甚至在某些预算下超越完整模型。在 GSM8K 上，EvolKV 在两个模型上都始终优于基线方法，使用减少的缓存预算就能达到接近完整模型的性能。在 NIAH 和 RULER 上，EvolKV 相比基线方法取得显著提升，有效探索并利用了模型在长上下文检索中的潜在层级 KV cache 分配，且优化的 KV 预算可有效迁移到其他基准评估中，证明了其强大的泛化能力、长上下文检索和推理能力。