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

ChapterPal

Siddharth Gollapudi$^{*}$
UC Berkeley

Nilesh Gupta
UT Austin

Prasann Singhal
UC Berkeley

Sewon Min
UC Berkeley

$^{*}$ Correspondence to [email protected].

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.

Executive Summary: Language models have long been positioned as potential alternatives to traditional vector-based retrieval, where a query is matched against a separate index of document embeddings. The paper investigates whether an LM can instead receive an entire corpus directly in its context and generate the identifier of the relevant document, an approach called in-context retrieval. This could simplify retrieval-augmented generation pipelines and support richer notions of relevance, but prior evaluations had been limited to small candidate sets or proprietary systems, leaving performance at realistic million-token scales unknown.

The work set out to determine whether a modest-sized LM can perform true corpus-scale retrieval, generalize to contexts much larger than those seen in training, and match or exceed specialized dense retrievers. It also aimed to diagnose any failures and test targeted remedies.

The authors built BlockSearch, a 0.6-billion-parameter model based on Qwen3, incorporating random per-document identifier codes, in-batch negative sampling, and an on-policy auxiliary loss to reduce exposure bias during decoding. They trained on filtered hard-negative data from BEIR and evaluated on standard benchmarks (MS MARCO, Natural Questions, HotpotQA) plus the lexical-similarity benchmark LIMIT. Experiments scaled the corpus from hundreds to 10,000 documents (roughly one million tokens), with systematic attention measurements to isolate where the retrieval signal was lost.

At moderate scales the model performed competitively, matching a same-sized dense retriever and a much larger concurrent in-context model. At extreme scale, however, retrieval collapsed even though at least one attention head continued to rank the gold document highest; the normalized attention mass on the correct document dropped sharply because distractors increasingly dominated the softmax denominator. Length-aware softmax rescaling and document-level routing at an intermediate layer largely restored performance, bringing million-token recall close to the dense baseline on standard tasks and tripling it on LIMIT.

These results indicate that in-context retrieval is now practical at the scales demanded by real retrievers, provided attention dilution is controlled. The approach can therefore serve as a drop-in or complementary option to dense retrieval, particularly when similarity is not well captured by a single embedding. At the same time, the residual performance gap at the largest sizes and the reintroduction of an internal retrieve-then-read step via routing show that attention mechanisms remain a limiting factor.

Organizations considering deployment should adopt the proposed softmax and routing adjustments when using generative retrievers at scale, and should run targeted evaluations on their own notion of relevance rather than relying solely on MS MARCO-style benchmarks. Further work on stronger length-adaptive attention and on closing the remaining gap between internal attention signals and final output is needed before fully replacing two-stage pipelines. The main limitation is that all results are confined to the evaluated datasets and a 0.6 B backbone; larger-scale or domain-specific corpora may surface additional degradation modes.

1. Introduction

Section Summary: Recent work has explored whether language models could retrieve documents by directly generating their identifiers from a large set of passages placed in context, potentially simplifying retrieval-augmented systems. This paper introduces a compact model called BlockSearch that improves such in-context retrieval through architectural tweaks and specialized training, yet still fails at million-token scales because attention scores spread too thinly across irrelevant documents. The authors show that targeted fixes, including length-aware softmax adjustments and selective document attention, can raise accuracy substantially and make the approach competitive with traditional dense retrievers on standard benchmarks.

Retrieval—identifying the relevant document(s) in a corpus ${D_1, \dots, D_N}$ for a query $Q$ —has long been dominated by vector-based methods [1]. Recent advances in language modeling (LMs) suggest a more radical alternative: casting retrieval as conditional generation, where the model directly decodes the identifier of the relevant document from the corpus in-context ([2]). Such in-context retrieval (ICR) could collapse retrieval and generation into a single model, replacing the two-stage retrieval-augmented generation pipeline ([3]) and enabling complex retrieval behavior beyond inner-product similarity ([4]). In other words, rather than relying on external harnesses (retrievers), the model itself decides which parts of its context are relevant to the task at hand. However, this possibility remains largely untested at realistic scale: prior work either relies on proprietary systems without controlled evaluation ([2]) or studies reranking over small candidate sets ([5, 6]), rather than true corpus-scale retrieval. As a result, it remains unclear whether LMs can reliably retrieve relevant documents from large in-context corpora ([7, 6, 8]).

ICR differs fundamentally from standard long-context modeling. Rather than processing a single coherent sequence, it operates over large collections of independent documents, creating both opportunities and challenges. On one hand, document-level independence naturally enables parallel encoding and caching. On the other hand, effective retrieval requires two capabilities that current LMs struggle with: (C1) scaling to corpora containing millions of tokens, and (C2) generalizing to corpus sizes far beyond those seen during training. Both are difficult in practice: million-token attention remains computationally expensive, while extrapolation to substantially longer contexts is known to be brittle in modern LMs ([7, 9, 10]).

To this end, we present the first systematic study of whether an LM can act as a corpus-scale retriever under these conditions. We first introduce $\textsc{BlockSearch}$ (Section 3), a 0.6B-parameter LM retriever that builds on prior ICR architectures using block-sparse attention ([5]). For scale and length generalization, it adds randomized per-document identifier codes, in-batch negative training, and an on-policy auxiliary loss that trains on the model's own rolled-out digit prefixes. On widely studied retrieval benchmarks (e.g., MS MARCO and NQ), $\textsc{BlockSearch}$ matches a dense retriever at smaller corpus sizes ($95.8%$ vs. $95.5%$ on MS MARCO at $\sim$ 45k tokens), and maintains non-trivial accuracy beyond the LM baseline, generalizing up to $10\times$ beyond its training context.

Despite these improvements, $\textsc{BlockSearch}$ still collapses at large scale: recall approaches zero near million-token contexts, remaining far below a simple dense retrieval baseline. We trace this failure not to a ranking collapse, but to an attention dilution effect, a degradation that has also been observed in other settings ([9, 11, 12, 13]). Leveraging the natural alignment between queries and their gold documents, we find that the transformer usually assigns the highest pre-softmax attention score to the gold document even at million-token scale, but the aggregate contribution of irrelevant documents to the softmax denominator grows faster with corpus size, causing the normalized attention mass on the gold document to collapse (Section 4).

Motivated by this analysis, we consider two techniques to mitigate attention dilution at extreme context lengths (§ 5). First, length-aware sinks reshape the effective softmax denominator, reducing the influence of diffuse attention without modifying QK scores. Second, document-level sparse attention reduces the number of documents participating in attention at intermediate layers ([14, 15, 10, 16, 17]). Together, these methods mitigate the effects of attention dilution. On well-studied retrieval benchmarks (e.g. MSMARCO and NQ from BEIR [18]), they improve million-token Recall@1 from as little as $0.2%$ to $20.5%$, recovering the gap to dense retrieval. They also match or exceed MSA ([16]), a concurrent LM trained on much longer contexts with $7\times$ more parameters. On benchmarks requiring more complex notions of similarity, where dense retrieval tends to fail, such as LIMIT ([4]), our method exceeds Qwen3-dense by nearly $3\times$ (Section 6).

Our results establish corpus-scale in-context retrieval as a viable alternative to dense retrieval under both scale and length extrapolation. More broadly, they identify attention dilution as the primary bottleneck at million-token scale: while simple interventions can substantially mitigate its effects, it remains a fundamental challenge for scalable in-context retrieval.

2. Related Work

Section Summary: This section reviews prior research on using language models for retrieval from long contexts, noting that most evaluations rely on synthetic tasks and that earlier approaches either rerank small document sets or embed the entire corpus into model parameters, requiring retraining for updates. It contrasts these with the current focus on corpus-scale, in-context retrieval that supports easy updates without retraining and builds on techniques like sparse or routed attention to handle very long inputs efficiently. The review also discusses recent observations about attention “sinks” and the flattening of attention distributions as context grows, positioning these issues as central to retrieval degradation and highlighting how the proposed methods address them more efficiently than concurrent large-scale efforts.

oldparagraphIn-context and generative retrieval. Most long-context evaluation work focuses on synthetic retrieval-style tasks such as Needle-in-a-Haystack and RULER ([19, 7, 8]), potentially overstating real-world retrieval capability. The premise that an LM can subsume dense retrieval is first articulated by [2], primarily evaluating proprietary Gemini systems. Subsequent work studies LMs as rerankers over small candidate pools using block-sparse attention and relevance heads in intermediate Transformer layers ([6, 5]). In contrast, we study LMs as corpus-scale retrievers, training with small contexts and evaluating up to $10^4$ documents ($\sim$ 1M tokens) in an effort to elicit length-generalization.

Generative retrieval ([20, 21, 22]) also decodes with identifiers, but stores corpus information in model parameters, requiring retraining when the corpus changes. In contrast, ICR encodes the corpus explicitly in context, allowing corpus updates without retraining.

Sparse and routed attention.

A large body of work makes long-context attention tractable by routing or sparsification; our document-level sparse attention proposed in Section 4 builds on this line of work. Prior approaches reduce attention cost through inference-time pruning ([23, 24]), block-sparse kernels ([17]), learned block routing ([14, 15]), or hierarchical sparse attention ([10, 25]). Our work is grounded in this literature, but the retrieval setting introduces unique opportunities and challenges.

The most directly comparable work is concurrent: MSA-4B ([16]), which, like $\textsc{BlockSearch}$, learns to perform generative document citation from an in-context corpus. However, the two works focus on different questions: MSA primarily scales the retrieval recipe through larger models and extensive training, whereas our work studies the fundamental failure modes of ICR as corpus size grows. In fact, our proposed methods targeting these failure modes match or exceed MSA despite using $7\times$ fewer parameters and significantly shorter training-time contexts (§ 5).

Attention sinks and softmax dilution.

