ACL2020论文阅读笔记-FastBERT: a Self-distilling BERT with Adaptive Inference Time

0. 背景

题目：
FastBERT: a Self-distilling BERT with Adaptive Inference Time
机构：北大、腾讯、北师大
作者：Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Haotang Deng, Qi Ju
论文地址：https://arxiv.org/abs/2004.02178
收录会议：ACL2020
代码：https://github.com/autoliuweijie/FastBERT

摘要

预训练模型如BERT所取得的成果举世瞩目。但是如何在确保模型性能的前提下提高模型效率，是近年来的研究热点。为此本文提出 FastBERT。FastBERT是一种速度可调具有自适应推理的模型。推理速度可根据不同需求灵活调整，避免样本冗余计算。此外，该模型在微调方面采用了独特的自精馏机制，进一步提高了计算效率，使性能损失最小。FastBERT在12个中英文数据集中取得了很好的结果。如果给定不同的加速阈值以进行速度性能折衷，那么FastBERT能够比BERT加速1-12倍。

介绍

预训练模型的成功众所周知，从最初的BERT，不断升级进化，做大做强，包括GPT2、XLNet、RoBERTa、T5、UER等等。然而，在许多实际的场景中，这些预训练模型通常需要较多的计算资源。针对这一问题，常见的解决方案包括有
(1)量化，把FP32改成FP16甚至INT8。
(2)剪枝，剪掉多余的连接、多余的注意力头甚至是去掉部分的网络层。
(3)蒸馏，用一个较小的Student模型来学习较大Teacher模型的知识。
通过查验许多NLP数据集，发现这些样本有不同的难度级别。较大的模型可能会对简单的样本输入过多计算，而较小的模型难以应付复杂的样本。最近的研究表明预训练模型中确实存在冗余，因此设计一个适合所有样本难度的单一化模型是有用的，该模型可以满足不同复杂度的样本，并在精度损失最小的情况下获得计算效率。

基于这种目的，本文提出FastBERT：一个具有样本自适应机制的预训练模型。FastBERT 可以动态调整执行层的数量，以减少计算量。换句话说，就是自适应调整每个样本的计算量，难度低的样本通过一两层就可以预测出来，难度较高的样本则需要继续走到更深的层级甚至是整个全程。此外，该模型也有一个独特的自蒸馏处理机制，以尽量少改变模型结构，在单一框架内实现更快且准确的结果。与笨重的BERT相比，FastBERT加速了2到11倍，且结果精度也极具竞争优势(精度只是略有损失而已)。
小结下：
(1)本文改进了BERT，提出了一种实用的速度可调模型，即FastBERT。
(2)首次将样本自适应机制和自蒸馏机制相结合，并提高了模型的推理速度。同时在12个NLP数据集上验证了其有效性。

模型

模型架构

FastBERT有2部分组成：主干和分支。主干由12层的 Transformer encoder 和一个额外的teacher-classifier组成；而分支是附加到每个Transformer输出的student-classifiers。具体模型结构如Figure 2所示：

主干由三部分组成：embedding层、encoder(由Transformer堆叠起来的blocks组成)和teacher classifier。embedding层和encoder没什么特别的，与BERT相同，都是word、position和segment embedding的和。跟在encoder最后的是teacher-classifier，它提取域内特征用于下游推理。

分支，即每个student classifier，是接到每个Transformer block的输出，使得难度低的样本可以提前结束计算输出结果。为了平衡模型的准确性和推理速度。

本文精心设计了student classifier，因为过于简单的网络可能会影响推理的准确性，而过于复杂的分类器会严重降低了推理速度。

模型训练

FastBERT分别训练主干和分支，在训练另一个模块时，一个模块中的参数是固定住的。模型训练包括三个步骤:主干预训练、主干微调和student classifier的自蒸馏。

主干预训练：
与BERT的各种版本模型一样，比如BERT-WWM、RoBERTa 等等模型的预训练方法和结果都可以直接拿来用。所以FastBERT可以不需要自己进行预训练，直接加载各种已经公布的高质量预训练模型。需要注意的是，由于teacher-classifier仅用于推理，因此预训练过程不受影响。

