论文阅读_知识蒸馏_MobileBERT

时间:2022-09-26 12:01:08

英文题目: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所示:

论文阅读_知识蒸馏_MobileBERT

与之前的模型相比,文中模型只在预训练阶段使用知识转移,在下游任务中不需要微调教师模型和数据增强;另外,它不修改模型层数,只改变了每层的宽度。

方法

图-1展示了模型参数的具体大小:

论文阅读_知识蒸馏_MobileBERT

瓶颈和逆瓶颈

从表-1中可以看到,每块维度仅为128,另外,将全流程的层输入输出都调整为更窄的512。
在训练过程中,首先训练老师模型 IB-BERT ,再将知识转移到学生模型MobileBERT。老师模型设计如图-1(b)所示,它是一个逆瓶颈结构,师生模型的Feature map大小均为512,这样可以直接对比层输出。同时使用瓶颈和逆瓶颈使老师模型保持性能,学生模型足够紧凑。

多层前馈网络

注意力层和前馈网络功能不同:前者连接了不同空间数据,后者提升了模型的非线性表达能力。一般BERT的注意力层和前馈网络参数比例为1:2,而由于MobileBERT模型比较窄,注意力参数相对更多,为解决此问题,设计了多个叠加的前馈层,如图-1©所示。

操作优化

通过延迟分析发现,归一化层和 gelu 激活函数占延迟的很大比例,因此对其进行优化。

  • 去掉归一化层:用针对元素的线性变换替换n通道的隐藏层归一化。

论文阅读_知识蒸馏_MobileBERT

  • 使用relu代替glue

嵌入分解

嵌入词表占了模型的很大比例,如表-1所示,MobileBERT将嵌入词表压缩到128维,再用核为3的卷积层将其扩展到512维。

训练目标

转移目标包含:feature map转移和注意力转移。

Feature map转移
BERT中每层的输出是下层的输入,因此,需要让学生的层输出尽量与老师结果一致。这里使用MSE作为损失函数。

论文阅读_知识蒸馏_MobileBERT

公式中l为层索引,T为序列长度,N为feature map大小。

注意力转移
注意力提升了BERT模型的语言识别能力,因此也通过学习注意力层将知识更好地转移到MobileBERT模型。使用KL散度来计算注意力层的差异,作为损失函数:

论文阅读_知识蒸馏_MobileBERT

预训练蒸馏
除了针对各Transformer层的转移,在预训练时,还使用了蒸馏的loss函数,最终损失函数由三部分组成:遮蔽的损失(MLM),下一句预测的损失(NSP)和针对MLM的蒸馏损失:

论文阅读_知识蒸馏_MobileBERT

训练策略

下面讨论三种训练策略,如图-2所示:

论文阅读_知识蒸馏_MobileBERT

辅助知识转移(AKT)
将中间层的知识转移作为知识蒸馏的辅助任务。损失函数是所有层的知识转移损失和预测蒸馏损失的线性组合。

联合知识转移(JKT)
将训练分成两步:先训练中间层知识转移,然后训练预训练的知识蒸馏。

递进知识转移(PKT)
下层的错误可能影响到上层的学习,因此设计了先训练下层,再训练上层的知识转移方法。一般在训练上层时冻结下层参数,也可以使用软化方法,训练上层时对下层参数使用较小的学习率。实验证明该方法效果最好。

实验

表-2展示了在SQuAD数据集上,使用不同参数大小的蒸馏结果:

论文阅读_知识蒸馏_MobileBERT

图-3展示了堆叠前馈层的效果:

论文阅读_知识蒸馏_MobileBERT

精调下游任务的方法与普通BERT一致,略有不同的是MobileBERT常需要更大的学习率和更多次迭代。主实验结果如表-8所示:

论文阅读_知识蒸馏_MobileBERT

可以看到,由于是对BERT LARGE的蒸馏,MobileBERT与BERT BASE模型效果类似,而大小比BERT小4.3倍,速度是BERT的5.5倍。

表-8对比了三种不同策略的训练方法,可以看到递进训练效果更好。

论文阅读_知识蒸馏_MobileBERT

文中用的每个优化技术都在实验中做了效果比较,这里只罗列了最重要的部分,其它详见论文正文。