PaperlayerView exact revision on arXiv

Source-linked AI summary

DREAM: Dense Retrieval Embeddings via Autoregressive Modeling

Yixuan Tang, Yi Yang

arXiv:2606.24667v1cs.CL

TL;DR

Dense retrievers typically rely on costly labeled positive and negative pairs, motivating whether autoregressive language-model prediction can supervise retrieval. DREAM connects a dense retriever to a frozen LLM through attention and consistently outperforms existing baselines across BEIR and RTEB at multiple model scales.

  • Problem

    Dense retriever training often requires costly labeled relevance pairs, leaving the usefulness of autoregressive next-token prediction as retrieval supervision to be investigated.

  • Method

    DREAM injects retriever query-document similarity scores into selected attention heads of a frozen LLM, allowing next-token prediction loss to train a standalone dense retriever.

  • Results

    Across BEIR and RTEB and embedding backbones from 0.5B to 3B parameters, DREAM consistently outperforms existing baselines, with gains over Revela of 0.015 to 0.081 NDCG@10 on BEIR and 0.068 to 0.102 on RTEB.

  • Takeaways & Limitations

    The results support autoregressive next-token prediction as a practical and scalable alternative to contrastive supervision for training dense retrievers.

Abstract

from arXiv · show

Dense retrieval embedding models are a fundamental component of modern retrieval-based AI systems. Most dense retrievers are trained with contrastive objectives, which require labeled positive and negative document pairs that are often costly and difficult to obtain. In this work, we investigate whether the autoregressive next-token prediction objective of a large language model (LLM) can provide supervision for dense retrieval. The intuition is simple: if a document contains information relevant to a query, conditioning on that document should make the target output easier for the LLM to predict. A key challenge is that the next-token prediction loss is computed inside the LLM, while the retriever is a separate embedding model. To address this challenge, we propose DREAM (Dense Retrieval Embeddings via Autoregressive Modeling), which injects retriever-generated query-document similarity scores into selected attention heads of a frozen LLM. During training, these scores determine how much attention each candidate document receives while the LLM predicts the target output. The resulting prediction loss provides gradients for retriever training through the attention mechanism. We evaluate DREAM on retrieval benchmarks BEIR and RTEB using embedding backbones ranging from 0.5B to 3B parameters. DREAM consistently outperforms existing baselines across different model scales. These results demonstrate that DREAM provides a promising approach for training dense retrievers through autoregressive modeling.

1 Introduction

DREAM addresses the difficulty of training dense retrievers from costly contrastive examples by connecting retriever scores to a frozen LLM’s attention, allowing next-token prediction loss to supervise retrieval. Across 0.5B–3B embedding backbones, DREAM consistently outperforms RePlug and Revela on BEIR and RTEB.

  • Motivation: Retrievers determine which context an LLM receives, making retriever training important for retrieval-augmented and agentic systems.These systems retrieve external documents, revisit memory, call tools, and plan subsequent actions.
  • Problem: Contrastive dense-retriever training increases similarity to positive documents and decreases similarity to sampled negatives, but constructing these examples is often the bottleneck.Positive documents are desired retrieval targets, whereas sampled negatives are lower-relevance alternatives.
  • Supervision signal: DREAM uses a candidate document’s usefulness for target-output prediction as retrieval supervision: useful documents should reduce the LLM’s next-token prediction loss.The signal evaluates retrieval by measuring its effect on the LLM’s prediction loss.
  • Method: DREAM injects retriever-computed query-document similarity scores into selected attention heads of a frozen LLM, enabling prediction-loss gradients to update the retriever.Candidate documents, the query, and target output are concatenated, while retrieval scores control attention assigned from the query to each candidate document.
  • Results: DREAM consistently outperforms RePlug and Revela across 0.5B–3B embedding backbones, gaining 0.015–0.081 NDCG@10 on BEIR and 0.068–0.102 on RTEB.The evaluations use the BEIR and RTEB retrieval benchmarks.

2 Related Work

