PaperlayerView exact revision on arXiv

Source-linked AI summary

The Optimal BERT Surgeon: Scalable and Accurate Second-Order Pruning for Large Language Models

Eldar Kurtic, Daniel Campos, Tuan Nguyen, Elias Frantar, Mark Kurtz, Benjamin Fineran, Michael Goin, Dan Alistarh

arXiv:2203.07259v3cs.CLcs.LG

TL;DR

Large language models are accurate but costly to deploy, motivating better accuracy-compression trade-offs for BERT pruning. The paper introduces scalable second-order pruning with oBERT and combines it with other compression methods, reporting strong results across pruning stages and deployment metrics.

  • Problem

    Large Transformer models can be too large and computationally intensive for standard deployments, motivating improved accuracy-compression trade-offs for BERT pruning.

  • Method

    oBERT uses approximate second-order information for scalable unstructured and block pruning, and the paper compounds pruning with layer dropping and quantization.

  • Results

    oBERT improves on prior pruning methods across pre-training and fine-tuning, while compressed models improve accuracy, model size, and inference speed metrics.

  • Takeaways & Limitations

    The compound compression pipeline produces sparse models that can be deployed for CPU inference with order-of-magnitude speedups at low accuracy loss.

  • Takeaways & Limitations

    oBERT relies on approximate second-order information and requires non-trivial additional storage cost; the study also focuses on fine-grained sparsity types.

Abstract

from arXiv · show

Transformer-based language models have become a key building block for natural language processing. While these models are extremely accurate, they can be too large and computationally intensive to run on standard deployments. A variety of compression methods, including distillation, quantization, structured and unstructured pruning are known to decrease model size and increase inference speed, with low accuracy loss. In this context, this paper's contributions are two-fold. We perform an in-depth study of the accuracy-compression trade-off for unstructured weight pruning of BERT models. We introduce Optimal BERT Surgeon (oBERT), an efficient and accurate weight pruning method based on approximate second-order information, which we show to yield state-of-the-art results in both stages of language tasks: pre-training and fine-tuning. Specifically, oBERT extends existing work on unstructured second-order pruning by allowing for pruning blocks of weights, and by being applicable at the BERT scale. Second, we investigate the impact of this pruning method when compounding compression approaches to obtain highly compressed but accurate models for deployment on edge devices. These models significantly push boundaries of the current state-of-the-art sparse BERT models with respect to all metrics: model size, inference speed and task accuracy. For example, relative to the dense BERT-base, we obtain 10x model size compression (in MB) with < 1% accuracy drop, 10x CPU-inference speedup with < 2% accuracy drop, and 29x CPU-inference speedup with < 7.5% accuracy drop. Our code, fully integrated with Transformers and SparseML, is available at https://github.com/neuralmagic/sparseml/tree/main/research/optimal_BERT_surgeon_oBERT.

1 Introduction

The paper studies compression of large Transformer models through unstructured and block weight pruning, introducing oBERT as a scalable second-order method. It also evaluates combining pruning with other compression techniques for efficient deployment.

  • Method: oBERT applies approximate second-order information to unstructured and block pruning at BERT scale.The method builds on the Optimal Brain Surgeon framework and targets improved accuracy-compression trade-offs.
  • Study scope: The study evaluates lottery-ticket, movement, magnitude, and second-order pruning approaches for LLMs.
  • Results: More than 2% absolute F1 improvement over Movement Pruning is reported at the same sparsity for BERTBASE.oBERT can also match Movement Pruning accuracy with 3x fewer parameters.
  • Deployment: Compound compression combines pruning with layer dropping and/or quantization to produce sparse models suitable for CPU inference.The paper reports order-of-magnitude improvements over other compound-compressed models.

2 Background and Related Work

The background distinguishes major LLM compression strategies and prior pruning methods, while positioning oBERT as a second-order approach that works across pre-training and fine-tuning. The paper also frames pruning as complementary to other efficiency techniques.

  • Compression approaches: LLM compression includes knowledge distillation, quantization, layer dropping, unstructured pruning, and semi-structured block pruning.Semi-structured pruning zeros rectangular groups of weights and can benefit from efficient computational support.
  • Pruning methods: Unstructured pruning removes individual weights, while magnitude, gradient-based, and uncertainty-based methods use different saliency signals.Movement Pruning removes weights moving toward zero, whereas PLATON estimates uncertainty in importance scores.
  • Second-order pruning: Earlier second-order pruning methods relied on loss-curvature approximations developed primarily for image classification.The paper extends this line of work to LLM pre-training, fine-tuning, and block compression.
  • Comparative results: oBERT consistently improves on MvP, Prune OFA, and Lottery Tickets across upstream and downstream pruning stages.The resulting compressed models are also reported as smaller, faster, and more accurate than several prior compressed models.
  • Related efficiency methods: Token pruning and early exiting are described as orthogonal approaches to the paper’s weight-pruning focus.

