CoKV: Optimizing KV Cache Allocation via Cooperative Game

📄 arXiv: 2502.17501v1 📥 PDF

作者: Qiheng Sun, Hongwei Zhang, Haocheng Xia, Jiayao Zhang, Jinfei Liu, Kui Ren

分类: cs.LG, cs.AI

发布日期: 2025-02-21

💡 一句话要点

CoKV：通过合作博弈优化KV缓存分配，提升LLM长文本处理能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 KV缓存 缓存分配 合作博弈 注意力机制 长文本处理 模型推理 Shapley值

📋 核心要点

  1. 现有KV缓存分配方法忽略了注意力头之间的合作关系，导致缓存分配与实际性能贡献不符。
  2. CoKV将注意力头之间的合作建模为合作博弈，评估每个头在合作中的贡献，从而更有效地分配缓存预算。
  3. 实验结果表明，CoKV在LongBench基准测试中，使用LLama-3-8B-Instruct和Mistral-7B模型时，达到了最先进的性能。

📝 摘要（中文）

大型语言模型(LLMs)在人类生活的各个方面都取得了显著的成功。然而，部署这些模型的主要挑战之一是存储键值对(KV)所需的大量内存消耗，这带来了巨大的资源需求。最近的研究主要集中在KV缓存预算分配上，一些方法提出通过评估各个注意力头的相对重要性来进行head级别的预算分配。然而，这些方法独立地评估head的重要性，忽略了它们在模型中的合作贡献，这可能导致与它们对模型性能的真实影响产生偏差。鉴于此限制，我们提出CoKV，一种新颖的方法，将模型推理中head之间的合作建模为一个合作博弈。通过评估合作博弈中每个head的贡献，CoKV可以更有效地分配缓存预算。大量实验表明，在使用LLama-3-8B-Instruct和Mistral-7B模型时，CoKV在LongBench基准测试中实现了最先进的性能。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）部署时，KV缓存占用大量内存的问题。现有方法，如head级别的预算分配，通常独立评估每个注意力头的重要性，忽略了它们之间的合作关系，导致缓存分配效率低下，无法充分利用有限的资源。这种独立评估方式无法准确反映每个head对模型性能的真实贡献。

核心思路：CoKV的核心思路是将模型推理过程中注意力头之间的合作关系建模为一个合作博弈。通过合作博弈理论，可以量化每个注意力头在整体合作中的贡献，从而更准确地评估其重要性。基于这种合作贡献的评估，CoKV能够更有效地分配KV缓存预算，提升模型性能。

技术框架：CoKV的技术框架主要包含以下几个阶段：1) 合作博弈建模：将LLM推理过程中的注意力头之间的交互建模为一个合作博弈，每个注意力头作为一个参与者。2) 贡献评估：利用合作博弈中的Shapley值等概念，评估每个注意力头在合作中的贡献。Shapley值能够公平地衡量每个参与者对整体合作的贡献。3) 缓存预算分配：根据每个注意力头的贡献，动态地分配KV缓存预算。贡献越大的注意力头，分配的缓存预算越多。4) 模型推理：利用分配好的缓存预算进行模型推理。

关键创新：CoKV的关键创新在于将合作博弈理论引入到KV缓存分配中。与现有方法独立评估注意力头的重要性不同，CoKV考虑了注意力头之间的合作关系，从而更准确地评估每个头的贡献。这种基于合作的评估方式能够更有效地利用有限的缓存资源，提升模型性能。

关键设计：CoKV的关键设计包括：1) 合作博弈的定义：如何将LLM推理过程中的注意力头交互转化为一个合适的合作博弈模型。2) 贡献评估方法：选择合适的合作博弈理论工具（如Shapley值）来评估每个注意力头的贡献。3) 缓存预算分配策略：设计一种有效的缓存预算分配策略，将注意力头的贡献转化为实际的缓存分配方案。具体参数设置和损失函数等细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CoKV在LongBench基准测试中，使用LLama-3-8B-Instruct和Mistral-7B模型时，实现了最先进的性能。实验结果表明，CoKV能够显著提升LLM在长文本处理任务上的性能，证明了其在KV缓存优化方面的有效性。具体的性能提升数据在论文中进行了详细展示。

🎯 应用场景

CoKV具有广泛的应用前景，可应用于各种需要部署大型语言模型的场景，例如智能客服、机器翻译、文本生成等。通过优化KV缓存分配，CoKV可以降低LLM的内存需求，使其能够在资源受限的设备上运行，并提高推理效率。未来，CoKV可以进一步扩展到其他类型的模型和任务中，为LLM的部署和应用提供更有效的解决方案。

📄 摘要（原文）

Large language models (LLMs) have achieved remarkable success on various aspects of human life. However, one of the major challenges in deploying these models is the substantial memory consumption required to store key-value pairs (KV), which imposes significant resource demands. Recent research has focused on KV cache budget allocation, with several approaches proposing head-level budget distribution by evaluating the importance of individual attention heads. These methods, however, assess the importance of heads independently, overlooking their cooperative contributions within the model, which may result in a deviation from their true impact on model performance. In light of this limitation, we propose CoKV, a novel method that models the cooperation between heads in model inference as a cooperative game. By evaluating the contribution of each head within the cooperative game, CoKV can allocate the cache budget more effectively. Extensive experiments show that CoKV achieves state-of-the-art performance on the LongBench benchmark using LLama-3-8B-Instruct and Mistral-7B models.