UNREAL: Unifying Retrieval and Long-Context with a Single Model

AI summary

Overview

Research area: Natural Language Processing — retrieval-augmented generation (RAG), long-context inference, and model-internal evidence selection.

Technical level: Intermediate to Advanced. It assumes familiarity with decoder-only LLMs, dense retrieval, late-interaction scoring, and contrastive training.

Scope: The paper introduces UNREAL, a framework that trains fewer than 500K parameters on top of a frozen LLM so that the same internal mechanism retrieves evidence from a 21M-chunk corpus and from within a long prompt.

What This Paper Is About

Long-context models and retrieval-augmented generation both solve the same underlying problem: picking out the small amount of text that actually answers a query. Long-context models do this implicitly and internally but are vulnerable to distractors, while RAG uses a separate external retriever, forcing production systems to train, serve, and align two models. The paper asks whether one pretrained decoder-only LLM can explicitly select evidence at both scales using a single model-internal mechanism, and shows that it can.

Key Contributions

1. Intrinsic corpus-scale retrieval for modern decoder-only LLMs. UNREAL extends INTRA (Hoffer et al., 2026) to decoder-only dense and hybrid models by using a single frozen LLM to encode candidates and extract retrieval queries. With fewer than 500K trainable parameters, all four tested backbones outperform the strongest dedicated retriever-reranker systems on a 21M-chunk Wikipedia index.

2. Long-context selection with the same mechanism. The retrieval-trained module also ranks chunks inside a long prompt and returns only the selected text to the generator, outperforming full-context inference and external-retriever baselines on NoLiMa, HELMET, LV-Eval, and LOFT, while reducing FLOPs and time-to-first-token at practical context lengths.

3. One selector across six orders of magnitude. The same model-native module selects evidence from an 8K-token context up to a 3B-token corpus without architectural changes or backbone finetuning.

4. A shared formulation of retrieval and long-context reading. The authors argue that the two are the same query-conditioned evidence-selection problem at different scales, and recommend that sparse-evidence long-context tasks be trained and evaluated as retrieval problems, with evidence recall reported alongside answer accuracy.

Main Findings

• Full-corpus retrieval gains. On the 3B-token, 21M-chunk Wiki-2018 index, the best model raises recall from 49.1% to 73.2% on HotpotQA, from 31.7% to 60.1% on 2WikiMultiHopQA, and from 8.8% to 14.4% on MuSiQue. All four dense and hybrid UNREAL backbones outperform state-of-the-art retriever-reranker systems. Metrics are complete-evidence recall@10, where an example counts only if all annotated oracle chunks are retrieved.

• Generation quality improves end to end. With Nemotron-3.5-Lightning fixed as generator and only the retriever varied at top-5 context, UNREAL-Nemo3.5-Lightning reaches EM/F1 of 52.6/65.8 on HotpotQA, 42.0/50.1 on 2Wiki, and 16.9/25.4 on MuSiQue, the highest of any method including BM25, BGE-large-en-v1.5, Qwen3-Emb-0.6B, Qwen3-Emb-4B, hybrid RAG, LateOn + reranker, and Qwen3-Emb-4B + reranker. With Qwen3.5-35B-A3B as generator, UNREAL-Qwen3.5-35B-A3B reaches 48.2/61.6, 40.2/47.2, and 17.8/27.1 respectively.

• Long-context accuracy recovers from near-total failure. Applying the same selection mechanism to remove distractors before generation raises NoLiMa accuracy from 1.0% to 24.83% at its maximum context length of 128K tokens, and LV-Eval's F1 score from 49.97% to 54.66% at 256K.

• Selection also wins on standard long-context benchmarks. On the RAG subset of HELMET, with context lengths of 8K–128K tokens and scores averaged across lengths, UNREAL achieves the highest exact-match score for both Nemotron-3.5-Lightning-30B-A3B and Qwen3.5-35B-A3B generators, beating full-context inference and other retrieval baselines. On LOFT, all selective budgets outperform full-context reading.

• Efficiency improves at long context. UNREAL uses fewer FLOPs than full-context inference for context lengths over 18K tokens, and measured stage-wise time-to-first-token with vLLM on a single H100 shows speedup growing with context length, starting roughly at 32K tokens. Unity retrieval encodes chunks independently, avoiding most inter-chunk attention, and generation attends only to selected chunks.

• Selection budget has an inverted-U optimum. On NoLiMa, increasing the number of selected chunks first improves accuracy by raising evidence recall, then degrades it as extra distractors interfere with generation.

• Frozen representations already retrieve, weakly. Untrained mean-pooled attention keys and queries from Nemotron-3-Nano retrieve evidence above a random baseline on NoLiMa, motivating the training of retrieval tokens to strengthen a signal that is already present.

• Distractor identity matters. On NoLiMa, with the gold chunk always present, retrieved distractors reduce the score more than the same number of random distractors, and the gap grows with the retrieval budget.

• Layer choice matters. Chunk embeddings are taken from an intermediate layer, selected on a development set. For UNREAL-Nemotron-3.5-Lightning-30B-A3B, layers 30, 38, and 42 give the strongest recall, with layer 42 best on HotpotQA (Avg R@10 81.28, All R@10 72.93) and MuSiQue (32.70, 13.79).

• More pooled vectors per chunk helps. Recall rises monotonically as the number of pooled vectors per chunk L_p increases over 1, 3, 7, and 11, at higher memory and compute cost.

• Multi-layer read-out is the single most important design choice. Aggregating all six global-attention blocks instead of a single one costs 8.26 Avg R@10 on HotpotQA and 11.92 All R@10 when removed, the largest drop among the ablated knobs. Removing BM25 initial context entirely costs 10.20 Avg R@10 on HotpotQA.

Methodology in Plain English

