高精度压缩Transformer，NNI剪枝一站式指南

无论在学术界还是产业界，今年人工智能大模型都是爆款话题。但面对这些动不动就数十亿级别参数的模型，使用传统方法微调，宛如水中捞月、海底捞针。作为微软亚洲研究院为科研人员和算法工程师量身定制的一站式 AutoML（自动机器学习）工具， NNI（Neural Network Intelligence）在过去的三年间不断迭代更新，加强了对各种分布式训练环境的支持，成为了最热门的 AutoML 开源项目之一。

近日，微软亚洲研究院对 NNI 进行了更新。在最新的版本中，NNI 集成了大量前沿的剪枝算法，如 TaylorFO Weight、Movement 等。基于现有的经典预训练模型，研究员们通过大量实验，发现了既能降低模型参数量和计算量，又能保持模型较高精度的剪枝步骤与算法组合，获得超越 SOTA 的模型剪枝效果。

今天我们就以 Transformer 系列的预训练模型和数据集 GLUE-MNLI 为例，为大家介绍一下 NNI 的 pruner 剪枝流程和使用的剪枝算法组合。

剪枝流程

在正式介绍剪枝流程前，我们需要先了解什么是 pruner，mask 和 SpeedUp。

pruner：使用具体的剪枝算法实例化的剪枝器。

mask：在剪枝过程中，pruner 会生成一个和目标子模块大小相同的 mask（全1）矩阵，并在 mask 矩阵中将目标子模块中需要剪掉的部分的对应位置置为0。最后通过将目标子模块和对应的 mask 矩阵相乘，即可得到模拟剪枝后的模型效果。

SpeedUp：从上述描述可以看出，在剪枝过程中，实际上只是将需要剪枝的部分用0进行了替换，因此使用 SpeedUp 模块是修剪上述目标子模块中需要剪掉的参数，而不是用0替代，从而实现真正意义上的减少参数量。

在使用 NNI Compression 模块中的 pruner 进行剪枝操作时，用户只需完成数据/模型等的准备、pruner 的构建，以及模型剪枝和再训练，即可为模型构建一个剪枝的 pipeline。

以 Transformer 系列的预训练模型为例，其剪枝流程共包含4步：首先准备数据/模型等，接着针对多头自注意力机制（Multi-head Attention）、嵌入层（embedding）和前馈神经网络（FFN）分别剪枝和再训练模型。

图1：Transformer 系列模型的剪枝流程示意图

1. 准备数据/模型等

在正式构建剪枝过程之前，用户需要加载预训练模型，对数据预处理并创建相应的 dataloader，同时设计相应的训练/评估函数，以用于后期对模型的训练和评估。其流程如图2所示，共包含5步：

图2：数据/模型准备过程的流程示意图

具体来说，首先需要从 Transformers 库中加载预训练模型，然后对数据 GLUE-MNLI 进行处理，并得到相应的 dataloader。随后，针对模型和数据集 GLUE-MNLI，构建相应的训练/评估函数。最后将模型在 GLUE-MNLI 数据集上进行微调。

完成以上步骤就相当于完成了数据/模型等的准备工作，可以得到预训练模型在 MNLI 数据集上微调后的模型。考虑到 Transformer 系列预训练模型的模型参数中的大头为嵌入层，且编码层/解码层中包含了多头自注意力机制和前馈神经网络。因此，在之后的步骤中需要分别对多头自注意力机制、嵌入层和前馈神经网络剪枝，并引入动态蒸馏机制对剪枝后的模型再训练。

2. 多头自注意力机制的剪枝和基于动态蒸馏机制的模型再训练

多头自注意力模块的剪枝和模型再训练分为3步，如图3所示：首先要构建 pruner，接着对多头自注意力模块进行剪枝，最后使用动态蒸馏机制再训练模型。

图3：多头自注意力机制的剪枝和再训练流程示意图

在进行剪枝前，用户需要选定一个剪枝算法并实例化相应的 pruner。所有的剪枝算法均需向模型中传入 config_list 参数，因为其定义了需要剪枝的运算名、运算类别及稀疏度等。具体到 Movement 剪枝算法，还需要设置其他的一些参数，如：evaluator 参数，用于训练感知的模型压缩过程；movement_mode 参数，共有“soft“和”hard“两种模式，若为”soft”，则难以精确地控制模型剪枝后的稀疏度，但是可以得到性能更好的模型。参数 regular_scale 用于控制剪枝的稀疏度，regular_scale 越大，模型剪枝后的稀疏度越高。更多其他参数可参阅https://nni.readthedocs.io/zh/stable/reference/compression/pruner.html#movement-pruner

接下来，要使用构造的剪枝算法实例 pruner 对多头自注意力模块进行剪枝。用户只需调用 pruner.compress() 即可执行对模型的剪枝过程，并得到剪枝后的模型和 attention_mask。其中 attention_mask 给出了需要剪枝的子模块的参数剪枝范围，0代表该位置被剪掉，1代表该位置被保留。

