关于人工智能:深度学习实践篇17模型压缩技术模型蒸馏算法DynaBERTTinyBERT

深度学习实际篇[17]：模型压缩技术、模型蒸馏算法：Patient-KD、DistilBERT、DynaBERT、TinyBERT

1. 模型压缩概述

1.2 模型压缩原有

实践上来说，深度神经网络模型越深，非线性水平也就越大，相应的对事实问题的表达能力越强，但相应的代价是，训练老本和模型大小的减少。同时，在部署时，大模型预测速度较低且须要更好的硬件反对。但随着深度学习越来越多的参加到产业中，很多状况下，须要将模型在手机端、IoT 端部署，这种部署环境受到能耗和设施体积的限度，端侧硬件的计算能力和存储能力绝对较弱，突出的诉求次要体现在以下三点：

• 首先是速度，比方像人脸闸机、人脸解锁手机等利用，对响应速度比拟敏感，须要做到实时响应。
• 其次是存储，比方电网周边环境监测这个利用场景中，要图像指标检测模型部署在可用内存只有 200M 的监控设施上，且当监控程序运行后，残余内存会小于 30M。
• 最初是耗能，离线翻译这种挪动设施内置 AI 模型的能耗间接决定了它的续航能力。

以上三点诉求都须要咱们依据终端环境对现有模型进行小型化解决，在不损失精度的状况下，让模型的体积更小、速度更快，能耗更低。

但如何能产出小模型呢？常见的形式包含设计更高效的网络结构、将模型的参数质变少、将模型的计算量减少，同时进步模型的精度。可能有人会提出疑难，为什么不间接设计一个小模型？要晓得，理论业务子垂类泛滥，工作复杂度不同，在这种状况下，人工设计无效小模型难度十分大，须要十分强的畛域常识。而模型压缩能够在经典小模型的根底上，稍作解决就能够疾速拔高模型的各项性能，达到“多快好省”的目标。

上图是分类模型应用了蒸馏和量化的效果图，横轴是推理耗时，纵轴是模型准确率。图中最上边红色的星星对应的是在 MobileNetV3_large model 根底上，应用蒸馏后的成果，相比它正下方的蓝色星星，精度有显著的晋升。图中所标浅蓝色的星星，对应的是在 MobileNetV3_large model 根底上，应用了蒸馏和量化的后果，相比原始模型，精度和推理速度都有显著的晋升。能够看出，在人工设计的经典小模型根底上，通过蒸馏和量化能够进一步晋升模型的精度和推理速度。

1.2. 模型压缩的根本办法

模型压缩能够通过以下几种办法实现：

• 剪裁：相似“化学结构式的减肥”，将模型构造中对预测后果不重要的网络结构剪裁掉，使网络结构变得更加”瘦身“。比方，在每层网络，有些神经元节点的权重十分小，对模型加载信息的影响微不足道。如果将这些权重较小的神经元删除，则既能保障模型精度不受大影响，又能减小模型大小。
• 量化：相似“量子级别的减肥”，神经网络模型的参数个别都用 float32 的数据表示，但如果咱们将 float32 的数据计算精度变成 int8 的计算精度，则能够就义一点模型精度来换取更快的计算速度。
• 蒸馏：相似“老师教学生”，应用一个成果好的大模型领导一个小模型训练，因为大模型能够提供更多的软分类信息量，所以会训练出一个成果靠近大模型的小模型。
• 神经网络架构搜寻（NAS）：相似“化学结构式的重构”，以模型大小和推理速度为束缚进行模型构造搜寻，从而取得更高效的网络结构。

除此以外，还有权重共享、低秩合成等技术也可实现模型压缩。

2.Patient-KD 模型蒸馏

2.1. Patient-KD 简介

论文地址：Patient Knowledge Distillation for BERT Model Compression

<center> 图 1: Vanilla KD 和 PKD 比拟 </center>

