SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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:
$$
\mathcal {L} \left(\theta^ {\mathrm {s}}, \theta^ {\mathrm {t}} | x\right) = \lambda_ {1} \mathcal {L} _ {\text {t a s k}} \left(\theta^ {\mathrm {s}} | x\right) + \lambda_ {2} \mathcal {L} _ {\text {l o g i t}} \left(\theta^ {\mathrm {s}}, \theta^ {\mathrm {t}} | x\right) + \lambda_ {3} \mathcal {L} _ {\text {t o k e n}} \left(\theta^ {\mathrm {s}}, \theta^ {\mathrm {t}} | x\right), \tag {5}
$$
where  $\theta^{\mathrm{s}}$  and  $\theta^{\mathrm{t}}$  represent student and teacher models respectively,  $x$  are the inputs,  $\mathcal{L}_{\mathrm{task}}$  is the loss associated with the task (e.g. cross-entropy for text-classification),  $\mathcal{L}_{\mathrm{logit}}$  is the KL-divergence between output logits as described in [16], and  $\mathcal{L}_{\mathrm{token}}$  is our token-level distillation loss. Hidden tensors passed between consecutive transformer layers are of constant shape  $\mathbf{H} \in \mathbb{R}^{B \times seq \times 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 \times seq$  vectors  $\mathbf{h} \in \mathbb{R}^{H}$  each carrying intermediate model representations of input tokens  $x$ . We define the loss  $\mathcal{L}_{\mathrm{token}}$  as an Euclidean distance  $\Delta$  between vectors  $\mathbf{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
$$
\mathcal {L} _ {\mathrm {t o k e n}} ^ {k} = \frac {1}{\sum_ {j = 1} ^ {B \times s e q} \mathbb {1} [ j \notin \mathbf {P} ]} \sum_ {j = 1} ^ {B \times s e q} \mathbb {1} [ j \notin \mathbf {P} ] \cdot \Delta \left(\mathbf {h} ^ {\theta_ {\mathrm {s}}}, \mathbf {h} ^ {\theta_ {\mathrm {t}}}\right), \tag {6}
$$
where  $\mathbf{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.
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  $\mathrm{BERT}_{\mathrm{base}}$  and  $\mathrm{BERT}_{\mathrm{large}}$  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 OpenTextCorpus [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.
Figure 1: Illustration of the ZipLM pipeline: 1) inference specifications, 2) runtime benchmarking of candidates for pruning, 3) gradual structured pruning until all speedup targets are met.
Baselines. In the gradual pruning setting, we explore the performance of ZipLM pruning of BERT and 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
$\mathrm{BERT}_{\mathrm{base}}$  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  $\mathrm{BERT}_{\mathrm{base}}$  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.
Figure 2: Structured compression of  $\mathrm{BERT}_{\mathrm{base}}$  (left) and  $\mathrm{BERT}_{\mathrm{large}}$  (right) on the SQuADv1.1 task. Dashed horizontal lines represent full and  $99\%$  accuracy recovery of the uncompressed model.
Figure 3: Structured compression of  $\mathrm{BERT}_{\mathrm{base}}$  on QNLI, MNLI, SST-2, and QQP tasks. Dashed horizontal lines represent full and  $99\%$  accuracy recovery of the uncompressed model.
$\mathrm{BERT}_{\mathrm{large}}$  results. To verify that our approach does not pertain only to the  $\mathrm{BERT}_{\mathrm{base}}$  model, we apply ZipLM structured pruning to the  $3\mathrm{x}$  larger  $\mathrm{BERT}_{\mathrm{large}}$  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  $\mathrm{BERT}_{\mathrm{large}}$  up to  $4\mathrm{x}$  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  $\mathrm{BERT}_{\mathrm{base}}$  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  $\mathrm{BERT}_{\mathrm{large}}$  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 OpenTextCorpus dataset, followed by zero-shot evaluations, without any fine-tuning, 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.5\mathrm{x}$ . Thus, we run ZipLM pruning to  $1.5\mathrm{x}$ ,  $2\mathrm{x}$ ,  $2.5\mathrm{x}$ , and  $3\mathrm{x}$  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 1: Zero-Shot perplexity (PPL) of compressed GPT2 in two regimes: pruning for throughput and pruning for latency. *GPT2 was trained by OpenAI [41] on a much larger closed-source dataset and for significantly longer. The only direct comparison is between DistilGPT2 and ZipGPT2.
Model	Pruning for throughput	Pruning for latency
Speedup	Decoder size	Wiki Text-103 PPL ↓	Speedup	Decoder size	Wiki Text-103 PPL ↓
GPT2*	1.0x	85.0M	28.5	1.0x	85.0M	28.5
DistilGPT2	1.6x	42.5M	43.0	1.9x	42.5M	43.0
ZipGPT2 (ours)	1.5x	47.3M	35.4	1.6x	48.7M	37.8
2.1x	26.5M	41.5	2.0x	39.2M	41.2
2.7x	14.0M	50.4	2.2x	26.6M	49.0
3.3x	5.7M	72.1	2.5x	20.7M	55.0
Table 2: One-shot (posttraining) structured pruning of  $\mathrm{BERT}_{\mathrm{base}}$  on three downstream datasets and two speedup targets.
Speedup	Kwon et al. [26]	ZipBERT base
SQuAD, F1
1.5x	86.2	87.1
2.0x	76.5	84.1
QQP, acc.
1.5x	89.5	89.7
2.0x	83.9	84.8
MNLI, acc.
1.5x	82.8	83.0
2.0x	78.1	78.2
4.2 On the Importance of Inference-Awareness
Depth vs. width pruning. A particularly interesting illustration of the importance of inference-awareness 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  $12\mathrm{x}$  speedup on a V100 is only  $5\mathrm{x}$  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  $\approx 90\%$  sparsity  $(3072\rightarrow 302)$  gives  $\approx 7\mathrm{x}$  speedup on a V100 but only  $\approx 3\mathrm{x}$  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.
Figure 4: Ablation study for the impact of the pruning target: pruning for sparsity (like prior approaches) versus pruning for speedup (the ZipLM approach).
Table 3: Speedups from shrinking the intermediate size of MLPs in the FFN section of a Transformer layer, on different GPUs.
MLP size	Speedup
V100	A100
3072	1.0x	1.0x
1814	1.6x	1.1x
1322	2.0x	1.4x
302	6.9x	3.1x
130	11.8x	4.4x
76	13.1x	4.4x
33	14.8x	4.4x
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  $\mathrm{BERT}_{\mathrm{base}}$  on the SQuADv1.1 task for two speedup targets:  $1.5\mathrm{x}$  and  $2.0\mathrm{x}$ . 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  $3\mathrm{x}$  to  $13\mathrm{x}$ , and at the largest compression ratio from  $30\mathrm{x}$  to  $50\mathrm{x}$ . 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  $\mathrm{BERT}_{\mathrm{base}}$  models with speedup targets ranging from  $2\mathrm{x}$  to  $15\mathrm{x}$ . 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  $\mathrm{BERT}_{\mathrm{base}}$  models on a single RTX A6000 GPU in  $\sim 35$  hours on larger datasets (e.g. MNLI) and only  $\sim 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.
Figure 5: Scaling laws of structured pruning vs. distillation on the standard BERT architecture.
Table 4: Sensitivity to the number of calibration samples.
Method	Num samples	F1 score at
1.5x	2.0x
ZipLM	4	82.3	48.4
32	86.8	82.6
128	86.8	83.6
512	86.8	84.1
2048	87.1	84.1
4096	87.6	84.7
Kwon et al.	2048	86.2	76.5
Scaling laws for structured pruning. To further understand the accuracy-speedup trade-offs, we run ZipLM on larger speedup ratios, up to  $55\mathrm{x}$  for  $\mathrm{BERT}_{\text{large}}$  and  $75\mathrm{x}$  for  $\mathrm{BERT}_{\text{base}}$ . 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  $\mathrm{BERT}_{\text{large}}$  or  $\mathrm{BERT}_{\text{base}}$ , 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:  $\mathrm{F1}_{\text{large}} \approx 92.1 - 0.3 \times \mathrm{speedup}_{\text{large}}$ , and  $\mathrm{F1}_{\text{base}} \approx 90.3 - 0.6 \times \mathrm{speedup}_{\text{base}}$ . Thus, the rate of decrease in accuracy for  $\mathrm{BERT}_{\text{base}}$  is twice as large as that of  $\mathrm{BERT}_{\text{large}}$ , 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.