引言

上一节介绍了$\text{DistilBERT}$模型，本节将继续介绍优化性更强的知识蒸馏$\text{BERT}$模型——$\text{TinyBERT}$模型。

回顾：DistilBERT模型

$\text{DistilBERT}$模型是一种基于$\text{BERT-base}$的知识蒸馏版本，其模型结构表示如下。单从模型结构的角度观察，学生模型神经元的维度没有发生变化$(768)$，仅是$\text{Encoder}$层数减少为$\text{BERT-base}$的一半；并且各层的初始化继承了一部分$\text{BERT-base}$：从教师模型的$\text{Encoder}$层中每两层选择一层作为学生模型$\text{Encoder}$层的初始化。

$\text{DistilBERT}$训练出的学生模型($\text{param:66M}$)依然可以达到$\text{BERT-base}$模型几乎$97$%的准确度。能够达到这个效果离不开$\text{DistilBERT}$的三个核心策略：

• 掩码语言模型策略$(\text{Masked Language Model})$：根据$\text{RoBERTa}$中的描述，摒弃掉下句预测$(\text{Next Sentence Prediction,NSP})$策略，并使用动态掩码替代静态掩码作为$\text{BERT}$模型的训练策略；
• 蒸馏策略$(\text{Distillation loss})$：通过使用$\text{Softmax}$温度函数将教师模型$\text{BERT-base}$与学生模型$\text{DistilBERT}$输出层的解空间尽可能地相似：
  其中$N$表示教师模型和学生模型的输出层维度，在$\text{DistilBERT}$模型中，两者的维度相同，均为$768$。
  $$\begin{array}{rl}& \mathcal{T}(x)=(t_1,t_2,\cdots ,t_N)\\ & \mathcal{S}(x)=(s_1,s_2,\cdots ,s_N)\\ & {\mathcal{L}}_{cross}=-\sum _{i=1}^N{t}_i\ast \log (s_i)\end{array}$$
• 余弦嵌入策略$(\text{Cosine Embedding loss})$：通过计算输出层分布向量之间夹角的余弦值$\cos [\mathcal{T}(x),\mathcal{S}(x)]$，当该值为$1$时，对应的${\mathcal{L}}_{cosine}$达到最小。此时两向量的方向为同一方向，教师和学生模型输出的解空间已被对齐：
  $${\mathcal{L}}_{cosine}=1-\cos [\mathcal{T}(x),\mathcal{S}(x)]$$

总结：
在$\text{ALBERT}$模型中介绍过，虽然$\text{ALBERT}$也是$\text{BERT}$的简化版本，但它们的解空间并不相同；

与此相反，$\text{DistilBERT}$中除了继承了$\text{BERT}$中的掩码语言模型策略外，剩余的两条策略均是围绕牢牢绑定教师模型与学生模型的解空间而设计。

重新观察学生模型$\text{DistilBERT}$，它能够达到如此精炼的模型结构$(\text{param:110M -> 66M})$，但依然保持极高的准确性，没有出现欠拟合的情况。这至少意味着：$\text{DistilBERT}$模型中的神经元被利用得更加充分。

在这种情况下，是否可以百尺竿头更进一步呢$?\text{ TinyBERT}$模型给了我们一个更精进的答案。

TinyBERT模型结构

相比于$\text{DistilBERT}$模型中$\text{Encoder}$层数减半的严肃操作，$\text{TinyBERT}$模型可以自定义学生模型的层数。并且还可以设置隐藏层单元中神经元的维度，从而使模型更加精简。那么它是如何实现在如此精简的模型结构下，不仅没有欠拟合，而且还能保持优秀的训练结果呢$?$ 自然是依靠更加严苛的策略作为约束。

$\text{TinyBERT}$模型的教师-学生模型结构表示如下：

其中索引$0$表示嵌入层，$1$表示第一个$\text{Encoder}$，以此类推。最后$\text{N+1,M+1}$分别表示教师、学生模型的预测层。

该蒸馏结构与$\text{DistilBERT}$之间没有太大区别，只不过没有$\text{DistilBERT}$中的初始化操作。教师与学生模型中各层的迁移过程可以表示为如下式子：
$$n=\mathcal{G}(m)$$
其表达的含义是：将教师模型中的第$n$层迁移到学生模型的第$m$层。例如：

• $0=\mathcal{G}(0)$表示将教师模型的嵌入层知识迁移到学生模型的嵌入层；
• $N+1=\mathcal{G}(M+1)$表示将教师模型的预测层知识迁移到学生模型的预测层；
• $n=\mathcal{G}(m)$表示将教师模型的第$n$个$\text{Encoder}$层知识迁移到学生模型的第$m$个$\text{Encoder}$层。

