DREAM : Embeddings de récupération dense par modélisation autorégressive

Résumé

Les modèles d'embedding pour le retrieval dense sont un composant fondamental des systèmes modernes de recherche basés sur l'IA. La plupart des retrieveurs denses sont entraînés avec des objectifs contrastifs, qui nécessitent des paires de documents positifs et négatifs étiquetées, souvent coûteuses et difficiles à obtenir. Dans ce travail, nous examinons si l'objectif de prédiction autorégressive du prochain jeton d'un grand modèle de langage (LLM) peut fournir une supervision pour le retrieval dense. L'intuition est simple : si un document contient des informations pertinentes pour une requête, le conditionnement sur ce document devrait faciliter la prédiction de la sortie cible par le LLM. Un défi clé est que la perte de prédiction du prochain jeton est calculée à l'intérieur du LLM, tandis que le retrieveur est un modèle d'embedding séparé. Pour relever ce défi, nous proposons DREAM (Dense Retrieval Embeddings via Autoregressive Modeling), qui injecte les scores de similarité requête-document générés par le retrieveur dans des têtes d'attention sélectionnées d'un LLM gelé. Pendant l'entraînement, ces scores déterminent l'attention que chaque document candidat reçoit lorsque le LLM prédit la sortie cible. La perte de prédiction résultante fournit des gradients pour l'entraînement du retrieveur via le mécanisme d'attention. Nous évaluons DREAM sur les benchmarks de retrieval BEIR et RTEB en utilisant des backbones d'embedding allant de 0,5B à 3B paramètres. DREAM surpasse constamment les lignes de base existantes à différentes échelles de modèles. Ces résultats démontrent que DREAM offre une approche prometteuse pour entraîner des retrieveurs denses via la modélisation autorégressive.

English

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.