DREAM: Dense Retrieval Embeddings via Autoregressive Modeling

Traditional dense retrieval models rely on contrastive learning, which requires high-quality positive and negative document pairs. These pairs are often expensive to curate and prone to noise. This paper investigates whether the autoregressive next-token prediction (NTP) objective of a large language model (LLM) can provide a more effective, label-free supervision signal for training dense retrievers.

The authors propose DREAM (Dense Retrieval Embeddings via Autoregressive Modeling). Instead of treating the retriever and LLM as separate, disconnected components, DREAM uses the LLM as a judge. During training, the retriever's similarity scores for candidate documents are injected into specific, query-focused attention heads of a frozen LLM. As the LLM attempts to predict a target passage, the attention mechanism is modulated by these scores. If a document is relevant, it receives more attention, helping the LLM predict the target tokens more accurately. The resulting prediction loss provides gradients that directly update the retriever, effectively training it to prioritize documents that aid the LLM's generation.

DREAM consistently outperforms existing baselines, including contrastive methods (InfoNCE) and other LLM-supervised approaches (RePlug, Revela), across various embedding backbone scales (0.5B to 3B parameters). The authors demonstrate that the choice of attention heads is critical; injecting scores into heads already specialized for query-document retrieval yields significantly better performance than using random heads. This confirms that the LLM's internal attention mechanism can serve as a powerful, natural interface for guiding retrieval training.