O1 Embedder：让检索器思考后再行动

25年2月来自中科大和北京智源研究院的论文“O1 Embedder: Let Retrievers Think Before Action”。

大语言模型 (LLM) 的功能日益强大，彻底改变人们获取和利用信息的方式。值得注意的是，LLM 擅长执行细粒度数据表示，这有助于精确检索信息。它们还可以根据外部参考生成高质量的答案，从而产生有用的知识。最近推出的推理模型（如 OpenAI O1 和 DeepSeek R1）标志着又一次飞跃，凸显 LLM 在提供最终答案之前进行渐进式思考的能力。这一突破显著提高处理复杂任务（例如编码和数学证明）的能力。

本文旨在为检索模型开发类似的功能，这对解决该领域的关键挑战大有裨益，包括多任务检索、零样本检索以及需要对复杂关系进行深入推理的任务。基于此，提出 O1 Embedder，在对目标文档进行检索之前，会先为输入查询生成有用的想法。为了实现这一目标，首先设计一个数据合成工作流程，通过从 LLM 专家那里生成初步想法，然后使用一个检索委员会对其进行细化，为 O1 Embedder 创建训练信号。其次，优化训练过程，使预训练模型能够联合微调，通过行为克隆生成检索想法，并通过对比学习执行密集检索。该方法通过全面的实验进行评估，在 12 个流行数据集上取得实质性的改进，涵盖域内和域外场景。

信息检索 (IR) 是许多重要应用的基础，例如搜索引擎和问答系统 [23, 24]。最近，它引起更高的关注，因为它在用外部知识增强大语言模型 (LLM) 方面发挥着关键作用 [36, 67, 69]，这种范例称为检索增强生成 (RAG) [12, 65, 66]。在过去十年中，IR 技术取得巨大进步。一个重要的突破是密集检索，其中可以通过向量搜索有效地检索相关数据。随着该领域开源模型的流行 [40, 55, 61]，密集检索已成为现实中实现检索应用的首选。

密集检索

密集检索在基础模型和训练技术的进步推动下，在检索精度方面取得重大进展。早期的突破包括对密集检索的初步预训练模型（如 BERT 和 RoBERTa [8, 34]）进行微调，这些模型已经表现出与 BM25 等传统方法相比具有竞争力的性能。同时，由于采用多语言 [4, 20] 和多模态预训练模型 [56, 68]，密集检索的范围得到大幅扩展。更先进的训练策略的引入，如面向检索的自适应 [35, 53, 60]、难负样本挖掘 [62]、批次量扩展 [42] 和来自交叉编码器的知识蒸馏 [17]，不断促进密集检索性能的提高。

除了提高检索准确率外，开发用于通用检索应用的多任务检索器，也越来越受到重视。最近的研究表明，通过扩大训练规模 [45] 和模型架构 [41] 可以大幅提升检索器的泛化能力。基于这些启发，人们对预训练和微调任务进行显著扩展，出现一系列流行的通用检索器，如 BGE、E5 和 GTE [33, 52]。同时，人们还引入大语言模型 (LLM) 作为检索器的骨干，带来了检索性能的显著提升。例如，RepLLaMA 通过直接微调预训练的 Llama [55] 提出一个强大的密集检索器。Llama2Vec 通过结合对预训练的 Llama 的无监督自适应进一步增强 RepLLaMA [28]。 Promptriver [58] 建立在 RepLLaMA 之上，为检索模型提供遵循指令的能力。NV-Embed 和 ICL-Embedder 等方法通过使用大量微调数据进行持续训练，实现进一步的改进 [26, 29]。如今，基于 LLM 的检索器几乎主导 IR 相关评估的所有主要基准测试。

尽管如此，密集检索在现阶段仍然受到许多限制。首先，现有方法在未见过的场景中难以完成零样本检索任务，这些场景与源域有很大不同。例如，来自一般数据集训练的嵌入模型在应用于专门问题（如医疗或法律案件检索）时容易出现性能受限的情况 [31, 39, 50]。其次，现有模型不足以区分复杂的关系，因为它们无法直接从语义含义中识别出来。例如，检索计算机程序的有用代码片段，或检索多跳推理问题的证据 [19, 27, 32, 63]。

上述挑战可以从先前的深度推理过程中受益，而不是直接判断相关性。最近，随着推理 LLM 的引入，例如 OpenAI O1、O3 和 DeepSeek R1 [14]，在这个方向上取得显着的进步。具体来说，当出现复杂任务时，LLM 首先会被提示对问题产生长时思考。通过这一过程，LLM 可以逐步接近正确解决方案，从而最终产生高质量的答案。最近的研究将这一操作概念化为测试-时间规模化 [44]，这推动解决复杂问题（例如编码和数学证明）的重大改进。

LLM 的推理能力

