Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling

1. Introduction

• The central challenge is achieving length generalization while retaining long-range access: expanding attention windows raises costs, whereas sliding windows miss dependencies outside their field.The introduction notes that expanding windows during post-training increases training and inference costs because of self-attention's quadratic complexity.Source: §1 · p. 1
• GCA makes chunk relevance scores participate differentiably in next-token prediction, allowing the retriever to learn past chunks that reduce subsequent-token autoregressive loss.It applies chunk-wise cross-attention and uses relevance scores to weight information from past chunks.Source: §1 · pp. 1–2
• DRT builds on GCA, using top-k past chunks with fixed-size sliding-window self-attention for linear-complexity attention operations over the input sequence.During inference, past-chunk hidden states can be offloaded to CPU memory and reloaded when retrieved, reducing memory footprint at the cost of additional memory swaps.Source: §1 · p. 2
• 100% performance on the single NIAH test is maintained up to 16 million tokens, while DRT significantly outperforms baselines on summarization tasks.The introduction also reports comparable pre-training costs and much lower inference costs.Source: §1 · p. 2
• ArXiv-math case studies show DRT retrieving distant lemmas, variants, or functions that are used in the next chunk.The paper presents these cases as validation of DRT's causal retrieval ability.Source: §1 · p. 2

2. Related works

• RPT uses BM25-selected chunks and a reference language model to identify training examples for its retriever, requiring prior reference-model pre-training on the target dataset.GCA instead integrates retrieval into attention for joint end-to-end pre-training without external data or models.Source: §2 · p. 2
• Landmark Attention selects top-k chunks per token and layer, adding inference costs, and the paper reports that it fails to extrapolate to longer context lengths.GCA combines chunk retrieval and grouped softmax while seeking to balance training efficiency and inference performance.Source: §2 · p. 2
• Sliding-window attention addresses differences between training and inference windows but sacrifices long-range dependencies; other methods extend windows or compress past information with different limitations.The related-work discussion notes additional training costs and length-generalization issues for extended windows, and reduced flexibility to access arbitrary past tokens for fixed-size memory methods.Source: §2 · pp. 2–3
• Figure 1 contrasts chunked cross-attention, where relevance scores only select chunks, with GCA, which separately applies cross-attention to retrieved chunks.The accompanying prose says GCA's softmaxed relevance scores weight the separately gathered chunk information and receive back-propagation from the loss.Source: Figure 1 · p. 3, §2 · p. 3

3.1. Model Architecture

• DRT divides sequences into chunks, appends a landmark token to each, and uses grouped upper decoder layers to retrieve relevant past chunks for the next chunk.Lower layers are standard Transformer layers; upper layers add GCA after self-attention.Source: §3.1 · pp. 3–4
• GCA computes cross-attention separately for retrieved chunks and fuses the outputs using softmax-normalized relevance scores.The retrieved chunks’ representations are shared across upper layers, with shared key and value transformations to reduce parameter and memory costs.Source: §3.1 · pp. 4–5

3. Methodology

• GCA replaces chunked cross-attention with per-chunk attention and softmaxed relevance weights, allowing retrieval scores to receive gradients from next-token prediction.This makes the retriever learn to select past chunks that help predict subsequent tokens.Source: §3 · p. 3, §3.1 · p. 5

3.2. Training & Inference

• During training, Gumbel noise is added to relevance scores before top-k selection, balancing exploration with exploitation in chunk retrieval.Chunks with higher relevance remain more likely to be selected, while lower-scoring chunks can still be considered.Source: §3.2 · p. 5
• Sliding-window attention and top-k retrieval reduce self-attention and GCA costs while retaining random-access flexibility.The method also encodes chunks to compute relevance scores.Source: §3.2 · p. 5
• Offloading past chunk representations to CPU memory reduces GPU memory usage, with retrieval-related memory transfers occurring only G times every S tokens.The authors describe the resulting memory-exchange cost as minimal.Source: §3.2 · p. 5

4.1.1. DATASETS

• The evaluation compares GCA with other attention mechanisms in language modeling, real-world downstream tasks, and the RULER benchmark, training all models from scratch with similar configurations.PG19 contains full-length Project Gutenberg books, while arXiv-math is a corpus of mathematical papers characterized by sustained long-distance coherence.Source: §4 · p. 5, §4.1.1 · p. 5

4.1.2. MODELS

• DRT uses 12 decoder layers, a 512-token sliding window, 64-token chunks, and retrieves eight chunks; baselines are configured for comparable attention-field coverage or parameter count.The experiments also include ablations for Triton, Gumbel top-k, and learned versus off-the-shelf or random retrieval.Source: §4.1.2 · pp. 5–7

4.2. Long-range Language Modeling

• Beyond 16K evaluation length, DRT outperforms all baselines, and it surpasses Landmark Attention on longer inputs despite retrieving less frequently.The authors attribute the gains in part to pre-training with random access to distant contexts under a fixed attention field.Source: Table 1 · p. 6, §4.2 · p. 7