BERT 预训练模型对资源的高需要导致其很难被利用在理论问题中，为缓解这个问题，论文中提出了 Patient Knowledge Distillation（Patient KD）办法，将原始大模型压缩为等同无效的轻量级浅层网络。同时，作者对以往的常识蒸馏办法进行了调研，如图 1 所示，vanilla KD 在 QNLI 和 MNLI 的训练集上能够很快的达到和 teacher model 相媲美的性能，但在测试集上则很快达到饱和。对此，作者提出一种假如，在常识蒸馏的过程中过拟合会导致泛化能力不良。为缓解这个问题，论文中提出一种“急躁”师生机制，即让 Patient-KD 中的学生模型从老师网络的多个中间层进行常识提取，而不是只从老师网络的最初一层输入中学习。

2.2. 模型实现

Patient-KD 中提出如下两个常识蒸馏策略：

1. PKD-Skip: 从每 k 层学习，这种策略是假如网络的底层蕴含重要信息，须要被学习到（如图 2a 所示）
2. PKD-last: 从最初 k 层学习，假如老师网络越靠后的层蕴含越丰盛的常识信息（如图 2b 所示）

<center> 图 2a: PKD-Skip 学生网络学习老师网络每两层的输入 图 2b: PKD-Last 学生网络从老师网络的最初六层学习 </center>

因为在 BERT 中仅应用最初一层的 [CLS] token 的输入来进行预测，且在其余 BERT 的变体模型中，如 SDNet，是通过对每一层的[CLS] embedding 的加权平均值进行解决并预测。由此能够推断，如果学生模型能够从任何老师网络中间层中的[CLS] 示意中学习，那么它就有可能取得相似老师网络的泛化能力。

因而，Patient-KD 中提出非凡的一种损失函数的计算形式：

$$
L_{PT} = \sum_{i=1}^{N}\sum_{j=1}^{M} \left \| \frac{h_{i,j}^s}{\left \| h_{i,j}^s \right \|_{2}} – \frac{h_{i, I_{pt}(j)}^t}{\left \| h_{i, I_{pt}(j)}^t \right \|_2}\right \|_2^2
$$

其中，对于输出 $x_i$，所有层的 [CLS] 的输入示意为：

$$
h_i = [h_{i,1}, h_{i,2},…, h_{i,k}] = BERT_{k}(x_i) \in \mathbb{R}^{k\times d}​
$$

$I_{pt}$ 示意要从中提取常识的一组中间层，以从 $BERT_{12}$ 压缩到 $BERT_6$ 为例，对于 PKD-Skip 策略，$I_{pt} = {2,4,6,8,10}$；对于 PKD-Last 策略，$I_{pt} = {7,8,9,10,11}$。M 示意学生网络的层数，N 是训练样本的数量，上标 s 和 t 别离代表学生网络和老师网络。

同时，Patient-KD 中也应用了 $L_{DS}$ 和 $L_{CE}^S$ 两种损失函数用来掂量老师和学生网络的预测值的间隔和学生网络在特定上游工作上的穿插熵损失。

$$
L_{DS}=-\sum_{i \in [N]} \sum_{c \in C}[P^t(y_i = c|x_i;\hat{\theta}^t)\cdot log P^s(y_i = c |x_i; \theta^s)]
$$

$$
L_{CE}^s=-\sum_{i \in [N]} \sum_{c \in C}\mathbb{1}[y_i=c]\cdot log P^s(y_i = c|x_i;\theta^s)]
$$

最终的指标损失函数能够示意为：

$$
L_{PKD} = (1-\alpha)L_{CE}^S+\alpha L_{DS} + \beta L_{PT}
$$

2.3. 试验后果

<center> 图 3: results from the GLUE test server</center>

作者将模型预测提交到 GLUE 并取得了在测试集上的后果，如图 3 所示。与 fine-tuning 和 vanilla KD 这两种办法相比，应用 PKD 训练的 $BERT_3$ 和 $BERT_6$ 在除 MRPC 外的简直所有工作上都体现良好。其中，PKD 代表 Patient-KD-Skip 办法。对于 MNLI- m 和 MNLI-mm，六层模型比微调（FT）基线进步了 1.1% 和 1.3%，

