Research map / Architectures and efficiency / Attention and transformers
Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes REFORM, a training-free inference framework that compresses long inputs into cross-layer QKV retrieval embeddings via recurrent chunked KV compression with early exit, then gathers relevant tokens by similarity and recomputes their KV cache.
Problem
Processing extremely long contexts beyond pre-trained limits: recurrent KV-compression methods lose information ("forgetting"), while random-access methods that keep full KV caches need heavy memory and CPU offloading.
Approach
Use lightweight cross-layer QKV embeddings from a few selected heads (built during recurrent compressed forwarding and exited early) to retrieve query-relevant tokens by cosine similarity, then recompute full KV cache only for those tokens, combining the memory efficiency of recurrence with random-access recall.
Key results
- On RULER at 1M context with Mistral-NeMo-Instruct-2407, REFORM scores 75.5 vs. best baseline InfLLM 23.3 (over 52% gain as claimed in abstract).
- On infinity-Bench average with Mistral-NeMo, REFORM reaches 50.2 vs. InfLLM 37.6.
- At 256k input, REFORM reduces inference time to 27.24s vs. InfLLM 129.14s and peak memory to 35.00 GB vs. 37.06 GB for TOVA/InfiniPot (InfLLM 51.62 GB).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use REFORM instead of the strongest reported baseline | Mistral-NeMo-Instruct-2407 · RULER 1M | average performance: 23.3 -> 75.5 | Helped |
| use REFORM instead of the strongest reported baseline | Mistral-NeMo-Instruct-2407 · BABILong 1M | average performance: 8.8 -> 48.8 | Helped |
| use REFORM instead of the strongest reported baseline | Qwen2.5-7B-Instruct · RULER 1M | average performance: 23.7 -> 75.1 | Helped |
| use REFORM instead of the strongest reported baseline | Qwen2.5-7B-Instruct · BABILong 1M | average performance: 40.2 -> 58.8 | Helped |
| replace H2O compression with StreamingLLM (ablation) | REFORM · RULER 300k | average performance: 84.6 -> 82.7 | Hurt |
| replace H2O compression with StreamingLLM (ablation) | REFORM · BABILong 512k | average performance: 47.6 -> 44.6 | Hurt |
| replace H2O compression with TOVA (ablation) | REFORM · RULER 300k | average performance: 84.6 -> 81.4 | Hurt |
| replace H2O compression with TOVA (ablation) | REFORM · BABILong 512k | average performance: 47.6 -> 46.8 | Hurt |
| use four selected retrieval heads instead of four random heads (ablation) | REFORM · RULER 300k | average performance: 80.3 -> 84.6 | Helped |
| use four selected retrieval heads instead of four random heads (ablation) | REFORM · BABILong 512k | average performance: 43.0 -> 47.6 | Helped |
| use four selected retrieval heads instead of four worst-performing heads (ablation) | REFORM · RULER 300k | average performance: 44.7 -> 84.6 | Helped |
| use four selected retrieval heads instead of four worst-performing heads (ablation) | REFORM · BABILong 512k | average performance: 22.8 -> 47.6 | Helped |
Surprising findings
- Cosine similarity on value/QKV states from a few selected heads outperforms attention scores and full 5120-dim hidden-state similarity for retrieval.
- Best retrieval heads are not necessarily in the final layers, enabling early exit without losing embedding quality.
Limitations
- Gather stage's token selection may fail on edge cases; selection quality bounds performance.
- Current implementation recomputes attention scores via extra matmul because Flash Attention does not output attention weights, causing redundant computation.
- Uses H2O eviction criterion directly; more advanced compression not explored.
Experimental setup
Datasets: RULER BABILong infinity-Bench RepoEval MM-NIAH HotPotQA WikiText Paul Graham essays (NIAH)
Benchmarks: Needle-In-A-Haystack RULER BABILong infinity-Bench RepoEval MM-NIAH
Metrics: accuracy edit similarity normalized performance inference time peak memory MNR
Baselines: StreamingLLM TOVA H2O InfiniPot InfLLM Truncation
Builds on: H2O StreamingLLM InfLLM InfiniPot
Closest papers by meaning
- Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.93)
- Compressed Context Memory for Online Language Model Interaction (ICLR 2024 · similarity 0.91)
- ChunkKV: Semantic-Preserving KV Cache Compression for Efficient Long-Context LLM Inference (NeurIPS 2025 · similarity 0.90)
- SCBench: A KV Cache-Centric Analysis of Long-Context Methods (ICLR 2025 · similarity 0.90)
- Long Context Compression with Activation Beacon (ICLR 2025 · similarity 0.90)
- Efficient Prompt Compression with Evaluator Heads for Long-Context Transformer Inference (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.