主干微调：
对于每个下游任务，用对应的任务数据输入到模型，对主干进行和teacher-classifier微调。注意，此时尚未启用student classifier。

分支自蒸馏：
经过上述预训练和微调的主干对知识具有良好的表征和抽取，其输出是高质量的soft-label，既包含原始embedding，又包含通用知识。在此基础上蒸馏出student classifiers。由于student classifier之间是相互独立的，他们的预测结果分别与teacher-classifier的soft-label进行比较，其差异用KL散度来衡量。FastBERT中一共有L-1个student classifier，所以自蒸馏的总损失即为它们KL散度的总和。

预训练和微调阶段都只更新主干参数，蒸馏阶段固定住主干参数，用student classifier蒸馏主干中teacher classifier 的概率分布。由于在这个阶段仅仅需要teacher classifier的输出，那么可以自由地使用数量无限的未标记数据，而无需局限于有标记的数据。这意味着只要teacher classifier允许，总是可以提升student classifier，以尽可能的拟合teacher。此外，以前的蒸馏方法是用两个模型去做，一个模型学习另一个模型的知识，而FastBERT的自蒸馏与之不同，因为teacher和student在同一个模型。此所谓"自蒸馏"便是这个意思，自己(分支)蒸馏自己(主干)的知识。

模型自适应推理

在推理阶段，FastBERT可以根据样本的复杂性调整模型中执行的encoding层的数量。在每一Transformer层，用归一化的熵衡量每个样本的当前推理是否可靠到足以终止。归一化熵定义如下：
$$\text{Uncertainty}=\frac{{\sum }_{i=1}^N{p}_s(i)\log p_s(i)}{\log \frac{1}{N}}$$
其中$p_s$表示输出的概率分布，N表示label的类别。
通过不确定性的定义，提出一个重要的假设：
假设1。LUHA：不确定度越低，准确度越高。
定义1。Speed：区分高低不确定性的阈值。
当样本的Uncertainty得分低于Speed则输出，无需再后传，否则再进入下一层。直观上看，速度越快，发送到更深层的样本应该就越少，整体推理速度越快，反之亦然。因此，Speed可以作为权衡推理精度和效率的阈值。

实验结果

在12个NLP数据集(6个英文和6个中文)上对FastBERT进行评测。

分析FLOPs：
所谓的FLOPs(floating point operations)是指每秒所执行的浮点运算次数，可以用来衡量模型计算复杂度。FLOPs与模型的运行环境(CPU、GPU或TPU)无关，反映模型计算的复杂性。一般来说，模型的FLOPs越大，推理时间就会越长。在同样的精度下，更低FLOPs的模型效率更高，更适合生产环境。
Table 1中列出了Transformer和分类器这两种结构的FLOPs，从中可以看出，分类器的计算负荷(FLOPs)要比Transformer轻得多。这是FastBERT加速的基础，虽然它增加了额外的分类器，但减少Transformer的计算量来实现加速。

性能对比：
BERT、DistilBERT和FastBERT在12个数据集上对比结果如 Table 2所示。

从Table 2可以看出，Speed=0.1时，FastBERT在多数数据集上比BERT加速2-5倍且无损精度。如果允许精确度稍有损失，FastBERT比BERT快7到11倍。FastBERT甚至在部分任务上(如Weibo)还有细微提升。相比之下DistillBERT的性价比就不是这么高了，损失的性能换得的速度提升不如FastBERT多。

总结

本文提出一个快速版BERT，即FastBERT。具体来说，FastBERT在训练阶段采用一种自蒸馏机制，在推理阶段采用一种自适应机制，提高了推理效率且准确度损失极少。此外，FastBERT在工业场景中有一个非常实用的特性，即推理速度可调。实验表明，在不降低性能的情况下，FastBERT可以比BERT快2到3倍。若允许精度略有损失，模型推理加速在1到12倍。此外，FastBERT仍然与其他类BERT模型的参数设置兼容，即可以自由加载其他如BERT-WWM、ERNIE和RoBERTa等预训练的模型。