DocReLM：自建数据集训练检索模型和重排序模型提高 RAG 效果

背景

根据 Crossref 的数据，目前学术界已经积累了 1.4 篇文献，每年还以几百万的速度增加。这种惊人的增长速度，让研究人员跟踪业内动态变得越来越困难。因此，大家对高效的语义检索系统需求越来越强烈。

这类系统的开发面临非常多的困难：

• 首先，学术论文使用高度专业化的语言，需要深刻理解上下文才能提取关键信息。

• 其次，理解一篇论文不仅要阅读其内容，还要理解引用及其相互关系，这要求系统能够把握论文内引用间的联系。

• 最后，用户可能无法用精确关键词描述信息需求，系统需能洞悉用户意图，提供最相关文献。尽管已有尝试采用通用嵌入模型构建检索系统，但模型训练与实际搜索之间存在差距，导致性能不佳。

基于以上背景，今天这篇论文的作者推出了 DocReLM，一个利用语言模型强化的文档检索系统。

在 DocReLM 中，LLM 以两种独特方式发挥作用：

• 首先，作为自动数据标注器，生成高质量的训练数据以提升检索模型性能。使得密集检索器和重排模型在性能上超越了竞争对手。其他模型由于训练于通用语义相似性任务，性能不及 DocReLM 的模型。通过 LLM 生成的数据训练，DocReLM 更适合文档检索任务，同时领域专业知识也在这一过程中被融入模型。

• 另外，LLM 作为搜索 Agent，根据检索文档的文本信息和引用关系来优化论文选择。类似于人类通过阅读论文引用来寻找最相关文献的推理过程。LLM 理解上下文和引用关系的能力使其能够精准识别相关论文，从而提升系统性能。在自然科学领域，这一功能尤为有用，因为论文间高度互联，理解引用对于找到相关文献至关重要。这也是一种新颖的在检索系统中应用大型语言模型的方法。

系统架构

图片

DocReLM 系统整体结构如上图所示，主要分为三个部分：检索器、重排器和引用提取器。

• • 用户提出查询时，检索器利用嵌入技术迅速从文档库中筛选出候选段落。
• • 接着，重排器进一步精准排序这些段落。
• • 最后，引用提取器分析重排器提供的前 k 个结果，根据文档内容生成恰当的引用。

检索器

检索器的作用是高效地从庞大的文档库中提取精选文档集。

一般来说，检索器会用到两种模型：稀疏检索器和密集检索器。

例如，著名的 BM25 稀疏检索器利用词袋模型向量匹配候选项。而密集嵌入模型则将句子明确编码为密集向量，并通过余弦相似度来评估查询与段落嵌入之间的匹配可能性。

在 DocReLM 中，作者采用了密集检索器，并基于开源的 jina-embedding-v2-base 模型进行定制训练。这个检索器充当编码器角色，已在科学领域多种任务中得到应用。

通过对比学习，密集检索器生成的嵌入能够突出关键段落的细节，即便术语不同，也能连接语义相似的句子。然而，预训练的嵌入模型在文档检索任务中常常缺乏特异性，与特定领域的论文存在知识鸿沟。为此，作者引入大型语言模型自动用学术论文库标注伪查询，以此优化嵌入模型，使其更贴合学术领域的任务需求。

构建这一检索系统时，将长文档切分为段落，并将它们转换为嵌入向量，存储于 Faiss 向量数据库中。用户输入查询后，DocReLM 利用检索模型提取查询的语义嵌入，并基于文档和查询嵌入之间的余弦距离挑选候选段落。

重排器

快速检索之后，重排器被用来提升结果的准确性。虽然它通常更精确，但与检索器相比，重排器的处理速度较慢。这决定了它通常在检索后使用，检索器首先通过排除大部分不相关段落来缩小选择范围。

在此任务中，作者采用了交叉编码器，它处理由查询字符串和检索子系统提供的候选段落组成的串联序列，以[SEP]标记分隔。

