[Paper Review] 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 Applications36 citations

TL;DR

Introduces the Optimal Brain Compressor (OBC), a unified, exact OBS-based framework for one-shot post-training pruning and quantization that improves accuracy under compression and supports joint pruning-quantization.

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.

Motivation & Objective

• Address the challenge of compressing trained DNNs in a one-shot post-training setting using limited calibration data.
• Develop a unified framework that handles both pruning and quantization efficiently and accurately.
• Provide exact OBS-based methods with practical time and space efficiency for modern networks.
• Enable compounding pruning and quantization to achieve high speedups with minimal accuracy loss.

Proposed method

• Formulate layer-wise compression as minimizing the layer output change under a compression constraint.
• Adapt the Optimal Brain Surgeon (OBS) framework to the squared-error layer-wise loss to obtain an exact greedy pruning algorithm.
• Develop ExactOBS to prune one weight at a time with O(d_row * d_col^2) time and O(d_col^2) memory by row-wise Hessian handling and row/column removal updates.
• Extend OBS to weight quantization, creating the Optimal Brain Quantizer (OBQ) that selects weights to quantize based on loss impact and applies closed-form updates to remaining weights.
• Unify pruning and quantization into the Optimal Brain Compressor (OBC) and implement practical extensions for N:M and block sparsity, with optional group updates.
• Provide efficient, exact implementations and a public repository for reproducibility.

Experimental results

Research questions

• RQ1 Can a single, exact OBS-based framework be effectively applied in a post-training setting for both pruning and quantization?
• RQ2 Does a layer-wise, one-shot compression approach yield competitive accuracy under practical FLOP/latency constraints without retraining?
• RQ3 Can pruning and quantization be composed to achieve greater speedups with minimal accuracy loss in both GPU and CPU environments?
• RQ4 How can second-order information be efficiently computed and updated to enable exact pruning and quantization at DNN scale?

Key findings

• The ExactOBS algorithm achieves an exact greedy pruning solution with significantly reduced complexity to prune layers at modern DNN scale on a single GPU.
• The OBS-based approach can be extended to iterative one-at-a-time weight quantization, producing the Optimal Brain Quantizer (OBQ) and unifying with pruning into the Optimal Brain Compressor (OBC).
• OBC yields state-of-the-art accuracy versus compression trade-offs for post-training pruning and quantization across image classification, object detection, and language modeling tasks.
• Compound compression (pruning plus quantization) delivers substantial speedups, e.g., a 12x reduction in theoretical operations with only 2% accuracy loss in GPU scenarios and 4x runtime speedup with 1% loss on CPU scenarios.
• Layer-wise compression, followed by stitching, with calibration data, can closely match or exceed results of more expensive global optimization methods in post-training setups.
• The framework supports non-uniform compression and practical sparsity patterns such as N:M, block sparsity, and can adapt to per-layer constraints via a DP-based layer-wise solver.