《TinyBERT: Distilling BERT for Natural Language Understanding》（ICLR-2020在审）论文阅读

前言

• 论文地址：https://arxiv.org/abs/1909.10351
• 代码地址：https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/TinyBERT

Abstract

• 目的：加快速度，减少内存，同时尽可能保持效果。

1、Introduction

• 提出了一种新的 Transformer 蒸馏方法。
• 提出了一个两段式的学习框架，该框架可以在 pre-training 和 fine-tuning 阶段都能使用上述蒸馏的方式，从而学习到更丰富的知识。

2、Preliminaries

这部分就介绍了 transformer 和知识蒸馏，都是些先验知识。有兴趣的也可以分别去看我的 transformer 和 fastbert 的博客。

3、Method

（1） 模型结构

先看总体的模型图：

从图中可以看出，从上图可知，就是一个普通的蒸馏模型，Student 模型也是 transformer 结构，这没啥好说的。

模型定了，就看怎么来蒸馏了，也就是损失函数怎么定的问题，具体来说，一共包括以下四个部分：

• 1、Embeding 层。
• 2、预测层。
• 3、transformer 中的multi-head attention。
• 4、transformer 层的输出。
  这四部分来挨个说。

Ⅰ、embedding 层

embedding 层的蒸馏是下面函数：

E 代表 Embedding 矩阵，这里加一个W个人觉得应该是变换维度，比如 $E^s$ 的词向量维度小一些，否则，直接用老师的词向量多好。MSE 是平方损失函数。

Ⅱ、预测层

这里是用两个模型预测概率的交叉熵，这里不放公式了，就是常规的模型蒸馏方式。

III、transformer 层

这里首先要确定一些对应的关系，比如说 bert 有12层 transformer，但是我的 Student 模型只有三层，这样就需要确定怎样对每一层对应起来了，作者给出的方法是搞一个对应的映射，比如说这里原始 bert 的第一层对应 Student 模型的第一层，然后 bert 隔四层对应 Student 的第二层，以此类推。然后再来看细节：

其中 MHA 是multi-head attention，从图中可以看出来，是对权重矩阵和隐输出层计算损失：

其实看图已经很明了，这里的 W 跟Embedding的中W用法是一个意思。

总的来说，模型的损失函数，就是上面讲的四个损失的和。

（2）两段式学习框架

感觉这里主要是借鉴bert的预训练和微调思路，大概就是长下面这样：

具体步骤如下:

• 1、使用 Bert 作为 teacher 模型，而且不对其进行微调，利用大规模本语料库作为学习数据，这样做的好处是相当于获得一个通用的 TinyBert。
• 2、然后使用微调后的 Bert 作为 teacher 模型，针对特定的任务，再进行一次蒸馏，得到最终的 TinyBert。也就是相当于把第一步当作一个 TinyBert 的初始化操作。

Experiments

过。