通过合作博弈优化注意力头缓存分配，提升大模型推理效率
CoKV: Optimizing KV Cache Allocation via Cooperative Game
- 将注意力头间的协作关系建模为合作博弈，动态评估其联合贡献
- 在LongBench上使用Llama-3-8B-Instruct和Mistral-7B实现最佳性能
- 适合关注大模型推理优化与缓存管理的研究者和工程师
大语言模型在人类生活的多个方面取得了显著成功，但部署时面临存储键值对（KV）所需的大量内存消耗问题。现有方法多通过独立评估各注意力头的重要性来分配缓存预算，忽略了头之间的协同作用，可能导致对实际性能影响的误判。为此，我们提出CoKV，将模型推理中注意力头的协作关系建模为合作博弈，通过分析每个头在博弈中的贡献，更精准地分配缓存预算。大量实验表明，CoKV在LongBench基准上使用LLama-3-8B-Instruct和Mistral-7B模型均达到当前最优表现。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。