Press "Enter" to skip to content

TinyBert: 模型蒸馏的全方位应用

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

Overall

 

《Bert与模型蒸馏:PKD和DistillBert》 一篇中,我们介绍了模型蒸馏的知识以及如何在Bert上做模型蒸馏。而本篇将会介绍模型蒸馏在Bert上全方位的应用。

 

之所以说全方位,体现在如下几个方面:

 

在预训练阶段和微调阶段都应用知识蒸馏

 

在Transformer块的几个部分都使用知识蒸馏

 

下图中可以看到TinyBert和之前方法的区别,可以看到,TinyBert的知识蒸馏的应用非常的全面。

 

 

我们注意到TinyBert那一行有两个地方没有勾上,一个Init,一个预训练阶段的Pred。Init是因为TinyBert不再像PKD或者DistillBert那样用Teacher模型的前几层或者隔一层抽一层去初始化Student,而是随机初始化。这样可以:

 

 

保证灵活性,如果使用Teacher模型的参数,则会限制Student的宽度,导致结构上受限。

 

而因为有Embed,Attn,Hidn等各个部分上的蒸馏,所以不用Teacher模型的参数初始化也没有损失。

 

 

问题建模

 

假设Student模型有M层,Teacher模型有N层,那幺在做蒸馏的时候就需要作出一些选择,即Student的某一层去模仿Teacher模型哪一层。这里,可以把这种选择当成一个函数,即 n = g(m) ,其中 0 = g(0)是embedding层, N+1 = g(M+1) 是prediction层。如下图所示:

 

 

我们先不管怎幺对应的问题,整个模型蒸馏的损失应该是:

 

 

Transformer的知识蒸馏

 

Bert是由多个Transformer块堆叠而成的,在每个Transformer块上,都计算两个蒸馏损失:attn蒸馏和hidn蒸馏。如下图所示:

 

 

其中,attn蒸馏是在没有归一化的attn矩阵上计算均方差。这里隐含要求多头注意力的头的个数是相同的。

 

 

而对于隐含状态而言,因为Student模型和Teacher模型的隐含状态大小可能不同,因而直接计算均方差是不现实的。所以,在Student模型上用一个矩阵先做变换再去和Teacher模型去做均方差计算,如下图所示。在蒸馏的时候,这个单独的矩阵也会被梯度下降,但蒸馏完成后,这个矩阵是用不到的。

 

 

类似的,Embedding层上的蒸馏也是这样做,如下图所示:

 

 

而prediction层的蒸馏则是使用交叉熵来计算,如下图:

 

 

因此,Student模型每一层上的蒸馏损失为:

 

 

两步蒸馏

 

正如上面所说,TinyBert不仅在预训练阶段,还发生在微调阶段。所以TinyBert的训练过程如下,在微调阶段,Teacher模型是微调后的Bert。

 

 

在做微调阶段的蒸馏的时候,还会做数据增强,增强的方式就是选择一些位置,用Teacher模型去Top-N的预测,然后用预测值替换该位置上原来的词语。

 

具体算法如下:

 

 

实验

 

实验中训练的TinyBert为:

 

4层

 

hidden_size = 312

 

全连接层 = 1200

 

注意力头数 = 12

 

共14.5M参数

 

如下面两个表所示,TinyBert相对于之前的蒸馏方法能够在带来更大速度提升的条件下保持更好的效果。

 

 

同时,TinyBert也具有很好的延展性,Student模型越大,保持的Teacher模型的程度就越好。

 

 

而在消融实验中,可以看到三点:

 

 

GD(在预训练时候的蒸馏)、TD(在微调时的蒸馏)、DA(数据增强)缺一不可

 

模型蒸馏的诸多损失函数中,起决定作用的是Transformer块上的蒸馏,其中的注意力和全连接两个部分上,注意力上的蒸馏作用更大一些。

 

Student模型和Teacher模型的层次对应上,使用平均策略更好。

 

 

 

总结与思考

 

TinyBert中的蒸馏比较全,从Embedding到Transformer块再到prediction层都会被用来做蒸馏,这样全面的方式可以替代初始化,即可以用随机的方式初始化Student模型。

 

而在微调阶段的蒸馏,数据增强则是不可或缺十分重要的因素。

 

勤思考, 多提问是Engineer的良好品德。

 

问题如下:

 

这里的数据增强是用Bert的预测结果来替换一些原始值。论文中控制着一定的比例,那幺问题是:能不能完全用模型生成的数据去做蒸馏呢?

 

 

参考文献

 

[1]. Jiao, Xiaoqi, et al. “Tinybert: Distilling bert for natural language understanding.” arXiv preprint arXiv:1909.10351 (2019).

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注