Title: ZipLM: Inference-Aware Structured Pruning of Language Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: LLMs

pruning

compression

inference
Abstract: 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-runtime 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 BERT-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 BERT-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.

ZipLM: Inference-Aware Structured Pruning of Language Models

Abstract

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-runtime 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 BERTbase 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 BERTlarge 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.

1 Introduction

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.

Dan Alistarh

IST Austria & Neural Magic

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-runtime 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 summarize our contributions as follows:

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.

2 Related Work

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.

3 Method

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 X and the original layer weights W , we aim to find compressed weights ̂ W respecting the compression constraint 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:

This objective can be decomposed across the rows of 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 ̂ W ∈ C is equivalent to finding: 1) the optimal structure S of the desired shape to be removed, which we assume to be applied across all rows, with corresponding pruning mask M S , where pruned indices have value 1 in the mask, and others are 0 ; and 2) the corresponding update δ S to all of the remaining weights, optimally compensating for the error caused by the removal of weights in S .

Saliency scores and weight update. Let H = XX ⊤ be the Hessian matrix for the ℓ 2 -minimization problem in Equation 1, which is independent of the weights. Define W i, M S to be the subset of weights under the mask M S in row i , and by ( H -1 ) M S , M S the submatrix of the inverse Hessian corresponding to the entries under the mask M S . Then, we can obtain the optimal mask and weight update as follows:

We obtain this by extending the Optimal Brain Surgeon [13, 24] formulas for solving Equation 1 to cover all d row weight matrix rows simultaneously. Importantly, the subselection of the inverse

Hessian (( H -1 ) M S , M S ) -1 is shared between all rows. Further, since we generally consider only non-overlapping sets S of the same size, we pay just O ( d col · | M S | 2 ) total cost for all extra inversions. Since the number of structures in the mask | M 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 δ S and fully recomputing 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 3 col ) 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 M S directly in the inverse with one step of Gaussian elimination [8], applied block-wise to cover larger structures, as follows:

which takes only O ( | M S | · d 2 col ) time. We provide complete pseudocode in Algorithm 1.

Algorithm 1 The ZipLM pruning algorithm. Given inverse Hessian H -1 = (2 XX ⊤ + λ I ) -1 , we remove exactly k structures from the corresponding weight matrix W .

