1. 文章简介

• 标题：TinyBERT: Distilling BERT for Natural Language Understanding
• 作者：Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, Qun Liu
• 日期：2019
• 期刊：arxiv preprint

2. 文章概括

  文章提出了一种两阶段的BERT蒸馏模型TinyBERT。TinyBERT在GLUE上击败了所有当前的SOTA蒸馏BERT模型[1]，且参数量仅为SOTA的38%，推理时间仅为SOTA的31%。此外TinyBERT在所有GLUE任务中平均表现约为96.8%，几乎完美还原BERT的能力。
  TinyBERT的整体学习步骤如下

3 文章重点技术

3.1 Transformer Distillation

  所谓Transformer Distillation(TD)，即对Transformer架构的蒸馏。假设教师模型和学生模型的层数分别为$N$和$M$，则首先定义一个映射函数$n=g(m)$表示用学生模型的第$m$层去学习教师模型的第$n=g(m)$层的信息。文章通过数值实验选用了$g(m)=3m$。定义第$0$层为嵌入层，第$M+1$层为预测层，则我们可以将模型的损失函数写作$${\mathcal{L}}_{model}=\sum _{x\in \mathcal{X}}\sum _{m=0}^{M+1}{\lambda }_m{\mathcal{L}}_{layer}(f_m^S(x),f_{g(m)}^T(x))\text{\hspace{1em}(1)}$$，其中${\mathcal{L}}_{layer}$表示$layer$层的损失函数，$f_m^S(x),f_{g(m)}^T(x)$分别表示学生和教师模型在第$m$或$g(m)$层的函数，${\lambda }_m$为超参数，表示第$m$层的重要性。下面为针对不同层的蒸馏方式

• Transformer-layer Distillation：
  
  如上图所示，Transformer-layer Distillation包含以下两种蒸馏方法
  • Attention based distillation：蒸馏注意力机制矩阵，损失函数为$${\mathcal{L}}_{attn}=\frac{1}{h}\sum _{i=1}^{h}MSE(A_i^S,A_i^T)\text{\hspace{1em}(2)}$$，其中$h$为多头注意力机制的head数目，$MSE$表示Mean Squared Error，$A_i^S,A_i^T$分别表示学生模型和教师模型的注意力矩阵。
  • hidden tsates based distillation：蒸馏隐藏层（即FFN的输出层）状态，蒸馏的损失函数为$${\mathcal{L}}_{hidn}=MSE(H^S{W}_h,H^T)\text{\hspace{1em}(3)}$$，其中$H^S,H^T$分别表示学生模型和教师模型的隐藏层状态，$W_h$为可学习的参数，旨在将学生模型的隐藏向量映射到和教师模型隐藏状态相同的高维空间
• Embedding-layer Distillation：对嵌入层进行蒸馏，损失函数为$${\mathcal{L}}_{embd}=MSE(E^S{W}_e,E^T)\text{\hspace{1em}(4)}$$，其中$E^S,E^T$分别表示学生模型和教师模型的嵌入层向量，$W_e$和上述$W_h$作用相同，旨在将学生模型的嵌入向量映射到和教师模型嵌入向量相同的高维空间
• Prediction-layer Distillation：采用损失函数$${\mathcal{L}}_{pred}=CE(z^T/t,z^S/t)\text{\hspace{1em}(5)}$$，其中$z^S,z^T$分别表示学生模型和教师模型的输出logits，$t$表示蒸馏的温度。此设置参考原始蒸馏论文中的设置。
    最后，将上述所有损失函数进行统一，得到$(1)$式中的损失函数可表示为$${\mathcal{L}}_{layer}=\{\begin{array}{}\end{array}$$