ACL 2022 | 普林斯顿陈丹琦组：模型剪枝的加速方法

2022-05-07 04:05

©PaperWeekly 原创 · 作者 | BNDSBilly

单位 | 中科院软件所

研究方向 | 自然语言处理

Abstract

大规模预训练语言模型规模不断扩大，导致模型压缩受到了越来越多的关注。当下最主流的两种模型压缩方式是模型剪枝和模型蒸馏。前者可以显著降低模型大小，但很难像后者一样大幅提升运行速度；而后者的缺点在于需要大量 unlabeled data 和昂贵的成本进行训练。

所以在本文中，作者提出了一个 task-specific 结构化剪枝方法 CoFi，可以在不使用任何 unlabeled data 的前提下达到与模型蒸馏方法基本相当的准确性和速度提升。在 GLUE 和 SQuAD 数据集上的实验表明，CoFi 生成的压缩模型可以提升超过倍的速度，并且达到与原模型相当的精度。

论文标题：

Structured Pruning Learns Compact and Accurate Models

收录会议：

ACL 2022

论文链接：

https://arxiv.org/abs/2204.00408

代码链接：

https://github.com/princeton-nlp/CoFiPruning

Introduction

当下，预训练的语言模型已成为自然语言处理的中流砥柱。但这些模型在存储和计算时间方面的成本很高，推动了大量针对模型压缩相关的工作，帮助这些模型能够变得更小、更快。

模型剪枝致力于在一个大规模预训练模型中搜索一个更小、但可以与原模型表现相当的子网络，一般针对于解决特定任务。其整体趋势是移除一些细粒度单元，从而使得最终的子网络有较为灵活的结构。但到目前为止，模型剪枝很难实现对速度的大幅提升（最多达到 倍的提升）。

模型蒸馏则一般提前设定好一个较小的学生模型结构，并将原模型作为教师模型。通过在大量 unlabeled data 上进行蒸馏或 fine-tune 对学生模型参数进行训练。尽管这一方法可以大幅提升模型速度，但是需要大量的 unlabeled data 和昂贵的成本进行训练。 如下表所示，为一般的蒸馏和剪枝方法对比。

本文中，作者提出了一个针对特定任务的结构化剪枝方法 CoFi（Coarse and Fine-grained Pruning），主要思路是结合了对粗粒度单元（如自注意力层和前馈层）的剪枝和对细粒度单元（例如 heads 和 hidden dimensions）的剪枝。同时，作者还提出了一种分层蒸馏方法，可以动态的学习教师模型和学生模型之间的层映射关系，从而引入蒸馏的目标函数，帮助提升模型性能。

根据实验表明，CoFi 压缩的模型在 GLUE 和 SQuAD v1.1 数据集上的表现超过了其他的剪枝和蒸馏模型表现，实现了超过倍的模型加速和的参数剪枝，并且保持了原模型以上的准确率。这说明，本文提出的 CoFi 方法是一种可以生成更小更快模型的剪枝方法，并且不需要额外的 unlabeled data 进行一般蒸馏。

Methods

3.1 Coarse- and Fine-Grained Pruning

模型剪枝一般逐步从原模型中移除冗余参数，之前的工作分别侧重于对 Transformer 结构中的不同组件进行剪枝。例如移除整个 Transformer 块（ MHA 和 FFN 层）、移除多头注意力中的某些头、移除 FFN 层中的某些中间层等。

在本文中，作者提出了一种结合粗粒度剪枝和细粒度剪枝的方法，使得剪枝更加灵活。例如，作者加入了两个新的 masks （和）分别控制对整个 MHA 层和 FFN 层的剪枝，而不用再将一对 MHA 和 FFN 看作一个整体：

上式中，和分别为之前工作中对多头注意力中头的剪枝 mask 和对 FFN 层的中间层的剪枝 mask。为注意力机制的输出，分别为前馈层的两个投影矩阵。

同时，作者还引入了来对 MHA 和 FFN 的输出层维度进行剪枝，在对 BERT 的实验中，作者发现输出维度从原来的降为了，但这仍然可以帮助提升模型效率。

通过如上方式，单个参数的剪枝策略不再受限于单个 mask，而是由多个 mask 共同控制，进一步提高了模型的灵活度。新提出的 mask 值参考 [1] [2] 进行学习。

3.2 Distillation to Pruned Models

