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

2026-05-09

AI Summary

This work addresses the high memory overhead of KV caching in large language models during long-sequence inference, a limitation exacerbated by existing pruning methods that overlook the attention redistribution dynamics and spatiotemporal characteristics of the cache. The authors formulate KV cache pruning as an optimization problem and introduce, for the first time, layer-wise output reconstruction to explicitly minimize the impact of pruning on model outputs. To enhance robustness, they integrate exponential moving average and an adaptive window mechanism to capture temporal continuity and spatial locality, respectively. Evaluated on LongBench and RULER, the proposed method outperforms current state-of-the-art approaches by 2.58% and 15.2%, respectively, while achieving a 10.61× reduction in decoding latency at a context length of 128k tokens.

Abstract

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.

Problem

• KV cache eviction
• attention redistribution
• spatial-temporal dynamics
• long-context inference
• memory efficiency

Innovation

• KV cache eviction
• layer-wise reconstruction
• spatial-temporal smoothing
• attention redistribution
• long-context inference