ACL 2022杰出论文：为什么生成式大规模预训练模型难以量化？

2022-05-17 05:05

©PaperWeekly 原创 · 作者 | BNDSBilly

研究方向 | 自然语言处理

Abstract

随着大规模预训练模型（PLMs）越来越大，对其进行模型压缩的需求也越来越高。尽管已经有较多研究针对 BERT 进行了压缩，但是针对一般的生成式 PLMs 的尝试都效果不大，且原因不明。本文提出了可能的两个原因：词嵌入同质化以及权重分布差异，并分别提出了 token 级别的对比学习蒸馏以及模块级别的动态缩放方法针对性的解决这两个问题。实验表明本文提出的压缩方法分别将 GPT-2 和 BART 压缩了 14.4 倍和 13.4 倍，且在多个数据集上均取得了与原模型相当的表现。

论文标题：

Compression of Generative Pre-trained Language Models via Quantization

收录会议：

ACL 2022

论文链接：

https://arxiv.org/abs/2203.10705

Introduction

基于 Transformer 架构的生成式 PLMs 展现了其少样本学习的能力以及在多种任务上的 SOTA 表现。然而过多的参数量以及逐个 token 生成的模式导致模型并行效果不好，计算量和内存需求过大。现有的模型压缩方法主要针对于如 BERT 等的理解型任务，然而，针对生成式 PLMs 的模型压缩一直没有取得如 BERT 一样的成功，其潜在的失败原因也并不清楚。

本文中，作者研究了生成式 PLMs 的量化压缩方法，发现主要有以下两点失败原因：（1）模型精度下降导致的词嵌入同质化；（2）不同模块和 Transformer 层的权重分布差异。由于生成式 PLMs 的从左到右预测的性质，量化误差会逐步累积，导致这两个问题被进一步的放大。

为了解决以上问题，本文提出了 token 级别的对比学习蒸馏来解决词嵌入同质化问题，并提出了模块级动态缩放方法来解决权重分布差异问题。作者将量化得到的 GPT 和 BART 模型（分别命名为 QuanGPT 和 QuanBART）在三种生成式下游任务上分别进行了实验，结果表明，将权重量化到时，可以实现与原模型相当的性能；将权重量化到时，实现了超过倍的压缩，也仅有轻微的表现下降。

‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍

Methods

3.1 Quantization

量化是指将模型权重从全精度（例如位 float）量化到低精度（例如位 int），从而实现模型压缩。本文中，作者使用了[1]中的训练方式：设全精度模型权重向量化为，首先通过一个裁剪因子将其修剪到范围内（不在此范围内的分量会被设为或）并等比例缩放到。然后通过量化函数将其映射到离散集合

中，其中，为量化的目标数，也即权重位宽。量化后权重为，计算其损失函数，并用其梯度来更新全精度模型下的参数 [2]。

从上式可以看出，裁剪因子的选取非常重要，好的裁剪因子应该能够覆盖权重最密集分布的部分，以减少量化误差。现有的一些研究（PACT [3]、LSQ [4] 等）提出了学习的不同参数化方法，但当权重位宽逐渐减小时（也即精度下降时），效果会越来越差。所以需要提出新的学习裁剪因子的方法，具体内容会在 3.3 中提到。

在本文中，作者参考了针对 BERT 的量化方法 [5]，对 Transformer 层的权重采取按层量化（每个权重矩阵一个裁剪因子），对 embedding 层的权重采取按行量化（每个 embedding 一个裁剪因子），对自注意力层和GeLU的激活层采取非对称量化，对其余激活层采取对称量化。

本文利用一些具有代表性的量化方法（包括用于量化 BERT 的 LAQ、常用于计算机视觉的 PACT 和 LSQ）量化了 GPT-2，并比较了其在不同权重位宽下的性能。如下图所示，随着位宽的减小，困惑度上升，性能急剧下降，特别是对于 PACT 和 LSQ。在下文中，作者对其原因进行了分析。

3.2 Token-level Contrastive Distillation

作者首先将全精度模型和量化后模型的词嵌入分布进行了可视化，如下图所示，全精度模型的词嵌入是分散可区分的，而之前的量化方法（如 PACT，LSQ 等）学习到的词嵌入是同质的（聚在一起且不易区分）。作者推测这是由 GPT 的顺序计算性质引起的。具体来说，与 BERT 并行计算不同，GPT 需要以从左到右的顺序计算每个 token，量化误差逐步累积，使得学习噪声逐步增大，导致了最终学习到的词嵌入信息量较少，不同 token 之间的依赖性越低。

本文提出的 token 级别对比学习蒸馏受 [5] 的启发，在学习量化模型中一个 token 的 embedding 时，利用该 token 在全精度教师模型中的 embedding 作为正例，所在序列中不同 token 的 embedding 作为负例，进行对比学习。同时，作者还设置了一个 memory bank 来存储量化模型中各 token 的 embedding，以提升计算 loss 和更新 embedding 的速度。

具体来说，作者使用上标分别表示量化的学生网络和全精度教师网络。将长度为的输入序列记为。假设来自量化和全精度网络的最后一个 Transformer 层的隐藏状态线性投影到了。设为的负例索引集合，则该序列的对比学习蒸馏损失可表示为：

上式中，表示的余弦相似度，为温度。除此之外，作者还考虑了两个模型的输出概率交叉熵作为另一部分的损失函数：

通过设置一个平衡因子，得到总的损失函数。

