Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning
Abstract
We consider the problem of model compression for deep neural networks (DNNs)in the challenging post-training setting, in which we are given an accurate trainedmodel, and must compress it without any retraining, based only on a small amountof calibration input data.
This problem has become popular in view of the emergingsoftware and hardware support for executing models compressed via pruning and/orquantization with speedup, and well-performing solutions have been proposedindependently for both compression approaches.
In this paper, we introduce a newcompression framework which covers both weight pruning and quantization in aunified setting, is time- and space-efficient, and considerably improves upon thepractical 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] extendedto also cover weight quantization at the scale of modern DNNs, and is enabledby a series of algorithmic developments which may be of independent interest.
From the practical perspective, our experimental results show that it can improvesignificantly upon the compression-accuracy trade-offs of existing post-trainingmethods, and that it can even enable the accurate compound application of bothpruning and quantization in a post-training setting.
Introduction
The impressive recent progress of deep learning for solving challenging tasks across several domainshas been accompanied by a significant increase in parameter counts and computational costs forexecuting such models.
A natural consequence has been a growing effort to reduce such costs viamodel compression , and the two most popular approaches for model compression are pruning —removing neural network weights by setting them to zero—and quantization , reducing the precisionat which neural network weights and activations are stored and manipulated.
Hundreds of suchpruning and quantization approaches have been proposed and analyzed [17, 11] with the general goalof obtaining efficient deep neural nets (DNNs) which would preserve accuracy while maximizingcompression.
Despite impressive progress, compression is still a laborious process, as the pruning andquantization stages are often done independently, and recovering model accuracy after compressionoften requires partial or even full retraining of the compressed model.
By contrast, an alternative but challenging scenario is the post-training compression setup [31, 24,19, 25], in which we are given a trained but uncompressed model, together with a very small amountof calibration data , and must produce an accurate compressed model in one shot , i.e. a singlecompression step, without any additional retraining.
This is the setting we focus on in this paper.
Compression via weight pruning started with seminal work by LeCun et al.
[23], complemented by
Hassibi and Stork [13], who proposed a mathematical framework called the Optimal Brain Surgeon(OBS) , for choosing the set of weights to remove from a trained neural network, by leveraging secondorder information.
(We describe their approach in Section 3.)
Recent advances, e.g. [6, 39, 38, 10]showed that OBS can lead to state-of-the-art compression at DNN scale, by introducing numericalmethods which can approximate the second-order information needed by OBS at the massiveparameter counts of modern models.
However, these approaches do not apply to the post-trainingsetting, as they require gradual pruning, as well as significant retraining, to recover good accuracy.
An alternative approach, which is standard in the context of post-training compression, has been tobreak the task into layer-wise sub-problems, essentially finding a compressed weight approximationfor each layer, given a sub-sample of the layer’s inputs and outputs obtained from calibration data.
This line of work, e.g. [41, 31, 19], introduced elegant solvers for the resulting layer-wise weightquantization problem, which achieve state-of-the-art results for post-training quantization.
Recently,
AdaPrune [18] showed that this approach can also be effective for post-training weight pruning.
In this context, a natural question is whether existing approaches for pruning and quantization can beunified in order to cover both types of compression in the post-training setting, thus making DNNcompression simpler and, hopefully, more accurate.
This question is also of practical importance,since both GPU and CPU platforms now jointly support sparse and quantized formats [30, 35], and,as we illustrate experimentally, the resulting models could be executed with compound speedups.
Contribution.
In this paper, we provide a mathematical framework for compression via pruning orquantization, which leads to state-of-the-art accuracy-versus-compression trade-offs in the challengingpost-training compression setup.
Our framework starts from the layer-wise compression problemdescribed above, by which the global compression task, defined either for pruning or quantization,is first split into layer-wise sub-problems, based on the layer behavior on the calibration data.
Specifically, given a layer ℓ defined by weights W ℓ , and layer inputs X ℓ , we aim to find a compressedversion of the weights (cid:99) W ℓ which minimizes the output difference relative to the uncompressed layer,measured via the squared error between the original and compressed layer, acting on the sample input|| W ℓ X ℓ − (cid:99) W ℓ X ℓ || 22 , under a fixed compression constraint on (cid:99) W ℓ .
Although solving this problem optimally for sparsity or quantization constraints is NP-hard [2, 31], itis a key step in all state-of-the-art post-training compression methods, both for pruning [18, 9] andfor quantization [31, 18, 24].
Once this is solved per layer, a solution to the global problem can beobtained by combining layer-wise solutions, which is handy especially for non-uniform compression,e.g. [15, 9].
Thus, several approximations for this problem have been proposed [31, 19, 18].
We show that there still is significant room for improvement when solving the layer-wise compressionproblem.
Roughly, our approach is to specialize the OBS framework to the squared error formulationabove: in this case, the framework can in theory produce an exact greedy solution, but a directimplementation would have infeasible Θ( d 4 ) computational cost, where d is the layer dimension.
Ourmain technical contribution is a series of algorithms which reduce this computational cost, withoutany approximations , to O ( d · d 2 col ) where d col is the column dimension of the weight matrix.
Inpractice, these improvements are significant enough to allow us to implement the exact OBS greedysolution, which prunes one weight at a time , and updates all remaining weights after each step, at thescale of modern DNNs with tens of millions of parameters, within reasonable time, on a single GPU.
We provide an efficient PyTorch implementation of our algorithms in the Supplementary Material.
In turn, this algorithmic development allows us to apply the OBS approach to quantization .
Theresulting algorithm, called the Optimal Brain Quantizer (OBQ) , quantizes weights iteratively one-ata-time, depending on their impact on the loss increase, after which it applies a closed-form update tothe remaining unquantized weights, further reducing the loss.
This solves the two problems efficiently,and in a unified manner—we call the unified framework the Optimal Brain Compressor (OBC) .
Experimental Results.
We apply OBC to standard tasks and models covering image classification,object detection, and language modelling applications.
We first show that our framework yieldssignificantly better solutions for the layer-wise compression problem, which leads to higher-accuracyend-to-end compressed models for both pruning and quantization, relative to the corresponding stateof-the-art techniques, often by significant margins.
Second, we show that our pruning and quantizationapproaches can be compounded, with surprisingly strong results: we obtain a 12 × reduction intheoretical operations with a 2% accuracy drop for GPU-supported compound compression [30], anda 4 × speedup in actual runtime with only 1% accuracy drop for a CPU-based runtime [35].
Theseresults show for the first time that post-training compression can be competitive with full retraining.
Related Work
Optimal Brain Surgeon (OBS).
The classic OBS framework [23, 13] was originally applied tonetworks with hundreds of weights; more recently, methods such as WoodFisher [38] rendered theapproach computationally feasible for DNNs by using a block-diagonal Fisher approximation ofthe Hessian, while a follow-up method called M-FAC [10] introduced more efficient and generalalgorithms for handling the inverse Fisher matrix.
Earlier work called Layer-wise OBS (L-OBS) [6]:inspired by the K-FAC approximation [29, 12], L-OBS approximates the OBS framework not for the global objective, but for a quadratic per-layer loss, while also pruning all weights based on a single Hessian computation.
At a high level, our approach is similar, in that we apply OBS layer-wise;however, we apply OBS exactly , that is, pruning one weight at a time, and exactly recomputing the Hessian after every pruning step.
This is made computationally tractable by several new algorithmicideas, and yields significantly improved results relative to L-OBS.
This prior work on pruningconsidered settings with extensive finetuning.
By contrast, we will focus on the post-training setting,where only a small amount of calibration data is available.
Post-Training Quantization.
This setting has been primarily considered for quantization, andmost state-of-the-art methods work by performing layer-wise compression.
Specifically, BitSplit[6] optimizes the quantized weights bit by bit, while AdaRound [31] finds a weight rounding policythrough gradient based optimization with an annealed penalty term that encourages weights to movetowards points on the quantization grid.
AdaQuant [19] relaxes the AdaRound constraint, allowingweights to change during quantization-aware optimization, via straight-through approximation [33].
BRECQ [24] suggested that accuracy can be improved further by integrating second-order informationinto the layer-wise losses and by jointly optimizing hand-crafted blocks of related layers.
A key step of AdaRound, AdaQuant and BRECQ is to quantize layers incrementally, in sequentialorder, so that errors accumulated in earlier layers can be compensated by weight adjustments in laterones.
This significantly improves performance, but reduces flexibility, as the entire process may needto be re-done whenever we wish to change compression parameters of one layer.
We instead targetindependent compression of each layer, allowing the end model to be simply “stitched” together fromlayer-wise results.
Despite operating independently on each layer, we find that, after correcting basicstatistics such as batchnorm, our method performs on par to sequential ones for uniform quantization.
Post-Training Sparsification.
The layer-wise approach was shown to also be effective for posttraining pruning by AdaPrune [18], which pruned weights to the GPU-supported N:M pattern [44].
AdaPrune first drops parameters according to their magnitude [45] and then reoptimizes the remainingweights to reconstruct the pre-compression calibration set output.
This is similar to [16, 8] whichalso perform layer-wise reoptimization of the remaining weights.
Follow-up work [10] noted that theresults of AdaPrune can be improved further by performing more frequent pruning/optimization steps.
Our algorithm pushes this idea to the limit, performing full reoptimization after every single prunedweight, while remaining computationally tractable.
We further use a more sophisticated weightselection metric which incorporates second-order information.
Finally, [10] also introduces global
AdaPrune , a more expensive global optimization step applied on top of the layer-wise AdaPruneresults, which can bring additional accuracy gains.
This can also be applied to our pruned models.
Non-Uniform Compression.
An orthogonal practical question is how to compress different layersto maximize accuracy under a given resource constraint, such as latency or energy consumption.
Existing methods can be roughly categorized into search-based and solver-based approaches.
Theformer, e.g. AMC [15] or HAQ [40], search for a layer-wise compression policy directly via, forexample, reinforcement learning or genetic programming [42], whereas the latter, e.g. HAWQv3 [43]or AdaQuant [19], construct a relaxed version of the overall problem that is then solved exactly.
We focus here on solver-based approaches, as they can rapidly adapt to different scenarios whencombined with accurate independent layer-wise compression schemes; however, our techniques couldbe of interest for search-based methods as well.
Concretely, we use the problem formulation of
AdaQuant [19] to which we apply the DP algorithm of SPDY [10] to achieve fast solving times evenwith a large number of possible choices per layer.
Problem Definition and Background
The Layerwise Compression Problem.
We begin by defining the general problem that is the focusof this work, in line with prior work on post-training compression, e.g. [31, 19].
Mathematically, wemodel a layer ℓ as a function f ℓ ( X ℓ , W ℓ ) acting on inputs X ℓ which is parametrized by weights W ℓ .
The goal of layer-wise compression is to find a “compressed” version of W ℓ that performs as similarlyas possible to the original weights.
More formally, the compressed weights (cid:99) W ℓ should minimizethe expected layer output change as measured by some loss L while at the same time satisfying acertain compression constraint, which we generically denote by C ( (cid:99) W ℓ ) > C here, but which will becustomized depending on the compression type:
[Equation]
The expectation over the layer inputs X ℓ is usually approximated by taking the mean over a smallset of N input samples.
This low-data setting is one of the primary applications of layer-wise compression.
Further, most works [41, 31, 19] focus on compressing linear and convolutional layers,which can be unfolded into linear ones, as these are prevalent in practice, and use the squared lossto measure the approximation error.
This definition of the loss can be motivated, via a sequence ofapproximations, from second-order information: please see [31] for a precise derivation.
Furthermore,this approximation approach has been shown to work well in many applications [31, 19, 9].
We follow these conventions as well, and work with the specific layer-wise compression problemstated formally below, where the weights W ℓ are a d row × d col matrix (for conv-layers d col correspondsto
the total number of weights in a single filter), and the input X ℓ has dimensions d col × N .
[Equation]
The Optimal Brain Surgeon (OBS) Framework.
The OBS framework [23, 13] considers theproblem of accurately pruning a trained dense neural network.
It starts from the Taylor approximationat the given point (assumed to have negligible gradient), and provides explicit formulas for theoptimal single weight to remove, as well as the optimal update of the remaining weights which wouldcompensate for the removal.
More precisely, let H denote the Hessian matrix of the loss at thegiven (dense) model.
Then the weight to prune w p which incurs the minimal increase in loss and thecorresponding update of the remaining weights δ p can be calculated as follows:
[Equation]
where [ H − 1 ] pp denotes the p th diagonal entry of the inverse Hessian, and H − 1 : ,p is its p th column.
OBS for Layer-Wise Pruning.
We will now instantiate this framework for the layer-wise pruningproblem, defined above.
First, the loss in equation (2) is quadratic and since our starting point is givenby the dense weights achieving the minimal loss of 0, the assumptions of the OBS framework arefully met, meaning that its formulas are exact for this specific problem formulation.
Thus, iteratingthe OBS framework to remove one weight at a time would yield an exact greedy solution for thelayer-wise pruning problem, as it takes the (locally) optimal decision at each step.
While this greedyapproach does not guarantee convergence to a global optimum, such approaches can be very effectivefor dealing with problem instances that are too large to be handled by exact methods.
An Optimal Greedy Solver for Sparsity
The obvious challenge is that applying the OBS framework in its true form, i.e. pruning a single weightat a time using the exact formulas in (3), is computationally very demanding.
The Hessian H is a d × dmatrix where d = d row · d col , which is already expensive to store and compute with.
Additionally, thismatrix needs to be updated and inverted at each of the O ( d ) steps with a computational complexityof Θ( d 3 ) .
Clearly, an O ( d 4 ) total runtime is too inefficient for pruning most layers of modern neuralnetworks, as d is usually ≥ 10 5 or even ≥ 10 6 for several layers.
However, as we will now show, it isactually possible to reduce the overall costs of this process to O ( d row · d 3 col ) time and Θ( d 2 col ) memory,making it efficient enough to prune e.g. all layers of a medium-sized model such as ResNet50 in abit more than one hour on a single NVIDIA RTX 3090 GPU.
We emphasize that the techniques weintroduce are exact; unlike prior work [6, 38], we do not rely on any approximations.
The ExactOBS Algorithm.
In the following, we introduce our efficient instantiation of the OBSframework, for the layer-wise compression problem, which we call ExactOBS, in step-by-step fashion.
We start by rewriting the matrix squared error in (2) as the sum of the squared errors for each row inthe weight matrix.
As we are always dealing with a fixed layer ℓ , we drop the subscript ℓ to simplifynotation.
The objective is then equivalent to (cid:80) d row i =
1 || W i , : X − (cid:99) W i , : X
|| 22 .
This way of writing the error makes it clear that removing a single weight [ W ] ij only affects theerror of the corresponding output row Y i , : =
W i , : X .
Hence, there is no Hessian interaction betweendifferent rows
and so it suffices to work only with the individual d col × d col Hessians correspondingto each of the d row rows.
Further, as the dense layer output Y = WX is fixed, the objective for eachrow has standard least squares form and its Hessian is given by H =
2XX ⊤ .
Although this observation already reduces the computational complexity, two key challenges remain:(a) applying OBS to each row still costs O ( d col · d 3 col ) time, which is too slow for large layers, and (b)we need fast access to the Hessian inverses of all d row rows, since we want to prune the minimumscore weight of the whole matrix rather than just per row in each step.
In particular, (b) requires
O ( d row · d 2 col ) GPU memory, which is likely to be infeasible.
Step 1: Handling a Single Row.
We first describe how to efficiently prune weights from a singlerow with d col parameters.
For simplicity, we denote such a row by w with corresponding Hessian
H .
The full algorithm for this procedure is given in Algorithm 1; in the following, we provide adetailed description.
The key idea is to avoid having to do the full Θ( N · d 2 col ) calculation and Θ( d 3 col )inversion of H in each step.
The former is easy, as the weights themselves do not enter the calculationof H = 2 XX ⊤ , and the Hessian for the weights with pruning mask M denoted by H M is thussimply comprised of the corresponding rows and columns in the fully dense version H .
Hence, weonly have to compute H (which is actually the same for all rows) once, from which we can thenextract the rows and columns corresponding to M as needed.
Critically, this trick is not applicable to the inverse, as ( H M ) − 1 ̸
= ( H − 1 ) M .
However, using thefact that the removal of one parameter p simply drops the corresponding row and column from H ,we can actually update the inverse to remove parameter p directly using a single step of Gaussianelimination, with cost Θ( d 2 col ) .
The following result, whose proof is in the Appendix, formalizes this.
Lemma 1 (Row & Column Removal) .
Given an invertible d col × d col matrix H and its inverse H − 1 ,we want to efficiently compute the inverse of H with row and column p removed, which we denote by
H − p .
This can be accomplished through the following formula:
[Equation] [Equation]
which corresponds to performing Gaussian elimination of row and column p in H − 1 followed bydropping them completely.
This has Θ( d 2 col ) time complexity.
The resulting pseudocode is shown in Algorithm 1, where we avoid constantly resizing H − 1 (andcorrespondingly changing indices) by utilizing the fact that row and column p have no effect onany future calculations after they have been eliminated by Lemma 1 as they are 0 (and the non-zerodiagonal element is never accessed again).
One can check that this algorithm applies OBS to a singlerow of W with a per-step cost of Θ( d 2 col ) , and thus Θ( k · d 2 col ) overall time for pruning k weights.
Algorithm 1 Prune w with inverse Hessian H − 1 = (2 XX ⊤ ) −according to OBS in O ( k · d 2 col ) time.
M = { 1 , . . .
, d col } for i = 1 , . .
, k do p ← argmin p ∈ M
[ H − 1 ] pp · w 2 p w ← w − H − 1 : ,p 1 [ H − 1 ] pp · w p H − 1 ← H − 1 − 1 [ H − 1 ] pp H − 1 : ,p H − 1 p, :
M ← M − { p } end for
[Figure]
Step 2: Jointly Considering All Rows.
Applying the OBS framework to the full weight matrix Wrather than just to each row independently requires fast access to all d row row-wise inverse Hessians,in order to select the weight with the smallest overall pruning score in each step.
However, storingd row matrices of size d col × d col each in GPU memory can be too expensive; while it would be possibleto offload some Hessians to main memory, this could result in a large number of expensive memorytransfers.
However, since there is no Hessian interaction between rows, the final compressed weightsof each row only depend on the total number of parameters that were pruned in it.
Similarly, thechange in loss incurred by pruning some weight only depends on the previously pruned weights inthe same row, which also means that the order in which weights are pruned in each row is fixed.
The consequence of these insights is that we can process each row independently, pruning all weightsin order while always recording the corresponding change in loss δ L p =
w 2 p /
[ H − 1 ] pp .
At the end,we know δ L p for all d weights and can then simply determine the global mask that would be chosenby OBS on the full matrix by selecting the weights with the lowest values in order, requiring onlyΘ( d ) extra memory.
We note that once the per-row masks M i are known, we can directly solvefor the optimal update of the remaining weights via the corresponding group OBS formula [21]δ − M i
= H − 1 : ,
− M i
( H − 1 ) − M i ) − 1 w − M i where the index set − M i denotes all indices not in mask
M i .
This will be considerably faster in practice than simply rerunning the iterative pruning process in
Algorithm 1.
Alternatively, if enough CPU memory is available, one can keep the full pruning traceof each row, that is, the full weight vector after every individual pruning step, in CPU memory andultimately simply reload the entries corresponding to the global mask.
This requires O ( d row · d 2 col )extra CPU memory but avoids a second computation pass to reconstruct the not pruned weights andwill therefore be faster.
Figure 1 visualizes both options just discussed.
Implementation Details.
In practice, the matrix H might not always be invertible for reasons such asusing too few data samples or dead / linearly dependent inputs.
The former can usually be addressedby extending the calibration dataset with augmentations (additional augmented samples only needto be accumulated into the Hessian once and are thus very cheap to include) and the latter can beprevented by adding a small diagonal dampening term to the Hessian before inverting it.
Second,a direct GPU implementation of Algorithm 1 will perform a large number of small CUDA calls,which can be expensive.
This overhead can be removed by using batch operations to process multiplematrix rows simultaneously—for more details please see our sample implementation provided inthe Supplementary.
Finally, when applied to an already sparse weight matrix, the complexity of ouralgorithm can scale cubicly with the row-density by working with a dense version of the weights /
Hessians consisting only of the non-zero elements and mapping the pruning result back at the end.
N:M Sparsity.
Our method can be easily extended to various forms of semi-structured sparsity.
Thisincludes, for example, the N:M sparsity pattern [44], which enforces exactly N non-zero values ineach block of M consecutive weights, and is becoming popular due to support on newer NVIDIAhardware [30].
Adapting our algorithm to this pattern requires only one simple change: instead ofselecting the weight with the smallest change in loss, we select the weight with the smallest change inloss that is in a block with < N pruned weights.
We note that all rows have exactly the same sparsity1 − N/M in the N:M pattern and so we can terminate per-row pruning as soon as this target sparsityvalue is reached.
For the same reason, there is no need for the global mask selection step describedearlier.
Thus, our method will be even more efficient in this case.
Block-Sparsity.
Another practically relevant pruning pattern, particularly in the context of CPUacceleration [7, 22], is block-pruning , where zeros appear only in consecutive blocks of size c , whichis typically a small number like 4 or 8.
We follow recent work [21] that extends the OBS frameworkto pruning small groups of connected weights in order to account for the correlation between them,using the following formulas for the target block and weight update, respectively:
[Equation]
where P denotes the set of indices corresponding to one block.
Algorithm 1 can easily be adaptedto operate on blocks using the above equations and applying the update of H − 1 via Lemma 1successively for all p ∈ P .
Although there are now only d col /c steps per row, each update of H − 1also takes O ( c · d 2 col ) time and so the overall asymptotic runtime stays the same.
Additional practicaloverhead only comes from the extra O ( c 2 · d 2 col ) terms that are the result of computing and multiplyingwith the c × c matrices (( H − 1 ) P ) − 1 .
The Optimal Brain Quantizer (OBQ)
Although the classical OBS framework [23, 13] has inspired a long line of work on pruning methodsfor DNNs [38, 10, 27], so far it has not been used for quantization.
We now show that our resultsfrom the previous section can in fact be extended to quantization in an effective and accurate way, viaa new method which we call the Optimal Brain Quantizer (OBQ).
The Quantization Order and Update Derivations.
Under the standard assumption that the gradientat the current point w is negligible, the OBS formulas for the optimal weight to be pruned w p and thecorresponding update δ p can be derived by writing the locally quadratic problem under the constraintthat element p of δ p is equal to − w p , which means that w p is zero after applying the update to w .
This problem has the following Lagrangian:
[Equation]
where H denotes the Hessian at w and e p is the p th canonical basis vector.
The optimal solution isthen derived by first finding the optimal solution to δ p via setting the derivative ∂L/∂
δ p to zero andthen substituting this solution back into L and solving for λ ; please see e.g. [13, 38] for more details.
Assume a setting in which we are looking to quantize the weights in a layer on a fixed grid of width∆ while minimizing the loss.
To map OBS to a quantized projection, we can set the target of the
Lagrangian constraint in (6) to ( quant ( w p ) − w p ) , where quant ( w p ) is the weight rounding given byquantization; then w p = quant ( w p ) after the update.
Assuming we wish to quantize weights iteratively, one-at-a-time, we can derive formulas for the “optimal” weight to quantize at a step, in terms of minimizing the loss increase, and for the correspondingoptimal update to the unquantized weights, in similar fashion as discussed above:
[Equation]
In fact, since − w p is a constant during all derivations, we can just substitute it with ( quant ( w p ) − w p )in the final result.
We note that the resulting formulas are a generalization of standard OBS forpruning, if quant ( · ) always “quantizes” a weight to 0, then we recover the original form.
Quantizing Full Layers.
At first glance, OBQ might appear curious since one usually quantizesall weights in a layer, leaving no more weights to update.
At the same time, the weight selectionmetric influences only the quantization order, but not the quantization value.
However, this viewchanges when considering OBQ in the context of our efficient one-weight-at-a-time pruning algorithmdescribed in the previous section.
Specifically, using OBQ, we can greedily quantize the currently“easiest” weight by the above metric, and then adjust all the remaining unquantized weights tocompensate for this loss of precision, thus changing their value.
We then choose the next weight toquantize, and so on.
This can result in quantization assignments that are different from the ones thatwould have been chosen by rounding initially, and in better overall quantization results.
Concretely, torealize this, we can plug (7) into Algorithm 1 in order to iteratively quantize weights for a given layer,leading to the similar Algorithm in the Appendix, thus essentially unifying pruning and quantization.
Quantization Outliers.
One practical issue with this greedy scheme can occur especially whenapplied to quantization grids that permit some outliers in order to achieve a lower error on the majorityof weights, which are currently standard [4, 34].
Since these outliers can have high quantization error,they will usually be quantized last, when there are only few other unquantized weights available thatmay be adjusted to compensate for the large error incurred by quantizing the outliers.
This effectcan become worse when some weights are pushed even further outside the grid by intermediateupdates.
We prevent this with a simple but effective heuristic: we quantize outliers, e.g. weightswith a quantization error > ∆ / 2 where ∆ is the distance between quantized values, as soon as theyappear (which typically happens only a few times per layer).
With this heuristic, OBQ yields a highlyeffective layer-wise quantization scheme, as our experiments in the next section demonstrate.
Finally,we note that the OBQ version of the techniques discussed in Section 4 has all the same runtime andmemory characteristics (barring the global step in Figure 1, which is unnecessary for quantization).
Experiments
Objectives, Models & Datasets.
To demonstrate the effectiveness and flexibility of our method, weconsider several different standard post-training compression scenarios [31, 19, 18].
We begin withsettings where only a single type of compression is applied: concretely, we consider unstructuredpruning for given FLOP targets, global 2:4 and 4:8 pruning, as well as uniform weight quantization.
Additionally, we also study two practical tasks that feature joint pruning and quantization: a GPUscenario where quantization and N:M pruning are combined, as well as a CPU scenario combiningquantization and block pruning.
We work with variants of the following models and tasks: ResNet[14] for image classification on Imagenet [37], YOLOv5 [20] for object detection on COCO [26] and BERT [5] for question answering on SQuAD [36].
Our smaller BERT models denoted by BERT3 and BERT6 correspond to the smaller 3 and 6 layer variants of BERT-base, respectively, trained by [21].
The Appendix contains additional experiments as well as runtime information of our algorithms.
Experimental Setup.
All of our calibration datasets consist of 1024 random training samples.
For
ImageNet, where we use roughly 0 . 1% of the training data, we additionally apply standard flippingand cropping augmentations to artificially increase the size of this dataset by 10 × ; other tasks do notuse any augmentations.
While the effect of augmentations is typically minor, they are very cheap toinclude for our method.
For ResNet models, batchnorm statistics are reset using 100 batches of 128samples from the calibration set with standard augmentations.
For other models, we apply mean andvariance correction [32, 1] after all normalization layers (so that the correction parameters can beeasily merged and incur no extra cost) on a single batch of samples of size 128 (for YOLO) and 512(for BERT).
We found this to be more effective than batchnorm tuning for YOLO, and the BERTmodels have no batchnorm layers.
When compressing to a given FLOP or timing constraint, we need to solve the problem of identifyingper-layer compression targets, which match the constraint, while maximizing accuracy.
To identifythese non-uniform targets, we follow the approach of [10]: we first collect a “model database” containing for each compression level (e.g. bit-width or sparsity setting) the corresponding (independently) compressed version of each layer.
For building a joint sparse and quantized database we simplysparsify layers first and then apply quantization to the remaining weights.
Next, similarly to [19],we compute the layer-wise calibration losses (without augmentations) for all compression levels,corresponding to the models with exactly one layer compressed to a certain level.
Then, givenlayer-wise FLOP or timing information, we set up a constrained layer-wise compression problemof the form described in AdaQuant [19] and solve it with the dynamic programming algorithm of
SPDY [10].
This returns an optimal per-layer assignment of compression levels, for which we canthen easily produce the corresponding model, via a two-step process: we first stitch together layers atthe corresponding compression levels from the database, and then perform the discussed statisticscorrection to recover some extra accuracy [19].
Unstructured Sparsity.
We begin our experiments with unstructured sparsity, comparing againstglobal magnitude pruning (GMP) [45], the approximate layer-wise OBS method L-OBS [6], andthe post-training pruning state-of-the-art method AdaPrune [18].
As a sanity check, we examine in
Figure 1 whether our method provides better results in terms of layer-wise squared error, pruning thefirst layer of a ResNet18 (RN18) model to several sparsities.
In this metric, ExactOBS performs bestby a wide margin ahead of AdaPrune, which significantly outperforms the other two methods.
Next, in Table 1, we turn our attention to the practical problem of pruning various models to achievea given FLOP reduction of 2 × – 4 × , applying the per-layer target sparsity optimization techniquedescribed above.
Our ExactOBS generally performs best (except for YOLOv5l 2 × where all methodsperform similarly in terms of mAP@0.5) and at
× FLOP reduction even with a > 1% gap to thenext best method.
Interestingly, on the hard-to-prune BERT model, ExactOBS appears to be theonly method which still produces reasonable results at higher reduction targets.
For BERT 3 × and4 × , where the performance drop of all methods is > 3% , we additionally assess the compatibilityof our results with the more powerful (but also more expensive) post processing method global
AdaPrune [10].
While this global optimization technique is able to recover lost accuracy, the
ExactOBS models still maintain a > 1% and > 2 . 5% F1 advantage, respectively (see Table 5).
[Figure]
[Table]
N:M Sparsity.
Next, we study the performance of our method for semi-structured sparsity via the
N:M pattern.
Specifically, we compare against the 4:8 results of AdaPrune with batchnorm tuning [18]on ResNet models (see Table 2) in addition to a 2:4 comparison on BERT models (see Table 3).
We highlight that ExactOBS matches or even slightly exceeds the 4:8 results of AdaPrune withthe considerably more stringent 2:4 pattern, which is already well supported on NVIDIA hardware.
Furthermore, in a 2:4 comparison on BERT models, ExactOBS achieves 1 – 2% higher F1 scores.
[Table] [Table]
Quantization.
Further, we compare OBQ’s independent performance (after batchnorm tuning)
withthe state-of-the-art sequential post-training methods AdaQuant [19], AdaRound [31] and BRECQ [24].
We perform standard asymmetric per-channel quantization of all weights, using the authors’ implementations.
We rerun all methods on Torchvision [28] ResNets to ensure a uniform baseline.
The quantization grids for OBQ as well as AdaRound are determined with the same LAPQ [34]procedure that is used by BRECQ.
Surprisingly, we find that, despite optimizing layers independently,
OBQ achieves very similar (sometimes even slightly better) accuracies as existing non-independentmethods for 4 and 3 bits.
This suggests that it should be well-suited for mixed precision applicationswhere one needs to quickly generate many non-uniform models optimized for different constraints.(However, we note that ExactOBS can also be applied sequentially; see Appendix.)
[Table] [Table]
[Figure]
[Table]
[Figure]
BOP-Constrained Mixed GPU Compression.
We now consider a practical setting where we aregiven a trained model together with some calibration data and want to compress this model forefficient inference on an NVIDIA GPU which supports 8-bit and 4-bit arithmetic, also in combinationwith 2:4 sparsity.
Thus, there are 4 possible compression choices per layer: 8bit weights + 8bitactivations (8w8a), 4w4a, 8w8a + 2:4 and 4w4a + 2:4.
Unlike in the previous section, we do symmetricper-channel quantization of the weights as it has better hardware support; activations are quantizedasymmetrically per-tensor.
We then generate mixed precision configurations for various BOP (numberof bits times FLOPs) reduction targets and visualize the resulting compression-accuracy trade-offcurves in Figure 2.
In summary, at the cost of a ≈ 2 . 5% relative performance drop, we can achievea 12 − 14 × BOP reduction for ResNets and a 7 − 8 × reduction for the more challenging YOLOand BERT models (relative to the compute in compressible layers).
To the best of our knowledge,we are the first to consider joint N:M pruning and quantization in a post-training setting.
Recentwork [3] also studies joint 4w4a + 2:4 compression for ResNet18 but with 90 epochs of (sparse)
Quantization-Aware Training (QAT) on the full dataset and report 67.33% accuracy.
Althoughnot perfectly comparable (we keep the first layer dense and their dense baseline has 0 .
94% higheraccuracy and uses 4:8 sparse activations), we achieve similar 67.20% accuracy for 4w4a
+ 2:4 posttraining , which emphasizes the effectiveness of our methods for joint sparsification and quantization.
Time-Constrained CPU Compression.
Lastly, we explore a similar scenario, but targeting actual
CPU inference speedup on a 12-core Intel Xeon Silver 4214 CPU using the DeepSparse inferenceengine [35, 22], which provides acceleration for joint 8-bit quantization and block-sparsity withblocksize 4.
In this case, we work with real layer-wise timing data (for batchsize 64), as in [9].
Thereare 30 available block-sparsity targets per-layer, in steps of pruning 10% of the remaining weights,all of which are further quantized to 8 bits.
The base acceleration of the dense 8 bit model is ≈ 2 .
7 ×on top of which sparsity speedup acts roughly multiplicatively.
Figure 2d shows results for ResNet50and several (real-time) speedup targets—we achieve 4 × and 5 × (actual) speedup with 1% and 2%accuracy loss, respectively.
These are the first full post-training results in this setting (the authorsof [10] only performed 4-block pruning post-training, followed by 5 epochs of QAT on the entire
ImageNet dataset), and they show very encouraging accuracy-speedup trade-offs.
Conclusions & Future Work
We have presented a new efficient and accurate approach for solving the layer-wise compressionproblem, and built on it to obtain state-of-the-art post-training compression solutions for both pruningand quantization.
Our framework should be naturally extensible to structured pruning, which infact should allow for further optimizations, and should also be compatible with further compressionvia unstructured pruning and quantization.
Our results suggest that post-training compression maybe able to reach comparable accuracies to much more expensive retraining methods.
We plan toinvestigate this in future work, in particular in the context of more resource-intensive models, such asvery large-scale language models.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
Assets are open-source with standard permissive licenses.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]