MobileBERT: 一种适用于资源有限设备的紧凑型任务无关BERT

摘要

近年来，自然语言处理（NLP）通过使用具有数亿参数的巨大预训练模型取得了巨大成功。然而，这些模型存在模型体积庞大和延迟高的问题，使得它们无法部署到资源有限的移动设备上。在本文中，我们提出了MobileBERT，用于压缩和加速流行的BERT模型。与原始BERT一样，MobileBERT是任务无关的，即它可以通过简单的微调广泛应用于各种下游NLP任务。基本上，MobileBERT是BERTLARGE的瘦身版，同时配备了瓶颈结构，并在自注意力和前馈网络之间进行了精心设计的平衡。为了训练MobileBERT，我们首先训练了一个特别设计的教师模型，即一个结合了倒置瓶颈的BERTLARGE模型。然后，我们从这个教师模型向MobileBERT进行知识转移。实证研究表明，MobileBERT比BERTBASE小4.3倍，快5.5倍，同时在知名基准测试中取得了竞争性的结果。在GLUE的自然语言推理任务上，MobileBERT取得了77.7的GLUE分数（比BERTBASE低0.6），并在Pixel 4手机上实现了62毫秒的延迟。在SQuAD v1.1/v2.0问答任务上，MobileBERT取得了90.0/79.2的开发F1分数（比BERTBASE高1.5/2.1）。

1 引言

NLP社区见证了预训练自监督模型的革命。这些模型通常具有数亿参数（Peters等，2018；Radford等，2018；Devlin等，2018；Radford等，2019；Yang等，2019）。在这些模型中，BERT（Devlin等，2018）显示出显著的准确性提升。然而，作为NLP中最大的模型之一，BERT存在模型体积庞大和延迟高的问题，使得在资源有限的移动设备上部署BERT的力量变得不切实际，例如在移动机器翻译、对话建模等方面。

已有一些努力将BERT特定任务地蒸馏成紧凑模型（Turc等，2019；Tang等，2019；Sun等，2019；Tsai等，2019）。据我们所知，目前还没有任何工作用于构建一个任务无关的轻量级预训练模型，即一个可以像原始BERT一样在不同下游NLP任务上进行通用微调的模型。在本文中，我们提出了MobileBERT来填补这一空白。实际上，BERT的任务无关压缩是可取的。特定任务的压缩需要首先将原始的大型BERT模型微调为特定任务的教师模型，然后进行蒸馏。这样的过程比直接微调一个任务无关的紧凑模型要复杂得多（Wu等，2019）且成本更高。

乍一看，获得一个任务无关的紧凑BERT似乎很简单。例如，可以只取一个更窄或更浅的BERT版本，并通过最小化预测损失和蒸馏损失的凸组合来训练它直到收敛（Turc等，2019；Sun等，2019）。不幸的是，实证结果表明，这种直接的方法会导致显著的准确性损失（Turc等，2019）。这可能并不令人惊讶。众所周知，浅层网络通常没有足够的表示能力，而窄而深的网络难以训练。

我们的MobileBERT设计为与BERTLARGE一样深，同时通过采用瓶颈结构并在自注意力和前馈网络之间进行平衡，使每一层变得更窄（图1）。为了训练MobileBERT，一个深而薄的模型，我们首先训练了一个特别设计的教师模型，即一个结合了倒置瓶颈的BERTLARGE模型（IB-BERT）。然后，我们从IB-BERT向MobileBERT进行知识转移。在我们的实证研究中，仔细研究了各种知识转移策略。

实证评估1表明，MobileBERT比BERTBASE小4.3倍，快5.5倍；同时它仍然可以在知名的NLP基准测试中取得竞争性的结果。在GLUE的自然语言推理任务上，MobileBERT可以取得77.7的GLUE分数，仅比BERTBASE低0.6，在Pixel 4手机上的延迟为62毫秒。在SQuAD v1.1/v2.0问答任务上，MobileBERT取得了90.3/80.2的开发F1分数，甚至比BERTBASE高1.5/2.1。