通过在输入中整合查询和段落，重排器增强了两者的交互，从而提高了其表达力，并在效果上超越了检索器。交叉编码器的一个局限是它无法预处理文档库，需要对每个新查询进行实时推理，因此不能单独作为系统使用，它需要检索器提供一小批相关候选。

在 DocReLM 中，使用了开源的 XLM-RoBERTa-large 模型作为基础，并用作者的自定义数据进一步训练，采用 LCE 方法。检索器确定的前 200 个段落随后输入到重排器中，为每个串联对输出一个分数，然后根据这些分数对候选段落进行重新排序。

引用提取器

经过精确的段落检索和交叉编码器的重排，相关论文根据与查询的相关性被依次展示。

但是，通过一个额外步骤可以显著提升搜索结果的质量：DocReLM 系统中的引用提取器，它负责阅读段落内容并从中抽取最佳引用文献。

这种做法更贴近实际的搜索习惯。虽然检索器和重排器挑选出的论文可能并非最佳答案，但其引用列表中很可能包含了正确答案。为了提高搜索的准确性，引导大语言模型来处理检索结果，并从引用文献中识别出更加合适的论文。

具体操作中，通过在论文段落中嵌入参考文献的标识符来调整文本。当这段经过修改的文本被大型语言模型处理时，它会根据上下文信息抽取出这些标识符。如果模型认定段落本身足以回答问题，它将输出引用文献的标识符。系统将根据提供的标识符搜索相应的论文。

图片

如上图展示的，检索器迅速从完整文献库中筛选出 n 个候选文段。紧接着，重排器对这些文段进行排序。继而，引用提取器细读排名最前的 10 篇论文，并为每篇提取三条引用，这些引用被巧妙地嵌入到对应原文之后。

模型训练

作者分别对检索器和重排序模型进行了训练。

先利用对比学习在生成的训练集上训练检索器。随后利用这个检索器筛选出难度较高的负样本，用于对重排器进行更加精细的对比学习微调。

利用 LLM 生成训练数据

为了通过对比学习训练检索器和重排器，首先需要搜集一系列的用户查询和文档段落。尽管市面上有 MS MARCO 这类带有人工标注的数据集，但人工标注的成本非常高，使得获取领域特定的检索数据集不是很容易。

在本项研究中，运用 LLM 来从文档中生成伪查询。选用了 unarXive 上的论文来构建数据集，对他们的原始数据进行处理，剔除了语义上不相关的文本，并加入了数学 LaTeX 格式和引用标识符。这些引用标识符按照 Ref.X of IDFp 的格式进行编排，其中 IDFp 代表的是原论文的标识符。将整篇论文切分成多个段落，并动用 vicuna-7b-v1.5-16k 根据段落、标题和摘要为每个句子生成查询。借鉴思维链方法，首先采用零样本方法引导大语言模型 Agent 应用基于标题、摘要和内容草拟大纲，继而生成查询。

作者为实验建立了两个子集。其一是量子物理领域，包含 56927 篇论文，搜集了 2.8⁢𝑀 的训练数据。其二是计算机视觉领域，包含 37390 篇论文，搜集了 1.1⁢𝑀 的训练数据。这些数据被用于训练密集型段落检索器与重排器。

在对比训练过程中，检索器和重排器的负样本生成策略各有不同。检索器的负样本是通过随机抽样来获得的。在每次训练批次中，我们会抽取 N+1 组正样本对，每组正样本对将伴随 N 组负样本对。增加批次大小可以提升检索器的性能，因为它增加了引入困难负样本的可能性。为了在对比学习中获得更佳表现，采用了梯度缓存技术来扩大训练批次的规模。

至于重排器的训练，我们采用已训练好的检索器，按照文献[8]中提出的 LCE 方法来挑选负样本。对于每一个查询检索器会返回与之相关性最高的 N 个段落。这些段落与查询紧密相连，表明挑选出了除了正确答案之外最具迷惑性的错误选项。因此，这些误导性强的段落成为了训练重排器时使用的高质量困难负样本。

效果如何

