预训练模型BERT在自然语言处理的各项任务中都展现了不错的效果。但在移动手机普及的今天，如何在移动端或者资源受限的设备上使用BERT模型，一直是个挑战。最近，Google Brain团队提出了MobileBERT模型，对BERT模型进行压缩，在损失效果不大的前提下，使得模型大小减少了3-4倍，而速度提升了4-5倍，让各种NLP应用程序都能在移动端轻松部署。相关论文已被ACL 2020接收。本文是AI前线第108篇论文导读，我们将对这项研究工作进行详细解读。

概览

近几年，NLP社区见证了一场预训练监督模型的革命。这些模型通常有数亿个参数，在这些模型中，BERT的精度提升最为显著。然而，作为NLP中最大的模型之一，BERT的模型容量大、时延高，使得资源有限的移动设备无法在基于移动的机器翻译、对话建模等方面发挥BERT的能力。

目前业界已经做了一些努力，特别是针对特定任务对BERT模型进行蒸馏。据我们所知，目前还没有任何工作能构建出一个与任务无关的轻量级预训练模型，也就是说，它需要跟原始BERT一样能在不同的下游NLP任务上进行一般的微调。在本文中，我们提出了MobileBERT来填补这一空白。在实际应用中，我们通常需要对BERT进行任务无关的压缩。特定任务的压缩需要先将原始的大BERT模型微调为特定任务的教师，然后再蒸馏。与直接微调任务无关的压缩模型相比，这样的过程要复杂得多，成本也更高。

乍一看，获得任务无关的压缩BERT似乎很简单。例如，可以取一个较窄或较浅版本的BERT，通过最小化预测损失和蒸馏损失的凸组合来训练它直到收敛。不幸的是，实验结果表明，这种直接的方法会导致精度损失严重。这也许并不奇怪，众所周知，浅层网络通常没有足够的表示能力，而窄层和深层网络则很难训练。

我们的MobileBERT被设计成和BERT（LARGE）一样深，而每一层都通过采用bottleneck结构和平衡注意力机制和前馈网络进而使其变得更加狭窄（图1）。为了训练MobileBERT，我们首先训练了一个特殊设计的教师模型，一个包含BERT（LARGE）模型（IB-BERT）的inverted-bottleneck模型。然后，我们将知识从IB-BERT迁移到MobileBERT。我们在实证研究中仔细研究了各种知识迁移策略。

实证评估表明，MobileBERT比BERT（BASE）小4.3倍、快5.5倍，并且可以在知名的NLP基准上获得可观的结果。在GLUE的自然语言推理任务中，MobileBERT的GLUE得分为77.7，仅比BERT（BASE）低0.6，在Pixel 4手机上的延迟为62ms。在SQuAD v1.1/v2.0问答任务中，MobileBER获得的dev F1得分为90.3/80.2，甚至比BERT（BASE）高出1.5/2.1。

MobileBERT原理详解

在本节中，我们将详细介绍MobileBERT的体系结构设计和高效训练MobileBERT的训练策略。具体的模型配置汇总在表1，这些配置是通过大量的结构搜索实验获得的，将在第4.1节中介绍。

瓶颈层和倒瓶颈层

MobileBERT的体系结构如图1（c）所示。它和BERT（LARGE）一样深，但每个构建块都小得多。如表1所示，每个构建块的隐状态大小仅为128。另一方面，我们为每个构建块引入两个线性变换，以将其输入和输出维度调整为512。根据（He等人，2016）中的术语，我们将这种架构称为瓶颈层。

训练这样一个深并且瘦小的网络是很大的挑战。为了解决这一问题，我们首先构建了一个教师网络，并对其进行训练直至收敛，然后将知识从这个教师网络迁移到MobileBERT。我们发现这比从头开始直接训练MobileBERT的效果要好得多。各种训练策略将在后面的章节中讨论。这里，我们介绍教师网络的架构设计，如图1（b）所示。事实上，教师网络就是BERT（LARGE），同时增加了倒瓶颈结构（Sandler等人，2018），将其feature map大小调整为512。在下文中，我们将教师网络称为IB-BERT（LARGE）。请注意，IB-BERT和MobileBERT具有相同的feature map大小，即512。因此，我们可以直接比较IB-BERT和MobileBERT的layer-wise输出的差异。我们的知识迁移战略需要这样一种直接的比较。

值得指出的是，同时引入瓶颈和倒瓶颈结构使得体系结构设计具有相当的灵活性。我们可以只使用MobileBERT的瓶颈（相应地，教师变为BERT）或者只使用IB-BERT的倒瓶颈（MobileBERT中没有瓶颈）来对齐它们的特征映射。然而，当使用这两种方法时，我们可以允许IB-BERT（LARGE）在保证MobileBERT足够紧凑的同时保持BERT（LARGE）的性能。

