Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution
Alessio Devoto, Maximilian Jeblick, Simon Jégou
TL;DR
This work tackles the memory bottleneck of KV caches in long-context LLMs by introducing Expected Attention, a training-free method that estimates future attention from the distribution of future queries. By deriving a closed-form, Gaussian-based model for expected attention scores, it computes an expected KV contribution to the residual stream and uses this to prune KV entries with minimal performance loss. The approach, including head-adaptive compression, yields strong memory savings across prefilling and decoding on multiple models and benchmarks, without architectural changes or retraining. The authors accompany their method with KVPress, a PyTorch-based library and public leaderboard to enable standardized benchmarking of KV cache compression methods in long-context inference.
Abstract
Memory consumption of the Key-Value (KV) cache represents a major bottleneck for efficient large language model inference. While attention-score-based KV cache pruning shows promise, it faces critical practical limitations: attention scores from future tokens are unavailable during compression, and modern implementations like Flash Attention do not materialize the full attention matrix, making past scores inaccessible. To overcome these challenges, we introduce $\textbf{Expected Attention, a training-free compression method}$ that estimates KV pairs importance by predicting how future queries will attend to them. Our approach leverages the distributional properties of LLM activations to compute expected attention scores in closed form for each KV pair. These scores enable principled ranking and pruning of KV pairs with minimal impact on the residual stream, achieving effective compression without performance degradation. Importantly, our method operates seamlessly across both prefilling and decoding phases, consistently outperforming state-of-the-art baselines in both scenarios. Finally, $\textbf{we release KVPress, a comprehensive library to enable researchers to implement and benchmark KV cache compression methods, already including more than 20 techniques}$.