4.3. Downstream task evaluation & Efficiency Analysis

• 100% performance on single NIAH tests at context lengths up to 16 million tokens; DRT also significantly outperforms all baselines on summarization.In 2-hop NIAH, DRTretrieval×2 extrapolates beyond 64K tokens and significantly outperforms DRTretrieval×1.Source: §4.3 · p. 7, Table 2 · p. 8
• DRT significantly outperforms Landmark Attention in inference speed and memory footprint; its chunk-wise retrieval runs once every 64 tokens.The retrieval operation is 1/(12 × 64) of Landmark Attention’s corresponding operation.Source: §4.3 · p. 7, Table 3 · p. 8

4.4. Case Studies

• In arXiv-math case studies, retrieved chunks include referenced definitions and lemmas needed in the next chunk, illustrating retrieval beyond semantic similarity.The paper presents this as evidence for causal retrieval of information that better predicts the next chunk.Source: §4.4 · p. 8

5. Conclusion & Future works

• Perfect accuracy in passkey retrieval on 16M contexts is reported as a first in the field.The conclusion describes GCA as maintaining efficient long-range information access.Source: §5 · p. 8
• Figure 4 illustrates retrieval of a proposition appearing in the current chunk and adjacent chunks introducing a lemma used in the next chunk.Source: Figure 4 · p. 9
• GCA makes relevance scores learnable by using retrieved information for next-chunk prediction, with Gumbel top-k sampling enabling pre-training up to 64K tokens.Source: §5 · p. 9
• Future work will explore self-supervised causal retrieval from tokens outside the context and combine structured representations for multi-granular retrieval.Source: §5 · p. 9

Impact Statement

• The authors identify potential societal consequences associated with their work on LLMs.Source: Impact Statement · p. 9
• The statement says no consequences beyond LLMs need to be highlighted.Source: Impact Statement · p. 9
• The work is framed as addressing LLM length generalization.Source: Impact Statement · p. 9

A. Hyper-parameters

• The long-range language-modeling setup uses a 12-layer decoder-only Transformer with 12 heads per layer, 768-dimensional embeddings, and an FFN size of 2048.Source: A. Hyper-parameters · p. 14
• Training uses AdamW, a 2 × 10−3 base learning rate, 2% warmup, cosine scheduling, and a final learning rate of 4 × 10−4.Source: A. Hyper-parameters · p. 14
• Models train for 60K steps on 32.2 billion tokens using eight A100 GPUs; Base LM, RPT, and DRT use 16K context, while Landmark and Block Recurrent Transformer use 768 and 4K, respectively.Source: A. Hyper-parameters · p. 14

B. Hardware-aware GCA psuedo-code

• The FLASHGCA forward pass takes Q, K, V, weights, and block sizes as inputs, then returns output O and log-sum-exp L.Source: B. Hardware-aware GCA psuedo-code · p. 14
• The forward pass partitions Q into row blocks and K and V into tiled blocks, then processes those blocks to compute outputs.Source: B. Hardware-aware GCA psuedo-code · p. 14
• The backward pass computes and returns gradients for Q, K, V, and the weights using tiled intermediate computations.Source: B. Hardware-aware GCA psuedo-code · p. 15

• In Figure 5, the top-ranked chunk supplies a definition, the second-ranked chunk supplies a theorem and corollary, and the third- and fourth-ranked chunks cover target variants.Source: Figure 5 · p. 16
• In Figure 6, the top-ranked retrieved chunk introduces a lemma used in the target chunk.Source: Figure 6 · p. 16
• In Figure 7, the top-ranked and third-ranked chunks are adjacent and share the target equation, while the second- and fourth-ranked chunks mention log G(n, x).Source: Figure 7 · p. 17, C. More case studies · p. 17

• NIAH models are fine-tuned from PG19-derived checkpoints using the same synthetic data for all evaluations.Source: D. The details for the NIAH test · p. 17
• Fine-tuning uses one-tenth of the pre-training steps, while all other hyperparameters remain constant.Source: D. The details for the NIAH test · p. 17
• Table 5 gives task examples for the two NIAH tests; inputs are padded so a landmark token can be inserted before “is” in the question.Source: D. The details for the NIAH test · p. 17, Table 5 · p. 17

• Original RPT uses a reference language model to pre-prepare target chunks for each chunk.Source: E. Discussions about the RPTContriver… · p. 17
• The original method offers stronger causal retrieval than Contriever, but its unreleased distillation code, cost, and lower flexibility motivate using Contriever instead.Source: E. Discussions about the RPTContriver… · p. 17
• RPTcontriever combines RETRO-like retrieval with RPT-style integration, retrieving past chunks from 64-token chunks encoded by Contriever and ranked by cosine similarity to the current chunk.Source: E. Discussions about the RPTContriver… · p. 17