ACL 2022 | DQ-BART：基于联合蒸馏和量化的高效Seq2Seq模型

©PaperWeekly 原创 · 作者 |

BNDSBilly

单位 | 中科院软件所

研究方向 | 自然语言处理

Abstract

大规模预训练端到端模型如 BART 和 T5 已经在各类 NLP 任务上取得了 sota 表现。然而，由于它们的大内存需求和高延迟，这些模型在资源受限的场景中的应用受到了巨大的挑战。

为解决这个问题，本文同时使用了模型蒸馏和模型量化两种方式，将 BART 模型压缩了 16.5 倍，在多个摘要和 QA 数据集上达到了与原模型相当的表现。

论文标题：

DQ-BART: Efficient Sequence-to-Sequence Model via Joint Distillation and Quantization

论文链接：

https://arxiv.org/abs/2203.11239

Introduction

预训练端到端模型如 BART 和 T5 在各类 NLP 任务（如文本摘要，机器翻译，QA，信息抽取等）上取得了很大的成功。然而，这些大规模预训练模型参数量已经达到了几亿甚至数十亿，并且还在不断增加。这导致推理期间的计算和内存资源需求很大，且很难部署到实际场景，尤其是实时及资源受限的场景。

以上问题促进了对模型压缩的研究，模型压缩可以使大规模预训练模型变快、变小，且保持与原模型相当的表现。近期模型量化收到了较多的关注，因为它不需要改变精心设计的模型结构，只需要用较少的位数来表示模型权重，从而降低模型规模。然而，对编码器-解码器结构的 transformers 模型量化的研究较少。Prato 等人利用 8-bit 量化压缩了端到端的 transformer 模型。但因为精度损失问题，不能继续压缩到 4-bit，且并不是针对大规模预训练模型的量化，也仅仅面向了机器翻译；Shleifer 等人对 BART 在文本摘要的应用上采用了模型蒸馏，但他们的工作并没有显著减少模型空间。

本文针对以下两个问题进行了研究：

量化后的端到端模型能否在各类NLP任务上都取得较好表现？
是否可以结合模型量化和模型蒸馏更进一步的压缩端到端预训练模型？

因此，作者利用了一个高精度的端到端教师模型进行模型蒸馏及量化，训练出了一个层数较少，且低位数的学生模型，缩小了原模型的 16.5 倍。本文还进一步将模型压缩到了 27.7 倍，并展示了预训练模型做生成任务的性能-效率权衡效果。本文也是第一篇探究端到端预训练模型采用模型蒸馏和模型压缩在语言生成任务表现的研究。

Method

3.1 Quantization

模型量化的一般思想是将高精度的实数映射到低精度的整数，从而达到降低模型大小的目的，同时还能加快模型运算速度。目前已经有很多针对神经网络的量化方法研究。包括线性量化、非线性量化、基于近似的量化方法和损失感知量化等。本文结合了线性量化和近似量化方法，对隐藏层和大部分嵌入层的参数进行了量化，没有量化 positional embeddings，激活层量化到了 8bit，以提升矩阵运算的速度。

假设为第次迭代时需要量化的权重矩阵所有列堆叠得到的向量，我们通过如下方式计算量化因子和低精度数：

其中，为我们量化的目标位数。计算出后，我们可以直接计算。由此，我们就可以将表示为的形式，从而完成量化。

3.2 Distillation

模型蒸馏的一般思想是训练较小的学生模型，来模拟较大的教师模型的表现，包括下游任务的概率输出、attention 层以及隐藏层的输出等。我们通过从教师模型的最大间隔层复制权重来初始化学生模型，例如用层教师模型训练层学生模型时，我们分别选取层教师模型中的第层初始化学生模型的层，从而达到“最大间隔”，能够获取更多信息。如果学生模型只有层，我们就选取教师模型的最后一层而不是第一层，因为已有实验证明选取最后一层，模型表现更好。

模型蒸馏的目标函数包括下游任务的损失函数以及蒸馏损失，其中蒸馏损失包含三个部分，均为教师模型和学生模型在三种表现上的均方误差（MSE）值：输出概率、注意力层输出（包括编码器和解码器分别的注意力层输出误差以及 cross attention）以及隐藏层输出（包含编码器和解码器分别的隐藏层状态误差）。