咱们将模型预测提交给官网 GLUE 评估服务器以取得测试数据的后果。后果总结在表 1 中。与间接微调和一般 KD 相比，咱们应用 BERT3 和 BERT6 学生的 Patient-KD 模型在除 MRPC 之外的简直所有工作上都体现最好。此外，6 层的 $BERT_{6}-PKD$ 在 7 个工作中有 5 个都达到了和 BERT-Base 类似的性能，其中，SST-2（与 BERT-Base 老师相比为 -2.3%）、QQP（-0.1%）、MNLI-m（-2.2%）、MNLI-mm（-1.8%）和 QNLI (-1.4%))，这五个工作都有超过 6 万个训练样本，这表明了 PKD 在大数据集上的体现往往更好。

<center> 图 4: PKD-Last 和 PKD-Skip 在 GLUE 基准上的比照 </center>

只管这两种策略都比 vanilla KD 有所改进，但 PKD-Skip 的体现略好于 PKD-Last。作者揣测，这可能是因为每 k 层的信息提炼捕捉了从低级到高级的语义，具备更丰盛的内容和更多不同的示意，而只关注最初 k 层往往会捕捉绝对同质的语义信息。

<center> 图 5: 参数量和推理工夫比照 </center>

图 5 展现了 $BERT_3$、$BERT_6$、$BERT_12$ 的推理工夫即参数量, 试验表明 Patient-KD 办法实现了简直线性的减速，$BERT_6$ 和 $BERT_3$ 别离提速 1.94 倍和 3.73 倍。

3.DistilBERT 蒸馏

3.1. DistilBERT 简介

论文地址：DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter

<center> 图 1: 几个预训练模型的参数量统计 </center>

近年来，大规模预训练语言模型成为 NLP 工作的根本工具，尽管这些模型带来了显著的改良，但它们通常领有数亿个参数（如图 1 所示），而这会引起两个问题。首先，大型预训练模型须要的计算成本很高。其次，预训练模型一直增长的计算和内存需要可能会妨碍语言解决利用的宽泛落地。因而，作者提出 DistilBERT，它表明小模型能够通过常识蒸馏从大模型中学习，并能够在许多上游工作中达到与大模型类似的性能，从而使其在推理时更轻、更快。

3.2. 模型实现

学生网络结构

学生网络 DistilBERT 具备与 BERT 雷同的通用构造，但 token-type embedding 和 pooler 层被移除，层数减半。学生网络通过从老师网络中每两层抽取一层来进行初始化。

Training loss

$L_{ce}$ 训练学生模拟老师模型的输入散布：

$$
L_{ce} = \sum_i t_i * log(s_i)
$$

其中，$t_i$ 和 $s_i$ 别离是老师网络和学生网络的预测概率。

同时应用了 Hinton 在 2015 年提出的softmax-temperature：

$$
p_i = \frac{exp(z_i/T)}{\sum_j exp(z_j/T)}
$$

其中，$T$ 管制输入散布的平滑度，当 T 变大时，类别之间的差距变小；当 T 变小时，类别间的差距变大。$z_i$ 代表分类 $i$ 的模型分数。在训练时对学生网络和老师网络应用同样的 temperature $T$，在推理时，设置 $T=1$，复原为规范的 softmax
最终的 loss 函数为 $L_{ce}$、Mask language model loss $L_{mlm}$（参考 BERT）和 cosine embedding loss $L_{cos}$（student 和 teacher 暗藏状态向量的 cos 计算）的线性组合。

3.3. 试验后果

图 2：在 GLUE 数据集上的测试后果、上游工作测试和参数量比照

依据上图咱们能够看到，DistilBERT 与 BERT 相比缩小了 40% 的参数，同时保留了 BERT 97% 的性能，但进步了 60% 的速度。

4.DynaBERT 蒸馏

4.1. DynaBERT 简介

论文地址：DynaBERT: Dynamic BERT with Adaptive Width and Depth

预训练模型，如 BERT，在自然语言解决工作中的弱小之处是毫无疑问，然而因为模型参数量较多、模型过大等问题，在部署方面对设施的运算速度和内存大小都有着极高的要求。因而，面对理论产业利用时，比方将模型部署到手机上时，就须要对模型进行瘦身压缩。近年的模型压缩形式基本上都是将大型的 BERT 网络压缩到一个固定的小尺寸网络。而理论工作中，不同的工作对推理速度和精度的要求不同，有的工作可能须要四层的压缩网络而有的工作会须要六层的压缩网络。DynaBERT（dynamic BERT）提出一种不同的思路，它能够通过抉择自适应宽度和深度来灵便地调整网络大小，从而失去一个尺寸可变的网络。

