​NeurIPS 2022 | 仅需3分钟！开源Transformer快速训练后剪枝框架来了

©PaperWeekly 原创 · 作者 | An.

单位 | 中科院自动化所

研究方向 | 计算机视觉、模型压缩

引言

近年来，Transformer 已成为自然语言处理的标准模型结构，并在计算机视觉、语音识别等领域也取得了许多成果。然而，存储空间占用大、推理延迟高等问题阻碍了其实际应用。因此，针对 Transformer 的模型压缩方法得到了广泛的研究，结构化剪枝就是其中非常重要的一类方法。

过往的 Transformer 结构化剪枝方法虽然可以对模型参数量和计算量进行压缩，但由于下面的三类原因，在实践中通常难以应用：

1. 重新训练和/或联合学习剪枝配置方案会使训练时间增加多达 10 倍，显著增加了计算开销；

2. 复杂的剪枝框架包含了许多辅助模块和超参数，这些都为工程实现和调试带来了很大的困难；

3. 过去针对 Transformer 结构化剪枝方法往往不能适应用户的差异化需求，它们要么依赖正则化超参数模糊控制模型的稀疏性，要么使用与用户设置无关的固定模型结构，这些都可能导致次优剪枝模型。

为解决上述问题，本文提出了一个快速的 Transformer 训练后剪枝框架。训练后剪枝（Post-training pruning）是训练后压缩（Post-training compression）的一种，过去的研究主要集中在训练后量化（Post-training quantization）。训练后剪枝是指在不进行任何再训练的情况下对模型进行剪枝，从而避免额外的训练成本和用户干预。

本文提供了一个开箱即用的训练后剪枝框架，在单个 GPU 上仅需不到 3 分钟即可实现端到端的结构化剪枝，比过往的 Transformer 剪枝方法快了 100 多倍。

论文标题：

A Fast Post-Training Pruning Framework for Transformers

论文链接：

https://arxiv.org/pdf/2204.09656.pdf

代码链接：

https://github.com/WoosukKwon/retraining-free-pruning

该框架将 Transformer 模型、样本数据集和 FLOPs / 延迟约束作为输入，不进行任何再训练的情况下最终输出一个用于部署的剪枝模型。如图 1 所示，该框架有三个阶段：

1. 基于费雪信息矩阵（Fisher information matrix）对角化的掩码搜索（Mask Search），以确定各层修剪比例；

2. 基于层内相互作用的掩码重排列（Mask Rearrangement），以确定各层的修剪位置，找到最优二值掩码；

3. 基于线性最小二乘的掩码微调（Mask Tun