CoKV：通过合作博弈优化 KV 缓存分配
机器学习
2025-02-26 v1 人工智能
摘要
大型语言模型(LLM)在人类生活的各个方面取得了显著成功。然而，部署这些模型的主要挑战之一是存储键值对 (KV) 需要大量内存消耗，这带来了巨大的资源需求。最近的研究集中在 KV 缓存预算分配上，有几种方法提出通过评估单个注意力头的重要性来进行头级别的预算分配。然而，这些方法独立评估头的重要性，忽略了它们在模型内的合作贡献，这可能导致其偏离对模型性能的真实影响。鉴于这一局限性，我们提出了 CoKV，这是一种将模型推理中头之间的合作建模为合作博弈的新方法。通过评估合作博弈中每个头的贡献，CoKV 能够更有效地分配缓存预算。大量实验表明，使用 LLama-3-8B-Instruct 和 Mistral-7B 模型，CoKV 在 LongBench 基准测试上取得了 SOTA 性能。
引用
@article{arxiv.2502.17501,
title = {CoKV: Optimizing KV Cache Allocation via Cooperative Game},
author = {Qiheng Sun and Hongwei Zhang and Haocheng Xia and Jiayao Zhang and Jinfei Liu and Kui Ren},
journal= {arXiv preprint arXiv:2502.17501},
year = {2025}
}