今天介绍的这篇文章是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模型,提升其拟合能力。
4、实验结果
本文的对比实验中,将文中提出的方法与两类模型进行了对比。第一类方法是不带知识蒸馏的稠密或稀疏检索模型,包括BM25、DeepCT、docT5query等;第二类方法是使用知识增强后的稠密检索模型,包括Margin-MSE、TCT-ColBERT等。从实验结果可以看出,文中提到的方法在各项指标上都有明显的效果提升。