Fast KV Compaction via Attention Matching

Abstract: Scaling LLMs to long contexts is often bottlenecked by the size of the key-value (KV) cache. In deployed settings, long contexts are typically managed through compaction in token space via summarization. However, summarization can be highly lossy, substantially harming downstream performance. Recent work on Cartridges has shown that it is possible to train highly compact KV caches in latent space that closely match full-context performance, but at the cost of slow and expensive end-to-end optimization. This work describes an approach for fast context compaction in latent space through Attention Matching, which constructs compact keys and values to reproduce attention outputs and preserve attention mass at a per-KV-head level. We show that this formulation naturally decomposes into simple subproblems, some of which admit efficient closed-form solutions. Within this framework, we develop a family of methods that significantly push the Pareto frontier of compaction time versus quality, achieving up to 50x compaction in seconds on some datasets with little quality loss.

• The paper introduces attention matching objectives to compact KV pairs while preserving per-head attention mass and output.
• It employs HAT and OMP strategies with closed-form solutions to optimize speed-quality trade-offs in compaction.
• Empirical benchmarks demonstrate that AM methods outperform token-level summarization and Cartridges, achieving up to 50× compaction.

Fast KV Compaction via Attention Matching: Technical Summary

Problem Statement and Motivation

Modern LLMs encounter substantial memory bottlenecks in long-context settings due to the exponential growth of key-value (KV) cache requirements. Traditional management via token-level summarization is lossy and triggers rapid degradation in downstream performance for tasks demanding precise retention and reasoning across extended histories. Methods such as Cartridges achieve high compaction ratios but incur prohibitive computational costs, rendering them impractical in real-world deployments. This paper introduces a family of attention-matching (AM) objectives and algorithms for fast latent-space KV compaction, explicitly targeting the preservation of per-head attention mass and attention outputs for sets of reference queries while enabling closed-form solutions for most subproblems, thereby expanding the speed–quality Pareto frontier.

Attention Matching Methodology

The central AM objective seeks to compact the original $(\bm{K},\bm{V})$ pairs—of length $T$—to $(\bm{C}_k, \bm{C}_v, \bm{\beta})$—of length $t \ll T$—such that, for any query $\bm{q}$, the attention output and mass closely approximate those of the full cache. The decomposition leverages a mixture model interpretation, ensuring that compaction remains robust to arbitrary future cache concatenations. Critical to this construction is the introduction of per-token bias terms ($\bm{\beta}$), enabling accurate attention mass preservation, which is unattainable by simple token selection.

Closed-form solutions for $\bm{C}_v$ and $\bm{\beta}$ are derived via least squares and non-negative least squares (NNLS), respectively, contingent on the selection of $\bm{C}_k$. The selection strategy for compacted keys employs two principal approaches: highest attention keys (HAT), aggregating reference query attention scores, and orthogonal matching pursuit (OMP), which greedily minimizes mass error via NNLS. Empirically, restricting $\bm{C}_k$ to subsets of original keys is adequate, providing both efficiency and accuracy.

Reference Query Generation and Integration

Reference queries are sourced from repeat-prefill, context-prefill, and self-study mechanisms. The latter involves synthetic prompting to broaden the query distribution, substantially improving performance in high-compaction regimes. On-policy mechanisms mitigate distribution drift across layers during sequential compaction. The robust integration of multiple query-generation methods offers consistent results across datasets and model families, as detailed in Figure 1.

Figure 1: Reference query sampling comparison. Self-study-based methods yield optimal performance at high compaction ratios; repeat and context-prefill are competitive, and random vectors lag.

Nonuniform Per-Head Compaction

AM exposes layer- and head-specific sensitivity to compaction. Extensive empirical measurement (Figure 2 and Figure 3) uncovers stable, input-invariant patterns of head importance. These findings motivate model-specific, nonuniform allocation schedules for KV head budgets, derived via greedily optimized exchange algorithms leveraging sensitivity curves.

Figure 2: Head sensitivity curves in Qwen3-4B. Certain heads show negligible loss with stricter compaction, while others require substantially more capacity.

Figure 3: Additional head sensitivity curves; layer and head dependency persists across models and datasets.

Optimized head allocations depart from uniform scheduling and, unlike PyramidKV, prioritize later layers (Figure 4), highlighting the necessity for tailored compaction strategies.

Figure 4: Optimized head budgets assign more capacity to late model layers, deviating from PyramidKV's linear decrease.

Empirical Evaluation

Benchmarks on QuALITY and LongHealth, using Qwen3-4B, Llama3.1-8B, and Gemma3-12B, demonstrate that AM methods outperform both token-selection baselines and summarization in accuracy versus compaction ratio, especially in extreme reduction regimes (e.g., $50\times$). Compaction time is reduced to seconds or minutes rather than hours, represented in Figure 5 and Figure 6.

Figure 5: Accuracy vs. compaction time trade-off; AM methods dominate the speed–quality Pareto frontier, outclassing prior baselines and Cartridges.

Figure 6: Accuracy vs. compaction ratio across methods; AM approaches maintain high quality at extreme compression, surpassing baselines.

In sliding-window architectures, AM retains effectiveness on global attention layers, solidifying its general applicability.

Ablation Analysis and Practical Extensions

Leave-one-out experiments (Figure 7) and further ablations confirm that post-compaction attention biases, value fitting, nonuniform head scheduling, and query-sampling are all integral for optimal reconstruction and downstream fidelity.

Figure 7: Ablations reveal that all components of AM-OMP are necessary for minimizing log-perplexity reconstruction error.

Furthermore, the AM pipeline can be stacked over summarization, achieving compound compaction rates exceeding $200\times$ with only marginal further loss—useful in settings where high recall is unnecessary.

Theoretical and Practical Implications

The essential theoretical contribution is the principled disentanglement of logical cache length from physical size, facilitated by attention mass/output preservation. Practically, this enables plug-and-play compaction modules for inference engines, compatible with high-throughput kernels (e.g., FlashAttention) and suited for online compaction in long-horizon agentic contexts (e.g., coding agents, tool execution traces). The methodology preserves reasoning capability through repeated compaction, empirically validated via performance scaling in online settings.

Limitations and Future Directions

While AM scales efficiently and achieves substantial reduction ratios, it is not competitive with Cartridges in ultra-high compaction ($>100\times$) regimes. Future directions include direct optimization of compact keys outside subset selection, integration of compaction-aware architectures, and the confluence of token-space retrieval and latent-space compaction.

Conclusion

Attention Matching constitutes a robust, efficient, and principled approach for fast latent-space KV cache compaction in LLMs. By matching per-head attention dynamics, it bridges the gap between speed and quality, providing practical primitives for scaling memory-intensive inference without sacrificing downstream performance.