4.2. 模型实现

DynaBERT 的训练阶段包含两局部，首先通过常识蒸馏的办法将 teacher BERT 的常识迁徙到有自适应宽度的子网络 student $DynaBERT_W$ 中，而后再对 $DynaBERT_W$ 进行常识蒸馏失去同时反对深度自适应和宽度自适应的子网络 DynaBERT。训练过程流程图如图 1 所示。

<center> 图 1：DynaBERT 的训练过程 </center>

宽度自适应 Adaptive Width

一个规范的 transfomer 中蕴含一个多头注意力（MHA）模块和一个前馈网络（FFN）。在论文中，作者通过变换注意力头的个数 $N_h$ 和前馈网络中中间层的神经元个数 $d_{ff}$ 来更改 transformer 的宽度。同时定义一个缩放系数 $m_w$ 来进行剪枝，保留 MHA 中最右边的 $[m_wN_H]$ 个注意力头和 FFN 中 $[m_wd_{ff}]$ 个神经元。

为了充分利用网络的容量，更重要的头部或神经元应该在更多的子网络中共享。因而，在训练宽度自适应网络前，作者在 fine-tuned BERT 网络中依据注意力头和神经元的重要性对它们进行了排序，而后在宽度方向上以降序进行排列。这种选取机制被称为 Network Rewiring。

<center> 图 2: Network Rewiring</center>

那么，要如何界定注意力头和神经元的重要性呢？作者参考 P. Molchanov et al., 2017 和 E. Voita et al., 2019 两篇论文提出，去掉某个注意力头或神经元前后的 loss 变动，就是该注意力头或神经元的重要水平，变动越大则越重要。

训练宽度自适应网络

首先，将 BERT 网络作为固定的老师网络，并初始化 $DynaBERT_W$。而后通过常识蒸馏将常识从老师网络迁徙到 $DynaBERT_W$ 中不同宽度的学生子网络。其中，$m_w = [1.0, 0.75, 0.5, 0.25]$。

模型蒸馏的 loss 定义为：

$$
L = \lambda_1l_{pred}(y^{(m_w)}, y) + \lambda_2(l_{emb}(E^{(m_w)},E) + l_{hidn}(H^{(m_w)}, H))
$$

其中，$\lambda_1, \lambda_2$ 是管制不同损失函数权重的参数，$l_{pred}, l_{emb}, l_{hidn}$ 别离定义为：

$$
\begin{align}
l_{pred}(y^{(m_w)}, y) &= SCE(y^{(m_w)}, y), \\
l_{emb}(E^{(m_w)}, E) &= MSE(E^{(m_w)}, E), \\
l_{hidn}(H^{(m_w)}, H) &= \sum^{L}_{l=1} MSE(H^{(m_w)}_l, H_l) \\
\end{align}
$$

$l_{pred}$ 代表预测层的 loss，SCE 代表穿插熵损失函数。$l_{emb}$ 代表嵌入层的 loss，MSE 代表均方差损失函数。$l_{hidn}$ 则为暗藏层的 loss。

训练深度自适应网络

训练好宽度自适应的 $DynaBERT_W$ 后，就能够将其作为老师网络训练同时具备宽度自适应和深度自适应的 DynaBERT 了。为了防止宽度方向上的灾难性忘记，在每一轮训练中，仍对不同宽度进行训练。深度调节系数 $m_d$ 对网络层数进行调节，在训练中定义 $m_d = [1.0, 0.75, 0.5]$。深度方向上的剪枝依据 $mod(d+1, \frac{1}{m_d}) = 0$ 来去掉特定层。

模型蒸馏的 loss 定义为：

