论文

梦想:通过自回归建模进行密集检索嵌入

DREAM: Dense Retrieval Embeddings via Autoregressive Modeling

上下文与知识检索增强

摘要

密集检索嵌入模型是现代基于检索的人工智能系统的基本组成部分。大多数密集的 检索器 都是通过对比目标进行训练的,这需要标记的正负文档对,而这些文档对通常成本高昂且难以获得。在这项工作中,我们研究了 大语言模型 (LLM) 的自回归下一个标记预测目标是否可以为密集检索提供监督。直觉很简单:如果文档包含与查询相关的信息,则对该文档的条件应该使 LLM 更容易预测目标输出。一个关键的挑战是下一个词元预测损失是在 LLM 内部计算的,而 检索器 是一个单独的嵌入模型。为了应对这一挑战,我们提出了 DREAM(通过自回归建模的密集检索嵌入),它将 检索器 生成的查询文档相似性得分注入到冻结 LLM 的选定注意头中。在训练期间,这些分数决定了 LLM 预测目标输出时每个候选文档受到的关注程度。由此产生的预测损失通过注意力机制为 检索器 训练提供梯度。我们使用 0.5B 到 3B 参数范围的嵌入主干在检索基准 BEIR 和 RTEB 上评估 DREAM。 DREAM 在不同模型规模上始终优于现有基线。这些结果表明,DREAM 提供了一种通过自回归建模训练密集 检索器 的有前途的方法。