TinyBERT模型策略

那么$\text{TinyBERT}$是如何制定策略的呢$?$ 主要围绕三个部分制定策略：

• $\text{Transformer}$层$(\text{Encoder})$
• 嵌入层$(\text{Embedding Layer})$
• 预测层$(\text{Predict Layer})$

Transformer层蒸馏

在$\text{Transformer}$层也就是编码器层，需要使用多头注意力机制计算注意力矩阵，再使用$\text{FeedForward Network}$进行一个前馈计算，并将最终计算得到的隐藏状态特征作为该编码器的输出。在$\text{TinyBERT}$中除了将教师模型中$\text{Encoder}$内的注意力矩阵迁移到学生模型相应的$\text{Encoder}$中，也同时将相应的隐藏状态特征迁移到学生模型中。因而$\text{Transformer}$层蒸馏包括两次知识蒸馏：

• 基于注意力的蒸馏
  通过最小化对应学生$\text{Encoder}$和教师$\text{Encoder}$内注意力矩阵的均方误差来训练对应学生$\text{Encoder}$层：
  • 其中$h$表示注意力机制头的数量;${\mathcal{A}}_i^{\mathcal{S}}$表示学生$\text{Encoder}$内第$i$个头的注意力矩阵;${\mathcal{A}}_i^{\mathcal{T}}$表示教师$\text{Encoder}$内第$i$个头的注意力矩阵;$\text{MSE}$表示均方误差操作。
  • 个人疑问：当学生模型隐藏层维度变化的时候${\mathcal{A}}_i^{\mathcal{S}},{\mathcal{A}}_i^{\mathcal{T}}$是一样大的吗？但书中并没有解释。
    $${\mathcal{L}}_{attn}=\frac{1}{h}\sum _{i=1}^h\text{MSE}({\mathcal{A}}_i^{\mathcal{S}},{\mathcal{A}}_i^{\mathcal{T}})$$
    需要注意的是，这里的注意力矩阵${\mathcal{A}}_i^{\mathcal{S}},{\mathcal{A}}_i^{\mathcal{T}}$使用的是执行$\text{Layer Norm}$映射前的矩阵，这样做的目的是保证信息的完整性，并且更快地收敛。
• 基于隐藏状态的蒸馏
  隐藏状态是当前$\text{Encoder}$的输出，我们同样需要将教师$\text{Encoder}$的隐藏层知识迁移到学生$\text{Encoder}$的隐藏层状态中：
  其中${\mathcal{H}}_{\mathcal{S}}$表示学生$\text{Encoder}$内的隐藏层状态;${\mathcal{H}}_{\mathcal{T}}$表示教师$\text{Encoder}$内的隐藏层状态。同样使用均方误差使${\mathcal{H}}_{\mathcal{S}}$向${\mathcal{H}}_{\mathcal{T}}$方向拟合。
  $${\mathcal{L}}_{hidn}=\text{MSE}({\mathcal{H}}_{\mathcal{S}},{\mathcal{H}}_{\mathcal{T}})$$
  但需要注意的是：当学生$\text{Encoder}$隐藏层维度发生变化时，${\mathcal{H}}_{\mathcal{S}}$和${\mathcal{H}}_{\mathcal{T}}$两者之间的维度之间存在差异，因而需要训练一个新的权重矩阵${\mathcal{W}}_h$使两者处于同一级别的维度空间：
  相当于作用在损失函数上的权重矩阵，反向传播过程中同样存在梯度更新。
  $${\mathcal{L}}_{hidn}=\text{MSE}({\mathcal{H}}_{\mathcal{S}}{\mathcal{W}}_h,{\mathcal{H}}_{\mathcal{T}})$$

嵌入层蒸馏

关于嵌入层的蒸馏与隐藏状态的蒸馏相似，当学生模型设置的隐藏层维度与教师模型维度不同时，两者对应的$\text{Embedding}$也不同。同样在损失函数中添加一个新的权重参数${\mathcal{W}}_{\mathcal{E}}$，使两个$\text{Embedding}$处于同一级别的维度空间：
其中${\mathcal{E}}_{\mathcal{S}}$表示学生模型的$\text{Embedding}$矩阵;${\mathcal{E}}_{\mathcal{T}}$表示教师模型的$\text{Embedding}$矩阵。$\text{MSE}$表示均方误差。

$${\mathcal{L}}_{emb}=\text{MSE}({\mathcal{E}}_{\mathcal{S}}{\mathcal{W}}_{\mathcal{E}},{\mathcal{E}}_{\mathcal{T}})$$

预测层蒸馏

