PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Edan Kinderman, Elad Hoffer, Yochai Blau, Brian Chmiel, Ron Banner, Daniel Soudry, Boris Ginsburg

arXiv:2610.08463v1cs.CLcs.IRcs.LG

TL;DR

Long-context inference and RAG select evidence at different scales, but whether a single model-internal mechanism can span them remains an open question. UNREAL uses frozen LLM representations to rank evidence and outperforms baselines across corpus retrieval and long-context evaluations.

  • Problem

    Long-context inference selects evidence internally from a prompt, while corpus-scale RAG uses an external retriever; the paper asks whether one pretrained LLM can select evidence at both scales.

    Source: §1 · p. 2

  • Method

    UNREAL encodes chunks and derives retrieval queries from frozen LLM representations, training a small selector without changing the backbone.

    Source: §2.1 · p. 3, §7 · p. 9

  • Results

    UNREAL outperforms retrieval and long-context baselines across both settings; with Nemotron-3.5-Lightning fixed, it achieves the highest EM and F1 on all three datasets.

    Source: §3.1 · p. 5, §7 · pp. 9–10

  • Takeaways & Limitations

    The findings support viewing corpus retrieval and long-context inference as evidence selection at different scales, with a path toward using relevant information without a separate retrieval model.

    Source: §7 · p. 10

    Training originates from Wikipedia QA, so transfer to heterogeneous domains and languages remains future work; long-context experiments focus on sparse-evidence tasks.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-context inference and RAG both select query-relevant evidence, but use separate selection mechanisms at prompt and corpus scales. UNREAL uses one model-internal selector across both settings, with a frozen LLM backbone.

1 INTRODUCTION

  • Long-context and corpus retrieval both perform query-conditioned evidence selection, differing in scale and whether selection is internal or external.The paper asks whether a single pretrained LLM can explicitly select evidence at both scales.Source: §1 · p. 2

1.1 MODEL-INTERNAL EVIDENCE SELECTION ACROSS SCALES

  • UNREAL encodes and ranks candidate chunks using the frozen LLM’s representations, adding fewer than 500K trainable parameters.The same mechanism is designed to operate across long-context and corpus scales.Source: §1.1 · p. 2
  • Recall@10 rises from 49.1% to 73.2% on HotpotQA and from 31.7% to 60.1% on 2WikiMultiHopQA.The index contains 3B tokens and 21M Wikipedia chunks; all four tested backbones outperform the strongest dedicated retriever-reranker systems.Source: §1.1 · p. 2
  • At 128K tokens, NoLiMa accuracy rises from 1.0% to 24.83%, and at 256K, LV-Eval F1 rises from 49.97% to 54.66%.The same selector reduces FLOPs and time-to-first-token versus full-context inference from roughly 32K tokens onward, with larger gains as context grows.Source: §1.1 · p. 2

1.2 CONTRIBUTIONS

  • One model-native module selects evidence from an 8K-token context to a 3B-token corpus without architectural changes or backbone finetuning.The contribution frames retrieval and long-context reading as two scales of the same selection mechanism.Source: §1.2 · p. 2

2 METHOD

UNREAL uses a decoder-only LLM’s own residual-stream representations to encode chunks and form query-conditioned retrieval scores. It trains a small selector while leaving the backbone frozen, then generates from selected chunks.

2 METHOD

  • UNREAL extends intrinsic retrieval to decoder-only and hybrid LLMs, replacing encoder-derived chunk embeddings and cross-attention queries with representations available inside the model.Its components are chunk embeddings, residual-state queries, and generation.Source: §2 · p. 3

2.1 CHUNK EMBEDDINGS

  • Each corpus chunk is encoded independently with the frozen LLM, using token-level residual-stream representations from an intermediate layer selected by development-set performance.Each chunk is represented as a sequence of token vectors.Source: §2.1 · p. 3

2.2 RESIDUAL STATE QUERIES

  • UNREAL forms query representations from retrieval-token states conditioned on the query and BM25-ranked initial context, then scores chunks with late-interaction MaxSim.Training uses a multi-positive InfoNCE loss over oracle chunks and sampled hard negatives; only retrieval tokens and layer-mixing coefficients are updated.Source: §2.2 · p. 4

