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

2023-02-18 15:02

©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 Tuning），以找到最优实值掩码，尽可能恢复模型性能。

▲ 图1. 本文剪枝框架示意图

方法

2.1 问题分析

问题定义：本文的剪枝问题可以定义为在稀疏约束（计算量或推理速度）下寻找最优掩码。该剪枝问题可以描述为掩码上的约束优化问题。

在不进行再训练的基础上，该问题可分解为三个子问题：1）确定各层注意力头和中间层神经元的修剪比例；2）确定各层注意力头和中间层神经元的修剪位置，即获得最优二值掩码；3）确定最优实值掩码，尽可能恢复模型性能。三个子问题分别由 2.2, 2.3, 2.4 三个小节依次解决。

问题简化：由于式（1）中的通常是掩码的 -范数的函数，它是不可微的，因此需要引入几个假设和近似来简化问题。

作者使用了初始掩码附近的二阶泰勒展开来近似损失函数：

当模型收敛到局部最小值时，式（2）的梯度项接近 0，由于是常数，因此你优化目标可以改写为：

由于准确计算海森矩阵是不可行的，作者利用掩码变量的（经验的）费雪信息矩阵来近似海森矩阵：

2.2 掩码搜索

费雪信息矩阵的对角近似：用完整的费雪信息矩阵求解式（3）的优化目标仍存在困难。因此，作者先忽略不同掩码变量的相互关系，假设是对角矩阵。掩码搜索的目标是确定各层的修剪比例，掩码变量限制为 0 或 1，因此该阶段的目标函数可表示为：

求解 FLOPs 约束问题：头/神经元的数量和 FLOPs 成线性关系，因此可以采用网格搜索和贪心算法进行求解，搜索算法代码如下：

for num_heads in range(1, num_hidden_layers * num_attention_heads + 1): # 遍历所有的可能的头修剪数量
        heads_mac = mac_per_head(num_patches, hidden_size, attention_head_size) * num_heads # 计算修剪头减少的 FLOPs
        neurons_mac = max_mac - heads_mac   # 按照 FLOPs 约束, 计算神经元需要减少的 FLOPs
        num_neurons = int(neurons_mac / mac_per_neuron(num_patches, hidden_size))   # 计算最少需要修剪的神经元数量
        num_neurons = max(num_neurons, 0)

        # 贪心策略修剪重要性得分最小的元素
        total_importance = sorted_head_importance[:num_heads].sum() + sorted_neuron_importance[:num_neurons].sum()
        if total_importance > max_importance:   # 记录全局最优的头/神经元修剪比例
            max_importance = total_importance
            head_indicies = sorted_head_indicies[:num_heads]
            neuron_indicies = sorted_neuron_indicies[:num_neurons]

求解 Latency 约束问题：基本思想与求解 FLOPs 约束问题相同，不同点在于 Latency（推理延迟）与头/神经元不是线性关系。当同一层剩余的头/神经元足够多时，并行核等硬件资源可以得到充分利用，推理延迟大致与头/神经元数量成正比，然而，如果剩余头/神经元数量过少时，硬件未得到充分利用，基础开销主导着推理延迟。在提供推理延迟查找表后，分别处理的常数部分和线性部分。

▲ 图2. 推理延迟和头/神经元数量的大致关系

2.3 掩码重排列

费雪信息矩阵的块对角近似：掩码搜索阶段忽略了不同掩码变量间的相互作用，这虽然简化了问题，但也导致了次优解。例如，同一层的两个注意力头有类似的作用，只修剪其中一个对性能的影响不大，但当两者都被修剪时，模型的性能显著降低。因此作者在这一阶段考虑了每个 MHA 层或 FFN 层内部的相互作用，以找到更好的修剪位置。

掩码重排列将固定掩码搜索阶段得到的各层头/神经元修剪比例，逐层利用贪心搜索来重新排列每一层的二值掩码变量，以找到最优的修剪位置，获得最优二值掩码变量，具体代码如下：

# Greedy search
masked_indicies = indicies[:num_pruned] # 单一 MHA/FFN 层已修剪的元素索引
for index in indicies[num_pruned:]: # 遍历未修剪的元素
    masked_indicies.append(index)   # 将其加入已修剪元素列表
    grad_vectors = grads[masked_indicies]
    grad_sum = grad_vectors.sum(dim=0)

    complement = grad_sum - grad_vectors
    grad_sum_length = complement.pow(2).sum(dim=1)

    removed = grad_sum_length.argmin()  # 选出最重要的元素移出已修剪列表
    del masked_indicies[removed]

2.4 掩码微调

在这一阶段，前两个阶段获得的掩码变量中的非零变量将被微调至任何实值，以使得修剪后的模型尽可能恢复性能。

基于线性最小二乘的逐层重构：类似 [1]，求解最优实值掩码可视为从第一层到最后一层，用修剪模型的剩余头/神经元重构原始模型的输出激活，具体形式如下：

式（6）可以简化为

的线性最小二乘问题。由于矩阵的尺寸很大，直接求解会导致数值不稳定的结果。为解决这一问题，作者采用了带有正则化超参数（即阻尼）的 CuPy [2] 中的 LSMR 求解器。

实验

1. FLOPs 和准确率比较：如图 3 所示，在没有任何再训练和大幅降低修剪成本的情况下，本文的方法取得了与之前方法相当或更好的结果。

▲ 图3. 与过往方法的压缩性能对比

2. 修剪成本比较：如表 1 所示，本文的方法在修剪成本上降低了 2~3 个数量级（1/500~1/3300）。

▲ 表1. 与过往方法的修剪成本对比

3. 消融实验：如表 2 所示，本文提出的三个阶段对于性能的恢复都有贡献，掩码微调尤为关键。图 4 进一步研究了掩码搜索和重排列的重要性，从图中可知，本文提出的掩码搜索和重排列是获得最优二值掩码的必要条件，掩码微调需要在二值掩码保持较高精度时才有效。

▲ 表2. 三阶段消融实验

▲ 图4. 与其他二值掩码选择策略的比较，以分析掩码搜索和掩码重排列的作用

参考文献

[1] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. InProceedings of the IEEE international conference on computer vision, pages 1389–1397, 2017.

[2] ROYUD Nishino and Shohei Hido Crissman Loomis. Cupy: A numpy-compatible library for nvidia gpu calculations. 31st confernce on neural information processing systems, 151, 2017.

更多阅读