WWW'23|递进式知识蒸馏提升检索模型效果

今天介绍的这篇文章是WWW 2023中微软发表的一篇工作，应用场景是检索，主要提出了一种新的蒸馏学习方法，用来实现将复杂的student模型的知识，蒸馏到简单的双塔dense检索模型上。核心提出的优化点是模型层面的aggressive蒸馏，以及数据层面的aggressive蒸馏。这里的aggressive，指的是teacher模型的复杂度逐渐提升，或者样本的预测难度逐渐增加。下面为大家详细介绍下这篇工作。

论文标题：PROD: Progressive Distillation for Dense Retrieval

下载地址：https://arxiv.org/pdf/2209.13335.pdf

1、出发点

知识蒸馏已经有比较长的发展历史了，其基本目标是，将一个模型学到的知识（teacher模型）迁移到另一个模型上（student模型）。知识蒸馏的应用场景很广泛，其中最重要的一个场景是模型压缩。例如，这篇文章的场景就是线上的检索模型。为了检索效率，线上模型一般采用的是简单的双塔结构，query和document使用一个训练好的模型生成表示向量，直接计算二者的内积计算相似度，实现topK检索。

这种方法的效率虽然很高，但是预测准确度肯定不如那些单塔模型（即query和document拼接到一起输入BERT）。并且线上的模型由于追求效率，不能使用太多参数。为了提升这种双塔模型的效果，基于teacher-student的知识蒸馏方法被提出：先训练一个强大的teacher模型，再把teacher模型的知识蒸馏到student模型中。

然而，随着进一步的研究发现，并不是越准确、越复杂对的teacher模型，蒸馏出来的student效果越好，有时还不如一个弱一些的teacher模型。为了解决这个问题，本文提出了一种递进蒸馏的方法，主要分为模型侧的递进（Teacher Progressive Distillation）和数据侧的递进（Data Progressive Distillation）。

2、模型侧递进

模型侧的递进主要体现在teacher模型的从易到难。假设我们给一个6层双塔模型蒸馏知识，整个模型递进分为3个过程：使用12层的双塔模型作为teacher->使用6层的单塔模型作为teacher->使用12层的单塔模型作为teacher。这个过程中，teacher模型的复杂度和预测准度是逐渐增强的（双塔模型复杂度小于单塔；浅层网络复杂度小于深层网络），给student模型一个渐进式的学习过程。整个蒸馏过程的示意图如下。

在双塔蒸馏双塔的过程中，主要是2个损失函数：一个是拟合hard label的损失，另一个是student模型的预测打分分布和teacher模型的预测打分分布之间的差异。后者公式如下：

在双塔蒸馏单塔的过程中，除了上述两个loss外，还引入了一个正则化loss，让模型的预测分布不要偏离第一阶段的蒸馏后的模型太远，防止catastrophic forgetting问题。将第一阶段的模型备份一个并freeze参数，让后续的蒸馏打分分布不要偏离freeze的参数太远，损失函数如下：

3、数据侧递进

上述过程从模型结构的角度，从易到难，将知识蒸馏到student模型中，但是仍然会有一些样本在teacher和student模型中的预测结果存在差异。为了进一步提升蒸馏效果，文中从数据侧也进行了递进式的蒸馏过程。

数据侧递进是在模型侧递进过程训练完成之后。数据侧递进的核心目标是，随着student模型的训练过程，逐步提升student模型样本的难度，让其能够达到接近teacher模型的水平。如何从样本中选择困难样本呢？文中设计了如下循环过程进行每个阶段student模型输入样本的采样：

• 获取student模型打分topK的负样本，这部分样本是当前student预测效果较差的部分；
• 获取query集合，teacher模型能将正样本预测成top1，但是student样本只能预测成topk'，利用这部分query构造数据集，这部分query的正样本就是标注样本，负样本是上一步中挖掘到的困难样本；
• 使用这部分样本训练teacher模型；
• 使用上一步训练好的teacher模型蒸馏student模型。

重复上述过程，不断挖掘困难样本，并利用这部分样本训练student模型，提升其拟合能力。

4、实验结果

本文的对比实验中，将文中提出的方法与两类模型进行了对比。第一类方法是不带知识蒸馏的稠密或稀疏检索模型，包括BM25、DeepCT、docT5query等；第二类方法是使用知识增强后的稠密检索模型，包括Margin-MSE、TCT-ColBERT等。从实验结果可以看出，文中提到的方法在各项指标上都有明显的效果提升。