ReST-KV：具有逐层输出重建和时空平滑功能的鲁棒 KV 缓存驱逐

由于键值 (KV) 缓存的内存需求不断增加，尤其是长序列，大语言模型 (LLM) 在高效生成推理方面面临着越来越大的挑战。现有的驱逐方法通常保留具有高注意力权重的 KV 对，但忽略了 token 删除引起的注意力重新分配的影响，以及 KV 选择中的时空动态。在本文中，我们提出了 ReST-KV，一种鲁棒的 KV 驱逐方法，结合了分层输出重建和时空平滑，为 KV 缓存驱逐任务提供了更全面的视角。具体来说，ReST-KV 将 KV 缓存驱逐作为一个优化问题，通过有效的逐层重建来最小化输出差异。通过直接建模每个标记的删除如何影响模型输出，我们的方法自然地捕获注意力重新分配效应，超越对原始注意力权重的简单依赖。为了进一步增强鲁棒性，我们设计了指数移动平均平滑来处理时间变化，并设计了基于自适应窗口的机制来捕获空间模式。我们的方法 ReST-KV 显着提高了长上下文基准测试的性能。它在 LongBench 上超过了最先进的基线 2.58%，在 RULER 上超过了最先进的基线 15.2%。此外，ReST-KV 在 Needle-in-a-Haystack 和 InfiniteBench 上始终优于现有方法，同时在 128k 上下文长度下实现了 10.61$\times$ 的解码延迟显着降低。该代码可在 https://github.com/an-yongqi/rest-kv 上公开获取，以促进可重复性和进一步研究。