模型蒸馏的目标函数一般为教师模型和学生模型的输出概率交叉熵，其中下标分别代表学生模型和教师模型。

除此之外，在目标函数中加入中间层的蒸馏对模型表现也有帮助 [3]：设学生模型的第层学习教师模型中第层的输出，则中间层的损失函数为：

其中为投影矩阵。但在模型剪枝方法中，最终得到的学生模型结构是一开始不知道的，所以很难建立层映射关系。本文中，作者提出了一种动态建立层映射的方法：将教师模型的第层去匹配最近的学生模型中的层：

这一过程可以高度并行化，不会过多影响模型效率。同时，为了防止交叉匹配（例如教师模型的高层匹配了学生模型的低层，而教师模型的低层匹配了学生模型的高层），作者还加入了额外限制，规定教师模型某一层只能匹配比上一次匹配的学生模型层更低的层。

综上所述，最终的蒸馏目标函数即为

。

Experiments

4.1 Experiments Setup

本文在 GLUE 和 SQuAD v1.1 两个数据集上进行实验，首先利用蒸馏目标函数对模型进行 finetune，然后利用剪枝逐步增加模型稀疏度。最终得到了从的不同稀疏度的模型。本文使用了 DistillBERT、TinyBERT、DynaBERT 等模型和 FLOP、Movement Pruning 等方法进行结果对比。

4.2 Main Results

如上图所示，作者分别从加速和模型大小两个角度将 CoFi 与其他模型、方法进行了对比，可以看到 CoFi 在各个层面上都比现有压缩模型有更好的表现。从下表也可以看出，CoFi 生成的剪枝模型提供了超过倍的加速比，并且在大部分任务上取得了优于四层 TinyBERT 的表现。

4.3 Ablation Study

作者首先探究了节提出的结合粗细粒度剪枝方法的有效性。如下表所示，作者分别删除了和、两类掩码，并探究了和两种稀疏度下模型的表现。可以看到，去除隐藏层维度掩码后，模型会有微弱的速度提升，但性能全面下降；此外，去除 MHA 和 FFN 层掩码会导致高稀疏度模型的速度显著下降。

其次，作者探究了不同目标函数对模型表现的作用，如下表所示，作者发现去掉蒸馏目标函数会导致模型性能全面大幅下降，这表明加入蒸馏的目标函数是必要的。同时，作者比较了动态建立层映射方法和提前固定好层映射方法的表现。对比发现，固定层映射的表现不如 CoFi 的动态方法。

4.4 Structures of Pruned Models

作者还探究了最终模型的结构，发现了以下几种现象：

（1）FFN 层维度有了显著的减小，且剪枝比例远高于 MHA 层的剪枝结果。例如在稀疏度模型中，FFN 层的剪枝比例高达，而 MHA 层的比例仅为，说明 FFN 层较 MHA 层更加高度冗余。

（2）CoFi 更倾向于对上层的单元进行剪枝，例如上层的 MHA 层头数剩余更少。

（3）综合各个数据集上的 CoFi 剪枝模型，作者发现 MHA 层被保留的更多（如下如所示），这可能说明 MHA 层更有助于解决下游任务。

Conclusion

本文提出了 CoFi，这是一种结合了粗细粒度级别单元的结构化剪枝方法。根据实验表明，CoFi 可以在不使用大量 unlabeled data 的前提下，得到与标准蒸馏模型完全不同结构的压缩模型，能在超过倍的加速下取得有竞争力的表现。

作者认为，尽管由于上游剪枝选择的复杂性，本文并没有讨论将 CoFi 用于任务无关模型的结构化剪枝。但未来的研究可以对该方法进行进一步探讨，因为与一般蒸馏相比，从大型预训练模型中剪枝可能会产生更少的计算量，并得到更灵活的模型结构。

参考文献

[1] C Louizos, M Welling, and DP Kingma. 2018. Learning sparse neural networks through l0 regularization. In International Conference on Learning Representations (ICLR).

[2] Ziheng Wang, Jeremy Wohlwend, and Tao Lei. 2020b. Structured pruning of large language models. In Empirical Methods in Natural Language Processing (EMNLP), pages 6151–6162.

[3] Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019. Patient knowledge distillation for BERT model compression. In Empirical Methods in Natural Language Processing (EMNLP), pages 4314–4323.

更多阅读