PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast KV Compaction via Attention Matching

Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim

arXiv:2602.16284v2cs.LG

TL;DR

Long-context language models face KV-cache growth, while existing compaction methods can lose downstream performance or require expensive optimization. This paper uses Attention Matching to construct compact keys and values that preserve attention outputs and mass, achieving fast compaction with little quality degradation at ratios up to 50×. Its methods improve the trade-off between compaction cost and quality, though chunked approximation and extreme compaction remain practical boundaries.

  • Problem

    KV-cache reduction methods and summarization can degrade downstream behavior at high compaction ratios, while latent-space Cartridges requires expensive end-to-end optimization.

  • Method

    Attention Matching directly optimizes compact keys, values, and scalar biases to match per-KV-head attention outputs and attention mass using efficient subroutines.

  • Results

    The methods achieve little performance degradation at compaction ratios up to 50× while making compaction orders of magnitude faster than gradient-based optimization.

  • Takeaways & Limitations

    Attention Matching significantly improves the Pareto frontier between compaction cost and quality for latent-space KV compaction.

  • Takeaways & Limitations

    Text-based chunking is an approximation, and Cartridges can outperform Attention Matching at 100× compaction on some LongHealth benchmarks.

Abstract

from arXiv · show

Scaling language models 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.

1. Introduction

Long-context language models are constrained by KV-cache growth, while existing reduction methods and summarization can lose important information. The paper introduces Attention Matching as a faster latent-space alternative to Cartridges that preserves quality at high compaction ratios.

  • Motivation: KV caches can reach many gigabytes per request, making efficient context compaction important for long-horizon language-model applications.Existing token eviction, token merging, and head sparsification methods degrade rapidly at high reduction ratios.
  • Motivation: Token-space summarization remains common in deployed systems but can discard information needed for downstream model behavior.The paper frames compaction as reducing KV-cache size in a single pass while preserving downstream behavior.
  • Prior Work: Cartridges reaches compaction ratios as high as 50× with minimal performance loss, but training a cache for one context can require several GPU-hours.Its latent-space approach uses prefix-tuning on synthetic self-study data and end-to-end gradient optimization.
  • Contribution: Attention Matching directly optimizes compact keys and values to reproduce attention outputs and attention mass for every KV-head using reference queries.The objective decomposes into simple subroutines, some with efficient closed-form solutions, avoiding gradient descent at compaction time.
  • Contribution: Attention Matching provides a family of methods that reaches ratios up to 50× with little performance degradation while reducing compaction time from hours to minutes.Different design choices produce a speed–performance frontier.

2. KV Compaction via Attention Matching

Attention Matching replaces a full KV cache with a shorter cache whose local attention outputs and attention mass match those of the original block. Matching both quantities preserves the block’s contribution when arbitrary fixed or future tokens are concatenated.

  • Problem Setup: KV compaction replaces K and V with shorter Ck and Cv caches containing t < T entries while targeting similar behavior for arbitrary queries.The formulation is stated per KV-head, with keys and values represented in R^T×d and R^t×d.
  • Problem Setup: The compacted prefix must remain compatible with concatenation with arbitrary uncompacted or future tokens.This requirement makes preservation of the original prefix’s attention output under concatenation central to the formulation.
  • Attention Matching: Attention over concatenated blocks decomposes into locally normalized attention outputs weighted by each block’s attention mass.This decomposition motivates matching two block-level quantities rather than enforcing behavior for every possible future token sequence.
  • Attention Matching: Attention Matching uses reference queries to match the compacted block’s local attention output and attention mass.The method optimizes compact keys, values, and per-token scalar biases to make both quantities agree with the original block.
  • Implementation: Scalar biases compensate for the impossibility of exact mass matching with fewer keys, adding negligible memory overhead and negligible-to-zero attention-runtime change.The biases are supported in common attention implementations.
  • Attention Matching: Matching attention outputs and attention mass together preserves the compacted block’s contribution under concatenation, unlike unweighted token dropping or merging.Without biases, methods that only drop or merge tokens systematically underestimate the block’s contribution during future decoding.
  • Implementation: The compacted cache retains logical length T despite storing only t entries, preserving position IDs and RoPE phases for newly appended tokens.This separates the cumulative length seen by the cache from its physical size.

