英文题目:MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices
中文题目:MobileBERT:面向资源有限设备的任务无关的压缩模型
论文地址:https://arxiv.org/pdf/2004.02984
领域:自然语言处理,模型蒸馏
发表时间:2020
作者:Zhiqing Sun,卡内基梅隆大学,谷歌
出处:ACL
被引量:162
代码和数据:https://github.com/google-research/google-research/tree/master/mobilebert
阅读时间:22.06.16
读后感
使用了瓶颈结构,同时微调模型实现加速,最终实现了任务无关的蒸馏BERT模型。
介绍
大模型效果好,但受资源限制。文中提出了MobileBERT模型对BERT模型进行压缩和加速,它和原始BERT一样,是任务无关的,可以泛化到不同下游任务。MobileBERT是对BERT LARGE的优化,使用瓶颈结构。具体实现分两步:先训练一个倒瓶颈的特殊的老师模型,然后再将知识转移到学生模型MobileBERT。它比BERT BASE模型小4.3倍,快5.5倍,在GLUE数据集上达到与BASE BASE类似的效果。
文中提出的MobileBERT模型与BERT LARGE的深度相同,用瓶颈结构使每一层更窄,如图-1所示:
与之前的模型相比,文中模型只在预训练阶段使用知识转移,在下游任务中不需要微调教师模型和数据增强;另外,它不修改模型层数,只改变了每层的宽度。
方法
图-1展示了模型参数的具体大小:
瓶颈和逆瓶颈
从表-1中可以看到,每块维度仅为128,另外,将全流程的层输入输出都调整为更窄的512。
在训练过程中,首先训练老师模型 IB-BERT ,再将知识转移到学生模型MobileBERT。老师模型设计如图-1(b)所示,它是一个逆瓶颈结构,师生模型的Feature map大小均为512,这样可以直接对比层输出。同时使用瓶颈和逆瓶颈使老师模型保持性能,学生模型足够紧凑。
多层前馈网络
注意力层和前馈网络功能不同:前者连接了不同空间数据,后者提升了模型的非线性表达能力。一般BERT的注意力层和前馈网络参数比例为1:2,而由于MobileBERT模型比较窄,注意力参数相对更多,为解决此问题,设计了多个叠加的前馈层,如图-1©所示。
操作优化
通过延迟分析发现,归一化层和 gelu 激活函数占延迟的很大比例,因此对其进行优化。
- 去掉归一化层:用针对元素的线性变换替换n通道的隐藏层归一化。
- 使用relu代替glue
嵌入分解
嵌入词表占了模型的很大比例,如表-1所示,MobileBERT将嵌入词表压缩到128维,再用核为3的卷积层将其扩展到512维。
训练目标
转移目标包含:feature map转移和注意力转移。
Feature map转移
BERT中每层的输出是下层的输入,因此,需要让学生的层输出尽量与老师结果一致。这里使用MSE作为损失函数。
公式中l为层索引,T为序列长度,N为feature map大小。
注意力转移
注意力提升了BERT模型的语言识别能力,因此也通过学习注意力层将知识更好地转移到MobileBERT模型。使用KL散度来计算注意力层的差异,作为损失函数:
预训练蒸馏
除了针对各Transformer层的转移,在预训练时,还使用了蒸馏的loss函数,最终损失函数由三部分组成:遮蔽的损失(MLM),下一句预测的损失(NSP)和针对MLM的蒸馏损失:
训练策略
下面讨论三种训练策略,如图-2所示:
辅助知识转移(AKT)
将中间层的知识转移作为知识蒸馏的辅助任务。损失函数是所有层的知识转移损失和预测蒸馏损失的线性组合。
联合知识转移(JKT)
将训练分成两步:先训练中间层知识转移,然后训练预训练的知识蒸馏。
递进知识转移(PKT)
下层的错误可能影响到上层的学习,因此设计了先训练下层,再训练上层的知识转移方法。一般在训练上层时冻结下层参数,也可以使用软化方法,训练上层时对下层参数使用较小的学习率。实验证明该方法效果最好。
实验
表-2展示了在SQuAD数据集上,使用不同参数大小的蒸馏结果:
图-3展示了堆叠前馈层的效果:
精调下游任务的方法与普通BERT一致,略有不同的是MobileBERT常需要更大的学习率和更多次迭代。主实验结果如表-8所示:
可以看到,由于是对BERT LARGE的蒸馏,MobileBERT与BERT BASE模型效果类似,而大小比BERT小4.3倍,速度是BERT的5.5倍。
表-8对比了三种不同策略的训练方法,可以看到递进训练效果更好。
文中用的每个优化技术都在实验中做了效果比较,这里只罗列了最重要的部分,其它详见论文正文。