在蒸馏之后对学生模型再进行微调，进一步提升能力。

随着NLP模型的大小增加到数千亿个参数，创建这些模型的更紧凑表示的重要性也随之增加。知识蒸馏成功地实现了这一点，在一个例子中，教师模型的性能的96%保留在了一个小7倍的模型中。然而，在设计教师模型时，知识的提炼仍然被认为是事后考虑的事情，这可能会降低效率，把潜在的性能改进留给学生。

此外，在最初的提炼后对小型学生模型进行微调，而不降低他们的表现是困难的，这要求我们对教师模型进行预训练和微调，让他们完成我们希望学生能够完成的任务。因此，与只训练教师模型相比，通过知识蒸馏训练学生模型将需要更多的训练，这在推理的时候限制了学生模型的优点。

如果在设计和训练教师模型的过程中，将知识的提炼放在首要位置和中心位置，会有什么可能呢？我们能否设计并成功地训练一个“假定”要被蒸馏的模型，蒸馏后的版本能否成功地在任何下游任务上进行微调？这是我们在本文中总结的MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices中提出的一些问题。据作者说目前还没有与任务无关的轻量级预训练模型，他们要填补这个领域的空白。实验结果上看，这个空白填补的还是很成功的。

MobileBERT 结构

MobileBERT 采用的和 BERT-large 一样深的层数，在每一层中的 transformer 中加入了 bottleneck 机制使得每一层 transformer 变得更窄，但是并没有丢失其与 self-attention 和 Feed-Forward 之间的平衡。

Transformer结构可视化 (a) BERT， (b) MobileBERT 教师，© MobileBERT 学生。用“Linear”标记的绿色梯形称为bottlenecks。

线性层

知识的蒸馏要求我们比较老师和学生的表示，以便将它们之间的差异最小化。当两个矩阵或向量维数相同时，这是很直接的。因此，MobileBERT在transformer块中引入了一个bottleneck 层。这让学生和老师的输入在大小上是相等的，而他们的内部表示可以不同。这些bottleneck在上图中用“Linear”标记为绿色梯形。在本例中，共享维度是512，而教师和学生的内部表示大小分别是1024和128。这使得我们可以使用BERT-large(340M参数)等效模型来训练一个25M参数的学生。

此外，由于两个模型的每个transformer块的输入和输出尺寸是相同的，因此可以通过简单的复制将嵌入参数和分类器参数从教师传递给学生！

多头注意力

细心的读者会注意到，多头注意块(MHA)的输入不是先前线性投影的输出。相反，使用初始输入。这篇论文的设计选择没有任何动机，只能让我们去推测。我相信其原因是它所允许的自由度增加了。基本上，我们将迫使模型处理信息的方式分离为两个单独的流，一个流入MHA块，另一个作为跳跃连接。(使用线性投影的输出并不会因为初始的线性变换而改变MHA块的行为，这也是很容易说服自己的。)

堆叠FFN

为了在这个小的学生模型中实现足够大的容量，作者引入了他们所谓的stacked FFN，如图中学生模型概述中的虚线框所示。Stacked FFN只是简单的将Feed Forward + Add & Norm blocks块重复了4次，选择这一方式来得到MHA和FFN block之间的良好的参数比例。本工作中的消融研究表明，当该比值在0.4-0.6范围内时，性能最佳。

操作优化

由于其目标之一是在资源有限的设备上实现快速推理，因此作者确定了他们的架构可以进一步改进的两个方面。

  1. 把smooth GeLU的激活函数更换为ReLU

  2. 将normalization操作转换为element-wise的线性变换

建议知识蒸馏目标

为了实现教师和学生之间的知识转移，作者在模型的三个阶段进行了知识蒸馏：

  1. 特征图迁移[Feature Map Transfer (FMT)] — 允许学生模仿老师在每个transformer层的输出。在上面的架构图中，它表示为模型输出之间的虚线箭头。

  1. 注意力图迁移[Attention Transfer (AT)] —这让老师在不同层次上关注学生，这也是我们希望学生学习的另一个重要属性。这是通过最小化每一层和头部的注意力分布(KL散度)之间的差异而实现的。

  1. 预训练蒸馏[Pre-training Distillation (PD)] — 也可以在预训练中使用蒸馏，通过组合Masked语言建模和下一个句子预测任务的线性组合。

有了这些目标，我们就有了不止一种方法来进行知识的提炼。作者提出了三种备选方案：

  1. 辅助知识迁移。分层的知识迁移目标与主要目标(Masked语言建模和下一句预测)一起最小化。这可以被认为是最简单的方法。

  2. 联合知识迁移。不要试图一次完成所有的目标，可以将知识提炼和预训练分为两个阶段。首先对所有分层知识蒸馏损失进行训练直到收敛，然后根据预训练的目标进行进一步训练。

  3. 进一步的知识转移。两步法还可以更进一步。如果所有层同时进行训练，早期层没有很好的最小化的错误将会传播并影响以后层的训练。因此，最好是一次训练一层，同时冻结或降低前一层的学习速度。

知识迁移技术。(a) 辅助知识迁移，(b) 联合知识迁移， © 进一步知识迁移

实验结果

作者在三种配置中评估了他们提出的MobileBERT：有25M参数的主模型(MobileBERT)，没有操作优化的同一个模型(MobileBERT w/o OPT)，以及只有15M参数的模型(MobileBERT-tiny)。这些模型都与基线算法进行了比较，如ELMo，GPT和BERT-base以及相关蒸馏工作：BERT-PKD, DistilBERT。

研究发现，通过渐进式知识转移，训练这些不同的MobileBERT是最有效的，其效果始终显著优于其他两个。

在GLUE benchmark上的结果

我们发现，MobileBERT w/o OPT比大得多的BERT-base的平均GLUE得分高出0.2，而比它小4倍。另一方面，MobileBERT仅落后于BERT-base 0.6分，但推理时间快得多 —— 在Pixel 4的手机上对128个tokens序列进行推理只需要62毫秒！然而，它的性能仍然具有竞争力，因为它的性能远远超过GTP和ELMo。

因此，可以肯定地说，我们可以创建一个蒸馏模型，它可以在资源有限的设备上达到既高性能又快速

与TinyBERT相比，MobileBERT-tiny的性能略好一些。然而，当你考虑到TinyBERT是如何为GLUE任务进行微调时，这一点就更令人印象深刻了。记住，在此工作之前，不可能去微调学生模型，因为他们的容量太小了。因此，在把它的知识提炼成TinyBERT之前，必须对它的老师BERT-base进行微调！但对MobileBERT来说，情况并非如此。

MobileBERT已经在GLUE上被自己进行了微调，这证明了它可以通过建议的蒸馏过程创建一个任务不可知的模型。

总结

MobileBERT在transformer模块中引入了bottlenecks，这使得我们可以更容易地将知识从大尺寸的教师模型中提取到小尺寸的学生模型中。这种技术允许我们减少学生的宽度，而不是深度，这是已知的，以产生一个更有能力的模型。这个模型强调了这样一个事实，它可以创建一个学生模型，它本身可以在最初的蒸馏过程后进行微调。

此外，结果还表明，这在实际应用中也是正确的，因为在Pixel 4的手机上，MobileBERT在GLUE上的参数减少了4倍，推理速度加快了5.5倍，却能达到BERT-base性能的99.2%。

论文地址：https://arxiv.org/abs/2004.02984

pytorch代码地址: https://github.com/lonePatient/MobileBert_PyTorch