$$
L = \lambda_1l^{‘}_{pred}(y^{(m_w,m_d)},y^{(m_w)}) + \lambda_2(l^{‘}_{emb}(E^{(m_w,m_d)},E^{(m_w)}+l^{‘}_{hidn}(H^{(m_w,m_d)},H^{(m_w)}))
$$

4.3. 试验后果

依据不同的宽度和深度剪裁系数，作者最终失去 12 个大小不同的 DyneBERT 模型，其在 GLUE 上的成果如下：

<center> 图 3: results on GLUE benchmark</center>

<center> 图 4:Comparison of #parameters, FLOPs, latency on GPU and CPU between DynaBERT and DynaRoBERTa and other methods.</center>

能够看到论文中提出的 DynaBERT 和 DynaRoBERTa 能够达到和 $BERT_{BASE}$ 及 $DynaRoBERTa$ 相当的精度，然而通常蕴含更少的参数，FLOPs 或更低的提早。在雷同效率的束缚下，从 DynaBERT 中提取的子网性能优于 DistilBERT 和 TinyBERT。

5.TinyBERT 蒸馏

5.1. TinyBERT 简介

论文地址：[TinyBERT: Distilling BERT for Natural Language Understanding
](https://arxiv.org/pdf/1909.10351.pdf)

预训练模型的提出，比方 BERT，显著的晋升了很多自然语言解决工作的体现，它的弱小是毫无疑问的。然而他们普遍存在参数过多、模型宏大、推理工夫过长、计算低廉等问题，因而很难落地到理论的产业利用中。TinyBERT 是由华中科技大学和华为诺亚方舟实验室联结提出的一种针对 transformer-based 模型的常识蒸馏办法，以 BERT 为例对大型预训练模型进行钻研。四层构造的 $TinyBERT_{4}$ 在 GLUE benchmark 上能够达到 $BERT_{base}$ 96.8% 及以上的性能体现，同时模型放大 7.5 倍，推理速度晋升 9.4 倍。六层构造的 $TinyBERT_{6}$ 能够达到和 $BERT_{base}$ 同样的性能体现。

<center> 图 1: TinyBERT learning</center>

TinyBERT 次要做了以下两点翻新：

1. 提供一种新的针对 transformer-based 模型进行蒸馏的办法，使得 BERT 中具备的语言常识能够迁徙到 TinyBERT 中去。
2. 提出一个两阶段学习框架，在预训练阶段和 fine-tuning 阶段都进行蒸馏，确保 TinyBERT 能够充沛的从 BERT 中学习到个别畛域和特定工作两局部的常识。

5.2. 模型实现

5.2.1 常识蒸馏

常识蒸馏的目标在于将一个大型的老师网络 $T$ 学习到的常识迁徙到小型的学生网络 $S$ 中。学生网络通过训练来模拟老师网络的行为。$f^S$ 和 $f^T$ 代表老师网络和学生网络的 behavior functions。这个行为函数的目标是将网络的输出转化为信息性示意，并且它可被定义为网络中任何层的输入。在基于 transformer 的模型的蒸馏中，MHA（multi-head attention）层或 FFN（fully connected feed-forward network）层的输入或一些两头示意，比方注意力矩阵 $A$ 都可被作为行为函数应用。

$$
L_{KD} = \sum_{x \in X}L(f^S(x), f^T(x))
$$

其中 $L(\cdot)$ 是一个用于评估老师网络和学生网络之间差别的损失函数，$x$ 是输出文本，$X$ 代表训练数据集。因而，蒸馏的关键问题在于如何定义行为函数和损失函数。

5.2.2 Transformer Distillation

假如 TinyBert 有 M 层 transformer layer，teacher BERT 有 N 层 transformer layer，则须要从 teacher BERT 的 N 层中抽取 M 层用于 transformer 层的蒸馏。$n = g(m)$ 定义了一个从学生网络到老师网络的映射关系，示意学生网络中第 m 层网络信息是从老师网络的第 g(m)层学习到的，也就是老师网络的第 n 层。TinyBERT 嵌入层和预测层也是从 BERT 的相应层学习常识的，其中嵌入层对应的指数为 0，预测层对应的指数为 M + 1，对应的层映射定义为 $0 = g(0)$ 和 $N + 1 = g(M + 1)$。在模式上，学生模型能够通过最小化以下的指标函数来获取老师模型的常识：

$$
L_{model} = \sum_{x \in X}\sum^{M+1}_{m=0}\lambda_m L_{layer}(f^S_m(x), f^T_{g(m)}(x))
$$

其中 $L_{layer}$ 是给定的模型层的损失函数（比方 transformer 层或嵌入层），$f_m$ 代表第 m 层引起的行为函数，$\lambda_{m}$ 示意第 m 层蒸馏的重要水平。

TinyBERT 的蒸馏分为以下三个局部：transformer-layer distillation、embedding-layer distillation、prediction-layer distillation。

Transformer-layer Distillation

Transformer-layer 的蒸馏由 attention based 蒸馏和 hidden states based 蒸馏两局部组成。

<center> 图 2: Transformer-layer distillation</center>

其中，attention based 蒸馏是受到论文 Clack et al., 2019 的启发，这篇论文中提到，BERT 学习的注意力权重能够捕捉丰盛的语言常识，这些语言常识包含对自然语言了解十分重要的语法和共指信息。因而，TinyBERT 提出 attention based 蒸馏，其目标是使学生网络很好地从老师网络处学习到这些语言常识。具体到模型中，就是让 TinyBERT 网络学习拟合 BERT 网络中的多头注意力矩阵，指标函数定义如下：

$$
L_{attn} = \frac{1}{h}\sum^{h}_{i=1}MSE(A^S_i, A^T_i)
$$

其中，$h$ 代表注意力头数，$A_i \in \mathbb{R}^{l\times l}$ 代表学生或老师的第 $i$ 个注意力头对应的注意力矩阵，$l$ 代表输出文本的长度。论文中提到，应用注意力矩阵 $A$ 而不是 $softmax(A)$ 是因为试验结果显示这样能够失去更快的收敛速度和更好的性能体现。

hidden states based 蒸馏是对 transformer 层输入的常识进行了蒸馏解决，指标函数定义为：

$$
L_{hidn} = MSE(H^SW_h, H^T)
$$

其中，$H^S \in \mathbb{R}^{l \times d^{‘}},\quad H^T \in \mathbb{R}^{l \times d}$ 别离代表学生网络和老师网络的隐状态，是 FFN 的输入。$d$ 和 $d^{‘}$ 代表老师网络和学生网络的暗藏状态大小，且 $d^{‘} < d$，因为学生网络总是小于老师网络。$W_h \in \mathbb{R}^{d^{‘} \times d}$ 是一个可训练的线性变换矩阵，将学生网络的暗藏状态投影到老师网络暗藏状态所在的空间。

Embedding-layer Distillation

$$
L_{embd} = MSE(E^SW_e, E^T)
$$

Embedding loss 和 hidden states loss 同理，其中 $E^S,\quad E^T$ 代表学生网络和老师网络的嵌入，他呢和暗藏状态矩阵的形态雷同，同时 $W_e$ 和 $W_h$ 的作用也雷同。

Prediction-layer Distillation

$$
L_{pred} = CE(z^T/t, z^S/t)
$$

其中，$z^S, \quad z^T$ 别离是学生网络和老师网络预测的 logits 向量，$CE$ 代表穿插熵损失，$t$ 是 temperature value，当 $t = 1$ 时，体现良好。

对上述三个局部的 loss 函数进行整合，则能够失去老师网络和学生网络之间对应层的蒸馏损失如下：

$$
\begin{equation}
L_{layer} =
\left\{
\begin{array}{lr}
L_{embd}, & m=0 \\
L_{hidn} + L_{attn}, & M \geq m > 0 \\
L_{pred}, & m = M + 1
\end{array}
\right.
\end{equation}
$$

5.3. 试验后果

<center> 图 3: Results evaluated on GLUE benchmark</center>

作者在 GLUE 基准上评估了 TinyBERT 的性能，模型大小、推理工夫速度和准确率如图 3 所示。试验结果表明，TinyBERT 在所有 GLUE 工作上都优于 $BERT_{TINY}$，并在均匀性能上取得 6.8% 的晋升。这表明论文中提出的常识整顿学习框架能够无效的晋升小模型在上游工作中的性能。同时，$TinyBERT_4$ 以~4% 的幅度显著的晋升了 KD SOTA 基准线（比方，BERT-PKD 和 DistilBERT），参数放大至~28%，推理速度晋升 3.1 倍。与 teacher $BERT_{base}$ 相比，TinyBERT 在保持良好性能的同时，模型放大 7.5 倍，速度晋升 9.4 倍。

更多优质内容请关注公号：汀丶人工智能