We utilize the fact that the values corresponding to pruned weights in W and in the inverse Hessian 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, (( H -1 ) M S , 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 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-ofthe-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 , . . . , N 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 , . . . , 42 ; in relative steps of 10% up until ≈ 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 SPDY approach [5].

Structured SPDY search. The SPDY 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 SPDY 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 = || ̂ W s X -WX || 2 / || WX || 2 , which simply has a value of 1 for a fully dropped layer. Furthermore, the original SPDY 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.

Our approach. We address these challenges differently, by leveraging the fact that ZipLM preserves the hidden dimension size, and propose to use distillation of intermediate token representations across the entire model. The resulting minimization objective consists of three components:

where θ s and θ t represent student and teacher models respectively, x are the inputs, L task is the loss associated with the task (e.g. cross-entropy for text-classification), L logit is the KL-divergence between output logits as described in [16], and L token is our token-level distillation loss. Hidden tensors passed between consecutive transformer layers are of constant shape H ∈ R B × seq × H , where B stands for the batch-size, seq for the sequence length, and H for the hidden size defined by the model architecture. This tensor can be interpreted as a collection of B × seq vectors h ∈ R H , each carrying intermediate model representations of input tokens x . We define the loss L token as an Euclidean distance ∆ between vectors h corresponding to each non-padded token in the input sequence, averaged over all unpruned layers. Formally, for a layer k , it is defined as

✶ where P stands for the set of padding tokens. This formulation encourages the student model to generate vector representations for each token that are similar to those produced by the teacher model. In Appendix B, we present ablation studies and comparisons for ZipLM and CoFi, with and without their respective distillation objectives.

4 Experiments

Setup. Given a pre-trained model, a dataset, and a set of desired speedups in a target inference environment, we iteratively fine-tune and prune the model in a structured way such that in the end we obtain a set of accurate compressed models, one for each speedup target. We consider pruning of the standard BERTbase and BERTlarge architectures, evaluating on dev-sets of established benchmarks: SQuADv1.1 [42], and a subset of GLUE [54] tasks: SST-2 [48], QNLI [54], MNLI [57], and QQP [46], selected to match publicly-available checkpoints from prior work. For a precise comparison to prior work [59], our inference environment is a single NVIDIA V100 16GB GPU, batch size of 128, and sequence lengths of 384 and 128 for SQuAD and GLUE tasks, respectively. In addition to encoder-based BERT models, we also consider pruning of the decoder-based GPT2 model on the OpenWebTextCorpus [10], for which we consider two inference environments: pruning for throughput (batch-size=16, sequence-length=1024), and pruning for latency (batch-size=1, a set of prompts with varying lengths). For illustration, our pipeline is depicted in Figure 1. In Appendix H and I, we report exact values for all results, as well as hyper-parameters for reproducibility.

#heads

Concat

Linear

Scaled Dot-Product Attention

Linear

Linear

Linear

Q

K

V

Concat

Linear

Scaled Dot-Product Attention

Linear

Linear

Linear

Q

K

V

#heads - N

Scaled Dot-Product Attention

Linear

Linear

Linear

1

Concat

Linear

Q

K

V

On-Device Multi-Head Attention timing

h

h inter

x

h inter  - N

1

x

x

On-Device Feed-Forward Network timing

model: BERT-large

task: SQuADv1.1

batch-size: 128

seq-length: 384

device: NVIDIA V100 16GB GPU

speedups: {2x, 3x, ..., 49x, 50x}

Fine-tune

BERT-large

ZipLM prune

to 2x speedup

Steps

Learning rate

ZipLM prune

to 3x speedup

ZipLM prune

to 50x speedup

Fine-tune

BERT-large 2x

Fine-tune

BERT-large 3x

Fine-tune

BERT-large 50x

ZipLM pruning and fine-tuning

Inference environment specifications

Baselines. In the gradual pruning setting, we explore the performance of ZipLM pruning of BERTand GPT2-family models, across a wide range of inference speedup targets, ranging from 2x to 15x, in unit increments. This allows us to compare the effectiveness of our approach against a diverse set of structured pruning and distillation-based techniques, including state-of-the-art CoFi pruning, competitive Block Movement Pruning, and distillation approaches including TinyBERT, DistilBERT,

DistilGPT2, MobileBERT, MiniLM, and DynaBERT. Additionally, we include comparisons with other relevant methods. For fairness, we follow [59] and report TinyBERT and DynaBERT results without data augmentations. In the post-training/one-shot setting, which does not allow retraining, we demonstrate that ZipLM outperforms the prior state-of-the-art approach of [26]. We evaluate inference speedups of all models in the same environment, unless the models are not publicly available, in which case we report speedups from their respective papers. We refer to ZipLM compressed BERT models as ZipBERT, and to ZipLM compressed GPT2 models as ZipGPT2.

4.1 Gradual Structured Pruning

BERTbase results. In Figure 2 we compare structured compression methods on the SQuADv1.1 task. ZipLM outperforms both CoFi and TinyBERT, prior state-of-the-art techniques, by 3 points in the F1 score at the same speedup factor, while at the same F1 score it is able to improve inference speedups by at least 60%. In Figure 3, we extend this comparison to a subset of GLUE tasks and provide an exhaustive overview of various structured compression techniques. Results on the other four remaining GLUE tasks are provided in Appendix Figure 7. As can be observed, distillation-based methods usually provide either one or a few structurally-compressed models, due to the massive costs associated with training from scratch for each new model. Relative to the most competitive approaches, such as TinyBERT, CoFi, and MiniLM, ZipLM provides consistent improvements in terms of both, accuracy and speedup, while providing guarantees for each compressed model in terms of the expected speedup in the target inference environment. Interestingly, on tasks like QQP and SST-2, ZipLM is able to compress the BERTbase model up to 6x and 10x speedups, respectively, while maintaining the accuracy of the dense model. In Appendix D, we provide additional comparisons against CoFi on test-set results from the official GLUE evaluation server.

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Speedup factor

82

84

86

88

F1 score (%)

BERTbase

99% BERT base

SQuADv1.1

ZipBERT base

(ours)

CoFi

Block Movement Pruning

TinyBERT

DistilBERT

FLOP-RoBERTa

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Speedup factor

88

89

90

91

F1 score (%)

BERTlarge

99% BERT large

SQuADv1.1

ZipBERT large

(ours)

Block Movement Pruning

MobileBERT

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

86

87

88

89

90

91

92

Accuracy (%)

BERTbase

99% BERT base

QNLI

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

79

80

81

82

83

84

85

Accuracy (%)

BERTbase

99% BERT base

MNLI

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

89.5

90.0

90.5

91.0

91.5

92.0

92.5

93.0

93.5

Accuracy (%)

BERTbase

99% BERT base

SST-2

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

90.0

90.2

90.4

90.6

90.8

91.0

91.2

Accuracy (%)

BERTbase

99% BERT base

QQP

ZipBERT base

(ours)

CoFi

Block Movement Pruning

DynaBERT

DistilBERT

TinyBERT

MiniLM

Low-Rank BERT

LayerDrop

BERT-PKD

Well-Read Students

Poor Man's BERT

BERT-of-Theseus

FLOP-RoBERTa

BERTlarge results. To verify that our approach does not pertain only to the BERTbase model, we apply ZipLM structured pruning to the 3x larger BERTlarge model on the SQuADv1 task. In this setup, we compare against the only two approaches that attempted to structurally compress this larger model, Block Movement Pruning and distillation-based MobileBERT. As can be seen in Figure 2, ZipLM is able to compress BERTlarge up to 4x faster inference while maintaining the F1 score of the uncompressed model. At the same F1 score as the fastest Block Movement Pruning model (3x), ZipLM doubles the inference speedup (6x). A result worth emphasizing is that ZipLM is even able to match the performance of the highly optimized MobileBERT model by simply compressing the baseline BERT architecture, without the many additional optimizations and custom components

used by MobileBERT. Specifically, some of the module- and operator-level optimizations used by MobileBERT include: bottleneck structures and carefully-balanced self-attention and feed-forward modules, embedding layer factorization, a bespoke closed-source teacher model, replacement of LayerNorm layers with lower-latency NoNorm layers, and replacement of GELU activation functions with ReLU activations.

99% recovery. The MLPerf Benchmark [36] targets recovery of > 99% of the baseline accuracy. At this industry-defined threshold, ZipLM models set new state-of-the-art performance across all of the considered datasets with the following BERTbase inference speedups: 5x on the SQuADv1 task, 6x on QNLI and MNLI, and, surprisingly, 13x and 15x on SST-2 and QQP, respectively. When compressing BERTlarge on the SQuADv1 task, ZipLM produces a 6x faster model at 99% recovery.

GPT2 results. To validate that our approach does not only apply to encoder-based models, we apply ZipLM structured pruning to the decoder-based GPT2 model. In addition to this, to further demonstrate the inference-awareness property of our approach and its importance for real-world applications, we consider two different regimes: pruning for throughput and pruning for latency. An example application for the former regime is a server-side deployment where the model processes many queries at the same time, while an application for the latter regime is a text-generation scenario where the model is used in an online fashion to auto-complete user's text.

For a fair comparison, we follow the DistilGPT2 setup [44] and prune the 124M parameters GPT2 variant on the OpenWebTextCorpus dataset, followed by zero-shot evaluations, without any finetuning, on the test-split of the WikiText [37] dataset. Because of the enormous vocabulary size, the maximum achievable speedup in the throughput regime for this model is roughly 3.5x. Thus, we run ZipLM pruning to 1.5x, 2x, 2.5x, and 3x speedup targets. For the latency regime, we report the median time to process sequences of various lengths when generating text with Top-K sampling [4]. In Table 1, we present zero-shot evaluations of the uncompressed GPT2 model which serves as a baseline relative to the competing DistilGPT2 approach, and four variants of our ZipLM pruned GPT2. In the pruning for throughput scenario, at similar speedup and decoder size (1.6x-vs-1.5x and 42.5M-vs-47.3M), ZipGPT2 achieves significantly lower perplexities relative to DistilGPT2. Further, at slightly better (lower) perplexities, ZipGPT2 reduces the decoder size from 42.5M to only 26.5M parameters (60% reduction) and improves speedup from 1.6x to 2.1x (30% faster). In the pruning for latency scenario, at a similar speedup of 1.9x-vs-2.0x, ZipGPT2 reduces the decoder size by 3M params while providing almost 2 points improvement in the zero-shot perplexity.

Table 2: One-shot (posttraining) structured pruning of BERTbase on three downstream datasets and two speedup targets.

4.2 On the Importance of Inference-Awareness

Depth vs. width pruning. A particularly interesting illustration of the importance of inferenceawareness in the pruning algorithm is given by our GPT2 models running directly in the PyTorch-HuggingFace framework, which can be used in two different modes: batch-prediction (throughput-constrained) and text-generation (latency-constrained). For the former, inputs are typically large, and shrinking weight matrices is an effective way to achieve speedups. However, for the latter, the inputs are much smaller, and the size of weight matrices is no longer the primary bottleneck.

In this scenario, the only way to achieve substantial speedups is to completely drop some modules, which prior methods cannot account for as they solely optimize for overall model sparsity. However, with ZipLM, runtime measurements from the target inference environment guide pruning decisions, allowing it to learn the best way to compress the model for an optimal speedup-accuracy trade-off. Our GPT2 compression results in Table 1 clearly illustrate and support these statements. Even though pruned for the same speedup target, the final architectures of ZipGPT2 models are drastically different. For the throughput-constrained scenario, the model's depth was preserved but the matrix dimensions were significantly reduced (roughly by a factor of 10) making the corresponding multiplications with large input tensors much faster. In contrast, for the latency-constrained scenario, the model's width (shapes of weight matrices) was mostly preserved but the depth was shrunk almost by a factor of 4, making the forward pass with small inputs faster by reducing the effective number of modules.

Inference device capabilities. Incorporating capabilities of the inference device is another important aspect for effective structured pruning which prior methods do not account for as they solely optimize for higher sparsities. As noted in Section 3.2, this reflects in larges discrepancies between speedups obtained on different devices, e.g. a compressed model with 12x speedup on a V100 is only 5x faster on an A100 GPU. This arises because the A100 GPU is significantly more powerful and thus faster on the dense model; at the same time, it is highly underutilized for small matrices, which significantly limits the speedups for very high sparsity. To illustrate this, we have measured the speedup from reducing the MLP size for both GPU types (see Table 3). As can be seen, pruning to ≈ 90% sparsity (3072 → 302) gives ≈ 7x speedup on a V100 but only ≈ 3x speedup on an A100. Such differences are automatically captured by ZipLM, where pruning for sparsity is replaced by pruning for speedup.

Pruning for speedup vs. pruning for sparsity. In Figure 4 we compare results with ZipLM pruning when the target for pruning is sparsity (like prior approaches) and when the target for pruning is speedup (the ZipLM approach). Pruning for speedup brings significant improvements, up to 10 points, especially at higher speedups where inference-awareness is very important as the algorithm does not remove components that do not bring any further speed and therefore helps preserving accuracy.

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Speedup factor

75

80

85

F1 score (%)

BERTbase

SQuADv1.1

Pruning for sparsity

Pruning for speedup

4.3 Post-training/One-shot Structured Pruning

We now study the performance of ZipLM when applied purely in one-shot , without any retraining. In this setting, we compare against the state-of-the-art method of Kwon et al. [26] which combines several heuristics: Fisher-based mask search, mask rearrangement, and mask tuning. Instead of heuristics, our pruning framework utilizes direct end-to-end loss information to find the optimal sparsity configuration. During the warm-start phase, [26] utilizes a diagonal Fisher matrix to estimate the significance of heads and filters, which discards correlations caused by off-diagonal elements. Although the approach attempts to address this limitation by approximating correlations within a single layer, it will not capture global dependencies. Furthermore, the weights are adapted for layerwise reconstruction at the very end of the compression step, whereas our method does it continuously during the pruning (please see Section 4 for the significance of doing this). For a fair comparison, we apply the authors' own implementation in latency-constrained mode on the exact same model weights. Table 2 presents results on several datasets and speedups, showing that ZipLM is even more accurate than the approach designed and optimized specifically for the post-training/one-shot pruning.

Sensitivity to calibration data. Additionally, we have found that ZipLM is very robust to the amount of calibration data. In Table 4 we present a sensitivity analysis with respect to the number of calibration samples. We one-shot prune BERTbase on the SQuADv1.1 task for two speedup targets: 1.5x and 2.0x. In this setup, we compare results against Kwon et al. [26], which uses 2048 samples by default. As can be seen from the table, ZipLM outperforms prior state-of-the-art starting at only 32 samples. As we increase the number of samples, the results improve, up to 2 points in F1 score.

5 Discussion and Extensions

CPU as an LLM-inference environment. In Section 4 we have focused on various GPU-based inference environments as it enabled us to conduct fair comparisons against prior structural compression techniques. However, CPUs present another compelling inference environment focused on edge deployment of LLMs. Therefore, we target the recently proposed compound compression pipeline of [24], which involves three steps: structured pruning, unstructured pruning, and quantization. We replace their structured pruning approach based on layer dropping with ZipLM. As a result, at full accuracy recovery, we are able to improve speedup from 3x to 13x, and at the largest compression ratio from 30x to 50x. Due to space constraints, we provide full results in Appendix A.

Computational efficiency. Relative to distillation-based methods, structured pruning is an order of magnitude more efficient in terms of GPU hours due to the massive costs associated with pretraining from scratch for each compressed model [59, 51, 21]. For efficiency comparisons to CoFi, we consider the task of producing a full family of compressed BERTbase models with speedup targets ranging from 2x to 15x. In this setup, ZipLM requires only 115 epochs in total, whereas CoFi would require 560 epochs. Therefore, ZipLM is 4.87 times more efficient than CoFi. In terms of end-to-end runtime, ZipLM produces the entire family of compressed BERTbase models on a single RTX A6000 GPU in ∼ 35 hours on larger datasets (e.g. MNLI) and only ∼ 10 hours on smaller ones (e.g. SST2). Finally, it is worth emphasizing that we have not taken into account the cost of hyper-parameter tuning in the above comparisons, but that this is very favorable to ZipLM: it uses a single set of hyper-parameters to produce an entire family of compressed models while other methods require hyper-parameter tuning for each model independently.

1

10

20

30

40

50

60

70

80

Speedup factor relative to BERT

large

75

80

85

90

F1 score (%)

mini

small

medium

base

large

SQuADv1.1

BERT distillation

ZipBERT base

(ours)

ZipBERT large

(ours)

Scaling laws for structured pruning. To further understand the accuracy-speedup trade-offs, we run ZipLM on larger speedup ratios, up to 55x for BERTlarge and 75x for BERTbase . To the best of our knowledge, this is the first result in literature demonstrating that such extreme compression ratios are achievable with structured pruning without model collapse. In Figure 5, we compare these results against distillation-based downscaling of the BERT architecture [52]. The results clearly demonstrate that each of the pruned models, based either on BERTlarge or BERTbase , significantly outperforms comparable pre-trained variants. An emergent behavior that can be observed is that structurally pruned models tend to follow a linear scaling law, meaning that the accuracy decreases linearly with the increase of the speedup ratio, at a slope given by the original model. Fitting linearly via least squares produces the following expressions for the accuracy-speedup relationship: F1 large ≈ 92 . 1 -0 . 3 × speedup large , and F1 base ≈ 90 . 3 -0 . 6 × speedup base . Thus, the rate of decrease in accuracy for BERTbase is twice as large as that of BERTlarge , which can be attributed to the presence of more redundant representations in the larger model, making it more resilient to pruning. In Appendix G we provide additional analysis of the structure of pruned models.

References

Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once-for-all: Train one network and specialize it for efficient deployment. arXiv preprint arXiv:1908.09791 , 2019.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. In International Conference on Learning Representations , 2019.

Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. arXiv preprint arXiv:1805.04833 , 2018.

Elias Frantar and Dan Alistarh. SPDY: Accurate pruning with speedup guarantees. arXiv preprint arXiv:2201.13096 , 2022.

Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323 , 2022.

Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-fac: Efficient matrix-free approximations of second-order information. Advances in Neural Information Processing Systems , 34, 2021.

Elias Frantar, Sidak Pal Singh, and Dan Alistarh. Optimal Brain Compression: A framework for accurate post-training quantization and pruning. arXiv preprint arXiv:2208.11580 , 2022. Accepted to NeurIPS 2022, to appear.

Amir Gholami, Sehoon Kim, Zhen Dong, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. A survey of quantization methods for efficient neural network inference. arXiv preprint arXiv:2103.13630 , 2021.

Aaron Gokaslan and Vanya Cohen. Openwebtext corpus, 2019.

Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao. Knowledge distillation: A survey. International Journal of Computer Vision , 129(6):1789-1819, 2021.

Zichao Guo, Xiangyu Zhang, Haoyuan Mu, Wen Heng, Zechun Liu, Yichen Wei, and Jian Sun. Single path one-shot neural architecture search with uniform sampling. In Computer VisionECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XVI 16 , pages 544-560. Springer, 2020.

Babak Hassibi and David Stork. Second order derivatives for network pruning: Optimal brain surgeon. Advances in neural information processing systems , 5, 1992.

Yang He, Ping Liu, Ziwei Wang, Zhilan Hu, and Yi Yang. Filter pruning via geometric median for deep convolutional neural networks acceleration. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4340-4349, 2019.

Yihui He, Ji Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. Amc: Automl for model compression and acceleration on mobile devices. In Proceedings of the European conference on computer vision (ECCV) , pages 784-800, 2018.

Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2(7), 2015.

Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research , 22(241):1-124, 2021.

Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. Advances in Neural Information Processing Systems , 33:9782-9793, 2020.

Itay Hubara, Brian Chmiel, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. Accelerated sparse neural training: A provable and efficient method to find n: m transposable masks. Advances in Neural Information Processing Systems , 34:21099-21111, 2021.

Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithmetic-only inference. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2704-2713, 2018.

Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020 , pages 4163-4174, 2020.

Aran Komatsuzaki. One epoch is all you need. arXiv preprint arXiv:1906.06669 , 2019.

Eldar Kurtic and Dan Alistarh. Gmp*: Well-tuned global magnitude pruning can outperform most bert-pruning methods. arXiv preprint arXiv:2210.06384 , 2022.

Eldar Kurtic, Daniel Campos, Tuan Nguyen, Elias Frantar, Mark Kurtz, Benjamin Fineran, Michael Goin, and Dan Alistarh. The optimal bert surgeon: Scalable and accurate second-order pruning for large language models. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 4163--4181, 2022.

Mark Kurtz, Justin Kopinsky, Rati Gelashvili, Alexander Matveev, John Carr, Michael Goin, William Leiserson, Sage Moore, Bill Nell, Nir Shavit, and Dan Alistarh. Inducing and exploiting activation sparsity for fast inference on deep neural networks. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 5533-5543, Virtual, 13-18 Jul 2020. PMLR.

Woosuk Kwon, Sehoon Kim, Michael W Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gholami. A fast post-training pruning framework for transformers. arXiv preprint arXiv:2204.09656 , 2022.

François Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush. Block pruning for faster transformers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 10619-10629, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.

Quentin Lhoest, Albert Villanova del Moral, Yacine Jernite, Abhishek Thakur, Patrick von Platen, Suraj Patil, Julien Chaumond, Mariama Drame, Julien Plu, Lewis Tunstall, Joe Davison, Mario Šaško, Gunjan Chhablani, Bhavitvya Malik, Simon Brandeis, Teven Le Scao, Victor Sanh, Canwen Xu, Nicolas Patry, Angelina McMillan-Major, Philipp Schmid, Sylvain Gugger, Clément Delangue, Théo Matussière, Lysandre Debut, Stas Bekman, Pierric Cistac, Thibault Goehringer, Victor Mustar, François Lagunas, Alexander Rush, and Thomas Wolf. Datasets: A community library for natural language processing. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 175-184. Association for Computational Linguistics, November 2021.

Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. arXiv preprint arXiv:1608.08710 , 2016.

Yawei Li, Kamil Adamczewski, Wen Li, Shuhang Gu, Radu Timofte, and Luc Van Gool. Revisiting random channel pruning for neural network compression. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 191-201, 2022.

Lucas Liebenwein, Cenk Baykal, Harry Lang, Dan Feldman, and Daniela Rus. Provable filter pruning for efficient neural networks. arXiv preprint arXiv:1911.07412 , 2019.

Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978 , 2023.

Liyang Liu, Shilong Zhang, Zhanghui Kuang, Aojun Zhou, Jing-Hao Xue, Xinjiang Wang, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Group fisher pruning for practical network compression. In International Conference on Machine Learning , pages 7021-7032. PMLR, 2021.

Zhuang Liu, Jianguo Li, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang. Learning efficient convolutional networks through network slimming. In Proceedings of the IEEE international conference on computer vision , pages 2736-2744, 2017.

Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning , pages 22137-22176. PMLR, 2023.

Peter Mattson, Vijay Janapa Reddi, Christine Cheng, Cody Coleman, Greg Diamos, David Kanter, Paulius Micikevicius, David Patterson, Guenther Schmuelling, Hanlin Tang, et al. Mlperf: An industry standard benchmark suite for machine learning performance. IEEE Micro , 40(2):8-16, 2020.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.

Markus Nagel, Rana Ali Amjad, Mart Van Baalen, Christos Louizos, and Tijmen Blankevoort. Up or down? adaptive rounding for post-training quantization. In International Conference on Machine Learning , pages 7197-7206. PMLR, 2020.

NeuralMagic. Deep sparse: A fast cpu inference engine, 2021.

Matan Ben Noach and Yoav Goldberg. Compressing pre-trained language models by matrix decomposition. In Proceedings of the 1st Conference of the Asia-Pacific Chapter of the Association for Computational Linguistics and the 10th International Joint Conference on Natural Language Processing , pages 884-889, 2020.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. In EMNLP , 2016.

Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. Poor man's bert: Smaller and faster transformer models. arXiv preprint arXiv:2004.03844 , 2020.

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108 , 2019.

Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in Neural Information Processing Systems , 33:20378-20389, 2020.

S. Shankar. Identifying quora question pairs having the same intent. 2017.

Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. Advances in Neural Information Processing Systems , 33, 2020.

Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing , pages 1631-1642, 2013.

Yang Sui, Miao Yin, Yi Xie, Huy Phan, Saman Aliari Zonouz, and Bo Yuan. Chip: Channel independence-based pruning for compact neural networks. Advances in Neural Information Processing Systems , 34:24604-24616, 2021.

Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pages 4323-4332, 2019.

Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobilebert: a compact task-agnostic bert for resource-limited devices. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics , pages 2158-2170, 2020.

Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: The impact of student initialization on knowledge distillation. ArXiv , abs/1908.08962, 2019.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. ArXiv , abs/1804.07461, 2018.

Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. Advances in Neural Information Processing Systems , 33:5776-5788, 2020.

Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 6151-6162, 2020.

Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers) , pages 1112-1122. Association for Computational Linguistics, 2018.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-theart natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, October 2020. Association for Computational Linguistics.

Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 1513-1528, 2022.

Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, and Ming Zhou. Bert-of-theseus: Compressing bert by progressive module replacing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 7859-7869, 2020.

