降低预测过程计算成本，这些NLP模型压缩方法要知道

编译 | 凯隐

出品 | AI科技大本营（ID:rgznai100）

近年来，基于谷歌Transformer的语言模型在神经机器翻译，自然语言推理和其他自然语言理解任务上取得了长足进展。

通过多种语言模型的平均损失进行自我监督预训练，使得在大范围的语料库上训练的模型可以在许多任务中提高下游性能。然而，大量的参数和计算量仍然是阻碍BERT和其衍生模型部署的难点。

值得庆幸的是，在过去的两年里，我们已经看到了各种各样的技术，可以缩短模型在进行实际预测时消耗的时间。因此，这篇文章主要着眼于在基本模型预训练后可以用于降低预测过程计算成本的方法，主要包含以下方法：

• 数值精度约简：通过降低计算过程中使用的浮点数精度（浮点约简）和量化，来加速计算。
• 计算融合：在计算图中选择节点并进行合并的技巧。
• 网络修剪: 识别和删除网络中不重要的部分。
• 知识提炼: 训练更小的，效率更高的模型来模仿表现力更强，同时计算成本更高的大模型。
• 模块替换：通过替换部分模块来降低模型的深度和复杂度。

数值精度约简

数值精度约简可能是为模型带来预测加速最通用的方法。在过去几年GPU对16位浮点操作的支持性很差，这意味着降低权重和激活值的精度往往不会带来加速，甚至造成减速。英伟达（Nvidia）Volta和图灵张量核架构的引入在一定程度上解决了这个问题，让GPU能够更好的实现高效的16位浮点精度运算。

1、浮点数的表示

浮点类型数据主要存储三种类型的数值信息：符号，指数，分数。传统的32位浮点表示法分别用8位和23位来表示指数和分数（剩下一位应该表示正负），而传统的16位表示法(用于NVIDIA硬件的格式)将32位表示法中的指数和分数部分大致减半。相对于GPU而言，TPU（张量处理单元）则使用一种称为bfloat16的变体，它可以选择将一些位从分数移动到指数，即牺牲一定的精度来换取表示更大范围的值的能力。

三种不同的浮点数表示方法

Transformer网络的大部分结构都可以直接转换为16位浮点数的权重和激活值，且不会带来计算准确率的下降。而网络的一小部分，特别是softmax操作部分必须保留为32位浮点精度。这是因为大量的小数值(我们的logits)累积起来就可能成为错误的来源。由于这样做同时使用了16位和32位精度的浮点数值，因此这种方法通常称为“混合精度”。

低精度的数值表示主要可以从两方面实现加速：

（1）机器自带的半精度指令（速度更快）

（2）更大的batch size（得益于更紧凑的表示）

NVIDIA已经发布了一套相当广泛的与浮点精度降低相关的基准测试——实际上，这种方法可以使速度提高3倍。

链接：https://github.com/NVIDIA/DeepLearningExamples

2、整型量化

32位浮点到8位整数型值的量化也能实现加速，但需要更细节的实现。特别的，为了确保8位整型的计算结果与32位浮点计算结果尽可能相同，还需要在训练后对输入进行校准。

如果已经知道网络激活值可能所占用的范围，可以将该范围划分为256个离散块，并将每个块分配给一个整数。只要存储了比例因子和占用的范围，就可以使用整数近似来进行矩阵乘法，并将乘法输出恢复为浮点值。

两种不同的整型量化方法

一种简单的整型量化方法，如上图右侧，首先选择一个范围和偏移量，这样就不会将校准后的输入浮点激活值激活映射到unit8值（-128,127）两端的整数。但这种做法在适应两端极值的同时牺牲了一些精度。如上图左侧，为了解决这个问题，TensorRT之类的框架会选择合适的scale和offset值，以最小化模型的32位浮点版本和8位整型版本的输出激活值之间的KL差异。这使我们能够以合理的方式在取值范围和计算精度之间权衡取舍。由于KL散度可以被看作激活值在不同编码（不同精度可以看做不同编码）下的信息损失的度量，因此是一种自然合理的度量方式。

