用强化学习学会预测令牌未来价值,动态优化大模型缓存淘汰
Learning to Evict from Key-Value Cache

- 将缓存淘汰转化为强化学习问题,按未来使用价值排序令牌
- 在128K上下文任务中超越强基线,零样本泛化到更长序列
- 无需修改模型,轻量级代理可适配不同场景的缓存管理
大型语言模型(LLMs)规模增大导致推理效率受限,主要因自回归键值(KV)缓存的内存开销。现有淘汰或压缩方法依赖启发式规则(如访问频次或历史注意力分数),仅间接反映令牌未来价值且引入额外计算开销。本文将KV缓存淘汰重构为强化学习(RL)问题:学习对令牌按其未来解码有用性进行排序。为此提出轻量级每头强化学习代理框架KVP,基于预计算生成轨迹,仅使用键和值向量训练。每个代理学习专用淘汰策略,奖励函数综合未来效用,评估全缓存预算下的排名质量,无需修改底层模型或额外推理。在两个模型族上测试,涵盖128K令牌长上下文基准RULER及多轮对话基准OASST2-4k,KVP显著优于强基线。零样本测试在标准下游任务(BoolQ、LongBench段落检索、GovReport)中也表现优异,证明其泛化能力与超长序列适应性。结果表明,学习预测未来令牌价值是高效可扩展的自适应缓存管理新范式。
原文摘要 · Abstract (English)
The growing size of Large Language Models (LLMs) makes efficient inference challenging, primarily due to the memory demands of the autoregressive Key-Value (KV) cache. Existing eviction or compression methods reduce cost but rely on heuristics, such as recency or past attention scores, which serve only as indirect proxies for a token's future utility and introduce computational overhead. We reframe KV cache eviction as a reinforcement learning (RL) problem: learning to rank tokens by their predicted usefulness for future decoding. To this end, we introduce KV Policy (KVP), a framework of lightweight per-head RL agents trained on pre-computed generation traces using only key and value vectors. Each agent learns a specialized eviction policy guided by a holistic reward, derived from future utility, that evaluates the quality of the ranking across all cache budgets, requiring no modifications to the underlying LLM or additional inference. Evaluated across two model families on the long-context benchmark RULER (up to 128K tokens) and the multi-turn dialogue benchmark OASST2-4k, KVP significantly outperforms strong baselines. Zero-shot tests on standard downstream tasks (BoolQ, LongBench passage retrieval, GovReport) further show that KVP generalizes beyond its training distribution and to considerably longer sequence lengths. These results demonstrate that learning to predict future token utility is a powerful and scalable paradigm for adaptive KV cache management.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。