NeurIPS 2022 | 最优脑压缩，训练后剪枝又一力作！

科技

2023-04-04 05:04

@Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

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

论文标题：

Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning

论文链接：

https://arxiv.org/pdf/2208.11580.pdf

代码链接：

https://github.com/IST-DASLab/OBC

本文提出了一种新的训练后压缩框架 Optimal Brain Compression（OBC），该框架在统一范式下涵盖了权重剪枝和量化。OBC 在无需任何再训练的情况下，只需少量的校准输入数据，即可获得空间和时间效率大大提升且模型性能良好的压缩模型。该方法基于最优脑手术（OBS）[1] 框架进行扩展，将其推广到了包括 CNN、Transformer 在内的现代 DNN 模型中。

实验证明，OBC 在 GPU 支持的复合压缩中，实现了精度仅下降 2% 的情况下，理论计算量减少 12 倍，以及在 CPU 稀疏感知运行时，精度仅下降 1% 的情况下，实际运行加速 4 倍。这些结果首次表明，训练后压缩可以获得与完全再训练压缩相近的性能表现。

问题定义

分层压缩问题：分层压缩是训练后压缩的标准方法，它将压缩任务分解为分层子问题，根据校准数据给出层输入和输出的子样本，为每一层确定压缩权重的近似重建。数学上，我们可以定义层的函数为，该层的输入为，权重为，分层压缩的目标是在约束条件 下找到权重的近似重建 以最小化损失 ，具体公式如下：

层输入的期望通常是通过统计个输入样本的 mini-batch 的平均值来近似估计的。这种低数据设置也是分层压缩的一大优势。与过往的分层压缩工作 [2,3] 一致，本文也只对线性层和卷积层进行压缩，这些层均使用平方误差来衡量重建误差。详细的推导可以参阅 [3]。式（1）更具体的表示如下：

最佳脑手术（OBS）框架：OBS 框架 [1] 针对密集神经网络的精确修剪问题，从给定点处的泰勒近似出发，给出了最优去除权重选择的显式公式，以及补偿剩余权重的最优更新。具体来说，设为给定模型的海森矩阵，引起最小损失增加的修剪和剩余权重的更新量如下：

OBS 分层剪枝：有了上述基础，我们现在可以利用分层修剪问题实例化本文的框架。首先，式（2）的损失是二次的，完全满足 OBS 框架的假设，这意味着式（3）对于分层压缩问题是精确的，不引入额外的误差。

因此，迭代 OBS 框架能以一次去除一个权重并对其他权重进行补偿的方式，为分层压缩问题提供一套精确的贪心解决方案。虽然这种贪心方法并不能保证收敛到全局最优值，但实验证明，这种方法对于处理太大而无法通过精确方法处理的问题实例非常有效。

稀疏性的最优贪心求解器

迭代 OBS 的贪心解决方案存在一个巨大的挑战，即使用式（3）修剪单个权重的计算量很高。海森矩阵是一个的矩阵，其中，该矩阵需要在步中的每一处更新和求逆，计算复杂度为。显然，的总体运行时间对于修剪如今规模的神经网络来说效率太低。

因此，本文的核心创新是在不引入额外误差的情况下，将此过程的整体成本降低至 的时间和 的内存。

2.1 ExactOBS 算法

我们首先将式（2）改写成行表示形式

，从中可以看出，删除权重只影响对应行输出。因此，不同行之间没有海森矩阵的交互，因此仅需使用对应每行的单个的海森矩阵就足够了。这一观察降低了大量的计算复杂度。为了进一步降低时间和空间复杂度，作者还做了下面两个改进。

第1步 - 处理单行：这一改进的关键思想在于避免在每一步中对重新的统计和的求逆。前者很容易，因为权重的变化不会对带来改变，我们只需在修剪前对各层的海森矩阵进行统计，在后续修剪时根据每行的掩码进行选择即可。

然而，这一技巧并不适用于求逆操作，因为。本文提出了一种基于高斯消元法的等价求解方法，来根据当前海森矩阵的逆快速计算下一步迭代时新的海森矩阵逆。当我们去除某一个权重时，海森矩阵对应行列的元素也将删除，基于这一事实，在更新逆时我们可以直接使用高斯消元法将对于行列的元素去除，具体公式如下：

由此可以得到单行修剪的算法伪代码（见图 1），其中海森矩阵逆的更新的时间复杂度为。

▲ 图1. 单行 OBS 修剪的伪代码

第2步 - 联合考虑所有行：为了在每一步中选择具有最小整体修剪损失的权重（即每行最小损失的最小值），需要将 OBS 框架应用于完整的权重矩阵，而不仅仅是独立应用于每行。

图 1 的算法已经解决了计算复杂度的问题，但这种多行联合求解带来了过于昂贵的内存占用（每行都需要存储）。作者巧妙地发现了一个性质，即每行内的修剪顺序是固定的，最终剪去多少仅取决于修剪的参数总数。这为降低存储复杂度提供了可能。