2.3 SCORING AND GENERATION

  • UNREAL selects the top-n chunks, concatenates them with the query, and reencodes them during a second, ordinary generation pass.The first LLM pass forms retrieval queries; the second generates from the selected chunks.Source: §2.2 · p. 4, §2.3 · p. 4

2.2 RESIDUAL STATE QUERIES

  • Mean-pooling chunk and query representations into Lp and G groups trades retrieval accuracy against index size and MaxSim compute cost.Chunk representations are compressed from Ti × d to Lp × d, and query representations from R × d to G × d.Source: §2.2 · p. 4

2.3 SCORING AND GENERATION

  • Across four architectures, UNREAL outperforms baselines on complete-evidence recall@10 over 21M Wiki-2018 chunks, especially on multi-hop datasets.Figure 3 reports results across eight datasets.Source: Figure 3 · p. 5

3 UNREAL AS A RETRIEVAL MECHANISM

UNREAL evaluates frozen-LLM retrieval over a large Wikipedia corpus and across diverse benchmarks. It performs strongly in both evidence retrieval and end-to-end question answering while adding few trainable parameters.

3 UNREAL AS A RETRIEVAL MECHANISM

  • The evaluation uses a 3B-token, 21M-chunk Wiki-2018 corpus and measures complete-evidence recall@n against retrieval and reranking baselines.The benchmark spans eight datasets; partial evidence receives no credit for multi-hop tasks.Source: §3 · p. 5

3.1 RESULTS

  • UNREAL demonstrates architectural generality across four tested dense and hybrid backbones.The evaluated models include Qwen3.5-4B, Muse-Glimmer-30B, Qwen3.5-35B-A3B, and Nemotron-3.5-Lightning-30B-A3B.Source: §3.1 · p. 5
  • Each UNREAL backbone leads on most of eight datasets for complete-evidence recall@10, with the largest gains on four challenging multi-hop tasks.The largest gains occur on HotpotQA, 2WikiMultiHopQA, MuSiQue, and HoVer.Source: §3.1 · p. 5
  • UNREAL attains the highest EM and F1 on all three datasets using top-5 retrieved contexts and the same LLM for retrieval and generation.Table 1 reports end-to-end QA with Nemotron-3.5-Lightning as the fixed generator.Source: §3.1 · p. 5, Table 1 · p. 6
  • UNREAL turns a frozen LLM into a full-corpus retriever by training retrieval-token embeddings and layer-mixing weights, adding fewer than 0.5M parameters.The added parameters are under 0.005% of each tested backbone’s parameter count.Source: §3.1 · p. 6

4 UNREAL AS A LONG CONTEXT MECHANISM

UNREAL treats a long prompt as chunks to retrieve from, using the same frozen-LLM selection mechanism as corpus retrieval. Across long-context benchmarks, it outperforms full-context inference and retrieval baselines, while its computational savings grow with context length.

4.1 RESULTS

  • UNREAL chunks the prompt, ranks chunks against the query, selects the top-n, and passes their unchanged text to the same frozen LLM for answering.The models are those presented in Section 3; no special long-context training data was involved.Source: §4.1 · p. 6
  • UNREAL outperforms all baselines on a constructed QA benchmark spanning 8K to 100M tokens and enables generation from evidence selected at 100M tokens.The benchmark pads queries and oracle chunks with random Wiki-2018 distractor passages.Source: §4.1 · p. 6
  • UNREAL consistently achieves the highest accuracy across NoLiMa lengths up to 128K tokens and LV-Eval lengths up to 256K words, as full-context performance degrades.Figure 4 compares NoLiMa accuracy and LV-Eval F1 across context lengths.Source: Figure 4 · p. 7, §4.1 · p. 7
  • UNREAL achieves the highest exact-match score on HELMET’s RAG subset for both generators, while full-context inference remains a strong baseline.The results are averaged across context lengths of 8K–128K tokens.Source: Figure 5 · p. 7, §4.1 · p. 7
  • On NoLiMa, increasing the selected chunk count first improves accuracy through higher evidence recall, then reduces accuracy as added distractors interfere with generation.The resulting accuracy curve follows an inverted-U pattern.Source: §4.1 · p. 7
  • Untrained frozen-LLM layer representations retrieve NoLiMa evidence above random, motivating UNREAL’s training of retrieval tokens to strengthen that signal.The evaluation scores chunks using pooled attention keys and query representations.Source: §4.1 · p. 7, Figure 6 · p. 8

