随着NLP模型的规模增加到数千亿个参数，对这些模型进行高效的压缩也越来越重要。知识蒸馏成功地实现了这一目标，典型例子：student模型压缩了1/7却达到了模型96%的效果（https://medium.com/dair-ai/tinybert-size-does-matter-but-how-you-train-it-can-be-more-important-a5834831fa7d）。但是，在设计teacher模型时仍会考虑知识蒸馏，这可能会降低有效性，使student模型潜在的表现提高。
此外，在初始蒸馏后微调小型student模型而又不降低其性能的是非常困难的，要求我们针对希望student模型能够解决的任务进行预训练和微调。因此，与仅训练teacher模型相比，通过知识蒸馏来训练student模型将需要更多的训练次数和时间，相对而言，student模型优势就是，inference的时间大大缩短。
相反，如果在teacher模型的设计和训练过程中知识蒸馏放在首位并成为中心，这将有可能吗？我们是否可以设计并成功训练应该被提炼的模型，并且提炼的版本可以在任何下游任务上成功进行微调吗？这些是MobileBERT中解决的一些问题：我们在本文中尝试回答了这个问题：MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices 。
MobileBERT架构
（a）BERT，（b）MobileBER Tteacher模型和（c）MobileBERT student模型中的Transformer架构可视化。标有“linear”的绿色梯形称为bottlenecks。
线性（linear）部分
知识蒸馏要求我们比较teacher和student模型的表现形式，以使两者之间的差异最小。当两个矩阵或向量的维数相同时，这很简单。因此，MobileBERT 在转换模块中引入了bottleneck层（bottleneck）。这样，student模型和teacher模型的输入内容就可以相等，而内部表示形式也可以不同。这些bottleneck显示为绿色梯形，在上图中标有“线性”。在此特定情况下，共享维度为512，而teacher和student模型的内部表示大小分别为1024和128。这使我们可以使用BERT大型（340M参数）等效模型来训练2500万参数student模型。
此外，由于两个模型的每个Transformer块的输入和输出尺寸都相同，因此只需复制它们，就可以将embedding参数和分类器参数从teacher模型迁移到student模型！
多头注意力（Multi-Head Attention）
细心的读者会注意到，多头注意模块（MHA）的输入不是先前线性投影的输出。而是使用初始输入。本文中没有提到选择这种设计的动机，这使我们不得不猜测。我认为原因是它可以增加的模型的自由度。基本上，我们将模型被强制处理为两个独立的流，一个流送入MHA块，另一个流作为skip-connection。（也很容易说服自己，使用线性投影的输出不会由于其初始线性变换而改变MHA块的行为。）
堆叠式FFN
为了在简答的student模型中获得足够的容量，作者介绍了他们所谓的堆叠式FFN，在上图的student模型概述中以虚线框显示。堆叠的FFN只需将前馈+add&norm重复4次，即可选择此选项以在MHA和FFN块之间实现良好的参数比率。这项工作中的消融（Ablation）研究表明，当该比率在0.4-0.6范围内时，可获得最佳性能。
操作优化（Operational optimization）
由于目标之一是在资源受限的设备上实现快速inference，因此作者确定了可以进一步改进其体系结构的两个领域。
1. 替换ReLU的平滑GeLU激活功能
2. 替换normalization操作，实用简单的元素线性转换（element-wise linear transformation）
提出的知识蒸馏目标
为了实现建议的teacher和student之间的知识迁移，作者在模型的三个阶段应用知识蒸馏：
1. Feature map 迁移 -允许student模型在每个transformer层输出处模仿teacher模型。在上面的架构图中，这显示为模型输出之间的虚线箭头。
2. Attention map迁移 -teacher要在不同层次和不同head上加入的记号是我们希望student模型学习的另一个重要属性。通过最小化每一层和head的注意力分布（KL散度）之间的差异，可以启用此功能。
3. pre-training蒸馏 —通过将屏蔽语言建模和下一句预测任务以线性组合的方式进行组合，也可以在预训练期间使用蒸馏。
通过这些目标，我们可以进行知识蒸馏的方式不止一种。作者提出了三种选择：
1. 辅助知识迁移。分层知识迁移目标与主要目标（mask语言建模和下一句预测）一起被最小化。这可以被认为是最简单的方法。
2. 联合知识迁移。不必一次实现所有目标，而是可以将知识蒸馏和预训练分为两个训练阶段。首先，对所有分层知识蒸馏损失进行训练，直到收敛为止，然后执行具有预训练目标的进一步训练。
3. 渐进的知识迁移。两步方法甚至可以更进一步。如果同时训练所有层，则早期层中尚未正确最小化的错误将传播并影响后面层的训练。因此，在冻结或降低前一层的学习率的同时，一次训练一层可能会更好
知识迁移技术。
（a）辅助知识迁移，（b）联合知识迁移，（c）渐进知识迁移。
实验结果
作者以三种场景下评估了他们提出的MobileBERT：具有25M参数的主模型（MobileBERT），未经操作优化的相同模型（MobileBERT w / o OPT）以及仅有1500万参数的模型（MobileBERT-tiny）。将这些模型与基准算法（例如ELMo，GPT和BERT-base）以及相关的蒸馏工作（BERT-PKD，DistilBERT和TinyBERT）进行了比较。
结果表明，通过逐步的知识迁移过程来训练MobileBERT的这些变体是最有效的，该过程始终以显着优势领先于其他两个。
GLUE基准上的实验结果。
我们发现，不带OPT的MobileBERT的性能比基于BERT的大得多，平均GLUE得分要高出0.2，模型要小4倍。另一方面，MobileBERT仅比BERT基准落后0.6点，而推理时间则快得多-Pixel 4手机上128个token序列的推理时间为62毫秒！但是，由于其性能大大优于GTP和ELMo，因此其性能仍然具有竞争力。
因此，可以肯定地说，可以创建一个精简模型，该模型在资源受限的设备上既高效又快速！
与TinyBERT相比，MobileBERT-tiny可获得更好的性能。但是，当你考虑如何针对TinyBERT对GLUE任务进行微调时，这确实变得更加令人印象深刻。请记住，在进行这项工作之前，由于student模型的能力有限，因此无法对其进行微调。因此，在将TinyBERT的teacher模型的BERT知识基础提炼成TinyBERT之前，必须对其进行微调！这是不是对MobileBERT的情况。
MobileBERT已在GLUE上进行了微调，这证明可以通过本文提出的蒸馏过程创建任务不可知模型！
结论
MobileBERT 在转换模块中引入了bottleneck，这使我们可以更轻松地将知识从大teacher模型传授给小student模型。这种技术使我们可以减小student模型的宽度而不是深度，这可以产生更强大的模型。该模型强调了这样一个事实，即可以创建一个student模型，该模型本身可以在初始蒸馏过程之后进行微调。
此外，结果还表明，实践也是如此，因为MobileBERT能够在GLUE上达到BERT基于99.2％的性能，而Pixel 4手机的参数减少4倍，推理速度提高5.5倍。
往期精品内容推荐
给点提示，就可以自动续写故事！NLP最强文本生成模型GPT2！
2020年AI统计学新书-《数据科学、机器学习、人工智能基础》免费pdf分享
自然语言处理任务相关经典论文、免费书籍、博客、tf代码整理分享