PaperlayerView exact revision on arXiv

Source-linked AI summary

ZipLM: Inference-Aware Structured Pruning of Language Models

Eldar Kurtic, Elias Frantar, Dan Alistarh

arXiv:2302.04089v2cs.LGcs.CL

TL;DR

Large language models impose substantial computational costs, while structured compression methods face accuracy and runtime-target challenges. ZipLM prunes structures using loss-runtime trade-offs and reports state-of-the-art results across one-shot and gradual settings, with runtime speedup guarantees.

  • Problem

    Structured compression can damage accuracy, complicate distillation, and fail to meet runtime targets when pruning uses sparsity or FLOPS goals alone.

    Source: §1 · p. 1

  • Method

    ZipLM iteratively removes structures using loss-runtime trade-offs and supports specified inference configurations; it also uses layer-wise token distillation without manual layer matching.

    Source: §1 · p. 2, §3.1 · p. 4, §3.2 · p. 5

  • Results

    ZipLM achieves state-of-the-art results in both post-training/one-shot and gradual compression settings while providing speedup guarantees for compressed models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors expect lower inference costs and use without powerful hardware, while noting that reduced costs can also facilitate detrimental uses.

    Source: J Broader Impact and Limitations · p. 19

    The benchmarks use English-language datasets and therefore do not establish compression effects for low-data languages.

    Source: J Broader Impact and Limitations · p. 21

Abstract

from arXiv · show

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.

1 Introduction

Structured pruning can make model savings usable across hardware, but accuracy loss and the lack of runtime guarantees make compression difficult to deploy. ZipLM balances loss and runtime during pruning and targets specified speedups.

1 Introduction

  • ZipLM achieves state-of-the-art results in both post-training/one-shot and gradual compression settings, while providing speedup guarantees for compressed models.Its inference-aware algorithm balances the loss-runtime trade-off at each pruning step.Source: §1 · p. 2
  • ZipLM combines weight magnitude, activation impact, and linear redundancy criteria while accounting for local and global correlations.It targets desired latency or throughput in a specified inference configuration.Source: §1 · p. 2
  • ZipLM adds layer-wise token-level distillation that avoids manual layer matching and boosts accuracy on small datasets.Source: §1 · p. 2
  • In one run with shared hyperparameters, ZipLM produces compressed models for multiple speedup targets and is compatible with unstructured pruning and quantization.The paper reports state-of-the-art results in GPU- and CPU-based inference environments.Source: §1 · p. 2

2 Related Work

Prior compression work includes student-teacher distillation and structured pruning of pretrained models. Other compression techniques, including dynamic sparsity and quantization, can complement structured pruning.

2 Related Work

  • MiniLM and TinyBERT distill teacher representations using specialized mechanisms, while MobileBERT uses a custom teacher and architecture-search-like component design at significant computational cost.Source: §2 · p. 2
  • Structured pruning methods reduce pretrained-model dimensions through approaches such as block removal, low-rank decomposition, submodule replacement, layer dropping, or CoFi’s multi-granularity masks.CoFi is described as the state-of-the-art structural pruning method, outperforming MiniLM and TinyBERT but not MobileBERT in accuracy-vs-speedup.Source: §2 · p. 2
  • Dynamic sparsity and lower-bit-width representations are complementary to ZipLM, and the paper applies quantization for higher compression on commodity CPUs.Source: §2 · p. 3

3 Method

ZipLM formulates structured pruning as preserving layer outputs while removing shared structures and compensating for their removal. It extends this local procedure with inference-aware search and token-level distillation.

3.1 The ZipLM Structured Pruning Algorithm (Local Correlations)

  • ZipLM seeks compressed layer weights that approximate original outputs on calibration inputs while meeting a compression constraint.Structured pruning makes row-wise choices dependent because the same weight indices must be removed across rows, with updates compensating for the removed structures.Source: §3.1 · p. 3
  • ZipLM extends Optimal Brain Surgeon formulas to prune shared structures jointly across matrix rows, accounting for correlations and updating remaining weights after each removal.Its structured approach also matches consecutive matrix shapes and has computational complexity lower by a factor of the hidden dimension than the cited layer-wise approach.Source: §3.1 · pp. 3–4
  • The method prunes attention heads, feed-forward intermediate dimensions, or entire attention and FC modules, choosing columns so Hessian information can guide pruning.Corresponding rows in preceding layers can then be removed without changing outputs.Source: §3.1 · p. 4

3.2 Inference-Aware Structured Pruning (Global Correlations)

  • Inference-aware pruning ranks component removals by loss-versus-speedup trade-offs for specified batch size, sequence length, hardware, and runtime targets.ZipLM precomputes per-layer pruned variants and searches for a configuration meeting the speedup constraint; its adapted structured SPDY search uses a squared-error prior and fixed 1,000-step search.Source: §3.2 · pp. 4–5