在预测层蒸馏部分，迁移的是输出层的知识信息。这里和$\text{DistilBERT}$模型关于预测层的损失类似。对于教师模型的输出${\mathcal{Z}}^{\mathcal{T}}$和学生模型的输出${\mathcal{Z}}^{\mathcal{S}}$：

• 使用$\text{Softmax}$温度函数分别获取对应的软目标${\mathcal{P}}^{\mathcal{T}}$和软预测${\mathcal{P}}^{\mathcal{S}}$结果：
  同理,${\mathcal{Z}}^{\mathcal{T}}$对应的软目标结果${\mathcal{P}}^{\mathcal{T}}$不再赘述。
  $$\{\begin{array}{l}\begin{array}{r}{\mathcal{P}}_i^{\mathcal{S}}=\frac{\exp ({\mathcal{Z}}_i^{\mathcal{S}}/\mathcal{T})}{{\sum }_j\exp ({\mathcal{Z}}_j^{\mathcal{S}})/\mathcal{T}}\end{array}\\ \\ {\mathcal{P}}^{\mathcal{S}}=({\mathcal{P}}_1^{\mathcal{S}},{\mathcal{P}}_2^{\mathcal{S}},\cdots ,{\mathcal{P}}_N^{\mathcal{S}})\end{array}$$
• 再使用交叉熵损失函数对${\mathcal{P}}^{\mathcal{S}}$与${\mathcal{P}}^{\mathcal{T}}$进行描述：
  $${\mathcal{L}}_{pred}=-{\mathcal{P}}^{\mathcal{T}}\cdot \log \left({\mathcal{P}}^{\mathcal{S}}\right)$$

最终，$\text{TinyBERT}$包含所有层的损失函数表示如下：
这里$[{\mathcal{S}}_m,{\mathcal{T}}_{\mathcal{G}(m)}]$表示学生模型的第$m$层与教师模型第$\mathcal{G}(m)$之间的迁移关系。
$$\mathcal{L}\left[{\mathcal{S}}_m,{\mathcal{T}}_{\mathcal{G}(m)}\right]=\{\begin{array}{l}{\mathcal{L}}_{emb}({\mathcal{S}}_0,{\mathcal{T}}_0)m=0\\ {\mathcal{L}}_{hidn}({\mathcal{S}}_m,{\mathcal{T}}_{\mathcal{G}(m)})M\ge m>0\\ {\mathcal{L}}_{pred}({\mathcal{S}}_{M+1},{\mathcal{T}}_{N+1})m=M+1\end{array}$$
可以看出：

• $\text{TinyBERT}$损失函数数量是不确定的。它取决于设计学生模型$(\text{TinyBERT})$的层的数量；
• 相比于$\text{DistilBERT}$，$\text{TinyBERT}$需要为削减隐藏层状态维度和层数付出相应的代价——设计的策略需要与教师模型关系更加紧密，并精确到注意力矩阵和隐藏层状态，从而得到一个与教师模型关联更加紧密的、学生模型的解空间。

TinyBERT模型的训练

在文章中作者描述的训练流程表示如下：

在$\text{TinyBERT}$模型中，使用两个阶段进行训练：

• 通用蒸馏：在该阶段，使用$\text{BERT-base}$预训练模型作为教师，并使用$\text{BERT-base}$的训练集对学生模型$(\text{TinyBERT})$进行蒸馏。并将该模型称作通用$\text{TinyBERT}$模型。
• 特定任务蒸馏：在微调阶段，将基于一项具体任务对通用$\text{TinyBERT}$模型进行微调$(\text{fine-tuning})$。具体微调过程方式为：
  • 使用预训练$\text{BERT-base}$模型针对具体任务进行微调，并将这个微调后的$\text{BERT-base}$模型作为教师；
  • 将上述经过通用蒸馏得到的通用$\text{TinyBERT}$模型作为学生，经过蒸馏，得到的$\text{TinyBERT}$模型称作微调的$\text{TinyBERT}$模型。

效果展示

论文中关于$\text{TinyBERT}$对于各下游任务中，与各模型比较结果如下：

其中，${\text{DistilBERT}}_4$表示学生模型包含$4$层$\text{Encoder}$，其他同理。可以发现：

• 相比于${\text{DistilBERT}}_4$，${\text{TinyBERT}}_4$使用不到其$30$%，但准确率却远高于$\text{DistilBERT}$模型；
• ${\text{TinyBERT}}_6$参数数量是$\text{BERT-base}$的$60$%左右，但其准确性基本与$\text{BERT-base}$持平。

$\text{Reference}$：
论文链接
《BERT基础教程——Transformer大模型实战》