PaperlayerView exact revision on arXiv

Source-linked AI summary

Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution

Alessio Devoto, Maximilian Jeblick, Simon Jégou

arXiv:2510.00636v1cs.AIcs.CL

TL;DR

KV-cache memory limits efficient long-context LLM inference, while future attention needed for principled pruning is unavailable during compression. Expected Attention estimates future attention from query distributions to rank and prune KV pairs without training, and it outperforms state-of-the-art methods across prefilling and decoding while KVPress supports benchmarking.

  • Problem

    KV-cache memory grows with sequence length, while future-token attention scores needed to determine KV-pair importance are unavailable during compression.

  • Method

    Expected Attention uses the distribution of future queries to estimate each cached KV pair’s contribution and evict pairs with the lowest expected contributions.

  • Results

    Expected Attention outperforms state-of-the-art KV-cache compression methods across several benchmarks and both prefilling and decoding scenarios.

  • Takeaways & Limitations

    The method enables training-free KV-cache compression without architectural modifications or additional training, and KVPress provides a library for benchmarking over 20 techniques.

  • Takeaways & Limitations

    The training-free method does not match the performance of trainable methods, and its PyTorch implementation is not optimized for efficiency.

Abstract

from arXiv · show

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}$.

1. Introduction

Long-context LLM applications make KV-cache memory a major deployment bottleneck, creating demand for training-free compression that preserves transformer architectures. Expected Attention estimates future attention to prune less important KV pairs, while KVPress supports benchmarking compression methods.

  • Motivation: Long-context applications such as large-codebase analysis and short-video processing can involve hundreds of thousands of tokens, making KV-cache memory a critical deployment issue.The KV cache becomes a major bottleneck when extensive contextual information must be processed.
  • Motivation: A one-million-token KV cache requires approximately 320 GB of GPU memory for a medium-sized 70B model.KV-cache memory grows linearly with sequence length and can exceed typical GPU capacities.
  • Motivation: Training-free KV-cache compression is needed to mitigate memory growth while preserving transformer architectures for pretrained modern LLMs.Alternative architectural approaches may reduce cache size but require implementation changes and do not remove the attention bottleneck.
  • Approach: Expected Attention estimates each KV pair’s importance from the distribution of future queries and prunes up to 60% of the cache without architectural modifications or additional training.The method targets each pair’s global effect on the transformer’s residual stream rather than relying on local attention metrics or heuristics.
  • Resources: KVPress is a research library for implementing, testing, and benchmarking more than 20 KV-cache compression methods.The library is intended to support researchers evaluating compression methods on long-context benchmarks.

2. Expected Attention

Expected Attention estimates future-query attention from Gaussian activation distributions, then ranks KV pairs by their expected contribution to the residual stream for compression.

  • 2.1. Key-Value Cache in Autoregressive Transformers: KV caches store keys and values from prior inference steps, reducing recomputation but increasing memory linearly with sequence length.The cache supports autoregressive attention over previously processed tokens.
  • 2.1. Key-Value Cache in Autoregressive Transformers: Each KV pair’s residual-stream contribution combines its attention weight with the transformed value magnitude, providing an output-impact measure for pruning.The contribution is represented by Δh_ti = a_ti W_o v_i, whose norm captures both factors.
  • 2.1. Key-Value Cache in Autoregressive Transformers: Future attention weights are unavailable during compression, and Flash Attention also prevents access to complete past attention matrices.These constraints make direct evaluation of the contribution score impractical at inference time.
  • 2.2. Expected Attention: Estimating Attention From Future Queries: Expected Attention models hidden states as approximately Gaussian, so query activations inherit Gaussian distributions through linear projection and RoPE.The method averages RoPE matrices over the next T positions to obtain a tractable future query distribution.
  • 2.2. Expected Attention: Estimating Attention From Future Queries: The method analytically computes expected unnormalized attention for each key, applies softmax, and substitutes the expected weight into the contribution score.The Gaussian moment-generating function yields a closed-form expectation, avoiding the need for future queries.
  • 2.2. Expected Attention: Estimating Attention From Future Queries: Expected Attention evicts the r% of cached pairs with the lowest expected contributions, while adaptive per-layer compression retains more pairs for more important heads.Figure 2 plots LongBench scores against compression ratio, with an uncompressed baseline shown as a horizontal line.