堆叠式前馈网络

MobileBERT的瓶颈结构带来的一个问题是，多头注意力（MHA）模块和前馈网络（FFN）模块之间的平衡被打破。MHA和FFN在Transformer结构中扮演着不同的角色：前者允许模型共同处理来自不同子空间的信息，而后者增加了模型的非线性。在原始的BERT中，MHA和FFN中的参数之比总是1:2。但在瓶颈结构中，MHA的输入来自更宽的特征映射（块间大小），而FFN的输入来自更窄的瓶颈（块内大小）。这导致MobileBERT中的MHA模块相对包含更多的参数。

为了解决这个问题，我们建议在MobileBERT中使用堆叠式前馈网络来重新平衡MHA和FFN之间的相对大小。如图1（c）所示，每个MobileBERT层包含一个MHA，但多个堆叠的FFN。在MobileBERT中，我们在每个MHA之后使用4个堆叠FFN。

操作优化

通过模型延迟分析，我们发现，layer层归一化（Ba等人，2016）和gelu激活函数（Hendrycks和Gimpel，2016）占总延迟的相当大比例。因此，我们提议在MobileBERT中用新的操作来取代它们。

移除层归一化 我们将n-channel的隐状态h的层归一化替换为element-wise的线性变换：

式中，γ，β∈Rn和o表示哈达玛积。请注意，即使在测试模式下，NoNorm也具有与LayerNorm不同的属性，因为原始层归一化不是一个批向量的线性操作。

使用relu激活 我们用更简单的relu激活代替了gelu激活（Nair和Hinton，2010）。

Embedding因子分解

BERT模型中的嵌入表占模型大小的很大一部分。为了压缩嵌入层，如表1所示，我们将MobileBERT中的嵌入维度减少到128。然后，我们对原始token嵌入应用核大小为3的一维卷积来产生512维的输出。

训练目标

我们提出使用以下两个知识迁移目标，即特征映射迁移和注意力迁移，来训练MobileBERT。图1说明了提出的layer-wise的知识迁移目标。我们的最后对于第l层的知识迁移损失L(KT)(l)是以下两个目标的线性组合：

特征映射迁移（FMT） 由于BERT中的每一层只把前一层的输出作为输入，因此在layer-wise的知识迁移中最重要的是每一层的特征映射应该尽可能接近教师的特征映射。特别地，MobileBERT学生和IB-BERT教师的特征映射之间的均方误差被用作知识迁移目标：

其中，l是层的索引，T是序列长度，N是特征映射大小。在实际应用中，我们发现将这一损失项分解为归一化的特征映射差异和特征映射统计差异有助于训练的稳定性。

注意力迁移（AT） 注意机制极大地提高了NLP的性能，并成为Transformer和BERT中的一个关键构建块（Clark et al.，2019a；Jawahar et al.，2019）。这促使我们使用来自最优化教师的自注意力映射，以帮助MobileBERT在增强特征映射迁移方面的训练。特别是，我们最小化了MobileBERT学生和IB-BERT教师的每个注意力分布之间的KL距离：

其中A是注意力的数量。

预训练蒸馏 (PD) 除了layer-wise的知识迁移，我们也可以在训练前使用知识蒸馏损失。我们使用原始掩码语言建模（MLM）损失、下一句预测（NSP）损失和新的MLM知识蒸馏（KD）损失的线性组合作为预训练的蒸馏损失：

其中α是（0，1）中的超参数。

训练策略

鉴于上述目标，在训练中可以有各种组合策略。本文讨论了三种策略。

辅助知识迁移 在该策略中，我们将中间知识迁移作为知识蒸馏的辅助任务。我们使用单个损失，它是来自所有层的知识迁移损失和预训练蒸馏损失的线性组合。
联合知识迁移 然而，IB-BERT教师的中间知识（即注意力映射和特征映射）可能不是MobileBERT学生的最佳解决方案。因此，我们建议将这两个损失项分离开来，首先对MobileBERT进行layer-wise知识迁移损失联合训练，然后通过训练前蒸馏进一步训练。
渐进式知识迁移 人们还可能担心，如果MobileBERT不能完全模仿IB-BERT教师，来自较低层的错误可能会影响较高层的知识迁移。因此，我们建议在知识迁移的各个层面上逐步进行训练。渐进式知识迁移分为L个阶段，其中L是层的数目。

