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