PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Elias Frantar, Sidak Pal Singh, Dan Alistarh

arXiv:2208.11580v2cs.LG

TL;DR

The paper addresses one-shot post-training compression of trained DNNs using limited calibration data and no retraining. It develops an efficient OBS-based framework unifying pruning and quantization, and reports improved trade-offs plus accurate compound compression. The results suggest post-training compression can approach the accuracy of more expensive retraining methods.

  • Problem

    One-shot post-training compression must produce an accurate compressed DNN from limited calibration data without retraining and with limited computational costs.

  • Method

    The paper efficiently realizes OBS for exact greedy layer-wise pruning and extends it to iterative quantization through the unified OBC framework.

  • Results

    Post-training OBC improves compression-accuracy trade-offs and achieves 12× fewer theoretical operations with a 2% accuracy drop and 4× actual CPU speedup with a 1% accuracy drop.

  • Takeaways & Limitations

    The results suggest that post-training compression may reach accuracies comparable to much more expensive retraining methods while supporting compound pruning and quantization.

  • Takeaways & Limitations

    The layer-wise formulation relies on low-data calibration and squared-loss approximations, while direct OBS storage and computation can be infeasible for large layers.

Abstract

from arXiv · show

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.

1 Introduction

The paper targets one-shot post-training compression using limited calibration data, unifying pruning and quantization through an efficient OBS-based framework. Its experiments report improved compression-accuracy trade-offs and accurate compound compression without retraining.

  • Problem: Post-training compression produces an accurate compressed model from a trained model and limited calibration data in one step, without retraining.The setting is motivated by practical inference scenarios with limited computational costs.
  • Framework: The framework unifies weight pruning and quantization in a single post-training compression approach.It addresses a practical question raised by joint hardware support for sparse and quantized formats.
  • Framework: OBS specialization yields an exact greedy pruning solution in theory, while efficient algorithms make it practical at modern DNN scale.The direct implementation costs Θ(d^4), whereas the proposed algorithms reduce the cost without approximations.
  • Framework: OBQ quantizes weights iteratively according to loss impact and updates remaining unquantized weights in closed form.OBQ extends the OBS approach to quantization within the unified OBC framework.
  • Results: 12× fewer theoretical operations incurred a 2% accuracy drop, while 4× actual CPU runtime speedup incurred only a 1% accuracy drop.These results concern compound pruning and quantization in GPU-supported and CPU sparsity-aware settings, respectively.

2 Related Work

Prior work developed separate layer-wise approaches for post-training quantization and pruning, while this paper seeks independent, unified compression across both types. The related methods trade flexibility against accuracy and deployment considerations.

  • Pruning: OBS-based pruning methods use second-order information, but recent DNN-scale variants require gradual pruning and substantial retraining.These requirements prevent their direct application to the post-training setting.
  • Quantization: Post-training quantization methods commonly compress layers sequentially, allowing later layers to compensate for earlier errors.BitSplit, AdaRound, AdaQuant, and BRECQ represent this line of work.
  • Quantization: Sequential layer processing improves performance but reduces flexibility because changing one layer’s compression parameters may require repeating the entire process.The paper instead targets independent layer compression that can be stitched together.
  • Pruning: AdaPrune applies layer-wise calibration-output reconstruction to post-training pruning, including GPU-supported N:M sparsity patterns.Global AdaPrune adds a more expensive global optimization step that can further improve accuracy.
  • Non-Uniform Compression: Non-uniform compression methods either search directly for layer-wise policies or use solvers under resource constraints such as latency or energy.Examples include reinforcement-learning or genetic-programming search and solver-based approaches.

3 Problem Definition and Background

The paper formulates compression as a layer-wise constrained approximation problem using calibration inputs and output loss. Under squared loss, OBS provides exact greedy pruning updates, but its direct implementation is computationally impractical.

  • Problem Definition: Layer-wise compression minimizes expected output change while satisfying a generic compression constraint on the compressed weights.The expectation is over layer inputs and the constraint is customized for pruning or quantization.
  • Problem Definition: The expectation is approximated using a small calibration sample, and squared loss is commonly used for linear and convolutional layers.The squared-loss formulation is motivated through approximations involving second-order information.
  • Problem Definition: For a layer, the weight matrix has dimensions d_row × d_col and the input matrix has dimensions d_col × N.For convolutional layers, d_col corresponds to the number of weights in one filter.
  • OBS Framework: OBS uses the Hessian at a nearly stationary dense model to select a weight for removal and update remaining weights to compensate.The selection and compensation formulas depend on the inverse Hessian.
  • OBS Framework: For quadratic layer-wise loss starting from dense weights with zero loss, iterated OBS formulas provide an exact greedy pruning solution.The exactness applies to this specific problem formulation.

