Press "Enter" to skip to content

论文阅读_知识蒸馏_MobileBERT

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

英文题目: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对比了三种不同策略的训练方法,可以看到递进训练效果更好。

 

 

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

Be First to Comment

发表回复

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