百度提出动态自蒸馏方法，结合交互模型与双塔模型实现稠密段落检索

科技

2022-07-28 05:07

©PaperWeekly 原创 · 作者 | Maple小七

单位 | 北京邮电大学

研究方向 | 自然语言处理

论文题目：

ERNIE-Search: Bridging Cross-Encoder with Dual-Encoder via Self On-the-fly Distillation for Dense Passage Retrieval

论文链接：

http://arxiv.org/abs/2205.09153

Introduction

由无论是在工业界还是学术界，知识蒸馏 (Knowledge Distillation, KD) 都是一个很常见的模型压缩和涨点刷分手段，以 Hinton 范式为代表的传统知识蒸馏通常会假设 teacher 和 student 有着相同的模型结构，其中 teacher 通常有着更多的模型层数和更大的隐层维度。在知识蒸馏的过程中，蒸馏损失的设计通常与任务相关，但其基本思想是最小化 student 和 teacher 预测结果的差异。

一些早期的针对大规模预训练语言模型的知识蒸馏工作，比如 DistilBERT、TinyBERT 和 MobileBERT，通常会在最小化模型输出结果差异的基础上，加入额外的隐层特征一致性损失和注意力分布一致性损失，从而提升 teacher 模型的知识转化率，得到性能更优的 student 模型。

然而，在针对稠密检索模型的知识蒸馏的相关研究中，虽然目前已经有很多工作尝试过将表达能力更强的交互模型 (Cross-Encoder) 或延迟交互模型 (ColBERT) 蒸馏到推理速度更快的双塔模型 (Dual-Encoder) 中，实验结果也充分证明了知识蒸馏对于双塔模型训练的有效性，但是这些工作均采用的是跨结构知识蒸馏 (cross-architecture paradigm)，即 teacher (Cross-Encoder) 和 student (Dual-Encoder) 的模型结构有较大差异。

另外，模型结构的天然差异也导致了目前的工作均采用了仅蒸馏模型顶层 logits 输出的蒸馏损失，而没有考虑蒸馏 teacher 和 student 模型的隐层特征或注意力分布。

基于此，一个自然的问题是，当 teacher 和 student 的模型结构存在差异时，传统的知识蒸馏的效果是否会打折呢？如果会的话，我们应该如何去设计一个更有效的知识蒸馏范式呢？

针对这个问题，本文提出了两个面向稠密检索模型的蒸馏策略：交互蒸馏 (Interaction Disitllation) 和级联蒸馏 (Cascade Distillation)，如果将这两个蒸馏策略结合起来，就得到了该论文标题中提到的名词：动态自蒸馏 (self on-the-fly distillation)。由于是度厂出品，采用该蒸馏策略的训练得到的模型被命名为 ERNIE-Search，ERNIE-Search 在 MARCO 和 NQ 检索数据集上均取得了 SOTA 性能（虽然没出几天就被超越了）。

Methodology

在之前的工作中，Cross-Encoder 和 ColBERT 是两类常见的 teacher 模型结构，虽然前者表达能力更强，但由于其计算复杂度过高，人们通常采用样本间无交互的 MSE 或 MarginMSE 损失进行知识蒸馏，而后者与 Dual-Encoder 的接近程度更高，因此可以采用基于批内负样本的对比损失蒸馏策略。关于这部分背景知识更详细的介绍，读者可以参考 TAS-B 中针对多 teacher 知识蒸馏的讲解。

从数学层面上来讲，Cross-Encoder 是完全可以被 Dual-Encoder 所拟合的（可参考特征式匹配与交互式匹配有多大差距？中给出的数学证明）。但由于归纳偏差的存在，Cross-Encoder 中存储知识的结构与 Dual-Encoder 是完全不一样的，比如 Dual-Encoder 对于语义相关性的建模就是简单的点积，对语义相关性有着更强的连续性先验，而 Cross-Encoder 对于语义相关性的建模是蕴含在 token 级交互信息中的，我们无法知晓 Cross-Encoder 内部具体是通过什么机制来建模相关性的。

2.1 Interaction Distillation