3.3 Distillation-aware quantization

本文结合了模型量化和模型蒸馏，提出了基于蒸馏感知的量化方法：在每一轮循环中，我们首先将全精度学生模型（full-precision student model）量化为低精度模型（quantized student model），利用模型蒸馏的目标函数作为量化的损失函数，教师模型为全精度模型。得到损失函数后，我们回过头用其更新中的参数，得到新的，并进入下一轮循环。

Experiments

本文在三个 benchmark 上评估了蒸馏感知量化 BART 在文本摘要和长格式问答上的表现：CNN/DailyMail（文本摘要）, XSUM（文本摘要），ELI5 （长格式问答）。

4.1 Experimental Setup

本文选取的对比模型为基本大小的 BART（文本摘要和 QA）以及 mBART-large （机器翻译）。作者在仅使用了模型量化的 BART 上训练了个 epoch，在蒸馏感知的量化 BART 上训练了个 epoch。设置 batch-size 为，学习率为，的 lindar warmup，并根据开发集上的 rouge-L 分数选择最佳模型（rouge-L 分数是 2004 年 Chin-Yew Lin 提出的一种自动摘要评价方法，用于评估自动摘要以及机器翻译。是一种面向 n 元词召回率的评价方法，分数越高表明模型表现越好）。根据 Lewis 等人在 2020 年的工作设置了超参数，所有实验均在 A100 GPU 上进行。

4.2 DQ-BART 在文本摘要和问答上实验结果

本文提出的各个模型在 CNN/DailyMail benchmark 上的表现如下表所示：

将其可视化为模型表现与模型大小压缩比的关系图，如下所示：

实验分析如下：

仅使用量化方法压缩模型，在生成任务上表现很差。如图表所示，rouge-L分数（Recall-Oriented Understudy for Gisting Evaluation）与原模型相比降低了约。
位的蒸馏感知量化模型的性能与原模型相当甚至更好，说明对于生成模型如 BART 来说，量化到位是很容易的。
当量化到位时（），本文能够实现倍的模型大小压缩比，相比于原模型，仅摘要任务的性能略有下降，长格式 QA 任务的性能仍保持相当。当我们结合模型蒸馏后（减少编码器和解码器的层数），可以进一步将模型压缩到倍（），而不会过度牺牲性能。
当将模型压缩到极限时（），模型被压缩了倍，此时在三个 benchmark 上仍保持可接受的性能：CNN/DailyMail 的 rouge-L 下降了，XSUM 下降了，ELI5 下降了。这表明在实际应用中，选取的压缩比要基于具体选择的任务。

4.3 DQ-mBART 在机器翻译上的实验结果

由于普通 BART 没有多语言模型，所以本文在 mBART-large 上使用了模型蒸馏和模型量化，来探究其在机器翻译上的表现。在 WMT 英文-罗马尼亚语翻译任务上进行了实验，结果如下表所示：

如表所示，蒸馏感知的量化方法取得了较好的表现，但是当量化到位时，模型性能出现了大幅下降。这可能是因为位数较低时，深层模型的累积精度损失导致的。未来的工作可能会探索如何在低位设置下提高深度模型的联合蒸馏和量化的性能。

4.4 蒸馏感知的量化与仅采用模型蒸馏的对比

本文还将蒸馏感知的量化方法与仅采用模型蒸馏方法的模型表现进行了对比，如上表所示，本文提出的结合方法在文本摘要和问答任务上的表现更好。

Conclusion

基于 Transformer 的预训练端到端语言模型（如 BART）极大地提升了一系列 NLP 任务的最新技术水平。然而，这些极其大规模的模型在资源受限的情况下的应用面临巨大挑战。针对这个问题，本文提出了 DQ-BART，一种联合蒸馏和量化的 BART 模型。经验结果表明，在倍的模型压缩比下，DQ-BART 在三个生成基准测试中均取得了与原模型相当的表现。

此外，本文还研究了机器翻译任务中 mBART 的蒸馏和量化，并强调了结合这两种压缩方式在跨语言任务上更深层次模型的挑战。本文是第一个在预训练端到端语言模型上应用联合量化和蒸馏的工作。

更多阅读