4.2 EFFICIENCY

  • With fixed chunk size and selection budget, UNREAL scales linearly with context length, while standard-attention full-context inference scales quadratically.Independent chunk encoding avoids most inter-chunk prefill attention, and generation attends only to selected chunks.Source: §4.2 · p. 8
  • UNREAL uses fewer FLOPs beyond 18K tokens and reduces time-to-first-token from roughly 32K tokens onward, with speedups growing as context length increases.The reported speedup curves include FLOP-scaled projections beyond the last measured full-context runtime.Source: §4.2 · p. 8

5 RELATED WORK

UNREAL differs from retrieval and efficient-attention approaches by using a retrieval-trained mechanism built on a frozen decoder’s representations to select chunks before generation. It combines retrieval and generation within a shared frozen backbone rather than relying on a separate retriever or reader.

Retrieval models

  • UNREAL shares frozen-backbone representations across retrieval and generation while training a separate query head, unlike systems with distinct retriever and reader components.It adapts INTRA’s internal-representation retrieval to standard decoder-only LLMs.Source: §5 · p. 8

Efficient attention

  • Unlike efficient-attention methods that select blocks or route attention through summaries, UNREAL retrieval-selects chunks and gives only selected text to the LLM for generation.The passage distinguishes block selection, KV-block ranking, and hierarchical chunk-summary routing from UNREAL’s retrieval-trained selection.Source: §5 · pp. 8–9

RAG and long context

  • Unlike external retrieval pipelines, UNREAL needs no external retrieval model; unlike full-context baselines, it removes unselected chunks before generation.The paper frames removal of unselected chunks as preventing them from distracting the generator.Source: §5 · p. 9

6 LIMITATIONS

The evaluation is bounded by domain transfer, evidence type, and retrieval setup. Long-context results focus on sparse-evidence tasks, and retrieval adds indexing and search costs.

6 LIMITATIONS

  • Transfer beyond Wikipedia QA to heterogeneous domains and languages remains future work.Long-context experiments currently focus on sparse-evidence tasks.Source: §6 · p. 9
  • UNREAL stores Lp vectors per chunk and relies on a cheap BM25 initial context, adding storage, encoding, and search costs.The additional steps are inherent to retrieval-based approaches, not unique to UNREAL.Source: §6 · p. 9
  • Multi-pass agentic RAG systems are excluded because their full iterative pipelines are not directly comparable to UNREAL’s single retrieval stage.UNREAL could instead serve as a component in future agentic pipelines.Source: §6 · p. 9

7 DISCUSSION AND FUTURE DIRECTIONS

UNREAL treats corpus retrieval and long-context inference as evidence selection at different scales, using one model-internal mechanism. The authors argue for retrieval-aware training and evaluation while distinguishing sparse-evidence tasks from evidence-dense integration.

7 DISCUSSION AND FUTURE DIRECTIONS

  • UNREAL uses a frozen LLM’s representations to encode and rank chunks, with fewer than 500K trainable parameters, and outperforms state-of-the-art retrieval methods on retrieval and long-context benchmarks.The findings support viewing both settings as evidence selection at different scales.Source: §7 · pp. 9–10

7.1 FUTURE DIRECTIONS

  • Retrieval supervision makes the model’s internal retrieval ability trainable, improving recall on corpus-scale retrieval and in-prompt long-context selection.The authors describe this as learning to rank relevant chunks rather than relying on retrieval ability emerging from pretraining alone.Source: §7.1 · p. 10
  • The authors motivate training retrieval together with generation and evaluating sparse-evidence long-context tasks as retrieval problems, reporting evidence recall alongside answer accuracy.Models should identify needed evidence before generating.Source: §7.1 · p. 10
  • Benchmarks should distinguish sparse-evidence tasks from evidence-dense problems that require integrating information across much of the input.The paper explicitly cautions that not all long-context understanding reduces to selecting a few chunks.Source: §7.1 · p. 10

A.1 CORPUS-SCALE RETRIEVAL RESULTS