关于如何使用英伟达TensorRT来将8位整型量化应用到个人模型的更多详细细节，可以参考一下链接：

https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html

https://github.com/tensorflow/tensorrt

3、网络层融合与计算图优化

除了浮点数约简和整型量化外，操作融合提供了一个实用的，通用的选择来实现更高效的预测。操作融合的基本原理是将计算流程中若干层所执行的操作进行合并，以避免对全局设备存储器的冗余访问，进而提高执行效率。通过将多个计算操作合并到一个内核中，可以加速内存读写速度。

操作合并示例

如上图，我们可以将添加的skip连接与层规范化操作（LN）的比例（scale）和偏差（bias）合并在一起。

软件优化允许我们重构一些矩阵乘法操作以更好地进行计算并行话。如下图，我们可以将self-attention层的查询向量（Query）、键向量（Key）和值向量（Value）的计算投影合并到一个矩阵乘法中。

对自注意力层进行优化示意图

遗憾的是，从这类计算流程图优化中很少看到关于加速幅度的详细信息，但这种改进仍然是有效的，大约有10%的吞吐量提升。

网络修剪

除了纯粹的软件层面优化，还有其他许多方法可用来修剪网络结构，删除对最终模型贡献小的权重。许多剪枝方法（如Fan等人提出的"Reducing Transformer Depth on Command With Structured Dropout"）需要在训练前对网络进行修改，以生成足够稀疏的模型，并在训练后进行剪枝。剪枝文献中的其他论文聚焦于在没有具体预测目标的情况下，模型所学习的连接模型如何才能更加稀疏。（例如Gorden等人提出的"Compressing BERT: Studying the Effects of Weight Pruning on Transfer Learning"）

虽然这些方法本身都很有意义（而且结构化的删除部分层的方法在实际应用中有重要前景），但更有趣的是一种能以特定方式应用并且仍然能获得性能提升的方法。这类方法基于一个公认事实来对模型进行修剪，即解决特定任务只需要模型的一部分。

为了获得经验上的模型性能提升而进行修剪需要结构化的稀疏性。简单地将单个权值归零并不足以获得性能提升，因为我们没有实际的方法来利用这种稀疏性。因此，我们必须设法砍掉网络中的大部分结构，以获得实际的性能提升。

注意力头剪枝

在文献"Are Sixteen Heads Really Better than One?"中，作者通过迭代的方法从BERT模型中逐步去除注意力头（attention head）。他们使用了一种基于梯度检测的方法(对下游任务进行梯度估计)来估计每个注意力头的重要性，并通过绘制性能--去除的注意力头所占百分比函数来测试模型对注意力头剪枝的鲁棒性。

如下图，在实践中，作者发现20 - 40%的注意力头可以修剪，且对模型准确性的影响可以忽略不计。

基于门控的模型修剪

在J.S. McCarley，Rishav Chakravarti和Avirup Sil合著的《Structured Pruning of a BERT-based Question Answering Model》一书中，坐着探索了一种更通用的模型剪枝方法。不仅关注注意力头，还对每一层的输入以及每一BERT层的前馈层的激活值进行了门控。

他们探索了一些距离度量机制，从而有选择的对网络进行修剪，这包括Michel等人提出的重要措施。但最终确定了一个L0正则化项（添加到损失函数中），并且该项可以在超参调优时进行设置（作为超参数之一），来提高模型的稀疏性。为了使这个L0正则化项可微，他们使用了一个类似于在变分自动编码器中使用的重参数化技巧。

在他们的实验中他们发现通过正则化项进行稀疏惩罚，并通过参数调优得到的结果优于“Are 16 Heads Really Better than 1”使用的基于重要性的估计方法，并且他们发现可以另外删除近50%的前馈激活值，且在基准任务上对最终性能影响微乎其微。

为了获得更多提升，作者还使用了知识蒸馏技术。

知识蒸馏

1、发展历史

