DL-Paper精读：MobileBERT

MobileNERT: a Compact Task-Agnostic BERT for Resource-Limited Devices

https://arxiv.org/abs/2004.02984

Background

BERT在NLP领域的地位是举足轻重的，其预训练模型，在多种下游任务的迁移工作中都能给出非常好的效果。但于此同时，BERT也受困于其庞大的模型参数和较慢的运行速度，尤其是在于一些资源受限的移动端部署任务中，BERT的实际应用是非常受限的。

Related work and the limit

目前，针对BERT的压缩研究工作已有很多。比如distill BERT等工作，通过知识蒸馏，来获得一个高精度的小模型。但是这些工作一般都是基于具体下游任务场景的（“task-specifically”），这些工作的一般流程是首先将预训练的BERT模型在具体任务中进行fine-tune，然后再进行蒸馏来获得小模型。该过程是比较耗时且繁琐的，不如直接来fine-tune一个任务无关（“task-agnostic”）的小的预训练模型。但是如果直接预训练一个较窄或较浅的BERT，又会造成较严重的精度损失，无法达到原始BERT的精度。

Novel points

在预训练的阶段来进行知识蒸馏，获得一个压缩版的小BERT(mobileBERT)，能够达到可观的精度，同时是task-agnostic的。同时为了能够有效地蒸馏mobileBERT，本文设计了一个带反瓶颈结构（inverted-bottleneck）的近似BERT_large的网络（IB-BERT）。在实际训练中，首先充分训练IB-BERT作为teacher model，然后再对mobileBERT进行知识蒸馏。

Methodology

该部分主要分为结构设计和训练方法两个方面。

结构设计：

包括mobileBERT和IB-BERT两个网络，具体根据以下图表来呈现，图1分别为BERT/IB-BERT/mobileBERT的模块结构图，表1给出了整体的网络参数设计。可以看出IB-BERT和mobileBERT与BERT_L有着相同的层数，即深度相同，但mobileBERT每一层的宽度（width）则缩减了8倍，这也就造成了它相比于BERT_L有着接近20倍的参数缩减，同时还能保持着可观的精度。此外，具体的结构方面的改进主要包括以下4个方面：

bottleneck/inverted-bottleneck结构的使用。由图标可以看出，IB-BERT和MobileBERT的模块前后都添加了一个linear层。其中IB-BERT采用了一个Inverted-Bottleneck，将输入的512维度放大为1024，从而匹配BERT_L的模型尺寸，来获得相近的精度；而MobileBERT则采用了一个Bottleneck结构，用来将输入的512维度缩小为128，进一步缩减模型参数量和计算量。IB-BERT和MobileBERT的模块输入输出都是512，因此能够在每一层进行feature map的蒸馏。

堆叠的FFN操作。在MobileBERT的模块中，对FFN操作进行了多次堆叠。在传统的Transformer结构中，MHA（multi-head attention）和FFN（feed-forward network）都扮演着重要的角色，其中MHA能够提取获得各部分的注意力，而FFN能够增加模型的非线性能力，同时在传统BERT中，这两者的参数量总是保持1：2。但是在MobileBERT的结构设计中，因为宽度的变化，参数比例发了改变，MHA的参数相对变多了。因此文中提出通过堆叠FFN的方法来解决该问题，MobileBERT中，MHA之后有连续4个FFN。

文中通过分析发现LN和Gelu操作对于latency的影响较大，因此提出移除LN层，替换为文中所提出的一种element-wise的线性转换操作NoNorm，同时将Gelu换为Relu。（虽然文中这样提出了，但考察github上的几个实现代码，发现一般还是在使用LN和Gelu）

由表格可以看出，在embedding阶段，为了减少原BERT中embedding的庞大参数量，文中采用了一个 emb_128 + 512_conv的操作来代替。

网络训练

该部分主要针对IB-BERT -> MobileBERT的蒸馏训练设计。主要分为两部分。

Loss设计

由于两个网络中bottleneck/inverted-bottleneck结构的设计，所以文中提出可以在每一层都可以加入feature map transfer(FMT)和Attention Transfer(AT)。其中FMT是通过计算教师网络IB-BERT和学生网络MobileBERT每一层的输出的MSE来获得，AD是通过计算两者MHA中attention的分布的KL散度而获得，具体计算公式参考论文。两种transfer可以如图1中的虚线所示。因此，该预训练阶段的distillation训练的Loss函数，由三部分组成：传统的masked language modeling loss(MLM), next sentence prediction(NSP) loss和新的Knowledge Distillation loss(KD)所组成，公式如下：

训练策略

许多常见的蒸馏训练的技巧，包括 Auxiliary Knowledge Transer，Joint Knowledge Transfer, Progressive Knowledge Transfer等，具体见文中介绍。

Evaluation(benchmark, experiments design)

对NLP任务不够熟悉，因此不对实验部分进行过多的评估，GLUE上的结果如图所示。总的来说，相比于BERT_B，MobileBERT能够在保持相近精度的情况下，实现4.3×的压缩和5.5×的加速。

Thoughts:

1. Describe what the authors of the paper aim to accomplish, or perhaps did achieve.

对BERT进行压缩，给出一个精度较高且任务无关的压缩模型

2. If a new approach/ technique/ method was introduced in a paper, what are the key elements of the newly proposed approach?

通过蒸馏来获得一个高精度的预训练MobileBERT。为了能够获得较好的蒸馏效果，首先设计了一个接近的BERT_L的teacher模型，同时通过巧妙地bottleneck/inverted-bottleneck结构的使用，使得大小模型之间能够实现直接的信息transfer。

此外，详尽复杂的蒸馏策略设计等也对结果有着重要的影响。