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

Abstract: LLMs (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.

• The paper demonstrates that architectural modifications, including SSMax and sparse routing, enable moderately sized LMs to achieve retrieval performance comparable to dense models.
• It introduces BlockSearch, a 0.6B-parameter LM that uses block-sparse attention and randomized per-document codes to mitigate overfitting and exposure bias.
• The study identifies attention dilution as a core challenge and offers targeted interventions that recover performance even as corpus sizes exceed training scales.

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

Introduction and Motivation

This work provides a systematic evaluation of whether contemporary LMs can function as scalable corpus-level retrievers when conditioned on extremely large in-context corpora, significantly exceeding their training context lengths. While the promise of in-context retrieval (ICR)—where a LM is conditioned on a corpus and directly generates document identifiers instead of using an explicit retriever and re-ranker—has seen enthusiasm in the community, previous studies have largely concentrated on either proprietary systems without controlled comparisons, or on retrieval tasks formulated over small candidate pools (i.e., reranking), not true corpus-scale retrieval.

The essential hypotheses are: (1) Can a LM retrieve effectively when conditioned on corpora orders of magnitude larger than those seen during training? (2) What are the critical limitations as corpus size increases to the million-token regime?

BlockSearch: Architecture and Extrapolation Capabilities

The authors propose BlockSearch, a 0.6B-parameter retriever LM, and introduce several design choices for scalable ICR:

• Block-sparse attention: To make attention computationally feasible, document tokens attend within their own block, and the query block attends globally.
• Randomized per-document codes: Documents are assigned random codes at each step to prevent overfitting to absolute position, as sequential IDs leak position information.
• In-batch negative training and on-policy auxiliary loss: These mitigate exposure bias by including negative examples and training with the model’s own output sequences.

BlockSearch is trained on RLHN-filtered BEIR data, tested on standard retrieval datasets (MS MARCO, Natural Questions, HotpotQA), and benchmarked against strong dense retrievers and concurrently proposed long-context retriever LMs such as MSA-4B.

Strong numerical results: BlockSearch matches dense retrieval on these benchmarks at practical corpus sizes (e.g., Recall@1 of 95.8% vs 95.5% for dense retrieval on MS MARCO at 45k tokens), and generalizes up to 10× its training context, surpassing previous ICR-style models.

Attention Dilution as a Core Failure Mode

At larger corpus sizes, BlockSearch’s retrieval accuracy collapses precipitously, despite its strong performance at smaller scales. A detailed mechanistic analysis attributes this collapse to attention dilution: as the number of context tokens grows, the attention softmax distribution flattens. Even if a head assigns maximal (pre-softmax) score to the gold document, the sheer number of competing documents causes the normalized attention weight on the gold document to shrink towards zero, effectively “drowning” the signal in noise. This effect is reflected in the fact that the aggregate attention output at the relevant layer shifts from being dominated by gold-token values to being almost entirely composed of distractor values, even though the overall vector magnitude remains approximately constant.

Key empirical finding: The gap between per-head attention recall and actual output accuracy exposes the bottleneck at the readout step—at least one head retains the gold document as top-ranked via QK-MaxSim, but readout via attention aggregation fails due to dispersion.

Proposed Interventions for Attention Dilution

To address the attention dilution bottleneck at extreme scales, the authors consider two families of interventions:

• Length-aware attention normalization:
  • Additive attention sinks: Learning a per-layer scalar appended to the softmax denominator, so that diffuse attention distributions are “gated” and contribute less to the residual stream.
  • SSMax (Multiplicative score rescaling): Scaling the pre-softmax scores by $s \cdot \log N$ so that the effective sharpness of the softmax increases with $N$, counteracting the distributional flattening caused by large $N$.
• Document-level sparse attention/routing: At an intermediate layer, a routing step scores all documents and selects a top-$B$ shortlist for subsequent dense attention, thus constraining downstream computation to a tractable subset.

Performance results:

• SSMax dramatically improves large-scale retrieval; for MS MARCO at $N=10,000$, Recall@1 rises from $0.2\%$ to $16.5\%$ (versus $20.2\%$ for the dense baseline).
• Top-$B$ routing enables BlockSearch to closely match or exceed the performance of dense retrievers and the much larger MSA-4B model at similar scales while using $7\times$ fewer parameters.
• Combining SSMax and routing yields the strongest configuration on several datasets, with Recall@1 that matches or exceeds dense retrieval even at one million tokens.

Generalization Beyond Embedding Retrieval

The authors test their method on LIMIT, a benchmark diagnosing cases where dense retrieval based on semantic embeddings is fundamentally limited due to non-semantic similarity requirements (e.g., lexical matching). On LIMIT, BlockSearch with the proposed interventions dramatically outperforms dense retrieval (e.g., $N$0 vs $N$1 Recall@1 at $N$2), demonstrating a key theoretical advantage: in-context retrieval using LMs subsumes more complex, flexible similarity functions unattainable for single-vector (embedding) retrieval architectures.

Practical and Theoretical Implications

This study establishes that with moderate model sizes, LMs equipped with suitable architectural and training regimes can retrieve robustly from million-token in-context corpora on par with or better than strong dense retrievers. This blurs the line between retriever and generator, challenging the need for external retrieval modules in large-scale RAG architectures, especially as models become further length-generalized and more flexible at runtime.

However, despite these advances, attention dilution remains a principal open problem: as $N$3 grows, even state-of-the-art attention mechanisms struggle to maintain high specificity in identifying relevant tokens. While interventions like SSMax or sparse routing recover much of the lost performance, a residual gap persists and constitutes a fundamental limitation for future research in scalable in-context retrieval. Moreover, synthetic benchmarks (e.g., “needle-in-a-haystack”) are shown to be too lenient; real retrieval tasks surface stronger regime shifts at extreme context scales.

Outlook and Future Directions

Open research directions include:

• Architectural innovations in attention that provide sharper, more robust selection at scale without manually constructed or sampled routing steps.
• Unified models for RAG that further collapse retrieval and generation into a single, trainable, dynamically-adapting LM, exploiting the flexibility revealed by these results.
• Distributed and hardware-aware designs for long-context LMs, tailored not only for prefill efficiency but also for inference throughput at million-token scales.

Conclusion

This work demonstrates that with the right architectural and loss modifications, moderately sized LMs can perform effective in-context retrieval over million-token corpora, matching and frequently exceeding dense retrieval on both standard and out-of-distribution benchmarks. The analysis pinpoints attention dilution as the central bottleneck for length generalization, highlights interventions that partially overcome this challenge, and sets a new baseline for future exploration into scalable, end-to-end LM retrievers. Both the practical and theoretical implications signal a shift towards architectures where LMs act as universal retrievers and generators, subject to further advances in attention control and context scaling.