推理怎么又快又稳？且看我FastBERT

2020 年 7 月 22 日 学术头条

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 的输出，使得难度低的样本可以提前结束计算输出结果。为了平衡模型的准确性和推理速度。

模型训练

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

(1)主干预训练：
与 BERT 的各种版本模型一样，比如 BERT-WWM、RoBERTa 等等模型的预训练方法和结果都可以直接拿来用。所以 FastBERT 可以不需要自己进行预训练，直接加载各种已经公布的高质量预训练模型。需要注意的是，由于 teacher-classifier 仅用于推理，因此预训练过程不受影响。
(2)主干微调：
对于每个下游任务，用对应的任务数据输入到模型，对主干进行和 teacher-classifier 微调。注意，此时尚未启用 student classifier。
(3)分支自蒸馏：

经过上述预训练和微调的主干对知识具有良好的表征和抽取，其输出是高质量的 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 层，用归一化的熵衡量每个样本的当前推理是否可靠到足以终止。归一化熵定义如下：

实验结果

在 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 等预训练的模型。