3. Experiments

The experiments evaluate Expected Attention across prefilling and decoding settings, using long-context models, benchmarks, and selected compression baselines. Implementations measure peak memory and apply standardized configurations across hardware and phases.

  • Prefilling vs Decoding: Prefilling processes the entire prompt in parallel, whereas decoding generates tokens sequentially while appending key-value pairs iteratively.These phases have distinct computational characteristics and motivate separate evaluation settings.
  • Models and Datasets: Prefilling tests Llama3.1-8B, Qwen3-8B, and Gemma3-12B on LongBench, Ruler, and Needle in a Haystack.Decoding tests Qwen-15B-R1, Qwen-7B-R1, and OpenMath-Nemotron-14B on Aime25 and MATH-500.
  • Baselines: The study compares attention-based, embedding-based, and trainable KV-cache compression baselines selected for each use case.Prefilling baselines include SnapKV, TOVA, KeyDiff, and DuoAttention when its checkpoint is available.
  • Implementation details: Experiments use eight H100 GPUs with batch size 1, and implementation settings include epsilon=0.02 and averaging RoPE embeddings over the next T=512 positions.Needle in a Haystack uses epsilon=0; prefilling does not assume a question about the context.

4. Experimental Results

Expected Attention maintains strong task performance under KV-cache compression across prefilling and decoding benchmarks while reducing memory use. The reported results include stable long-context retrieval, strong high-compression reasoning performance, and increasingly large memory savings with longer contexts.

  • Prefilling: Expected Attention consistently achieves optimal compression-performance trade-offs across LongBench tasks and compression ratios.The result is reported for Llama3.1-8B and Qwen3-8B, with Gemma3-12B covered in Section D.
  • Prefilling: Expected Attention shows robust Needle-in-a-Haystack retrieval comparable to DuoAttention and more stable than other baselines through 125k-token contexts.The test varies needle depth and total context length.
  • Decoding: Expected Attention consistently matches or outperforms baselines on Aime25 and MATH-500, including particularly strong performance at 4× and 16× compression.Most methods show minimal degradation at 2× compression, while Expected Attention is strongest in high-compression scenarios such as 12× on reasoning traces.
  • Memory Savings and Efficiency: Memory savings become increasingly substantial as context length grows, with up to 15GB less memory reported for large Llama3.1-8B contexts.Figure 5a compares 50% and 90% compression against uncompressed vanilla attention through 120k tokens.
  • Memory Savings and Efficiency: Expected Attention achieves no accuracy loss at 50% compression on Qwen3-8B Needle-in-a-Haystack while reducing KV-cache size by 2×.Marker size indicates the actual KV-cache size in gigabytes.

5. KVPress: A Research Framework for KV Cache Compression

KVPress is a PyTorch-based research framework for implementing and benchmarking KV-cache compression methods within Hugging Face Transformers. It prioritizes readable, standardized evaluation over deployment optimization and provides a public leaderboard for reproducible comparisons.

  • Framework: KVPress streamlines development and benchmarking of KV-cache compression methods through a PyTorch-based library.It is designed to help researchers implement and test novel techniques rapidly.
  • Integration: The framework integrates natively with Hugging Face Transformers, allowing compression methods to run within the existing model pipeline.Forward hooks attached to attention layers calculate KV-pair importance after each attention-layer pass.
  • Design Trade-off: KVPress is not an optimized deployment engine because it prioritizes readability and consistent benchmarking.Production efficiency is left to later low-level optimization of individual methods.
  • Benchmarking: A public KVPress leaderboard provides standardized protocols for direct and reproducible comparisons across long-context benchmarks.The framework and benchmarking suite are intended to support validation of novel compression techniques.

6. Related Works

