大语言模型长上下文任务中的自适应 KV 缓存合并
计算与语言
2024-07-23 v2
摘要
如何高效地服务大型语言模型 (LLM) 已成为一个迫切的问题，因为其在自回归生成过程中的计算成本极大。为缓解计算成本，LLM 常采用 KV 缓存技术来提高生成速度。尽管提高了计算效率，但 KV 缓存的存储需求在长上下文场景下仍然显著，导致显存消耗大。现有的 KV 缓存驱逐方法常因驱逐引入的信息损失而在长上下文场景下降低 LLM 的性能。本文提出一种新颖的 KV 缓存合并方法，称为 KVMerger，通过自适应 KV 缓存压缩实现在受限内存预算下的长上下文任务，无显著性能下降。我们的方法灵感来自一个有趣的观察：在单个序列中，键状态在 token 级别上具有高度相似性。为促进合并，我们开发了一种有效且简洁的合并集合识别算法，以识别适合合并的 KV 状态。我们的合并集合识别算法激发了第二个观察：从相似性角度看，KV 缓存稀疏性独立于数据集，始终在模型级别保持持久。随后，我们提出了基于高斯核加权的合并算法，对每个合并集合内的所有状态进行选择性合并。我们进行了大量实验，以证明在受限内存预算下，KVMerger 对长上下文任务的有效性，适用于 Llama2-7B-chat 和 Llama2-13B-chat 等模型。通过使用 LongBench 和 ZeroScroll 基准测试，我们将方法与其他 KV 缓存压缩技术（包括 H2O 和 CaM）进行比较，显示在 50% 和 35% 的 KV 缓存预算下，我们的方法在各种任务中实现了优异的性能。
引用
@article{arxiv.2407.08454,
title = {Model Tells You Where to Merge: Adaptive KV Cache Merging for LLMs on Long-Context Tasks},
author = {Zheng Wang and Boxiao Jin and Zhongzhi Yu and Minjia Zhang},
journal= {arXiv preprint arXiv:2407.08454},
year = {2024}
}