A Compound Compression for Edge Deployment

Deploying large language models in edge environments requires running inference on low-power devices such as CPUs. Therefore, we follow the compound compression approach from [24] which bundles together structured, unstructured pruning, and quantization for efficient inference on CPUs. We start with ZipLM structurally pruned models, and apply on top the state-of-the-art oBERT unstructured pruning method [24] to 80% sparsity. After structured and unstructured pruning, we apply quantization-aware-training (QAT) [20] to quantize FP32 weights into INT8 representations. We benchmark these compound compressed models by running inference in the DeepSparse [39] engine, on a single-core of Intel Cascade Lake CPU. In this setting, we compare our results against the compound compression pipeline of [24] which applies layer dropping as a form of structured pruning. As can be seen from Figure 6, when we substitute layer dropping with a principled structured pruning via ZipLM, the resulting compound compressed models achieve very competitive latency-vs-accuracy performance in the edge-inference regime. At full accuracy recovery, ZipLM improves the speedup from 3x to 13x, while at the largest compression ratio ZipLM improves the speedup from 30x to 50x.

92

93

94

95

96

97

98

99

100

F1 recall (%)

1

10

20

30

40

50

CPU inference speedup

310 ms

60 ms

30 ms

20 ms

15 ms

batch-size=1, seq-len=384, #cores=1

