论文标题

RocketQAv2: A Joint Training Method for Dense Passage Retrieval and Passage Re-ranking

论文链接：

https://arxiv.org/pdf/2110.07367.pdf

引言

信息检索是当今时代人们获取知识的重要途径，信息检索的核心问题是，对于用户提出的问题，如何快速、准确地从海量文本中找出与该问题相关的文档或文字片段（答案）。现有的信息检索系统通常包括两部分：召回（检索）阶段和精排阶段。这种“召回-精排”的两阶段检索pipeline已经被应用在多种下游任务中，如问答系统、对话系统和实体链指等，同时在工业界也有着非常广泛的应用。

在两阶段pipeline中，检索和精排共同影响着最终的检索效果。基于这种关系，最近已经有一些工作对检索和精排模型之间的关系进行研究，它们指出，这两个模型可以基于它们之间的依赖关系进行分别训练，从而达到更好的效果。那么，顺着这个思路，我们是否可以直接对这两个模型进行联合训练，通过一次训练过程，同时得到处于最优性能的检索和精排模型呢？

在这篇EMNLP 2021的工作中，我们介绍了一种对稠密段落检索和基于预训练语言模型的段落精排模型的联合训练方法RocketQAv2，提出了一种两个模型统一的训练方式，并提出动态列表式蒸馏（dynamic listwise distillation）方法和融合数据增强（hybrid data augmentation）方法，提升了检索模型和精排模型的效果，使其在不损失检索效率的同时，拥有更高的准确率，在实时检索场景下具有更高的可用性和可扩展性。

背景与动机

段落检索（passage retrieval）任务旨在从一个庞大的文本段落库中，检索出和查询相关的段落。近年来，稠密段落检索已经成为了段落检索中的一种重要方法，它的核心思想是把文本查询（query）和文本段落（passage）分别编码成为低维的向量，通过计算向量相似度的方法，来估计query和passage间的相关性。段落精排（passage re-ranking）是段落检索的后续任务，通常通过对第一阶段的段落检索结果进行重排序，通过对query和passage进行更深层次的交互，来提高最终的检索效果。本文我们主要关注基于预训练语言模型的稠密段落检索以及段落精排。

随着基于预训练语言模型的稠密段落检索和段落精排受到学术界和工业界越来越多的关注，作为一个检索pipeline的两部分，段落检索和段落精排模型间的依赖关系也在很多研究中提及。

举例来讲，通常基于交互式模型（cross-encoder）的段落精排模型比基于双塔模型（dual-encoder）的段落检索模型具有更强的排序能力。受限于计算效率，cross-encoder通常只被应用于精排阶段，而dual-encoder由于Faiss等检索工具的出现，拥有更高的大规模检索效率，通常被应用于检索（召回）阶段。那么，我们可以蒸馏能力更强的cross-encoder的知识到dual-encoder中，从而提升dual-encoder的能力。同时，作为精排模型的cross-encoder需要根据一阶段检索的结果来构造训练数据并训练，否则由于检索阶段性能提高给候选集合带来更多的强负例段落，可能会导致精排模型性能处于次优状态。

虽然利用检索模型和精排模型的依赖关系进行训练，可以获得更高的pipeline检索性能，然而仍然需要对检索模型和精排模型进行分别甚至迭代的训练，这样并不是一个非常高效的方式。那么很自然地，一个直接的想法是对这两个模块进行联合的训练，让它们在训练可以互相提升，通过一次训练就分别得到两个处于最优状态的模型。

然而，要实现联合训练存在着比较大的挑战。主要挑战在于现有的稠密检索模型和精排模型的训练方式不同，通常检索模型采用批次内负采样（in-batch negative sampling）的方式进行训练，而精排模型采用基于pointwise或pairwise的训练方式。由于训练方式不同，二者的训练数据构造方式不同、输入形式不同、产出打分的分布不同，所以无法直接进行联合训练。

基于此，如何设计合理的训练方式，目标函数以及数据构造方式是本文需要解决的关键问题所在。

方法简介

我们将从训练方式设计、训练目标构造、训练数据构造以及训练流程分别介绍所提出的RocketQAv2方法，并介绍RocketQAv2方法和前序工作RocketQA的主要提升点。

训练方式设计

为了实现检索模型和精排模型的联合训练，首先需要统一各自的训练方式。这里，我们使用统一的listwise方式对两个模型进行训练。基于listwise的训练方式，对同一个query下的所有passage样本相似度打分进行优化，其通常包含个正样本和个负样本,其中为一个query 对应的passage 样本总数。

统一的listwise训练方式，为实现联合训练奠定了基础，基于此，我们就得以设计联合训练的优化目标。

训练目标构造

基于listwise的训练方式，我们提出了动态listwise蒸馏（dynamic listwise distillation）的方法，它可以同时更新检索模型和精排模型的参数，并且通过动态蒸馏的形式传递两个模型的信息。

我们做如下的形式化定义：给定一个query 和其相关的passage样本列表，我们可以得到基于dual-encoder的检索模型和基于cross-encoder的精排模型分别对与对应中passage的相关性打分，和。之后，我们用lisewise的方式对它们进行正则化，进而得到query与对应候选passage的相关性分布:

我们的主要想法是自适应地减少两个分布的差异，从而实现检索模型和精排模型的互相提升。因此，我们通过最小化两个相关性分布和的KL散度来进行优化：

