DREAM: Embeddings de Recuperação Densa via Modelagem Autorregressiva

Resumo

Modelos de embeddings de recuperação densa são um componente fundamental dos sistemas modernos de IA baseados em recuperação. A maioria dos recuperadores densos é treinada com objetivos contrastivos, que exigem pares de documentos positivos e negativos rotulados, muitas vezes caros e difíceis de obter. Neste trabalho, investigamos se o objetivo de predição autoregressiva do próximo token de um modelo de linguagem de grande porte (LLM) pode fornecer supervisão para a recuperação densa. A intuição é simples: se um documento contém informações relevantes para uma consulta, condicionar esse documento deve tornar mais fácil para o LLM prever a saída alvo. Um desafio chave é que a perda de predição do próximo token é calculada dentro do LLM, enquanto o recuperador é um modelo de embedding separado. Para enfrentar esse desafio, propomos o DREAM (Embeddings de Recuperação Densa via Modelagem Autoregressiva), que injeta pontuações de similaridade consulta-documento geradas pelo recuperador em cabeças de atenção selecionadas de um LLM congelado. Durante o treinamento, essas pontuações determinam quanta atenção cada documento candidato recebe enquanto o LLM prevê a saída alvo. A perda de predição resultante fornece gradientes para o treinamento do recuperador através do mecanismo de atenção. Avaliamos o DREAM nos benchmarks de recuperação BEIR e RTEB usando backbones de embedding variando de 0,5B a 3B parâmetros. O DREAM supera consistentemente as linhas de base existentes em diferentes escalas de modelo. Esses resultados demonstram que o DREAM fornece uma abordagem promissora para treinar recuperadores densos por meio de modelagem autoregressiva.

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.