Streaming-LLM ([26]) introduces explicit attention sinks for stable long-context decoding, and the gpt-oss model card ([27]) documents a null-attention mechanism that absorbs unused softmax mass. Whereas these works use sinks for streaming stability, we apply them to mitigate softmax dilution under large $N$. A number of results show that as the number of attended tokens grows, the softmax distribution flattens, removing attention peaks ([9, 11, 13]); this loss of sensitivity has been tied to representational collapse in longer contexts ([12, 28]). We show that the primary cause of retrieval degradation with growing context is softmax dilution, and that softmax-adjusting mechanisms such as sinks may hold the key to mitigating its effects.

3. BlockSearch: Model and Million-Token Evaluation

Section Summary: This section introduces BlockSearch, a set of training improvements to long-context language models that enable reliable retrieval of relevant documents from corpora far larger than the model’s native context window. The approach modifies the standard retrieval setup by randomizing document identifiers during training, applying block-sparse attention to handle up to millions of tokens efficiently, and using in-batch negatives plus an auxiliary loss to reduce exposure bias when generating document codes. Evaluations on challenging BEIR-derived benchmarks with 10,000-document corpora show that these changes allow the model to generalize effectively to scales well beyond training.

We first investigate several training-time improvements to the standard long-context retrieval recipe. After reviewing the problem setup (Section 3.1), we describe the changes that constitute $\textsc{BlockSearch}$ (Section 3.2), our evaluation setup (Section 3.3), and the results (Section 3.4). These results motivate the later analysis of possible failure modes (Section 4) and promising lines for even further improvements (Section 5).

3.1 Setup

Given a corpus of $N$ documents and a query, the model must generate the identifier of the document that best answers the query. The corpus is tokenized with at most $T_{\text{doc}}{=}300$ tokens per document and prefilled into the KV cache; the query is appended, and the model autoregressively decodes a four-digit code in ${0, \dots, 9999}$ that maps back to one document (that's hopefully relevant).

This concretizes challenges C1 and C2 from Section 1: just 10, 000 documents at $\sim$ 100 tokens each well exceeds the 32K-token native context of the Qwen3-0.6B backbone, and the model must generalize to corpora substantially larger than those seen during training: again, a requirement that distinguishes retrievers from rerankers and has been largely overlooked in prior LM-retrieval work.

3.2 Method: BlockSearch

We now describe $\textsc{BlockSearch}$: whenever possible, we follow prior ICR work, but our requirements C1 and C2 expose several limitations that motivate our proposed improvements.

Prompt format.

Prior in-context work assigns each document a sequential integer identifier from $1$ to $N$ ([5, 2]). In our regime, training-time IDs are a strict subset of inference IDs and risk overfitting to absolute position. We instead insert each document as <bos>Doc code: text (Doc code)<eos>, with code drawn uniformly at random per training step, breaking any association between code, semantics, or position. We also drop the query prefix from prior work ([5, 2]): while it improves performance, it prevents corpus reuse across queries, which is cost-prohibitive at our scale. For a full example of a corpus/query prompt, we refer the reader to Appendix A.

Block-sparse attention.

Causal prefill scales quadratically in tokens, which is prohibitive at our target $N$. Since documents are disjoint, we use a block-sparse mask ([5]): document tokens attend causally only within their own block, and the query block attends over the full corpus and causally to itself. Following common practice ([29, 5]), we reset RoPE ([30]) positions at each document start and shift the query to position 300. The mask is materialized via flex_attention ([31]).

Training data.

We train on the ReLabeled Hard Negatives (RLHN) version of BEIR ([32]), which uses an LLM judge to prune false negatives from mined hard negatives in datasets such as MS MARCO ([33]). Per query we keep one pruned positive and 15 hard negatives; per-query relevance scores for the auxiliary loss are from Qwen3-Embedding-8B ([34]). The training corpus totals $\sim$ 100M tokens; full mix and cleaning details are in Appendix C.

Training regime.

Inspired by in-batch contrastive training ([1, 35]), we adapt in-batch negatives to this setting. For a batch of $b$ (query, 16-document) tuples (1 positive + 15 hard negatives each), we prefill all $b{\times}16$ documents once (re-randomizing codes) and score every query against the shared corpus, converting one prefill into $b$ training signals. Every model is trained with 256 documents in its corpus, i.e. $b{=}16$.

Training objective.

We train with teacher forcing on the gold code, but this alon causes an exposure bias issue: at inference, each digit conditions on the model’s own potentially incorrect prefix, which lies off the training distribution. We therefore add an on-policy auxiliary loss. The model first rolls out a four-digit code from its own distribution with gradients disabled. For each digit position, we construct a teacher distribution from in-batch document scores, restricted to candidates whose prefix matches the rollout. We then replay the rollout with gradients enabled at the four answer positions and average the resulting cross-entropy losses. The total loss is $\mathcal{L} = \mathcal{L}{\mathrm{CE}} + \lambda\mathcal{L}{\mathrm{aux}}$ with $\lambda$ ramped in after a warmup; the full algorithm can be found in Algorithm 1 (Appendix F).

3.3 Evaluation

The model is judged by whether it can recover the relevant document from those placed in context; our metric is Recall@1 over the generated four-digit codes, decoded by a beam search over the digit sequence. We also evaluate Recall@5 using beam search over the top five predictions and observe trends consistent with Recall@1; see Appendix J for details.

Evaluation data.

We evaluate on three BEIR datasets ([18]): MS MARCO (dev), HotpotQA (test), and NQ (test). We choose these datasets because they are among the most widely studied retrieval benchmarks, on which dense retrievers have been extensively optimized, providing a strong calibration point for comparison. We later complement these evaluations with benchmarks requiring different notions of similarity, where dense retrieval struggles (§ 6).

For each dataset, we sample 400 queries without replacement from the its split. For each query, we take its gold document and 24 hard negatives (retrieved using Qwen3-Embedding-8B [36]); by taking a union of these documents across all 400 queries, we get a 10, 000 document corpus. In contrast with prior work [2], the substantial use of hard negatives makes the input corpus far more challenging.[^1]

[^1]: For NQ, we use fewer hard negatives per query due to longer documents; per-dataset sizes are in Appendix B (Table 4).

To evaluate across varying corpus sizes (e.g. $N{=}500$ to $N{=}10{,}000$), we also construct smaller versions of each dataset. We first include every query's gold document (400 documents), then fill out the rest of the dataset with documents randomly selected from the full 10, 000 document corpus.

Baselines.

We compare against two reference retrievers and two ablations of $\textsc{BlockSearch}$ itself:

$\textsc{BlockSearch}$-position: sequential codes ($0, \dots, 9999$ by document position) to isolate the contribution of random codes, and no auxiliary loss; this corresponds to the prior ICR recipe.
$\textsc{BlockSearch}$-offpolicy: $\textsc{BlockSearch}$ with the on-policy auxiliary loss removed (teacher-forced CE only); isolates the exposure-bias fix.
Qwen3-dense (0.6B): a dense retriever trained on the same RLHN data ([32]) with a contrastive objective ([1]) from the same Qwen3-0.6B backbone (Appendix E); our primary reference point and the dense-retrieval gold standard $\textsc{BlockSearch}$ must clear to be practical.
MSA-4B ([16]): a concurrent multi-million-token LM, $\sim 7\times$ larger and trained with a much larger long-context budget, sidestepping the C2 generalization requirement; an oracle reference.

3.4 Results: BlockSearch is competitive, but collapses as $N$ scales

$**Figure 1:** Recall@1 vs. approximate prefill size ($N \times$ mean doc length; see Appendix B) on NQ, MS MARCO, and HotpotQA. Solid: ICR variants ($\textsc{BlockSearch}$, $\textsc{BlockSearch}$-position, $\textsc{BlockSearch}$-offpolicy); dashed: baselines (`Qwen3-dense`, MSA-4B). Note that the Qwen3-0.6B native context limit is $\sim$ 32k tokens.$

We train Qwen3-0.6B into $\textsc{BlockSearch}$ (and variants) on $8\times$ NVIDIA A100s using the above changes; full hyperparameters are in Appendix D. The performance of the models is demonstrated in Figure 1. Note that since HotpotQA has multiple golds, a query is successful in recall@1 as long as one of the golds is the best result. This is primarily to ensure compatibility with MSA, which does not support more than recall@1.

BlockSearch improves extrapolation in LMs, up to 500,000 tokens.

All LMs perform strongly at small $N$ (

gt;95%$ on MS MARCO at $N{=}1{,}000$), but baselines degrade much faster. The position-coded variant, which is the prior ICR method, collapses to near-zero by $N{=}5{,}000$ on every dataset. $\textsc{BlockSearch}$-offpolicy trails $\textsc{BlockSearch}$ (MS MARCO $15.0$ vs. $18.8$ at $N{=}5{,}000$; HotpotQA $6.2$ vs. $13.0$), showing that random codes and the on-policy loss contribute distinct gains. Note that, while $\textsc{BlockSearch}$ substantially improves extrapolation, all variants eventually collapse beyond roughly $500{,}000$ tokens, corresponding to about $10\times$ the training context length.

BlockSearch vs. MSA-4B.

Despite having $\sim 7\times$ fewer parameters than MSA ([16]), $\textsc{BlockSearch}$ matches MSA at $N{=}500, 1000$ and $2500$, e.g., on MS MARCO, achieving $95.8%, 75.2%, 43.8%$ vs. $93.8%, 70.2%, 42.2%$ for MSA. At large corpus sizes, MSA-4B pulls ahead ($27.5%$ vs. $18.8%$ at $N{=}5{,}000$, $16.0%$ vs. $0.2%$ at $N{=}10{,}000$), which is expected given that MSA is explicitly trained for long-context tasks, whereas these settings require $\textsc{BlockSearch}$ to generalize to corpora $20$ – $40\times$ larger than those seen during training.

Despite near-perfect long-context performance on synthetic evaluation such as RULER's needle-in-a-haystack (NIAH) ([19]) as reported in the original paper ([16]), MSA degrades sharply on real retrieval tasks: substantially more than the dense retrieval baseline. This suggests that realistic retrieval settings are significantly more challenging than existing synthetic long-context benchmarks, highlighting the need for the long-context literature to adopt more semantically meaningful retrieval tasks in evaluation.

BlockSearch vs. dense retrieval.

Dense retrieval also degrades with $N$ but without the sharp LM collapse: $\textsc{BlockSearch}$ matches the Qwen3-dense gold standard at small and mid $N$ but trails it at large $N$ (MS MARCO $18.8$ vs. $38.5$ at $N{=}5{,}000$; near-zero vs. $20.2$ at $N{=}10{,}000$), with similar trends on NQ and HotpotQA. Closing this large- $N$ gap remains the central challenge for practical in-context retrieval and motivates the mechanistic analysis in Section 4.

4. Understanding the Large- $N$ Deterioration

Section Summary: The section explains that BlockSearch’s breakdown at very large document counts stems from two distinct failures: individual attention heads still tend to spotlight the correct document even at N=10,000, yet the model stops producing the right identifier token. At the same time, the fraction of each layer’s output that is driven by tokens from the gold document steadily shrinks, even though the layer’s overall output magnitude remains stable. This pattern isolates an “attention dilution” effect that accumulates across the retrieval and decoding stages without destroying the raw relevance signal.

Section 3 shows that $\textsc{BlockSearch}$ generalizes well past its training-time context length, but still collapses at $N{=}10{,}000$. In this section, we posit that this collapse can be decomposed into two parts: a per-head, per-layer measure of attention recall (does the relevant document receive the highest attention score?) remains deep into the large- $N$ regime, while generation recall (does the model decode the correct identifier?) collapses. We also find that the share of each attention layer's output that comes from gold tokens collapses while the layer's overall output magnitude is largely preserved, cleanly encapsulating the aformentioned attention dilution issue.

4.1 Setup

Decoder language models are not a homogeneous stack: layers and heads take on specific functional roles, with FFNs promoting concepts ([37, 38]), induction heads carrying in-context learning ([39]), and a sparse subset of retrieval heads carrying the long-context retrieval signal ([40]). For in-context retrieval specifically, BlockRank ([5]), ICR$^2$ ([6]), and QRHeads ([41]) build on those results, observing that per-head attention scores can serve directly as a relevance signal. This observation is the starting point for the measurements we report.

Preliminary.

We write $L \in {0, \dots, 27}$ for layer index, and $h \in {1, \dots, H{=}16}$ for attention-head index. We probe $\textsc{BlockSearch}$ at the final token of the query block, immediately before the first generated code digit (which we denote $d_1$). At this position let $q^{h}{L} \in \mathbb{R}^{d{\text{head}}}$ be the post-RoPE query vector and $k^{h}{L, t}, v^{h}{L, t}$ the key and value at prefill position $t$, where $d_{\text{head}}$ is the per-head dimension. The pre- and post-softmax attention scores assigned to position $t$ are:

$ s^{h}{L, t} ;=; \frac{q^{h}{L} \cdot k^{h}{L, t}}{\sqrt{d{\text{head}}}}, \qquad \alpha^{h}{L, t} ;=; \frac{\exp(s^{h}{L, t})}{\sum_{t'} \exp(s^{h}_{L, t'})},\tag{1} $