另外，我们还为联合训练提供了一个有监督的指导，采用一个交叉熵的损失函数，基于中的passage来优化精排模型：

其中和代表中的正样本和负样本passage。

最后，我们把两个损失函数结合起来，得到最终的损失函数：

下图对dynamic listwise distillation进行了细节解释：

训练数据构造

统一了两个模型的训练方式之后，需要对训练数据构造方式也进行统一，即为query 产出候选passage列表。我们希望候选列表包括多样化且高质量的passage，这样可能可以更好地表示整个语料库中的passage分布。因此，我们设计了一种融合数据增强方法（hybrid data augmentation），具体细节如下图所示：

融合数据增强的训练数据构造方式分为两部分，分别是非去噪样本和去噪样本。首先，我们使用了RocketQA中训练的检索模型（retriever），在语料库中检索回query的top-的passage作为样本备选集。由于标记数据中只包含正例passage，负样本需要通过设计采样策略获得。

对于非去噪样本，我们在样本备选集中随机采样条passage作为训练数据中的强负例。

对于去噪样本，我们使用RocketQA中训练的精排模型（re-ranker）对top-的passage进行打分，取得分最高且大于给定阈值的passage作为去噪正样本，并从前往后取个打分小于阈值的文档作为去噪负样本，从而构造出弱监督的去噪样本。

我们融合了这两部分训练数据作为最终的训练数据。

训练流程

下图是RocketQAv2的整体训练流程：

首先，我们使用RocketQA训练的retriever和re-ranker通过融合数据增强构造训练数据，并用它们进行模型的参数初始化。之后，基于动态listwise蒸馏来对retriever和re-ranker进行联合训练。经过训练阶段后，它们可以直接应用于两阶段的检索pipeline。

需要注意的是，方法中使用的RocketQA训练的retriever或者re-ranker为原文中步骤1训练的版本，非全步骤训练的强模型，且可以使用其他检索模型替代。

RocketQAv2 VS RocketQA

如上所述，RocketQAv2延续了部分RocketQA的网络结构和关键的训练技巧，并做出了重要改进，实现了检索模型和精排模型的联合训练。

在动态listwise蒸馏中，RocketQAv2设计了统一的listwise优化方式对两个模型进行优化，相比于RocketQA的硬标签蒸馏，RocketQAv2使用了软标签蒸馏的的方式来实现两个模型的互相提升。这种蒸馏机制带来了更简化的训练流程，提升了训练效率，同时也带来了端到端训练整个检索架构的可能性。

实验

数据集

我们使用了Google的Natural Questions和微软的MSMARCO两个大规模公开数据集进行实验。

检索（召回）性能

我们在两个数据集上进行了检索效果的测试并和现有的baseline进行了比较。首先可以看到，RocketQAv2的检索性能优于绝大部分的baseline，尤其是相比于RocketQA，在提高训练效率的情况下，还带来了性能的提升。PAIR是我们ACL 2021 Findings上的一个同期工作，它使用了额外数据进行了预训练，而RocketQAv2完全基于数据集内的数据，没有引入数据增广，同时整个流程也更高效。

精排性能

我们在MSMARCO数据集上，基于不同retriever进行了精排效果的测试，并和现有的baseline进行比较。

可以看到，RocketQAv2的re-ranker明显优于其他baseline模型，证明了精排效果从整个joint training的过程中得到了提升。

我们还尝试了基于RocketQA的retriever的检索结果，分别使用RocketQA和RocketQAv2的re-ranker进行精排，发现RocketQAv2的re-ranker精排性能领先0.9个百分点，这也证明了经过联合训练的re-ranker的能力能力更强。

另外，我们把retriever换成基于BM25的方法，最终的精排效果也优于其他使用BM25 retriever且base级别的精排模型，甚至不逊于large级别精排模型的性能。

分析

我们针对多个角度进行了更进一步的实验，并进行了分析。

动态蒸馏还是静态蒸馏？

我们使用一个训练好且固定参数的精排模型作为教师模型，在训练中进行静态的蒸馏，来查看检索模型的性能变化。我们发现使用了静态蒸馏后，模型性能有了比较大的下滑，证明了动态蒸馏的重要性。

Listwise还是Pointwise？

我们替换了监督信号的listwise指导方式，改为pointwise的形式，这样训练的时候候选passage列表就不再以相似性分布的形式存在，而是独立优化。可以看到替换了listwise优化方法后，检索效果也有了下降，证明了统一的listwise训练方式的有效性。

去噪样本的影响

我们在融合数据增强过程中，去掉了去噪样本部分的训练数据，这样同样会到来部分性能下滑。原因可能是训练样本的多样性降低，导致候选passage列表不能很好的表示passage语料库的分布。

强负例个数的影响

由于每个候选列表中包含一个正样本和个负样本，所以强负例个数也代表这候选列表的大小。可以看到，随着每个query候选样本的个数的增加，开始阶段检索和精排的性能都在逐渐上升，当样本数量超过128时，涨幅就变得很小。

总结

本文提出了一种对稠密段落检索和基于预训练语言模型的段落精排模型的联合训练方法（RocketQAv2），利用动态listwise蒸馏和融合数据增强的方法，高效且有效地实现了dense retriver和bert-based re-ranker的联合训练，是基于RocketQA的一个重要改进。