3.3 Layer-wise Token Distillation

  • ZipLM distills intermediate token representations across the model, combining task, logit, and token-level losses without manually matching teacher and student layers.The token loss compares teacher and student vectors for non-padded tokens and averages across unpruned layers.Source: §3.3 · pp. 5–6

4 Experiments

ZipLM is evaluated across BERT and GPT2 models, gradual and one-shot pruning, and inference settings with varied speedup targets. The experiments compare accuracy and runtime outcomes while examining how inference-awareness affects pruning decisions.

4 Experiments

  • ZipLM evaluates gradual pruning of BERTbase, BERTlarge, and GPT2 across speedup targets, and compares one-shot BERTbase pruning with prior work.In one-shot pruning, ZipLM outperforms Kwon et al. [49] starting with 32 calibration samples, with results improving by up to 2 F1 points as sample count rises.Source: §4 · pp. 6–7, Table 2 · p. 8, §4.3 · pp. 9–10

4.1 Gradual Structured Pruning

  • At the same speedup, ZipLM beats CoFi and TinyBERT by 3 F1 points on SQuADv1.1; at the same F1, it improves speedup by at least 60%.Across GLUE tasks, ZipLM also improves accuracy and speedup over competitive methods; it retains dense-model accuracy at up to 6x on QQP and 10x on SST-2.Source: §4.1 · p. 7, Figure 2 · p. 7, Figure 3 · p. 7
  • ZipLM matches uncompressed BERTlarge’s SQuADv1 F1 at up to 4x faster inference and doubles Block Movement Pruning’s speedup at the same F1.The comparison is against the fastest Block Movement Pruning model at 3x; ZipLM reaches 6x.Source: §4.1 · p. 7, Figure 2 · p. 7
  • At 99% accuracy recovery, ZipLM reaches 5x speedup on BERTbase SQuADv1, 6x on QNLI and MNLI, 13x on SST-2, and 15x on QQP.For BERTlarge on SQuADv1, it reaches 6x speedup at 99% recovery.Source: §4.1 · p. 8
  • At similar throughput speedup, ZipGPT2 improves on DistilGPT2’s perplexity and reduces decoder size by 60% while improving speedup by 30%.For latency pruning, ZipGPT2 has nearly 2 points better zero-shot perplexity and 3M fewer parameters at similar speedup.Source: §4.1 · p. 8, Table 1 · p. 8

4.2 On the Importance of Inference-Awareness

  • Pruning for speedup improves accuracy by up to 10 points over pruning for sparsity, especially at higher speedups.Speedups from the same MLP size reduction differ by device: shrinking 3072 to 302 yields 6.9x on V100 and 3.1x on A100.Source: §4.2 · p. 9, Figure 4 · p. 9

5 Discussion and Extensions

ZipLM extends structured pruning to CPU compound compression and extreme speedups while retaining competitive accuracy. It also reduces the cost of producing compressed-model families relative to CoFi.

5 Discussion and Extensions

  • At full accuracy recovery, ZipLM raises CPU speedup from 3x to 13x, and at the largest compression ratio from 30x to 50x.The CPU setting uses a compound pipeline combining structured pruning, unstructured pruning, and quantization.Source: §5 · p. 10
  • ZipLM needs 115 epochs to produce BERTbase models across 2x–15x speedup targets, versus 560 epochs for CoFi, making it 4.87 times more efficient.Producing the family takes about 35 hours on larger datasets and 10 hours on smaller ones using one RTX A6000; hyperparameter-tuning costs are excluded.Source: §5 · p. 10
  • ZipLM reaches structured-pruning speedups up to 55x for BERTlarge and 75x for BERTbase, with pruned models outperforming comparable pre-trained variants.The reported accuracy-speedup relationships are F1large ≈92.1 −0.3 × speeduplarge and F1base ≈90.3 −0.6 × speedupbase.Source: §5 · p. 10, Figure 5 · p. 10

A Compound Compression for Edge Deployment

ZipLM replaces layer dropping with structured pruning in a CPU compound-compression pipeline. The resulting models improve speedup at both full accuracy recovery and the largest compression ratio.

A Compound Compression for Edge Deployment

  • The pipeline applies ZipLM structured pruning, oBERT unstructured pruning to 80% sparsity, then QAT to convert FP32 weights to INT8 for single-core Cascade Lake CPU inference.Inference is benchmarked in DeepSparse.Source: A Compound Compression for Edge… · p. 15
  • At full accuracy recovery, ZipLM improves speedup from 3x to 13x, and at the largest compression ratio from 30x to 50x.The comparison substitutes ZipLM for layer dropping in a pipeline that also uses unstructured pruning and quantization.Source: Figure 6 · p. 15, A Compound Compression for Edge… · p. 15