and the layer's attention output is

$ a_L ;=; \bigl[, u^{1}{L};\ u^{2}{L};\ \dots;\ u^{H}{L} , \bigr], W^O_L, \qquad u^{h}{L} = \sum_t \alpha^{h}{L, t}, v^{h}{L, t},\tag{2} $

where $u^{h}{L} \in \mathbb{R}^{d{\text{head}}}$ is the per-head attention readout, $[, \cdot, ;, \cdot,]$ denotes concatenation along the head dimension into a vector in $\mathbb{R}^{H \cdot d_{\text{head}}}$, and $W^O_L \in \mathbb{R}^{(H \cdot d_{\text{head}}) \times d_{\text{model}}}$ is the layer's learned output projection.

Metrics.

For each query, partition the prefill tokens into the gold set $\mathcal{G}$ (the tokens of the gold document) and its complement $\bar{\mathcal{G}}$. We use three measurements throughout this section: AttnRank, GoldShare, and first-digit accuracy.

AttnRank measures whether attention identifies the right document. More precisely, for each document $d$ in the prefill, we collapse its tokens to a single per-head score using the MaxSim operator from late-interaction retrieval ([42]):

$ \mathrm{MaxSim}^{h}{L}(d) ;=; \max{t \in d} s^{h}_{L, t}.\tag{3} $

Writing $d^\star$ for the gold document, the per-head Recall@1 at layer $L$ is $R^{(h)}L$, the fraction of queries for which $\arg\max_d \mathrm{MaxSim}^{h}{L}(d) = d^\star$. We report two cross-head aggregators of $R^{(h)}_L$:

$R^{\text{sum}}L$ — sum-across-heads: rank documents by $\sum_h \mathrm{MaxSim}^{h}{L}(d)$ and check whether the top-ranked document is $d^\star$.
$R^{\text{any}}_L = $ fraction of queries for which at least one head puts $d^\star$ first.

GoldShare measures gold's contribution to the layer's attention output (before the residual add). By linearity of Equation 2 in the per-head readouts $u^{h}_{L}$, we decompose $a_L$ into a gold and a non-gold piece:

$ a_L ;=; a_L^{\mathcal{G}} + a_L^{\bar{\mathcal{G}}}, \qquad a_L^{\mathcal{G}} ;=; \bigl[, u^{1, \mathcal{G}}L;\ \dots;\ u^{H, \mathcal{G}}L , \bigr], W^O_L, \quad u^{h, \mathcal{G}}L = \sum{t \in \mathcal{G}} \alpha^{h}{L, t}, v^{h}{L, t},\tag{4} $

and report the gold-driven fraction $|a_L^{\mathcal{G}}| / |a_L|$ alongside the total magnitude $|a_L|$.

First-digit accuracy probes how decisively the model has committed to the first code digit at a given layer, via an LM-head probe ([43]): we project the layer- $L$ output through the model's final RMSNorm and lm_head and record the maximum probability assigned to any of the ten digit tokens at the $d_1$ position. We treat this as a coarse indicator of when the digit decision emerges, not as a calibrated probability at intermediate layers.

4.2 Layer Roles and the Recall–Generation Gap

Figure 2 reports $R^{\text{sum}}_L$ and first-digit accuracy on $\textsc{BlockSearch}$ at the $d_1$-emission position in the in-domain, small- $N$ regime ($N{=}500$). Two qualitatively distinct transitions emerge. First, $R^{\text{sum}}_L$ rises across $\mathrm{L}11$ – $\mathrm{L}20$: the model accumulates relevance information over many layers, consistent with prior reports of middle-layer retrieval in decoder LMs ([5, 6, 41, 40]). We refer to this contiguous range as the retrieval band. Second, first-digit accuracy rises sharply at $\mathrm{L}18{\to}\mathrm{L}19$, after which the model has effectively committed to $d_1$; we refer to $\mathrm{L}19$ onwards as the decode band.

$**Figure 2:** Per-layer $R^{\text{sum}}_L$ and LM-head first-digit accuracy on $\textsc{BlockSearch}$, MS MARCO, $N{=}500$. Retrieval signal emerges across $\mathrm{L}11$ – $\mathrm{L}20$; first-digit signal emerges sharply at $\mathrm{L}18{\to}\mathrm{L}19$.$

Heads no longer agree on the gold document.

Sweeping $N \in {500, 1000, 2500, 5000, 10000}$ on MS MARCO (Figure 3), we see these behaviors degrade. The retrieval band erodes under the sum aggregator: $R^{\text{sum}}_{19}$ drops from $0.97$ at $N{=}500$ to $0.01$ at $N{=}10{,}000$, matching the collapse in generation recall reported in Section 3.4. $R^{\text{any}}L$, however, stays at $1.00$ across $\mathrm{L}18$ – $\mathrm{L}25$ at every $N$: at every $N$, at least one head per query still ranks the gold document first. In other words, per-head retrieval signal persists with large $N$, while the agreement between heads does not. This attention-recall/readout gap is not specific to MS MARCO: on the out-of-distribution LIMIT benchmark ([4]), we get $R^{\text{any}}{19}{=}1.00$, while the code-generation readout recovers only $0.73$ (Section 6).

$**Figure 3:** Per-layer recall vs. $L$ across $N \in \{500, 1000, 2500, 5000, 10000\}$ on $\textsc{BlockSearch}$, MS MARCO. **(a)** $R^{\text{sum}}_L$ drops monotonically as $N$ grows, reaching $\approx 0$ across all layers at $N{=}10\mathrm{k}$. **(b)** $R^{\text{any}}_L$ remains at $1.00$ across $\mathrm{L}18$ – $\mathrm{L}25$ at every $N$.$

The pre-softmax score is preserved, but normalization breaks the final attention distribution.

The $R^{\text{any}}L$ result might suggest the per-head signal can be recovered by a better aggregator. But what the rest of the model sees is not a per-head ranking; it is the layer's attention output $a_L$, into which gold's value vectors enter in proportion to their post-softmax mass $\alpha^{h}{L, t}$, not the per-head ranking of gold over other documents. So even at L19, where $R^{\text{any}}_{19}{=}1.00$ at $N{=}10\mathrm{k}$, the relevant question is what $a_L$ contains.