大语言模型 (LLM) 的推理能力已通过模拟类似人类的问题解决过程的技术得到显著增强。该领域的一个重大突破是思维链 (CoT) [57]，它促使 LLM 通过将复杂问题分解为多个推理步骤来解决这些问题。在此基础上，自洽方法通过从 LLM 中抽取多条推理路径并通过多数表决选择最终答案来提高推理的鲁棒性 [9]。对于需要更多探索性推理的场景，思维树 (ToT) 方法 [64] 通过将问题解决过程构建为树来扩展 CoT。在每个节点，LLM 都会生成候选中间步骤，评估其可行性，并从死胡同回溯。进一步推进这一范式的是，思维图 (GoT) [2] 用有向无环图 (DAG) 取代树结构，使 LLM 能够根据需要合并或改进推理步骤。大语言模型 (LLM) 的推理能力或“三思而后行”工作流程代表了一种新范式，使其有别于传统语言模型。除了通常的扩展模型大小、数据集和训练计算的策略 [16, 22] 之外，推理计算规模化或测试-时间规模化[5, 43, 59] 成为推动 LLM 改进的另一个重要因素。最近具有推理能力的 LLM（例如 OpenAI 的 O1 和 O3、DeepSeek 的 R1 [14] 和 Gemini 2.03）显著增强并展示这种能力。这些模型在处理复杂问题时采用“慢思考”方法：它们不会立即提供答案，而是先生成冗长的结构化推理，然后再得出最终解决方案。这种方法使 LLM 在编码和数学证明等领域实现精英级的性能。

推理能力在应对传统检索方法带来的挑战方面也具有显著优势。然而，当前的嵌入模型主要侧重于生成有判别力的数据表示，这使得推理能力的开发在很大程度上尚未得到探索。

O1 Embedder，就是在目标文档检索中引入类似于 LLM 嵌入模型的慢速思考能力，如图所示：

密集检索基于查询和文档的嵌入相似性来衡量查询和文档之间的相关性。给定查询 𝑞 和文档 𝑑，使用嵌入模型（M）将它们编码为潜表示 𝒗_𝑞 和 𝒗_𝑑：𝒗_𝑞 ← M（𝑞），𝒗_𝑑 ← M（𝑑）。要从海量数据集 𝐷 中检索相关文档 𝑑∗，需要满足以下最近邻（NN）条件：𝑑∗ = argmax. {⟨𝒗_𝑞, 𝒗_𝑑⟩ | 𝑑 ∈ 𝐷}。

本文方法通过将思考操作引入嵌入模型来解决上述挑战。也就是说，嵌入模型 M 配备两个功能：思考 M.think(·) 和嵌入 M.embed(·)。对于输入查询 𝑞，嵌入模型首先生成满足查询信息需求的想法 (𝑡)：𝑡_𝑖 ← M.think(𝑞)，𝑖 = 1，…，𝑘。通过揭示与相关文档的关键语义匹配模式，生成的想法有望促进复杂查询的检索过程。在这里，总共针对查询独立生成 𝑘 个想法，这使得有用的模式能够得到全面覆盖。在嵌入模型 M 和聚合函数 Agg 的基础上，将查询及其想法联合转化为统一的嵌入，称为想法增强嵌入（𝒗ˆ_𝑞）：𝒗ˆ_𝑞 ←Agg.(𝑞, {𝑡_𝑖}𝑘; M.embed)。
因此，使用想法增强嵌入来计算查询和文档之间的相关性：⟨𝒗ˆ_𝑞, 𝒗_𝑑 ⟩。最后，问题被表述为模型 M 思维和嵌入能力的联合训练，从而优化端到端的检索性能。

O1 Embedder 的训练涉及两种类型的数据。一种用于嵌入功能，由查询及其相关文档组成，即 q-doc 元组。另一种用于思考功能，包括查询及其想法，即 q-thought 元组。与广泛存在的 q-doc 元组不同，现实中没有可用的 q-thought 元组。为了解决这个问题，提出一个数据合成流水线，利用 LLM 现成的推理能力来生成此类数据集。该方法遵循“探索-细化”工作流程，如图所示。首先，用 LLM 来探索给定查询 𝑞 的候选想法。为了方便从 LLM 进行正确的生成，系统提示采用以下模板制定，其中结合指令和示例：Prompt = Task：{Ins}；Example：{E}；Query：{q} 。

指令用于明确声明思考任务的需求；例如，“思考一个合理的响应来处理查询”。同时引入示例来演示理想输出的形式。在这里，从 q-doc 数据集的训练集中随机选择 𝑚 个样本：E = {Query：𝑞_𝑖，Response：𝑑_𝑖}_𝑚。

