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

2026-07-01

AI Summary

This work investigates the feasibility of performing in-context retrieval directly within language models operating on million-token-long contexts, addressing the failure of conventional dense retrieval methods under length generalization scenarios. The authors identify attention dilution as the primary cause of performance degradation in long-context retrieval and propose BlockSearch, a novel approach incorporating length-aware softmax reweighting and document-level sparse attention mechanisms. Evaluated on a 0.6B-parameter model, BlockSearch matches or exceeds state-of-the-art dense retrieval systems on standard benchmarks such as MS MARCO and Natural Questions, while achieving over three times higher scores on the LIMIT benchmark. Notably, it outperforms contemporary multi-stage attention (MSA) models with only one-seventh of their parameters, demonstrating both the effectiveness and efficiency of in-context retrieval at scale.

Abstract

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, a 0.6B LM retriever whose architectural and training modifications improve over prior LM baselines and length-generalize up to 10 times beyond its training regime. Nevertheless, 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, reducing the normalized mass on the gold document even when its pre-softmax score stays high. Motivated by this analysis, we introduce length-aware adjustments to the attention softmax and document-level sparse attention. With these modifications, at the million-token scale, our model matches dense retrieval on widely studied benchmarks (e.g, MS MARCO and NQ), while outperforming the concurrent model MSA despite being 7 times smaller. Furthermore, it significantly outperforms dense retrieval on tasks requiring entirely different notions of similarity, such as LIMIT, achieving a 3 times higher score. 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.

Problem

• in-context retrieval
• language models
• attention dilution
• length generalization
• million-token corpora

Innovation

• in-context retrieval
• attention dilution
• length generalization
• sparse attention
• language model retriever