DREAM: Dense Retrieval via Autoregressive Modeling

• The paper demonstrates that an autoregressive next-token prediction loss can supervise a dense retriever, eliminating the need for hand-labeled positive and negative pairs.
• DREAM injects retriever-generated query-document similarity scores into selected attention heads of a frozen LLM to enable gradient flow for optimizing the dual-encoder retriever.
• Empirical results on BEIR and RTEB benchmarks show that DREAM outperforms contrastive and lexical baselines, with notable gains at increased embedding scales.

Searching arXiv for the specified DREAM paper and closely related retrieval context. Dense retrieval embedding models are a fundamental component of modern retrieval-based AI systems, yet most dense retrievers are trained with contrastive objectives that require labeled positive and negative document pairs that are often costly and difficult to obtain. DREAM, short for Dense Retrieval Embeddings via Autoregressive Modeling, investigates whether the autoregressive next-token prediction objective of a LLM can provide supervision for dense retrieval by using a frozen decoder-only LLM as a training-time judge and a trainable dual-encoder retriever as the object of optimization (Tang et al., 23 Jun 2026). Its central intuition is that if a document contains information relevant to a query, conditioning on that document should make the target output easier for the LLM to predict; DREAM operationalizes this intuition by injecting retriever-generated query-document similarity scores into selected attention heads of the frozen LLM and using the resulting prediction loss to train the retriever through the attention mechanism (Tang et al., 23 Jun 2026).

1. Motivation and problem setting

Traditional dense retrievers rely on contrastive losses over hand-labeled positive and negative document pairs. In practice, obtaining reliable relevance labels and hard negatives is expensive and often noisy. DREAM replaces this supervision regime with an autoregressive one: if candidate documents are fed into a frozen LLM as additional context, then documents that truly help predict a query’s answer should lower the LLM’s next-token prediction loss, and that drop in loss can be turned into a training signal for a standalone dense retriever (Tang et al., 23 Jun 2026).

This design eliminates the need for explicit relevance labels and hard-negative mining. The supervision source is instead the LLM’s own cross-entropy loss. The paper emphasizes a structural difficulty, however: the next-token prediction loss is computed inside the LLM, while the retriever is a separate embedding model. DREAM addresses this by coupling retriever scores to the LLM’s attention distribution so that gradients from autoregressive prediction can flow back into the retriever (Tang et al., 23 Jun 2026).

A plausible implication is that DREAM reframes retrieval supervision as conditional utility estimation under an autoregressive LLM rather than pairwise relevance discrimination. The paper itself describes the resulting effect more concretely: because the retriever assigns a normalized distribution over candidate documents, increasing weight on helpful documents necessarily decreases weight on distractors, thereby creating implicit negatives without explicit sampling (Tang et al., 23 Jun 2026).

2. Training-time architecture

DREAM consists of two modules at training time. The first is a trainable dual-encoder dense retriever $f_\phi$ that maps queries $q$ and documents $d_j$ into a shared embedding space and computes similarity scores. The second is a frozen autoregressive LLM, specifically a decoder-only model, that acts as a judge by attending over a concatenation of candidate documents, the query, and the target passage, then computing next-token cross-entropy loss on the target (Tang et al., 23 Jun 2026).

The retriever embeds and L2-normalizes query and document representations:

$$e_q = \frac{f_\phi(q)}{\|f_\phi(q)\|_2}, \qquad e_{d_j} = \frac{f_\phi(d_j)}{\|f_\phi(d_j)\|_2}.$$

Similarity is cosine:

$s_\phi(q,d_j) = e_q^\top e_{d_j}.$

Over $K$ candidates, these scores are converted into a softmax distribution with learnable temperature $\tau$:

$$p_\phi(d_j\mid q) = \frac{\exp\bigl(s_\phi(q,d_j)/\tau\bigr)} {\sum_{k=1}^K \exp\bigl(s_\phi(q,d_k)/\tau\bigr)}.$$

This distribution supplies the cross-document weighting used inside the LLM’s attention mechanism (Tang et al., 23 Jun 2026).