注：尽管查询 𝑞 的相关文档 𝑑 看起来像是一个微不足道的解决方案，但它不适合用作想法。这是因为生成的想法将进一步用于嵌入任务，其目标是基于想法增强嵌入来区分相关文档 𝑑。如果使用 𝑑（或其任何重新表述的版本），训练过程将被规避，最终导致嵌入任务崩溃。

由于 LLM 可能产生幻觉，生成的想法可能并不总是能提高检索性能。为了确保包含有用的想法，探索过程会重复多次，为给定的输入查询生成几个独立的想法。为了识别最有用的想法，引入质量控制机制来过滤生成的候选。具体来说，使用一组多样化的检索器，表示为 𝑅。
对于每个检索器 𝑟 ∈ 𝑅，计算相关文档 𝑑 和想法 𝑡_𝑖 之间的相似度得分：𝜎^𝑟（𝑡_𝑖，𝑑）。每个检索器选择相似度得分最高的想法，即 𝑡_𝑟∗ ← argmax（{𝜎^𝑟（𝑡_𝑖，𝑑）}_i = 1，…，k）。最后，进行多数投票过程以确定最有用的想法。从检索器获得最多提名的想法被选为最终结果：𝑡 ← voting{𝑡_𝑟∗}_𝑟 ∈ 𝑅。

通过将上述数据合成工作流程应用于现有的 q-doc 数据集：𝐷 = {(𝑞_𝑖,𝑑_𝑖)}_𝑁 ，可以获得由 q-thought-doc 三元组组成的想法增强数据集：𝐷ˆ = {(𝑞_𝑖, 𝑡_𝑖, 𝑑_𝑖)}_𝑁 ，它提供最佳检索效用的长时思考。

如下图是数据生产过程中概述的算法流程：对于数据集中的每个查询-文档对（𝑞_𝑖 ，𝑑_𝑖 ），算法首先从 𝐷 中采样 𝑚 个示例以创建提示，其中包含具体指令。然后，它通过使用采样示例格式化提示来生成 𝑘 个候选想法 𝑡_𝑗。接下来，对于每个检索模型 𝑟 ，它计算候选想法与文档 𝑑_𝑖 之间的相似度得分，选择得分最高的想法作为 𝑡_𝑟∗。最后，该算法通过多数投票机制，汇总这些顶级想法以确定最终想法 𝑡_𝑖 。然后构建增强数据集 𝐷ˆ，包括原始查询、生成的想法及其对应的肯定文档。

O1 embedder 建立在预训练的 LLM 之上，利用模型固有的生成和推理能力。此外，LLM 还显示出作为判别嵌入模型进行微调的强大潜力 [37, 70]。对预训练的 LLM 主干应用以下的多任务训练，该主干通过监督行为克隆建立其思维能力，并通过对比学习建立嵌入能力。

行为克隆。通过监督微调，训练基础模型为输入查询生成想法。给定 q-thought-doc 三元组的数据集：D^ = {(q_i, t_i, d_i)}_N，训练样本 x_i 用以下模板制定：x_i = q_i t_i ，其中 、、 分别是 token 查询、思想和生成完成的特殊 token。通过上述训练样本的制定，模型针对以下生成损失进行微调：

其中，从思考的开头开始，每个 token 的下一个 token 预测损失都最小化（即 𝑗 ≥ |𝑞_𝑖 |）。

对比学习。预训练的 LLM 也经过微调，根据其生成的嵌入将相关文档与查询区分开来。传统上，基于 LLM 的embedder 使用 token 进行文本嵌入。然而，在本文方法中， token 已被指定为指示生成过程的完成。因此，加入额外的嵌入任务可能会导致训练过程崩溃。为了避免两个并行训练任务之间的相互干扰，用另一个特殊 token 并将其附加到输入 𝑥 的末尾（在 之后）来计算文本嵌入：v_x ← LLM(𝑥; )[−1]。

这个简单的修改大大提高兼容性，这对于维持联合训练过程的成功运行至关重要。考虑到人们可能希望利用想法增强嵌入来处理复杂的检索任务，同时仍然依赖基本查询嵌入来处理简单的检索任务，并行生成两种形式的查询嵌入来识别相关文档 𝑑_𝑖。因此，执行以下复合对比学习，其中分别基于每个训练样本的查询和想法增强查询计算两个对比损失：

此处，𝑞ˆ_𝑖 表示想法增强查询：𝑞ˆ_𝑖 ← 𝑞_𝑖 + 𝑡_𝑖 ，而 𝐷′ 代表负样本集合，包括批次内负样本和由预训练 embedder 引入的难负样本。