The approach builds on the observation that an intermediate layer of an LLM already holds rich meaning about the text it is reading. Instead of training a separate retriever model, the authors freeze the LLM and attach two small learnable pieces to it.

First, each candidate chunk is encoded on its own by running it through the frozen LLM and reading the residual stream — the running internal state — at one chosen middle-to-late layer. This produces a sequence of token-level vectors per chunk rather than a single embedding, and these are stored in the index. To keep the index small, the token vectors can be mean-pooled into a fixed number of groups.

Second, for a query, the authors prepend the top BM25-ranked chunks as an initial context, then append a small set of learned retrieval tokens at the end of the input. Because of the causal mask, those tokens' internal states see both the question and the BM25 context. Reading the model's residual stream at those token positions across multiple layers produces the query representation, and learned layer weights combine those layers into one. Retrieval scoring is a late-interaction MaxSim between the query representation and the stored chunk vectors.

Training uses a multi-positive InfoNCE contrastive loss that contrasts chunks containing annotated evidence against hard negatives mined with BM25 — up to 512 per query, lexically close to the gold evidence but guaranteed not to be part of it. Only the retrieval tokens and layer-mixing weights are updated; the LLM stays frozen. A single soft prompt token is also learned.

At inference, retrieval takes one forward pass to form the queries, and generation takes an ordinary second pass over the selected chunks, which are re-encoded rather than passed through cross-attention as in INTRA. For long-context tasks, the same procedure is applied to the prompt itself: the prompt is split into chunks of about 140 tokens, each is encoded separately, the chunks are ranked against the query, and only the top-n unchanged text is handed back to the same frozen LLM. The long-context experiments use the same modules trained on retrieval data, with no long-context training data.

Evaluation uses the DPR Wiki-2018 corpus mapped through a KILT-style procedure, retaining an oracle only if it comes from an exact title match or has an alignment score above 0.6. Four backbones are tested: dense Qwen3.5-4B and Muse-Glimmer-30B, linear-attention hybrid Qwen3.5-35B-A3B, and Mamba-attention hybrid Nemotron-3.5-Lightning-30B-A3B.

Why This Matters

Impact on research. The paper reframes long-context inference and corpus retrieval as one problem at two scales rather than competing paradigms, and shows that retrieval supervision — not just a bigger context window — is what teaches a model to locate evidence. It also provides a decoder-only alternative to INTRA, which required an encoder-decoder architecture and cross-attention. The authors argue benchmarks should annotate required evidence and distinguish sparse-evidence tasks from evidence-dense tasks that genuinely require integrating most of the input.

Real-world applications:

• Production RAG systems that currently serve and maintain separate embedding, reranking, and generation models could collapse to one frozen backbone.
• Long-document analysis — legal, medical, or financial review — where a single prompt contains one relevant passage among hundreds of thousands of tokens.
• Multi-hop question answering over large corpora, such as Wikipedia-scale or enterprise knowledge bases.
• Cost-sensitive long-context serving, where time-to-first-token and KV-cache size dominate, since UNREAL selects chunks before generation.

Industry relevance. The method adds fewer than 500K trainable parameters, under 0.005% of each tested backbone's parameter count, and leaves the backbone frozen, so it can be bolted onto existing pretrained models without retraining them. The demonstrated hardware setup is a single H100 with vLLM, and the reported savings begin at context lengths already common in deployment.

Future Directions

• Domain and language transfer. Training data came from Wikipedia QA, so transfer to heterogeneous domains and languages remains open.
• Beyond sparse evidence. Long-context experiments focus on sparse-evidence tasks; the authors call for benchmarks that separate sparse-evidence from evidence-dense problems requiring integration across most of the input.
• Reducing index and pipeline cost. UNREAL stores L_p vectors per chunk, costing more to build and store than a single-vector dense index, and it still relies on a cheap BM25 initial context, without which recall drops substantially.
• Integration with agentic pipelines. Multi-pass agentic RAG systems interleave reasoning with repeated retrieval and are excluded from the baselines because they evaluate a whole iterative pipeline; UNREAL is positioned as a single-pass component that could serve inside such systems.

Target Audience

Researchers and engineers working on retrieval-augmented generation, long-context inference, and efficient LLM serving; practitioners who want retrieval capability without maintaining a second model; and readers interested in the debate over whether long context subsumes retrieval. The paper is best suited to those already comfortable with dense retrieval, late-interaction scoring, and contrastive objectives.

Authors’ abstract

Long-context inference and Retrieval-Augmented Generation (RAG) handle evidence selection at vastly different scales, from a single long prompt to an entire corpus. We ask whether a single model-internal mechanism can select evidence across this range. We introduce UNifying REtrieval And Long-Context with a Single Model (UNREAL), a model-native evidence selection framework to span corpus retrieval and long-context inference. UNREAL encodes chunks and derives retrieval queries directly from the frozen LLM's internal representations. It adds fewer than 500K trainable parameters and leaves the backbone unchanged. On a 3B-token, 21M-chunk Wikipedia index, all four dense and hybrid UNREAL backbones outperform state-of-the-art retriever-reranker systems. The best model raises recall from 49.1% to 73.2% on HotpotQA, from 31.7% to 60.1% on 2WikiMultiHopQA, and from 8.8% to 14.4% on MuSiQue. Applied to long-context tasks, the same selection mechanism removes distractors before generation, raising NoLiMa accuracy from 1.0% to 24.83% at its maximum context length of 128K tokens, and LV-Eval's F1 score from 49.97% to 54.66% at 256K. UNREAL also reduces FLOPs and time-to-first-token relative to full-context inference from roughly 32K tokens onward, with larger gains as context grows. Together, these results establish model-internal evidence selection as a common foundation for corpus retrieval and evidence-sparse long-context inference.