论文阅读_知识蒸馏_MobileBERT

Posted 2022-10-03

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

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