Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution (2510.00636v1)

Abstract: Memory consumption of the Key-Value (KV) cache represents a major bottleneck for efficient LLM 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}$.

• The paper introduces a training-free technique that compresses KV caches by analytically estimating future query attention from Gaussian distributed activations.
• It employs a principled ranking strategy for KV pairs based on expected contribution, achieving substantial memory savings, including reductions up to 15GB in long-context settings.
• The method preserves residual stream integrity with minimal performance loss and supports both prefilling and decoding phases for efficient long-context inference.

Expected Attention: Training-Free KV Cache Compression via Future Query Distribution Estimation

Introduction

The paper "Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution" (2510.00636) addresses the critical challenge of Key-Value (KV) cache memory consumption in transformer-based LLMs during long-context inference. The authors propose a training-free compression method, Expected Attention, which estimates the future importance of cached KV pairs by analytically modeling the distribution of future queries. This approach enables principled ranking and pruning of KV pairs, achieving substantial memory savings with minimal impact on model performance, and is compatible with both prefilling and decoding phases of LLM inference.

Theoretical Foundations

KV Cache in Transformers

In autoregressive transformers, the KV cache stores key and value vectors for each processed token, facilitating efficient attention computation. The memory footprint of the KV cache grows linearly with sequence length, quickly becoming a bottleneck for long-context inference. The optimal strategy for cache compression is to remove KV pairs with minimal impact on the residual stream, quantified by the contribution $\|\Delta h_{ti}\| = a_{ti} \|W_o v_i\|$, where $a_{ti}$ is the attention weight and $W_o v_i$ is the transformed value vector.

Distributional Properties of LLM Activations

The authors leverage the empirical observation that hidden states and query activations in modern LLMs are approximately Gaussian distributed (Figure 1). This property enables closed-form estimation of expected attention scores for future queries, using the moment-generating function of the Gaussian distribution.

Figure 1: Hidden states and queries in Llama3.1-8B are well-approximated by Gaussian distributions, supporting analytical estimation of expected attention.

Expected Attention Score

Given the Gaussian distribution of future queries, the expected unnormalized attention score for a key $k_i$ is computed as:

$$\hat{z}_i = \exp\left(\frac{\bar{\mu}_q^T k_i}{\sqrt{d}} + \frac{k_i^T \bar{\Sigma}_q k_i}{2d}\right)$$

where $\bar{\mu}_q$ and $\bar{\Sigma}_q$ are the mean and covariance of the averaged future query distribution. The expected attention weight is then:

$$\hat{a}_i = \frac{\hat{z}_i}{\sum_{j=1}^t \hat{z}_j}$$

The expected contribution magnitude for each KV pair is:

$$\|\widehat{\Delta h}_i\| = (\hat{a}_i + \epsilon) \|W_o v_i\|$$

where $\epsilon$ is a small regularization constant.

Compression Algorithm

The Expected Attention algorithm ranks all cached KV pairs by their expected contribution and evicts the lowest-scoring pairs according to a user-specified compression ratio. The method is compatible with both prefilling (one-shot compression) and decoding (streaming compression) phases, and supports head-adaptive compression to account for heterogeneous head importance.

Experimental Evaluation

Prefilling Phase

On LongBench, Expected Attention achieves optimal trade-offs between compression ratio and task performance across diverse datasets and models (Figure 2). The method consistently retains critical information even under aggressive compression.

Figure 2: Expected Attention maintains high scores across LongBench tasks and compression ratios, outperforming baselines.

On the Ruler benchmark, Expected Attention demonstrates robust performance in retrieval, multi-hop reasoning, and aggregation tasks, especially at high compression ratios. In the Needle in a Haystack (NIAH) test, Expected Attention exhibits stable retrieval accuracy across context lengths and needle positions, matching or exceeding trainable baselines (Figure 3).

Figure 3: Needle in the Haystack retrieval accuracy for Llama3.1-8B at 50% compression; Expected Attention retains critical tokens regardless of position.

Decoding Phase

For reasoning models evaluated on Aime25 and MATH-500, Expected Attention outperforms or matches state-of-the-art streaming-compatible baselines, particularly at high compression ratios (Figure 4). The method demonstrates that a substantial fraction of tokens in reasoning traces are redundant and can be pruned without degrading performance.

Figure 4: Decoding performance on Aime25; Expected Attention achieves superior scores at reduced KV cache sizes.

Memory Efficiency

Expected Attention delivers significant memory savings, with up to 15GB reduction for 120k-token contexts at 90% compression (Figure 5, left). At 50% compression, the method maintains baseline performance on NIAH while halving the KV cache size (Figure 5, right).

Figure 5: Left: Peak memory usage vs. sequence length for Llama3.1-8B; Right: NIAH score vs. compression ratio for Qwen3-8B, with marker size indicating cache size.

Residual Stream Integrity

Expected Attention achieves the lowest reconstruction error in the residual stream compared to competing methods, indicating minimal perturbation of hidden states and supporting downstream performance preservation.

Implementation and Framework

The authors release KVPress, a PyTorch-based research library for standardized implementation and benchmarking of KV cache compression methods. KVPress integrates with Hugging Face transformers via forward hooks, enabling non-invasive, layer-wise compression and supporting over 20 techniques. The framework prioritizes readability and reproducibility over deployment efficiency, facilitating rapid experimentation and fair comparison.

Limitations and Future Directions

Expected Attention, as a training-free method, does not match the performance of trainable approaches but offers practical deployment advantages. The requirement for manual compression ratio specification and the lack of automated adaptation to task or context are noted limitations. Future work may explore hybrid approaches combining analytical estimation with lightweight fine-tuning, as well as optimized CUDA implementations for production deployment.

Implications and Prospects

Expected Attention provides a principled, theoretically grounded solution to KV cache compression, compatible with modern transformer implementations (e.g., Flash Attention) and applicable to both prefilling and decoding. The method enables scaling LLM inference to longer contexts within hardware constraints, supporting advanced applications in agentic workflows, code analysis, and multimodal reasoning. The analytical framework may inform future research on activation modeling, adaptive compression, and integration with quantization or sparse attention techniques.

Conclusion

Expected Attention introduces a robust, training-free algorithm for KV cache compression in LLMs, leveraging the Gaussianity of activations to estimate future query attention. The method achieves strong empirical results across benchmarks and models, with substantial memory savings and minimal performance loss. The accompanying KVPress library establishes a standardized platform for research and benchmarking in KV cache management, advancing the state of the art in efficient long-context LLM inference.