DREAM: Dense Retrieval Embeddings via Autoregressive Modeling

2026-06-23

AI Summary

This work proposes a novel approach that integrates autoregressive language modeling with dense retrieval, circumventing the need for costly annotated positive and negative sample pairs traditionally required by contrastive learning. Specifically, the method injects query-document similarity scores produced by a retriever into designated attention heads of a frozen large language model (LLM) and leverages the LLM’s next-token prediction loss to backpropagate gradients for end-to-end training of the retriever—without explicit supervision from labeled relevance pairs. Evaluated on the BEIR and RTEB benchmarks, the approach consistently outperforms existing methods across embedding backbones ranging from 0.5B to 3B parameters, demonstrating both its effectiveness and scalability.

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

Problem

• dense retrieval
• autoregressive modeling
• large language models
• contrastive learning
• retrieval supervision

Innovation

• dense retrieval
• autoregressive modeling
• large language model
• attention mechanism
• self-supervised training