Before training, DREAM identifies a small set $H$ of query-focused retrieval heads in the frozen LLM via the method of Zhang et al. (2025). At each LLM layer $\ell$ and head $q$0, the model re-weights the post-softmax attention from query-token positions $q$1 to document tokens $q$2. If $q$3 denotes the original attention, DREAM first performs within-document normalization,

$q$4

then applies retriever-guided cross-document weights,

$q$5

with $q$6 if $q$7 is not in any document span. The final attention is a gated mixture,

$q$8

where $q$9 is a small learnable gate. All other heads, or non-query rows, retain the original attention $d_j$0 (Tang et al., 23 Jun 2026).

At inference time, the LLM is discarded and only $d_j$1 is used as a standard dual-encoder retriever. This separation between training-time supervision and inference-time deployment is a defining property of the method (Tang et al., 23 Jun 2026).

3. Objective function and gradient pathway

After injecting the modified attentions $d_j$2 into the frozen LLM’s forward pass, DREAM computes the standard next-token cross-entropy loss over the target passage tokens $d_j$3:

$d_j$4

where $d_j$5 denotes the frozen LLM parameters (Tang et al., 23 Jun 2026).

The crucial mechanism is the gradient flow:

$d_j$6

Only the retriever parameters $d_j$7, along with the gate $d_j$8 and temperature $d_j$9, are updated; the judge LLM remains frozen throughout training (Tang et al., 23 Jun 2026).

The paper argues that this configuration yields stable supervision. In the ablations, keeping the LLM frozen yields better retriever performance than also fine-tuning LoRA adapters in the judge, and the stated interpretation is that a fixed judge provides a stable supervision signal (Tang et al., 23 Jun 2026). This suggests that DREAM depends less on co-adaptation between retriever and generator than on preserving a stable mapping from retrieval-conditioned attention to autoregressive loss.

4. Experimental configuration

Training uses Wikipedia chunked into sets of $$e_q = \frac{f_\phi(q)}{\|f_\phi(q)\|_2}, \qquad e_{d_j} = \frac{f_\phi(d_j)}{\|f_\phi(d_j)\|_2}.$$0 passages, with one chosen as the target and paired with a synthetic query generated by Qwen3-14B. The embedding backbones are Qwen2.5-0.5B, Llama-3.2-1B, and Llama-3.2-3B, with an additional 8B-scale scaling check. The frozen judge LLM is Llama-3.1-8B-Instruct (Tang et al., 23 Jun 2026).

The baselines are BM25, InfoNCE, RePlug, and Revela. BM25 serves as the lexical baseline. InfoNCE is contrastive training on the same data using positives versus in-set negatives. RePlug is described as LLM-likelihood distillation with scores outside the forward pass. Revela is joint LM+retriever training on raw chunks (Tang et al., 23 Jun 2026).

Evaluation is conducted on BEIR, described as 9 tasks for zero-shot retrieval, and RTEB, described as 14 tasks under the Massive Text Embedding Benchmark. The metric is NDCG@10 averaged over tasks (Tang et al., 23 Jun 2026).

Method	BEIR average NDCG@10	RTEB average NDCG@10
BM25	0.4122	0.3176
InfoNCE	0.2993 / 0.3268 / 0.3339	0.2950 / 0.3658 / 0.3405
RePlug	0.2593 / 0.2535 / 0.2705	0.2782 / 0.2855 / 0.3250
Revela	0.4011 / 0.4075 / 0.4315	0.4107 / 0.4499 / 0.4945
DREAM	0.4163 / 0.4888 / 0.5074	0.4788 / 0.5514 / 0.5892

The three values in each multi-scale row correspond to 0.5B, 1B, and 3B embedding backbones, exactly as reported in the paper (Tang et al., 23 Jun 2026).

5. Empirical performance and scaling behavior

DREAM consistently outperforms all baselines at each embedding scale. On BEIR, its gains over Revela range from +0.015 up to +0.081. On RTEB, the gains range from +0.068 to +0.102 (Tang et al., 23 Jun 2026). These margins are notable because Revela is also an LLM-supervised baseline, which means the reported advantage is not merely over lexical or contrastive retrieval but over an alternative form of language-model-based supervision.

