EvolKV: Evolutionary KV Cache Compression for LLM Inference
EvolKV: Evolutionary KV Cache Compression for LLM Inference
发表时间: 2025-11 · EMNLP 2025 Findings
原文: https://aclanthology.org/2025.findings-emnlp.88
Bohan Yu (School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences; The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, CAS), Yekun Chai (ETH Zurich)
速读
一句话结论 提出 EvolKV 框架,将大语言模型各层的 KV Cache 预算分配转化为多目标优化问题,利用进化算法根据下游任务表现进行动态搜索,在极低显存预算下大幅超越了现有的启发式压缩基线。
要解决什么问题 现有的 KV Cache 压缩方法主要依赖静态的启发式规则,比如在所有层保留固定位置的 Token(如 StreamingLLM)、在所有层分配相同的 Cache 预算(如 SnapKV),或者按照预设的衰减比例从浅层到深层递减(如 PyramidKV)。这种做法的卡点在于,它忽略了 Transformer 不同层在处理信息时扮演的异质性角色,也没有建立缓存分配与下游任务实际表现之间的动态联系。先前的研究表明,大模型的不同层在信息处理的粒度和重要性上存在显著差异,例如某些中间层可能是上下文推理的计算瓶颈。强制套用统一或单调递减的规则,会导致模型在长上下文推理或复杂推理任务中,错误地丢弃了对当前任务至关重要的特征,从而在压缩率较高时出现严重的精度掉点和泛化能力失效。
怎么做的 核心思路是摒弃人工设计的分配规则,将每一层的 KV Cache 预算作为可学习变量,利用黑盒进化算法(CMA-ES,一种基于协方差矩阵自适应的随机优化算法)直接以下游任务的评测指标(如准确率、F1 分数)为导向,搜索出最优的逐层预算配置。这种方法能绕开启发式规则的盲区,自动挖掘出高度非均匀且打破常规金字塔假设的缓存分配模式(例如在中间层出现预算峰值),从而精准保留对任务最有用的上下文。为了解决搜索空间过大导致的优化不稳定问题,EvolKV 引入了分组机制。假设模型有 $L$ 层,目标平均每层缓存预算为 $c$,方法将相邻层划分为大小为 $n_g$ 的组。优化过程自底向上逐组进行,在优化当前组时,固定已优化组的配置,并保持未优化组为初始值。搜索的核心在于平衡任务表现与显存开销,其定义性目标函数为:
其中,$f(S)$ 是候选分配方案 $S$ 在下游任务上跑出的真实得分;$\lambda$ 是平衡系数;$\mathrm{CACHESCORE}(S, c)$ 是一个惩罚项,当方案的平均层预算 $\bar{k}$ 超过目标预算 $c$ 时输出低分,低于或等于 $c$ 时给予平滑奖励:
效果如何 实验在 Mistral-7B-Instruct(32K 上下文)和 Llama-3-8B-Instruct(8K 上下文)上进行。对比基线代表了三大主流路线:固定位置保留路线的 StreamingLLM、全局等量分配结合注意力淘汰路线的 SnapKV,以及逐层递减的金字塔分配路线的 PyramidKV。在 LongBench 的 16 个子任务中,EvolKV 在 128 到 2048 的目标预算下全面超越所有基线。最突出的量化结果是,在代码补全任务上,EvolKV 仅用原模型 1.5% 的 KV Cache 预算(平均每层 128 个 Token),其表现甚至反超了使用完整缓存的全量模型。在 GSM8K 数学推理任务中,EvolKV 在 128 预算下比最强基线准确率高出 7 个百分点,在 512 预算下保留了全量模型 95.7% 的性能(基线最高仅为 84.5%)。在 Needle-in-a-Haystack 检索任务中,比最优基线提升最高达 13%。该方法的一个显著工程优势是极低的搜索代价与极强的泛化性:仅需随机抽取 30 条样本进行搜索,且在极低预算(如 128)下搜索出的层级分配比例,可以直接等比放大应用到大预算(如 1024 或 2048)设置中,效果甚至优于直接在大预算下重新搜索。作者承认的局限性在于,当前的预算分配仅停留在层级别,尚未深入探索注意力头级别的细粒度预算分配。
主要贡献
在大型语言模型(LLMs)的推理过程中,键值(KV)缓存已成为提高效率的基石,它允许模型重用先前计算的隐藏状态从而减少冗余计算。然而,KV缓存的内存占用随输入序列长度呈线性增长,自注意力的二次复杂度使得在保留完整缓存时,长上下文推理变得极其缓慢。
现有的KV缓存压缩方法主要依赖于基于规则的启发式策略,例如跨层统一分配缓存或静态驱逐策略。这些方法未能考虑到特定层特征模式与任务性能之间的关键相互作用,从而导致泛化能力下降。为了解决这一问题,本文提出了EvolKV,这是一个用于逐层、任务驱动的KV缓存压缩的自适应框架,旨在联合优化内存效率和下游任务性能。
本文的核心创新点包括:
1. 首次提出将逐层KV缓存预算分配形式化为一个多目标黑盒优化问题。
2. EvolKV利用进化搜索算法动态配置各层的缓存预算,同时直接最大化下游任务的性能表现。
3. EvolKV在冻结参数的LLMs上运行,支持任意评估指标,无需进行模型微调或架构修改。
4. 广泛的实验表明,任务感知的KV缓存分配模式始终偏离传统的启发式规则,倾向于打破固定或金字塔规则的非均匀分布。在11个任务中,EvolKV在广泛的KV缓存预算下均优于所有基线方法,在代码补全任务中仅使用原始预算的$1.5\%$即可超越完整KV缓存设置的性能。
背景知识与设计原则
KV缓存压缩的现有范式:现有的KV缓存压缩方法主要分为三类:
1. 固定位置保留:在所有层中保留相同位置的KV缓存【11,Generating long sequences with sparse transformers+2019】【5,Longformer: The long-document transformer+2020】【47,Efficient streaming language models with attention sinks+2024】。
2. 相同预算分配:每层保留相同预算的KV缓存,但保留的位置因层而异,通常驱逐累积注意力权重最低的缓存【51,H2O: Heavy-hitter oracle for efficient generative inference of large language models+2023】【35,Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time+2023】【33,Snapkv: Llm knows what you are looking for before generation+2024】。
3. 金字塔形分配:KV缓存预算分配呈现金字塔模式,从底层到高层逐渐减少【6,Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling+2024】【48,Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference+2024】。
现有方法的局限性与动机:先前的研究表明,LLMs的不同层在信息处理中的重要性和处理粒度存在差异。然而,大多数现有的压缩方法忽略了这种异质性,并采用基于规则或启发式的策略进行KV缓存压缩(例如设置衰减系数来控制每层的预算,忽略了每层实际需要的缓存预算并不一定呈现单调递减的模式),这通常会导致次优的推理性能。这些观察结果凸显了为每一层单独自适应调整KV缓存预算的必要性。
方法细节
应对启发式策略的局限性:为了解决基于规则或启发式分配策略的局限性,引入了EvolKV,这是一个动态的、任务驱动的进化框架,通过利用来自下游任务的性能反馈,自适应地为每一层分配KV cache预算。
进化压缩的优化目标:进化算法生成候选解并评估其适应度(fitness),根据适应度反馈迭代改进搜索策略,从而逐步引导种群向更好的解演化。EvolKV将下游任务的性能反馈视为适应度,并利用进化算法指导逐层的KV cache压缩。具体而言,在具有$L$个transformer层的语言模型中,将第$i$层的KV cache预算表示为$k_i \in \mathbb{N}, \forall i \in \{1, \dots, L\}$。给定进化算法为下游任务$f(\cdot)$生成的一组候选压缩方案$\mathbb{S}$,目标是找到最优方案$\boldsymbol{S}^*$,在最大化任务性能的同时,最小化与目标平均KV cache预算$\boldsymbol{c}$的偏差:
其中$f(\boldsymbol{S})$是使用压缩方案$\boldsymbol{S} \in \mathbb{S}$获得的下游任务性能,超参数$\lambda > 0$用于平衡原始性能与缓存效率。由于下游性能指标(如准确率、F1、ROUGE)种类繁多且值域不同,采用了一个直接根据任务性能加权的缓存效率项,以确保可比性。缓存效率项$\mathrm{CACHESCORE}(\boldsymbol{S}, \boldsymbol{c}) \in [0, 1]$为平均每层缓存预算$\bar{k} = \frac{1}{L} \sum_{i=1}^L k_i^{(\boldsymbol{S})}$超过目标预算$\boldsymbol{c}$的方案分配较低的值,同时对保持在目标内的方案应用平滑折扣:
KV Cache预算的分组:为了提高优化效率,引入了组大小参数$n_g$,将KV cache预算$\mathbf{K} = \{k_1, k_2, \dots, k_L\}$划分为$J = \lceil L / n_g \rceil$个组,表示为$G = \{g_1, g_2, \dots, g_J\}$。每个组$g_j$包含一个连续的缓存预算子集,定义为$g_j = \{k_{(j-1)\cdot n_g + 1}, k_{(j-1)\cdot n_g + 2}, \dots, k_{\min(j\cdot n_g, L)}\}, \forall j \in \{1, 2, \dots, J\}$。为简单起见,假设总层数$L$能被组大小$n_g$整除,即$L = J \cdot n_g$。在此公式下,候选压缩方案$\mathbb{S}$在组级别应用,并表示为$\mathbb{S}_g$。基于下游任务性能为每组选择的最优方案表示为$S_g^*$。这种分组公式显著减少了搜索空间,并在进化搜索过程中促进了更稳定的优化动态。
进化压缩的迭代过程:KV cache预算优化是以分组方式进行的,具体过程如下述伪代码所示,从底层到顶层顺序进行。在优化每一组时,先前优化组的KV cache预算固定为其各自的最优方案$S_g^*$,而其余组保留其初始值。如果候选方案$S_g$达到比当前最佳更高的适应度分数$r$,则相应地更新当前组的KV cache预算。此过程迭代重复,直到所有组都被优化。
# 算法1:EvolKV优化过程
Require: 目标平均KV cache预算 c; 缓存预算效率权重 \lambda; 平滑因子 \gamma; 组大小 n_g; 模型层数 L; 最大迭代次数 M; CACHESCORE函数; 下游任务打分器 f(\cdot); 进化优化器 \mathcal{A}
Ensure: 全局最优的组KV cache预算 G^*
1: 初始化KV cache预算 \mathbf{K} = (c, \dots, c) \in \mathbb{N}^L
2: 将 \mathbf{K} 划分为 J = \lceil L / n_g \rceil 个组 G = \{g_1, \dots, g_J\}
3: G^* \gets G # 初始化组KV cache预算
4: F_{best} \gets -\infty # 初始化全局最佳适应度
5: for j \gets 1 to J do # 一次优化一个组
6: \mathcal{A}.\mathrm{INITIALIZE}(g_j) # 初始化优化器 \mathcal{A} 的参数
7: for m \gets 1 to M do
8: 从 \mathcal{A} 获取候选组压缩方案 \mathbb{S}_g
9: 评估每个 S_g \in \mathbb{S}_g 的适应度 r:
10: \tilde{G} = G^*,其中 g_j 被替换为 S_g
11: r \gets f(\tilde{G}) (1 + \lambda \mathrm{CACHESCORE}(S_g, c))
12: # 评估时,G^*中的g_j被S_g替换,其他组保持固定
13: 如果 r > F_{best},则使用 r 更新 F_{best},使用 \tilde{G} 更新 G^*
14: 使用 r 和 S_g 更新进化优化器 \mathcal{A}
15: end for
16: end for
17: return G^
KV Cache预算补全:为了确保评估的公平性,需要对总大小偏离目标的任何KV cache预算优化结果进行补全。具体来说,首先计算已达到的总KV cache预算$A = \sum_{i=1}^L k_i$与目标总预算$T = c \cdot L$之间的差异,表示为$\Delta_{\mathrm{cache}} = T - A$。然后,根据各层原来在$A$中的份额,将该差异按比例重新分配到各层。补全后的KV cache预算$B = \{b_1, b_2, \dots, b_L\}$,其中$b_i = \left\lceil k_i + \frac{k_i}{A} \cdot \Delta_{\mathrm{cache}} \right\rceil, i \in \{1, 2, \dots, L\}$。
实验环境
- 模型:使用两个开源模型:具有32K上下文长度的Mistral-7B-Instruct和具有8K上下文长度的Llama-3-8B-Instruct。
-
数据集:
- LongBench:包含16个代表性子数据集,涵盖6个主要任务类别:单文档QA、多文档QA、摘要、少样本学习、合成推理和代码补全。
- GSM8K:用于评估数学逻辑推理能力。
- Needle-in-a-Haystack (NIAH):用于评估长上下文检索能力。
- RULER:评估11个子数据集,涵盖检索、聚合和多跳追踪任务。
-
基线方法:StreamingLLM(固定位置)、SnapKV(统一预算分配)、PyramidKV(金字塔形规则分配)。
- 软件与超参数配置:采用协方差矩阵适应进化策略(CMA-ES【20,Reducing the time complexity of the derandomized evolution strategy with covariance matrix adaptation (cma-es)+2003】)作为进化优化器。窗口大小设为32,卷积核大小设为7并应用最大池化。EvolKV超参数固定为$\lambda = 0.3$,$\gamma = 0.2$,CMA-ES学习率$\sigma = 0.3$,组大小$n_g = 8$。CMA-ES的种群大小根据经验公式计算:$4 + \lfloor 3 \cdot \ln(n_g) \rfloor$。
实验结果
LongBench上的实验
- 实验内容:在Mistral-7B-Instruct上,仅使用NarrativeQA的30个随机样本在目标预算$c=128$下优化,然后将分配结果外推至256、512、1024和2048。在Llama-3-8B-Instruct上,从6个子数据集中各抽取5个样本在$c=128$下优化,外推至256、512、1024(2048进行单独优化)。评估时去除了所有训练样本。
- 实验结果:在Mistral-7B-Instruct上,EvolKV在所有评估的KV缓存预算下均达到最高平均性能。在MultiFieldQA-en、TriviaQA等任务上,EvolKV在特定预算下甚至超越了未压缩的完整模型。在Llama-3-8B-Instruct上,EvolKV同样在所有预算下表现卓越,特别是在$c=128$时,在TREC子集上比最强基线高出7.69个百分点。
- 分析结论:EvolKV在低预算(128、256)和高预算(512-2048)下均保持优势。在$c=128$(仅占上下文长度1.5%)的极端压缩下,EvolKV仍能超越完整模型,而其他基线均失败。在$c=128$下优化的预算能平滑泛化到更大的预算,表明进化搜索捕获了稳定、与任务对齐的重要性模式,而不是过度拟合。
GSM8K上的实验
- 实验内容:量化模型在不同KV缓存预算下的逻辑推理能力。随机抽取30个GSM8K训练实例,在$c=128$下以准确率为目标进行优化,并放大至256和512。
- 实验结果:在Llama-3-8B-Instruct上,EvolKV在$c=128, 256, 512$时分别比最强竞争对手提高至少7.28、2.05和7.58的准确率。在$c=512$时,EvolKV保留了完整模型$95.7\%$的性能,而最强基线仅达到$84.5\%$。
- 分析结论:仅在$c=128$处优化的预算能有效转移到更大的预算。StreamingLLM在此任务上表现不佳,表明固定位置策略对于面向推理的任务是次优的。进化分配揭示了固定启发式忽略的特定层缓存需求。
NIAH和RULER上的实验
- 实验内容:在NIAH上评估长上下文检索能力。优化时使用不超过35个平均得分低于60的实例,目标预算$c=128$。随后将NIAH中优化的分配方案应用于RULER基准测试(外推至1024)。
- 实验结果:在NIAH上,EvolKV在Llama-3-8B-Instruct上比基线提高了4个百分点以上,在Mistral-7B-Instruct上提高了13个百分点以上(如图7)。在RULER上,EvolKV在128和1024预算下的平均得分均优于所有基线,在Mistral上提升达0.99分,在Llama-3上提升达3.6分。
- 分析结论:EvolKV有效地探索并利用了模型在长上下文检索中的潜在逐层分配策略,且优化的预算可以有效转移到其他基准评估中,展示了强大的泛化能力。
补充细节
优化过程中的下游任务性能:在Mistral-7B-Instruct上随机抽取30个NarrativeQA实例进行实验。结果表明,随着迭代次数的增加,模型在训练数据上的性能稳步提升。这表明原始的均匀分配留有大量改进空间,简单的启发式规则不足以找到最优分布。
优化预算分配的讨论:实验揭示EvolKV发现了与启发式方法完全不同的分配模式。在模型中间层出现了一致的预算峰值,表明这些层是上下文推理的计算瓶颈。这种非直觉模式在不同任务和预算中持续存在。在低预算时(如$c=128$),优化倾向于将资源集中在少数层;而在高预算时,分配在模型中更为分散。任务优化的分配始终偏离固定或金字塔规则,并未呈现单调递减模式,有些底层获得的分配极少,而高层获得了更大的预算。
组大小的影响:在优化过程中,测试了组大小$n_g \in \{2, 4, 8, 16, 32\}$。下游任务性能通常随着组大小的增加而提高,在$n_g=8$时达到峰值,之后显著下降。组太小会导致过拟合有限的优化数据,组太大则阻碍各层之间的有效预算分配。因此,选择$n_g=8$作为最佳权衡。
EvolKV的鲁棒性分析:在Llama-3-8B-Instruct上进行三次独立优化,总平均分为35.88,各次偏差在0.1以内,标准差为0.078,表现出一致的稳定性。此外,测试了不同的训练数据组合,在剔除训练数据后,EvolKV在$c=512$和$1024$等设置下仍优于基线,突显了其跨任务的强适应性。
EvolKV的泛化分析:比较了直接缓存扩展(即前述的补全方法)与直接在对应目标预算下优化。当$c \geq 512$时,基于扩展的方法始终优于直接优化。这表明严格预算下的优化能有效揭示各层的实际需求,并能推广到更高预算。此外,将NIAH数据集上$c=128$优化的预算应用于LongBench,EvolKV在$c=128$和256下均优于基线,证明了其跨数据集的泛化能力。
不同系列模型的实验:在Qwen系列模型(Qwen2.5-1.5B-Instruct和Qwen2.5-3B-Instruct【2,Qwen technical report+2023】)上使用30个NarrativeQA实例在$c=128$下进行优化。剔除训练数据后,EvolKV在单文档和多文档QA任务中均优于所有压缩方法,而PyramidKV表现明显较差,证明了EvolKV跨模型系列的一致优势。
专门任务优化的实验:为了评估EvolKV在专门任务上的优化性能,分别在单文档和多文档QA任务上进行了KV cache预算优化(各抽取30个实例)。在剔除训练数据后,EvolKV在$c=128$时在两项任务上均取得了最高平均分,超越了所有基线。
推理时间和内存成本评估:在Mistral-7B-Instruct上使用FlashAttention【13,Flashattention: Fast and memory-efficient exact attention with io-awareness+2022】评估推理时间和峰值内存($c=128$)。与其他压缩方法相比,EvolKV在不同生成长度下的推理时间(prefill + decoding)变化微乎其微。同时,其峰值内存使用量与其他压缩方法相当,且比完整缓存显著降低了内存消耗。
结论
本文介绍了EvolKV,这是一个任务驱动的框架,利用进化算法优化LLMs中逐层的KV缓存预算。与基于规则或启发式的方法不同,EvolKV在不修改模型参数且仅需少量标记示例的情况下,直接最大化下游性能。广泛的实验表明,任务感知的进化缓存预算分配揭示了现有方法忽略的潜在层重要性模式,并在严格的缓存限制下提供了最先进的性能,同时能优雅地扩展到更大的目标缓存预算。因此,EvolKV为下游任务中的高效推理提供了一个实用的即插即用解决方案。未来的工作将研究不同分词方案以及无分词方法下KV缓存压缩的鲁棒性,并探索在注意力头级别上的预算分配。
💬 评论讨论
欢迎在这里分享您的想法和见解!