LayerDrop+oBERT+QAT

ZipLM+oBERT+QAT

B Ablation Studies

In Table 5, we present ablation results for ZipLM and CoFi, with and without their respective layerwise distillation techniques. ZipLM outperforms CoFi in all tasks when both methods use distillation, and in three out of four when distillation is not used. For example, ZipLM outperforms CoFi with a significant 3 point increase in F1 score on the SQuAD task in both setups. Furthermore, when comparing ZipLM results with and without layer-wise distillation, it can be observed that benefits are pronounced for low data tasks, where accuracy improvements reach up to 2 points.

C Additional GLUE results

Due to space constraints, in Figure 3 we present results only on four GLUE tasks. Therefore, for completeness, in Figure 7 we present results on the remaining four GLUE tasks, namely: CoLA, MRPC, STS-B, and RTE. Results show the same trends as the other four GLUE tasks, with large improvements for ZipLM, especially at higher compression rates.

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

35

40

45

50

55

Matthews corr.

BERTbase

CoLA

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

76

78

80

82

84

86

88

Accuracy (%)

BERTbase

MRPC

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

83

84

85

86

87

88

89

Spearman corr.

BERTbase

STS-B

2

3

4

5

6

7

8

9 10 11 12 13 14 15

Speedup factor

60

62

