论文

行为保留 KV 缓存压缩

Behavior-Preserving KV Cache Compression

模型推理KV Cache

摘要

KV 缓存是大语言模型长上下文推理和长格式生成的主要瓶颈。现有的无需训练驱逐策略在很大程度上依赖于代理重要性信号,例如注意力质量,来决定保留哪些过去的token。我们认为缓存压缩应该保留全缓存模型的预测行为,保留其删除将大大改变模型输出分布的条目。我们提出了行为保留 KV 缓存压缩,这是一个无需训练框架,它通过估计由删除引起的压缩缓存 logits 并评估生成的 KL 到全缓存下一个token分布来对候选逐出进行评分。使用预驱逐前向统计,该方法避免为每个候选者运行单独的屏蔽前向传递。在不同的架构以及预填充时间和生成时间压缩中,我们的方法在匹配的保留 KV 预算下比轻量级基于注意力的启发式方法在下游任务质量方面带来了显着的收益,并且在积极压缩下收益最大。它通过额外的压缩时间计算实现了这些收益,同时在我们评估的设置中保留了相对于全缓存推理的端到端加速。