arXiv:2605.08840cs.CL2026-05中稿 · ICLR被引 9

提出ReST-KV,通过重建与时空平滑提升长文本生成的缓存效率。

ReST-KV: Robust KV Cache Eviction with Layer-wise Output Reconstruction and Spatial-Temporal Smoothing

论文配图:ReST-KV: Robust KV Cache Eviction with Layer-wise Output Reconstruction and Spatial-Temporal Smoothing
图 1 · 摘自论文原文
  • 基于逐层输出重建优化缓存淘汰,捕捉注意力重分配效应
  • 在LongBench和RULER上分别提升2.58%和15.2%,延迟降低10.61倍
  • 适合长序列生成、低延迟推理场景,代码开源可复现

大语言模型在生成推理中面临日益增长的键值(KV)缓存内存压力,尤其在长序列场景下。现有淘汰方法通常仅保留高注意力权重的KV对,却忽略了令牌删除引发的注意力重分配影响,以及KV选择中的时空动态特性。本文提出ReST-KV,一种结合层间输出重建与时空平滑的鲁棒缓存淘汰方法。该方法将缓存淘汰建模为最小化输出差异的优化问题,通过直接建模每个令牌移除对模型输出的影响,自然捕捉注意力重分配效应,突破了对原始注意力权重的简单依赖。为进一步增强鲁棒性,设计指数移动平均平滑以应对时间波动,并引入自适应窗口机制捕捉空间模式。ReST-KV在长上下文基准测试中显著提升性能,在LongBench上超越当前最优基线2.58%,在RULER上提升15.2%。同时在Needle-in-a-Haystack和InfiniteBench上持续优于现有方法,且在128k上下文长度下实现10.61×的解码延迟降低。代码已公开于https://github.com/an-yongqi/rest-kv,支持复现与进一步研究。

原文摘要 · Abstract (English)

Large language models (LLMs) face growing challenges in efficient generative inference due to the increasing memory demands of Key-Value (KV) caches, especially for long sequences. Existing eviction methods typically retain KV pairs with high attention weights but overlook the impact of attention redistribution caused by token removal, as well as the spatial-temporal dynamics in KV selection. In this paper, we propose ReST-KV, a robust KV eviction method that combines layer-wise output Reconstruction and Spatial-Temporal smoothing to provide a more comprehensive perspective for the KV cache eviction task. Specifically, ReST-KV formulates KV cache eviction as an optimization problem that minimizes output discrepancies through efficient layer-wise reconstruction. By directly modeling how each token's removal affects the model output, our method naturally captures attention redistribution effects, going beyond simplistic reliance on raw attention weights. To further enhance robustness, we design exponential moving average smoothing to handle temporal variations and an adaptive window-based mechanism to capture spatial patterns. Our method, ReST-KV, significantly advances performance on long-context benchmarks. It surpasses state-of-the-art baselines by 2.58% on LongBench and 15.2% on RULER. Additionally, ReST-KV consistently outperforms existing methods on Needle-in-a-Haystack and InfiniteBench, all while achieving a remarkable 10.61$\times$ reduction in decoding latency at 128k context length. The code is publicly available at https://github.com/an-yongqi/rest-kv to facilitate reproducibility and further research.

KV缓存长文本生成推理加速注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。