[论文解读] Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning

Elias Frantar, Singh, Sidak Pal+1 more|arXiv (Cornell University)|Aug 24, 2022

Medical Imaging Techniques and Applications被引用 36

一句话总结

提出 Optimal Brain Compressor (OBC)，一个统一的、基于 OBS 的精确框架，用于一次性后训练剪枝与量化，在压缩下提升准确性并支持联合剪枝-量化。

ABSTRACT

We consider the problem of model compression for deep neural networks (DNNs) in the challenging one-shot/post-training setting, in which we are given an accurate trained model, and must compress it without any retraining, based only on a small amount of calibration input data. This problem has become popular in view of the emerging software and hardware support for executing models compressed via pruning and/or quantization with speedup, and well-performing solutions have been proposed independently for both compression approaches. In this paper, we introduce a new compression framework which covers both weight pruning and quantization in a unified setting, is time- and space-efficient, and considerably improves upon the practical performance of existing post-training methods. At the technical level, our approach is based on an exact and efficient realization of the classical Optimal Brain Surgeon (OBS) framework of [LeCun, Denker, and Solla, 1990] extended to also cover weight quantization at the scale of modern DNNs. From the practical perspective, our experimental results show that it can improve significantly upon the compression-accuracy trade-offs of existing post-training methods, and that it can enable the accurate compound application of both pruning and quantization in a post-training setting.

研究动机与目标

• 在有限的校准数据下，解决以一次性后训练设置对训练好的 DNN 进行压缩的挑战。
• 开发一个统一框架，能够高效且准确地处理剪枝与量化。
• 为现代网络提供在实际时间和空间效率下的精确 OBS 基方法。
• 使剪枝与量化的复合实现成为可能，以在最小的准确性损失下获得高加速。

提出的方法

• 将逐层压缩表述为在压缩约束下最小化层输出变化。
• 将 Optimal Brain Surgeon (OBS) 框架适应到平方误差的逐层损失，以获得精确的贪婪剪枝算法。
• 开发 ExactOBS，通过逐行处理 Hessian 及行/列移除更新，以 O(d_row * d_col^2) 时间和 O(d_col^2) 内存逐次剪枝一个权重。
• 将 OBS 扩展到权重量化，创建 Optimal Brain Quantizer (OBQ)，基于损失影响选择需要量化的权重，并对剩余权重应用闭式解更新。
• 将剪枝与量化统一为 Optimal Brain Compressor (OBC)，并实现对 N:M 和块稀疏的实际扩展，提供可选的分组更新。
• 提供高效、精确的实现以及用于可重复性的公开代码库。

实验结果

研究问题

• RQ1 一个单一的、精确的 OBS 基框架是否可以在后训练设置中同时有效应用于剪枝和量化？
• RQ2 逐层的一次性压缩方法是否在不重新训练的前提下，在实际的 FLOP/延迟约束下仍能获得具有竞争力的准确性？
• RQ3 剪枝与量化是否可以组合实现更高的加速，在 GPU 和 CPU 环境下尽量降低准确性损失？
• RQ4 如何高效地计算和更新二阶信息，以实现 DNN 规模的精确剪枝和量化？

主要发现

• ExactOBS 算法在单个 GPU 上，以显著降低的复杂度实现对现代 DNN 规模的逐层剪枝的精确贪婪解。
• 基于 OBS 的方法可以扩展到逐步一次一个权重量化，产生 Optimal Brain Quantizer (OBQ)，并与剪枝统一为 Optimal Brain Compressor (OBC)。
• OBC 在图像分类、目标检测和语言建模任务中，实现了在压缩权衡上的一流准确性。
• 复合压缩（剪枝加量化）带来显著加速，例如在 GPU 情况下理论运算减少 12x，准确度仅损失 2%，在 CPU 情况下运行时间加速 4x，损失 1%。
• 分层压缩，随后结合，使用校准数据，在后训练设置中可以接近甚至超过更昂贵的全局优化方法的结果。
• 该框架支持非均匀压缩和实用的稀疏模式，如 N:M、块稀疏，并且可以通过基于 DP 的逐层求解器适应每一层的约束。