DREAM：透過自回歸建模的密集檢索嵌入

摘要

密集檢索嵌入模型是現代基於檢索的人工智慧系統中的基本組件。大多數密集檢索器是透過對比學習目標訓練的，這需要標註的正負文檔對，而這些標註資料往往成本高昂且難以取得。在本研究中，我們探討大型語言模型（LLM）的自回歸下一個詞預測目標是否能為密集檢索提供監督訊號。直覺很簡單：若某文檔包含與查詢相關的資訊，則以該文檔為條件應能讓LLM更易於預測目標輸出。一個關鍵挑戰在於，下一個詞預測的損失是在LLM內部計算的，而檢索器是獨立於LLM的嵌入模型。為解決此挑戰，我們提出DREAM（透過自回歸建模實現的密集檢索嵌入），該方法將檢索器生成的查詢-文檔相似度分數注入凍結LLM的特定注意力頭。在訓練過程中，這些分數決定了LLM在預測目標輸出時，每個候選文檔會獲得多少注意力。由此產生的預測損失通過注意力機制為檢索器訓練提供梯度。我們在檢索評測基準BEIR和RTEB上評估DREAM，使用的嵌入骨幹模型參數量從0.5B到3B不等。DREAM在不同模型規模下持續優於現有基準方法。這些結果顯示，DREAM為透過自回歸建模訓練密集檢索器提供了一種有前景的途徑。

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.