64

66

68

70

72

74

Accuracy (%)

BERTbase

RTE

ZipBERT base

(ours)

CoFi

DistilBERT

TinyBERT

Poor Man's BERT

MiniLM

Well-Read Students

Low-Rank BERT

LayerDrop

BERT-PKD

BERT-of-Theseus

D Additional Validation

Evaluating and comparing compressed models on the development set (dev-set) is standard practice, as it enables comparisons with off-the-shelf results from the literature. However, an implicit assumption behind such comparisons is that all methods tune their hyper-parameters only on a subset of the dev-set before evaluating and reporting results on all samples, which is not always the case. Moreover, specifically-tuned hyper-parameters can lead to large performance differences, especially when compressing LLMs [23]. To ensure that there is no such 'overfitting' on the dev-set, in Table 6 we compare ZipLM against the prior state-of-the-art CoFi approach on unseen test-set , obtained by submitting predictions to the official GLUE evaluation server. The results show consistent improvements over CoFi, on both dev- and test-sets.

E Latency table used for ZipLM pruning

As described in Section 3.2, we record the time to run an attention block, including all overheads, with 0 , . . . , N heads -1 heads pruned (pruning everything with runtime 0 is also considered) and similarly for the fully-connected block with the intermediate dimension shrunk by a factor of 0 . 9 i , for i = 0 , . . . , 42 ; in relative steps of 10% up until ≈ 99% sparsity, following [5]. In Table 7 we present an example of such a latency table used in ZipLM pruning approach.

