FastBERT——自蒸馏模型与自适应调整推断时间技术

前言

这几天被分配到的任务是调研现有的几种基于 BERT 的蒸馏模型，所以花了些时间了解了下 DistilBERT，TinyBERT 以及 FastBERT.

自从 BERT 发布以来，各种基于 BERT 的改良版本（如 RoBERTa）层出不穷，模型效果也有着不断的提升，但碍于这些模型过于巨大，需要大量的计算资源，在工业中的应用受限。所以，大家都在想方设法地对 BERT 进行“瘦身”，期望经瘦身后的模型能带来更快的推断速度，并大部分地保留原始 BERT 的推断性能。为了达到此目的，研究者们分成了多个方向：量化（quantinization）、权重剪枝（weights pruning）和知识蒸馏（knowledge distillation，KD）。

知识蒸馏（以下简称为 KD）一般需要两个模型，一大一小，大的充当老师的角色，而小的则作为学生学习老师的行为，目标就是希望学生的小脑瓜子能够容纳住老师的毕生所学，独当一面（虽然能力可能会比老师差点，但够用就行）。对于现今 NLP 领域中的 KD，老师和学生模型大多都是基于 Transformers 的模型，一般将 BERT 作为老师，而另一个小的基于 Transformers 的模型作为学生。

通过了解，我发现 FastBERT 更有趣些，它采用 自蒸馏 （self-distilling）技术，即整个蒸馏过程仅需要单个模型，该模型即是老师又是学生，与以往的需要两个模型的蒸馏方案不同（如DistilBERT、TinyBERT），而且它能根据样本的难易程度 适应性地调整推断时间 （adaptive inference time）。所以本文接下来重点介绍 FastBERT.

首先，通过论文的名字 《FastBERT: a Self-distilling BERT with Adaptive Inference Time》就能看出 FastBERT 的两大特性：self-distilling 和 adaptive inference time.

需要注意的是：该论文只使用文本分类任务做实验，没有尝试 token-level 的分类任务（如命名实体识别）或多标签分类任务，所以将 FastBERT 应用到 token-level 或多标签的分类不一定适用或适用效果不好（纯猜测）。

模型结构

Figure 1. 模型结构与推断过程

 

FastBERT 由两个成分组成：主干（backbone） 和 分支（branches）。如图 1，主干结构与传统的 BERT 一致，包含 12 层的 Transformer block，此外再加上一个分类器，取名为 老师分类器（teacher-classifier）；分支由多个 学生分类器（student-classifiers）组成，其结构与老师分类器一致，这些分类器被加在除了最后一层的每一层 Transformer block 之上，即第 $i$ 层 Transformer block 的输出作为第 $i$ 个 student-classifier 输入，$i=0,\dots ,L-2$。

在 FastBERT 中，主干和分支（student-classifiers）是分开训练的，其 训练过程 如下：

1. 主干（除了 teacher-classifier）部分的预训练
2. 整体主干（包含 teacher-classifier）的微调
3. 分支（student-classifiers）的自蒸馏

前两步的训练过程与传统 BERT 一样，所以你也可以采用 BERT 家族中的其他模型（如 BERT-WWM、RoBERTa 和 ERNIE）作为主干。对于第一步，为了方便，我们可以直接在网上下载已公开的预训练好的 BERT 家族的模型。然后根据任务在自己的标注语料上执行第二步，在这一阶段，所有的 student-classifier 还未被用到。

自蒸馏

在自蒸馏阶段（训练的第三步），模型主干（包括 teacher-classifier）需要被 freeze，即不再进行参数更新。使用 teacher-classifier 的输出作为高质量的 软标签（soft-label）训练 student-classifier. 由于各个 student-classifier 是相互独立的，所以可以分别计算它们的预测 $p_s$ 与 teacher-classifier 输出的软标签 $p_t$之间的差距，论文使用 KL 散度度量这两者之间的距离：
$$D_{KL}(p_s,p_t)=\sum _{i=1}^N{p}_s(i)\cdot \text{log}\frac{p_s(i)}{p_t(j)}$$

