RAT+ Deep Dive: A New KV Cache Compression Paradigm with Dense Pretraining and Dilated Inference
> Paper: RAT+: Train Dense, Infer Sparse — Recurrence Augmented Attention for Dilated Inference > Authors: Xiuying Wei, Caglar Gulcehre > arXiv: https://arxiv.org/abs/2602.18196 > Code: https://github.com/wimh966/rat-plus > Tags: ICML2026, KV Cache, Dilated Attention, Efficient LLM
---
1. The Problem: Why Has Dilated Attention Never Worked Well?
Dilated attention is attractive in theory: by having tokens skip a fixed interval when computing attention, both FLOPs and KV cache can be compressed by a factor of D (the dilation size), while retaining long-range connectivity.
In practice, however, there is a stubborn failure mode: directly sparsifying a pretrained dense model into a dilated pattern collapses accuracy.
Wei and Gulcehre provide the data: a standard Transformer at D=64 sees perplexity explode from 7.44 to 129 — this is not a performance drop, it is total failure.
There are two root causes:
1. Distribution shift: during pretraining the model saw fully connected attention; at inference the connectivity suddenly becomes sparse, and the output distributions no longer match. 2. Lazy recurrence: even with a recurrence module, if full-sequence attention already provides complete connectivity, the recurrence module has no incentive to learn long-range dependencies — it takes the lazy path and converges to a very short effective length.
RAT+ is a systematic solution targeting both problems.
---
2. Core Design: Three Key Techniques
2.1 Full-Sequence Recurrence
Original RAT splits the sequence into blocks of size L=D, and each token only recurses within its own block. This means tokens in different positions see different recurrence window lengths, and long-range information between blocks depends entirely on attention.
RAT+ instead uses a **fixed-length recurrence window L*=64** (looking back 64 tokens from the current one), computed in parallel during prefill via associative scan. Every token thus has a fixed recurrence context length independent of dilation.
Furthermore, ARL (Active Recurrence Learning, below) allows simplifying L* to L=T (the full sequence), so each token's recurrence state accumulates information from the start of the sequence.
The recurrence uses a forget-gate mechanism:
where \(g_l\) is an input-dependent forget gate, providing LSTM-like decaying memory.
2.2 Active Recurrence Learning (ARL)
This is RAT+'s most crucial innovation.
The authors found that simply adding a recurrence module during training leads to "lazy recurrence" — full-sequence attention already provides complete token-to-token connectivity, so the recurrence module has no pressure to learn long-range aggregation and converges to a very short effective length (e.g., aggregating only 2-4 tokens).
The solution is a joint training strategy:
Within each batch, half the samples are trained with D=1 (dense attention) and half with D=64 (dilated attention).
| Configuration | Purpose | |------|------| | D=1 (dense) | Preserves full attention capability, preventing the model from "forgetting" how to do normal attention | | D=64 (sparse) | Forces the recurrence module to learn to compensate for long-range information when attention is sparsified |
The effect is dramatic:
- Without ARL: PPL at D=64 degrades from 7.66 to 9.96
- With ARL: PPL stays at 7.4-7.66 across all D settings
- Data: 1B randomly sampled FineWeb-Edu tokens
- Learning rate: 0.1× the original learning rate
- No warmup
- Pretrain once (dense + ARL)
- Choose D=2/4/8/16/32/64 at inference as needed
- Only a lightweight 1B-token adaptation phase required
- Full-sequence recurrence solves the information-loss problem of dilated attention
- Active Recurrence Learning solves the "laziness" problem of recurrence modules
- Resolution Adaptation lets one model flexibly switch compression ratios
- Hybrid strategies enable fine-grained layer/head-level tuning
- Wei, X., & Gulcehre, C. (2026). RAT+: Train Dense, Infer Sparse — Recurrence Augmented Attention for Dilated Inference. *arXiv preprint* arXiv:2602.18196.
- Code: https://github.com/wimh966/rat-plus
- SnapKV: Li, Y., et al. (2024). SnapKV: LLM Knows What You are Looking for Before Generation.
- KIVI: Liu, Z., et al. (2024). KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache.
- H2O: Zhang, Z., et al. (2023). H2O: Heavy-Hitter Oracle for Accurate Generative Inference of LLMs.
2.3 Resolution Adaptation
After pretraining, the model can flexibly switch dilation settings at inference — but this requires a small "adaptation" phase.
Key observation: switching from continuous dense attention to discrete dilated connectivity requires recalibration. But the adaptation is very fast:
| Configuration | 0 tokens | 500M | 1B | 2B | 3B | |------|---------|------|-----|-----|-----| | D=16,W=16 | 11.51 | 9.54 | 9.37 | 9.22 | 9.11 | | D=64,W=64 | 9.20 | 8.63 | 8.57 | 8.51 | 8.47 |
RAT+'s advantage: only 1B tokens of adaptation are needed, rather than retraining the whole sparse model from scratch as in prior methods.
Adaptation setup:
---
3. Performance: Near-Lossless at D=16, Only 2-3 Points Lost at D=64
3.1 Perplexity (1.5B model, 100B tokens)
| D | FLOPs | PPL | Relative to D=1 | |---|-------|-----|----------| | 1 | T | 7.40 | 100% | | 2 | T/2 | 7.42 | 100.3% | | 4 | T/4 | 7.45 | 100.7% | | 8 | T/8 | 7.48 | 101.1% | | 16 | T/16 | 7.52 | 101.6% | | 32 | T/32 | 7.58 | 102.4% | | 64 | T/64 | 7.66 | 103.5% |
Compared to a standard Transformer:
| D | Standard Transformer | RAT+ | |---|------------------|------| | 1 | 7.44 | 7.40 | | 16 | 87.9 (complete collapse) | 7.52 (near-lossless) | | 64 | 129 (complete collapse) | 7.66 (slight degradation) |
Core conclusion: RAT+ is nearly lossless at D=16 and only 0.26 PPL above dense at D=64 — meaning KV cache and attention FLOPs are simultaneously compressed 64x while the model remains usable.
3.2 Downstream Tasks (Commonsense Reasoning)
| D | ARC-C | ARC-E | Hella. | LMB. | PIQA | Wino. | Avg | |---|-------|-------|--------|------|------|-------|-----| | 1 | 40.78 | 73.23 | 59.81 | 50.34 | 73.56 | 57.54 | 59.21 | | 16 | 40.44 | 73.06 | 58.89 | 47.72 | 72.85 | 55.41 | 58.06 | | 64 | 39.76 | 72.90 | 57.93 | 44.05 | 72.91 | 57.22 | 57.46 |
The average drops only 1.15 points at D=16 and 1.75 points at D=64.
3.3 Long-Context Tasks (LongBench)
| D | NarrativeQA | Qasper | MultiField | HotpotQA | 2Wiki | GovReport | Avg | |---|-------------|--------|------------|----------|-------|-----------|-----| | 1 | 13.87 | 15.59 | 27.45 | 15.38 | 17.53 | 26.42 | 19.37 | | 16 | 14.72 | 14.54 | 25.99 | 14.68 | 17.40 | 23.68 | 18.50 | | 64 | 13.89 | 13.99 | 21.57 | 13.18 | 14.70 | 22.03 | 16.56 |
Long-context tasks are more sensitive to dilated attention, but D=16 still retains 95% of performance.
---
4. Comparison with SnapKV: Why RAT+ Wins Decisively
SnapKV is one of the mainstream KV cache compression methods; its core idea is to keep tokens with the heaviest attention (heavy hitters) and discard the rest.
But on NIAH (Needle In A Haystack) tasks, SnapKV performs extremely poorly:
| Method | S-1 | S-2 | S-3 | MK-1 | MK-2 | MK-3 | MV | MQ | |------|-----|-----|-----|------|------|------|-----|-----| | RAT+ Dilation (D=4) | 100 | 100 | 97.8 | 99.8 | 97.4 | 88.8 | 99.8 | 99.5 | | SnapKV | 57.6 | 64.2 | 13.2 | 56.4 | 41.2 | 7.2 | 46.3 | 50.9 |
SnapKV's failure mode: it only keeps tokens deemed "important," but if the needed information is not among them, it is simply lost. This is fatal for tasks requiring precise retrieval of specific information.
RAT+ differs: dilated attention keeps a regular connectivity pattern — no region is skipped, the spacing just grows. Combined with full-sequence recurrence, long-range information is preserved via recursive aggregation, so critical content is not lost.
Interestingly, RAT+ can also be combined orthogonally with top-k block selection:
| Configuration | MK-2 | MK-3 | |------|------|------| | D=64, K=8 | 57.4 | 11.0 | | D=64, K=8 + dilation | 97.4 | 79.6 |
This shows dilated attention and top-k selection are complementary — dilation provides global coverage while top-k increases local density.
---
5. Scalability: The Bigger the Model, the Smaller the Dense-Sparse Gap
Wei and Gulcehre ran scaling experiments from 1.5B → 2.6B → 7.6B and found an interesting trend:
| Model | D=1 PPL | D=64 PPL | Gap | |------|---------|----------|------| | 1.5B | 7.40 | 7.66 | 0.035 | | 2.6B | 7.20 | 7.50 | 0.029 | | 7.6B | 6.95 | 7.30 | 0.022 |
The larger the model, the smaller the performance loss from dilation. This fits an intuition: larger models have stronger "compression recovery" capability — even with sparsified attention, richer parameters can compensate for information loss.
Detailed 7.6B results:
| Configuration | CommonSense | LongBench | NIAH-S | NIAH-MK | NIAH-MV/MQ | |------|-------------|-----------|--------|---------|-----------| | D=1 | 64.4 | 20.8 | 100.0 | 99.8 | 100.0 | | D=16 | 63.6 | 20.1 | 98.8 | 90.5 | 99.4 | | D=64 | 63.1 | 20.0 | 97.6 | 79.1 | 98.6 |
At D=64, CommonSense drops only 1.3 points; NIAH-MK drops ~20 points (but remains usable).
---
6. Hybrid Strategies: Layer-Wise and Head-Wise Combinations
RAT+ supports flexibly mixing dilation settings across layers and heads at inference.
Layer-Wise Hybrid
| Configuration | LongBench Avg | |------|---------------| | Even layers D=16+W=1024, odd layers StreamingLLM | 17.69 | | Layers 0,12 D=1 (dense), rest D=16,W=256 | 18.69 |
Key finding: inserting a small amount of dense attention at critical layers (e.g., layer 0 input, layer 12 middle) significantly improves performance with minimal KV cache overhead.
Head-Wise Hybrid
| Configuration | LongBench Avg | |------|---------------| | First 2 heads D=1 (dense), rest D=4 | 19.27 | | All D=4 | 18.15 |
Key finding: keeping only the first 2 heads dense while the rest use D=4 actually outperforms all-D=4. This suggests heads divide labor differently — some handle local detail (need dense), others handle long-range aggregation (can be dilated).
---
7. Engineering Implications
7.1 "Train Once, Infer Many"
RAT+'s core value is one model, multiple inference modes:
This is far more efficient than training a separate model per compression ratio (as with H2O, SnapKV).
7.2 A "Third Paradigm" for KV Cache Compression
KV cache compression has gone through three generations:
| Paradigm | Representative Methods | Compression Approach | Problem | |------|----------|----------|------| | Gen 1: Head/Dim compression | MQA, GQA, MLA | Reduce KV head count or dimension | Limited expressiveness | | Gen 2: Token selection | H2O, SnapKV, StreamingLLM | Discard "unimportant" tokens | May lose critical information | | Gen 3: Structured sparsity | RAT+ | Regular dilated pattern + recurrence compensation | Global coverage retained while compressing |
RAT+'s dilated attention offers a deterministic compression that skips no region — a fundamental advantage over the uncertainty of token selection.
7.3 Relation to SSM/Mamba
RAT+'s recurrence module resembles Mamba/SSM, with key differences:
| Feature | RAT+ | SSM (Mamba) | |------|------|-------------| | Long-range access | ✓ (via recurrence) | ✗ (state compression loses detail) | | Parallel prefill | ✓ (associative scan) | ✓ | | Compatible with standard attention | ✓ (same model can switch) | ✗ (separate architecture needed) | | Pretraining data efficiency | Same as Transformer | Requires more data |
SSMs have advantages at extremely long sequences (>100K), but RAT+ offers a better accuracy-efficiency trade-off at medium-long sequences (4K-64K) and is compatible with the existing Transformer ecosystem.
---
8. Limitations and Future Directions
8.1 Current Limitations
1. Long-context tasks remain sensitive: on LongBench, performance at D=64 drops noticeably (19.37 → 16.56) 2. Requires joint training: cannot be applied post-hoc to existing models the way SnapKV can 3. Recurrence overhead: although FLOPs decrease, recurrence adds sequential dependency, potentially limiting very large-scale parallelism
8.2 Future Directions
1. Combination with quantization: RAT+ compresses KV cache token count; combined with KIVI's 2-bit quantization, compression could go further 2. Dynamic D selection: choose D based on input length and task — D=1 for short text, D=16/64 for long text 3. Larger-scale validation: 7.6B already shows a scaling trend; validation on 70B+ models is needed 4. Combination with speculative decoding: dilated attention reduces decoding FLOPs and could stack with speculative decoding
---
9. Conclusion
RAT+ is an important advance in KV cache compression. Its core value is not any single technique but a systematic design:
The experimental data is convincing: near-lossless at D=16, only 2-3 points lost at D=64 — meaning KV cache and attention FLOPs are simultaneously compressed 64x, with major implications for long-context inference and edge deployment.
More deeply, RAT+ represents a structured sparsity paradigm — rather than cleverly selecting which tokens to keep (SnapKV), regularly skip positions and compensate with recurrence. This "dumb" approach is more reliable in engineering and easier to combine with hardware optimization.
---