🤖 machine learning

Fast KV Compaction via Attention Matching

This paper introduces "Attention Matching," a fast and efficient method for compressing language model KV caches in latent space by solving subproblems with closed-form solutions to achieve up to 50x compaction with minimal quality loss, overcoming the speed limitations of previous optimization-based approaches.

Original authors: Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a very long story so you can answer questions about it later. In the world of Artificial Intelligence (AI), this "memory" is called a KV Cache (Key-Value Cache). As the story gets longer, this memory file grows huge, filling up the computer's hard drive and slowing everything down.

Usually, when the memory gets too big, AI systems try to fix it by summarizing the story. They throw away the details and keep only a short summary. But this is like trying to remember a complex mystery novel by only reading the back cover: you lose the clues, the plot twists, and the ability to answer specific questions.

Another method, called "Cartridges," tries to create a tiny, perfect version of the memory by doing a massive amount of math training for every single story. It works well, but it's so slow and expensive that it's like hiring a team of architects to redesign a house every time you want to move a piece of furniture.

This paper introduces a new, faster way called Attention Matching. Here is how it works, using simple analogies:

1. The Problem: The "Too-Long" Bookshelf

Think of the AI's memory as a bookshelf with thousands of books (tokens). When you ask a question, the AI looks at all the books to find the relevant ones. If the shelf is too full, the AI gets overwhelmed.

Old Way (Summarization): Throw away 90% of the books and keep only a summary note. You save space, but you can't find specific details anymore.
Old Way (Cartridges): Try to rewrite the entire library into a single, perfect, tiny book. It's accurate, but it takes days to write.

2. The Solution: The "Highlighter and Translator" (Attention Matching)

Instead of throwing books away or rewriting the whole library, this new method acts like a smart librarian who does two things instantly:

Step A: The Highlighter (Selecting Keys)
The librarian looks at the story and asks, "If I were to ask a question about this, which pages would I look at?" They identify the most important pages (keys) and keep only those.
Step B: The Translator (Adjusting Values & Biases)
Here is the magic trick. If you just keep a few pages, the story feels "lighter" because you removed the weight of the missing pages. To fix this, the librarian adds a special bias (a little weight adjustment) to the kept pages.
- Analogy: Imagine you have a backpack with 100 heavy stones. You need to carry it, but you can only hold 5 stones. If you just pick 5 stones, the bag is too light. So, you attach a "magic weight" to each of the 5 stones so that, in total, they feel just as heavy and important as the original 100.

3. How It Works Without Slow Training

The paper claims that instead of spending hours training a new model (like the "Cartridges" method), this approach uses math shortcuts (closed-form solutions).

It's like solving a puzzle by using a formula rather than trying every possible piece combination.
It calculates exactly how to adjust the "weights" (biases) and the "values" (the content) so that when the AI looks at the small, compacted memory, it gets the exact same "feeling" or result as if it had seen the whole original story.

4. The Results: Fast and Accurate

The authors tested this on long documents (like medical records or long articles) and compared it to other methods.

Speed: They can shrink the memory by 50 times in just a few seconds.
Quality: Unlike summarization, which loses accuracy, this method keeps the AI's ability to answer questions almost as good as if it had the full memory.
The Trade-off: It sits on the "Pareto frontier," meaning it offers the best possible balance between speed and quality. It is much faster than the slow training methods and much more accurate than the fast summarization methods.

5. A Special Feature: "Non-Uniform" Compaction

The paper also notes that not all parts of the AI's brain (called "heads") are equally important.

Analogy: In a library, some shelves hold the most critical books, while others hold reference manuals you rarely need.
This method figures out which shelves need to stay full and which can be emptied more aggressively. It doesn't treat every part of the memory the same; it gives more space to the parts that matter most.

Summary

This paper presents a way to shrink an AI's memory file fast without losing the details. Instead of throwing away information (summarization) or spending hours retraining (Cartridges), it uses a mathematical trick to keep the most important pieces and "weight" them correctly so the AI behaves as if it still remembers everything. It allows AI to handle very long conversations or documents without running out of memory or getting confused.

Technical Summary: Fast KV Compaction via Attention Matching

1. Problem Statement

Scaling large language models (LLMs) to long contexts is currently bottlenecked by the memory requirements of the Key-Value (KV) cache. In autoregressive Transformer architectures, the KV cache must retain keys and values for all previous tokens, often reaching gigabytes per request in long-horizon settings (e.g., multi-session dialogue, agentic coding).

Existing approaches to reduce KV cache size generally fall into two categories, both with significant limitations:

Token-space methods: Techniques like token eviction (H2O, SnapKV), token merging (KVMerger), and summarization degrade rapidly at high compression ratios. Summarization, in particular, is highly lossy and harms downstream performance on tasks requiring precise information extraction.
Latent-space methods: Recent work like Cartridges (Eyuboglu et al., 2025) demonstrated that compact KV caches can be trained in latent space to match full-context performance. However, this approach relies on end-to-end gradient-based optimization (prefix-tuning) for every context, which is prohibitively slow and expensive (requiring several GPU-hours per context).

The open problem is achieving fast, high-quality context compaction in latent space that preserves downstream model behavior without the computational cost of per-context gradient optimization.

2. Methodology: Attention Matching

The authors propose a family of methods based on Attention Matching, which constructs a compact KV cache $(C_k, C_v)$ to reproduce the attention outputs and attention mass of the original cache $(K, V)$ for a set of reference queries.

Core Objective

The goal is to find a compact cache $(C_k, C_v)$ of size $t \ll T$ such that for any query $q$ and any appended fixed/future tokens $(K_{fixed}, V_{fixed})$ , the attention output remains approximately preserved. The authors leverage the decomposition of attention over concatenated blocks, showing that it suffices to match two quantities over a set of reference queries $Q_{ref}$ :

