DREAM:通过自回归建模实现密集检索嵌入
DREAM: Dense Retrieval Embeddings via Autoregressive Modeling
摘要
密集检索嵌入模型是现代基于检索的AI系统中的重要组成部分。大多数密集检索模型都是通过对比约束来训练的,这种训练方式需要标注好的正样本和负样本文档对,而获取这些样本往往非常困难且成本较高。在本研究中,我们探讨大型语言模型的自回归下一个词预测机制是否能够作为密集检索的引导机制。其逻辑很简单:如果某个文档包含与查询相关的信息,那么以该文档为条件的话,语言模型就能更容易地预测出目标输出。一个主要挑战在于,下一个词预测的损失计算是在语言模型中进行的,而检索器则是一个独立的嵌入模型。为了解决这一问题,我们提出了DREAM(通过自回归建模的密集检索嵌入)方法,该方法将检索器生成的查询-文档相似性分数引入到冻结的语言模型的注意力模块中。在训练过程中,这些分数决定了每个候选文档所获得的注意力权重,同时语言模型则负责预测目标输出。最终得到的预测损失可以通过注意力机制为检索器的训练提供梯度信息。我们在BEIR和RTEB等检索基准上对DREAM进行了测试,所使用的嵌入模型参数范围从0.5B到3B不等。结果显示,DREAM在不同模型规模下都优于现有的基线方法。这些结果表明,DREAM是一种有效的手段,可以通过自回归建模来训练高效的密集检索模型。
English 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.