Press "Enter" to skip to content

DistilBERT — 更小更快的BERT模型

使用蒸馏(Distillation)技术来帮BERT瘦身

 

介绍

 

今天要介绍的是 Distilled BERT,是由hugging face在今年(2019)所提出的一篇论文,相信用pytorch来做BERT的朋友一定对 hugging face 不陌生,他们在 Github上开源的 transformers
专案

 

回到今天的主题,Distill的意思是蒸馏,我们可以从字面上猜测,我们要从一个很大的模型,蒸馏成比较小的模型,也可以用一种角度想,我们让大的模型当作小的模型的老师,而小模型这个学生,只会尽可能的学老师的每个动作

 

大致上Distilled BERT的思想就是这样简单,根据作者的实验数据,DistilBERT的参数大约只有BERT的40%,而速度快了60%,并保有一定一精準度

 

训练

 

很多人刚看到这个的想法,可能会直觉想说直接拿现成的BERT来训练小的DistilBERT模型应该不太耗费效能,不过由于Distillation Loss需拿BERT跟DistilBERT的结果相互比较,因此仍需要对BERT做forward pass,再计算出Triple Loss后进行Backward Propagation。

 

方法

 

所以我们实际上到底怎幺训练的,作者提出的DistilBERT的架构,层数为BERT的一半,这边的重点其实是Loss Function,训练採用的Loss Function为三种的结合(Triple Loss)

Distillation Loss
Softmax-temperature

 

    1. Distillation Loss

 

    1. : 从图中可以看出,当BERT预测(ti)越高,而DistilBERT预测(si)越低,产生的Loss就会越高,而机率是通过 softmax-temperature (Hinton et al., 2015) 计算

 

    1. Masked Language Modeling Loss

 

    1. : 参考BERT

 

    1. Cosine Embedding Loss

 

    1. : 用于让DistilBERT能趋向产生跟BERT更像的hidden vector

 

 

Student Architecture

 

 

    1. Linear layer and Layer Normalization

 

    1. : 已被高度优化且证明有效(Highly optimized),因此不更动

 

    1. The size of hidden vector of last layer

 

    1. : 根据实验发现减少 hidden vector size 并不太影响效能

 

    1. Number of layers

 

    1. : 影响效能、推理速度较多,因此作者注重在此参数

 

 

Student Initialization

 

因爲层数为BERT的一半,直接使用BERT的参数当作初始值

 

Compute Power

 

8 16GB V100 GPUs for approximately 90hrs on Toronto Book Corpus(原先BERT使用的资料集)

 

Evaluation

 

GLUE: DistilBERT retrain 97% of BERT performance,有些甚至超越BERT Base

 

IMDb accuracy: BERT(93.46) DistilBERT(92.82)

 

Speed(A full pass of GLUE task): BERT(668s) ELMo(895s) DistilBERT(410s)

 

On device computation

 

将容量缩小后,也可以直接将模型放在设备端(i.e. APP),论文中的DistilBERT大约 207MB

 

 

    1. Building a mobile application for QA

 

    1. 207 MB, could be further reduced with quantization

 

 

运行在下游任务(Downstream tasks)

 

We also studied whether we could add another step of distillation during the adaptation phase by fine-tuning DistilBERT on SQuAD using a BERT model previously fine-tuned on SQuAD as a teacher for an additional term in the loss (knowledge distillation). In this setting, there are thus two successive steps of distillation, one during the pre-training phase and one during the adaptation phase.

 

从上面的段落看来,在问答任务(属于一种下游任务)中,将DistilBERT再跟Fine-tuned BERT进行一次蒸馏,就可以达到不错效果

 

结论

 

DistilBERT整体来说很适合使用在应用层面,我们只需完成蒸馏的动作,就可以大幅增进推理速度,以及减少模型大小,可以减低伺服器的成本

 

Reference

 

Distilling the Knowledge in a Neural Network

DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注