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

Forum topic · 小凯 · 2026-03-06

Summary

Researchers at MIT (Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim) propose AM-OMP, a training-free KV cache compaction method that reformulates compression from gradient-based optimization into a closed-form linear algebra problem. Instead of optimizing output likelihood like Cartridges, AM-OMP directly fits compact Keys and Values to reproduce attention outputs and Attention Mass, using Orthogonal Matching Pursuit for key selection, non-negative least squares for scalar bias fitting, and ordinary least squares for value reconstruction. A per-token scalar bias compensates for discarded tokens at negligible memory overhead. On Qwen3-4B with the QuALITY benchmark, AM-OMP achieves 50x compression with 0.67 accuracy versus 0.72 for the full cache, in roughly 30 seconds on one H100 — two orders of magnitude faster and more accurate than Cartridges (~5 GPU-hours, 0.60). It works with FlashAttention/FlexAttention, supports non-uniform per-head budgets, extends to MoE models, and combines with summarization for up to 200x total compression. Code is open-sourced.

MIT AM-OMP: Fast KV Compaction via Attention Matching — Deep Dive

Paper: Fast KV Compaction via Attention Matching Authors: Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim (MIT) arXiv: https://arxiv.org/abs/2602.16284 Code: https://github.com/adamzweiger/compaction

Key points

AM-OMP reframes KV cache compression as a closed-form linear algebra problem rather than gradient-based optimization, cutting compaction time from GPU-hours (~5h for Cartridges) to ~30 seconds per document on a single H100.
Core insight: an attention mixture identity — concatenated attention output decomposes into a weighted mix of local attention outputs, with weights given by Attention Mass: Mass(q; K) = Σⱼ exp(q·Kⱼᵀ).
A per-token scalar bias β ∈ ℝᵗ (fitted via NNLS) lets retained keys represent the mass of discarded tokens, adding memory overhead of only (2d+1)/(2d).
Three-step pipeline: OMP key selection → NNLS bias fitting → OLS value reconstruction. No gradient descent anywhere.
Results (Qwen3-4B, QuALITY, 50x compression): AM-OMP 0.67 accuracy vs. 0.72 full cache, beating Cartridges (0.60), H2O+ (0.48), and text summarization (0.51).
Composable with summarization for 200x total compression; supports online compaction (AIME 2025: 13/30 at effective length 8192 with physical length 2048).

Technical Principles

Problem

Attention Matching

OMP Key Selection

Three-Step Closed-Form Solve

| Step | Algorithm | Objective | |------|-----------|-----------| | Key selection | OMP | Greedy residual reduction | | Bias fitting | NNLS | Match Attention Mass | | Value reconstruction | OLS | Match attention outputs |

Comparison with Existing Methods

| Feature | AM-OMP | Cartridges | SnapKV | H2O | |---------|--------|------------|--------|-----| | Training | None | Pretrained dictionary | None | None | | Compression | 50x (98%) | 50x (98%) | High | Medium | | Time | ~30 s | ~5 GPU-hours | Fast | Fast | | Head budget | Non-uniform | Uniform | Head-level | Uniform |

Method Family & Reference Queries

Variants trade speed for quality: AM-Highest (fastest), AM-OMP (best), AM-OMP-Batch (compromise). Reference query strategies range from context-prefill (cheapest) to repeat-prefill, self-study (synthetic interactions), random vectors, and on-policy sampling (best quality, highest cost).

Engineering & Practicality

Plug-and-play: works with FlashAttention (via flat varlen KV segments), FlexAttention, PyTorch SDPA; no model weight changes.
Non-uniform per-head budgets: sensitivity-curve calibration done once; greedy swap assigns per-head compression ratios.
Long documents: chunked compression; MoE models cheap (Qwen3-235B costs only ~30% more than 4B).

Applications & Limitations

Use cases: long-running agents, multi-turn dialogue (replacing turn-dropping), long-document QA, real-time/streaming compression.

Limitations: depends on reference query quality; per-layer errors may accumulate; performance can drop at >100x compression; still requires extra compute.

Future Directions

Adaptive compression ratios, cross-layer joint optimization, better reference query generation
Hardware-aware and incremental/distributed compression
Multimodal extension, retrieval synergy, task-adaptive head budgets

Conclusion

AM-OMP turns KV compaction from an hours-long optimization into seconds of linear algebra: 50x compression with only ~5 points of accuracy loss, no training, and drop-in compatibility with standard attention kernels. Fully open-sourced.

Links: arXiv · GitHub