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

Technical Summary: Can Language Models Actually Retrieve In-Context?

Problem Statement
While language models (LMs) offer a potential alternative to vector-based retrieval by conditioning on an in-context corpus to directly generate relevant document identifiers, this capability remains largely unexplored at realistic scales. Prior work has focused on proprietary systems or small-scale reranking tasks, leaving two critical challenges unaddressed: (C1) scaling to corpora containing millions of tokens, and (C2) generalizing to corpus sizes far beyond those seen during training. Current LMs struggle with both million-token attention costs and the brittleness of extrapolation to longer contexts. The central question is whether an LM can reliably act as a corpus-scale retriever under these conditions without collapsing.

Methodology: BLOCKSEARCH and Analysis
The authors introduce BLOCKSEARCH, a systematic study and model architecture designed to tackle corpus-scale in-context retrieval (ICR).

• Architecture & Training: Built on a 0.6B parameter backbone (Qwen3-0.6B), BLOCKSEARCH employs block-sparse attention to handle disjoint documents efficiently. Key architectural and training modifications include:

  • Randomized Document Identifiers: Instead of sequential IDs, documents are assigned random four-digit codes per training step to prevent overfitting to absolute positions.
  • In-Batch Negatives & On-Policy Loss: The model is trained with 1 positive and 15 hard negatives per query. To mitigate exposure bias (where the model conditions on its own potentially incorrect generated prefixes during inference), an on-policy auxiliary loss is added. This involves rolling out a code from the model's own distribution and training the model to correct its path based on in-batch scores.
  • Evaluation Setup: The model is evaluated on BEIR datasets (MS MARCO, NQ, HotpotQA) with corpus sizes ranging from 500 to 10,000 documents (up to ~1 million tokens), far exceeding the model's native 32k token context window.

• Mechanistic Analysis of Failure: The study identifies that while BLOCKSEARCH improves extrapolation, it still collapses at large scales (near-zero recall at 1M tokens). The authors trace this not to a failure in identifying the relevant document (ranking), but to an attention dilution effect:

  • Retrieval vs. Generation Gap: Even at 10,000 documents, individual attention heads often still assign the highest pre-softmax score to the gold document ($R_{any} \approx 1.0$).
  • Softmax Collapse: However, as the corpus ($N$) grows, the aggregate contribution of irrelevant documents to the softmax denominator increases faster than the gold document's score. Consequently, the normalized attention mass (post-softmax weight) on the gold document collapses, causing the layer's output to become an average of distractors rather than the relevant content.

• Proposed Mitigations: To counteract attention dilution, the authors propose two length-aware adjustments:

  1. Length-Aware Softmax Adjustments:
     • Additive Sinks: Appending a learned constant to the softmax denominator to absorb diffuse attention mass (inspired by Streaming-LLM).
     • Multiplicative Rescaling (SSMax): Scaling pre-softmax scores by $s \cdot \log N$ to widen the gap between the gold score and distractors, preventing the gold weight from vanishing as $N$ increases.
  2. Document-Level Sparse Attention: A routing mechanism at an intermediate layer (L16) that scores all documents and restricts dense attention to the top-$B$ (e.g., 256) documents for subsequent layers, effectively pruning the context before the final generation.

Key Results

• Baseline Performance: The initial BLOCKSEARCH model matches dense retrieval performance on standard benchmarks (MS MARCO, NQ) at smaller corpus sizes (~45k tokens) and generalizes up to ~500k tokens (10x training context), but collapses near 1M tokens.
• Impact of Modifications:
  • SSMax: Significantly improves performance across the full N sweep. On MS MARCO at 1M tokens, it achieves 16.5% Recall@1 (vs. 0.2% for the unmodified model), closing most of the gap to the dense baseline (20.2%).
  • Routing: Document-level routing further improves performance, achieving 18.8% Recall@1 on MS MARCO at 1M tokens, nearly matching the dense baseline.
  • Combined Approach: The combination of SSMax and routing (BLOCKSEARCH-SSMax-routing) achieves 20.5% Recall@1 on MS MARCO at 1M tokens, slightly edging out the dense baseline (20.2%) and outperforming the concurrent 7x larger MSA-4B model (16.0%).
• Out-of-Distribution (LIMIT): On the LIMIT benchmark, which requires lexical similarity rather than semantic embedding similarity (where dense retrieval struggles), the modified BLOCKSEARCH significantly outperforms the dense baseline. At 850k tokens, BLOCKSEARCH-SSMax-routing achieves 14.9% Recall@1 compared to 3.5% for the dense retriever (a ~3x improvement).
• Attention Ceiling: Across all experiments, the "attention ceiling" (the fraction of queries where at least one head ranks the gold document first) remains near 100%, confirming that the retrieval signal is present internally but lost during the readout/generation phase due to attention dilution.

Significance and Claims
The paper claims to provide the first systematic study of in-context retrieval at the million-token scale, establishing it as a viable alternative to classical dense retrieval under specific conditions.

• Viability: With architectural tweaks (SSMax and routing), LMs can match or exceed dense retrieval on standard benchmarks and significantly outperform them on tasks requiring non-semantic similarity.
• Efficiency: The proposed 0.6B model achieves these results with 7x fewer parameters than concurrent large-context models (MSA-4B) and without the need for retraining when the corpus changes (unlike generative retrieval).
• Fundamental Bottleneck: The work identifies attention dilution (softmax dispersion) as the primary bottleneck for scalable in-context retrieval. While simple interventions can mitigate this, the paper emphasizes that controlling attention under extreme context growth remains a fundamental challenge for future retrieval architectures.

The authors remain modest, noting that while their methods extend functional retrieval to larger corpora, they do not fully "solve" the problem of retrieval at extreme scales, as performance still declines as $N$ increases, and a residual gap to dense retrieval remains on some datasets.