3 The Optimal BERT Surgeon (oBERT)

oBERT formulates pruning as minimizing the loss increase from zeroing selected weights, using approximate second-order information and optimal compensating updates. To scale this approach to BERT, it combines groupwise scoring, block-diagonal Fisher approximations, and recursive inverse updates.

  • Second-order pruning formulation: oBERT starts from an optimized dense model and seeks a sparse model whose remaining weights can be updated to preserve the loss.The pruning process can remove weights gradually through zero-masking.
  • Second-order pruning formulation: The method approximates pruning-induced loss with a Hessian-based quadratic form and uses the resulting loss increase as a saliency score for candidate weight groups.The Hessian can be approximated with a dampened empirical Fisher information matrix.
  • Block pruning: For a candidate block, oBERT imposes zeroing constraints and solves a Lagrangian optimization problem to obtain the optimal compensating weight update.The generalized formulation supports arbitrary groups of weights rather than only individual weights.
  • Scalable pruning: Direct inversion of the empirical Fisher is infeasible for LLMs, so oBERT ignores correlations between different groups and ranks groups by their within-group saliency scores.Selecting many groups jointly is combinatorial; smaller pruning sub-steps can improve Taylor accuracy, but the authors do not implement recomputation between sub-steps.
  • Scalable pruning: oBERT uses Woodbury/Sherman-Morrison recursive updates to avoid explicit matrix inversions, update the inverse as gradients arrive, and reduce gradient-storage requirements.The authors report B = 50 and m = 1024 as practical settings fitting on a 24GB RTX 3090 GPU.
  • Scalable pruning: The block-diagonal approximation reduces inverse-Fisher memory from O(d^2) to O(Bd) and runtime from O(md^2) to O(mBd).Here, B is the block width, d the number of weights, and m the number of gradient outer products.
  • Scalable pruning: The investigated matrix-free alternative has complexity Θ(dm), but required high m for accuracy and therefore incurred excessive memory cost for BERT models.This limitation motivated the block-wise approximation used by oBERT.

4 Experimental Validation

Experiments evaluate oBERT across downstream pruning, upstream pruning, compound CPU compression, and GPU-compatible 2:4 sparsity. Results show improved accuracy–sparsity trade-offs and substantial CPU speedups, while preserving dense accuracy in the reported GPU setting.

  • Experimental setup: Experiments use BERTBASE with 110M parameters, pruning 85M encoder weights, and evaluate SQuAD, QQP, and MNLI downstream tasks.The study compares oBERT with Movement Pruning and Lottery Ticket methods.
  • Experimental setup: oBERT approximates the inverse Hessian with 1,024 gradients, prunes in blocks of 50 weights, and applies gradual sparsity schedules.The first pruning step removes 70% of weights before cubic interpolation reaches the target sparsity.
  • Downstream unstructured pruning: More than 2 F1 points separate oBERT from Movement Pruning at the same sparsity, while 97% oBERT sparsity matches Movement Pruning at 90%.The 90%-sparse Movement Pruning model has roughly three times more weights.
  • Upstream pruning: Upstream pruning on BookCorpus and English Wikipedia reduces downstream computation because only remaining unpruned weights require fine-tuning.The setup compares oBERT with Prune OFA and LT-BERT on selected GLUE tasks.
  • Compound compression for CPUs: Compound compression combines pruning, layer dropping, quantization, and a sparsity-aware CPU runtime without model collapse, although accuracy losses compound.Structured layer removal and fine-grained pruning are reported as complementary.
  • Compound compression for CPUs: 29x CPU-inference speedup occurs below a 7.5% accuracy drop, while 10x and 15x speedups occur below 2% and 3% drops, respectively.The reported trade-off uses dense BERTBASE with 88.54 F1 as the accuracy baseline.
  • Pruning for GPU speedups: oBERT’s one-shot 2:4 pruning recovers dense accuracy after one fine-tuning epoch and achieves 1.85x speedup on NVIDIA Ampere devices.The reported SQuAD scores are (F1, EM) = (88.58, 81.16).

5 Discussion

Without knowledge distillation, oBERT outperforms PLATON across GLUE tasks and improves upon competing compressed BERT models across accuracy, model size, and inference speed. These advantages extend to BERTLARGE and benchmark-oriented compression settings.

  • Comparison with concurrent work: oBERT outperforms PLATON across all GLUE tasks without knowledge distillation.The comparison uses an upstream-pruned model and sweeps only the number of fine-tuning epochs, with early stopping on smaller tasks.
  • Broader comparison: Compressed oBERT models improve upon DistilBERT, TinyBERT, and block MvP across accuracy, model size, and inference speed.The paper presents these models as competitive baselines across all three metrics.
  • BERTLARGE results: At 95% sparsity, oBERT-pruned BERTLARGE outperforms Prune OFA at 90% sparsity on SQuADv1.1.This tests whether the approach extends beyond the BERTBASE model.
  • MLPerf Inference Benchmark: The MLPerf experiments apply the full compound compression pipeline to BERTLARGE and MobileBERT, achieving order-of-magnitude model-size and inference-speed improvements while maintaining over 99% of dense BERTLARGE accuracy.The benchmark targets compressed and fast models under MLPerf guidelines.