既然我们难以将 Cross-Encoder 中的知识完美迁移到 Dual-Encoder 当中，那如果我们从一开始就让 Cross-Encoder 和 Dual-Encoder 共享同一套模型结构和模型参数呢？据此，本文提出了一种新的蒸馏方式：Interaction Distillation，该方式的核心思想就是将 teacher 和 student 集成到同一个模型中，共享同一套参数，从而有效避免了 teacher 和 student 表示知识的方式不一致的问题。

作者发现，当我们将 ColBERT 作为 teacher 时，这一思路有一种很简单的实现方式。如上图所示，由于 ColBERT 和 Dual-Encoder 的相关性建模均发生在顶层，我们可以自然而然地共享它们的底层参数，当输入的 query 和 document 均被编码为低维稠密向量后，我们可以分别计算 Dual-Encoder 的 Dot-Product 相关性和 ColBERT 的 MaxSim 相关性打分：

然后分别计算候选文档的概率分布：

本文提出的 Interaction Distillation，其实就是让这里的和分布尽量一致，这里作者选择采用 KL 散度来衡量分布的一致性：

但是训练模型光有一致性损失还不够，当我们共享 Dual-Encoder 和 ColBERT 参数时，仅仅是一个正则项而已，我们依旧需要引入监督信号用于 Dual-Encoder 和 ColBERT 打分的训练：

最终训练采用的损失函数则为上述三个损失的和：

综上，作者提出的 Interaction Distillation 实际上就是将 ColBERT 的 MaxSim 打分作为了一个辅助信号来引导 Dual-Encoder 的训练，这一结构有效地避免了 Dual-Encoder 和 ColBERT 内部知识表示不一致的问题。另外，这一结构也不需要像传统的知识蒸馏那样，需要先训好一个 teacher，再训 student，该结构可以通过一次正反向传播同时训练 teacher 和 student。

2.2 Cascade Distillation

在知识蒸馏的相关研究中，人们发现当 teacher 和 student 的模型尺寸差异过大的时候，student 的性能反倒会下降。也就是说，BERT-xxlarge→BERT-small 的效果通常比 BERT-largeBERT-small 的效果差。这一现象虽然匪夷所思但也不是完全没有道理，就好比大学教授给小学生上课，很多在 teacher 眼中理所应当的知识，在 student 眼中可能就不知所云了。

因此，本文在 Interaction Distillation 的基础上，进一步提出了 Cascade Distillation。顾名思义，该策略的思想就是将 ColBERT 当作助教（teacher assistant），将 Cross-Encoder 当作教师，执行如下的两阶段蒸馏策略：Cross-Encoder→ColBERT→Dual-Encoder。

那么 Cross-Encoder 蒸馏 ColBERT，是不是也可以采用上面提到的 Interaction Distillation 策略呢？很遗憾，作者并没有想到一个可以兼容这两者的结构，因为 Cross-Encoder 和 ColBERT 的结构有着本质的区别，若将底层的编码器换成交互式结构，那么 ColBERT 顶层的 late inteaction 结构是没有意义的。因此，作者提出的 Cascade Distillation，本质上其实就是额外引入 Cross-Encoder 的蒸馏损失来进一步辅助 Dual-Encoder 的训练，使得知识可以从 Cross-Encoder 逐步传递到 Dual-Encoder。

首先，我们计算 Cross-Encoder 的打分和 ColBERT 打分的 KL 散度：

其中：

另外，作者认为除了蒸馏 Cross-Encoder 的分数分布，Cross-Encoder 的 token 级交互信息也是很重要的知识来源，因此作者提出了一种蒸馏注意力分布的一致性损失：

其中：

这里的和分别表示的是 Cross-Encoder 和 ColBERT 顶层的注意力分布，的计算表示为，下标表示第个注意力头，和分别表示 query 和 passage 的长度。

最后，我们将之前的 Interaction Distillation 结合进来，就有了如下的损失函数：

其中、和为对应模型的监督损失，、和为蒸馏损失，注意这里还额外包括了传统的 Cross-Encoder 到 Dual-Encoder 的蒸馏损失：

