Attention Matching for KV-Cache Compaction

AI & Machine Learning Foundations advanced How AI Works ~35 min read

Overview

Attention Matching is a latent KV-cache compaction method for long-context language model inference. Instead of summarizing old text in token space or evicting tokens with a heuristic, it constructs a smaller set of keys and values that reproduces the behavior of the original attention cache. The goal is not to preserve every token; it is to preserve the attention outputs that future queries would obtain from the full cache.

The method was introduced in the 2026 paper Fast KV Compaction via Attention Matching. The authors frame the KV cache as a serving bottleneck: cache size grows with context length, limiting concurrency, batch size, and maximum practical context. Prior latent compaction work such as Cartridges showed that highly compressed caches can preserve quality, but required slow gradient-based optimization. Attention Matching reformulates the objective so much of the work decomposes into efficient subproblems with closed-form or simple least-squares solutions.

This makes Attention Matching a distinct member of the KV-cache compression family. TriAttention predicts which tokens to keep using offline attention geometry; Attention Matching instead fits compacted key/value representations to match attention mass and attention output for reference queries.

Core Concepts

KV Cache as Latent Context

During autoregressive decoding, each previous token contributes key and value vectors to the cache. Future tokens use those vectors to compute attention weights and then a weighted value sum. If a compact cache can make future attention produce the same weighted sums, it can preserve behavior without storing all original positions.

Attention Output

The attention output is the value-space result produced by a query attending over cached keys and values. Matching this output means the compact cache should provide the same information the model would have extracted from the full cache for representative queries.

Attention Mass

Attention mass is the total normalized weight assigned across the original cache. When many removed keys are represented by one retained key, a scalar bias can preserve the mass that would otherwise disappear. This lets a retained key stand in for a group of removed keys.

Key Concept: Match Behavior, Not Tokens

The compact cache is optimized to reproduce how attention heads use memory. It is allowed to differ from the original token sequence as long as the resulting attention behavior stays close.

How It Works

Choose Reference Queries

The compactor needs representative queries that approximate how the model will consult the context later. The paper discusses methods such as repeat-prefill prompts that ask the model to revisit the previous context, and self-study prompts that ask for synthetic tasks like extracting facts, dates, or structured summaries.

reference_queries = build_reference_queries(
    context=document_tokens,
    strategies=[
        "repeat_prefill",
        "extract_key_facts",
        "structure_dates_and_numbers"
    ],
)

Select Compact Keys

The system chooses a smaller set of keys to retain or represent the original cache. Selection can use signals such as high attention under the reference queries. The retained keys define the support over which compact values and bias terms will be fitted.

Fit Values and Biases

Given reference queries and selected keys, the value fitting problem can be solved with algebraic methods rather than end-to-end gradient optimization. Ordinary least squares or nonnegative least squares can fit compact values and scalar biases that reproduce attention outputs and preserve mass.

def fit_compact_cache(full_keys, full_values, selected_keys, queries):
    target_outputs = attention(queries, full_keys, full_values)
    target_mass = attention_mass(queries, full_keys)

    compact_values = solve_least_squares(
        design=attention_weights(queries, selected_keys),
        target=target_outputs,
    )
    compact_bias = solve_nonnegative_mass_bias(
        selected_keys=selected_keys,
        queries=queries,
        target_mass=target_mass,
    )
    return selected_keys, compact_values, compact_bias

Compact in Chunks

For long contexts, compaction can process contiguous chunks independently and concatenate the compacted caches. Chunking bounds computation and makes the method practical for very long documents or repeated mid-generation compaction.

Practical Applications

Long-Document Question Answering

A system can prefill a long document, compact its KV cache, and then answer multiple questions against the compact latent context. This avoids repeatedly paying full attention cost over the original sequence while retaining more information than a text summary.

Long-Horizon Agent Sessions

Agents that run for many steps accumulate large working histories. Attention Matching can periodically compact older context while keeping recent tokens uncompressed, allowing the model to continue reasoning without dropping all early evidence.

High-Concurrency Serving

KV memory often limits how many long-context requests fit on a GPU. A fast compactor that reduces cache size in seconds can improve concurrency and reduce offloading pressure for workloads where long contexts are common.

Lossy Alternatives to Text Summaries

Text summaries are human-readable but discard details. Latent cache compaction is not interpretable in the same way, but it can preserve model-usable information that would be hard to summarize faithfully.

Implementation Patterns

Keep a Recent Window Uncompressed

Recent tokens often carry local syntax, current task state, and conversation continuity. A practical runtime can compact older cache blocks while leaving the newest window intact.

def maybe_compact_cache(cache, max_tokens, recent_window):
    if cache.length <= max_tokens:
        return cache

    old_cache = cache[:-recent_window]
    recent_cache = cache[-recent_window:]
    compact_old = attention_match_compact(old_cache, ratio=0.02)
    return concat(compact_old, recent_cache)

Separate Offline and Online Work

Reference-query construction and fitting can run when a long document is loaded, while the interactive serving path reuses the compacted cache. For streaming agents, compact only at memory thresholds to avoid adding latency to every token.

Evaluate Per Head and Per Task

Attention Matching targets per-KV-head behavior. Compression ratios that work for retrieval-heavy tasks may not work for math reasoning, code editing, or instruction-following. Tune the ratio, reference-query strategy, and chunk size per workload.

Performance Considerations

Common Pitfalls

Further Reading