MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices（2020-4-6）

模型介绍

近年来，自然语言处理(NLP)通过使用具有数亿个参数的巨大预训练模型取得了巨大的成功。然而，这些模型受到沉重的模型尺寸和高延迟的影响，因此无法部署到资源有限的移动设备上。因此这里提出了MobileBERT来压缩和加速流行的BERT模型。与原始的BERT一样，MobileBERT是任务不可知论的，也就是说，它可以通过简单的微调通用地应用于各种下游NLP任务。

基本上，MobileBERT是${\mathrm{BERT}}_{\mathrm{LARGE}}$的精简版，同时配备了瓶颈结构和精心设计的在自注意和前馈网络之间的平衡。为了训练MobileBERT，首先训练了一个专门设计的教师模型，一个包含反转瓶颈的${\mathrm{BERT}}_{\mathrm{LARGE}}$模型。然后，将这个老师的知识转移到MobileBERT。实证研究表明，MobileBERT比${\mathrm{BERT}}_{\mathrm{BASE}}$小4.3倍，快5.5倍，同时在知名基准测试中取得了具有竞争力的结果。

模型结构

MobileBERT被设计成与${\mathrm{BERT}}_{\mathrm{LARGE}}$一样深，而每一层都通过采用瓶颈结构和平衡自关注和前馈网络而变得更窄。为了训练一个又深又瘦的模型MobileBERT，首先训练一个专门设计的教师模型，它是一种结合了反瓶颈（Inverted-Bottleneck）的${\mathrm{BERT}}_{\mathrm{LARGE}}$模型(IB-BERT)。然后，将知识从IB-BERT转移到MobileBERT。

上图中，(a) 为BERT，(b)为反向瓶颈 BERT (IB-BERT)，( c)为MobileBERT。在 (b) 和 ( c) 中，红线表示块间流动，蓝线表示块内流动。MobileBERT是通过层层模仿 IB-BERT 来训练的。

模型改进

Bottleneck and Inverted-Bottleneck

MobileBERT 模型的深度和 ${\mathrm{BERT}}_{\mathrm{LARGE}}$ 一样，但是每个 block 都更小。如上表所示，每一个 block 的隐藏层维度为128。另一方面，为每个block引入两个线性变换，以将其输入和输出维度调整为512。把这样的结构称为瓶颈（Bottleneck）。

训练这样一个又深又薄的网络是一项挑战。为了克服训练问题，首先构建一个教师网络并对其进行训练直到收敛，然后将知识从这个教师网络转移到MobileBERT。实验发现这比直接从零开始训练MobileBERT要好得多。事实上，教师网络就是${\mathrm{BERT}}_{\mathrm{LARGE}}$，只是用反向瓶颈结构进行了增强，以调整其特征映射大小为512。在接下来的内容中，将教师网络称为${IB-BERT}_{\mathrm{LARGE}}$。注意，IB-BERT和MobileBERT的特征映射大小相同，都是512。因此，可以直接比较IB-BERT和MobileBERT的分层输出差异。在知识转移策略中需要这样的直接比较。