3. Methods

The method constructs compact KV caches by matching attention behavior over reference queries, then allocates keys, mass biases, and values through efficient subroutines. It also supports query-generation variants, sequential on-policy optimization, and nonuniform head-wise budgets.

  • Reference-query sampling: Reference queries approximate the queries the model is likely to produce when attending to the context.
  • Reference-query sampling: Repeat-prefill extracts reconstruction queries from a sequence containing the context twice, while self-study broadens the query distribution with synthetic interactions.Repeat-prefill is faster than self-study while performing nearly as well; self-study gives the best downstream performance.
  • Reference-query sampling: Queries extracted after earlier layers are compacted provide on-policy references and yield slight but consistent improvements.
  • Constructing β and Cv: Given compact keys, nonnegative least squares fits mass weights and ordinary least squares fits compact values to minimize attention-matching error.The mass weights become β through β_j = log(w_j), and each weight represents how many original keys’ attention mass a compact key accounts for.
  • Selecting Ck: Compact keys are restricted to a subset of original keys, with highest-attention selection offering a fast heuristic and OMP greedily matching attention mass.OMP directly constructs the selected keys and mass weights, then fits compact values by least squares.
  • Selecting Ck: OMP is empirically strongest but slower; selecting multiple keys per step and refitting periodically reduces compaction time by 4–8× with little degradation.
  • Nonuniform compaction: Nonuniform compaction assigns different ratios across heads and layers, prioritizing capacity for sensitive heads through a reusable model-specific schedule.The schedule is reported as robust across instances and datasets, while variable-length packing avoids padding overhead.
  • Implementation choices: KV-based chunking is used by default because it more faithfully preserves model behavior than text-based chunking without cross-chunk interactions.

4. Results

The evaluation compares Attention Matching variants across benchmarks, models, compaction strategies, and efficiency settings. Attention Matching generally improves the quality–time trade-off, including high-compaction and hybrid-model settings, while retaining known limitations at extreme ratios and in runtime optimization.

  • Evaluation setup: Attention Matching variants trade off compaction time and downstream quality through different reference-query and key-selection strategies.AM-OMP and AM-OMP-fast use OMP, while AM-HighestAttnKeys variants select keys by attention and differ in reference-query generation.
  • Main results: Attention Matching traces the Pareto frontier at 50× compaction, bridging fast heuristic baselines and optimization-heavy Cartridges on Qwen3-4B QuALITY.Fast baselines degrade significantly at this ratio, while Cartridges represents the slower optimization-heavy end of the comparison.
  • Main results: Attention Matching consistently outperforms token eviction and summarization, particularly at 20×–100× compaction across QuALITY and LongHealth.Performance degrades faster on information-dense LongHealth; summarization matches the no-context baseline there, while KVzip sometimes matches Attention Matching through a non-uniform budget.
  • Generalization to other benchmarks: Attention Matching is evaluated beyond multiple-choice comprehension on QASPER, LongBench v2, and the intensive RULER benchmark.The evaluation covers free-form generative QA, code repositories, tables, multilingual documents, and long structured contexts; the fastest variant surpasses KVzip and six public kvpress methods on RULER.
  • Generalization to other architectures: Attention Matching remains effective on Gemma-3-12B sliding-window models when only global-attention layers are compacted.The reported ratios apply only to the global-attention portion of the cache and require slightly more conservative ratios than full-attention models.
  • Compaction efficiency: The method’s efficiency profile is constrained by query generation, which dominates runtime on a 60k-token LongHealth context.Timings use a single H200 GPU, FP32 computation followed by BF16 storage, and chunked prefill; quantization is not evaluated.
  • Compaction efficiency: AM-OMP applied after summarization reaches 200× compaction, from 6340 to 31 effective tokens on average, with performance comparable to summarization alone.Summarization alone provides approximately 20× compaction in this comparison.

5. Related Work