这一性质让我们可以独立处理每一行，按照顺序修剪所有权重，并记录损失变化量。通过这种预处理，我们只需要的额外内存，即可按照预先记录的在完整矩阵中选择全局掩码。

如果有足够的 CPU 内存可用的话，我们还可以保留每行的修剪轨迹，即每个单独修剪步骤之后的完整权重矩阵，这是一种用空间换时间的权衡，在额外的 CPU 内存占用下，可以节省第二次调用图 1 的算法。图 2 可视化了刚刚讨论的两个选项（Less compute 和 Less memory）。

▲ 图2. 利用逐行结果的高效全局 OBS 算法

至此，我们已经介绍完了本文的核心创新——在不引入额外误差的情况下，将此过程的整体成本降低至的时间和的内存。为了让读者更清楚地理解整个修剪过程，我再阐述一次整体的算法流程：

（a）将整个模型的剪枝问题分解为各层（线性层或卷积层）的修剪和权重重建子问题；

（b）统计每层的海森矩阵，并初始化海森矩阵的逆；

（c）对于每层的权重，分行利用式（3）计算每个权重剪去带来的误差，依次记录最小的索引 和对应的误差 ，剪去对应权重，利用式（3）对剩余权重进行更新，并利用式（4）更新海森矩阵的逆和行掩码【这一步并不会真正更新模型权重，只是为了记录各行修剪的顺序和每次修剪的误差】；

（d）根据上一步各行的记录结果，选择全局的最小索引，进行完整权重矩阵的修剪和更新。

2.2 扩展：N:M 和块稀疏

N:M 稀疏性：该方法可以很容易地扩展到各种形式的半结构化稀疏性。例如 N:M 稀疏模式 [4]，它在 M 个连续权重中恰好选保留 N 个非零值。只需简单的调整即可实现：不在选择损失变化最小的权重，而是选择在具有 < N 个修剪权重的块中损失变化最小的权重。当某行达到的稀疏比时就终止修剪。因此也不再需要前面介绍的全局掩码选择步骤。

块稀疏性：块修剪是另一种实践中常用的修剪模式，特别是在 CPU 加速的情况下。它要求仅在大小为的连续块中出现零，本文的方法可以扩展到修剪小的权重连续组，具体公式如下：

最佳大脑量化器（OBQ）

上一节的算法可以通过一种有效且准确的方式扩展到量化，作者将其称为最优脑量化器（OBQ）。

将 OBS 映射到量化投影，是量化的权重舍入，更新后。我们可以根据式（3）推导出类似的量化更新公式：

从 OBS 的视角看量化时，我们可以想象，我们先找到“最简单”量化的权重，调整所有剩余未量化权重来补偿精度损失，以此类推。我们只需要在式（6）的基础上使用图 1 提到的算法，即可迭代式的量化给定层的权重。

实验

1. 非结构化稀疏性：表 1 展示了不同 FLOP 的非结构化剪枝结果。在 ResNet、YOLOv5 和 BERT 上，本文提出的 ExactOBS 算法的表现通常是最好的，尤其是在大稀疏率的情况下。

▲ 表1. 不同 FLOP 目标下的非结构化剪枝方法对比

2. N:M 稀疏性：如表 2 所示，本文提出的方法在多个 ResNet 变体中均取得了最优的效果。

▲ 表2. 半结构化剪枝的性能对比

3. 量化：表 3 对比了 OBQ 与其他最先进的训练后量化方法的性能结果。从中可以看出，尽管 OBQ 独立优化了每个层，但在性能表现上与其他联合优化的方法相比非常相近。

▲ 表3. 训练后量化方法的性能对比

4. BOP 约束混合 GPU 压缩：图 3 展示了OBC 将剪枝和量化结合起来的性能效果。在相对性能下降约 2.5% 的代价下，ResNets 的 BOP（比特数乘以 FLOPs）降低了 12~14 倍。

▲ 图3. 混合量化+2:4剪枝的 BOP 与准确率变化关系（BOP 指比特数乘以 FLOPs）

参考文献

[1] Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In IEEE International Conference on Neural Networks, 1993.

[2] Peisong Wang, Qiang Chen, Xiangyu He, and Jian Cheng. Towards accurate post-training network quantization via bit-split and stitching. In International Conference on Machine Learning (ICML), 2020.

[3] Markus Nagel, Rana Ali Amjad, Mart Van Baalen, Christos Louizos, and Tijmen Blankevoort. Up or down? Adaptive rounding for post-training quantization. In International Conference on Machine Learning (ICML), 2020.

[4] Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. Learning N:M fine-grained structured sparse neural networks from scratch. InInternational Conference on Learning Representations (ICLR), 2021.

更多阅读