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

2023-06-16
浙江
本文字数：7625 字
阅读完需：约 25 分钟

深度学习实践篇[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

图 1: Vanilla KD 和 PKD 比较

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

2.2. 模型实现

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

PKD-Skip: 从每 k 层学习，这种策略是假设网络的底层包含重要信息，需要被学习到（如图 2a 所示）
PKD-last: 从最后 k 层学习，假设教师网络越靠后的层包含越丰富的知识信息（如图 2b 所示）

图 2a: PKD-Skip 学生网络学习教师网络每两层的输出图 2b: PKD-Last 学生网络从教师网络的最后六层学习

因为在 BERT 中仅使用最后一层的[CLS] token 的输出来进行预测，且在其他 BERT 的变体模型中，如 SDNet，是通过对每一层的[CLS] embedding 的加权平均值进行处理并预测。由此可以推断，如果学生模型可以从任何教师网络中间层中的[CLS]表示中学习，那么它就有可能获得类似教师网络的泛化能力。

因此，Patient-KD 中提出特殊的一种损失函数的计算方式：

L_{P T} = i = 1 \sum N j = 1 \sum M ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ \frac{h _{i, j}^{s}}{∣ ∣ ∣ h _{i, j}^{s} ∣ ∣ ∣ 2} - \frac{h i , I _{p t} ( j ) ^{t}}{∣ ∣ ∣ ∣ h _{i, I_{p t} (j)}^{t} ∣ ∣ ∣ ∣ _{2}} ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣_{2}^{2}

其中，对于输入 $x_{i}$ ，所有层的[CLS]的输出表示为：

h_{i} = [h_{i, 1}, h_{i, 2}, . . ., h_{i, k}] = B E R T_{k} (x_{i}) \in R^{k \times d} ​

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

同时，Patient-KD 中也使用了 $L_{D S}$ 和 $L_{C E}^{S}$ 两种损失函数用来衡量教师和学生网络的预测值的距离和学生网络在特定下游任务上的交叉熵损失。

L_{D S} = - i \in [N] \sum c \in C \sum [P^{t} (y_{i} = c ∣ x_{i}; \hat{θ}^{t}) \cdot l o g P^{s} (y_{i} = c ∣ x_{i}; θ^{s})]

L_{C E}^{s} = - i \in [N] \sum c \in C \sum 1 [y_{i} = c] \cdot l o g P^{s} (y_{i} = c ∣ x_{i}; θ^{s})]

最终的目标损失函数可以表示为：

L_{P K D} = (1 - α) L_{C E}^{S} + α L_{D S} + β L_{P T}

2.3. 实验结果

图 3: results from the GLUE test server

作者将模型预测提交到 GLUE 并获得了在测试集上的结果，如图 3 所示。与 fine-tuning 和 vanilla KD 这两种方法相比，使用 PKD 训练的 $B E R T_{3}$ 和 $B E R T_{6}$ 在除 MRPC 外的几乎所有任务上都表现良好。其中，PKD 代表 Patient-KD-Skip 方法。对于 MNLI-m 和 MNLI-mm，六层模型比微调（FT）基线提高了 1.1%和 1.3%，

我们将模型预测提交给官方 GLUE 评估服务器以获得测试数据的结果。结果总结在表 1 中。与直接微调和普通 KD 相比，我们使用 BERT3 和 BERT6 学生的 Patient-KD 模型在除 MRPC 之外的几乎所有任务上都表现最好。此外，6 层的 $B E R T_{6} - P K D$ 在 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 在大数据集上的表现往往更好。

图 4: PKD-Last 和 PKD-Skip 在 GLUE 基准上的对比

尽管这两种策略都比 vanilla KD 有所改进，但 PKD-Skip 的表现略好于 PKD-Last。作者推测，这可能是由于每 k 层的信息提炼捕获了从低级到高级的语义，具备更丰富的内容和更多不同的表示，而只关注最后 k 层往往会捕获相对同质的语义信息。

图 5: 参数量和推理时间对比

图 5 展示了 $B E R T_{3}$ 、 $B E R T_{6}$ 、 $B E R T_{1} 2$ 的推理时间即参数量, 实验表明 Patient-KD 方法实现了几乎线性的加速， $B E R T_{6}$ 和 $B E R T_{3}$ 分别提速 1.94 倍和 3.73 倍。

3.DistilBERT 蒸馏

3.1. DistilBERT 简介

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

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