Table 1 reports $|a_{19}|$ and GoldShare across the $N$ sweep. The total magnitude $|a_{19}|$ shrinks by only $\sim 36%$ from $N{=}500$ to $N{=}10\mathrm{k}$ — the layer keeps writing into the residual at roughly its original amplitude. GoldShare, however, drops from $0.91$ to $0.01$, a factor of about $130$. The L19 attention output is rewritten from a gold-token average to a non-gold-token average of comparable size; by the time the LM head reads the residual at L21, the slot that carried gold-derived information at small $N$ now carries an aggregate of distractors. Additional per-head softmax statistics behind this swap (gold's pre-softmax score eroding while the non-gold log-sum-exp grows) are reported in Appendix G.

\begin{tabular}{l rrrr}
\toprule
$N$ & $\|a_{19}^{\mathcal{G}}\|$ & $\|a_{19}^{\bar{\mathcal{G}}}\|$ & $\|a_{19}\|$ & $\|a_{19}^{\mathcal{G}}\| / \|a_{19}\|$ \\
\midrule
$500$ & $43.03$ & $17.47$ & $47.48$ & $0.91$ \\
$1\mathrm{k}$ & $30.99$ & $21.11$ & $45.36$ & $0.68$ \\
$2.5\mathrm{k}$ & $\phantom{0}7.64$ & $33.64$ & $43.03$ & $0.18$ \\
$5\mathrm{k}$ & $\phantom{0}2.10$ & $34.61$ & $36.90$ & $0.06$ \\
$10\mathrm{k}$ & $\phantom{0}0.21$ & $29.88$ & $30.27$ & $0.01$ \\
\bottomrule
\end{tabular}

Conclusion: softmax attention dispersion as $N$ grows.

The vector-level result of Table 1 is driven by a single mechanism inside each head. As $N$ grows, gold's largest pre-softmax score $s^{h}{L, t^\star}$ (with $t^\star = \arg\max{t \in \mathcal{G}} s^{h}{L, t}$) no longer dominates over the larger pool of competing distractor scores, so the post-softmax weight assigned to gold's value vector shrinks and the per-head readout $u^{h}{L}$ becomes a non-gold average at comparable magnitude. This length-driven softmax dispersion is a generic phenomenon ([9, 11, 12]) and an empirical driver of length-generalization failure across tasks ([44, 45]); the contribution here is connecting it to a vector-level swap at the residual stream rather than treating it as an abstract probability-mass collapse. Section 5 then evaluates two avenues of improvement that target this mechanism.

5. Addressing the Attention Dilution Problem

Section Summary: The section explores fixes for attention weights becoming diluted across too many tokens as the corpus grows larger. Researchers test length-aware tweaks to the softmax operation itself, such as adding a learned bias term or rescaling scores by a factor involving log N, along with a separate method that first scores whole documents and restricts later attention to only the top-ranked ones. Experiments show the rescaling approach and the document-routing method largely restore strong retrieval performance even at very large corpus sizes, while the additive bias helps little.

Having identified softmax dispersion under growing $N$ as the guilty party, we study two potential directions that address it. We first consider two length-aware modifications to attention softmax that mitigate the dispersion issue. We also consider document-level sparse attention, thereby pruning the set of tokens that enter the softmax.

5.1 Method

oldparagraphLength-aware softmax. We first describe two methods that change how softmax converts the pre-softmax scores $s^{h}{L, t}$ to attention weights. Both are length-aware: each carries a corpus-size signal that lets the softmax adapt to the large $N$ setting. An additive sink appends a learned constant to the softmax denominator, so a head whose largest scores fall below that constant leaks probability mass into the appended slot and emits a small attention output ([26, 27]); the per-token weights $\tilde\alpha^{h}{L, t}$ no longer sum to one. The length-awareness is implicit: training under a wide range of effective corpus sizes $N_{\text{eff}} \sim \mathrm{LogUniform}(128, 5000)$) lets the learned scalar absorb a length signal, so the gate fires more often at large $N$ where attention is diffuse. Multiplicative score rescaling keeps a proper softmax but sharpens it, and is length-aware by construction: SSMax ([9]) multiplies the pre-softmax scores by $s \cdot \log N$ (for a tuned scaling parameter $s$) so that the gap between the largest score and the bulk grows directly with $N$.

Following GPT-OSS ([27]) we instantiate the sink with a learned scalar $b_L \in \mathbb{R}$ per attention layer:

$ \tilde{\alpha}^{h}{L, t} ;=; \frac{\exp(s^{h}{L, t})}{\sum_{t'} \exp(s^{h}_{L, t'}) + \exp(b_L)}.\tag{5} $

The appended slot has no value vector, so $\sum_t \tilde{\alpha}^{h}_{L, t} \le 1$ and a layer with diffuse attention writes a smaller residual update (gate-form derivation in Appendix H). We refer to the resulting model as $\textsc{BlockSearch}$-sink. We then produce $\textsc{BlockSearch}$-SSMax following [9] with $s$ initialized to $0.43$ at every attention layer. Both models compose with the block-sparse attention pattern and train seamlessly with the changes in Section 3.

Document-level sparse attention.

On the flip side, we also consider scoring documents and keeping a top- $B$ shortlist before the code generation is performed; this is a popular approach in long-context research that prunes the context down into a manageable state [16, 14, 17, 23, 24, 25]. We perform a doc-level routing step at $\mathrm{L}16$, immediately upstream of the retrieval band (Section 4.2): a forward pass through $\mathrm{L}0$ – $\mathrm{L}15$ produces the per-document score $R^{\text{sum}}_{16}(d)$ from Section 4.2; only the tokens of the top $B{=}256$ documents participate in dense attention at $\mathrm{L}17$ onward. We call this $\textsc{BlockSearch}$-routing, and we pick doc-level over token-level because each document is a conveniently coherent semantic unit. We refer the reader to Appendix I for more information on the choice of $B$.

5.2 Results

::: {caption="Table 2: Recall@1 (× 100) vs. corpus size N on Natural Questions, MS MARCO, and HotpotQA. SSMax (applied at all 28 attention layers), top- B=256 routing at L16, and their composition are strong for the full N sweep, recovering most of the gap to Qwen3-dense and matching or exceeding MSA-4B at ∼ 1/7 the parameters. NQ caps at N=8,607 to stay under a ∼ 1M-token prefill budget. The top row is the attention ceiling $R^a$ny_19 (Section 4): the fraction of queries (n=400) for which at least one head ranks the gold document first at L19 by pre-softmax QK-MaxSim."}

:::

In the top row of Table 2, the attention ceiling $R^{\text{any}}_{19}$ stays near perfect across all three datasets and the full range of $N$, confirming that the collapse in the rows below reflects a readout failure rather than a loss of the underlying retrieval signal. The proposed modifications aim to recover this latent retrieval signal at the model output.

The additive sink barely moves the large- $N$ collapse.

$\textsc{BlockSearch}$-sink yields only modest mid- $N$ gains and a small lift at MS MARCO $N{=}10\mathrm{k}$ ($0.2 \to 2.5$, from a near-floor base); its absolute level stays low across all three datasets and tracks the no-modification curve at large $N$. A learned constant cannot restore the gold token's post-softmax weight once the softmax denominator has grown by orders of magnitude: it shifts the denominator uniformly and cannot rescale its $N$-dependence.

SSMax holds up across the full $N$ sweep.

SSMax closes most of the gap to Qwen3-dense on every dataset: MS MARCO $33.8$ at $N{=}5\mathrm{k}$ and $16.5$ at $N{=}10\mathrm{k}$ ($82\times$ over no-modification, against $20.2$ for the dense baseline); HotpotQA $56.8$ at $N{=}10\mathrm{k}$ where the additive sink reaches $1.0$. Scaling pre-softmax scores by $\log N$ widens the gap between gold and distractor scores so the post-softmax gold weight stays bounded as $N$ grows: with $\Delta = s^{h}{L, t^\star} - \bar{s}^{h}{L, \text{distractor}}$, $\alpha_{\text{gold}} \approx 1/\bigl(1 + (N-1), N^{-s\Delta}\bigr)$, and the $\log N$ schedule cancels the $(N-1)$ growth in the denominator whenever $s\Delta > 1$.

Top- $B$ routing matches the dense baseline at large $N$.

Routing applied to vanilla $\textsc{BlockSearch}$ reaches $18.8$ at MS MARCO $N{=}10\mathrm{k}$, within $1.4$ points of the dense baseline; on HotpotQA it reaches $78.5$, essentially matching the dense baseline ($79.5$) and exceeding MSA-4B ($75.5$) at roughly $1/7$ the parameter count. We note that while routing addresses the same mechanism as SSMax, it reintroduces a retrieve-then-read (e.g. RAG) decomposition inside the model: the very structure in-context retrieval is intended to remove.

SSMax plus routing.

Stacking the two ($\textsc{BlockSearch}$-SSMax-routing) gives the strongest configuration on MS MARCO at $N{=}10\mathrm{k}$ ($20.5$, edging the dense baseline) and matches $\textsc{BlockSearch}$-routing on Natural Questions and HotpotQA, indicating they may be complementary rather than redundant. A residual gap to Qwen3-dense at $N{=}10\mathrm{k}$ remains on HotpotQA and Natural Questions, and is the subject of future work.

Gap to attention ranking

While the proposed modifications substantially narrow this gap, long-context performance remains well below the near-perfect attention ceiling $R^{\text{any}}_{19}$. This suggests that attention dilution, not attention ranking, remains the primary bottleneck for future work.

6. Out-of-Distribution Generalization: LIMIT

Section Summary: When standard benchmarks favor dense retrieval, a key open question is whether in-context methods like BlockSearch can actually surpass them on tasks requiring lexical rather than semantic matching. On the LIMIT biography retrieval benchmark, which adds distractors to scale the corpus from dozens to thousands of documents, the SSMax-plus-routing variant of BlockSearch maintains substantially higher recall than a matched dense retriever even at roughly 850k tokens, while attention patterns continue to surface the correct documents. These gains demonstrate stronger out-of-distribution behavior, yet performance still declines with scale, indicating that the modifications extend capability without fully eliminating readout and length-related bottlenecks.

