学习从键值缓存中淘汰

随着大型语言模型（LLMs）规模的增长，高效的推理变得具有挑战性，主要由于自回归键值（KV）缓存的内存需求。现有的淘汰或压缩方法虽然降低了成本，但依赖于启发式方法，如近期性或过去的注意力分数，这些方法仅作为未来令牌效用的间接代理，并引入了计算开销。我们将KV缓存淘汰重新定义为强化学习（RL）问题：通过预测未来解码中令牌的有用性来对令牌进行排序。为此，我们引入了KV策略（KVP），一种轻量级的每头RL代理框架，使用预计算的生成轨迹进行训练，仅使用键和值向量。每个代理学习一个由未来效用引导的专业化淘汰策略，评估所有缓存预算下的排序质量，无需修改底层LLM或额外推理。在长上下文基准RULER和多轮对话基准OASST2-4k上进行评估，KVP显著优于基线。此外，在标准下游任务（如LongBench、BOOLQ、ARC）上的零样本测试表明，KVP在其训练分布之外以及更长的上下文长度上具有良好的泛化能力。这些结果表明，学习预测未来令牌效用是适应性KV缓存管理的强大且可扩展范式。