B Ablation Studies

ZipLM generally outperforms CoFi in the ablations, and layer-wise distillation provides its largest gains on low-data tasks.

B Ablation Studies

  • With distillation, ZipLM outperforms CoFi on all tasks; without distillation, it wins on three of four.The comparison reports results with and without each method’s layer-wise distillation.Source: B Ablation Studies · p. 15
  • On SQuAD, ZipLM exceeds CoFi by 3 F1 points in both the with-distillation and without-distillation setups.Source: B Ablation Studies · p. 15
  • Layer-wise distillation improves ZipLM by up to 2 points, with pronounced benefits on low-data tasks.Source: B Ablation Studies · p. 15

C Additional GLUE results

The additional GLUE tasks show gains especially at higher compression, and unseen test-set comparisons retain ZipLM’s improvements over CoFi. The appendix also describes how ZipLM’s latency table is constructed.

C Additional GLUE results

D Additional Validation

  • ZipLM consistently improves over CoFi on both dev and unseen GLUE test sets evaluated through the official server.The test-set comparison addresses the risk that dev-set hyperparameter tuning affects reported results.Source: D Additional Validation · p. 16

E Latency table used for ZipLM pruning

  • ZipLM’s latency table records attention-block runtimes across head counts and fully connected-block runtimes across intermediate-dimension shrinkage levels.The fully connected block is sampled in relative 10% steps up to approximately 99% sparsity.Source: E Latency table used for ZipLM pruning · p. 16

F Speedup Evaluations

ZipLM uses runtime measurements of higher-level Transformer modules to guide pruning decisions. Across the models reported in Table 8, achieved speedups deviate from targets by at most 5.28%.

F Speedup Evaluations

  • ZipLM measures runtimes of Transformer layers with different head counts and fully connected matrix sizes to guide pruning.The runtime lookup table measures higher-level modules rather than low-level operators, making the approach independent of inference-engine and framework optimizations such as operator fusion.Source: F Speedup Evaluations · pp. 16–17
  • Table 6 compares ZipBERTbase and CoFi on development and test sets at comparable speedups.Source: Table 6 · p. 16
  • The maximum deviation between target and on-device measured speedups is 5.28%.Table 8 compares the desired and achieved speedups for different models.Source: F Speedup Evaluations · p. 17, Table 8 · p. 17

G Structure of Pruned Models

ZipLM prunes intermediate size more heavily than attention heads as it targets higher speedups. The figures show the resulting component sizes across speedups and datasets.

G Structure of Pruned Models

  • At 2x speedup, roughly 60% of intermediate size and 40% of attention heads are removed.The intermediate size sets the dimensions of the Transformer block’s two large feed-forward linear layers.Source: G Structure of Pruned Models · p. 17, Figure 8 · p. 18
  • At 15x speedup, models retain on average 2% of intermediate size and 6% of attention heads, totaling 2.9M parameters.Figure 9 visualizes encoder size against speedup for BERTbase models averaged over the considered datasets.Source: G Structure of Pruned Models · p. 17, Figure 9 · p. 18
  • Figures 10–13 visualize the remaining number of attention heads and intermediate size across all Transformer layers at various speedup targets, using a subset of GLUE datasets.Source: G Structure of Pruned Models · p. 17

H Experiments - Additional Results

Additional results include accuracy and model-size reporting and layerwise views of remaining components across several GLUE datasets. The appendix also records hyperparameters for gradual ZipLM runs.

H Experiments - Additional Results

I Hyper-parameters for Reproducibility

  • Table 10 lists hyperparameters for gradual ZipLM runs; GPT2 pruning omits knowledge distillation because of excessive memory overhead.The authors hypothesize that distillation could further improve GPT2 results, following insights from DistilGPT2.Source: I Hyper-parameters for Reproducibility · p. 19

J Broader Impact and Limitations

The paper presents lower inference costs as a potential benefit while noting that cheaper inference also lowers costs for detrimental uses. Its benchmarks are English-only and do not establish compression effects for low-data languages.

J Broader Impact and Limitations

  • Lower inference costs may reduce energy and monetary costs and enable use without powerful hardware, but may also make detrimental uses such as spam generation cheaper.The authors identify watermarking and alignment as methods to ensure safe usage.Source: J Broader Impact and Limitations · p. 19
  • The accompanying tables report accuracy and model size for pruned models and hyperparameters for gradual ZipLM runs.Source: Table 9 · p. 20, Table 10 · p. 20
  • Because benchmarks are English-only, results do not establish compression effects for low-data languages.The authors hypothesize ZipLM may still perform satisfactorily there, but report no supporting data.Source: J Broader Impact and Limitations · p. 21

Related papers

Loading 2302.04089v2…