图2展示了这三种策略的示意图。对于联合知识迁移和渐进知识迁移，在layer-wise知识迁移阶段，初始嵌入层和最终分类器没有知识迁移。它们被从IB-BERT老师复制了到MobileBERT学生。此外，对于渐进式的知识迁移，当我们训练第l层时，我们会冻结下面各层中所有可训练的参数。在实践中，我们可以将训练过程软化如下。当训练一个层时，我们进一步调整较低层的学习速度，而不是完全冻结。

实验

在这一部分中，我们首先介绍了我们的结构搜索实验，这些提供了表1中的模型设置，然后介绍了MobileBERT和各种基线的基准测试的经验结果。

模型设置

我们进行了广泛的实验，为IB-BERT老师和MobileBERT学生寻找良好的模型设置。我们从SQuAD v1.1 dev F1分数开始，作为搜索模型设置的性能指标。在这一部分中，我们只训练每个模型125k步，2048个批次，这是原始BERT训练计划的一半。

IB-BERT的体系结构搜索 我们对于教师模型的设计理念是尽可能使用小的块间隐藏大小（特征映射大小），只要不存在精度损失。在这个指导下，我们设计实验来操作BERT（LARGE）大小的IB-BERT的块间大小，结果如表2所示，带有标签（a）-（e）。我们可以看到，减小块间隐状态大小不会损害BERT的性能，直到它小于512。因此，我们选择块间隐状态大小为512的IB-BERT（LARGE）作为教师模型。

人们可能会怀疑我们是否也可以缩小教师的块内隐大小。我们进行了实验，结果如表2所示，带有标签（f）-（i）。我们可以看到，当块内隐大小减小时，模型的性能会大大降低。这意味着块内隐大小在BERT中起着至关重要的作用，它代表了非线性模块的表示能力。因此，与块间隐大小不同，我们不缩减教师模型的块内隐大小。

MobileBERT的架构搜索 我们为BERT（BASE）寻求4倍的压缩比，为了选择一个好的MobileBERT 学生模型，因此我们设计了一组MobileBERT模型，这些模型的近似为25M参数，但MHA和FFN中的参数个数之比不同。表3显示了我们的实验结果。它们在MHA和FFN之间有不同的平衡。从表中可以看出，当MHA和FFN中的参数之比为0.4∼0.6时，模型的性能达到峰值。这也许可以解释为什么原始Transformer选择MHA和FFN的参数比为0.5。

考虑到模型的准确性和训练效率，我们选择128个块内隐大小和4个层叠FFN的结构作为MobileBERT学生模型。我们还相应地将教师模型中的注意力头数量设置为4，为layer-wise知识迁移做准备。表1展示了IB-BERT（LARGE）教师和MobileBERT学生的模型设置。

人们可能会怀疑，减少头的个数是否会损害教师模型的表现。通过比较表2中的（a）和（f），我们可以看出，将头数量从16个减少到4个并不影响IB-BERT(LARGE)的性能。

实现细节

继BERT（Devlin et al，2018）之后，我们使用了BookCorpus（Zhu et al.，2015）和英文的 Wikipedia作为我们的预训练数据。为了使IB-BERT（LARGE）教师达到与原始BERT(LARGE)相同的精度，我们在256 TPU v3芯片上对IB-BERT（LARGE）进行500步训练，批量大小为4096，采用LAMB优化器（You等人，2019）。

为了与最初的BERT进行公平的比较，我们不在其他BERT变体中使用训练技巧（Liu et al，2019b；Joshi et al，2019）。对于MobileBERT，我们在训练前蒸馏阶段使用相同的训练计划。此外，我们使用渐进式的知识迁移来训练MobileBERT，它在24层上额外增加了240k步数。在消融研究中，我们将MobileBERT的预训练蒸馏计划减半，以加速实验。此外，在知识迁移策略的消融研究中，为了公平比较，联合知识迁移和辅助知识迁移也采取了额外的240k步。

对于下游任务，所有报告的结果都是通过简单的微调MobileBERT获得的，就像原始BERT所做的那样。为了对预先训练的模型进行微调，我们在搜索空间中搜索优化超参数，包括不同的批大小（16/32/48）、学习率((1-10)*e-5)和epoch数（2-10）。搜索空间不同于原始的BERT，因为我们发现MobileBERT通常需要更高的学习率和更多的微调训练时间。我们根据它们在验证集上的性能来选择测试模型。

在GLUE的结果