近年来，大规模预训练语言模型成为 NLP 任务的基本工具，虽然这些模型带来了显著的改进，但它们通常拥有数亿个参数（如图 1 所示），而这会引起两个问题。首先，大型预训练模型需要的计算成本很高。其次，预训练模型不断增长的计算和内存需求可能会阻碍语言处理应用的广泛落地。因此，作者提出 DistilBERT，它表明小模型可以通过知识蒸馏从大模型中学习，并可以在许多下游任务中达到与大模型相似的性能，从而使其在推理时更轻、更快。

3.2. 模型实现

学生网络结构

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

Training loss

$L_{c e}$ 训练学生模仿教师模型的输出分布：

L_{c e} = i \sum t_{i} * l o g (s_{i})

其中， $t_{i}$ 和 $s_{i}$ 分别是教师网络和学生网络的预测概率。

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

p_{i} = \frac{e x p ( z _{i} / T )}{\sum _{j} e x p ( z _{j} / T )}

其中， $T$ 控制输出分布的平滑度，当 T 变大时，类别之间的差距变小；当 T 变小时，类别间的差距变大。 $z_{i}$ 代表分类 $i$ 的模型分数。在训练时对学生网络和教师网络使用同样的 temperature $T$ ，在推理时，设置 $T = 1$ ，恢复为标准的 softmax 最终的 loss 函数为 $L_{c e}$ 、Mask language model loss $L_{m l m}$ （参考 BERT）和 cosine embedding loss $L_{c o s}$ （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 $D y n a B E R T_{W}$ 中，然后再对 $D y n a B E R T_{W}$ 进行知识蒸馏得到同时支持深度自适应和宽度自适应的子网络 DynaBERT。训练过程流程图如图 1 所示。

图 1： DynaBERT 的训练过程

宽度自适应 Adaptive Width

一个标准的 transfomer 中包含一个多头注意力（MHA）模块和一个前馈网络（FFN）。在论文中，作者通过变换注意力头的个数 $N_{h}$ 和前馈网络中中间层的神经元个数 $d_{f f}$ 来更改 transformer 的宽度。同时定义一个缩放系数 $m_{w}$ 来进行剪枝，保留 MHA 中最左边的 $[m_{w} N_{H}]$ 个注意力头和 FFN 中 $[m_{w} d_{f f}]$ 个神经元。

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

图 2: Network Rewiring

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

训练宽度自适应网络

首先，将 BERT 网络作为固定的教师网络，并初始化 $D y n a B E R T_{W}$ 。然后通过知识蒸馏将知识从教师网络迁移到 $D y n a B E R T_{W}$ 中不同宽度的学生子网络。其中， $m_{w} = [1.0, 0.75, 0.5, 0.25]$ 。

模型蒸馏的 loss 定义为：

L = λ_{1} l_{p r e d} (y^{(m_{w})}, y) + λ_{2} (l_{e m b} (E^{(m_{w})}, E) + l_{h i d n} (H^{(m_{w})}, H))

其中， $λ_{1}, λ_{2}$ 是控制不同损失函数权重的参数， $l_{p r e d}, l_{e m b}, l_{h i d n}$ 分别定义为：

$l_{p r e d}$ 代表预测层的 loss，SCE 代表交叉熵损失函数。 $l_{e m b}$ 代表嵌入层的 loss，MSE 代表均方差损失函数。 $l_{h i d n}$ 则为隐藏层的 loss。

训练深度自适应网络

训练好宽度自适应的 $D y n a B E R T_{W}$ 后，就可以将其作为教师网络训练同时具备宽度自适应和深度自适应的 DynaBERT 了。为了避免宽度方向上的灾难性遗忘，在每一轮训练中，仍对不同宽度进行训练。深度调节系数 $m_{d}$ 对网络层数进行调节，在训练中定义 $m_{d} = [1.0, 0.75, 0.5]$ 。深度方向上的剪枝根据 $m o d (d + 1, \frac{1}{m _{d}}) = 0$ 来去掉特定层。

模型蒸馏的 loss 定义为：

