Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution
This addresses a critical efficiency problem for deploying large language models in resource-constrained environments, offering a novel solution to KV cache compression.
The paper tackles the memory bottleneck of Key-Value (KV) cache in large language model inference by introducing Expected Attention, a training-free compression method that estimates KV pair importance based on future query distributions, achieving effective compression without performance degradation and outperforming state-of-the-art baselines in both prefilling and decoding phases.
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}$.