其中，$N$ 为分类类别的个数。将所有 student-classifier 与 teacher-classifier 的 KL 散度之和作为自蒸馏的训练损失：
$$Loss(p_{s_0},\dots ,p_{s_{L-2}},p_t)=\sum _{i=0}^{L-2}{D}_{KL}(p_{s_i},p_t)$$
其中，$L$ 为主干 Transformer block 的层数。

由于自蒸馏使用 teacher-classifier 的输出作为标签，因此该阶段也可以使用 未标注的语料 进行训练，这是大多数蒸馏模型（如 TinyBERT）的优点。

 

适应性推断

为除去最后一层的每层 Transformer block 配置一个 student-classifier 的好处就是 FastBERT 可以根据样本的复杂度动态地调整 Transformer block 的计算层数。对于每一层，分别为每个样本判断此时的推断结果是否足够可信以终止往后的运算。给定一个输入，student-classifier 输出的 不确定性（uncertainty）计算方式如下：
$$Uncertianty=\frac{{\sum }_{i=1}^N\text{log}{p}_s(i)}{\text{log}\frac{1}{N}}$$

对于该不确定性，论文作者作了一个假设：不确定性越低，准确率越高（the Lower the Uncertainty, the Higher the Accuracy）.

对于一个样本，如果其在第 $i$ 个 student-classifier 上得出的不确定性高于所设定的 阈值 时，那么该样本会继续传入到往后的 Transformer block 和 student-classifier 中进行运算，反之该样本则不需要涉及进一步的运算，即就拿当前第 $i$ 个 student-classifier 的输出作为该样本最终的输出。

论文中将这里的阈值定义为 Speed，其与 Uncertainty 的取值均在 0 到 1 之间。很明显，Speed 越高，能一直坚持到更高层的样本就越少，则总体的推断速度就会更快，反之亦反。因此，可将 Speed 看作为推断速度与准确率之间的权衡。根据 Speed 的不同取值，理论上 FastBERT 的速度可以提升至 BERT 的 1~12 倍。

 

一些分析

你可能会有疑问，给每层 Transformer block（除了最后一层）都加上一个 student-classifier，难道不会增加模型的计算量，使得模型更加缓慢吗？论文中展示了 Transformer block 和 student-classifier 各自需要的浮点计算数量，如图 2，可看出，相较于 Transformer block，student-classifier 的计算需求小到可以忽略。

Figure 2. 浮点计算量

 

再看一下在 Book review 数据集上，样本被执行的层数的分布图：

Figure 3. 样本执行层数的分布

 

虽然 Speed 取值的高低会造成不同的分布情况，但可以确定的是，大部分样本都不会坚持到最后几层，这也说明了 FastBERT 确实可以提升推断速度。

对于 FastBERT 的性能可看下表：

Figure 4. 性能对比

不仅提升了速度，FastBERT 的性能对比 BERT 几乎没有损失，这是很难得的。

 

思考

其实了解了 FastBERT 的结构后就能明白为什么论文作者只在文本分类这种简单的分类问题上做实验，主要是考虑到 uncertainty 的计算问题。对于文本分类，每个样本有且仅对应一个分类任务的 uncertainty，但对于 token-level 的分类问题（如 命名实体识别），每个样本对应着 $L$ （token 的个数）个分类任务的 uncertainty，这时如何计算样本层面上的 uncertainty 就可以有不同选择了，例如，可以每个 token 的 uncertainty 取平均或取最大。或许 FastBERT 本质上并不适合 token-level 的分类问题，这还有待实验证实，多标签分类也是类似。

参考源

• FastBERT: a Self-distilling BERT with Adaptive Inference Time
• FastBERT：又快又稳的推理提速方法