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

• URL: http://arxiv.org/abs/2607.01538v1
• Date: Wed, 01 Jul 2026 23:38:25 GMT
• Title: Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale
• Abstract summary: In this work, we present the first systematic study of in-context retrieval on two scales practical retrievers demand.<n>We first introduce BlockSearch, a 0.6B LM retriever whose architectural and training modifications improve over prior LM baselines.<n>Motivated by this analysis, we introduce length-aware adjustments to the attention softmax and document-level sparse attention.
• Score: 22.13186135187399
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Surface-Form Neural Sparse Retrieval: Robust Fuzzy Matching for Industrial Music Search [6.820455959856271]
  Music search at the scale of Amazon Music presents a unique challenge: queries frequently deviate from indexed metadata due to misspellings, transpositions, and phonetic variations.<n>Our existing learning-to-retrieve system, the High Confidence Index (HCI), learns query-entity associations from customer behavior.<n>Traditional n-gram matching enables this exploration but suffers from poor semantic robustness and high noise.<n>We present a textbfrobust neural sparse retrieval system designed to maximize exploration efficiency.
  arXiv  Detail & Related papers  (2026-05-18T02:29:25Z)
• MM-Doc-R1: Training Agents for Long Document Visual Question Answering through Multi-turn Reinforcement Learning [74.07254720088926]
  We introduce MM-Doc-R1, a novel framework that employs an agentic, vision-aware workflow to address long document visual question answering.<n>We propose Similarity-based Policy Optimization (SPO), addressing baseline estimation bias in existing multi-turn reinforcement learning (RL) algorithms like GRPO.<n>Our experiments on the MMLongbench-Doc benchmark show that MM-Doc-R1 outperforms previous baselines by 10.4%.
  arXiv  Detail & Related papers  (2026-04-15T07:39:08Z)
• MARAG-R1: Beyond Single Retriever via Reinforcement-Learned Multi-Tool Agentic Retrieval [50.30107119622642]
  Large Language Models (LLMs) excel at reasoning and generation but are inherently limited by static pretraining data.<n>Retrieval-Augmented Generation (RAG) addresses this issue by grounding LLMs in external knowledge.<n>MarAG-R1 is a reinforcement-learned multi-tool RAG framework that enables LLMs to dynamically coordinate multiple retrieval mechanisms.
  arXiv  Detail & Related papers  (2025-10-31T15:51:39Z)
• LLM-guided Hierarchical Retrieval [54.73080745446999]
  LATTICE is a hierarchical retrieval framework that enables an LLM to reason over and navigate large corpora with logarithmic search complexity.<n>A central challenge in such LLM-guided search is that the model's relevance judgments are noisy, context-dependent, and unaware of the hierarchy.<n>Our framework achieves state-of-the-art zero-shot performance on the reasoning-intensive BRIGHT benchmark.
  arXiv  Detail & Related papers  (2025-10-15T07:05:17Z)
• SitEmb-v1.5: Improved Context-Aware Dense Retrieval for Semantic Association and Long Story Comprehension [77.93156509994994]
  We show how to represent short chunks in a way that is conditioned on a broader context window to enhance retrieval performance.<n>Existing embedding models are not well-equipped to encode such situated context effectively.<n>Our method substantially outperforms state-of-the-art embedding models.
  arXiv  Detail & Related papers  (2025-08-03T23:59:31Z)
• FrugalRAG: Learning to retrieve and reason for multi-hop QA [10.193015391271535]
  Large-scale fine-tuning is not needed to improve RAG metrics.<n>Supervised and RL-based fine-tuning can help RAG from the perspective of frugality.
  arXiv  Detail & Related papers  (2025-07-10T11:02:13Z)
• ELITE: Embedding-Less retrieval with Iterative Text Exploration [5.8851517822935335]
  Large Language Models (LLMs) have achieved impressive progress in natural language processing.<n>Their limited ability to retain long-term context constrains performance on document-level or multi-turn tasks.
  arXiv  Detail & Related papers  (2025-05-17T08:48:43Z)
• SAGE: A Framework of Precise Retrieval for RAG [9.889395372896153]
  Retrieval-augmented generation (RAG) has demonstrated significant proficiency in conducting question-answering tasks.<n>RAG methods segment the corpus without considering semantics, making it difficult to find relevant context.<n>We introduce a RAG framework (SAGE) to overcome these limitations.
  arXiv  Detail & Related papers  (2025-03-03T16:25:58Z)
• How Does Generative Retrieval Scale to Millions of Passages? [68.98628807288972]
  We conduct the first empirical study of generative retrieval techniques across various corpus scales. We scale generative retrieval to millions of passages with a corpus of 8.8M passages and evaluating model sizes up to 11B parameters. While generative retrieval is competitive with state-of-the-art dual encoders on small corpora, scaling to millions of passages remains an important and unsolved challenge.
  arXiv  Detail & Related papers  (2023-05-19T17:33:38Z)
• Overcoming Classifier Imbalance for Long-tail Object Detection with Balanced Group Softmax [88.11979569564427]
  We provide the first systematic analysis on the underperformance of state-of-the-art models in front of long-tail distribution. We propose a novel balanced group softmax (BAGS) module for balancing the classifiers within the detection frameworks through group-wise training. Extensive experiments on the very recent long-tail large vocabulary object recognition benchmark LVIS show that our proposed BAGS significantly improves the performance of detectors.
  arXiv  Detail & Related papers  (2020-06-18T10:24:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.