BERT模型蒸馏有哪些方法？

会员服务 ·

BERT模型蒸馏有哪些方法？

2020 年 10 月 12 日 PaperWeekly

©PaperWeekly 原创 · 作者｜蔡杰

学校｜北京大学硕士生

研究方向｜问答系统

我们都知道预训练模型的标准范式：

pretrain-利用大量的未标记数据通过一些自监督的学习方式学习丰富的语义和句法知识。例如：Bert 的 MLM，NSP 等等。
finetune-将预训练过程中所学到的知识应用到子任务中，以达到优异的效果。

预训练模型在各个领域虽然带来了巨大的提升，但是也有一些致命的问题：

预训练模型高计算复杂度-不可能在实时系统中运行。
大存储需求——预训练模型一般都很大，少则几百 M，大则几 G，无法在有限资源的设备上部署。

所以模型压缩和加速技术的研究迫在眉睫！

Logit Distillation

知识蒸馏最早是 Hinton 在 15 年提出的一个黑科技技术，核心思想是通过迁移知识，从而通过训练好的大模型得到更加适合推理的小模型。

论文标题：Distilling the Knowledge in a Neural Network

论文来源：NIPS 2014

论文链接：http://arxiv.org/abs/1503.02531

神经网络通常使用 “softmax” 输出层生成每个类别的概率：

输出层通过将与其他 logit 进行比较，以表示为每个类计算出的 logit。T 是温度系数，通常设置为 1。

softmax 是一种 “soft max”，在使用 cross-entropy 计算 loss 的时候，相对的差别规模较小的 logit 值被压缩甚至丢失。

论文中用了对 softmax 的公式进行了推导，如果 T 远高于对数的量级，且对数为零均值，则上图公式（4）和公式（2）和几乎相同。在 T 非常小的情况下，梯度会接近于 qi-pi，所以当多个类别的 pi 接近于 0 时，最终输出的相似性信息是没有体现在梯度中的，所以网络会忽略掉这些信息；

所以，单纯的用logits来蒸馏的方法是存在一些 bias 的，那能否有更准确以及更快速的蒸馏方法呢？

Beyond Logit Distillation

论文标题：TinyBERT: Distilling BERT for Natural Language Understanding

论文来源：ICLR 2020

论文链接：http://arxiv.org/abs/1909.10351

代码链接：https://github.com/huawei-noah/Pretrained-Language-Model

TINYBERT 提出了新的迁移蒸馏方法，蒸馏的 loss 可以从 4 个方面来构造：

Embedding Layer
Transformer Layer（Hidden States 和 Attention Matricies）
Prediction Layer

整个模型的 loss 是学生模型的每一层 loss 的求和：

构造的 Loss 都是清一色的 MSE 函数：

Embedding-layer Distillation

Hidden-layer Distillation

Self-Attention Distillation

Target Classifier Distillation

最终每一层的 loss 如下:

通过论文中的实证研究表明了 TinyBERT 的有效性，在 GLUE 上达到了与 BERT 相当（下降 3 个百分点）的效果，同时模型大小只有 BERT 的 13.3%（BERT 是 TinyBERT 的 7.5 倍），Inference 的速度是 BERT 的 9.4 倍。

此外，TinyBERT 还显著优于当前的 SOTA 基准方法（BERT-PKD），但参数仅为为后者的 28％，推理时间仅为后者的 31％。

Curriculum Distillation

论文标题：FitNets : Hints for Thin Deep Nets

论文来源：ICLR 2015

论文链接：https://arxiv.org/abs/1412.6550

代码链接：https://github.com/adri-romsor/FitNets

蒸馏时采用的中间层匹配本质上是一种正则化形式，Transformer 的分层蒸馏可能会导致过度正则化。

本文将 Hinton 的蒸馏 output 的 logits 的 idea 进行了扩展，不仅使用输出的 logits，同时也用模型的中间层信息来作为 hints 来提高 student 模型的性能。通过这样的方法可以训练出一个 deeper 和 thinner 的 student 模型。

代表 teacher 网络中的第一层到 hint 层的函数映射，Whint 表示其中的参数；

代表 student 网络中的第一层到 hidden 层的函数映射，Wguided 表示其中的参数；

r 代表 hidden 层的一个额外的映射关系，Wr 是其中的参数，这是为了使得 hidden 层与 hint 层的神经元大小一致。

Fig. 1 (a) 作者选择 FitNet 的一个隐藏层，作为学习层，去学习 teacher model 的某一层（引导层）。我们希望学习层能够预测引导层的输出。

Fig. 1 (b) 作者从一个训练过的 teacher 网络和一个随机初始化的 FitNet 开始，在 FitNet 学习层的顶部加入一个由 Wr 参数化的回归因子，将 FitNet 参数 WGuided 训练到学习层，使 Eq.(3) 最小化 (Fig. 1 (b))。

最后，从预训练的参数中，我们对整个 FitNet 的 Ws 的参数进行训练，使 Eq.(2) 最小化 (Fig. 1 (c))。

从实验结果上看，student 比 teacher 参数少，效果却比 teacher 还要好，可能是因为网络更深的原因，某种程度上说明了深层网络的有效性，深层网络的表示性能要更优一些。

实验也说明了 Hint 损失的重要性，Hint 损失在某种程度上弥补了交叉熵和 KD loss 损失的信息，而这种丢失的信息在某种程度上可以被更深层的网络来弥补。

Dynamic Early Exit

通过分辨简单样本和复杂样本，有针对性的去优化网络性能，加快 inference 的速度，具体可参考之前写的 FastBert。

论文标题：FastBERT: a Self-distilling BERT with Adaptive Inference Time

论文来源：ACL 2020

论文链接：https://arxiv.org/abs/2004.02178

代码链接：https://github.com/autoliuweijie/FastBERT

为了在保证模型性能的前提下提高其效率，本文提出了一种新的具有自适应推断时间的速度可调快速 bert。推理速度可根据不同需求灵活调整，同时避免了样本的冗余计算。此外，该模型采用了一种独特的自蒸馏机制进行微调，有更强的计算效能与最小的性能损失。

原 BERT 模型为主干（Backbone），每个分类器称为分支（Branch），其中分支 Classifier 都是由最后一层的分类器蒸馏而来，在预训练和微调阶段都只调整主干参数，finetune 之后主干参数 freeze，把最后一层classifier蒸馏到每一层 student classifier 中。

之所以叫自蒸馏，因为 student 和 teacher 都是由一个模型得到的，以往的 KD 都是两个模型，student 模型经过参数初始化，让 teacher 模型去优化 student 模型。

一共包含以下几个阶段：

Pre-training：和 BERT 预训练的流程一致，主要是要得到一个预训练模型。
Fine-tuning for Backbone：Backbone 的微调，也就是训练一个 Bert 的分类器，预训练模型后添加一层 classifier，classifier 的参数用于后期的指导训练。
Self-distillation for branch：分支（branch）的自蒸馏，每一层都有一个 student 的 classfier，由微调过的 Bert 的最后一层 classifier 蒸馏而来，每个 branch 都有一个 classifier。
Adaptive inference：自适应 inference，可以根据样本的难易程度决定样本要走几层分支分类器，简单的样本可以在底层直接给结果，困难的继续往高层走。