通用语言理解评估（GLUE）基准（Wang等人，2018）是9项自然语言理解任务的集合。我们比较了MobileBERT和BERT（BASE），以及GLUE排行榜上一些最先进的预训练BERT模型：OpenAI GPT（Radford等人，2018）和ELMo（Peters等人，2018）。我们还比较了最近提出的三种压缩BERT模型：BERT-PKD（Sun et al，2019）和DistilBERT（Sanh et al.，2019）。为了进一步显示MobileBERT模型相对于最近的small BERT模型的优势，我们还评估了MobileBERT（TINY）模型的一个较小的变量，该模型的参数约为15M，它减少了每个层中的ffn数量，并使用了较轻的MHA结构。此外，为了验证MobileBERT在实际移动设备上的性能，我们使用TensorFlow Lite API导出模型，并在固定序列长度为128的4个线程的 Pixel 4 phone 上测量推理延迟。结果见表4。

从表中我们可以看出，MobileBERT在GLUE基准上非常有竞争力。MobileBERT的GLUE总分为77.7分，仅比BERT（BASE）低0.6分，比BERT（BASE）小4.3倍，快5.5倍。此外，它在4.3×更小的模型尺寸下，比强大的OpenAI GPT基线性能好0.8个GLUE分数。它的性能也优于其他所有较小或相似大小的压缩BERT模型。最后，我们发现引入的操作优化对模型性能造成了一定的影响。如果没有这些优化，MobileBERT甚至可以比BERT(BASE)高出0.2个GLUE分数。

在SQuAD的结果

SQuAD是一个大型阅读理解数据集。SQuAD1.1（Rajpurkar等人，2016）仅包含给定上下文有答案的问题，而SQuAD2.0（Rajpurkar等人，2018）包含无法回答的问题。我们只在SQuAD 验证数据集上评估MobileBERT，因为在SQuAD测试排行榜上几乎没有单一的模型提交。我们将我们的MobileBERT与BERT(BASE)、DistilBERT和强大的基线DocQA进行了比较（Clark和Gardner，2017）。

如表5所示，MobileBERT的表现优于其他较小或相似的模型。

量化

我们将TensorFlow Lite中的标准训练后量化应用于MobileBERT。结果见表6。我们发现，虽然量化可以进一步压缩MobileBERT 4倍，但几乎没有性能下降。这表明MobileBERT在压缩方面还有很大的空间。

消融实验

操作优化 我们评估了第3.3节中介绍的两种操作优化的有效性，即用NoNorm替换层归一化（LayerNorm）和用relu替换gelu激活。我们报告推理延迟，使用与第4.6.1节相同的实验设置。从表7可以看出，NoNorm和relu在降低MobileBERT的延迟方面都非常有效，而这两种操作优化并没有减少FLOPS。这揭示了真实世界的推理延迟和理论计算开销（即FLOPS）之间的差距。

训练策略 研究了辅助知识迁移、联合知识迁移、渐进知识迁移等训练策略的选择对MobileBERT效果的影响。如表8所示，渐进式知识迁移策略明显优于其他两种策略。我们注意到辅助知识迁移与其他两种策略之间存在显著的效果差距。我们认为原因是教师的中间层知识（即注意力映射和特征映射）可能对学生不是最优的，所以学生需要一个额外的预训练蒸馏阶段来微调其参数。

训练目标 最后，我们进行了一组关于注意迁移（AT）、特征映射迁移（FMT）和预训练蒸馏（PD）的消融实验。在这些实验中，操作优化（OPT）被重新移除，以便在MobileBERT和原始BERT之间进行公平的比较。结果见表9。

我们可以看出，所提出的特征映射变换对MobileBERT的性能提高贡献最大，而注意迁移和预训练蒸馏也起到了积极的作用。我们还可以发现，我们的IB-BERT（LARGE）老师和原来的IB-BERT（LARGE）老师一样强大，而MobileBERT老师相比之下却退化了很多。因此，我们相信MobileBERT仍有很大的改进空间。

结论

我们介绍了MobileBERT，它是BERT的一个任务无关的压缩变种。对流行的NLP基准的实证结果表明，MobileBERT可以与BERT（BASE）相媲美，同时更小更快。MobileBERT可以使各种NLP应用程序轻松地部署在移动设备上。

本文的研究表明：1）保持MobileBERT的深度和厚度是关键；2）瓶颈/倒瓶颈结构能够实现有效的layer-wise知识迁移；3）渐进式知识迁移能够有效地训练MobileBERT。我们相信我们的发现是通用的，可以应用于其他模型压缩问题。

论文原文链接：

https://arxiv.org/pdf/2004.02984.pdf

创作场景

谷歌提出任务无关的轻量级预训练模型 MobileBERT：比 BERT 小 4 倍、速度快 5 倍