知识蒸馏首先由Geoffrey Hinton, Oriol Vinyals, 和 Jeff Dean 在2015年的工作"Distilling the Knowledge in a Neural Network"中提出。知识蒸馏主要通过损失修正的方法，将教师模型中包含的知识迁移到学生模型中。

首先，假设我们可以访问大量未标记的样本。如果我们相信教师模型的预测结果是足够准确的，但是教师模型在实际应用中部署使用太麻烦或计算成本太昂贵，我们可以使用教师模型来预测未标记样本池中的目标类，并将这些目标作为监督样本用于学生模型的训练。如果不是生成与最大可能性类相对应的硬目标，而是生成所有可能类的概率分布，那么学生模型将获得更多的信息丰富的监督样本。

直观来讲，学生模型的一些预测错误比其他错误更加合理，如下图，将哈士奇误认为勺子是荒谬的，但是将哈士奇误认为雪橇犬则是合理的。即设置的损失函数应该反映错误的严重程度。通过惩罚教师模型预测结果和学生模型预测之间的差异(即鼓励二者预测结果相同)，学生网络可以从教师网络给出的预测中学习更多有意义的信息。在语音识别任务中，教师网络的大部分网络性能可以通过仅约3%的训练数据来重现。

有证据表明高参数量可能对样本的有效学习至关重要，而且，就相同时间而言，以一个固定的复杂度训练大型模型可能比训练一个等价的紧凑模型效果更好。因此有效地把教师模型学到的知识传递给被压缩的学生模型是很有前途的。

2、相似模型的知识迁移

在上面讨论文章"Structured Pruning of a BERT-based Question Answering Model"中，通过知识蒸馏将未修剪的教师模型中包含的知识传递给修剪过的学生模型。在自然问题数据集上，教师模型在长答案和短答案问题上的F1分数分别为70.3和58.8。随着大约50%的注意力头和前馈激活值的削减，F1分说分别下降到67.8和55.5，即大约有2.5的下降。但如果用蒸馏损失来代替超参调优过程中的交叉熵损失，那么性能将恢复1.5到2个点，达到69.3和58.4。

不同的模型压缩方法带来的性能损失对比

由Victor Sanh，Lysandre starter，Julien Chaumond和Thomas Wolf提出的“DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter”，在一个语言建模任务的二级预训练步骤中，从基本BERT模型向一个6层的BERT压缩后的学生模型执行知识蒸馏。学生模型(以任务不可知的方式训练)在GLUE基准测试上保持了97%的模型性能，同时减少了60%的预测时间。

在文章“TinyBERT: Distilling BERT for Natural Language Understanding”中，作者采用了从BERT模型到一个4层，隐藏神经元个数为312的学生模型的知识蒸馏与迁移。他们在预训练和调参时都进行了迁移，得到的模型在GLUE基准测试上的达到了BERT-base性能的96%，且模型相对于BERT-base缩小了7.5倍，预测推断速度提高了近10倍。

在文章"Patient Knowledge Distillation for BERT Model Compression"中，作者将知识蒸馏损失函数应用于12层BERT教师模型以及6层学生模型，这在大约5/6的GLUE任务中得到了精度提升（相对于仅将知识蒸馏应用于模型本身）。

3、不同结构模型的知识迁移

在目前讨论的论文中，教师模型和学生模型具有相同的基本架构，学生模型通常使用教师模型的权重来进行初始化。然而，即使在教师和学生模型结构差异很大的情况下，也可以应用知识蒸馏损失来让二者的预测结果相近，从而将教师模型学习到的知识迁移到学生模型。

在文章"Training Compact Models for Low Resource Entity Tagging using Pre-trained Language Models"中，作者首先在命名实体识别任务上训练了一个BERT教师模型（参数量约330M），然后将其迁移到一个更紧凑更高效的CNN-LSTM学生模型上（参数量约3M），这样做使得他们在CPU硬件上以最小的精度损失实现了高达2个数量级的速度提升。

