DREAM: Dense Retrieval Embeddings via Autoregressive Modeling

Abstract: 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 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.

• The paper introduces a novel dense retriever training method that leverages next-token prediction losses from a frozen LLM to eliminate the need for contrastive pairs.
• It incorporates candidate score injection into selected query-focused attention heads, thereby enhancing retrieval efficacy across benchmarks like BEIR and RTEB.
• Empirical evaluations demonstrate improved NDCG@10 and superior embedding uniformity, affirming the scalability and robustness of the autoregressive approach.

DREAM: Dense Retrieval Embeddings via Autoregressive Modeling

Motivation and Background

Dense retrieval models are a crucial driver of retrieval-augmented generation and agentic search in modern LLM architectures. Traditionally, dense retrievers are trained with supervised contrastive objectives over query-positive-negative document pairs, relying on curated relevance annotations and hard-negative mining. This paradigm faces significant practical bottlenecks: relevance labels are labor-intensive, negative sampling may pollute supervision with false negatives, and the transferability of positive/negative designations is dataset-dependent.

Autoregressive next-token prediction (NTP) objectives, integral to LLM pretraining, offer an alternative avenue for supervision. The intuition underlying DREAM is that models conditioned on relevant documents should exhibit improved target passage generation, reflected in reduced NTP loss. Unlike contrastive objectives, NTP supervision requires no explicit relevance labels or negative sampling.

Methodological Formulation

DREAM introduces an architecture wherein a dense retriever is trained via gradients propagated from NTP loss in a frozen LLM. Candidate document scores computed by the retriever modulate selected attention heads in the LLM during prediction, directly coupling retriever decisions to LLM performance.

Figure 1: Overview of DREAM, coupling retriever scores to attention heads within a frozen LLM to supervise training via next-token prediction loss.

Key methodological aspects include:

• Retriever Score Injection: Retriever similarities $s_\phi(q, d_j)$ are normalized to weights $p_\phi(d_j|q)$ and injected into designated attention heads at query-token positions. These heads are selected pre-training using a ranking based on their query-to-candidate alignment, following query-focused retrieval head identification [zhang2025queryfocused]. Only heads with evidential relevance for retrieval are chosen, maximizing informative supervision.
• Attention Modulation: For each selected head $h$, retriever weights assign attention distribution across documents, while intra-document token selection adheres to frozen LLM preferences. A learnable gate determines the impact of retriever modulation.
• Training Objective: Cross-entropy NTP loss on target passage tokens is backpropagated, providing gradients solely to the retriever via the modified attention $\alpha'_h$. This induces competition among candidates, obviating the need for explicit negatives; increasing weight on non-informative candidates is directly penalized.

Empirical Evaluation

DREAM is benchmarked across BEIR and RTEB using embedding backbones from 0.5B to 8B parameters, with candidate and target passages constructed from Wikipedia chunks and queries synthesized by Qwen3-14B.

Retrieval Performance:

• DREAM consistently outperforms BM25, InfoNCE, RePlug, and Revela retrievers at each backbone scale in NDCG@10 across both BEIR and RTEB.
• Gains over contrastive InfoNCE baseline (trained on same data/candidate pools) isolate advantages arising specifically from the autoregressive training signal, not from candidate selection artifacts.

Attention Head Analysis:

• Injecting scores into query-focused retrieval heads yields markedly higher retrieval performance versus random or middle-layer heads, confirming that the interface between retriever and LLM is critical to signal efficacy.

  Figure 2: Retrieval heads selected via query-focused ranking enable significantly stronger supervision compared to non-specialized heads.

• Increasing the number of selected heads improves retrieval up to a threshold (Top 16 heads), with performance deteriorating when diluted by weaker heads.

  Figure 3: Average NDCG@10 as a function of the count of selected heads; optimal performance at Top 16 retrieval heads.

Embedding Geometry:

• DREAM's retrievers yield embeddings with superior uniformity and competitive alignment versus RePlug and Revela, achieving a less collapsed, more discriminative embedding space without contrastive objectives.

  Figure 4: Embedding-space analysis of query-positive pairs: DREAM achieves better uniformity with strong alignment.

Ablation Studies:

• Keeping the judge LLM frozen provides more stable supervision than jointly updating LoRA adapters, confirming that supervision should be anchored to fixed model computation.

  Figure 5: Frozen judge LLM yields stronger retriever training signal than LoRA-updated judge.

• Training with larger candidate document sets (e.g., 16 per sample) benefits retrieval performance, consistent with stronger inter-document competition.

  Figure 6: Retrieval improves as candidate set size grows, saturating at model scales and candidate counts used in DREAM.

• DREAM scaling experiments at 8B backbone show competitive or superior retrieval performance to strong off-the-shelf embedding models, validating the scalability of the method.

Theoretical and Practical Implications

The DREAM framework demonstrates that direct supervision for dense retrieval can be achieved through autoregressive losses, circumventing the established reliance on contrastive pair construction. By leveraging attention head specializations within frozen LLMs, DREAM effectively aligns retriever optimization with end-task utility, as measured by NTP performance.

This suggests future directions for:

• Self-supervised Retrieval Training: Exploiting rich NTP signals for retrieval in domains lacking relevance annotations or with weakly curated corpora.
• LLM-augmented Retriever Design: Harnessing head-level analyses and modular attention injections to integrate retrievers tightly with LLM inference architectures.
• Model Scaling and Adaptivity: Extending to larger embedding backbones and dynamic head selection for adaptive retrieval in evolving corpus domains.

The method challenges the notion that contrastive pairwise objectives are necessary for dense retriever training, proposing that autoregressive modeling, when correctly interfaced, is an effective, scalable alternative.

Conclusion

DREAM operationalizes dense retriever training directly via autoregressive NTP losses in a frozen LLM, using strategic attention head modulation to propagate supervision. The approach is empirically validated across standard benchmarks and scales, outperforming both classical lexical baselines and contemporary LLM-supervised methods. Theoretical implications support a shift in training paradigms for embedding-based retrieval, emphasizing autoregressive objectives and informed attention interfacing as core mechanisms.