Experiments in Section 3–Section 5 focused on widely studied retrieval benchmarks that dense retrieval has been heavily optimized for. $\textsc{BlockSearch}$, with the proposed modifications, matches dense retrieval, which is a notable milestone given that these benchmarks largely define the strengths of dense retrieval. However, a fundamental question remains: can in-context retrieval actually outperform dense retrieval, going beyond matching it? To answer this question, we evaluate $\textsc{BlockSearch}$ on LIMIT ([4]), a benchmark requiring a lexical notion of similarity, where dense retrieval struggles. Under a strict out-of-distribution evaluation setup, $\textsc{BlockSearch}$ significantly outperforms dense retrieval.[^2]

[^2]: For additional out-of-distribution results, we also provide figures for the OBLIQ evaluation ([46]) in Appendix L.

Evaluation setup.

In LIMIT, every query has two gold documents drawn from a corpus of $50{,}000$ short biographies. We scale the corpus from $N{=}46$ (the gold documents only, "LIMIT-small") up to $N{=}5{,}000$ by adding distractor biographies, which grows the prefill from $\sim$ 8k to $\sim$ 850k tokens while each query keeps its two golds. We evaluate the four $\textsc{BlockSearch}$ variants of Section 5 against a same-backbone pooled dense retriever, with $n{=}1{,}000$ queries and an the same version of Recall@1 for the multi-gold regime as HotpotQA (Section 3.3). We do this for consistency; for true Recall@2 results, we refer the reader to Appendix K.

::: {caption="Table 3: LIMIT length generalization: Recall@1 (multi-gold, n=1,000) as the corpus grows from N=46 to N=5,000 by adding distractor biographies (∼ 8k to ∼ 850k tokens). The top row is the attention ceiling $R^a$ny_19 (Section 4): the fraction of queries for which at least one head ranks a gold document first at L19, by pre-softmax QK-MaxSim under the corrected (live-decode-faithful) measurement."}

:::

Attention vs. readout.

The attention ceiling $R^{\text{any}}_{19}$ stays at $1.00$ across the entire sweep, from $N{=}46$ up to $N{=}5{,}000$ ($\sim$ 850k tokens): at every corpus size at least one head still ranks a gold document first. The base $\textsc{BlockSearch}$ readout, by contrast, collapses far faster: from $0.354$ to $0.000$ by $N{=}2{,}500$. The gap between the two is the same readout bottleneck identified in Section 4, now both out of distribution (a lexical task that $\textsc{BlockSearch}$ never trained on; its training mix is semantic, Appendix C) and length-generalized: the per-head retrieval signal persists while the code-generation readout cannot recover it.

Routing.

Of our modifications, SSMax+routing performs the best, with R@1 at $0.149$ at $N{=}5{,}000$ while $\textsc{BlockSearch}$ and $\textsc{BlockSearch}$-sink are at zero. SSMax also helps at mid- $N$, but reaches $0.054$ R@1 by $N{=}5{,}000$. The length-aware sink does not help on this lexical benchmark: $\textsc{BlockSearch}$-sink trails $\textsc{BlockSearch}$ at every $N$ (e.g. $0.252$ vs. $0.354$ at $N{=}46$), in contrast to its roughly neutral effect on MS MARCO. $\textsc{BlockSearch}$-SSMax-routing also stays above the same-backbone pooled Qwen3-dense retriever at every $N$ ($0.149$ vs. $0.035$ at $N{=}5{,}000$), suggesting that our methods do offer better out-of-distribution performance and potential extension to more complex tasks beyond basic semantic search.

Caveats.

Routing delays but does not prevent the inevitable: $\textsc{BlockSearch}$-SSMax-routing is still declining with $N$ ($0.149$ at $\sim$ 850k tokens). Thus, these modifications extend functional LIMIT retrieval to larger corpora, rather than solving it, which remains an open problem.

7. Conclusion

Section Summary: This research conducted the first large-scale test of whether language models can pull relevant information from collections of text as large as a million words, including sizes far beyond what they saw in training. With targeted training and simple changes to how the models focus, performance matched or exceeded standard retrieval methods on tough benchmarks, though attention became diluted and less effective as the text grew extremely long. Length-aware tweaks to the attention process eased this problem and point to both the promise and the core limits of using language models directly for large-scale retrieval.

In this work, we presented the first systematic study of corpus-scale in-context retrieval with long-context language models, focusing on two requirements expected of practical retrieval systems: scaling to million-token corpora and generalizing to corpus sizes far beyond those seen during training. We showed that, with appropriate training and architectural modifications, LMs can perform meaningful retrieval far beyond their nominal training regime, approaching dense retrieval performance at moderate scales. The story is a tale of two cities: on benchmarks built around single-vector retrieval, our methods match a strong dense baseline at million-token scale, while on LIMIT, designed to defeat embedding similarity, they exceed it by nearly $3\times$. At the same time, we identified attention dispersion as a fundamental scaling bottleneck: even when retrieval signals remain internally present, normalized attention mass collapses under extreme context growth. Motivated by this analysis, we demonstrated that length-aware tweaks to attention substantially improve retrieval performance at million-token scale. Together, our results suggest that scalable in-context retrieval is both promising and fundamentally constrained by attention dilution, motivating new directions for retrieval architectures and long-context modeling.

Acknowledgment

This research was supported by the NVIDIA Academic Grant Program. Siddharth Gollapudi is supported by the NSF (CSGrad4US award no. 2313998).

Appendix

Section Summary: The appendix details the technical setup for BlockSearch, including how documents and queries are formatted into prompts with block markers and attention rules so the model can identify relevant passages. It also describes the evaluation datasets drawn from sources like MS MARCO, the filtering process applied to RLHN training data to create balanced query-document sets, and the shared hyperparameters used to fine-tune the underlying language model. Additional notes cover modifications such as length-aware attention adjustments during training.

A. Prompt Format

To make the input format used by $\textsc{BlockSearch}$ concrete, Figure 4 shows a worked example for a single MS MARCO query against a four-document corpus. Each document is wrapped between BOS/EOS markers and bracketed by its randomly-assigned four-digit code (Section 3.2); RoPE [30] positions are reset to $0$ at every BOS so that the document tokens occupy positions $0$ – $T_{\text{doc}}$ regardless of where the document sits in the prefill. Document tokens attend causally only inside their own block.

After the corpus, the query block is appended at RoPE position $T_{\text{doc}}{=}300$ and attends to the entire prefilled corpus. The query block carries an instruction prefix, the query, and the four answer slots; the model is supervised to emit the four-digit code identifying the gold document and is decoded autoregressively at inference. The same template is used at training and at evaluation, with one difference: at training time the four answer digits are teacher-forced and the on-policy auxiliary loss (Appendix F) replays the model's own rollout at the same four positions.

$**Figure 4:** Worked example of the $\textsc{BlockSearch}$ prompt. Documents are prefilled with block-sparse attention (Section 3.2); the query is then appended and the model autoregressively decodes a four-digit code. The underlined digits are the supervision target at training time and the autoregressive output at inference.$

B. Evaluation Suite

\begin{tabular}{lrrrrr}
\toprule
\textbf{Dataset} & \textbf{Split} & \textbf{Queries} & $N_{\max}$ & \textbf{Avg. tokens / doc} & Tokens @ $N_{\max}$ \\
\midrule
MS\, MARCO & dev & 400 & 10{,}000 & 94.8 & 948{,}421 \\
HotpotQA & test & 400 & 10{,}000 & 116.3 & 1{,}163{,}144 \\
NQ & test & 400 & 8{,}600 & 139.6 & 1{,}201{,}153 \\
\bottomrule
\end{tabular}

Because NQ has far more tokens per document than MS MARCO or HotpotQA, we slightly truncate the dataset size, and reduce the number of hard negatives per query from 24 to either 20 or 21. This is the result of a round-robin algorithm that prunes out documents evenly from the dataset until we hit 8600 documents, or around 1.2M tokens.

C. RLHN Training-Data Composition

We train all $\textsc{BlockSearch}$ variants on a filtered version of the public RLHN-680K release ([32]).^3 Each query in the source ships with one or more positive passages and a pool of mined negatives. We score every (query, candidate) pair with Qwen3-Embedding-8B ([34]) (last-token pooling, $\ell_2$-normalized, max sequence length $512$), then apply two filters and a fixed-shape trim:

drop queries with fewer than $16$ candidates or fewer than $15$ negatives;
drop queries where the highest-scoring candidate is a negative, i.e. $\max \mathrm{score}(\text{pos}) \le \max \mathrm{score}(\text{neg})$;
retain the single best-scoring positive plus the top- $15$ highest-scoring negatives, giving exactly $16$ documents per query.

The resulting post-filter, post-trim dataset has $522{,}487$ training samples drawn from $7$ source corpora (Table 5); each sample is one query paired with its $16$ documents. Optimization steps are stratified per source: each batch contains queries from a single subset, so the gradient at any step reflects a single retrieval distribution. Documents are truncated to $T_{\text{doc}}{=}300$ tokens at training time. Token counts in Table 5 use the Qwen3-0.6B tokenizer; document-token averages are estimated from a stratified random sample of $8{,}000$ documents per subset (standard error

lt;0.5$ tokens at this $n$).

\begin{tabular}{lrrr}
\toprule
Source & # training samples & Avg. query tok & Avg. doc tok \\
\midrule
MS\, MARCO & 368, 961 & 7.0 & 84.6 \\
HotpotQA & 81, 551 & 24.2 & 100.6 \\
FEVER & 28, 561 & 11.7 & 265.1 \\
NQ & 27, 962 & 10.5 & 146.5 \\
SCIDOCS-RR & 11, 787 & 13.4 & 221.9 \\
FiQA & 2, 822 & 13.9 & 225.7 \\
ArguAna & 843 & 251.9 & 209.9 \\
\midrule
Total & 522, 487 & -- & -- \\
\bottomrule
\end{tabular}

D. Training Hyperparameters

All $\textsc{BlockSearch}$ variants in this paper share a single training recipe, fine-tuning Qwen/Qwen3-0.6B on the RLHN-filtered mix of Appendix C with the on-policy auxiliary loss of Appendix F. The shared hyperparameters are listed in Table 7; the modification variants of Section 5 ($\textsc{BlockSearch}$-sink, $\textsc{BlockSearch}$-SSMax) layer length-aware attention modifications on top of this recipe and otherwise inherit every other knob. The variant-specific differences are listed in Table 6.