The strongest average scores reported for the main scaling sweep are 0.5074 on BEIR and 0.5892 on RTEB with the 3B retriever. The 8B retriever scaling check yields 0.5531 on BEIR and 0.6417 on RTEB, which the paper describes as competitive with other off-the-shelf 7–8B models (Tang et al., 23 Jun 2026).

This performance profile suggests that DREAM benefits from retriever scale while retaining its training paradigm. The article’s evidence for that claim is empirical rather than theoretical: the same autoregressive supervision scheme continues to improve results from 0.5B through 3B and remains effective in the 8B-scale check (Tang et al., 23 Jun 2026).

6. Analyses and ablations

The head-selection results show that attention-head choice is not incidental. Injecting similarity into randomly chosen heads yields almost no retrieval improvement, with NDCG@10 approximately 0.06, whereas selecting the top 16 query-focused retrieval heads yields the full performance (Tang et al., 23 Jun 2026). This directly supports the paper’s claim that supervision must be routed through retrieval-relevant internal circuitry of the frozen LLM.

The number-of-heads ablation further refines this point. Performance rises from using the Top 1 head up to Top 16, then degrades when including weaker heads such as Top 32 and Top 64 (Tang et al., 23 Jun 2026). A plausible interpretation is that overly broad intervention dilutes the retrieval signal with less query-focused attention patterns.

Candidate set size also matters. Increasing $$e_q = \frac{f_\phi(q)}{\|f_\phi(q)\|_2}, \qquad e_{d_j} = \frac{f_\phi(d_j)}{\|f_\phi(d_j)\|_2}.$$1 from 4 to 16 improves performance, and the paper attributes this to sharper implicit competition among candidates (Tang et al., 23 Jun 2026). Because the document weights are normalized, a larger candidate pool produces a denser competitive distribution, strengthening the implicit-negative effect described in the training objective.

The embedding-geometry analysis compares DREAM with RePlug and Revela. DREAM achieves better uniformity of embeddings, meaning less collapse, while maintaining strong alignment of query-positive distances, despite never using an explicit contrastive loss (Tang et al., 23 Jun 2026). This is a notable empirical finding because it shows that contrastive geometry can emerge from autoregressive supervision when the retrieval scores are coupled to attention.

7. Limitations, scope, and future directions

The paper identifies compute overhead as a primary limitation: training requires running the full LLM forward pass over all $$e_q = \frac{f_\phi(q)}{\|f_\phi(q)\|_2}, \qquad e_{d_j} = \frac{f_\phi(d_j)}{\|f_\phi(d_j)\|_2}.$$2 candidates plus query and target, making it more expensive than pure contrastive pre-training (Tang et al., 23 Jun 2026). Head-selection dependence is another limitation, because success hinges on reliably identifying retrieval-relevant heads in the frozen LLM (Tang et al., 23 Jun 2026).

The reported experiments are focused on Wikipedia chunks and Qwen/Llama backbones. The paper states that applying DREAM to very large corpora or specialized domains may require further adaptation (Tang et al., 23 Jun 2026). This constrains how broadly the reported results can be generalized without additional evidence.

The future directions explicitly proposed include automatic or dynamic head selection, reduction of per-step LLM cost such as via distillation, integration with retrieval-generation loops, and combining autoregressive and contrastive objectives for even stronger retrievers (Tang et al., 23 Jun 2026). Taken together, these directions position DREAM less as a replacement for all retriever training paradigms than as a framework for label-free supervision using the predictive structure already present in frozen LLMs.

In summary, DREAM demonstrates that autoregressive next-token prediction can serve directly as a supervision signal for dense retrieval when retriever scores are injected into selected attention heads of a frozen LLM and trained through the resulting cross-entropy gradients (Tang et al., 23 Jun 2026). The resulting system is a standalone dual-encoder retriever at inference time, but one whose training signal is derived from retrieval-conditioned language modeling rather than explicit relevance annotation or hard-negative mining.