The related work spans one-shot token-space compaction, latent representations, token merging and eviction, head specialization, and attention biases. Attention Matching differs by operating post hoc on pretrained models and optimizing attention behavior rather than token or KV reconstruction alone.

  • Token-space compaction: One-shot compaction makes globally informed decisions over the full prefix and can be repeated to maintain a fixed maximum state size.This contrasts with online methods that continuously evict or merge tokens during decoding.
  • Token-space compaction: Token-space approaches combine retrieval augmentation, summarization or token dropping, and agentic memory management, while latent-space compaction is largely orthogonal.Summarization can rapidly degrade performance on knowledge extraction and long-context reasoning but may remain acceptable when retrieval or search is available.
  • Token pruning and merging: The paper positions KVzip and KVMerger among a broader literature on token eviction and token merging.KVPress benchmarks more than 20 such methods; the evaluation includes KVzip and KVMerger as baselines.
  • Latent representations: Soft-token and gist-token methods require training or finetuning, whereas Attention Matching operates post hoc on pretrained models without modification.This distinguishes the paper’s procedure from methods that train models to generate replacement representations.
  • Latent KV methods: Lexico reconstructs KV vectors with sparse codes over learned dictionaries, while Attention Matching optimizes compact KV vectors for attention behavior.Lexico uses a per-layer universal dictionary and OMP at inference time, but its objective is ℓ2 reconstruction of KV vectors.
  • Head specialization: DuoAttention learns retrieval-versus-streaming head partitions, while Attention Matching uses nonuniform compaction based on head variation.The comparison links both approaches to the idea that attention heads differ in their sensitivity to context.
  • Attention biases: Attention Matching uses per-token scalar attention biases to correct the attention mass received by the compacted cache.This use differs from the positional-encoding role of biases in T5 and ALiBi.

6. Discussion

KV compaction is presented as a practical primitive for reducing physical cache size while preserving logical context behavior. The discussion identifies systems integration, online use, and remaining efficiency and quality boundaries as important directions.

  • Practical implications: Attention Matching preserves attention outputs and attention mass, allowing compacted blocks to remain compatible with arbitrary appended or future tokens.This supports one-shot compaction without knowing future keys or values.
  • Limitations: At 50× compaction the approach can outperform Cartridges, but Cartridges can perform better at 100× on some LongHealth benchmarks.The paper attributes this boundary to gradient-based optimization searching a wider space of compact representations.
  • Limitations: Unoptimized OMP and self-study variants still require several minutes for compaction despite being much faster than end-to-end gradient-based methods.This leaves a practical time cost even after the paper’s speed improvements.
  • Future directions: Future work includes architectures or training procedures that treat compaction as a simple primitive or operate over a fixed set of keys and values.The paper also identifies moving beyond subset selection for compact keys as a promising direction.
  • Systems directions: Several systems directions remain unfinished, including inference-engine integration and combining token-space retrieval or summarization with latent-space compaction.The paper names RadixAttention-style prefix caching, variable-length KV packing, and disaggregated compaction as examples.
  • Future directions: Online compaction could maintain memory over long-horizon interactions by compacting the KV cache mid-trajectory under a fixed physical memory budget.A preliminary AIME result reports preserved reasoning performance under repeated compaction.

7. Conclusion

The paper formulates fast latent-space KV compaction around Attention Matching and decomposes it into objectives that preserve a compacted block’s contribution during concatenated attention. The method uses reference queries and compact parameters to match both local attention behavior and attention mass.

  • Conclusion: Attention Matching is proposed as an objective for fast latent-space compaction that improves the trade-off between compaction cost and quality.The conclusion frames the contribution as a Pareto-frontier improvement.
  • Formulation: The compact representation uses compacted keys, values, and per-token biases evaluated on reference query vectors.The supplied formulation identifies Ck, Cv, β, and Qref as the relevant compact parameters and query matrix.
  • Attention Matching: The method matches a block’s local attention output and attention mass over reference queries.These objectives are defined per compacted key-value block and support compatibility with appended context.
  • Attention Matching: Matching both quantities preserves the compacted block’s contribution when arbitrary fixed or future key-value blocks are appended.The mixture decomposition makes the compacted and original blocks comparable through their local contribution and mixture weight.
  • Formulation: Biases are needed because retaining fewer keys otherwise underestimates the original block’s attention mass.Multiplicative weights exp(βj) let retained keys represent the mass of removed keys.
  • Implementation: The attention operators are computed with per-query max-shifting for numerical stability without changing the stated identities.The implementation applies the same stabilization to mass and attention calculations.