This section situates DREAM among retrieval-model training, language-model supervision for retrieval, and attention-head research. It contrasts pairwise supervision with deriving retriever training signals from a frozen LLM’s next-token prediction objective.

  • Retrieval models: Dense dual-encoders separately embed queries and passages, ranking them by similarity within a shared space.
  • Retrieval models: Contrastive retrieval training relies on labeled or synthesized query-document pairs, whose construction can be costly and supervision noisy.Noise may arise from annotation errors or false negatives.
  • Language-model supervision for retrieval: DREAM derives supervision directly from a frozen LLM’s next-token prediction objective instead of pre-constructed positive and negative pairs.
  • Language-model supervision for retrieval: Language-model supervision has trained retrievers through masked-language-model objectives or frozen-LLM likelihoods.REALM jointly trains a retriever with a masked language model, whereas RePlug keeps retriever scores outside the LLM forward pass.
  • Attention heads in language models: Multihead attention supports specialized token-to-token computations, including induction heads for structured copying and query-focused retrieval heads linking queries to relevant context.Prior work mainly uses retrieval heads for analysis.

3 Method

DREAM trains a dense retriever by injecting its query-document weights into selected attention heads of a frozen decoder-only LLM, whose target-passage prediction loss directly supervises retrieval. The method preserves the LLM’s within-document token preferences while letting the retriever control document-level attention, with normalized candidate competition and no mined negatives.

  • Training architecture: DREAM feeds retriever query-document scores into selected attention heads of a frozen decoder-only LLM while it predicts a target passage.The frozen LLM evaluates how well the retriever’s selected documents support target prediction.
  • Retriever signal: Retriever scores are normalized into pϕ(dj | q), a document-level signal that enters the frozen LLM during training.The same similarity score is used by the final retriever at inference time, while a learnable temperature normalizes candidate scores for training.
  • Head selection: DREAM selects the top M retrieval-oriented heads using query-focused probes, query-independent baseline subtraction, and NDCG@10 against known relevant documents.Only the selected heads receive retriever-guided attention.
  • Retriever-guided attention: For selected heads and query-token rows, pϕ(dj | q) controls attention across documents while frozen LLM attention controls token allocation within each document.DREAM normalizes original attention over each document’s tokens, scales it by the retriever’s document weight, and mixes it with original attention using a learnable gate g = σ(γ).
  • Training objective: The frozen LLM applies standard next-token cross entropy to target-passage tokens, and gradients flow through modified attention into the retriever.Increasing a candidate’s similarity is rewarded when it helps target prediction and penalized when it does not.
  • Training objective: The objective directly updates the retriever through modified attention and creates normalized competition among candidates without mined negatives.Unlike likelihood distillation, it does not first generate judge scores for the retriever to imitate.

4 Experiments

DREAM is evaluated as a standalone retriever against lexical, contrastive, and LLM-supervised baselines on BEIR and RTEB. It consistently outperforms these baselines across embedding scales and retrieval tasks, with gains broadening across scientific, biomedical, community-question-answering, code, and structured-data domains.

  • The main experiment compares DREAM with lexical and LLM-supervised retrieval baselines to test whether it trains a stronger standalone retriever.
  • DREAM is compared with BM25, RePlug, Revela, and InfoNCE, with InfoNCE using the same data and candidate pool but a contrastive objective.No hard-negative mining is applied, so InfoNCE and DREAM differ only in their training objective.
  • Evaluations use NDCG@10 on BEIR’s nine tasks and RTEB’s fourteen tasks, spanning domains including scientific, biomedical, legal, code, and medical retrieval.Detailed per-task scores are provided in Section C.
  • DREAM outperforms InfoNCE at every scale, exceeds BM25 increasingly at larger scales on BEIR, and is better than BM25 on RTEB at all scales.Because InfoNCE and DREAM share data and candidate pools, the gap points to the training objective rather than the candidate set.
  • The gains are broad rather than dataset-specific, with especially large RTEB improvements on code and structured-data tasks such as Apps, MBPP, and WikiSQL.On BEIR, DREAM becomes strongest on nearly all tasks as the backbone grows, including scientific, biomedical, and community-question-answering tasks.

5 Analysis: Why Does the Signal Work?

DREAM’s signal depends on where retriever scores enter the frozen LLM and how many heads receive them. Selected heads and an intermediate head count provide stronger supervision, while DREAM also learns a less collapsed embedding space.

  • Effect of head selection: 0.4888 average BEIR NDCG@10 and 0.5514 average RTEB NDCG@10 are achieved by selected heads, versus 0.0637 and 0.0320 with fully random heads.Random middle-layer heads improve over fully random heads but remain far below selected heads.
  • Effect of head selection: Because the LLM is frozen, retrieval scores help only when they modulate heads already affecting how queries read candidate context.Increasing a useful candidate’s weight can lower next-token prediction loss and guide the retriever toward higher rankings.
  • Number of selected heads: Performance improves from Top 1 to Top 16 selected heads, then drops at Top 32 and Top 64, making Top 16 the best balance.One head is insufficient, whereas too many can include weaker retrieval heads and dilute the signal.
  • Embedding geometry: DREAM achieves the best uniformity while keeping alignment in the same range as Revela, indicating a less collapsed embedding space beyond improved alignment.RePlug places positives very close to queries but spreads embeddings poorly, while Revela improves spread with weaker alignment.