6 Broader Impact

The work aims to improve inference efficiency and reduce the computational and monetary costs of executing language models. Its speedup results target widely available CPUs, potentially broadening access beyond specialized computing clusters.

  • Broader Impact: The approach is intended to reduce model-execution costs and improve efficiency on widely available CPUs.The broader-impact discussion connects inference efficiency with reduced computational and monetary costs.

7 Limitations

The method relies on approximate second-order information and requires additional storage, limiting the devices on which it can be applied. Future work also targets its focus on relatively fine-grained sparsity types.

  • Method limitations: Approximate second-order information is required to scale the method to BERT-sized models.The paper treats the validity of these approximations as indirectly supported by empirical results.
  • Method limitations: Additional storage requirements limit the range of devices on which the technique can be applied.The experiments run on a single commodity GPU, while multi-GPU scaling is provided.
  • Future scope: Future work aims to remove the focus on relatively fine-grained sparsity types such as unstructured and semi-structured pruning.This is identified as a limitation of the current study rather than an inherent limitation of the method.

A.1 MLPerf Inference benchmark

The MLPerf study evaluates two compound-compression directions for maintaining over 99% of BERTLARGE F1 on SQuADv1.1: pruning and quantizing BERTLARGE, or compressing MobileBERT with architectural reductions.

  • MLPerf target: The benchmark targets compressed models that maintain over 99% of BERTLARGE F1 on SQuADv1.1.This follows the MLPerf benchmark guidelines.
  • oBERT-Large: oBERT-Large applies 4-block downstream pruning to 95% sparsity followed by quantization-aware training without changing BERTLARGE's architecture.This is the first compound-compression direction.
  • oBERT-MobileBERT: oBERT-MobileBERT seeks to recover BERTLARGE accuracy by compressing an already compact MobileBERT model with direct layer dropping.The passage describes this as the second benchmark direction.

A.2 Additional comparisons

The paper compares oBERT with inference-efficient methods that remove tokens or exit early rather than pruning weights. At the same accuracy loss, its compressed model achieves substantially higher throughput than Learned Token Pruning.

  • Learned Token Pruning removes unimportant input tokens adaptively and reports 2x higher throughput at < 1% accuracy drop.
  • At the same < 1% accuracy drop, the paper reports 8.4x higher throughput for its compressed model than Learned Token Pruning.
  • Early-exit methods such as DeeBERT and FastBERT are presented as orthogonal approaches to inference speedup.

A.3 Computational costs

oBERT’s computational cost is dominated by storing blockwise inverse-Hessian approximations, but the implementation fits on a 24GB GPU and supports distributed or asynchronous execution. Saliency scoring and optimal updates are comparatively quick.

  • The NB × B × B tensor can be split across additional GPUs or swapped with CPU memory when needed.
  • Inverse-Hessian updates run in negligible time asynchronously while the next gradient is fetched, while saliency scores and optimal weight updates take only a few seconds.

A.4 Optimal BERT Surgeon (oBERT) hyper-parameters

The appendix specifies oBERT’s shared pruning hyper-parameters, training recipes, compression procedures, evaluation metrics, and computational setup. It also documents reproducibility details, dataset use, and the compound-compression results summarized from the main experiments.

  • Hyper-parameters: oBERT exposes three tunable hyper-parameters: number of gradients m, block size B, and dampening λ.The reported runs use one hyper-parameter set across models and datasets, selected on BERTBASE with SQuAD v1.1.
  • Downstream setup: Downstream pruning uses knowledge distillation from fine-tuned BERTBASE teacher outputs on SQuAD2 and GLUE tasks.
  • Upstream setup: The upstream recipe pre-trains bert-base-uncased for 10 epochs on BookCorpus and English Wikipedia, focusing on masked language modeling.The batch size is 256 and the learning rate decays linearly from 1e-4 to zero.
  • Compound compression: Quantization-aware training quantizes selected embedding, encoder-linear, and output modules to 8 bits while preserving pruning masks.Softmax, LayerNorm, and GeLU operations are not quantized.
  • Evaluation: The appendix reports F1 for SQuAD v1.1, matched accuracy for MNLI, and accuracy for QQP, with additional metrics supplied in supplementary tables.
  • Results and reproducibility: Compound compressed models do not all improve inference or compression relative to retained performance, but some provide massive improvements.
Loading 2203.07259v3…