B.2. Datasets

The evaluation uses comprehension, long-healthcare, generative QA, retrieval, and diverse long-context benchmarks, with standardized one-shot compaction protocols and several baseline configurations. Dataset choices emphasize information extraction, reasoning, retrieval, and generalization across task types.

  • Primary datasets: QuALITY evaluates comprehension of 5–8k-token passages after article contexts are compacted and questions are answered by batched decoding.The evaluation covers the first 50 validation articles and 894 questions.
  • Long-context datasets: LongHealth concatenates five patient records into roughly 60k-token contexts and evaluates chunked compaction with 100 questions per context.Each context is compacted in five chunks.
  • Additional benchmarks: QASPER tests whether compaction preserves precise extractable information under generative decoding using 562 paper contexts and 2,010 questions.Performance is measured by token-level F1 against reference answers.
  • Additional benchmarks: RULER stresses long-context retrieval through needle-in-a-haystack, multi-hop tracing, and aggregation tasks.The evaluation uses 6,500 tasks and reports the fastest AM-HighestAttnKeys configuration.
  • Additional benchmarks: LongBench v2 covers code repositories, tables, multilingual documents, and long structured contexts, with analysis restricted to 47 contexts where full context helps.The selected contexts reach 97k tokens and average 42k tokens.
  • Baselines: Baseline comparisons include Cartridges and token-selection methods, with question-aware baselines adapted through a single synthetic mock question.The mock-question configuration makes originally question-aware methods comparable in a question-agnostic setting.
  • Baselines: DuoAttention performed poorly at the evaluated ratios retaining up to 0.2× of the original cache and was omitted from the main figure.Its implementation partitions heads into streaming and retrieval categories.

C. Algorithmic Implementation Details

The implementation combines per-head algorithms, least-squares and NNLS subroutines, bias stabilization, chunked compaction, and query-generation strategies. It preserves logical positional semantics while trading selection speed against quality.

  • Execution: Per-head algorithms are implemented sequentially, and batching across heads did not produce major speedups.The implementation therefore iterates through each KV-head in the model.
  • Key selection: Periodic-refit OMP selects keys greedily and periodically refits nonnegative weights, with faster settings selecting four keys per step and refitting every two iterations.The main algorithm uses k = 1 and τ = 1, while the fast variant uses k = 4 and τ = 2.
  • Numerics: The implementation computes compact parameters in FP32 before casting them to BF16 for storage and later use.This applies to Ck, β, and Cv.
  • Linear algebra: Least-squares quality ranked lstsq above cholesky and pinv, while cholesky was fastest; the experiments therefore use torch.linalg.lstsq.The differences in quality were minimal, and positive ℓ2 regularization degraded performance.
  • Linear algebra: The NNLS solver initializes from unconstrained least squares, clamps weights to a lower bound, and optionally refines them with projected gradient descent.The implementation supports optional upper bounds and uses power iteration to estimate the step-size denominator.
  • Bias stabilization: Bias stabilization constrains weights or biases because extremely negative β values can make selected keys ineffective for matching attention outputs.Highest Attention Keys uses βj ∈ [−3, 3], while OMP applies separate pruning and upper-bound safeguards.
  • Chunking: Chunked compaction independently compacts contiguous chunks per layer and head, then concatenates compacted segments while keeping prefix and suffix tokens fixed.The resulting layout is [prefix] + compacted chunks + [suffix].
  • Chunking: KV-based chunking preserves global positional semantics by continuing original positional indexing and retaining the original logical length and RoPE phases.Text-based chunking instead applies a uniform RoPE phase shift after local processing.

D. Impact of Reference Queries

Reference-query generation strongly affects compaction quality: self-study performs best, while repeat-prefill and context-prefill are close behind and cheaper. Random query vectors remain usable but lag these model-derived strategies.

  • Query-generation strategies: Randomly sampled query vectors work, but lag the other reference-query approaches.The evaluation compares log(Perplexity) on original-cache generations because it has lower variance and correlates well with downstream accuracy.
  • Query-generation strategies: Self-study-based reference-query sampling performs best, especially at the greatest compaction ratios.Subsampling reference queries preserves performance.