Note on the sink.

The per-layer sink scalars $b_L$ sit in their own parameter group with LR $1{\times}10^{-3}$ and zero weight decay (vs. $3{\times}10^{-5}$ / $0.01$ for the base weights): the higher LR compensates for the small gradients flowing through $\sigma(\mathrm{lse} - b_L)$ into a single scalar, and zero weight decay avoids pulling the gate toward an uninformative regime. The length signal is injected by sampling $N_{\text{eff}} \sim \log\mathcal{U}(N_0, 5\mathrm{k})$ per step ($N_0 = 128$ for our 8-GPU setup) and using $b_L + \alpha \log(N_{\text{eff}}/N_0)$ as the effective threshold, with the gate strength linearly ramped in over the first $\sim$ 2k steps. At evaluation the gate is disabled entirely: the learned $b_L$ stay in the checkpoint but do not affect attention. The mechanism only shapes what the rest of the model adapts to during training.

Note on SSMax.

The per-layer scalars $s_L$ are initialized to $0.43$ following [9] and trained in the shared parameter group — no separate LR or weight decay. Length conditioning is explicit through $\log N$, so the same scaling applies unchanged at evaluation.

::: {caption="Table 6: Variant-specific hyperparameters: differences between BlockSearch, BlockSearch-sink, and BlockSearch-SSMax. All other settings are shared (Table 7)."}

:::

::: {caption="Table 7: Shared training hyperparameters for all BlockSearch variants."}

:::

E. Training details for `Qwen3-dense`

We fine-tune Qwen3-0.6B (FlashAttention-2, gradient checkpointing) into a last-token-pooled, $\ell_2$-normalized embedding model on the same RLHN-filtered mix used for $\textsc{BlockSearch}$ training (Appendix C). Each training example pairs one query with $C{=}16$ teacher-scored candidates; queries are formatted with the same instruction prefix used at evaluation ("Given a web search query, retrieve relevant passages that answer the query"), and only the last hidden state of each sequence is pooled. The loss is a contrastive term plus a teacher-distillation term:

$ \mathcal{L} ;=; \mathcal{L}_{\text{CE}}

\lambda_{\text{KL}}, \mathcal{L}{\text{KL}}, \quad \mathcal{L}{\text{CE}} = \mathrm{CE}!\left(\tfrac{1}{\tau_s}, q, C_{\text{all}}^{!\top}, ; y\right), \quad \mathcal{L}_{\text{KL}} = \mathrm{KL}!\bigl(\sigma(s_t/\tau_t), |, \sigma(s_s/\tau_s)\bigr),\tag{6} $

where $C_{\text{all}}$ is the candidate matrix gathered (with gradient) across all DDP ranks, so each query is contrasted against $W \cdot B \cdot C$ documents in total; $y$ points at the gold candidate within the gathered block; and the KL term distills the teacher's softmax over the local 16 candidates into the student's softmax over the same 16. Hyperparameters are listed in Table 8; we train a single epoch over the RLHN-filtered triples on $8{\times}$ NVIDIA A100, with cosine decay to $0.1, \eta_0$ after a 500-step linear warmup, AdamW ($\beta$ defaults), and gradient clipping at $1.0$. The resulting checkpoint is the Qwen3-dense-0.6B dense baseline used throughout the paper.

::: {caption="Table 8: Qwen3-dense training configuration."}

:::

F. On-Policy Auxiliary Loss

Section 3.2 introduces an on-policy auxiliary loss to mitigate exposure bias when decoding the four-digit code. The procedure is given in Algorithm 1. For each query in the batch, we have $K$ candidate documents (1 positive, 15 hard negatives) with their full four-digit codes and Qwen3-Embedding-8B relevance scores, and the corpus has already been prefilled into the shared block-sparse cache (Section 3.2).