L = λ_{1} l^{^{'}} p r e d (y^{(m_{w}, m_{d})}, y^{(m_{w})}) + λ_{2} (l^{^{'}} e m b (E^{(m_{w}, m_{d})}, E^{(m_{w})} + l_{h i d n}^{^{'}} (H^{(m_{w}, m_{d})}, H^{(m_{w})}))

4.3. 实验结果

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

图 3: results on GLUE benchmark

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

可以看到论文中提出的 DynaBERT 和 DynaRoBERTa 可以达到和 $B E R T_{B A S E}$ 及 $D y n a R o B E R T a$ 相当的精度，但是通常包含更少的参数，FLOPs 或更低的延迟。在相同效率的约束下，从 DynaBERT 中提取的子网性能优于 DistilBERT 和 TinyBERT。

5.TinyBERT 蒸馏

5.1. TinyBERT 简介

论文地址：TinyBERT: Distilling BERT for Natural Language Understanding

预训练模型的提出，比如 BERT，显著的提升了很多自然语言处理任务的表现，它的强大是毫无疑问的。但是他们普遍存在参数过多、模型庞大、推理时间过长、计算昂贵等问题，因此很难落地到实际的产业应用中。TinyBERT 是由华中科技大学和华为诺亚方舟实验室联合提出的一种针对 transformer-based 模型的知识蒸馏方法，以 BERT 为例对大型预训练模型进行研究。四层结构的 $T i n y B E R T_{4}$ 在 GLUE benchmark 上可以达到 $B E R T_{b a s e}$ 96.8%及以上的性能表现，同时模型缩小 7.5 倍，推理速度提升 9.4 倍。六层结构的 $T i n y B E R T_{6}$ 可以达到和 $B E R T_{b a s e}$ 同样的性能表现。

图 1: TinyBERT learning

TinyBERT 主要做了以下两点创新：

提供一种新的针对 transformer-based 模型进行蒸馏的方法，使得 BERT 中具有的语言知识可以迁移到 TinyBERT 中去。
提出一个两阶段学习框架，在预训练阶段和 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_{K D} = x \in X \sum 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_{m o d e l} = x \in X \sum \sum M + 1 m = 0 λ_{m} L l a y e r (f_{m}^{S} (x), f_{g (m)}^{T} (x))

其中 $L_{l a y e r}$ 是给定的模型层的损失函数（比如 transformer 层或嵌入层）， $f_{m}$ 代表第 m 层引起的行为函数， $λ_{m}$ 表示第 m 层蒸馏的重要程度。

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

Transformer-layer Distillation

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

图 2: Transformer-layer distillation

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

L_{a t t n} = \frac{1}{h} i = 1 \sum h M S E (A_{i}^{S}, A_{i}^{T})

其中， $h$ 代表注意力头数， $A_{i} \in R^{l \times l}$ 代表学生或教师的第 $i$ 个注意力头对应的注意力矩阵， $l$ 代表输入文本的长度。论文中提到，使用注意力矩阵 $A$ 而不是 $s o f t m a x (A)$ 是因为实验结果显示这样可以得到更快的收敛速度和更好的性能表现。

hidden states based 蒸馏是对 transformer 层输出的知识进行了蒸馏处理，目标函数定义为：

L_{h i d n} = M S E (H^{S} W_{h}, H^{T})

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

Embedding-layer Distillation

L_{e m b d} = M S E (E^{S} W_{e}, E^{T})

Embedding loss 和 hidden states loss 同理，其中 $E^{S}, E^{T}$ 代表学生网络和教师网络的嵌入，他呢和隐藏状态矩阵的形状相同，同时 $W_{e}$ 和 $W_{h}$ 的作用也相同。

Prediction-layer Distillation

L_{p r e d} = C E (z^{T} / t, z^{S} / t)

其中， $z^{S}, z^{T}$ 分别是学生网络和教师网络预测的 logits 向量， $C E$ 代表交叉熵损失， $t$ 是 temperature value，当 $t = 1$ 时，表现良好。

对上述三个部分的 loss 函数进行整合，则可以得到教师网络和学生网络之间对应层的蒸馏损失如下：

5.3. 实验结果

图 3: Results evaluated on GLUE benchmark

作者在 GLUE 基准上评估了 TinyBERT 的性能，模型大小、推理时间速度和准确率如图 3 所示。实验结果表明，TinyBERT 在所有 GLUE 任务上都优于 $B E R T_{T I N Y}$ ，并在平均性能上获得 6.8%的提升。这表明论文中提出的知识整理学习框架可以有效的提升小模型在下游任务中的性能。同时， $T i n y B E R T_{4}$ 以 4%的幅度显著的提升了 KD SOTA 基准线（比如，BERT-PKD 和 DistilBERT），参数缩小至 28%，推理速度提升 3.1 倍。与 teacher $B E R T_{b a s e}$ 相比，TinyBERT 在保持良好性能的同时，模型缩小 7.5 倍，速度提升 9.4 倍。

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

发布于: 刚刚阅读数: 3

原文链接:【http://xie.infoq.cn/article/eb9bed498d96806d7d82985e4】。

汀丶人工智能

关注

本博客将不定期更新关于NLP等领域相关知识 2022-01-06 加入

本博客将不定期更新关于机器学习、强化学习、数据挖掘以及NLP等领域相关知识，以及分享自己学习到的知识技能，感谢大家关注！

发布

暂无评论

创作场景