6 Training Ablations

Training ablations show that freezing the judge LLM and using more candidate documents improve retriever training, while DREAM also scales to an 8B embedding backbone with strong BEIR and RTEB results. DREAM-8B matches E5-mistral-7b-instruct on BEIR but trails Qwen3-Embedding-8B.

  • Effect of updating the judge LLM: Freezing the judge LLM is better in all four comparisons than updating it with LoRA adapters.The fixed LLM provides a more stable training signal for the retriever, whereas updating the judge can weaken the signal reaching the retrieval model.
  • Number of candidate documents: Increasing the candidate set from 4 to 16 improves average NDCG@10 on both BEIR and RTEB.The larger set provides more documents for each training step to compare, sharpening the training signal; the experiments use the Llama-3.2-1B embedding backbone.
  • Scaling to larger models: DREAM-8B reaches 0.5531 average NDCG@10 on BEIR and 0.6417 on RTEB.DREAM-8B is trained from Llama-3.1-8B using the same recipe as the smaller models.
  • Scaling to larger models: DREAM-8B matches E5-mistral-7b-instruct on BEIR and achieves a 0.0386 higher score on RTEB, but does not surpass Qwen3-Embedding-8B.The Qwen3 comparison uses a different backbone model and a more carefully curated dataset, so it should be interpreted with caution.

7 Conclusion

DREAM trains standalone dense retrievers using autoregressive next-token prediction by injecting query-document scores into selected attention heads of a frozen LLM. Across BEIR and RTEB, it outperforms existing LLM-supervised retrieval baselines and benefits from query-focused retrieval heads.

  • 7 Conclusion: DREAM trains standalone dense retrievers through autoregressive next-token prediction.This enables retrieval training without manually labeled relevance pairs.
  • 7 Conclusion: Injecting query-document scores into selected attention heads of a frozen LLM enables retrieval training without manually labeled relevance pairs.
  • 7 Conclusion: Across BEIR and RTEB, DREAM consistently outperforms existing LLM-supervised retrieval baselines.
  • 7 Conclusion: Analyses show that the supervision signal works best through query-focused retrieval heads and naturally produces a better embedding space.

A Query-Generation Prompt

The query-generation prompt instructs the model to write one passage-grounded question that requires integrating multiple sentences. It constrains the response to question-only XML output using the target passage as input.

  • Usage: This prompt is used to generate a query from the target passage in each candidate chunk set.The target passage is inserted into the prompt as {target_passage}.
  • Prompt objective: The prompt asks for ONE question that can ONLY be answered from the target passage.The target passage is supplied through the placeholder {target_passage}.
  • Prompt objective: The answer must rely on passage information, avoid general knowledge, and require reading more than one sentence.These requirements are intended to keep generated queries passage-specific and multi-sentence.
  • Output format: The generated response must contain ONLY the question inside XML tags.The specified format is <question> Your question here </question>.

B Selected Attention Heads

DREAM’s default configuration selects 16 attention heads using a query-focused retrieval-head ranking from a frozen Llama-3.1-8B-Instruct judge, computed before retriever training.

  • Selected Attention Heads: The default DREAM setting uses the top 16 attention heads ranked for query-focused retrieval.The ranking is computed on 5,000 training examples before retriever training.
  • Selected Attention Heads: The retrieval-head ranking comes from the frozen Llama-3.1-8B-Instruct judge.
  • Selected Attention Heads: Table 3 lists each selected head as a zero-based (layer, head) pair.The indexing follows the detection code.

C Per-Task Retrieval Results

Tables 4 and 5 provide per-task NDCG@10 results for the main experiment on BEIR and RTEB, complementing the average scores in Table 1. Results are grouped by base embedding model, with BM25 listed without a learned backbone.

  • BEIR: BEIR per-task NDCG@10 results are grouped by the base embedding model.BM25 has no learned backbone.
  • RTEB: RTEB per-task NDCG@10 results are grouped by the base embedding model.BM25 has no learned backbone.
Loading 2606.24667v1…