F Speedup Evaluations

As shown in Figure 1, ZipLM is based on measuring runtimes of higher-level modules, such as attention heads and fully connected matrices, rather than low-level operators. This makes our

approach independent of underlying optimizations in different inference engines and frameworks, which usually perform further optimizations such as operator-fusion. Our runtime lookup table contains information about the runtime of a Transformer layer with different numbers of attention heads, and various dimensions of the fully connected matrices. This implies that we measure runtimes of a layer with 12 heads, 11 heads, 10 heads, and so on, as well as the runtimes of fully connected matrices with hidden sizes ranging from 3072 to 0. We utilize this information to guide pruning decisions.

To fully validate the ability of ZipLM to compress the model while satisfying desired speedup constraints via the described approach, we provide the timing results in Table 8, comparing the desired (target) speedup and the achieved (measured) speedup for different models.

As can be seen from the Table 8, the deviation between the desired (target) and the achieved (measured) speedup is at most 5.28%. This confirms that our approach indeed provides reliable runtime information to guide the pruning decisions.

G Structure of Pruned Models

Through a comprehensive examination of ZipLM pruned BERT models across all datasets considered in Section 4, we aim to identify trends in the pruning of key components of the Transformer layer, namely attention heads and intermediate size, needed to achieve a specific speedup target. As illustrated in Figure 8, we observe that the intermediate size is pruned at a higher rate relative to attention heads, which aligns with the fact that the intermediate size dictates the dimensions of the two large linear layers in the feed-forward part of the Transformer block. For instance, to attain a 2x speedup, roughly 60% of the intermediate size and 40% of the attention heads need to be removed. Additionally, in Figure 9, we visualize the entire encoder size needed to reach a specific speedup target. Interestingly, we find that 15x faster models retain on average only 2% of intermediate size and 6% of attention heads which amounts to only 2.9M parameters overall, while at the same time recovering more than 95% of the uncompressed model's accuracy (see Figure 3).

Additionally, in Figures 10, 11, 12, 13 we visualize the number of remaining heads and intermediate size across all Transformer layers and various speedup targets on a subset of GLUE datasets.

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Speedup factor

0

20

40

60

80

100

Percentage of pruned (%)

BERTbase  Attention heads

BERTlarge

Attention heads

BERTbase  Intermediate size

BERTlarge

Intermediate size

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

Speedup factor

0

20

40

60

80

Encoder-size (M)

85.1

38.1

23.7

17.0

13.1

10.3 8.5

7.2

6.1

5.2

4.5

4.0

3.5

3.2

2.9

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

2

4

6

8

10

#heads

3x

6x

9x

12x

15x

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

200

400

600

800

Intermediate size

3x

6x

9x

12x

15x

H Experiments - Additional Results

In Table 9 we report accuracy and model size of ZipLM pruned models visualized in Section 4, in Figures 2 and 3.

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

2

4

6

8

10

12

#heads

3x

6x

9x

12x

15x

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

200

400

600

800

Intermediate size

3x

6x

9x

12x

15x

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

2

4

6

