48小时单GPU训练DistilBERT！这个检索模型轻松达到SOTA

©PaperWeekly 原创 · 作者 | Maple小七

单位 | 北京邮电大学

研究方向 | 自然语言处理

论文标题: 

Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

收录会议：

SIGIR 2021

论文链接：

https://arxiv.org/abs/2104.06967

代码链接：

https://github.com/sebastian-hofstaetter/tas-balanced-dense-retrieval

基于 BERT 的稠密检索模型虽然在 IR 领域取得了阶段性的成功，但检索模型的训练、索引和查询效率一直是 IR 社区关注的重点问题，虽然超越 SOTA 的检索模型越来越多，但模型的训练成本也越来越大，以至于要训练最先进的稠密检索模型通常都需要 8×V100 的配置。而采用本文提出的 TAS-Balanced 和 Dual-supervision 训练策略，我们仅需要在单个消费级 GPU 上花费 48 小时从头训练一个 6 层的 DistilBERT 就能取得 SOTA 结果，这再一次证明了当前大部分稠密检索模型的训练是缓慢且低效的。

绪言

在短短的两年时间内，当初被质疑是 Neural Hype 的 Neural IR 现在已经被 IR 社区广泛接受，不少开源搜索引擎也逐渐支持了基于 BERT 的稠密检索（dense retrieval），基本达到了开箱即用的效果。其中，DPR 提出的 是当前最主流的稠密检索模型，然而众所周知的是， 的可迁移性远不如 BM25 这类 learning-free 的传统检索方法，想要在具体的业务场景下使用 并取得理想的结果，我们通常需要准备充足的标注数据进一步训练检索模型。

因此，如何高效地训练一个又快又好的 一直是 Neural IR 的研究热点。目前来看，改进 主要有两条路线可走，其中一条路线是改变 batch 内的样本组合，让模型能够获取更丰富的对比信息：

• 优化模型的训练过程：这类方法的代表作是 ANCE 提出的动态负采样策略，其基本思路是在训练过程中定期刷新索引，从而为模型提供更优质的难负样本，而不是像 DPR 那样仅从 BM25 中获取负样本。在此基础上，LTRe 指出目前的检索模型其实是按 learning to rank 来训练的，因为训练过程中模型仅能看到一个 batch 内的样本，但如果我们只训练 query encoder，冻结 passage embedding，我们就可以按照 learning to retrieve 的方式计算全局损失，而不是仅计算一个 batch 的损失。除此之外，RocketQA 提出了 Cross Batch 技巧来增大 batch size，由于检索模型采用对比损失训练，因此理论上增大 batch size 带来的基本都是正收益。

然而，这三种策略都在原始的 的基础上增加了额外的计算成本，并且实现都比较复杂。除此之外，我们也可以利用知识蒸馏（knowledge distillation）为模型提供更优质的监督信号：

• 优化模型的监督信号： 我们可以将表达能力更强但运行效率更低的 或 当作 teacher model 来为 提供 soft label。在检索模型的训练中，知识蒸馏的损失函数有很多可能的选择，本文仅讨论 pairwise loss 和 in-batch negative loss，其中 in-batch negative loss 在 pairwise loss 的基础上将 batch 内部其他 query 的负样本也当作当前 query 的负样本，这两类蒸馏 loss 的详细定义后文会讲。

本文同样是在上述两个方面对 做出优化，在训练过程方面，作者提出了 Balanced Topic Aware Sampling（TAS-Balanced）策略来构建 batch 内的训练样本；在监督信号方面，作者提出了将 pairwise loss 和 in-batch negative loss 结合的 dual-supervision 蒸馏方式。