这里列出的公式可能比较冗长，下面的图示会更直观一些：

2.3 Dual Regularization

另外，作者还借鉴了 SimCSE 和 R-Dropout 的思想，提出了针对双塔式模型的 Dual Regularization 正则项，其具体策略就是对 passage 做两次随机采样 dropout mask 的前向传播，得到 passage 的两个表示，然后分别和 query 计算打分得到和，并最小化这两个打分分布的双向 KL 散度。作者将该正则项运用到了 Dual-Encoder 和 ColBERT 上，并加到了之前的损失之中。

Experiments

3.1 Setup

作者在 MARCO 和 NQ 数据集上评测了模型的表现，其中 Dual-Encoder 的热启模型采用了 ERNIE 2.0-Base，Cross-Encoder 采用了 ERNIE 2.0-Large。另外作者采用了一套比较复杂的多阶段训练流程，大致可以分为一下三步：首先，在通用语料库上进行 coCondenser 预训练；然后，在下游任务上使用 PAIR 提供的训练数据和策略进行 Interaction Distillaion 训练；最后，在第二步得到的模型基础上，采用 Cascade Distillation 对模型做进一步的微调。

3.2 Results

实验结果如下表所示，ERNIE-Search 在 MARCO Dev 数据集上的 MRR@10 突破了 40，取得了 SOTA 结果，当我们将参数量增大到 24 亿时，指标还会有进一步的提升，但是单看指标我们无法得知模型的增益来自于什么地方。

3.3 Ablation

3.3.1 Post-Train

作者首先对 Post-Train 策略进行了消融实验，如下图所示，Vanilla Post-Train 表示 coCondenser 预训练，Vanilla Finetune 表示 DPR 原始的对比损失训练，ID 表示 Interaction Distillation，CB 表示 PAIR 中的 Cross-Batch 策略，ID' 表示 ID+CB，DualReg 表示前文提到的 Dual Regularization。该消融实验主要表明了将 ID' 策略加到 Post-Train 中是有效的，但需要注意这里作者并没有去掉 CB 策略单独对比 ID 策略。

3.3.2 Interaction Distillation

接下来，作者对 Interaction Distillation 策略进行了消融实验，如下表所示，可以看出作者提出的 Interaction Distillation 策略和 Dual Regularization 策略都能带来一定的性能提升，但是 ID 策略本身涨点并不显著，反倒是 CB 和 DualReg 看起来更有效。

3.3.3 Cascade Distillation

最后，作者对最终的损失函数进行了消融实验，可以发现，去掉任何一个损失都会导致性能下降。但是去掉性能下降最为严重，说明 Cross-Encoder 直接蒸馏 Dual-Encoder 在实践中其实是很有效的，而去掉作者提出的 Cascade 损失和 Attention Map 损失，MRR@10 仅下降了 0.2pp。

总结

总体来说，虽然本文的出发点较好，但损失设计和训练流程过于复杂，消融实验也做得不够诚恳，SOTA 性能的主要来源可能更多还是来自于 coConderser 预训练，PAIR 的微调策略以及 Large 版 Cross-Encoder 的使用，交互蒸馏和级联蒸馏的实用价值并没有满足本人的预期。

另外值得一提的是，本文的模型结构实际上和度厂之前发布的 RocketQAv2 很相似，即联合训练 Dual-Encoder 和 Cross-Encoder，而本文可以看作是在 RocketQAv2 的基础上加上了额外的 ColBERT 损失，但整体提升并不是很显著，有暴力调参的嫌疑。

因此，虽然本文的 Motivation 在于 Cross-Encoder 和 Dual-Encoder 的知识存储结构和表达能力存在较大差异，直接蒸馏 logits 可能会存在知识迁移效率低的情况，但实际上，在最终的实验结果中，作者提出的策略仅仅产生了不到 1pp 的波动，说明这个凭空假设的问题可能并不严重。所以，在做具体的模型优化工作的时候，还是要从具体的 badcase 入手，发现模型存在的具体问题，并提出针对性的解决策略，才是更好的模型迭代的方向。

更多阅读