8

10

12

#heads

3x

6x

9x

12x

15x

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

250

500

750

1000

1250

1500

Intermediate size

3x

6x

9x

12x

15x

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

2

4

6

8

10

12

#heads

3x

6x

9x

12x

15x

0

1

2

3

4

5

6

7

8

9

10

11

Transformer layer index

0

250

500

750

1000

1250

1500

Intermediate size

3x

6x

9x

12x

15x

I Hyper-parameters for Reproducibility

To facilitate reproducibility, we conduct experiments in the open-source Transformers library [58], and use publicly available datasets [28]. We plan to open-source our entire framework which supports one-shot and gradual structured pruning via SparseML [25], making it very easy to experiment with other models and datasets. In addition to our code, we plan to open-source all of our compressed models via the popular HuggingFace Hub. In Table 10 we report hyper-parameters used to produce our ZipLM pruned models in Section 4. Because of the excessive memory overhead, we don't make use of any kind of knowledge distillation when pruning the GPT2 model. Following insights from DistilGTP2, we hypothesize that this can further improve our results. We follow [22] and disable dropout regularization while pre-training ZipGPT2 models at OpenWebTextCorpus dataset.

J Broader Impact and Limitations

Our results contribute to the line of work on efficient language models. Thus, it should help reduce the energy and monetary cost of inference over such models, and allow them to be used without access to powerful hardware. While this is a mainly positive outcome, it also reduces the cost of employing these models for detrimental purposes, such as spam generation. Thus, this significant cost reduction for inference should also be seen as further motivation for methods to ensure safe usage of these models, such as watermarking or alignment.

As any academic study, our work is not without its limitations. All of our benchmarks are focused on English-language datasets and therefore our results do not provide insights into compression effects for low-data languages. Unfortunately, this limitation is inherent to all of the existing works on compression due to the lack of standardized benchmarks. Given that our structured pruning approach relies on a small sample of calibration data to perform pruning decisions, we hypothesize that our approach should be able to provide satisfying results in the low-data setup as well. At the moment we do not have data to support these claims, but we see it as an opportunity for future work.

Peer reviews:
comments: SUMMARY:
This paper proposed a novel structured compression approach for the inference of LLM, which reduces the model size by removing entire sub-components, like rows or columns from the model’s weight matrices. Based on their algorithm, they manage to drop out some attention heads and shrink the size of FFN layers. Besides, the proposed method considers the runtime speedup while deciding the sparsity of the pruning, which leads to a better trade-off between the loss and the performance.

STRENGTHS:
- This work comes up with a new structured pruning method for the transformer models and achieved reasonable speedup in the experiments with multiple models.
- The theoretical analysis is solid.

WEAKNESSES:
- The performance of the proposed method seems strongly dependent on the data distribution of inference, which may make it not that practical in real-world inference settings where the incoming inputs can vary a lot.
- In the latency-constrained scenario, the authors said "However, for the latter, the inputs are much smaller, and the size of weight matrices is no longer the primary bottleneck.", however, in my opinion, when the batch size is very small (like 1), loading weights from GPU HBM is exactly the bottleneck. Therefore, I guess sparsifying model parameters should lead to reasonable speedups?
- In the throughput-constrained (aka. large batch size), it will be interesting to see how ZipLM could support an even larger batch size since we can take advantage of the memory saved by pruning to hold more inputs. In this way, we may further increase the throughput?

QUESTIONS:
Please refer to the Weekness

