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

What it claims

UNREAL is a small retrieval head on a frozen decoder-only LLM. It adds R learned retrieval tokens (a soft prompt) and per-layer mixing weights α, totalling under 0.005% of each backbone's parameters. With these, the LLM ranks chunks from its own residual-stream states. The same trained module works at two scales: a 21M-chunk, 3B-token Wikipedia corpus, and chunks of a single long prompt. In both cases it selects evidence, then feeds only the selected chunks back to the same LLM for generation. It is from NVIDIA and Technion. It extends INTRA (Hoffer et al., 2026), which only worked for encoder–decoder models with cross-attention.

How it works

flowchart LR
  B[BM25 top chunks] --> P[Prompt: BM25 + query + retrieval tokens]
  Q[Query] --> P
  P --> L[Frozen LLM]
  L --> M[Mix layers with α]
  C[Each chunk] --> E[Frozen LLM, layer ℓc states]
  M --> S[MaxSim score]
  E --> S
  S --> T[Top-n chunks]
  T --> G[Same LLM generates answer]

1. Chunk encoding. Each chunk (mean 140 tokens) is run through the frozen LLM on its own. The token-level states at one intermediate layer ℓc are stored. ℓc is chosen on a dev set, because intermediate layers beat final ones.
2. Query formation. The input is BM25's top-ranked chunks, then the query, then the R retrieval tokens. The retrieval tokens go last, so under causal masking their states see both the query and the BM25 context. Their states are read at every layer.
3. Scoring. The per-layer query states are summed with the α weights and scored against each chunk using ColBERT-style late-interaction MaxSim. The top-n chunks are kept.
4. Training. Only ρ (the retrieval tokens) and α are trained. The loss is multi-positive InfoNCE over oracle chunks plus hard negatives.
5. Compression. Chunk token states are mean-pooled into L_p groups, and query states into G groups. This trades index size and MaxSim cost against accuracy.
6. Inference cost. Inference takes two passes: one to form the queries, and one ordinary generation pass over the selected chunks, which are re-encoded.

Because the method reads the residual stream, it works on dense, linear-attention and Mamba-hybrid backbones without changes.

Evidence

Corpus retrieval. The task is complete-evidence recall@10 on the 21M-chunk Wiki-2018 corpus. A query counts as correct only if all its oracle chunks are retrieved. Eight datasets were used.

• Backbones: Qwen3.5-4B, Muse-Glimmer-30B, Qwen3.5-35B-A3B and Nemotron-3.5-Lightning-30B-A3B.
• Baselines: BM25, BGE-large, Qwen3-Emb-0.6B/4B (with and without the Jina reranker), LateOn multi-vector + reranker, and hybrid RRF.
• Averages across the eight datasets: UNREAL variants score 45.7–49.6. Baselines score 23.3–36.5.
• Where the gains are: largest on multi-hop sets. 2Wiki rises from 31.7 to 60.1, and HoVer from 28.1 to 47.4.
• Where the gains are thin: on Natural Questions, UNREAL scores 24.9–30.6 against a best baseline of 26.7, so one UNREAL variant trails that baseline.

End-to-end QA (Table 1). Nemotron-3.5-Lightning is the generator throughout, with top-5 retrieved chunks. EM / F1:

Long context. Results are reported on NoLiMa, LV-Eval, HELMET and LOFT. Compute savings over full-context inference start at about 32K tokens. Below that length, full context is not more expensive.

Limits

• No adoption details. The page mentions no code, weights or licence, so there is nothing to install yet.
• Truncated text. The fetched text stops partway through Section 4, so per-benchmark long-context tables and the ablations (Tables 3–5) are not available here.
• BM25 dependency. Query formation needs a BM25 first pass, so a lexical index is still required.
• Indexing cost. Building the index means a forward pass of the generator LLM over the whole corpus. The index also stores token-level vectors per chunk, unless compression is used, which costs accuracy.
• Narrow corpus coverage. Corpus retrieval is evaluated only on Wikipedia.
• Low absolute scores. MuSiQue recall@10 reaches only 14.4%, and NoLiMa accuracy at 128K is 24.83%. MuSiQue EM remains far below the oracle (16.9 vs 47.5).