值得指出的是，同时引入瓶颈结构和反向瓶颈结构的结果是相当灵活的体系结构设计。可以只使用MobileBERT的瓶颈(相应地，教师变成了${\mathrm{BERT}}_{\mathrm{LARGE}}$，或者只使用IB-BERT的反向瓶颈(那么MobileBERT中就没有瓶颈了)来对齐他们的特征映射。当使用这两种方法时，可以允许${IB-BERT}_{\mathrm{LARGE}}$在保持MobileBERT足够紧凑的同时保持${\mathrm{BERT}}_{\mathrm{LARGE}}$的性能。

Stacked Feed-Forward Networks

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

为了解决这个问题，应该在MobileBERT中使用堆叠前馈网络来重新平衡MHA和FFN之间的相对大小。每个MobileBERT层包含一个MHA和几个堆叠的FFN，在MobileBERT中，每个MHA之后使用4个堆叠FFN。

Operational Optimizations

通过模型延迟分析，发现层归一化和gelu激活函数占总延迟的比例相当大。因此，建议用MobileBERT中的新操作替换它们。

Remove layer normalization：将n通道（维）隐藏状态h的层归一化替换为元素的线性变换。

$\mathrm{NoNorm}(h)=\gamma \circ h+\beta$

其中，$\gamma ,\beta \in R^n$，$\circ$ 表示Hadamard积。注意NoNorm与LayerNorm具有不同的属性，即使在测试模式下也是如此，因为原始的层归一化并不是对一批向量的线性操作。

Use relu activation：用更简单的relu激活函数来代替gelu激活函数。

Embedding Factorization

BERT模型中的嵌入表占模型大小的相当大的比例。为了压缩嵌入层，将MobileBERT中的嵌入维数减少到128。然后，对原始词嵌入应用内核大小为3的一维卷积，以生成512维的输出。

Training Objectives

建议使用以下两个知识转移目标：特征图转移（feature map transfer） 和注意力转移（attention transfer）。对第 $\ell$ 层的最终分层知识转移损失 $L_{KT}^{\ell }$ 是以下两个目标的线性组合：

1、Feature Map Transfer (FMT)

由于BERT中的每一层仅仅将前一层的输出作为输入，因此在分层知识转移中，最重要的是每一层的特征映射应该尽可能接近老师的特征映射。具体而言，使用MobileBERT学生和IB-BERT老师的特征图之间的均方误差作为知识转移目标：

其中，$\ell$ 为层索引，$T$ 为序列的长度，$N$ 为特征映射大小。在实践中发现将该损失项分解为归一化特征图差异和特征图统计差异有助于稳定训练。

2、Attention Transfer (AT)

注意机制极大地提高了NLP的性能，成为Transformer和BERT的最重要的 block。鉴于此，使用来自优化良好的老师的 self-attention 映射来帮助训练MobileBERT增强特征映射。具体的，通过最小化MobileBERT 学生和 IB-BERT 教师的self-attention 每一头的分布之间的KL散度：

其中，$A$ 为注意力的头数。

3、Pre-training Distillation (PD)

在对MobileBERT进行预训练时，除了分层的知识转移，还可以使用知识蒸馏损失。使用原始的屏蔽语言建模(MLM)损失、下句预测(NSP)损失和新的 MLM 知识蒸馏(KD)损失的线性组合作为训练前的蒸馏损失：

其中，$\alpha$ 是一个在（0，1）之间的超参数。

Training Strategies

根据上述目标，在培训中可以有各种组合策略。这里讨论了三种策略。

上图中，（a）为 auxiliary knowledge transfer (AKT)，（b）为 joint knowledge transfer (JKT)，（c）为 progressive knowledge transfer (PKT)。颜色较浅的 block 表示这个阶段被冻结了。

1、Auxiliary Knowledge Transfer

在该策略中，将中间知识转移作为知识蒸馏的辅助任务。然后使用单一损失，它是来自所有层的知识转移损失和训练前蒸馏损失的线性组合。

2、Joint Knowledge Transfer

然而，IB-BERT老师的中间知识(即注意力映射和特征映射)可能不是 MobileBERT 学生的最佳解决方案。因此，应该将这两个损失项分离，首先用所有分层的知识转移损失联合训练MobileBERT，然后通过预训练蒸馏进一步训练。

3、Progressive Knowledge Transfer

一个可能会关心的点是如果MobileBERT不能完美地模仿IB-BERT老师，来自下层的错误可能会影响高层的知识转移。因此，在知识转移中逐步训练每一层。将渐进式知识转移分为 L 个阶段，其中 L 为层数。

对于联合知识转移和渐进式知识转移，在分层知识转移阶段，最开始的嵌入层和最后的分类器不存在知识转移。它们从IB-BERT老师复制到MobileBERT学生。此外，对于渐进式知识转移，当训练第 $\ell$ 层时，将所有可训练的参数冻结在下面的层中。在实践中，我们可以通过以下方法来软化训练过程。当训练一个层时，以较小的学习率进一步调整较低的层，而不是完全冻结它们。

模型参考

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

代码地址：https://github.com/tensorflow/models/tree/master/official/nlp/projects/mobilebert