Across eight datasets on the full Wiki-2018 corpus, UNREAL outperforms state-of-the-art embedding models, rerankers, and late-interaction methods. Additional experiments examine long-context distractors and untrained intrinsic retrieval.

A.1 CORPUS-SCALE RETRIEVAL RESULTS

  • Across eight datasets on 21M Wiki-2018 chunks, UNREAL outperforms state-of-the-art embedding models, rerankers, and late-interaction methods on complete-evidence recall@20, average recall@10, and average recall@20.Complete-evidence recall@20 counts an example as correct only when all annotated evidence chunks are retrieved.Source: §A.1 · p. 16, Figure 8 · p. 16, Figure 9 · p. 16

A.2 LONG-CONTEXT RESULTS

  • With gold evidence present, retrieved distractors reduce NoLiMa scores more than the same number of random distractors, and the gap grows with distractor budget.Figure 11 compares retrieved and random distractors.Source: §A.2 · p. 17, Figure 11 · p. 18
  • All selective budgets outperform full-context reading on LOFT, while the dataset comparison at n = 10 uses gold-only context as its reference.The figure reports each dataset’s designated LOFT metric.Source: Figure 12 · p. 18

A.2.1 INTRINSIC RETRIEVAL

  • For intrinsic retrieval, frozen Nemotron-3-Nano chunks are scored against the question using attention keys and queries, without retrieval training.The experiment splits each context into chunks and evaluates retrieval on NoLiMa.Source: §A.2.1 · p. 18
  • Untrained model representations retrieve evidence above the random baseline, with recall@n computed over the top-n chunks.No retrieval tokens or learned layer weights are used.Source: §A.2.1 · p. 18

A.3 ABLATION STUDIES

The ablations examine recipe choices, embedding-layer choice, and the number of pooled vectors per corpus chunk. They report strongest retrieval from middle-to-late layers and monotonically increasing recall as the vector count rises.

A.3 ABLATION STUDIES

  • Middle-to-late layers yield the strongest corpus retrieval in the layer ablation, a finding reproduced across other architectures.Table 4 studies which layer produces corpus chunk embeddings.Source: §A.3 · p. 19, Table 4 · p. 19
  • The recipe ablation varies one design choice at a time, including retrieval tokens, query-side grouping, readout layers, hard negatives, BM25 context, and soft prompts.Table 3 reports Avg. R@10 and All R@10 for these departures from the main method.Source: Table 3 · p. 19, §A.3 · p. 19
  • Increasing pooled vectors per chunk (Lp) monotonically raises Average recall@10 and Complete-evidence recall@10, at the cost of higher evaluation memory and compute.Table 5 evaluates Lp values of 1, 3, 7, and 11.Source: §A.3 · p. 19

B.1 RETRIEVAL DATA

The appendices describe how retrieval data and extended-context benchmarks are constructed, then specify UNREAL’s encoding, training, and evaluation setup.

B.1 RETRIEVAL DATA

  • Recall@10 increases monotonically with pooled vectors per chunk, trading higher retrieval recall for more evaluation memory and compute.This comparison is for UNREAL-Nemotron-3.5-Lightning-30B-A3B; no table values are provided here.Source: Table 5 · p. 20
  • Gold evidence is mapped to a shared Wikipedia-2018 corpus by title matching and text overlap, with filtered evaluation and a direct DPR mapping for NQ.The default alignment quality threshold is 0.6; BM25 also supplies lexical context and up to 512 oracle-seeded hard negatives per query.Source: §B.1 · pp. 19–20

B.3 METHOD AND TRAINING CONFIGURATION

  • UNREAL modules are trained on standard retrieval data, not long-context data, then evaluated on both retrieval and long-context tasks.Training uses eight datasets mapped to Wikipedia-2018; evaluation scores each query exhaustively against all corpus chunks, without approximate indexing or candidate pre-filtering.Source: §B.2 · p. 20, §B.3 · p. 21

B.2 LONG-CONTEXT BENCHMARKS

  • The extended QA benchmark spans 8K–100M tokens, placing gold passages at fixed relative depths amid distractors drawn from a shared pool.It uses 2,000 questions sampled across four QA datasets; passages are split into sentence-aligned chunks averaging 141 tokens for retrieval evaluation.Source: §B.2 · pp. 20–21

