最近要开始使用Transformer去做一些事情了,特地把与此相关的知识点记录下来,构建相关的、完整的知识结构体系,
以下是要写的文章,本文是这个系列的第七篇:
在《Bert与模型蒸馏:PKD和DistillBert》一篇中,我们介绍了模型蒸馏的知识以及如何在Bert上做模型蒸馏。而本篇将会介绍模型蒸馏在Bert上全方位的应用。
之所以说全方位,体现在如下几个方面:
下图中可以看到TinyBert和之前方法的区别,可以看到,TinyBert的知识蒸馏的应用非常的全面。
我们注意到TinyBert那一行有两个地方没有勾上,一个Init,一个预训练阶段的Pred。Init是因为TinyBert不再像PKD或者DistillBert那样用Teacher模型的前几层或者隔一层抽一层去初始化Student,而是随机初始化。这样可以:
假设Student模型有M层,Teacher模型有N层,那么在做蒸馏的时候就需要作出一些选择,即Student的某一层去模仿Teacher模型哪一层。这里,可以把这种选择当成一个函数,即n = g(m),其中 0 = g(0)是embedding层,N+1 = g(M+1)是prediction层。如下图所示:
我们先不管怎么对应的问题,整个模型蒸馏的损失应该是:
Bert是由多个Transformer块堆叠而成的,在每个Transformer块上,都计算两个蒸馏损失:attn蒸馏和hidn蒸馏。如下图所示:
其中,attn蒸馏是在没有归一化的attn矩阵上计算均方差。这里隐含要求多头注意力的头的个数是相同的。
而对于隐含状态而言,因为Student模型和Teacher模型的隐含状态大小可能不同,因而直接计算均方差是不现实的。所以,在Student模型上用一个矩阵先做变换再去和Teacher模型去做均方差计算,如下图所示。在蒸馏的时候,这个单独的矩阵也会被梯度下降,但蒸馏完成后,这个矩阵是用不到的。
类似的,Embedding层上的蒸馏也是这样做,如下图所示:
而prediction层的蒸馏则是使用交叉熵来计算,如下图:
因此,Student模型每一层上的蒸馏损失为:
正如上面所说,TinyBert不仅在预训练阶段,还发生在微调阶段。所以TinyBert的训练过程如下,在微调阶段,Teacher模型是微调后的Bert。
在做微调阶段的蒸馏的时候,还会做数据增强,增强的方式就是选择一些位置,用Teacher模型去Top-N的预测,然后用预测值替换该位置上原来的词语。
具体算法如下:
实验中训练的TinyBert为:
如下面两个表所示,TinyBert相对于之前的蒸馏方法能够在带来更大速度提升的条件下保持更好的效果。
同时,TinyBert也具有很好的延展性,Student模型越大,保持的Teacher模型的程度就越好。
而在消融实验中,可以看到三点:
TinyBert中的蒸馏比较全,从Embedding到Transformer块再到prediction层都会被用来做蒸馏,这样全面的方式可以替代初始化,即可以用随机的方式初始化Student模型。
而在微调阶段的蒸馏,数据增强则是不可或缺十分重要的因素。
勤思考, 多提问是Engineer的良好品德。
问题如下:
回答后续公布,欢迎关注公众号【雨石记】