Research map / Architectures and efficiency / Attention and transformers
Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling
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 Grouped Cross Attention (GCA) and Differentiable Retrieval-based Transformers (DRT), which learn chunk retrieval end-to-end through the next-chunk LM loss, enabling length generalization to 16M-token passkey retrieval with a fixed attention window.
Problem
Transformers generalize poorly to longer contexts and full attention is quadratic; existing retrieval-based LMs use separately trained retrievers not optimized for the causal LM objective, and Landmark Attention fails to extrapolate far beyond training length.
Approach
GCA runs cross-attention separately against each retrieved chunk and fuses outputs with softmaxed chunk-relevance scores, so the retriever gets gradients from the auto-regressive loss ("causal retrieval"); with Gumbel top-k sampling, sliding-window self-attention and a Triton kernel, attention cost is near-linear.
Key results
- DRT (retrieval x2) reaches 100% single NIAH accuracy from 32K up to 256K and 16M tokens, while Landmark Attention drops to 0% at 64K, with models trained at 16K context.
- On PG19 and ArXiv-math at 16K eval length, DRT x2 achieves the lowest perplexity (13.16 PG19 test, 2.81 ArXiv-math test in the 128M setting), beating Landmark Attention and the Base LM.
- DRT x1 training throughput is about 2.32e5 tokens/s at 350M params and 32K context versus 4.44e4 for full attention, and CPU offloading gives about 1.5x memory reduction at 16K prompt length.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use DRT retrieval×2 instead of Landmark Attention | Long-range language modeling · PG19 test, eval length 16K, 128M | perplexity: 13.21 -> 13.18 | Helped |
| use DRT retrieval×2 instead of BaseLM +2 layers | Long-range language modeling · ArXiv-math test, eval length 16K, 128M | perplexity: 2.94 -> 2.85 | Helped |
| use DRT retrieval×2 instead of Landmark Attention | Long-context retrieval · Single NIAH, 128K | accuracy: 0.00 -> 100.00 | Helped |
| perform two retrievals instead of one | DRT · XSum | ROUGE-1: 30.30 -> 30.39 | Helped |
| perform two retrievals instead of one | DRT · CNN/DailyMail | ROUGE-1: 36.27 -> 36.39 | Helped |
| add Gumbel noise during top-k chunk selection (ablation) | DRT · PG19 test, eval length 16K, 128M | perplexity: 13.46 -> 13.18 | Helped |
| use learned relevance scores instead of Contriever (ablation) | DRT · ArXiv-math test, eval length 16K, 128M | perplexity: 3.06 -> 2.85 | Helped |
| use learned retrieval instead of random chunk selection (ablation) | DRT · ArXiv-math test, eval length 16K, 128M | perplexity: 3.05 -> 2.85 | Helped |
| implement GCA without Triton (ablation) | DRT · Long-range language modeling | perplexity | No clear effect |
Surprising findings
- Landmark Attention, despite comparable mechanism, fails to extrapolate beyond 32K context in single NIAH, while GCA stays stable up to 16M tokens.
- A 2-layer deeper Base LM is beaten by DRT on longer evaluation lengths, suggesting precise retrieval can be more useful than adding parameters.
Limitations
- Retrieval is performed only every chunk (64 tokens) and per group, so fine-grained token-level retrieval is not provided.
- Experiments are on small models (about 350M-760M params) trained on 32.2B tokens, with the 7B reference results only from other works.
- Causal retrieval is restricted to the context; self-supervised retrieval from outside the context is left for future work.
Experimental setup
Datasets: PG19 ArXiv-math XSum CNN/DailyMail
Benchmarks: PG19 ArXiv-math XSum CNN/DailyMail Single NIAH 2-hop NIAH RULER-style NIAH
Metrics: perplexity ROUGE-1 ROUGE-2 ROUGE-L NIAH accuracy training throughput inference time per token memory footprint
Baselines: Base LM (sliding window + ALiBi) RPT (Contriever retriever, reimplemented) Landmark Attention Block-Recurrent Transformer CEPE Yarn-base (7B)
Builds on: RETRO Retrieval-Pretrained Transformer (RPT) Landmark Attention Memorizing Transformers FlashAttention-2
Closest papers by meaning
- A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention (ICLR 2025 · similarity 0.89)
- A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.89)
- Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps (NeurIPS 2026 · similarity 0.89)
- Hardware-aligned Hierarchical Sparse Attention for Efficient Long-term Memory Access (NeurIPS 2025 · similarity 0.89)
- Training-Free Long-Context Scaling of Large Language Models (ICML 2024 · similarity 0.89)
- DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads (ICLR 2025 · similarity 0.88)
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.