最近要开始使用Transformer去做一些事情了,特地把与此相关的知识点记录下来,构建相关的、完整的知识结构体系,
以下是要写的文章,本文是这个系列的第五篇:
Bert模型虽然很火,但是模型太大,在device上基本不可用,要想更好的使用的话需要让模型变小。
而目前模型变小的技术大概有四种:
本文介绍的就是知识蒸馏这一部分。主要介绍的是参考论文[1]中的PKD方法,参考论文[2]中的DistillBert本身较为简单,本文也会简要介绍一下。
首先,我们先来了解一下知识蒸馏,知识蒸馏是让一个小模型去学习一个大模型,所以首先会有一个预训练好的大模型,称之为Teacher模型,小模型被称为Student模型。知识蒸馏的方法就是会让Student模型去尽量拟合。这个的动机就在于跟ground truth的one-hot编码相比,Teacher模型的输出概率分布包含着更多的信息,从Teacher模型的概率分布中学习,能让Student模型充分去模拟Teacher模型的行为。
在具体的学习Teacher模型概率分布这个过程中,知识蒸馏还引入了温度的概念,即Teacher和Student的logits都先除以一个参数T,然后再去做softmax,得到的概率值再去做交叉熵。温度T控制着Student模型学习的程度,当T>1时,Student模型学习的是一个更加平滑的概率分布,当T<1时,则是更加陡峭的分布。因此,学习过程中,T一般是一个逐渐变小的过程。Teacher模型经过温度T之后的输出被称之为soft labels。
Soft labels的计算过程如下,Pt中的t代表的是Teacher模型,同样的,Student模型的输出也要经过类似的计算。
得到Teacher模型和Student的模型后,就可以计算知识蒸馏这部分的损失,得到LDS
然后student模型除了向teacher模型学习外,还需要向ground truth学习,得到LCE
然后再让两个损失去做加权平均。得到:
以上就是知识蒸馏的思想,这方法可以直接应用在Bert上,但是作为一个深层模型,在中间层次上的信息也很丰富,如何利用这部分的信息呢?这就有了PKD中提出的方法。
首先,PKD论文中先做了对比,减少模型宽度和减少模型深度,得到的结论是减少宽度带来的efficiency提高不如减少深度来的更大,因此,PKD主要关注减少模型深度。即Student模型比Teacher模型要浅。
论文所提出的多层蒸馏,即Student模型除了学习Teacher模型的概率输出之外,还要学习一些中间层的输出。论文提出了两种方法,第一种是Skip模式,即每隔几层去学习一个中间层,第二种是Last模式,即学习teacher模型的最后几层。如下图所示:
如果是完全的去学习中间层的话,那么计算量很大。为了避免这个问题,我们注意到Bert模型中有个特殊字段[CLS],在蒸馏过程中,让student模型去学习[CLS]的中间的输出。
而对于中间层的学习,使用的损失函数是均方差函数:
最后的损失函数如下:
初始化的话就采用Teacher模型的前几层来做初始化。
在GLUE上的实验结果如下:
可以看到,多层的方法在除了MRPC之外的任务上都能达到比较最好的效果。而究其原因,可能是因为MRPC的数据较少,从而导致了过拟合。
而Last和Skip模式的对比如下:
可以看到,Skip一般会比Last模式要好。这是因为,Skip方式下,层次之间的距离较远,从而让student学习到各种层次的信息。
那么Student模型的计算量和参数数目是如何的呢?如下图所示:
由于需要每层计算损失,Student和Teacher模型每层的宽度都是一样(这其实是一个限制),Student模型的参数量主要少在层次少。
而更好的teacher模型会带来增长么?答案是不会的,可以看上图,把12层的Bert模型换成了24层的Bert模型,反而导致效果变差。究其原因,可能是因为在实验中,我们使用Teacher模型的前N层来初始化Student模型,对于24层模型来说,前N层更容易导致不匹配。而更好的方法则是Student模型先训练好,再去学Teacher模型。
DistillBert的做法就比较简单直接,同样的,DistillBert还是保证模型的宽度不变,模型深度减为一半。主要在初始化和损失函数上下了功夫:
具体结果就不赘述了,可以去参考原始论文。