4 An Optimal Greedy Solver for Sparsity

ExactOBS makes the OBS greedy pruning procedure practical at DNN scale by exploiting row-wise structure and efficient inverse-Hessian updates. It preserves exactness while addressing runtime and memory bottlenecks in global pruning.

  • Motivation: Exact OBS pruning is computationally demanding because it updates and inverts a d×d Hessian across O(d) single-weight pruning steps.The resulting total runtime is Θ(d4), which is impractical for layers with 10^5–10^6 parameters.
  • Row-wise decomposition: Writing the squared error row-wise shows that each weight affects only its corresponding output row, so each row can use an independent dcol × dcol Hessian.For the least-squares row objective, the Hessian is H = 2XX⊤.
  • Efficient updates: ExactOBS avoids repeated full inversions by computing the shared Hessian once and extracting masked rows and columns as pruning proceeds.The inverse-Hessian removal step uses Gaussian elimination to remove a parameter’s row and column directly.
  • Efficient updates: The resulting single-row algorithm prunes k weights in O(k · dcol^2) time while updating the remaining weights after each OBS step.Algorithm 1 maintains the pruning process without constantly resizing the inverse Hessian.
  • Global selection: Global pruning is recovered by recording row-wise loss changes and selecting the lowest-scoring weights across all rows.This requires only Θ(d) extra memory after the row-wise traces or masks are available.
  • Structured sparsity: ExactOBS extends to N:M and block sparsity by selecting eligible weights or groups with the smallest loss increase.Block pruning applies the inverse-Hessian update successively to all parameters in a selected block.

5 The Optimal Brain Quantizer (OBQ)

OBQ extends OBS from pruning to quantization by replacing zeroing constraints with quantization-rounding constraints. It quantizes weights greedily while updating the remaining unquantized weights to compensate for loss increase.

  • From OBS to OBQ: OBQ adapts OBS to quantization by constraining a selected weight’s update to its quantization error, quant(wp) − wp.After the update, the selected weight equals its quantized value.
  • Greedy quantization: The optimal quantization order minimizes the loss increase at each step, while the corresponding update adjusts the remaining unquantized weights.The derivation assumes the gradient at the current point is negligible and uses the Hessian-based OBS formulation.
  • Unified formulation: OBQ generalizes pruning OBS because setting quant(·) to zero recovers the original pruning formulation.The same final formulas therefore cover both zeroing and rounding operations.
  • Full-layer quantization: For full-layer quantization, OBQ repeatedly quantizes the easiest remaining weight and compensates by changing other unquantized weights.This can produce assignments different from initial rounding and improve overall quantization results.
  • Practical issue: Quantization grids with outliers can make greedy ordering problematic because high-error outliers may be quantized late, when fewer weights remain for compensation.Intermediate updates can push some weights farther outside the grid, worsening this effect.

6 Experiments

Experiments evaluate OBC across pruning, quantization, and compound compression settings using small calibration sets and per-layer target selection. OBC generally delivers strong accuracy–compression trade-offs, including practical GPU BOP and CPU speedup targets.

  • Experimental Setup: Calibration uses 1024 random training samples, with 10× augmentation for ImageNet and batchnorm resetting for ResNets.Other tasks use no augmentation.
  • Unstructured Sparsity: ExactOBS performs best by a wide margin on RN18 first-layer squared error across several sparsity levels.AdaPrune significantly outperforms global magnitude pruning and L-OBS.
  • Unstructured Sparsity: ExactOBS generally leads pruning models at 2×–4× FLOP reduction, including a greater than 1% gap at 4× and reasonable high-reduction results on BERT.For YOLOv5l at 2×, all methods perform similarly in mAP@0.5; global AdaPrune can recover lost accuracy for difficult BERT settings.
  • N:M Sparsity: ExactOBS matches or slightly exceeds AdaPrune’s 4:8 results using the more stringent 2:4 pattern, and achieves 1–2% higher F1 on BERT.The experiments cover N:M pruning with batchnorm tuning on ResNets and 2:4 pruning on BERT.
  • Quantization: OBQ achieves similar or slightly better accuracy than sequential post-training quantization methods at 4 and 3 bits while independently compressing layers.This supports rapidly generating mixed-precision models for different constraints.
  • BOP-Constrained Mixed GPU Compression: At approximately 2.5% relative performance loss, mixed GPU compression achieves 12–14× BOP reduction for ResNets and 7–8× for YOLO and BERT.The setting combines quantization with 2:4 sparsity and selects among four per-layer compression choices.
  • Time-Constrained CPU Compression: On a 12-core Intel Xeon CPU, joint 8-bit quantization and block sparsity achieve 4× and 5× actual speedups with 1% and 2% accuracy loss, respectively.The experiments use real layer-wise timing data and DeepSparse blocksize-4 acceleration.

