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

Language models (LMs) raise an alluring alternative to vector-based retrieval: generating a relevant answer, entirely based on an in-context corpus. In spite of this appeal, prior studies consider 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 LCLM retriever whose architectural and training modifications improve over prior LM baselines and length-generalize up to $10\times$ beyond its training regime. However, 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, minimizing the normalized mass on the gold document even as its pre-softmax score stays high. Motivated by this analysis, we introduce length-aware adjustments to the attention softmax and document-level sparse attention, improving retrieval at million-token scale, better performance than the concurrent $7\times$-larger MSA system, and comparable performance to dense retrieval. Together, our results position in-context retrieval as a viable alternative to classical retrieval, emphasizing attention control under extreme context growth as a new challenge.