SlowGuess/NeurIPS_2023_OCR huggingface dataset:
IST Austria & Neural Magic dan. alistarh@ist.ac.at
The breakthrough performance of large language models (LLMs) comes with major computational footprints and high deployment costs. In this paper, we progress towards resolving this problem by proposing a novel structured compression approach for LLMs, called ZipLM. ZipLM achieves state-of-the-art accuracy-vs-speedup, while matching a set of desired target runtime speedups in any given inference environment. Specifically, given a model, a dataset, an inference environment, as well as a set of speedup targets, ZipLM iteratively identifies and removes components with the worst loss-routine trade-off. Unlike prior methods that specialize in either the post-training/one-shot or the gradual compression setting, and only for specific families of models such as BERT (encoder) or GPT (decoder), ZipLM produces state-of-the-art compressed models across all these settings. Furthermore, ZipLM achieves superior results for a fraction of the computational cost relative to prior distillation and pruning techniques, making it a cost-effective approach for generating an entire family of smaller, faster, and highly accurate models, guaranteed to meet the desired inference specifications. In particular, ZipLM outperforms all prior  $\mathrm{BERT}_{\mathrm{base}}$  distillation and pruning techniques, such as CoFi, MiniLM, and TinyBERT. Moreover, it matches the performance of the heavily optimized MobileBERT model, obtained via extensive architecture search, by simply pruning the baseline  $\mathrm{BERT}_{\mathrm{large}}$  model. When compressing GPT2, ZipLM outperforms DistilGPT2 while being  $60\%$  smaller and  $30\%$  faster. Our code is available at: https://github.com/IST-DASLab/ZipLM.
The high accuracy of modern language models from the Transformer family [53] comes at the price of massive computational cost, which hinders their practical adoption in resource-constrained settings. This has motivated the development of model compression techniques, which can be categorized into pruning [17], quantization [9], and distillation [11]. In this paper, we focus on structural compression, whose goal is to reduce model size by removing entire sub-components, such as rows or columns from the model's weight matrices. The key advantage of structured pruning, relative to unstructured pruning of individual weights, is that the model can be reshaped to new dimensions, and the resulting computational savings can be leveraged on any hardware, without specialized computational support. At the same time, structured pruning introduces significant challenges. First, models are usually highly-sensitive to structured compression, and most methods require gradual compression, including retraining cycles designed to allow the model to recover accuracy. In addition, structural compression significantly complicates the use of knowledge distillation [16], which is usually done via manual or dynamic layer mapping [21, 59]. On the practical side, another challenge is that most existing techniques do not provide runtime speedup guarantees: the model is pruned to a fixed sparsity or FLOPS target, and then must be evaluated in the target inference environment. If the pruned model fails to meet the target inference specifications, the whole process must be repeated from scratch.
Overview. In this paper, we resolve these issues and provide a novel structured pruning approach called ZipLM, which achieves state-of-the-art performance, both in the post-training/one-shot setting, where retraining is not desirable, as well as in the popular gradual compression setting, where retraining is possible. We accomplish this via an inference-aware algorithm, which successfully balances the loss-routine trade-off at each pruning step. By taking runtime into account, we avoid removing components that do not bring significant speedup gains. Additionally, our algorithm provides speedup guarantees for compressed models, a highly-desirable property in practical applications.
- We introduce a novel structured pruning approach, which unifies the saliency criteria investigated by prior work-weight magnitude, activation impact, and removal of linearly-redundant structures, while considering local (layer-wise) and global correlations. We augment it to be inference-aware, ensuring desired latency or throughput in any given configuration.
- We complement the algorithm with a novel layer-wise token-level distillation, which consistently boosts accuracy on small datasets and does not require manual layer matching, circumventing a limitation of prior structured pruning techniques.
- ZipLM is the first structured pruning approach that achieves state-of-the-art results for both, post-training/one-shot compression and gradual pruning settings, while being applicable to both, BERT (encoder) and GPT (decoder) language models, without any modifications.
- ZipLM is practical and efficient. For a set of desired speedups (e.g. 2x, 5x, 10x) in the target inference environment (e.g. batch-size=128, sequence-length=384, device=V100), in a single run and under the same set of hyper-parameters, it produces the entire family of compressed models, one for each speedup target. Consequently, it leads to state-of-the-art results in GPU-based inference environments. Moreover, it is compatible with unstructured pruning and quantization, leading to state-of-the-art results even for CPU-based environments.
Distillation-based compression methods focus on training a smaller student model to mimic the representations of a larger teacher model. The "distance" between the representations of student and teacher is often architecture-specific. MiniLM [55] uses a deep self-attention mechanism to replicate the attention mechanism of the teacher, and TinyBERT [21] employs a bespoke distillation mechanism, for a manually-picked subset of layers. Both methods offer a very strong baseline, generally outperforming other approaches, except for MobileBERT. MobileBERT [51] involves first training a custom large BERT teacher model from scratch, and then deviates from the standard architecture [2] by introducing heavily-optimized components with reduced latency, whose combinations are decided in neural architecture search (NAS)-like fashion. It achieves strong results in terms of accuracy-per-parameter, at the cost of significant computational costs in the search process. DistilBERT and DistilGPT2 [44] involve training a fixed student obtained by removing every other layer from the teacher, while BERT-PKD [50] employs incremental knowledge extraction through the distillation of intermediate layers. Well-Read-Students [52] reduces the size of the standard BERT architecture through principled downscaling of internal dimensions. DynaBERT [18], on the other hand, distills knowledge to a student model that is both depth- and width-adaptive.
Structural pruning methods usually start from a large pre-trained model, and iteratively reduce the dimensions of weight matrices. Block Movement Pruning [27] identifies and removes redundant rectangular blocks of weights while following the movement pruning intuition [45] that weights moving towards zero during fine-tuning should be removed. FLOP [56] and Low-Rank [40] use matrix decomposition techniques to progressively remove rank-1 components from factorized weight matrices during training. BERT-of-Theseus [60] employs a similar approach, but replaces entire submodules with smaller counterparts. Methods like LayerDrop [3] and Poor Man's BERT [43] address structured compression through various layer-dropping techniques. LayerDrop uses structured layer-dropout regularization to train a model resilient to sub-network selection during inference, while Poor Man's BERT explores a wide range of layer-dropping strategies. The recent CoFi method [59] employs masks of different granularities to jointly prune coarse and fine-grained submodules during fine-tuning, combined with an optional customized distillation technique. CoFi is the state-of-the-art structural pruning method; relative to distillation methods, CoFi outperforms MiniLM and TinyBERT, but not MobileBERT, in terms of accuracy-vs-speedup.
Other compression methods such as the ones that exploit dynamic forms of sparsity which appear at runtime [35], or the ones that utilize lower bit-width representation of weights and/or activations [6, 32] are complementary to our approach. We demonstrate this in Section 5 where we apply quantization to obtain even higher compression ratios for edge deployment environments like commodity CPUs.
Removing large structures like entire matrix columns or attention heads from a language model quickly leads to severe accuracy degradation, from which it is often difficult to recover even with extensive finetuning. This is why current state-of-the-art approaches like Block Movement Pruning [27] or CoFi [59] opt for integrating pruning directly into training (via sampling or differentiable approximations), rather than performing it in the standard gradual pruning fashion of discrete steps with finetuning in between. However, as we will show, by designing a new highly accurate pruning algorithm which is able to account for both local correlations of structures within single layers as well as global correlations across layers, we can actually apply the gradual pruning paradigm, with all its advantages, to improve significantly over the current state-of-the-art.
3.1 The ZipLM Structured Pruning Algorithm (Local Correlations)
Most existing structured pruning criteria [33, 31] are based on one or two of the following assumptions about saliency: structures with lower (average) weight magnitude are easier to prune [15, 29], structures with small input activations can be removed at little loss [34], and structures that are close to a linear combination of other structures are the most redundant [14, 49]. We will now show how all these aspects can be jointly considered in a principled manner via our new ZipLM technique.
Problem formulation. Our approach starts from the idea of applying structured compression layerwise, in a way that allows the layer to preserve most of its output characteristics. This setup is popular in the post-training quantization and unstructured pruning literature [38, 19, 8], and can be implemented as follows. We are given a small amount of calibration data, which we run through the network, to obtain "reference" inputs and outputs for each layer. Then, for each layer, given the calibration inputs  $\mathbf{X}$  and the original layer weights  $\mathbf{W}$ , we aim to find compressed weights  $\widehat{\mathbf{W}}$  respecting the compression constraint  $\mathcal{C}$ , which best approximate the original output, measured via the squared error metric. If we assume that the input and weight matrices have an appropriate rectangular form, the problem can be formalized as:
$$
\operatorname {a r g m i n} _ {\widehat {\mathbf {W}}} \left\| \widehat {\mathbf {W}} \mathbf {X} - \mathbf {W} \mathbf {X} \right\| _ {2} ^ {2} \quad \text {s u b j e c t t o} \quad \widehat {\mathbf {W}} \in \mathcal {C}. \tag {1}
$$
This objective can be decomposed across the rows of  $\mathbf{W}$ , leading to a set of sparse linear regression problems, one per row. These row-wise problems are independent, which forms the basis of related work [8]; yet, since we do structured pruning, they become dependent, as we would like to prune the same weight indices across all rows, i.e. prune entire columns. Thus, finding the optimal weights  $\widehat{\mathbf{W}} \in \mathcal{C}$  is equivalent to finding: 1) the optimal structure  $\mathbf{S}$  of the desired shape to be removed, which we assume to be applied across all rows, with corresponding pruning mask  $\mathbf{M}_{\mathbf{S}}$ , where pruned indices have value 1 in the mask, and others are 0; and 2) the corresponding update  $\delta_{S}$  to all of the remaining weights, optimally compensating for the error caused by the removal of weights in  $\mathbf{S}$ .
Saliency scores and weight update. Let  $\mathbf{H} = \mathbf{X}\mathbf{X}^{\top}$  be the Hessian matrix for the  $\ell_2$ -minimization problem in Equation 1, which is independent of the weights. Define  $\mathbf{W}_{i,\mathbf{M_S}}$  to be the subset of weights under the mask  $\mathbf{M_S}$  in row  $i$ , and by  $(\mathbf{H}^{-1})_{\mathbf{M_S},\mathbf{M_S}}$  the submatrix of the inverse Hessian corresponding to the entries under the mask  $\mathbf{M_S}$ . Then, we can obtain the optimal mask and weight update as follows:
$$
\operatorname {a r g m i n} _ {\mathbf {S}} \sum_ {i = 0} ^ {d _ {\text {r o w}}} \mathbf {W} _ {i, \mathbf {M} _ {\mathbf {S}}} \cdot \left(\left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} _ {\mathbf {S}}, \mathbf {M} _ {\mathbf {S}}}\right) ^ {- 1} \cdot \mathbf {W} _ {i, \mathbf {M} _ {\mathbf {S}}} ^ {\top} \tag {2}
$$
$$
\boldsymbol {\delta} _ {S} = - \mathbf {W} _ {:, \mathbf {M} \mathbf {s}} \cdot \left(\left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} \mathbf {s}, \mathbf {M} \mathbf {s}}\right) ^ {- 1} \cdot \left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} \mathbf {s},:} \tag {3}
$$
We obtain this by extending the Optimal Brain Surgeon [13, 24] formulas for solving Equation 1 to cover all  $d_{\mathrm{row}}$  weight matrix rows simultaneously. Importantly, the subselection of the inverse
Hessian  $((\mathbf{H}^{-1})_{\mathbf{M}_{\mathbf{S}},\mathbf{M}_{\mathbf{S}}})^{-1}$  is shared between all rows. Further, since we generally consider only non-overlapping sets  $\mathbf{S}$  of the same size, we pay just  $O(d_{\mathrm{col}}\cdot |\mathbf{M}_{\mathbf{S}}|^{2})$  total cost for all extra inversions. Since the number of structures in the mask  $|\mathbf{M}_{\mathbf{S}}|$  is usually small, e.g. attention heads usually consist of 64 columns, the overall cost of these inversions is low.
Simply selecting the structures to prune according to the criterion in Equation 2 unifies the weight magnitude and activation influence criteria (via the Hessian), but still ignores any correlations between structures. We address this by pruning structures one-at-a-time, while always applying update  $\delta_{S}$  and fully recomputing  $\mathbf{H}^{-1}$  relative to the remaining structures. For example, if there exist two redundant structures  $S_{1}$  and  $S_{2}$ , we will first drop  $S_{1}$  and update  $S_{2}$  to compensate for this removal, at which point  $S_{2}$  is no longer easy to prune. Without this one-at-a-time removal, both structures would have been incorrectly removed as they each individually seem easy to prune according to Equation 2. Executing this strategy naively will require a full  $O(d_{\mathrm{col}}^3)$  recomputation of the inverse Hessian relative to the remaining structures at each step, which would be very slow. However, this can be avoided by removing the rows and columns corresponding to  $\mathbf{M}_{\mathbf{S}}$  directly in the inverse with one step of Gaussian elimination [8], applied block-wise to cover larger structures, as follows:
$$
\mathbf {H} ^ {- 1} - \mathbf {H} _ {:, \mathbf {M} _ {\mathbf {S}}} ^ {- 1} \cdot \left(\left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} _ {\mathbf {S}}, \mathbf {M} _ {\mathbf {S}}}\right) ^ {- 1} \cdot \mathbf {H} _ {\mathbf {M} _ {\mathbf {S}},:} ^ {- 1}, \tag {4}
$$
which takes only  $O(|\mathbf{M}_{\mathbf{S}}| \cdot d_{\mathrm{col}}^2)$  time. We provide complete pseudocode in Algorithm 1.
Algorithm 1 The ZipLM pruning algorithm. Given inverse Hessian  $\mathbf{H}^{-1} = (2\mathbf{X}\mathbf{X}^{\top} + \lambda \mathbf{I})^{-1}$ , we remove exactly  $k$  structures from the corresponding weight matrix  $\mathbf{W}$ .
$\mathbf{R}\gets$  set of all possible structures
for  $k$  times do
$$
\mathbf {S} \leftarrow \operatorname {a r g m i n} _ {\mathbf {S}} \sum_ {i = 0} ^ {d _ {\mathrm {r o w}}} \mathbf {W} _ {i, \mathbf {M} _ {\mathbf {S}}} \cdot \left(\left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} _ {\mathbf {S}}, \mathbf {M} _ {\mathbf {S}}}\right) ^ {- 1} \cdot \mathbf {W} _ {i, \mathbf {M} _ {\mathbf {S}}} ^ {\top}
$$
$$
\boldsymbol {\delta} _ {S} \leftarrow - \mathbf {W} _ {:, \mathbf {M} _ {S}} \cdot \left(\left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} _ {S}, \mathbf {M} _ {S}}\right) ^ {- 1} \cdot \left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} _ {S};}
$$
$$
\mathbf {W} \leftarrow \mathbf {W} + \delta_ {S}
$$
$$
\mathbf {H} ^ {- 1} \leftarrow \mathbf {H} ^ {- 1} - \mathbf {H} _ {:, \mathbf {M} _ {\mathbf {S}}} ^ {- 1} \cdot \left(\left(\mathbf {H} ^ {- 1}\right) _ {\mathbf {M} _ {\mathbf {S}}, \mathbf {M} _ {\mathbf {S}}}\right) ^ {- 1} \cdot \mathbf {H} _ {\mathbf {M} _ {\mathbf {S}},:} ^ {- 1}
$$
$$
\mathbf {R} \leftarrow \mathbf {R} - \{\mathbf {S} \}
$$
$\mathbf{W}\leftarrow \mathbf{W}\odot \mathbf{M}_{\mathbf{R}}$
We utilize the fact that the values corresponding to pruned weights in  $\mathbf{W}$  and in the inverse Hessian  $\mathbf{H}^{-1}$  do not affect any subsequent calculations and can therefore be ignored even if they are not exactly zero. However, in the end we have to prune them explicitly again by multiplying with the overall mask to ensure that they are exactly zero. In a practical implementation,  $(\mathrm{(H^{-1})_{Ms,M_s})^{-1}}$  should only be computed once and reused when computing the corresponding sum across all rows.
Pruned structures. Focusing on Transformers, we consider three types of structural removal: dropping attention heads, shrinking the expanded intermediate dimension of the fully-connected network (FC) layers, and removing entire residual parts, i.e. attention or FC-modules. We implement this by dropping  $d_{\mathrm{head}}$  consecutive columns in the out-matrix of the attention block and individual columns in the second linear layer of the feed-forward network. Once these column-structures are zeroed out, corresponding rows in previous layers can be safely removed without any output change. Crucially, by pruning e.g. columns in the FC2 layer rather than equivalent rows in FC1, we can utilize the input correlations via Hessian-information using the ZipLM pruner.
Novelty relative to existing Optimal Brain Surgeon (OBS) approaches. The original framework [13], as well as modern efficient versions [47, 7, 8], have been explicitly developed for unstructured pruning, i.e. removing individual weights. It is nontrivial to extend them to structured pruning, as this involves considering additional correlations, both within as well as across multiple blocks (such blocks are usually employed for computational tractability). For example, the state-of-the-art layer-wise approach of [8], performs unstructured pruning by handling weight matrix rows separately, and then greedily merging results. In contrast, we perform structured pruning jointly across multiple rows, which is not only necessary for correctness but additionally enables us to design an algorithm with a computational complexity that is lower by a full factor of the hidden dimension size. Additionally, structured pruning requires explicitly matching matrix shapes for consecutive layers and a dedicated strategy for utilizing weight updates even when entire blocks/rows are pruned.
3.2 Inference-Aware Structured Pruning (Global Correlations)
We now describe how to augment the algorithm to be inference-aware, in the sense that it accepts inference specifications, such as batch-size, sequence-length, and speedup on the target hardware, as additional inputs to optimize for.
Motivation. The main benefit of inference-aware structured pruning is the fact that pruning decisions are not guided purely by saliency scores, but instead by loss-vs-speedup trade-offs associated with the removal of each component in the model. Prior methods, e.g. [24, 27, 59] focus solely on pruning until a specific sparsity threshold is reached, without taking into account the real-world speedups corresponding to the compression threshold, which can vary significantly between settings. For example, a  $95\%$  sparse BERT produced by CoFi [59] has  $12x$  speedup on a V100 GPU, but only  $5x$  on an A100 GPU. With existing methods, if real-world timings fail to meet the inference requirements, the entire process has to be repeated with different sparsity values until the target speedup is achieved, which is both time-consuming and error-prone. An additional advantage of inference-awareness, which we showcase in our GPT experiments in Section 4, is that it enables optimizing for different real-world metrics, such as latency or throughput.
Runtime awareness. We integrate runtime constraints via a latency table [1] for our target inference environment, where we record the time to run an attention block, including all overheads, with  $0, \dots, N_{\mathrm{heads}} - 1$  heads pruned and similarly for the fully-connected block with the intermediate dimension shrunk by a factor of  $0.9^i$ , for  $i = 0, \dots, 42$ ; in relative steps of  $10\%$  up until  $\approx 99\%$  sparsity, following [5]. This allows rapid runtime estimation for different per-layer sparsity configurations. We provide an example of our latency table in Appendix E.
Finding the optimal sparsity configuration. Ultimately, our goal is to find a per-layer-sparsity configuration that satisfies a certain speedup-constraint while maximizing accuracy. A popular paradigm of doing this [15, 30] is to produce a large number of pruned models with different sparsity distributions across layers and then select the one, satisfying a target constraint, with the highest accuracy. To make this computationally feasible, it is crucial that pruning is cheap, yet accurate. ZipLM treats each layer independently, which makes it possible to precompute a database of several pruned versions with different sparsities for each layer. The entire database can be produced in a single run, utilizing the algorithm's one-at-a-time nature. While our algorithm is compatible with various search methods for finding layer-wise profiles [15, 12], we adapt the recent S&PY approach [5].
Structured SPY search. The SPY approach is designed for unstructured pruning and assigns a quadratic prior to per-layer sensitivity of different sparsity levels. This is not valid in our structured pruning scenario, since for instance it would suggest that dropping a full layer is only slightly more difficult than pruning it to  $99\%$  sparsity. Thus, using standard SPY would lead the algorithm to explore a large number of sub-optimal configurations, significantly wasting computational resources. To alleviate this problem, for a structured sparsity  $s$ , we introduce a better prior  $p_s$  as the relative layer-wise squared error incurred by pruning, defined as  $p_s = ||\widehat{\mathbf{W}}_s\mathbf{X} - \mathbf{W}\mathbf{X}||_2 / ||\mathbf{W}\mathbf{X}||_2$ , which simply has a value of 1 for a fully dropped layer. Furthermore, the original SPY approach uses shrinking neighborhood search, which has high variance in both runtime and solution quality for structured compression. Therefore, we perform a fixed number of 1000 steps, randomly mutating in expectation  $10\%$  of the layer-wise sensitivity coefficients. Finally, we note that any candidate evaluated by this procedure actually achieves the target speedup, leading to significantly decreased search time. We validate our approach in Appendix F, where we demonstrate that our speedup estimations are indeed very accurate in practice. Specifically, real-world on-device measurements deviate at most by  $5.28\%$  from their expected values.
3.3 Layer-wise Token Distillation
For structured pruning, it is common to apply layer-wise distillation objectives to transfer intermediate representations. However, structured pruning creates compatibility issues relative to the fixed teacher architecture, leading most methods to develop customized distillation strategies. A popular approach, introduced in [21] and improved by [59], solves the problem via static [21] or dynamic [59] mapping of a subset of teacher layers to a subset of student layers. Their main limitation is manual layer selection, where making the "optimal" choice would require evaluating all possible combinations, which can be very expensive. Another limitation is shape-matching between intermediate layers, which is solved by introducing a learnable linear transformation matrix attached to student outputs.