**Require:** model π(θ); batch of queries with prefilled document cache; per-query candidate codes c¹, …, c(K) (each a length- L digit string) with teacher relevance scores r¹, …, r(K).
Compute teacher weights w(k) ∝ exp(r(k) / τ(tch)) over the K candidates.
**(1) Rollout the model on its own.** Sample a code hata = hata₁ … hata(L) autoregressively from π(θ) with gradients off.
**(2) Build a per-position teacher.** For each digit position t:
Keep candidates whose first t-1 digits match the rollout: M(t) = (k : c(k)(<t) = hata(<t)).
Set the teacher distribution over the 10 digits to the candidate-weighted vote:
q(t)(d) ∝ ∑ over k ∈ M(t) w(k) · 1[c(k)(t) = d].
**(3) Replay with gradients.** Re-decode the four answer positions conditioned on hata(<t) (the rollout's prefix, not the gold prefix) with gradients on, obtaining the student distributions π(θ)(· | hata(<t)).
**Loss:** L(aux) = (1)/(L) ∑ over t=1..L CE(q(t), π(θ)(· | hata(<t))), added to the total loss as L = L(CE) + λ L(aux).

The rollout uses the model's own distribution rather than the gold code, so the prefixes that the four answer positions condition on are exactly the prefixes the model would visit at inference. The per-position teacher $q_t$ thus plays the role of a DAgger expert ([47]): at every state visited along the model's own trajectory, it specifies what the model should have done. We ramp $\lambda$ in linearly from $0$ over the first warmup steps so early training is dominated by the standard teacher-forced cross-entropy.

G. Per-head softmax statistics at $\mathrm{L}19$

Section 4 reports the vector-level outcome of the L19 failure ($|a_{19}|$ preserved while the gold-driven fraction collapses to $0.01$). Here we characterize the per-head softmax statistics that drive that vector-level swap, to show that the failure is a compound one: gold's pre-softmax score erodes and the competitor mass widens against it.

Gold's per-head softmax mass admits the closed form $\mathrm{gold_post}^{h}{L} = \sigma!\bigl(\mathrm{lse}{\mathcal{G}} - \mathrm{lse}{\bar{\mathcal{G}}}\bigr)$, a sigmoid of the gap between the gold-side and non-gold-side log-sum-exps. We track gold's largest attention logit $s^{\max}{\mathcal{G}}$ (which dominates $\mathrm{lse}{\mathcal{G}}$) and the noise gap $\mathrm{lse}{\bar{\mathcal{G}}} - \mathrm{lse}{\mathcal{G}}$ separately at L19 across the same $N$ sweep used in the main text (Table 9). Both move the wrong way: $s^{\max}{\mathcal{G}}$ drops by $\sim 3$ logit units (gold's query alignment itself erodes) and the noise gap widens by $\sim 3.5$ (competitors out-compete gold for log-mass). Together the per-head gold mass at L19 (median over heads, denoted $\mathrm{gold_post}_{19}$) collapses by $\sim 150\times$, against only the $\sim 20\times$ that pure $O(N)$ denominator dilution would predict ([11, 12]). The L19 collapse is therefore a compound failure of two independent effects, and either one alone would not produce the observed magnitude.

\begin{tabular}{l rrr}
\toprule
$N$ & $s^{\max}_{\mathcal{G}}$ & $\mathrm{lse}_{\bar{\mathcal{G}}} - \mathrm{lse}_{\mathcal{G}}$ & $\mathrm{gold\_post}_{19}$ \\
\midrule
$500$ & $14.60$ & $+4.63$ & $0.0320$ \\
$1\mathrm{k}$ & $14.27$ & $+5.44$ & $0.0152$ \\
$2.5\mathrm{k}$ & $13.81$ & $+6.60$ & $0.0039$ \\
$5\mathrm{k}$ & $12.95$ & $+7.31$ & $0.0010$ \\
$10\mathrm{k}$ & $11.53$ & $+8.06$ & $0.0002$ \\
\bottomrule
\end{tabular}

H. Equivalence of the additive-sink and sigmoid-gate forms of attention

We prove that the additive-sink softmax of Equation 5 is exactly equivalent to multiplying the standard softmax of Equation 1 by a sigmoid gate.

########## {caption="Proposition 1"}

Fix a layer $L$ and head $h$. Let $s = (s^{h}{L, 1}, \dots, s^{h}{L, T}) \in \mathbb{R}^T$ be the pre-softmax logits at the $d_1$-emission position, $b_L \in \mathbb{R}$ the learned sink scalar, $\alpha^{h}{L, t}$ the standard softmax weight from Equation 1, and $\tilde{\alpha}^{h}{L, t}$ the additive-sink weight from Equation 5. Define the layer-and-head-wise gate

$ g^{h}{L} ;=; \sigma!\bigl(\mathrm{lse}(s) - b_L\bigr), \qquad \mathrm{lse}(s) ;=; \log!\sum{t=1}^{T} \exp(s^{h}_{L, t}), \qquad \sigma(x) = \frac{1}{1 + e^{-x}}. $

Then for every $t \in {1, \dots, T}$,

$ \tilde{\alpha}^{h}{L, t} ;=; \alpha^{h}{L, t} \cdot g^{h}_{L}. $

Proof: Let $Z = \sum_{t'=1}^{T} \exp(s^{h}{L, t'}) = \exp(\mathrm{lse}(s))$ denote the standard softmax normalizer. By Equation 1, $\alpha^{h}{L, t} = \exp(s^{h}_{L, t}) / Z$. By Equation 5, the additive-sink weight is

$ \tilde{\alpha}^{h}{L, t} ;=; \frac{\exp(s^{h}{L, t})}{Z + \exp(b_L)}. $

Multiply numerator and denominator by $1/Z$:

$ \tilde{\alpha}^{h}{L, t} ;=; \frac{\exp(s^{h}{L, t})/Z}{1 + \exp(b_L)/Z} ;=; \alpha^{h}{L, t} \cdot \frac{1}{1 + \exp(b_L - \mathrm{lse}(s))} ;=; \alpha^{h}{L, t} \cdot \sigma!\bigl(\mathrm{lse}(s) - b_L\bigr) ;=; \alpha^{h}{L, t} \cdot g^{h}{L}, $

where the third equality uses $\sigma(x) = 1/(1 + e^{-x})$ with $x = \mathrm{lse}(s) - b_L$. $\blacksquare$

Interpretation.

The gate $g^{h}{L} \in (0, 1)$ depends only on the per-(layer, head) comparison between $\mathrm{lse}(s)$ — a measure of how much pre-softmax mass the layer's logits already concentrate — and the learned threshold $b_L$. When the layer's logits are sharp and concentrated relative to $b_L$ (large $\mathrm{lse}$), $g^{h}{L} \to 1$ and 5 reduces to the standard softmax. When the logits are diffuse and many tokens contribute small amounts (small $\mathrm{lse}$ relative to $b_L$), $g^{h}_{L} \to 0$ and the layer's contribution to the residual update $a_L$ is suppressed multiplicatively. This is the mechanism by which the sink "costs" a noisy layer its write to the residual stream without altering the relative ranking of the per-token weights.

I. Top- $B$ Routing Recall

The top- $B$ routing of Section 5 hinges on the assumption that the router can keep the gold document inside a small top- $B$ shortlist with high probability across our checkpoints. We measure exactly that: for $B \in {32, 64, 128, 256, 512, 1024}$ we report the fraction of queries whose gold document survives the QK-MaxSim router at layer $16$, the default routing configuration used in Section 5. Evaluation uses $N{=}10{,}000$ MS MARCO documents and $n{=}400$ queries (the same query slice for all three checkpoints). Recall@ $B$ is derived per query from the gold document's rank under the router's score relative to the full $N{=}10{,}000$ ranking.

Across all three checkpoints the gold document is retained for at least $93.5%$ of queries at the default $B{=}256$ ($\textsc{BlockSearch}$: $0.962$; $\textsc{BlockSearch}$-sink: $0.962$; $\textsc{BlockSearch}$-SSMax: $0.935$), and the curves flatten between $B{=}256$ and $B{=}512$ — going from $B{=}256$ to $B{=}1024$ recovers an additional $\le 5$ percentage points on every checkpoint (Figure 5). This is consistent with $B{=}256$ being a practical routing budget on this corpus: the remaining miss rate is dominated by the long tail of routing errors (gold ranked beyond the top- $1024$), which is unaffected by any reasonable enlargement of $B$. $\textsc{BlockSearch}$-SSMax's lower curve at small $B$ ($0.762$ at $B{=}32$ vs. $0.838$ for $\textsc{BlockSearch}$) is consistent with SSMax placing slightly less probability mass on the gold document at the routing layer, but the gap closes by $B{=}512$.

**Figure 5:** Routing recall@ $B$ for QK-MaxSim routing at layer $16$ on MS MARCO with $N{=}10{,}000$ and $n{=}400$ queries. Dashed line marks the default $B{=}256$ used throughout the paper.

J. Beam Search and Recall@5

The main text reports Recall@1 over the single highest-probability four-digit code (Section 3.3). Here we (i) specify the decoding procedure used to turn the model's per-digit distributions into a ranked list of document codes, and (ii) report the corresponding Recall@5 across the full length-generalization sweep, so that the main-text Recall@1 conclusions can be checked at a wider cutoff.

J.1 Decoding

A query is answered by a digit-by-digit beam search over the full $N$-document corpus (context-parallel across GPUs), returning a ranked list of complete codes. Recall@ $k$ is computed over this ranking.

J.2 Digit-by-digit beam search

Constrained beam search over a tree of identifier tokens is the standard decoder for generative and in-context retrieval, and we follow that practice rather than introduce a new procedure ([20, 21, 5]). Because each identifier here is exactly four digits drawn from a $10$-symbol vocabulary, the search tree has depth $4$ and branching factor $10$, and the beam is over digit strings rather than free-form text. Algorithm 2 shows the routine; our default configuration is beam width $B{=}25$, per-step pruning $k_{\text{step}}{=}5$, and return depth $k_{\text{ret}}{=}5$ (so Recall@5 reads the top five returned codes).

**Require:** model π(θ) with the document corpus prefilled into its (block-sparse) KV cache; digit token ids D=(id(0),…,id(9)); code length L=4; beam width B; per-step prune k(step); return depth k(ret).
Append the query and forward once; let ell₁ = logsoftmax of the next-token logits restricted to D.
Initialize B beams from the top- min(B,10) first digits, with beam scores =ell₁ of the chosen digits; broadcast the KV cache across the B beams.
**for** t = 2,…,L **do**
  Forward all live beams one step; let ell(t)(b) be the digit log-probs for beam b.
  Candidate scores: score(b,d) = beam(score)(b) + ell(t)(b)(d).
  Per beam, keep only its top- k(step) digits (set the rest to -∞).
  Flatten over (b,d) and keep the global top- B survivors; record each survivor's parent beam and digit.
  Reorder the KV cache by parent index; append the chosen digit to each surviving beam.
**end for**
Sort the B completed beams by score;
**return** the top- k(ret) digit strings.

J.3 Recall@5 across the length-generalization sweep

Table 10 reports Recall@5 for the same variants, datasets, and corpus sizes as the main-text Recall@1 (Figure 1 and Table 2). The qualitative picture is identical to Recall@1: the position-coded variant and the no-modification model collapse by $N{=}5$ k– $10$ k on every dataset, while SSMax, top- $B$ routing, and their composition hold up across the full sweep. Recall@5 stays low wherever Recall@1 collapses and rises with it wherever the modifications hold (e.g. MS MARCO $\textsc{BlockSearch}$ at $N{=}10$ k: R@1 $0.2$, R@5 $0.2$; $\textsc{BlockSearch}$-routing $18.8 \to 47.5$), confirming that the main-text collapse is not an artifact of the top-1 cutoff. For HotpotQA, whose queries have two golds, we also report Recall@2 in Table 11 for completeness.

::: {caption="Table 10: Recall@5 (× 100) vs. corpus size N, mirroring Table 2 (n=400). NQ caps at N=8,607. HotpotQA has two golds per query, so its columns report true recall (the fraction of the two golds that land in the top 5); Natural Questions and MS MARCO are effectively single-gold, where this matches the hit-based metric used in the main body."}

:::

\begin{tabular}{l rrrrr}
\toprule
Method & 0.5k & 1k & 2.5k & 5k & 10k \\
\midrule
\textsc{BlockSearch} & 63.4 & 61.2 & 39.8 & \phantom{0}8.9 & \phantom{0}0.0 \\
\textsc{BlockSearch}-position & 61.2 & 58.0 & 29.6 & \phantom{0}1.2 & \phantom{0}0.0 \\
\textsc{BlockSearch}-offpolicy & 61.2 & 58.5 & 36.5 & \phantom{0}3.8 & \phantom{0}0.0 \\
\midrule
\textsc{BlockSearch}-sink & 65.0 & 61.1 & 38.5 & \phantom{0}9.5 & \phantom{0}0.0 \\
\textsc{BlockSearch}-SSMax & 60.9 & 60.9 & 51.6 & 43.2 & 35.1 \\
\textsc{BlockSearch}-routing & 68.1 & 67.2 & 58.6 & 52.1 & 46.0 \\
\textsc{BlockSearch}-SSMax-routing & 64.6 & 64.0 & 57.4 & 50.2 & 44.6 \\
\bottomrule
\end{tabular}

K. LIMIT worked example

Figure 6 shows a worked LIMIT-small example for the lexical retrieval task evaluated in Section 6.

**Figure 6:** Worked LIMIT-small example. Relevance is lexical: the gold document is a person whose preference list contains the queried item (`Joshua Trees`), among the $46$ documents. Each query has two gold documents (we show one); either code is a correct answer. The query and gold text are from the evaluated corpus.

K.1 Recall@2 and Recall@5

The main-text Table 3 reports Recall@1 under the metric used throughout the body (a query counts if either gold is in the top 1). Here we additionally report standard Recall@ $k$: the fraction of a query's two golds that appear in the top $k$, at $k{=}2$ (Table 12) and $k{=}5$ (Table 13). The picture is the same: SSMax+routing holds up across the sweep while $\textsc{BlockSearch}$ and $\textsc{BlockSearch}$-sink collapse by $N{=}2.5$ k, and all baselines stay far below the $R^{\text{any}}_{19}$ attention ceiling, which surfaces both golds within the top $2$ at every $N$.

::: {caption="Table 12: LIMIT Recall@2 (fraction of a query's two golds in the top 2, n=1,000) vs. corpus size N, for the same runs as Table 3. Top row is the $R^a$ny_19 attention ceiling; the dense baseline is dense_v1."}

:::

::: {caption="Table 13: LIMIT Recall@5 (fraction of a query's two golds in the top 5, n=1,000) vs. corpus size N, for the same runs as Table 3."}

:::

L. OBLIQ results

We additionally evaluate on OBLIQ ([46]), a benchmark of "oblique" retrieval, where relevance is indirect rather than lexical or topical. We use three of its tasks. Math and Writing are analogues tasks: the query is a passage (a competition mathematics problem, or a paragraph of prose), and the gold documents are other passages that share an abstract reasoning pattern or argument with it despite entirely different surface content. Twitter is a descriptive task: the query is a natural-language description of a class of posts, and the gold documents are tweets matching that description (see Figure 7 for one example per task). Corpus size $N$ is set per task by feasibility: Math uses the full corpus, while Twitter and Writing use gold-preserving subsamples (drawn from the full $72$ k- and $10{,}388$-document corpora); per-dataset document counts, query counts, and average lengths are given in Table 14. We apply the OBLIQ self-match exclusion (each query's own source passage is dropped from its ranking). As in Section 6 we evaluate the four $\textsc{BlockSearch}$ variants and the same-backbone pooled dense retriever, and additionally a long-context generative-retrieval baseline, MSA-4B ([16]); for ICR the ranked list is the beam search output. Unlike the main body, here we report standard Recall@ $k$: the fraction of a query's gold documents that appear in the top $k$. Recall is inversely related to the number of golds per query — a top- $k$ list can hold at most $k$ of them — so with OBLIQ's many golds (a mean of ${\sim}9$ – $13$ per query) the values are correspondingly small. For the $R^{\text{any}}_{19}$ ceiling, a gold counts as retrieved at $k$ if some head ranks it within its top $k$.

\begin{tabular}{l r r r r}
\toprule
  {} & & & \multicolumn{2}{c}{Avg. tokens} \\
\cmidrule(lr){4-5}
Dataset & Documents & Queries & Doc & Query \\
\midrule
Math (analogue) & $3{,}507$ & $151$ & $116$ & $181$ \\
Twitter (descriptive) & $10{,}000$ & $281$ & $52$ & $53$ \\
Writing (analogue) & $4{,}000$ & $512$ & $498$ & $653$ \\
\bottomrule
\end{tabular}

**Figure 7:** Worked OBLIQ examples, one per task. Math and Writing are *analogues* tasks (the gold shares an abstract pattern with the query despite different surface content); Twitter is a *descriptive* task. Query and gold text are from the evaluated corpora, abbreviated; the Twitter handle is anonymized.

::: {caption="Table 15: OBLIQ results: Recall@1 / Recall@5 (the fraction of a query's gold documents that appear in the top k; with many golds per query these values are small), with the OBLIQ self-match exclusion applied. The top row is the attention ceiling $R^a$ny_19 (Section 4): the fraction of a query's golds that some head ranks within top- k at L19, by pre-softmax QK-MaxSim. MSA-4B emits ∼1 citation per query, so its R@5 is close to its R@1. On Twitter, $R^a$ny_19 surfaces nearly every gold within top-5 (R@5 ≈1) yet almost none at rank 1, because a fixed set of sink documents occupies the top positions for the relevant head."}

:::

References

Section Summary: This section compiles a bibliography of over three dozen research papers and preprints, primarily from conferences such as EMNLP, NeurIPS, and ACL, along with arXiv uploads dated between 2016 and 2026. The works focus on information retrieval techniques, retrieval-augmented generation, long-context language models, attention mechanisms, and related benchmarks or optimizations in AI systems. Many entries explore theoretical limits, efficiency improvements, and empirical evaluations of transformer-based architectures.

[1] Karpukhin et al. (2020). Dense passage retrieval for open-domain question answering. In Proceedings of the 2020 conference on empirical methods in natural language processing (EMNLP). pp. 6769–6781.

[2] Lee et al. (2024). Can long-context language models subsume retrieval, rag, sql, and more?. arXiv preprint arXiv:2406.13121.

[3] Lewis et al. (2020). Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in neural information processing systems. 33. pp. 9459–9474.

[4] Weller et al. (2025). On the theoretical limitations of embedding-based retrieval. arXiv preprint arXiv:2508.21038.

[5] Gupta et al. (2025). Scalable In-context Ranking with Generative Models. arXiv preprint arXiv:2510.05396.

[6] Qiu et al. (2025). Eliciting in-context retrieval and reasoning for long-context large language models. In Findings of the Association for Computational Linguistics: ACL 2025. pp. 3176–3192.

[7] Bai et al. (2025). Longbench v2: Towards deeper understanding and reasoning on realistic long-context multitasks. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). pp. 3639–3664.

[8] Hong et al. (2025). Context Rot: How Increasing Input Tokens Impacts LLM Performance. https://research.trychroma.com/context-rot.

[9] Nakanishi, Ken M (2025). Scalable-softmax is superior for attention. arXiv preprint arXiv:2501.19399.

[10] Leng et al. (2025). Understanding and Improving Length Generalization in Hierarchical Sparse Attention Models. arXiv preprint arXiv:2510.17196.

[11] Veličković et al. (2024). Softmax is not enough (for sharp size generalisation). arXiv preprint arXiv:2410.01104.

[12] Barbero et al. (2024). Transformers need glasses! information over-squashing in language tasks. Advances in Neural Information Processing Systems. 37. pp. 98111–98142.

[13] Duvvuri et al. (2026). LUCID: Attention with Preconditioned Representations. arXiv preprint arXiv:2602.10410.

[14] Lu et al. (2025). Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502.13189.

[15] Xiao et al. (2025). Optimizing mixture of block attention. arXiv preprint arXiv:2511.11571.

[16] Chen et al. (2026). MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens. arXiv preprint arXiv:2603.23516.

[17] Ohayon et al. (2025). Block Sparse Flash Attention. arXiv preprint arXiv:2512.07011.

[18] Thakur et al. (2021). Beir: A heterogenous benchmark for zero-shot evaluation of information retrieval models. arXiv preprint arXiv:2104.08663.

[19] Hsieh et al. (2024). RULER: What's the real context size of your long-context language models?. arXiv preprint arXiv:2404.06654.

[20] Tay et al. (2022). Transformer memory as a differentiable search index. Advances in neural information processing systems. 35. pp. 21831–21843.

[21] Wang et al. (2022). A neural corpus indexer for document retrieval. Advances in Neural Information Processing Systems. 35. pp. 25600–25614.

[22] Rajput et al. (2023). Recommender systems with generative retrieval. Advances in Neural Information Processing Systems. 36. pp. 10299–10315.

[23] Tang et al. (2024). Quest: Query-aware sparsity for efficient long-context llm inference. arXiv preprint arXiv:2406.10774.

[24] Wu et al. (2025). Tokenselect: Efficient long-context inference and length extrapolation for llms via dynamic token-level kv cache selection. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing. pp. 21275–21292.

[25] Hu et al. (2024). Efficient length-generalizable attention via causal retrieval for long-context language modeling. arXiv preprint arXiv:2410.01651.

[26] Xiao et al. (2023). Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453.

[27] Agarwal et al. (2025). gpt-oss-120b & gpt-oss-20b model card. arXiv preprint arXiv:2508.10925.

[28] Jacob et al. (2024). Drowning in documents: consequences of scaling reranker inference. arXiv preprint arXiv:2411.11767.

[29] Ma et al. (2024). Block-attention for efficient prefilling. arXiv preprint arXiv:2409.15355.

[30] Su et al. (2024). Roformer: Enhanced transformer with rotary position embedding. Neurocomputing. 568. pp. 127063.

[31] Dong et al. (2024). Flex attention: A programming model for generating optimized attention kernels. arXiv preprint arXiv:2412.05496. 2(3). pp. 4.

[32] Thakur et al. (2025). Hard Negatives, Hard Lessons: Revisiting Training Data Quality for Robust Information Retrieval with LLMs. arXiv preprint arXiv:2505.16967.

[33] Bajaj et al. (2016). MS MARCO: A human generated machine reading comprehension dataset. arXiv preprint arXiv:1611.09268.

[34] Yang et al. (2025). Qwen3 technical report. arXiv preprint arXiv:2505.09388.

[35] Qu et al. (2021). RocketQA: An optimized training approach to dense passage retrieval for open-domain question answering. In Proceedings of the 2021 conference of the North American chapter of the association for computational linguistics: human language technologies. pp. 5835–5847.

[36] Zhang et al. (2025). Qwen3 embedding: Advancing text embedding and reranking through foundation models. arXiv preprint arXiv:2506.05176.

[37] Geva et al. (2021). Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. pp. 5484–5495.

[38] Geva et al. (2022). Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 conference on empirical methods in natural language processing. pp. 30–45.

[39] Olsson et al. (2022). In-context learning and induction heads. arXiv preprint arXiv:2209.11895.

[40] Wu et al. (2024). Retrieval head mechanistically explains long-context factuality. arXiv preprint arXiv:2404.15574.

[41] Zhang et al. (2025). Query-focused retrieval heads improve long-context reasoning and re-ranking. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing. pp. 23802–23816.

[42] Khattab, Omar and Zaharia, Matei (2020). Colbert: Efficient and effective passage search via contextualized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval. pp. 39–48.

[43] nostalgebraist (2020). interpreting GPT: the logit lens. LessWrong. https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens.

[44] Li et al. (2025). On vanishing variance in transformer length generalization. arXiv preprint arXiv:2504.02827.

[45] Vasylenko et al. (2025). Long-context generalization with sparse attention. arXiv preprint arXiv:2506.16640.

[46] Tchuindjo et al. (2026). OBLIQ-Bench: Exposing Overlooked Bottlenecks in Modern Retrievers with Latent and Implicit Queries. arXiv preprint arXiv:2605.06235.

[47] Ross et al. (2011). A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics. pp. 627–635.

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

Abstract

1. Introduction

2. Related Work

3. BlockSearch: Model and Million-Token Evaluation

3.1 Setup

3.2 Method: BlockSearch

3.3 Evaluation

3.4 Results: BlockSearch is competitive, but collapses as $N$ scales

4. Understanding the Large- $N$ Deterioration

4.1 Setup

4.2 Layer Roles and the Recall–Generation Gap

5. Addressing the Attention Dilution Problem

5.1 Method

5.2 Results

6. Out-of-Distribution Generalization: LIMIT

7. Conclusion

Acknowledgment

Appendix

A. Prompt Format

B. Evaluation Suite

C. RLHN Training-Data Composition

D. Training Hyperparameters

E. Training details for Qwen3-dense

F. On-Policy Auxiliary Loss

G. Per-head softmax statistics at $\mathrm{L}19$

H. Equivalence of the additive-sink and sigmoid-gate forms of attention

I. Top- $B$ Routing Recall

J. Beam Search and Recall@5

J.1 Decoding

J.2 Digit-by-digit beam search

J.3 Recall@5 across the length-generalization sweep

K. LIMIT worked example

K.1 Recall@2 and Recall@5

L. OBLIQ results

References

E. Training details for `Qwen3-dense`