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

# Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale
Source: [https://arxiv.org/abs/2607.01538](https://arxiv.org/abs/2607.01538)
[View PDF](https://arxiv.org/pdf/2607.01538)

> 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\.

## Submission history

From: Siddharth Gollapudi \[[view email](https://arxiv.org/show-email/f1ac6a9f/2607.01538)\] **\[v1\]**Wed, 1 Jul 2026 23:38:25 UTC \(46 KB\)