NNI 的 SpeedUp 模块可以将被 mask 住的参数和计算从模型中删除，具体的删除逻辑如图4所示，以 Query Linear 层的 weight（记作Q）为例，其维度为[768,768]，那么 Q 的 weight 的 mask 矩阵维度也为[768, 768]，将其记作 mask。首先将该 mask 矩阵的维度进行变换，第一维是多头数目8，其余的则是第二维，将变换后的 mask 矩阵记作 reshaped mask 矩阵。接着，对 reshaped mask 矩阵在第二维度上求和，并判断求和后的值是否为0，此时的 mask 矩阵维度变为[8]，每个位置对应着一个多头。对于变换后的 mask 矩阵，若位置 i 的值为0，则代表在 Q 中的第 i 个多头需要被剪掉。在图中，位置0、3、7的值均为0，因此，在Q中的第0、3、7个多头需要被剪掉。最后，将[0,3,7]作为参数传入 prune_heads 函数中，对 Q 进行修剪。修剪后，Q 的维度为[576,768]。对 SpeedUp 更加全面的介绍可以参考发表于 OSDI 2022 的论文 SparTA。在即将发布的 NNI 3.0 中 SpeedUp 会对更多模型提供更加完善的支持。

图4：利用 prune_heads 函数修剪自注意力模块的过程示意图

在对多头自注意力模块剪枝后，以微调后的模型作为教师模型，以剪枝后的模型作为学生模型，然后借鉴 CoFi 中的动态蒸馏机制 [1] 对模型进行再训练，就可以得到新的模型。这里的动态蒸馏机制，是指教师模型的层和学生模型的层之间不是一个静态对应关系，每次蒸馏教师都可以选择从自身的高层动态蒸馏信息到学生模型低层中的一层里。

3. 嵌入层和前馈神经网络的剪枝，以及基于动态蒸馏机制的模型再训练

嵌入层和前馈神经网络的剪枝过程与多头自注意力模块的剪枝过程类似。此处使用 Taylor 剪枝算法（https://nni.readthedocs.io/zh/stable/reference/compression/pruner.html#taylor-fo-weight-pruner ）对嵌入层和前馈神经网络进行剪枝。同样地，研究员们定义了 config_list、evaluator 参数及 taylor_pruner_steps 参数。由于嵌入层的维度与后续模型中的维度具有相关性。因此，基于上述参数，在嵌入层的剪枝过程中研究员们将剪枝模式 mode 设置为了“dependency-aware”模式，并传入模型的输入 dummy_input，以帮助 pruner 捕捉和嵌入层维度具有依赖关系的子模型。

接下来，使用分别构造的 pruner 对前馈神经网络和嵌入层进行剪枝。和多头自注意力模块的剪枝不同的是，此处使用了迭代式剪枝法，即在模型基于动态蒸馏的再训练过程中，每2000步分别使用 pruner 对前馈神经网络和嵌入层剪枝一次，其中，前馈神经网络共剪枝19/24次，嵌入层共剪枝3次。每次剪枝后，使用 ModelSpeedUp 对前馈神经网络层进行剪枝，以实现真正意义上的修剪参数，而不是将需要修剪的参数用0替换。

实验结果

通过调整 regular_scale 参数的值和前馈神经网络的剪枝次数，研究员们得到了具有不同稀疏度和性能的模型。该过程使用了1张 A100 进行实验，并设置 batch_size 为32。

图5：实验结果

从上图实验结果可以看出：

随着 regular_scale 的增加，模型总的稀疏度有所增加。当 regular_scale 大于等于10时，模型总的稀疏度超过了69%，性能损失超过1%。

随着前馈神经网络剪枝次数的增加，模型总的稀疏度有所增加，同时模型的性能有所下降，且随着模型总稀疏度的增加，模型的性能下降程度逐渐增大。

对嵌入层剪枝3次，能够将模型的维度从768减小至561，在一定程度上提升了模型总的稀疏度。

实验结果与平台对比

进一步分析实验结果可以发现，使用 NNI 对 BERT 在 MNLI 数据集上剪枝后的性能好于 nn pruning 框架（图6(a)），且当模型总的稀疏度低于65%时，NNI 和 CoFi 对 BERT 在 MNLI 数据集上剪枝的性能差距较小，当模型总的稀疏度大于65%时，使用 NNI 对 BERT 在 MNLI 数据集上剪枝后的性能好于 CoFi。图6(b)和图6(c)分别展示了 NNI 在 T5 和 ViT 模型上的剪枝性能。从图中可以看出，当模型相应部分的稀疏度超过了75%后，模型性能下降约为3%，当模型相应部分的稀疏度低于50%时，模型性能下降较少。