[论文笔记]Rethink Training of BERT Rerankers in Multi-Stage Retrieval Pipeline

引言

今天带来远古时代一篇关于Reranker模型的论文笔记——Rethink Training of BERT Rerankers in Multi-Stage Retrieval Pipeline

为了简单，下文中以翻译的口吻记录，比如替换"作者"为"我们"。

经过深度LM微调的重排序器(Reranker)基于丰富的上下文匹配信号估计候选项的相关性。同时，深度LM也可以用于改进搜索索引，构建具有更好召回率的检索器。人们期望将两者直接结合在一个流程中能有累积的性能提升。然而，在本文中，我们发现这种流行的重新排序器无法充分利用改进的检索结果。因此，我们提出了一种局部对比估计(Localized Contrastive Estimation)方法用于训练重排序器。

代码开源在： https://github.com/luyug/Reranker

1. 总体介绍

启发式检索器(例如 BM25)能够高效地生成初步候选结果集合，随后由一个或多个复杂的重排序器对最有前景的候选项进行重新排序。像 BERT 这样的神经语言模型对这一架构产生了重大影响，它们为启发式检索器提供了更有效的索引术语和术语权重，并为重新排序器提供了丰富的上下文匹配信号。直观地说，更好的初步排序可以为后阶段的神经重排序器提供更多相关文档，从而将它们提升到最终排序的顶部。

在理想情况下，神经重排序器能够识别其候选池中的相关文档。然而，仅仅通过将 BERT 重新排序器附加到有效的第一阶段检索器上并不能保证最终排序的有效性。改进的候选列表有时会导致较差的重新排序。当候选列表改进时，假阳性可能更难以识别，因为它们往往与真实阳性具有混淆特征。一个有区分能力的重排序器应该能够处理检索器结果的t顶部部分，并避免依赖这些混淆特征。

在本文中，我们引入了局部对比估计(Localized Contrastive Estimation,LCE)学习。我们通过从目标检索器的顶部结果中采样来限定(localize)负样本分布。同时，我们使用了一种对比损失形式，它惩罚由混淆特征生成的信号，防止重排序器崩溃。

2. 背景知识

检索过程的阶段化(分阶段)自然地引入了由于不同排序模型之间的效率与效果权衡：快速但准确度较低的模型(例如 BM25)从整个语料库中检索，而较慢但更准确的模型(例如 BERT)在(排名)顶部候选列表中进行排序优化。

预训练的深度语言模型(如 BERT)在重排序任务上展示了强大的监督迁移性能。近年来的流行研究通过对 BERT 进行二分类目标的微调，显示它显著优于早期模型。然而，本文质疑这种简单的范式是否足以充分发挥 BERT 的潜力，特别是在高性能深度检索器生成包含更难负样本的候选项时。

相对于二分类目标，考虑负样本的对比学习目标是一种替代方案。流行的噪声对比估计(noise contrastive estimation,NCE)损失计算正实例和多个负实例的分数，将其归一化为概率，并训练模型将更高的概率分配给正实例。负样本的引入防止了模型的崩溃。尽管对比损失在表示学习中已被广泛研究，但在训练深度语言模型重排序器方面的相关工作仍然较少。

3. 方法论