E. Head Budgets

Head sensitivity and attention patterns support nonuniform KV budgets, because different heads benefit differently from retained capacity. Global attention selection induces stable allocations, while aggregation and reconstruction analyses clarify how methods should be compared.

  • Global allocation: Global highest-attention selection induces stable head proportions across articles and target compaction ratios.The observed stability supports a fixed, precomputed nonuniform schedule.
  • Optimized budgets: Uniform compaction uses one ratio per head and layer, whereas nonuniform compaction assigns potentially different ratios based on head and layer sensitivity.The allocation procedure uses sensitivity curves and greedy swaps to derive model-specific budgets.
  • Head sensitivity: Different heads show substantial sensitivity differences: some are largely insensitive to KV capacity, while others benefit significantly from retaining additional keys.This pattern appears in both Qwen3-4B and Llama-3.1-8B-Instruct.
  • Optimized budgets: The optimized head-budget schedule resembles the allocation induced by global highest-attention selection and assigns more budget to later layers than earlier layers.This differs from the direction associated with PyramidKV in the cited comparison.
  • Aggregation and evaluation: RMS aggregation offers a good balance among mean, RMS, and max-style choices, although no aggregation function is consistently best across methods.Keys are ranked after per-query attention weights are aggregated into one importance score per key.
  • Reconstruction and downstream accuracy: Reconstruction improvements generally translate into better QA accuracy, but summarization can preserve task-useful information despite poorer token-level reconstruction.Cartridges shows the opposite pattern, with slightly lower reconstruction loss than methods having similar downstream accuracy.
  • Reconstruction and downstream accuracy: Reconstruction loss is useful within attention-preserving method families but is not a universal predictor when methods fundamentally change objectives.The reported Qwen3-4B QuALITY evaluation did not show entropy blowup below 10k-token contexts, but did at longer context lengths.
  • Online compaction: Repeated mid-trajectory compaction preserves performance comparable to standard decoding at the same effective length, despite shrinking the KV cache up to six consecutive times.The evaluation uses a 50% reduction per compaction and reports effective versus physical sequence length.

F.4. Additional Benchmarks

Additional benchmarks reinforce the paper’s main result: Attention Matching remains strongest across evaluated settings. The fastest variant surpasses KVzip and leaderboard methods on retrieval-intensive RULER, while KVMerger trails Attention Matching on QuALITY.

  • Overall results: Across every additional benchmark, Attention Matching achieves the strongest performance.The benchmarks include KVMerger on QuALITY and the retrieval-intensive RULER benchmark.
  • QuALITY: KVMerger performs similarly to token-eviction baselines and well below Attention Matching across all evaluated QuALITY ratios.The comparison uses Qwen3-4B on 50 QuALITY contexts.
  • RULER: AM-HighestAttnKeys-fast surpasses KVzip and all leaderboard methods on the 6,500-task RULER benchmark.This fastest variant uses repeat-prefill queries without self-study or OMP.

G. Additional Analysis

Attention Matching supports selective, asynchronous, and architecture-aware compaction. Its per-layer, per-head formulation works with hybrid attention and GQA but not multi-head latent attention.

  • Selective compaction: Selective compaction can leave critical system-prompt or chat-template tokens uncompacted while evaluating alignment through QA and reconstruction metrics.Compacted and original generations differ little except at very high compaction ratios.
  • Architecture compatibility: The per-layer, per-KV-head formulation composes with sliding-window hybrid architectures and GQA under tensor parallelism.Gemma-3 experiments compact only global-attention layers.
  • Architecture compatibility: Attention Matching is incompatible with multi-head latent attention, which jointly caches a low-rank latent instead of separate keys and values.Sequence-dimension compaction is described as complementary to MLA’s per-token cache reduction on GQA models.
  • Serving design: Compaction can run off the request critical path because a worker processes a KV-cache snapshot before swapping in the compacted prefix.The serving engine continues answering from the original prefix until compaction completes.
Loading 2602.16284v2…