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

Abstract: In this paper, we consider the problem of sparsifying BERT models, which are a key building block for natural language processing, in order to reduce their storage and computational cost. We introduce the Optimal BERT Surgeon (oBERT), an efficient and accurate 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 second-order pruning by allowing for pruning weight blocks, and is the first such method that is applicable at BERT scale. Second, we investigate 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 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.

## introduction pre-trained transformer models @xcite provide robust language representations which can be specialized on various tasks. given their massive growth @xcite , techniques for reducing their computational overheads have become popular. one classic technique is knowledge 2020 ), @xcite , @xcite , in this order, of the bert base model on the squadv1.1 task. distillation (kd) @xcite , which transfers knowledge from a larger teacher to a smaller student model. other work has leveraged lowerprecision representations to produce quantized models. an orthogonal approach, which is our primary focus, has been to apply unstructured and block pruning, i.e. removing individual weights, to produce compressed but accurate language models. figure 1 provides a comparative overview of state-of-the-art results for unstructured pruning. in this paper, we introduce a method for improved unstructured and semi-structured (block) pruning, by leveraging the second-order approach pioneered by the optimal brain surgeon framework @xcite ), which we scale for the first time to llms. further, we put our results in the context of a compound compression approach, which combines several compression techniques to obtain sparse models which we execute on a sparsity-aware cpu-based runtime @xcite , showing order-ofmagnitude speedups at low accuracy loss. in summary, our contributions are as follows: for illustration, when pruning bert base , obert outperforms movement pruning (mvp), the most accurate prior approach, by more than 2% absolute f1 score at the same sparsity, and can match the accuracy of mvp models with 3x fewer parameters. we first examine the performance of prior pruning methods, notably mvp, prune ofa, and lottery tickets, relative to the new second-order obert method. the approach we propose consistently improves upon all these prior methods, both in pre-training (upstream) and fine-tuning (downstream) stages, and can be compounded with other compression techniques to obtain models that are smaller, faster and more accurate than models like distilbert, tinybert, and block mvp. additional approaches for efficient inference of llms exist, like token-pruning and early-exiting. these approaches are orthogonal to ours; therefore we discuss them in appendix a.2. ## experimental validation to ease reproducibility, we conduct our experiments in modified versions of the popular open-source libraries: transformers @xcite , and sparseml @xcite . all of our experiments are using publicly available datasets via @xcite and focus on the bert base model @xcite , one of the most commonly used llms, composed of 12 transformer layers with 110m parameters. following community standards, we prune encoder's weights (85m) and report sparsities relative to this number. all of our models, compression recipes and the full implementation will be made public. ## limitations as any academic study, our work is not without its limitations. we split their discussion into limitations that are inherent to our method, and limitations of our present study; the latter can be overcome by extensions of our work. in the first category, we begin by highlighting the fact that our second-order method relies on approximations, which are inherent in order to scale such methods to bert scale. prior studies, e.g. (singh and alistarh, 2020) have performed careful examinations of the validity of these approximations in the context of cnn models. the strength of our empirical results can be seen as indirect evidence that these approximations apply to bert models as well. a second, technical, limitation is the fact that our method requires non-trivial additional storage cost; while we have shown that our experiments can be executed on a single commodity gpu (nvidia rtx 3090), this limits the range of devices on which the technique may be applied. however, we provide an efficient and easy way to scale our approach with more gpus, which is automatically utilized in a multi-gpu environment. another limitation which we aim to remove in future work is the focus on relatively fine-grained sparsity types, such as unstructured and semistructured pruning. 4-block downstream pruning to 95% sparsity followed by the quantization aware training (qat). in the second direction we focus on recovering the bert large accuracy by compressing an already compact mobilebert model, dubbed obert-mobilebert. more specifically, we apply direct layer dropping, leaving only 14 transformer layers out of the original 24, followed by the 4-block pruning to 50% sparsity and quantization aware training. we present results in