预备知识 我们的目标是训练一个 BERT 重排序器来对查询-文档对进行打分，
$$s=\text{score}(q,d)=v_p^{\top }\text{cls}(\text{BERT}(\text{concat}(q,d)))\text{\hspace{1em}(1)}$$
其中，cls 提取 BERT 的 [CLS]向量，$v_p$是一个投影向量。我们将广泛采用的训练技术称为 Vanilla(原始) 方法。它独立地对每个查询-文档对进行采样，并使用基于查询($q$)和文档$d$及相应标签（+/-）的二分类交叉熵进行计算，
$${\mathcal{L}}_v:=\{\begin{array}{ll}\text{BCE}(\text{score}(q,d),+)& d\text{ 是正样本}\\ \text{BCE}(\text{score}(q,d),-)& d\text{ 是负样本}\end{array}\text{\hspace{1em}(2)}$$
Vanilla 方法将重排序器的训练视为一般的二分类问题。然而，重排序器的性质独特；它处理的是检索器结果的顶部部分，每个结果可能包含许多混淆特征。重新排序器需要，

• 擅长处理检索器结果的顶部部分；
• 避免陷入具有混淆特征的匹配中；

为此，我们引入了局部对比估计损失。对比损失防止了模型的崩溃，而局部负样本则使重排序器关注于顶部检索结果。

从目标检索器获取局部负样本 给定一个目标初始阶段检索器和一组训练查询样本，我们使用检索器从整个语料库中检索，生成查询的文档排名集。对于每个查询($q$)，从排名前$m$的文档集 $R_q^m$ 中采样$n$​个非相关文档作为负样本。所有采样文档一起形成负训练集。

对比损失 在从目标检索器采样所有负样本后，我们为每个查询$q$形成一个组$G_q$，其中包含一个相关的正样本 $d_q^{+}$和从 $R_q^m$​ 中采样的非相关负样本。将 BERT 打分函数视为深度距离函数：
$$\text{dist}(q,d)=\text{score}(q,d)=v_p^{\top }\text{cls}(\text{BERT}(\text{concat}(q,d)))\text{\hspace{1em}(3)}$$
定义一个查询$q$​的对比损失为，
$${\mathcal{L}}_q:=-\log \frac{\exp (\text{dist}(q,d_q^{+}))}{{\sum }_{d\in G_q}\exp (\text{dist}(q,d))}\text{\hspace{1em}(4)}$$
重要的是，这里的损失和梯度不仅依赖于相关对，还依赖于检索到的负样本。这有效地帮助防止了模型陷入简单的混淆匹配中。

LCE 批量更新 将所有内容汇总在一起，我们可以将局部对比估计(LCE)损失定义为一批训练样本集合$Q$​的损失，如下所示：
$${\mathcal{L}}_{LCE}:=\frac{1}{|Q|}\sum _{q\in Q,G_q\sim R_q^m}-\log \frac{\exp (\text{dist}(q,d_q^{+}))}{{\sum }_{d\in G_q}\exp (\text{dist}(q,d))}\text{\hspace{1em}(5)}$$
与标准的NCE损失相比，LCE 使用目标检索器来构建负样本，并将学习重点集中在顶部部分，而不是随机采样的噪声负样本上。

4. 实验方法

数据集和任务 使用 MSMARCO文档排序数据集。该数据集包含 300 万篇文档，每篇文档包括 3 个字段(标题、URL 和正文)，大约 900 个单词。模型在 37 万训练对的训练集上进行训练。使用开发集进行分析。

初始阶段检索器 对四种初始检索器进行了实验：Indri、未调优的 BM25、调优后的 BM25(记作 BM25*)和 HDCT。Indri 的搜索结果来自 MSMARCO 组织者。使用 Anserini 工具包构建了 BM25 索引，并生成了两个搜索结果集，一个是工具包的默认 BM25 参数，另一个是工具包作者建议的调优参数集。HDCT 是通过 BERT 重新估计术语权重来增强文档搜索索引的最先进方法；将前100 个候选列表输入到重排序器中。

实现 将拼接的文档标题、URL 和正文的前 512 个查询输入到重新排序器中。重排序器使用 PyTorch 在混合精度下构建和训练，基于 Huggingface 的 BERT 实现。从目标检索器排名前 100 的文档中采样负样本。在 4 个 RTX 2080 Ti GPU 上训练，每个 GPU 批处理 8 个文档。训练了 2 个 epoch，学习率为 1e-5，热身部分为 0.1。

5. 文档排名表现

在表 1 中，总结了在 MSMARCO 文档排序开发集和评估集查询上的排序性能。这里，Vanilla 方法和 LCE 都使用了来自目标检索器的负样本。

发现 LCE 显著提高了所有检索器的性能。同时，我们看到使用 LCE 的增益随着检索器的性能增强而扩大，表明它能够在改进的候选列表中捕捉到更复杂的匹配，而不会被更难的负样本混淆。

我们进一步结合了模型集成。集成提交使用了一个由 BERT、RoBERTa和 ELECTRA训练的 LCE 集成模型，对 HDCT 排名前 100 进行重排序。这一提交获得了第一名，达到了当时最先进的性能。

6. 分析

LCE 样本数量的影响 在图 1 中，研究了在 LCE 损失中变化每个查询采样的文档数量的效果。样本数量的增加可以带来一些改进。

负样本限定的影响 LCE 从目标检索器检索到的排名靠前的文档中采样负样本。在这里，我们定量评估其重要性。我们将训练中用于负样本采样的检索器记作训练检索器，将测试中用于候选生成的检索器记作测试检索器。使用第 5 节中的所有重排序器对所有检索器生成的候选列表进行排序，并在图 2 中绘制结果热力图。我们在水平轴上绘制使用不同训练检索器的重新排序器，在垂直轴上绘制测试检索器。每个 4 × 4 的子网格对应一个训练策略。

LCE 方法和 Vanilla 方法都受益于负样本的限定。

当负样本未通过测试检索器而是通过较弱的训练检索器进行限定时，Vanilla 训练的重新排序器的性能会严重下降。同样，使用 LCE 损失训练的重排序器在限定的情况下也表现更好。LCE 损失可以为重新序器带来一定的适应性，使其在测试检索器与训练检索器不同的情况下具有更强的鲁棒性。

7. 结论

作者提出了局部对比估计学习方法，用于通过目标检索器选择负样本，并使用对比损失来惩罚与混淆特征的匹配。

总结

⭐ 作者提出了一种局部对比估计方法用于训练重排序器，即首先通过一阶段的检索器来选择负样本，然后与正样本做对比学习，有效防止模型训崩。