为了检验 DocReLM 的效果，作者根据真实的研究场景构建了一个基准测试，邀请各领域的研究人员根据他们的研究方向提出查询，并附上论文作为标注的参考。为确保这些标注的准确性，研究人员必须提供调查研究中的论断作为依据。接着，选用 unarXive 的一个子集作为检索的文档库。最终的基准测试涵盖两个领域：量子物理领域包含 301 个查询和 56927 篇论文，计算机视觉领域则包含 204 个查询和 37390 篇论文。将 DocReLM 与 Google Scholar 进行了对比。此外，还进行了消融分析，以评估系统中三个独立组件的有效性。

图片

作者对比了三种基础检索模型：BM25、OpenAI 的 ada-002 以及 jina-base-v2。它们的准确度展示在商标中。稀疏型检索器 BM25 相较其他模型表现欠佳，其基于词频的方法不足以捕捉任务中的语义关联。DocReLM-retriever 在两个领域中均超越了这些模型。特别是在计算机视觉领域，DocReLM-retriever 的表现显著超越其他模型，其 top-1 准确度比其他最佳嵌入模型高出 5.89%，而 top-5 和 top-10 的准确度分别提高了 13.72%和 19.13%。值得一提的是，在计算机视觉领域，所有三种模型都展现出了较高的平均准确度。而在量子物理领域，嵌入模型的表现相近，top-1 结果并无太大差异；但在 top-5 结果上，作者的模型将准确度从 9.97%提升至 11.96%，超越了其他模型。

图片

为测试重排器的性能，挑选了三种不同的重排器进行比较：bge-reranker-large、Cohere 以及 DocReLM-reranker。这些重排器以 DocReLM-retriever 筛选出的前 200 个段落作为输入（如上表）。在计算机视觉领域，DocReLM-reranker 超越了其他模型，其 top-1 准确度比其他最佳模型高出 0.49%，而 top-5 和 top-10 的准确度分别提升了 1.97%和 2.94%。在量子物理领域，DocReLM-reranker 与 Cohere 表现相当，DocReLM-reranker 的模型在 top-20 的表现更佳，而 Cohere 在 top-1 和 top-10 的得分上更为突出。这些结果证明 DocReLM-reranker 在两个领域的效率均高于其他模型，凸显了重排器在提升检索系统准确度中的关键作用。

最后一项实验评估了引用提取模型的有效性。对比了开启和未开启引用提取功能的 DocReLM-reranker 的性能，结果列在上表中。在计算机视觉领域，引入了引用提取功能的模型 internLM 将 top-5 和 top-20 的准确度分别提升了 0.98%和 2.94%，证实了引用提取功能在提升检索系统准确度方面的有效性。在量子物理领域，同样的模型将 top-5 和 top-10 的准确度分别提升了 9.63%和 16.28%。

讨论

LLM 的精华提炼

神经密集型检索器和交叉编码重排器能够精准捕捉查询与文档间的语义联系，使得相关文档的检索更为精确。使用 LLM 生成的数据训练模型后，与基础模型相比，性能有了显著飞跃。这些数据由 LLM 生成，使得 LLM 的理解力得以凝练并注入到较小型的模型中，让通用模型能够适应特定任务。

LLM 与检索系统的融合创新

作者提出了一种新颖的大语言模型与检索系统集成方法。随着 LLM 的流行，人们一直在尝试将其与搜索引擎结合。这些尝试主要有两种方向：一是利用 LLM 扩展搜索查询，二是用 LLM 总结搜索结果。

而 DocReLM 则开辟了新径，让 LLM 理解搜索结果并继续搜索过程，推荐更佳的候选文档。此概念有进一步扩展的空间。该过程可以迭代执行，让 LLM 在引用网络中穿梭多个节点，直至锁定最相关的结果。这一任务要求 LLM 不仅要理解引用文献间的逻辑联系，还需具备足够的领域知识以精确提取引用。通过整合能够进行多次搜索并分析结果以优化搜索的 LLM，可以显著提升通用检索系统的性能。这种方法预示着检索系统未来发展方向的希望之光。

图片

图片

• 论文原文: https://arxiv.org/abs/2405.11461