DL-Paper精读：Tiny-BERT

TinyBERT: Distilling BERT for Natural Language Understanding

https://arxiv.org/pdf/1909.10351.pdf​arxiv.org

Backgroud and Issue

在NLP领域，预训练模型（BERT, XLNet, GPT-3）等已经取得了极大的成功。但是因为预训练模型极大的参数量和推理时间，因此对移动端部署带来了很大的困难。急需一种有效的方式，能够在保持预训练模型精度的同时，尽量地压缩其计算消耗和模型大小。

Method

本文采用知识蒸馏的方法来实现BERT的压缩。首先针对Transformer结构，提出了一种专门的蒸馏方法，并设计了模型专有的损失含糊；然后提出了一种两阶段的整流体系，分别在预训练阶段和下游任务训练阶段进行模型的蒸馏，此外，提出了一种数据增强的方式，来提升下游任务蒸馏的效果。

Transformer Distillation

文中所提出的Transformer Distillation是专门针对于Transformer网络提出的蒸馏方法。其主要解决两个问题：a、Transformer layer的匹配；b、损失函数的设计。

A、 Transformer layer的匹配。TinyBERT希望设计出的网络结构层数少于原BERT，但这为layer级别的蒸馏带来了问题：层数不匹配。例如BERT_base是12层，而所设计的TinyBERT则只有4层，本文提出从BERT的12层中挑选出4层来进行知识蒸馏。针对于这4层的挑选方法，采用了一个映射函数n = g(m)。文中的映射函数采用了uniform strategy，g(m) = 3×m，将BERT_base的12层等分映射到了TinyBERT的4层。

B、 损失函数。文中的蒸馏Loss，主要分为三部分：Embedding-layer Distillation，Prediction-layer Distillation, 以及Transformer每一层的Transformer-layer Distillation。

其中前者是对网络前段Embedding层的输出的学习，采用MSE作为loss函数；第二个是对网络输出层的匹配，采用交叉熵CE作为loss函数。最后是针对网络中间Transformer层的匹配，文中对于每一个Transformer层的输出hidden_state及中间的注意力权值attention进行了学习，采用MSE作为loss函数，如图所示：

2-stage learning framework

BERT的训练一般分为两个阶段：预训练和特定任务的fine-tune，因此文中提出对于TinyBERT的蒸馏，也采用这两个阶段，分别称为general distillation 和 task-specific distillation。首先采用预训练的BERT作为teacher来蒸馏TinyBERT，然后采用fine-tune后的BERT在特定任务数据集上再进行蒸馏该TinyBERT。值得注意的是，文中提出，在general disitillation阶段，prediction disitillation的作用并不明显，可以不用加。简单流程图如下：

此外，在task-specific阶段，为了增强模型能力，文中提出了一种数据增强的方式，结合预训练的BERT和GloVe word embedding实现了word-level的replacement。由于对语言任务的不熟悉，此部分未进行具体的研究。

Experiments

主要结果总结如下：1. TinyBERT4在GLUE数据集上，相比于直接训练的同等大小的模型BERT_tiny，有着明显的精度提升，相比于其他的模型压缩方法结果，也有着精度，加速比，模型参数量等各方面的提升；2. TinyBERT6能够得到与BERT_base相类似的结果，同时得到2倍的加速。

此外，文中还进行了一系列消融实验，充分证明了文中方法的有效性，主要结果如下：

1. 文中所提出的general-disitillation, task-specific disitillation 以及特有的数据增强方法 data augmentation都对结果产生了正面的影响，结果如下：

2. 文中对embedding的结果，transformer层的结果及其attention权值，以及最后prediction的输出都进行了蒸馏，它们也都对最终精度有所贡献，如下：

3. 文中还对于它们所采用的层序号匹配函数n = g(m)进行了分析，文中采用了均分映射的方法，进行了与头部映射及尾部映射的实验对比，结果发现，部分任务对于头部映射响应更好，而另一部分任务与尾部映射更加匹配，这反应出BERT网络的不同层本质上蕴含了不同的特征信息。但综合结果还是均分映射的方法更优秀。

Summary

本文其实并未提出较多original的东西，偏工程性更多。但BERT的压缩工作本质上就是一个工程问题，该方法能够将原本12层的BERT压缩到4层，同时保持可观的精度，确实值得学习。

但文中尚存在一个问题（或者我阅读不够仔细没有发现）。在工程解读过程中发现，TinyBERT同时还试图对于每一层Transformer layer的hidden_size进行压缩，这导致了student和teacher的中间层结果尺度不匹配，则无法直接进行MSE求loss。代码中采用了一个fit_dense将两者的尺度进行对等，然后进行MSE，而且对于所有层都采用的是同一个fit_dense操作。可以理解fit_dense的出现确实是为了尺寸匹配，但它也确实带来了参数，这是独立于网络之外的，对于student网络的影响不可忽略，而且所有层采用同一个操作，影响更加无法估量，这一点并未找到解释。