在文章"Distilling Transformers into Simple Neural Networks with Unlabeled Transfer Data"中，作者将BERT-Base和BERT-Large迁移到一个双向LSTM学生模型上，使得学生模型在4种分类任务（Ag News, IMDB, Elec, and DBPedia）上都能达到和教师模型相同的准确度，且参数量削减为13M。他们还发现蒸馏使得样本效率大大提高，每个任务只需要用大约500个带标签的样本进行训练，就能让学生模型达到和教师模型相同的水平（提供足够的未标记样本进行测试）。

在文章"Distilling Task-Specific Knowledge from BERT into Simple Neural Networks"中，作者Lin报告了在使用参数小于1M的单层BiLSTM执行各种句子配对任务(QQP、MNLI等)时的类似结果。

在文章"Attentive Student Meets Multi-Task Teacher: Improved Knowledge Distillation for Pretrained Models"中，作者将多任务学习与知识蒸馏相结合，将Transformer教师模型通过注意力机制迁移到deep LSTM学生模型上。文章指出，从知识蒸馏中得到的提升与多任务学习框架带来的泛化提升是一样的，并且预测速度是纯粹知识精馏的30倍，是TinyBERT的7倍。

知识蒸馏是最近很流行的方法，原因很明显——它很可能成为许多基于Transformer的语言模型应对逐渐增加的参数量的有效方法。如果我们想要尽可能地利用GPU，那么我们就需要通过知识蒸馏这样的方法来保持高预测速度。

模块替换

将要介绍的最后一篇论文有两个特点，一是采用了一种更新颖的模型压缩方法，二是与下面展示的现代艺术作品一起发表：

“BERT-of-Theseus:Compressing BERT by Progressive Module Replacing”是由Canwen Xu等提出的工作，不同于之前的工作仅训练一个单独的学生模型来最小化知识蒸馏损失，BERT-of-Theseus在调试阶段随机地用一个新的模块（后继模块）来替换原来的模块（前继模块），且在每个训练批次都随机替换。

BERT-of-Theseus是思想实验“Ship of Theseus”的一个版本，该实验探究的是一艘船经过一点一点的修复和升级后，是否仍然是原来的物体。BERT-of-Theseus将这种逐步替换思想应用到模型压缩的思想中。

后继模块通常是前继模块的廉价版本——在本例中，单个变压器层替换了一个由2个变压器层组成的块。与知识蒸馏不同，这里没有使用损失来鼓励后续模块模仿它们的前辈，而只是通过简单的让继任模块和前任模块可以互换使用，使得继任者学习模仿前任的行为。

这种隐式模仿行为的一个优点是，我们不再需要选择如何重量各种知识蒸馏损失目标模型的损失,不同于知识蒸馏使用的L0正则化损失 ,不存在二次训练的步骤,与下游整合压缩并发执行。最后，渐进式模块替换方法也适用于整个模型结构，且它在设计中没有利用Transformer模型的任何特定特性，因此泛化能力更强。

作者用线性学习率变化策略进行了实验，发现随着时间的推移，线性增加的模块替换率比恒定的替换率效果更好。

为了测试这种方法的鲁棒性，作者在将BERT-base应用到GLUE基准测试时使用了Theseus压缩，并且轻松地超过了其他几个基于知识蒸馏的方法，这些方法在将原始模型压缩到50%的大小时，常常落后于BERT-base不到1个点。

渐进式模块替换很有吸引力的部分原因是它提供了一种新的训练方法，可以用其他方法进行试验来增加模型吞吐量，而这些方法通常需要对scratch进行重新训练。对于独立的研究人员和较小的公司来说，从头开始重新培训transformer模型通常是很困难的，因此很难利用那些提出了提高模型效率的有用想法但没有发布预先培训过的模型的论文（即没有开源的论文）。

论文地址:

https://arxiv.org/pdf/2002.11985.pdf

相关链接地址:

http://mitchgordon.me/machine/learning/2019/11/18/all-the-ways-to-compress-BERT.html