7 Conclusions & Future Work

The paper concludes that its framework provides efficient, accurate post-training compression for pruning and quantization, with extensions and future investigation planned.

  • The framework delivers state-of-the-art post-training solutions for both pruning and quantization.
  • The approach may extend naturally to structured pruning and remain compatible with additional unstructured pruning and quantization.
  • The results suggest post-training compression may achieve accuracies comparable to more expensive retraining methods.
  • Future work will examine this possibility for more resource-intensive models, including very large-scale language models.

A.1 Proof of Lemma 1 (Row & Column Removal)

The proof shows that equivalent row and column transformations eliminate selected entries of a matrix and recover the inverse of the remaining block.

  • Equivalent matrix transformations zero [A]ij by subtracting a scaled column, while analogous row operations eliminate entries in a column.
  • Gaussian elimination of one row and column corresponds to zeroing the associated off-diagonal entries of H^-1.
  • Applying these transformations to H^-1H = I yields an equation in AB = C form.
  • Because eliminated-row and eliminated-column entries multiply zeros, the remaining A_i block is the inverse of the corresponding B_i block.

A.2 ExactOBS Global Step Pseudocode

The ExactOBS global step combines row-wise pruning orders and loss changes to determine a global mask with a specified number of pruned weights.

  • The procedure uses P, the per-row pruning order, and L, the corresponding loss changes, to determine the global OBS mask.
  • At each step, it selects the smallest next loss change in Q, removes that candidate, and inserts the following candidate from the same row.
  • After k selections, Q records the number of pruned elements per row, which together with P yields the mask.
  • A min-heap can implement Q efficiently, while selecting the k smallest entries in L usually gives essentially the same practical results.

A.3 OBQ-ExactOBS Algorithm Pseudocode

The OBQ variant applies the ExactOBS framework to quantize weights according to OBS, paralleling the pruning algorithm.

  • The OBQ version of ExactOBS is structurally similar to the pruning variant.
  • The algorithm quantizes k ≤ dcol weights from a row using the inverse Hessian H^-1 = (2XX^T)^-1 according to OBS.
  • The procedure initializes the available index set as M = {1, ..., dcol} and iterates for i = 1, ..., k.

A.4 Further Experiment Details

The experiments detail OBC’s implementation choices, runtime behavior, and comparisons across pruning, quantization, and compound compression. Results show strong accuracy advantages, practical runtimes, and low seed sensitivity across evaluated settings.

  • Non-Uniform Sparsity Choices: δ = 0.9 prunes 10% of the remaining weights at each sparsity choice, with options extending beyond 0.99 unstructured or 0.95 blocked sparsity.The grids are designed for flexibility, though fewer carefully selected options may often suffice.
  • Activation Quantization: Activation quantization optimizes each layer’s zero point and scale from one calibration batch using the same procedure as weight quantization.The procedure operates on tensors rather than channels and stores quantization information per layer.
  • Runtime: About one hour is required to quantize ResNet50 to 4 bits with BRECQ, AdaRound, or OBQ, while BitSplit takes about twice as long and AdaQuant is 3× faster.ExactOBS’s runtime is therefore in line with existing post-training methods, although AdaQuant is considerably less accurate.
  • Runtime: Quantization and unstructured pruning take about the same time, while 2:4 pruning and quantizing a 2:4-pruned model cost approximately half as much.Blocked pruning is most expensive for YOLO and BERT because of additional c × c block-matrix overhead.
  • Runtime: 75% of ResNet50’s compression runtime is spent in just three 3 × 3 convolutions in the last block, while most earlier layers finish within seconds.These layers have unfolded column dimensions of approximately 4500, making a few large layers dominate total runtime.
  • Multiple AdaPrune Iterations: After 16 AdaPrune iterations, its accuracy drop remains almost 2× larger than ExactOBS’s at comparable overall runtime.The results also show that a few AdaPrune recomputations rapidly reduce F1 drop.
  • Independent Quantization Comparison: At 3 bits, OBQ exceeds the other independent methods by several accuracy points, while at 2 bits it is the only method that avoids complete breakdown without statistics correction.This comparison uses symmetric per-channel quantization and raw performance immediately after independent compression.
  • Sequential Quantization with OBQ: Sequential OBQ matches independent OBQ at 4 and 3 bits, while sequential processing provides a noticeable benefit at 2 bits and catches up with BRECQ.Sequential application requires reoptimizing dense weights for compressed inputs before applying OBQ.
Loading 2208.11580v2…