MIT AM-OMP: Fast KV Compaction via Attention Matching — Full Technical Analysis

Forum topic · 小凯 · 2026-03-06

Summary

AM-OMP (Fast KV Compaction via Attention Matching) is a paper by researchers at MIT that reformulates KV cache compression as a linear algebra problem rather than gradient-based optimization. Built on an attention mixing identity, the method jointly matches local attention outputs and attention mass using three closed-form steps: greedy key selection via Orthogonal Matching Pursuit (OMP), scalar bias fitting via non-negative least squares (NNLS), and value reconstruction via ordinary least squares (OLS). The per-token scalar bias adds negligible memory overhead (~O(t)) and requires no changes to attention computation. On Qwen3-4B with the QuALITY benchmark, AM-OMP achieves 50x compression with 0.67 accuracy versus 0.72 for the uncompressed cache, while compressing in roughly 30 seconds on a single H100 — about two orders of magnitude faster than Cartridges (~5 GPU-hours) at higher accuracy. The method is training-free, plug-and-play compatible with FlashAttention, FlexAttention, and PyTorch SDPA, and supports non-uniform per-head compression budgets. Code is open-sourced at github.com/adamzweiger/compaction.

Overview

AM-OMP: Fast KV Compaction via Attention Matching is a deep-dive analysis of a paper from MIT researchers (Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim) proposing a training-free, closed-form method for compressing KV caches in long-context LLMs.

Paper: arXiv:2602.16284
Code: https://github.com/adamzweiger/compaction

Key Points

Problem: KV cache memory grows with context length, reaching multiple GB. Existing approaches (token eviction, token merging, head sparsification) degrade sharply at compression ratios above 10x. Prior learned methods like Cartridges work well but require hours of gradient-based optimization per document.
Core insight: Instead of optimizing output likelihood end-to-end, AM-OMP reframes KV compaction as linear algebra: directly optimize compact keys and values to reproduce both attention outputs and attention mass, where Mass(q; K) = Σⱼ exp(q·Kⱼᵀ). An attention mixing identity shows that concatenated attention blocks can be decomposed as a weighted mixture of local outputs, weighted by attention mass.
Scalar bias compensation: Because compressed length t ≪ original length T, subset selection alone cannot match attention mass exactly. AM-OMP introduces a per-token scalar bias β ∈ ℝᵗ (solved via NNLS) that acts as a multiplicative weight so each retained key represents the mass of discarded tokens. Memory overhead is negligible: (2d+1)/(2d), with near-zero impact on compute.

The Three-Step Closed-Form Pipeline

1. Key selection — Orthogonal Matching Pursuit (OMP): greedily select keys that maximally reduce residual attention mass, using the feature matrix Φᵢⱼ = exp(qᵢ·Kⱼᵀ). 2. Bias fitting — NNLS: fit scalar biases to match attention mass. 3. Value reconstruction — OLS: solve for compacted values to match attention outputs, avoiding bias from directly reusing original values.

Because the entire pipeline uses closed-form solutions with no gradient descent, compression takes seconds instead of GPU-hours.

Performance

Qwen3-4B, QuALITY dataset, 50x compression:

| Method | Accuracy (50x) | Time (single H100) | |---|---|---| | Uncompressed cache | 0.72 | — | | AM-OMP | 0.67 | ~30 s | | Cartridges | 0.60 | ~5 GPU-hours | | H2O+ | 0.48 | — | | Text summarization | 0.51 | — |

AM-OMP is roughly two orders of magnitude faster than Cartridges while achieving higher accuracy.

Other results:

LongHealth (60k tokens, 10x compression): ~0.70 accuracy vs ~0.80 uncompressed.
Combined with summarization: up to 200x total compression (6,340 → 31 effective tokens), matching summarization-only accuracy (55.7% vs 55.2%).
Online compaction (AIME 2025): physical length 2048 / effective length 8192 scored 13/30, on par with standard 8192 decoding.

Method Family and Trade-offs

The paper describes a family of variants along a speed–quality spectrum:

AM-Highest: fastest (highest-attention key selection), good quality
AM-OMP: best quality, moderate speed
AM-OMP-Batch: batched OMP selection, faster, near-best quality

Reference query strategies include context-prefill (cheapest), repeat-prefill, self-study (best), random vectors, and on-policy extraction (best but most expensive).

Engineering Practicality

Plug-and-play: natively compatible with FlashAttention, FlexAttention (supports scalar bias), and PyTorch SDPA; no model weight changes required.
Non-uniform per-head budgets: sensitivity curves rank head importance once; a greedy-swap algorithm assigns per-head compression ratios; varlen representations avoid padding overhead.
Long documents: chunk-wise compression concatenated back together.
MoE scalability: Qwen3-235B costs only ~30% more to compact than Qwen3-4B.

Applications and Limitations

Use cases: long-running agents (e.g., persistent coding assistants), multi-turn dialogue replacing turn-dropping, long-document QA (legal, scientific), and real-time/streaming compaction during inference.

Limitations: performance depends on the quality of reference queries; per-layer compression may accumulate errors across layers; accuracy can still degrade at extreme (>100x) ratios; compaction still requires non-trivial compute.

Conclusion

AM-OMP demonstrates that KV compaction can shift from hours-long gradient optimization to seconds-long closed-form linear algebra, retaining 50x compression with only ~5 points of accuracy loss on long-document QA, while remaining training-free and drop-in compatible with standard attention kernels.