Press "Enter" to skip to content

MobileBert: Pixel4上只需40ms

 

Overall

 

Bert模型虽然很火,但是模型太大,在device上基本不可用,要想更好的使用的话需要让模型变小。MobileBert综合运用多种技术,将Bert模型在手机上的运行降低到了40ms,让我看到了Bert在手机上运用的希望。

 

MobileBert中采用的技术如下:

Inverted-Bottleneck: 模型改动,让模型变窄而不是变浅。
Stacked Feed-Forward Network: 保证Attention部分和全连接部分的参数比,从而保证效果。
Operational Optimizations: 移除LayerNormalization,同时将Gelu激活函数换成更快的Relu。
Embedding Factorization: Embedding做矩阵分解。
目标函数:知识蒸馏目标函数,Feature Map Transfer,Attention Transfer。
知识蒸馏的不同策略。
量化

可以看到,优化方法仍然没有出我们在前面文章里讲的四条:量化,矩阵分解,结构改动,知识蒸馏。下面我们一项一项的进行讲解。

 

Inverted-Bottleneck(IB)

 

我们知道,Bert的体态其实很宽,Bert-base的宽度是768,Bert-large的宽度是1024。而Bert的宽度有几个特点:

等宽:从底层的embedding到中间的attention层再到最高层的输出,层次宽度是相等的。
例外:全连接部分的中间层的宽度是其他层次宽度的四倍。

这点如果不熟悉的欢迎去翻以前的文章,上面都有链接。

 

由于之前的研究表明,减少层次,让模型变浅,得到的小模型相对于大模型损失太大,所以MobileBert尝试把模型变窄。而直接全体变窄经过试验又会导致较大的损失,所以,MobileBert使用Inverted Bottleneck的形式来保证Bert模型中Transformer块之间的宽度不变。如下图所示:

a图是一个正常的Bert,有L个Transformer块。b图和c图则是带有Inverted-bottleneck的模型,可以看到,所谓的Inverted-Bottleneck,就是在每个Transformer块的前后各添加一个全连接层来保证Transformer块的输入和输出是同一个宽度。我们成这两个全连接层为 Linear入Linear 出这样宽度就变成两个: 块内宽度块间宽度

 

以此,b图可以把Transformer块变胖,就是一个更大的模型,c图是把中间的块变窄,就组成了MobileBert。

 

在b图和c图中,可以注意到,Multi-Head Attention的输入是Embedding,所以Multi-Head Attention和 Linear入 负责把输入embedding从 块间宽度 变成 块内宽度 。而输出的变化则由 Linear出 全权负责。

 

那幺就有个问题,为什幺输入上的做法要和输出不同?答案是为了保证模型的容量。MobileBertTiny就是全依靠Linear来变换输入输出。

在MobileBertTiny中,采用的Inverted-bottleneck就是只用 Linear入 来把embedding从 块间宽度 变成 块内宽度 ,然后 Linear入 的输出再输入给Multi-Head Attention。这样会带来参数量的大幅度降低,同时模型容量降低也可能伤害最后的效果。

 

由于c图从头训练达到较好效果比较难,所以MobileBert的做法是先训练b图的大模型,然后再去做模型蒸馏使得最后的MobileBert达到较好的效果。

 

Stacked Feed-Forward Network

 

我们知道,Transformer块分为Multi-Head Attention和Feed Forward两个部分,这两个部分有各自的功能:

Multi-Head Attention:负责捕捉不同位置的信息。
全连接层:增加模型的非线性表达能力。

而当用了Inverted-bottleneck后,Attention部分的输入是 块间宽度 的Embedding,而全连接层的输入则是 块内宽度 。因此就导致了一个问题,这两部分的参数矩阵此消彼长,Attention部分变多,而全连接层变少。导致它们的比例出了问题。而后面的实验会表明,这种比例的变化会导致模型的效果变差。

 

而解决方案则是把全连接层复制几份来恢复这个比例。还是看上面那张图,图中的c图的全连接层部分框起来乘以了一个F。这就是Stacked FFN。

 

Operational Optimizations(OPT)

 

对于中间的一些计算操作,MobileBert也做了优化。去掉了Layer Norm而用一个更简单的操作来替代:

其中,圆圈代表的是阿达玛矩阵乘积,即element-wise的矩阵乘法。

 

另个优化是Gelu激活函数被替换成Relu。

 

Embedding Factorization

 

embedding层占了很大一部分参数,为了解决这个问题,MobileBert使用了维度为128的embedding,然后用一维卷积来把embedding转化为512大小。

 

目标函数

 

MobileBert共采用了如下几种损失函数:

Feature Map Transfer(FMT): 即让Student模型去拟合Teacher模型上的Feature Map,Feature Map则是每个Transformer Block的输出。这一点来源于PKD算法。

Attention Transfer(AT): 即让Student模型去拟合Teacher模型上Attention矩阵。这点和TinyBert很像。

Pre-training Distillation(PD): 是蒸馏损失,MLM损失和NSP损失的线性组合:

知识蒸馏的不同策略

 

提出了三种策略:

Auxiliary Knowledge Transfer(AKT): 认为FMT和AT是辅助目标函数,因而用FMT,AT和PD线性组合直接训练。
Joint Knowledge Transfer(JKT): 认为中间层的损失函数不适合找到最后的最优模型,因而先用FMT和AT去训练子模型,然后再用PD去做第二阶段的训练。
Progressive Knowledge Transfer(PKT): 子模型可能不能完全学习Teacher模型,因为底层学的不对可能会影响上层,因此这个方法要逐层用FMT和AT进行学习,然后最后用PD做训练。

三种策略如下图所示:

实验

 

实验使用的一些模型结构参数如下表,可以看到,MobileBert和Bert-Large的深度相同,但是块内宽度只有128,块间宽度则是512。

IB-Bert Teacher模型的搜索

 

下表列举了一些探索的参数组合上的结果,可以看到:

从a-e可知,块间宽度降到512以下会有明显的损失,所以最后的Teacher模型的块间宽度是512.
从f-i可知,减少块内宽度也会有明显的损失,所以最后的Teacher模型的块内宽度是1024。
从a和f可知,减少head数目没有明显损失,所以Teacher模型的head数目是4。

Attention和FFN不同参数比例的结果

 

在跟Bert类似,接近1:2的参数比的时候会更好。

在下游任务上的对比以及latency

 

Glue任务上的结果如下:

可以看到,MobileBert without OPT效果出众,MobileBert-Tiny在Pixel 4上速度仅为40ms。

 

SQuAD上的结果如下:

量化

 

采用的post-training量化,会带来微弱的损失。

消融实验

 

Latency上的实验结果,去掉LayerNorm带来60ms提升,再去掉gelu又带来30ms提升。

不同策略对比,最后的PKT最好

损失函数上的消融实验,FMT起的效果最大。

总结与思考

 

勤思考,多总结是工程师的必备品德。

 

论文提出了Inverted-Bottleneck来从结构上对Bert进行优化,但小尺寸的IB难以训练,所以需要一个大尺寸的IB来做模型蒸馏。在这里可以看出模型蒸馏对于压缩的独特意义。

 

另一个独特的研究在于Multi-Head Attention和FFN的参数比的研究和实验,这里认为模型的容量体现在参数量上,那幺如何看待模型结构对于模型容量的影响?比如把FFN换成卷积也需要同样的参数量吗?

 

参考文献

[1]. Sun, Zhiqing, et al. “Mobilebert: a compact task-agnostic bert for resource-limited devices.” arXiv preprint arXiv:2004.02984 (2020).

Be First to Comment

发表回复

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