Related work spans trainable architectural or policy changes, training-free heuristics and attention-based methods, quantization, head-adaptive compression, and efficient KV-cache implementations. Expected Attention is positioned as a training-free method with a principled foundation based on transformer activation distributions.

  • Trainable KV-Cache Compression: Trainable approaches reduce cache size through shared key-value heads, latent projections, learned compression policies, masks, or alternative architectures.Examples include grouped representations, Multi-Head Latent Attention, learned policies, and State Space Models.
  • Training-Free KV Cache Compression: Training-free methods rank or retain KV pairs using attention scores, key norms, token positions, SVD projections, or head-specific strategies.Attention-based methods may require the full attention matrix, whereas heuristic measures avoid materializing it.
  • Training-Free KV Cache Compression: Expected Attention combines insights from heuristic approaches with a principled theoretical foundation based on transformer activation distributions.The paper presents this distributional foundation as its distinction within training-free compression.
  • Quantization: Quantization reduces cache precision rather than sequence length and is orthogonal to Expected Attention, allowing the methods to be integrated.Prior methods use block, per-layer, or per-token quantization schemes.
  • Efficient Implementations: Efficient inference systems optimize low-level KV-cache handling to improve performance, especially in multi-user serving settings.These implementation efforts complement, rather than replace, KV-cache compression methods.

7. Limitations

The training-free methodology underperforms trainable methods, requires manually specified compression ratios, and uses an implementation that is not optimized for efficiency.

  • Training-free Expected Attention does not match the performance of trainable methods.The authors identify this as an intentional trade-off that avoids the computational resources required for intensive training.
  • Compression ratios must be specified manually because the method lacks an automated mechanism for selecting optimal levels across scenarios.
  • The PyTorch implementation demonstrates the method’s principles but is not optimized for efficiency.The authors suggest custom CUDA kernels as a route to improving speed and practical utility.

8. Conclusion

The paper introduces Expected Attention as a training-free KV cache compression algorithm and reports that it outperforms state-of-the-art methods across benchmarks in both prefilling and decoding. It also releases KVPress as a research library for implementing, testing, and benchmarking compression methods.

  • Expected Attention is introduced as a training-free algorithm for KV cache compression.
  • Expected Attention outperforms state-of-the-art KV cache compression methods on several benchmarks in both prefilling and decoding scenarios.
  • KVPress enables researchers to implement, experiment with, and evaluate KV cache compression methods on popular long-context benchmarks.

A. Reconstruction Error Across Methods

The authors evaluate reconstruction error to assess how compression changes the residual stream and examine activation distributions underlying Expected Attention. Expected Attention preserves hidden states more effectively than competing methods, while the relevant activations are approximately Gaussian with some heavy-tailed outliers.

  • A. Reconstruction Error Across Methods: Expected Attention consistently achieves lower residual-stream reconstruction error than competing methods.The error is averaged over a long sequence of approximately 5K tokens and measures the difference between original and compressed hidden states.
  • A. Reconstruction Error Across Methods: The method focuses on pre-MLP hidden states and corresponding queries, which are predominantly unimodal and approximately Gaussian.
  • A. Reconstruction Error Across Methods: The relevant activations include a few heavy-tailed outliers despite their Gaussian-like distributions.
  • A. Reconstruction Error Across Methods: The paper leverages these activation distributions for KV cache compression rather than investigating their underlying origin.

C. Expected Attention Score

The paper validates Expected Attention by comparing its estimated scores with observed attention across layers and heads, while Figure 7 presents Qwen3-8B hidden-state and query distributions. The provided table caption indicates that Longbench compares average scores across compression ratios.

  • C. Expected Attention Score: Expected Attention scores are strongly correlated with original attention scores across different layers and attention heads.The comparison uses a 5K-token sequence, with the first 1K tokens used to compute query statistics.
  • C. Expected Attention Score: Figure 7 shows distributions of Qwen3-8B hidden states and queries.
  • C. Expected Attention Score: Table 3 reports Longbench average scores at increasing compression ratios across baselines.

D. Additional Results

Additional results report LongBench performance averaged across subsets, baseline selection on Ruler, activation distributions, and correlation between expected and observed attention scores.

  • LongBench results are reported as averages across all subsets.
  • Initial Ruler experiments searched more than 15 methods and selected the best-performing baselines.
  • KVZip was excluded from the selected Ruler baselines because it requires two forward passes and doubles the FLOPs cost relative to other baselines.
  • Figures 8 and 9 show the distributions of hidden states and queries for Llama3.1-8B and Gemma3-12B.
  • For Llama3.1-8B, expected attention scores show a strong correlation with observed attention scores over a 5K-token sequence using the first 1K tokens for statistics.
Loading 2510.00636v1…