# Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale

arXiv: https://arxiv.org/abs/2607.01538 · HTML: https://mlsanity.com/paper/2607.01538/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 48 references that are in the map.
- First posted on arXiv 1 Jul 2026.

## Builds on (most cited first, 30 of 45)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [gpt-oss-120b & gpt-oss-20b Model Card](https://mlsanity.com/paper/2508.10925/) (2025)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models](https://mlsanity.com/paper/2506.05176/) (2025)
- [BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models](https://mlsanity.com/paper/2104.08663/) (2021)
- [ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT](https://mlsanity.com/paper/2004.12832/) (2020)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space](https://mlsanity.com/paper/2203.14680/) (2022)
- [Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2406.10774/) (2024)
- [RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2010.08191/) (2020)
- [LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-context Multitasks](https://mlsanity.com/paper/2412.15204/) (2024)
- [SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model](https://mlsanity.com/paper/2502.02737/) (2025)
- [Transformer Memory as a Differentiable Search Index](https://mlsanity.com/paper/2202.06991/) (2022)
- [Retrieval Head Mechanistically Explains Long-Context Factuality](https://mlsanity.com/paper/2404.15574/) (2024)
- [Recommender Systems with Generative Retrieval](https://mlsanity.com/paper/2305.05065/) (2023)
- [MoBA: Mixture of Block Attention for Long-Context LLMs](https://mlsanity.com/paper/2502.13189/) (2025)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [A Neural Corpus Indexer for Document Retrieval](https://mlsanity.com/paper/2206.02743/) (2022)
- [On the Theoretical Limitations of Embedding-Based Retrieval](https://mlsanity.com/paper/2508.21038/) (2025)
- [Transformers need glasses! Information over-squashing in language tasks](https://mlsanity.com/paper/2406.04267/) (2024)
- [Can Long-Context Language Models Subsume Retrieval, RAG, SQL, and More?](https://mlsanity.com/paper/2406.13121/) (2024)
- [Softmax is not Enough (for Sharp Size Generalisation)](https://mlsanity.com/paper/2410.01104/) (2024)
- [Scalable-Softmax Is Superior for Attention](https://mlsanity.com/paper/2501.19399/) (2025)
- [TokenSelect: Efficient Long-Context Inference and Length Extrapolation for LLMs via Dynamic Token-Level KV Cache Selection](https://mlsanity.com/paper/2411.02886/) (2024)

## Built on by (most cited first, 2 of 2)

- [When Do Frozen VLMs Respond to Image-Free Object-Token Edits? An Answer-Key-Free Protocol and What It Reveals](https://mlsanity.com/paper/2609.03429/) (2026)
- [No More Free Lunch: Corpus Task Complexity Matters as Corpora Grow](https://mlsanity.com/paper/2609.29245/) (2026)

## Abstract (arXiv)

Language models (LMs) raise an intriguing alternative to vector-based retrieval: conditioning on an in-context corpus and directly generating a relevant answer. However, prior work has largely focused on proprietary systems or the smaller-scale reranking task, leaving corpus-scale in-context retrieval largely unexplored. In this work, we present the first systematic study of in-context retrieval on two scales practical retrievers demand: million-token corpora and length-generalization far beyond training-time sizes. We first introduce BLOCKSEARCH, an 0.6B LCLM retriever whose architectural and training modifications improve over prior LM baselines and length-generalize up to 10x beyond its training regime. Nevertheless, its retrieval still collapses under more extreme extrapolation. We trace this failure to an attention dilution effect: as the corpus grows, irrelevant documents dominate the softmax denominator and the normalized mass on the gold document collapses. Individual attention heads continue to locate the gold document more reliably than the model decodes it, even at million-token scale, though this signal also weakens as the corpus grows. Motivated by this analysis, we introduce length-aware adjustments to the attention softmax and document-level sparse attention, improving retrieval at million-token scale to performance comparable to a same-backbone dense retriever. On the lexical LIMIT benchmark, these gains also transfer out of distribution: at a million-token corpus, Recall@1 reaches 13.5% versus 2.9% for the dense baseline. Together, our results position in-context retrieval a promising alternative to classical retrieval while emphasizing attention control under extreme context growth as a new challenge.