2 相关工作

最近，BERT的压缩引起了广泛关注。Turc等人（2019）提出预训练较小的BERT模型以改进特定任务的知识蒸馏。Tang等人（2019）将BERT蒸馏成一个极小的LSTM模型。Tsai等人（2019）在序列标注任务上将多语言BERT蒸馏成较小的BERT模型。Clark等人（2019b）使用多个单任务BERT模型来教导一个多任务BERT。Liu等人（2019a）将来自多个BERT模型集合的知识蒸馏到单个BERT中。

与我们的工作同时，Sun等人（2019）通过知识蒸馏和多个中间层隐藏状态的额外知识转移，将BERT蒸馏成更浅的学生模型。Jiao等人（2019）提出了TinyBERT，它也在预训练和微调阶段使用了分层蒸馏策略。Sanh等人（2019）提出了DistilBERT，通过在预训练阶段和可选的微调阶段进行知识蒸馏，成功地将BERT模型的深度减半。

与这些现有文献相比，我们仅在预训练阶段使用知识转移，并且在下游任务中不需要微调的教师或数据增强（Wu等人，2019）。另一个关键区别是，这些先前的工作试图通过减少BERT的深度来压缩它，而我们则专注于通过减少BERT的宽度来压缩它，这已被证明更为有效（Turc等人，2019）。

3 MobileBERT

在本节中，我们介绍了MobileBERT的详细架构设计以及有效训练MobileBERT的训练策略。具体模型设置总结在表1中。这些设置是通过广泛的架构搜索实验获得的，将在第4.1节中介绍。

3.1 瓶颈与倒置瓶颈

MobileBERT的架构如图1©所示。它与BERTLARGE一样深；但每个构建块都变得更小。如表1所示，每个构建块的隐藏维度仅为128。另一方面，我们为每个构建块引入了两个线性变换，以将其输入和输出维度调整为512。根据（He等，2016）中的术语，我们将这种架构称为瓶颈。

训练这样一个深而薄的网络具有挑战性。为了克服训练问题，我们首先构建一个教师网络并训练它直到收敛，然后从这个教师网络向MobileBERT进行知识转移。我们发现这比直接从零开始训练MobileBERT要好得多。各种训练策略将在后面的章节中讨论。在这里，我们介绍教师网络的架构设计，如图1(b)所示。实际上，教师网络就是BERTLARGE，同时增加了倒置瓶颈结构（Sandler等，2018）以将其特征图大小调整为512。在下文中，我们将教师网络称为IB-BERTLARGE。请注意，IB-BERT和MobileBERT具有相同的特征图大小，即512。因此，我们可以直接比较IB-BERT和MobileBERT之间的逐层输出差异。这种直接比较在我们的知识转移策略中是必要的。

值得指出的是，同时引入的瓶颈和倒置瓶颈结构导致了相当灵活的架构设计。可以仅对MobileBERT使用瓶颈结构（相应地教师变为BERTLARGE）或仅对IB-BERT使用倒置瓶颈结构（那么MobileBERT中没有瓶颈）来对齐它们的特征图。然而，当同时使用它们时，我们可以让IB-BERTLARGE保持BERTLARGE的性能，同时使MobileBERT足够紧凑。

3.2 堆叠前馈网络

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

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

3.3 操作优化

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

移除层归一化 我们将n通道隐藏状态h的层归一化替换为逐元素的线性变换：

$$\mathbb{N}o\mathbb{N}orm(\mathbf{h})=\gamma \circ \mathbf{h}+\beta ,\text{\hspace{1em}(1)}$$

其中，$\gamma ,\beta \in {\mathbb{R}}^n$，o表示哈达玛积。请注意，即使在测试模式下，NoNorm与LayerNorm具有不同的性质，因为原始的层归一化对于一批向量来说不是线性操作。

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

3.4 嵌入分解