3.3 Module-dependent Dynamic Scaling

本文还进一步研究了全精度模型中的权重分布，由下图所示，全精度 GPT-2 的权重分布高度偏斜（图中的蓝色背景块），存在异常值，而 PACT 学习裁剪因子时，求取梯度过程中忽略了内权重的影响，这通常会使 PACT 学习的过大（例如（a）中 PACT 学习到的，然而本文提出方法学习到的小很多），进而导致量化后的权重没有良好的分辨率。

在本文中，作者转而学习一个新的因子，然后再根据计算裁剪因子。初始化，这样不仅简化了初始化过程，而且还确保了初始裁剪因子不会偏离全精度权重。本文还修正了 PACT 中求梯度的公式，将内权重的影响也考虑了进来：

其中为计算出的 training loss，为上文提到的，经过上式更新的参数相较于 PACT 对权重的变化分布更不敏感，因为缩放的梯度与当前网络平均权重大小成正比，从而能够缓解权重分布差异的问题。

Experiments

4.1 下游任务实验

本文首先从 huggingface 上获取现成的 GPT-2 和 BART 模型，分别在三个下游任务上进行微调，将得到的网络作为全精度教师模型，并进行量化得到学生模型。作者在 GPT 上进行了语言模型和 NUP (next utterance prediction) 两个下游任务，在 BART 上对摘要任务进行了实验。由于目前很少有面向生成式 PLMs 压缩的研究，所以作者实现了 PACT、LSQ、LAQ 三种量化方法作为对比。

语言模型的目标任务是预测单词序列的概率分布，本文在 Wiki-Text2, PTB 和 WikiText103 三个数据集上进行了实验，并利用困惑度（PPL）评估模型表现。如下表所示，本文量化的模型在下的性能与全精度模型相当，减少到时观察到较明显的性能下降（在 WikiText2 和 WikiText103 上下降了大约个 PPL 点，在 PTB 上下降了不到个 PPL 点）。当位宽进一步下降到 时，本文的量化模型平均下降了 个 PPL 点，但实现了 倍的模型压缩。

NUP 的任务是根据给定的句子生成其下一句，本文中使用了大规模对话数据集 Persona-Chat 做实验。如上表所示，尽管量化后的模型有较为明显的精度损失，但本文量化的模型得到了当前最高的准确率。

抽象摘要旨在生成一个简洁的摘要，以捕捉源文章的主要思想。本文在 XSum 上进行了实验，并用分数来评估此任务的性能。如下表所示，本文提出的方法在摘要任务上也优于其它量化方法。

4.2 消融实验

在对比蒸馏学习中，作者采用了全精度模型中相同 token 的 embedding 作为正例，其余 token 的 embedding 作为负例。为了说明这样选择负例的原因，作者比较了其它四种选取方式：（a）选取全精度模型和量化模型的不同 token 作为负例；（b）仅采用量化模型的不同 token 作为负例；（c）从外部词汇表中随机选取 token 作为负例（global 方法）；（d）使用不同 sequence 中所有 token 的平均 embedding 做负例。

如上表所示，加入了量化模型 token embedding 的方法表现较差，这说明来自未完全训练的量化网络的 embedding 是有噪声、不可靠的。global 方法的表现更差，这可能是因为同一 sequence 中的 token 与当前 token 上下文无关，能比随机标记提供更多信息。使用 sequence-level 负采样的表现最差，这可能是因为生成任务严重依赖更细粒度的标记表示。

如下图所示，本文作者还比较了对比学习中的 loss 函数，实验表明，如果仅考虑模型蒸馏而不加入对比学习，虽然训练时间能够加快，内存消耗能够减小，但是模型表现也会变差。

作者还对比了对比学习时使用的 token 表征的选取，表明选取解码器最后一层输出作为 token 表征效果最好。可能原因是最后一个解码器块的隐藏状态包含来自所有先前层的丰富信息。

Conclusion

本文研究了生成式 PLMs 的量化。首先作者研究发现，量化生成式 PLMs 的困难主要是由于词嵌入同质化和权重分布差异化引起的。为了解决这两个问题，本文提出了 token 级的对比蒸馏方法以学习可区分的词嵌入，提出了模块级动态缩放方法以实现更精确的权重量化。本文量化得到的模型 QuanGPT 和 QuanBART 在语言模型，NUP 和摘要任务上都取得了较好的表现，全部优于现有模型压缩方法。

参考文献

[1] Matthieu Courbariaux, Yoshua Bengio, and Jean-Pierre David. 2015. Binaryconnect: Training deep neural networks with binary weights during propagations. In Advances in neural information processing systems, pages 3123–3131.

[2] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 2013. Estimating or propagating gradients through stochastic neurons for conditional computation. Technical Report arXiv:1308.3432

[3] Jungwook Choi, Zhuo Wang, Swagath Venkataramani, Pierce I-Jen Chuang, Vijayalakshmi Srinivasan, and Kailash Gopalakrishnan. 2018. Pact: Parameterized clipping activation for quantized neural networks. Preprint arXiv:1805.06085.

[4] Steven K Esser, Jeffrey L McKinstry, Deepika Bablani, Rathinakumar Appuswamy, and Dharmendra S. Modha. 2020. Learned step size quantization. In International Conference on Learning Representations.

[5] Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. 2020. Ternarybert: Distillation-aware ultra-low bit bert. In Conference on Empirical Methods in Natural Language Processing

更多阅读