LIMITATIONS:
The potential social impact is not discussed in the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper looks into structure pruning (e.g. neurons, or attention heads) of Transformer based models.\nAuthors propose to formulate the pruning objective by requiring the output of each pruned layer to be as close as possible to that of an unpruned layer. Then they adopt optimal brain surgeon algorithm to come up with the weights to keep (mask) and the update to such weights. The proposed algorithm prunes each group (e.g. a neuron in a layer) one by one, ensuring that correlated groups are properly accounted for (which is hard to do when choosing a number of neurons to prune in one go). \nAdditional extensions to the method include inclusion of inference aware criteria into the algorithm for choosing the groups to keep (using pre-computed knowledge of latency table)\nSuggested approach can be used for one shot pruning and gradual pruning. Experimental results indicate good performance on attention head pruning, fully connected layers pruning and removing full attention modules.\n\nSTRENGTHS:\n- Really good experimental results\n- Well written and easy enough to follow (modulo my comments below)\n\nWEAKNESSES:\n- Not very clear novelty\n- seems computationally expensive\n-  Ablations are missing (is the improvement coming from the fact that you use formulation (1) for each layer? is it from line 155 of the algorithm? Is it the inference awareness inclusion?\n\nQUESTIONS:\n1) How does pruning two consecutive layers happens? First pruning first layer (Algorothm 1) and then pruning the next layer?\nIf yes you are essentially assuming that layers are independent (so assuming block diagonal structure of the hessian). \n2) In your formulation (1) you seem to suggest that X is the input to the layer (whose weights W you are considering to prune). In the original Optimal Brain Surgeon, X is actually based on the gradients (and the global objective is considered, instead of local l2 objective of matching each layer's output). I think this all needs to be made clearer \n2) for experiments, it would be nice to report flocs or any other time measurement for your method (not the inference) and competing methods. \n3) Potentially you can apply your idea of pruning only one structure at a time (line 155 in Algorithm 1)  to any of the existing methods, so ablation is needed (is your improvement comes from the fact that u do essentially more gradual pruning?)\n4) The method you propose does not seem to be LLM specific - why not to compare it with methods on structured pruning for non LLM models\n\n\n- I am confused about contribution on line 56 (produces a family of compressed models). How is it any different from say using gradual pruning with various levels of sparsity (and at each step updating the copy for that sparsity level). You are still saving/updating multiple copies of the models right?\n\n\nMinor:\n- Algorithm 1: Mask Mr is not defined and is not updated in the body of the algorithm. I think you want to compile it based on R at the end of k iterations (the same also applies to Ms - please define how this mask looks like and its shape)\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes ZipLM, a structured pruning method that can achieve desired target runtime speedups.​ The idea of ZipLM is to solve a layer reconstruction problem with a structural constraint which minimizes the output changes on a set of calibration examples if the layer is reconstructed as it. The proposed ZipLM is shown to be effective on both decoder-only models and encoder-only models, outperforming prior distillation approaches and structured pruning approaches.\n\nSTRENGTHS:\n* The paper focuses on structured pruning problems, where the pruned model can easily get real speedups. Moreover, the proposed approach takes as input a target speedup as well as the hardware environment and optimizes the model on this specific setup to ensure the model to achieve desirable speedup. I admire this realistic setting and believe that this can be practically impactful.\n* The paper has conducted extensive experiments studying the effectiveness of the proposed ZipLM approach, including on both encoder and decoder models, both one-shot pruning setting and gradual compression setting, and with different levels of sparsity. \n* The experimental results show the proposed approach is more effective compared to the existing baseline methods.\n\nWEAKNESSES:\n* I believe the paper (especially the methodology part) can be presented better and a lot more details should be added (even in the appendix). For example, how do you solve the reconstruction problem and obtain the optimal mask and weight update (equation (2) & (3))? How do you consider the constraint C in equation (1) in the solution equation (2) & (3)? What exactly does a latency table look like and how do you get the table given a hardware environment? How many calibration inputs did you use in your experiments and how sensitive the final results will be to the number/quality of the calibration inputs?\n* The paper is motivated by structured compression of large language models (LLMs). However, all experiments are conducted on models with hundreds of millions of parameters. Given the current state-of-the-art language models (e.g., LLAMA-7B/13B/65B) contain a couple orders of magnitude more parameters compared to BERT, it is not clear how the proposed method works in larger models.\n\nQUESTIONS:\n* Do you think the proposed approach can apply to larger language models such as LLAMA-7B? If not, what is the barrier?\n* How sensitive the proposed approach is to the calibration examples?\n\nLIMITATIONS:\nI didn’t see a clearly potential negative societal impact of this paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes ZipLM, a structured pruning and reconstructing + layer-wise distillation + inference-aware pruning algorithm. The authors first extend the pruning formula of OBS to structured pruning and utilize estimation of inference for each structure and structured SPDY search to achieve more accurate inference-aware pruning. Then, the authors distill the model in a token-wise manner. Experimental results demonstrate that ZipLM achieves state-of-the-art (SOTA) performance in both one-shot and pruning/knowledge distillation settings.\n\nSTRENGTHS:\n1. ZipLM demonstrates exceptional results among existing compression algorithms. It not only outperforms retraining-constrained approaches but also surpasses the performance of some pruning + knowledge distillation algorithms.\n2. The authors propose a systematic framework for structured pruning, encompassing aspects such as pruning and reconstruction, inference-aware structure search, and knowledge distillation for performance recovery. This comprehensive framework exhibits a well-designed structure that, in my opinion, contributes to the entire community.\n\nWEAKNESSES:\n1. The methodology is somewhat incremental. The core contributions of the authors revolve around extending previous methods (OBS, SPDY) to structured pruning since those methods couldn't be directly applied. While the authors have made these extensions, the core essence of the method remains largely based on the previous framework.\n2. The experimental section of this paper has some shortcomings. For instance, although the authors conducted a simple ablation experiment, they did not analyze which specific parts of the framework played a crucial role in improving performance. Additionally, in the majority of compression papers, the GLUE benchmark's eight datasets are evaluated as a whole since they represent the fundamental measure for assessing the performance of compressed models. However, the authors only evaluated four of these datasets.\n\nQUESTIONS:\n1. Could you provide a more comprehensive ablation experiment, specifically to identify which module within the entire framework contributes significantly to performance improvement?\n2. Since Kwon [1] also employed a pruning metric and the layer-wise reconstruction via LLS for structured pruning, have you directly compared your pruning metric and the performance after reconstruction of OBS with their method to determine which one performs better, apart from the effect of distillation?\n3. Line 192-193: For example, a 95% sparse BERT produced by CoFi has 12x speedup on a V100 GPU, but only 5x on an A100 GPU. Since this is an interesting observation that FLOPs-based calculation != inference speed, can you explain what is the potential reason of this?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper proposes a structured compression method to optimize inference efficiency for language models. The proposed method uses the accuracy-efficiency trade-off under specific inference objectives as the importance measure of the model component( attention head, MLP neurons, entire FC/attention blocks.) Extensive evaluation results are provided to demonstrate the effectiveness of the proposed method.

STRENGTHS:
I appreciate the thoroughness of the evaluation, in which the authors consider the encoder/decoder and decoder-only model,  retraining compression and zero-shot compression, throughput and latency objective, etc. Evaluation provides a better understanding of the proposed method.

WEAKNESSES:
This paper focuses on improving the language model's inference efficiency. Other than distillation and pruning, quantization is also a popular direction[1], which is not discussed/compared. Quantization generally reduces the model size. At the same time, a reduction in model size would require less number of GPUs for large model inference, leading to latency speed up.  Further, there is a line of work on dynamic sparsity[2] that improves inference latency, which is also not discussed/compared.

[1]Frantar, Elias, et al. "Gptq: Accurate post-training quantization for generative pre-trained transformers." arXiv preprint arXiv:2210.17323 (2022).
[2]Liu, Zichang, et al. "Deja vu: Contextual sparsity for efficient LLMs at inference time." (2023).

QUESTIONS:
Table 1 shows a significant drop in PPL compared to the original GPT2, even at the smallest speed-up. To better understand the trade-off, can the authors comment on at what speedup we can expect no loss in performance?

LIMITATIONS:
The authors didn't discuss limiation or negative social impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07