联合训练。该模型经过训练以最小化生成损失和对比损失的线性组合。为了在下游场景中实现精确检索，必须使用较大的训练批次进行对比学习。然而，两个训练任务的本机并行运行需要大量 GPU 内存，这严重限制可实现的批次大小。为了应对这一挑战，提出一种内存高效的联合训练来处理这两项任务（如图所示）。具体来说，对于每个训练样本 𝑥_𝑖 = (𝑞_𝑖, 𝑡_𝑖, 𝑑_𝑖)，对其进行一次编码，并在两个任务之间共享编码结果。生成损失是根据来自 𝑡_𝑖 token 的每个输出嵌入计算的；而对比损失是根据来自 token 的输出嵌入得出的。这使得生成任务在与对比学习任务一起训练时，几乎不消耗额外内存，从而能够大幅增加批量大小。

训练的 O1 embedder M 通过思考和嵌入应用于检索任务。首先，促使 O1 embedder 针对输入查询产生多个想法，全面揭示查询的隐藏信息需求：𝑡_𝑖 ←M.think(q), 𝑖 = 1,…,𝑘。接下来，对经过想法增强的查询进行独立编码和聚合。在这里，简单地采用均值池化作为聚合函数，它会产生以下经过想法增强的嵌入：

最后，根据与 𝒗ˆ_𝑞 的嵌入相似性检索前 N 个文档 𝐷∗：𝐷∗ ← top-N({⟨𝒗ˆ_𝑞, 𝒗_𝑑⟩|𝐷})。

数据集。O1 embedder 由从 MS MARCO（段落检索）数据集 [1] 创建的想法增强查询训练而成。训练的模型基于域内和域外数据集进行评估。对于域内评估，用 MS MARCO (dev)、TREC DL19 [7] 和 TREC DL20 [6] 数据集。对于域外评估，结合来自 BEIR [47] 以下八个问答数据集，包括 SciFact [51]、TREC-COVID [50]、DBpedia [15]、NQ [25]、HotPotQA [63]、FiQA [39]、Touche [3]、FEVER [48]，以及一个流行的代码搜索数据集：CosQA [18]。所有这些数据集都包含非对称检索任务，其中查询和文档以非常不同的形式呈现。因此，通过适当的推理，可以更有效地识别查询和文档之间的关系。排除常见的解释数据集，例如 Quora，因为它们仅涉及简单的相似性比较。根据先前的研究 [38、58]，用 MRR@10 和 Recall@1k 作为 MS MARCO 相关任务的指标，并使用 NDCG@10 作为其他数据集的指标。对 分布外（o.o.d. ）数据集的评估严格遵守 BEIR 协议，该协议禁止针对特定任务进行微调。

基线。选择各种流行的检索器作为基线，例如常用的稀疏检索方法 BM25，以及 ANCE [62]、TAS-B [17]、coCondenser [10]、SimLM [53]，它们使用 MS MARCO 数据集对基于 BERT 预训练模型进行微调。还介绍基于 LLM 的方法，包括 RepLLaMA [54] 和 Promptriver [58]。RepLLaMA 对 MS MARCO 上预训练的 LLM 进行微调，从而在各种下游任务中实现卓越的检索性能。而 Promptriever 则以 RepLLaMA 为基础，增强模型的指令跟随能力，从而在专门提示的基础上进一步提高检索性能。这两种基于 LLM 的方法都是从同一个预训练主干网络 (Llama-2-7B) 进行微调的，作为默认设置，从而确保在模型规模方面进行公平的比较。注：在评估中排除其他几种流行的检索器，包括 BGE [61]、E5 [52]、M3 [4] 和最近基于 LLM 的方法，如 E5-Mistral [55]、ICL-Embedder [30] 和 NV-Embed [26]。这些模型使用的训练数据远远超过 MS MARCO（本文方法和所包含基线使用的唯一训练数据集），其中许多与 BEIR 上的评估任务有很大的重叠。这种重叠使得难以评估分布外 (o.o.d.) 设置中的零样本检索性能。此外，这些方法依赖于不同的训练数据集，这使得无法进行公平的比较。

实施细节。在数据准备阶段，用强大的开源 LLM：Llama-3.1-70B-Instruct4 来生成候选想法。用 BM25、BGE-EN-large-v1.5、GTE-large 和 Stella-EN-1.5B-v5 为这个检索委员会服务。评估主要基于 Llama-2-7B 主干 [49]，并在扩展研究中分析其他替代 LLM。训练过程遵循 RepLLaMA 的方案 [38]，其中 LLM 的所有投影层 (q_proj k_proj v_proj o_proj gate_proj down_proj up_proj) 都通过 LoRA 进行微调，等级设置为 32，alpha 设置为 64。用 BF16 进行训练，学习率设置为 1 × 10−4。训练过程在 8xA800 GPU 上进行，批量大小设置为 64（每台设备 8 个）。为每个查询引入 15 个难负样本。最大查询长度设置为 32，最大段落长度设置为 192。最大 token 数设置为 256，用于想法生成。