Local Attention Output: The weighted sum of values.
Attention Mass: The sum of unnormalized attention logits (the denominator in the softmax).

To handle the discrepancy where the compacted cache has fewer tokens ( $t < T$ ) and thus cannot naturally match the total attention mass (e.g., for $q=0$ , mass is $T$ vs. $t$ ), the method introduces a per-token scalar bias $\beta$ . The optimization targets:
$\frac{\exp(qK^\top)V}{\sum \exp(qK^\top)} \approx \frac{\exp(qC_k^\top + \beta)C_v}{\sum \exp(qC_k^\top + \beta)}$
$\sum \exp(qK^\top) \approx \sum \exp(qC_k^\top + \beta)$

Algorithmic Pipeline

The framework decomposes the problem into efficient, closed-form subproblems, avoiding gradient descent at compaction time:

Reference Query Generation ( $Q_{ref}$ ):
- Repeat-prefill: Running a prefill pass on the context itself (or a repeated version) to extract query activations.
- Self-study: Generating synthetic interactions conditioned on the context to broaden the query distribution.
- On-policy: Compacting layers sequentially and extracting queries from the partially compacted model to reduce distribution shift.
Key Selection ( $C_k$ ):
- Highest Attention Keys: Selecting original keys that receive the highest attention scores across reference queries (aggregated via RMS).
- Orthogonal Matching Pursuit (OMP): Greedily selecting keys to best satisfy the attention mass matching objective, followed by Non-Negative Least Squares (NNLS) to fit weights.
Bias and Value Fitting ( $\beta, C_v$ ):
- Fitting $\beta$ : Solved via Non-Negative Least Squares (NNLS) to match the attention mass of the original keys.
- Fitting $C_v$ : Solved via Ordinary Least Squares (OLS) to match the attention outputs given the selected keys and fitted biases.

Non-Uniform Compaction

The authors observe that different attention heads exhibit varying sensitivity to compaction. They propose a non-uniform compaction schedule where different heads and layers are assigned different retention ratios. This schedule is precomputed once per model using a greedy exchange algorithm based on per-head sensitivity curves, allowing for robust application across different contexts.

Chunked Compaction

To support very long contexts (e.g., 60k+ tokens), the method applies compaction independently to contiguous chunks of the input. The KV-based chunking approach (prefilling the full context, slicing KV states, compacting, and merging) is used by default to preserve global positional semantics and cross-chunk interactions.

3. Key Contributions

Attention Matching Formulation: A novel objective for latent-space compaction that directly optimizes for attention behavior preservation rather than output likelihoods, naturally decomposing into efficient linear algebra problems.
Closed-Form Solutions: The method eliminates the need for end-to-end gradient-based optimization during compaction, replacing it with least squares and NNLS solvers.
Speed-Quality Pareto Frontier: The approach achieves compaction ratios up to 50× in seconds (on a single GPU), orders of magnitude faster than Cartridges, while maintaining comparable or superior downstream accuracy.
Non-Uniform Scheduling: A model-specific, precomputed budget allocation strategy that prioritizes capacity for attention heads most sensitive to compaction.
Bias Mechanism: The introduction of per-token scalar biases to correct attention mass discrepancies, enabling exact mass matching with a reduced set of keys.

4. Results

The methods were evaluated on QuALITY (long-document QA), LongHealth (dense clinical records), QASPER (generative QA), and LongBench v2.

Performance vs. Speed: On QuALITY with a 50× compaction ratio, Attention Matching methods trace the Pareto frontier, outperforming token-eviction baselines (H2O+, KVzip) and summarization. They match or exceed the accuracy of Cartridges while being two orders of magnitude faster (minutes vs. hours).
High Compression: At extreme ratios (up to 100×), the methods retain significant downstream accuracy (e.g., ~60% on LongHealth vs. ~30% for summarization baselines).
Generalization: The approach generalizes to generative tasks (QASPER) and retrieval-intensive benchmarks (RULER), where the fastest variant (AM-HighestAttnKeys-fast) surpassed the state-of-the-art KVzip on RULER.
Efficiency: Compaction of a 60k-token context on Gemma-3-12B takes approximately 100–600 seconds depending on the key selection strategy (OMP vs. Highest Attention), dominated by query generation and key selection steps.
Online Compaction: Preliminary results on AIME show that repeated mid-trajectory compaction preserves reasoning performance, decoupling reasoning depth from physical memory limits.

5. Significance and Claims

The paper claims that Attention Matching provides a practical, efficient primitive for KV cache compaction that bridges the gap between fast, lossy token-space methods and slow, high-quality latent-space methods.

Practicality: By avoiding per-context gradient descent, the method makes high-fidelity compaction feasible for real-world deployment where latency and compute cost are constraints.
System Integration: The approach is compatible with existing inference engines (e.g., FlashAttention, variable-length KV packing) and can be performed off the critical path (disaggregated compaction).
Limitations: The authors acknowledge that while their method is significantly faster than Cartridges, it is still slower than simple heuristics. Furthermore, at extreme compression ratios (e.g., 100× on LongHealth), gradient-based methods like Cartridges may still outperform due to their ability to search a wider representation space not restricted to original keys.
Future Directions: The paper suggests future work in optimizing query generation, moving beyond subset selection to direct optimization of compact keys, and integrating compaction into inference engines for online, long-horizon agent settings.

In summary, this work establishes that fast, high-quality context compaction is achievable by matching attention behavior directly, offering a scalable solution to the memory bottleneck in long-context LLMs.