B.3 METHOD AND TRAINING CONFIGURATION

  • The query uses 64 retrieval tokens pooled into four groups, while a learned soft prompt steers retrieval computation; all trainable components total fewer than 500K parameters.Corpus chunks are encoded once from a frozen-LLM layer into seven pooled vectors.Source: §B.3 · p. 21

C FLOP ANALYSIS

The FLOP analysis compares full-context inference with UNREAL’s chunk-then-select pipeline, identifying its savings, added costs, and approximate break-even context length.

C.2 FLOP DIFFERENCE: UNRAEL VS. FULL-CONTEXT INFERENCE

  • UNREAL saves work by encoding chunks only through layer ℓc and avoiding attention to unselected chunks, but adds retrieval and selected-chunk re-encoding passes.The analysis compares only work that differs between the two pipelines.Source: C FLOP ANALYSIS · p. 21, §C.2 · pp. 22–23

C FLOP ANALYSIS

  • During decoding, UNREAL attends to selected rather than full-context tokens, using fewer attention FLOPs and a smaller KV cache.The smaller cache benefits memory-bound decoding.Source: C FLOP ANALYSIS · p. 22

C.1 SETTING

  • The FLOP model assumes equal-length chunks and omits softmax, normalization, activations, and selection costs; omitted terms are small relative to chunk encoding.Sliding-window attention is also omitted, conservatively understating UNREAL’s savings.Source: §C.1 · p. 22

C.2.1 SIMPLIFYING ASSUMPTIONS

  • The simplified break-even equation assumes context is much longer than selected evidence; its neglected terms vanish in the long-context limit, but the result is an approximation, not a bound.The dropped terms have mixed signs, and diagonal corrections are smaller than their corresponding weight terms for evaluated backbones.Source: §C.2.1 · p. 23

C.2.2 BREAK-EVEN CONTEXT LENGTH

  • The break-even length falls when any ℓc>0 adds linear savings, and longer generations lower it toward nT from above.For last-layer readout, avoided cross-chunk attention balances the extra weight cost of tokens processed by UNREAL.Source: §C.2.2 · p. 24

C.3 N ∗FOR THE EVALUATED BACKBONES

  • For the evaluated backbones and settings, equation 10 is within 0.55% of the numerical break-even root of full equation 8.The comparison uses T=141, Tq=Tg=T, and |xret|=910; Table 6 reports approximate break-even context lengths.Source: §C.3 · p. 24, Table 6 · p. 24

D EFFICIENCY BENCHMARK: METHODOLOGY AND FULL RESULTS

The efficiency benchmark measures UNREAL’s chunk-then-select pipeline against full-context inference using vLLM and a FLOPs model. Measured TTFT scales roughly linearly for UNREAL, while the baseline grows quadratically and eventually becomes infeasible.

Setup

  • The benchmark uses randomly initialized models and random token inputs, served with vLLM on one NVIDIA H100, and measures performance rather than accuracy.It uses bf16 weights and activations, with the Qwen3.5-35B-A3B KV cache in fp8; the context sweep runs from 8K to 100M tokens.Source: D EFFICIENCY BENCHMARK: METHODOLOGY… · p. 24

FLOPs model

  • The FLOPs model sums per-layer projection and attention costs, yielding full-context cost O(L)+O(L^2) and UNREAL cost O(L)+O(LT) for chunk encoding.UNREAL generation reads a fixed 1,280 selected tokens in this benchmark.Source: D EFFICIENCY BENCHMARK: METHODOLOGY… · pp. 24–25

Measured wall-clock results

  • UNREAL’s measured TTFT grows roughly linearly with context length, while full-context TTFT grows quadratically and eventually becomes infeasible to run.The baseline is directly measured only up to its budget cutoff and then projected; UNREAL measurements omit MaxSim scoring and orchestration overhead.Source: Figure 13 · p. 25, D EFFICIENCY BENCHMARK: METHODOLOGY… · p. 25
  • Figure 14 reports prefill throughput for all four backbones using the same data and conventions as the TTFT measurements.The caption identifies throughput in tokens per second.Source: Figure 14 · p. 26
Loading 2610.08463v1…