Research map / Large language models / Retrieval and knowledge

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

NeurIPS 2026 · Large language models › Retrieval and knowledge · approach: Retrieval-augmented reasoning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Systematic study of corpus-scale in-context retrieval with a 0.6B LM (BlockSearch), diagnosing attention dilution as the cause of collapse at million-token scale and mitigating it with length-aware softmax (SSMax) and document-level top-B routing.

Problem

Whether language models can retrieve relevant documents directly from large in-context corpora (up to million tokens, and beyond training-length corpus sizes), where prior work relied on proprietary systems or small-candidate reranking.

Approach

Train a Qwen3-0.6B in-context retriever with random per-document codes, in-batch negatives, and an on-policy auxiliary loss; diagnose collapse as attention dilution (gold's softmax mass shrinks as corpus grows) and mitigate with log-N attention score scaling (SSMax) and document-level top-B routing at layer 16.

Key results

• BlockSearch matches the same-backbone dense retriever at small corpus size (95.8 vs 95.5 Recall@1 on MS MARCO at N=500).
• SSMax plus top-B=256 routing raises million-token (N=10k, ~1M tokens) MS MARCO Recall@1 from 0.2 to 20.0, close to dense baseline 20.2.
• On LIMIT at ~1.7M tokens, SSMax+routing reaches Recall@1 0.135 versus 0.029 for the dense baseline.

Surprising findings

• Individual attention heads still rank the gold document first more often than the model decodes it, even at N=10,000 (R_any 0.08 at L19; best layer 0.14).
• An additive learned attention sink barely helps at large N, while multiplicative log-N score scaling (SSMax) does.

Limitations

• Keeping the corpus in a KV cache requires roughly three orders of magnitude more memory than a dense index.
• Retrieval-only training reduces the backbone's general capability.
• Evidence comes from two small backbones and a single sample of 400 queries per corpus size.

Experimental setup

Datasets: RLHN-680K (filtered) MS MARCO HotpotQA Natural Questions LIMIT

Benchmarks: MS MARCO HotpotQA Natural Questions LIMIT

Metrics: Recall@1 Recall@5

Baselines: Qwen3-dense (0.6B) MSA-4B BlockSearch-position BlockSearch-offpolicy Gemini-3-flash BlockSearch-sink

Builds on: block-sparse attention for in-context reranking SSMax GPT-OSS attention sinks RLHN DAgger

Closest papers by meaning

• Scalable In-context Ranking with Generative Models (NeurIPS 2025 · similarity 0.90)
• Retrieval meets Long Context Large Language Models (ICLR 2024 · similarity 0.90)
• Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism (ICML 2025 · similarity 0.89)
• Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT (ICML 2024 · similarity 0.89)
• Making Retrieval-Augmented Language Models Robust to Irrelevant Context (ICLR 2024 · similarity 0.88)
• Efficient Length-Generalizable Attention via Causal Retrieval for Long-Context Language Modeling (ICML 2025 · similarity 0.88)