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

Eldar Kurtic  Corresponding author<EMAIL_ADDRESS>Institute of
Science and Technology Austria  Daniel Campos  Neural Magic Inc.  Department
of Computer Science, University of Illinois Urbana-Champaign  Tuan Nguyen
Neural Magic Inc.  Elias Frantar  Institute of Science and Technology Austria
Mark Kurtz  Neural Magic Inc.
Benjamin Fineran  Neural Magic Inc.  Michael Goin  Neural Magic Inc.  Dan
Alistarh  Institute of Science and Technology Austria  Neural Magic Inc.

###### Abstract

Pre-trained Transformer-based language models have become a key building block
for natural language processing (NLP) tasks. 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
be applicable to decrease model size and increase inference speed. In this
context, this paper’s contributions are two-fold. We begin with an in-depth
study of the accuracy-compression trade-off for _unstructured weight pruning_
in the context of BERT models, and introduce _Optimal BERT Surgeon (O-BERT-S)_
, an efficient and accurate weight pruning method based on approximate second-
order information, which we show to yield state-of-the-art results in terms of
the compression/accuracy trade-off. Specifically, Optimal BERT Surgeon extends
existing work on second-order pruning by allowing for pruning blocks of
weights, and by being applicable at BERT scale. Second, we investigate the
impact of this pruning method when _compounding_ compression approaches for
Transformer-based models, which allows us to combine state-of-the-art
structured and unstructured pruning together with quantization, in order to
obtain highly compressed, but accurate models. The resulting compression
framework is powerful, yet general and efficient: we apply it to both the
fine-tuning and pre-training stages of language tasks, to obtain state-of-the-
art results on the accuracy-compression trade-off with relatively simple
compression recipes. For example, we obtain 10x model size compression with <
1% relative drop in accuracy to the dense BERT-base, 10x end-to-end CPU-
inference speedup with < 2% relative drop in accuracy, and 29x inference
speedups with < 7.5% relative accuracy drop.

##  1 Introduction

Pre-trained Transformer models Vaswani et al. (2017); Devlin et al. (2019)
have become a standard building block for many natural language processing
(NLP) tasks, providing robust language representations which can be
specialized on “downstream” classification and generation tasks. Despite their
success, these models have large parameter counts, which limits their
usability. Techniques for reducing these parameter counts and the
corresponding computational overheads have become vital, especially given the
recent breakneck pace of model growth Radford et al. (2019); MTN .

Figure 1: Performance overview relative to current state-of-the-art
unstructured pruning methods Chen et al. (2020) Xu et al. (2021) Sanh et al.
(2020), in this order, on the 12-layer BERT-base-uncased model and the
question-answering SQuAD v1.1 dataset.

Several compression approaches are known for large language models (LLMs). One
example is Knowledge Distillation (KD) Hinton et al. (2015), which led to
smaller models like DistillBERT Sanh et al. (2019), MobileBERT Sun et al.
(2020), and TinyBERT Jiao et al. (2020) which approximate the original model
with minor accuracy changes. Other work has leveraged lower-precision
representations to produce quantized models such as Q8BERT Zafrir et al.
(2019), TernaryBERT Zhang et al. (2020), and Q-BERT Shen et al. (2020). An
orthogonal approach, which is our primary focus, has been to apply fine-
grained methods such as unstructured pruning to produce model families such as
SparseBERT Xu et al. (2021) and PruneBERT Sanh et al. (2020), which compress
by removing individual weights. Figure 1 provides a comparative overview of
state-of-the-art results for unstructured pruning on a standard model (BERT-
base).

We investigate improved methods for unstructured and semi-structured (block)
pruning approaches at BERT scale, by leveraging the second-order (curvature)
approach pioneered by the Optimal Brain Surgeon framework LeCun et al. (1989).
We put our results in the context of a _compound compression_ framework, which
combines unstructured pruning with structured pruning and quantization,
showing that these methods can be complementary. To realize these compression
gains in practice, we tailor the resulting compressed models to execute
efficiently on a sparsity-aware CPU-based runtime engine NeuralMagic (2021),
leading to order-of-magnitude speedups at tolerable accuracy loss.

In summary, our contributions are as follows:

  * •

We perform a thorough exploration of weight pruning approaches applied to
LLMs, including lottery-ticket, movement pruning, gradual magnitude and
second-order methods.

  * •

We introduce a general second-order pruning method called _Optimal BERT
Surgeon (O-BERT-S)_ , which also supports block pruning, and is the first
second-order method to be both highly-accurate, and scalable to the
dimensionality of BERT models.

  * •

We illustrate the benefits of this pruning approach in two practical
scenarios. The first is comparing accuracy-vs-parameter-count relative to
existing pruning methods for LLMs. In this context, Optimal BERT Surgeon sets
new state-of-the-art results, significantly improving upon existing pruning
methods. For illustration, when pruning BERT-base, Optimal BERT Surgeon
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 with 3x fewer parameters.

  * •

We investigate the applicability of this pruning approach in a framework which
_compounds_ popular compression approaches for LLMs, i.e. applying pruning in
combination with layer dropping and/or quantization. In this context, we show
that our method can provide order-of-magnitude speedups on a sparsity-aware
inference at tolerable accuracy loss, on commodity CPUs. For illustration, on
a popular question-answering benchmark, we observe 8.4x higher inference
throughput at less than 1% relative accuracy loss to the dense BERT-base
model, 10x speedup at < 2% accuracy loss, 15x speedup at < 3% drop, and 29x
speedup at < 7.5% relative accuracy loss. Our compression recipes are simple
and easily-extensible.

##  2 Background and Related Work

Transformer LLM Models are usually built using multiple transformer layers,
used to capture long-term input dependencies, using self-attention Vaswani et
al. (2017). Each transformer usually has some variation of two sub-components:
multi head attention (MHA) and fully connected feed forward networks (FFN).
MHA contains many self-attention heads, each of which having three sub-
components: queries, keys, and values. The output of the attention component
is the concatenation of each attention head and is fed into the FFN. The FFN
is a fully-connected feed forward network which includes linear
transformations and an activation function such as ReLU or GeLU.

Given the massive size of well-performing models, there has been growing
interest in LLM compression. Existing approaches build on compression for
computer vision (CV) models, and successful methods are usually able to scale
between tasks. While many successful compression approaches exist, Transformer
models have been shown to be fragile Kovaleva et al. (2021), as minor
perturbations can lead to model collapse. We now briefly overview existing
compression approaches.

Unstructured pruning techniques remove individual weights by setting them to
zero, which can be exploited for both storage and computational speedup.
Pruning schemes are often motivated by weight saliency metrics which minimize
the loss in accuracy due to pruning. It is common to perform pruning in
iterative steps, each of which removes weights from the network until a
certain desired sparsity level is reached.

A popular saliency metric for weight pruning is weight magnitude, also known
as _magnitude pruning_ Han et al. (2015); Gale et al. (2019). While popular
for CNNs, this approach is believed to falter for LLMs Xu et al. (2021); Sanh
et al. (2020). We will show that, with the right parametrization, magnitude
pruning can be competitive for LLMs.

Second-order unstructured pruning methods LeCun et al. (1989); Hassibi and
Stork (1993); Singh and Alistarh (2020); Frantar et al. (2021) were developed
in the context of vision models, and leverage complex approximations of model
curvature, which can lead to improved pruning decisions. However,
implementations of second-order pruning methods do not immediately scale to
BERT models, as they require an approximation of the inverse Hessian, which
are expensive to store and compute with. The approach we propose here is a
variant of the WoodFisher / M-FAC methods Singh and Alistarh (2020); Frantar
et al. (2021), carefully adapted to BERT models for both scale and accuracy.
Specifically, we propose the first such method for BERT dimensionality, and
extend existing second-order methods to semi-structured (block) compression.

Movement Pruning (MvP) Sanh et al. (2020); Lagunas et al. (2021) is
specifically designed for pruning pre-trained LLMs during fine-tuning;
intuitively, it removes the weights with the lowest rate of change during
finetuning, even though they may not be near zero. The resulting family of
unstructured pruned models is called PruneBERT, and was the first to achieve
high sparsity (80-90%) with relatively small accuracy loss. Our approach
outperforms MvP in terms of accuracy at the same sparsity levels.

Upstream pruning approaches compress the pre-trained models themselves, during
their training process, reducing the need for task-specific tuning of the
pruning process. Chen et al. (2020) specifically examined strategies for
generating sparse pre-trained “lottery tickets,” Frankle and Carbin (2019)
which would transfer well to a set of downstream datasets. Their key
limitation, which we illustrate later, is that their approach incurs major
accuracy loss even at moderate sparsities.

Recent work by Zafrir et al. (2021) also explored compressing _pre-trained_
models, and established that gradual magnitude pruning can be successful in
this context, with careful tuning. They established that the pre-trained
sparse “once-for-all” (OFA) models can be finetuned to be competitive with
downstream methods such as MvP.

Relative to the line of work, in this paper we examine the performance of the
above methods, notably MvP, OFA, and Lottery Tickets, relative to a new
second-order pruning implementation. The approach we propose consistently
improves upon the compression-accuracy trade-offs given by all these prior
methods, both in pre-training and finetuning regimes.

Structured Pruning focuses on removing logical groupings of parameters without
affecting accuracy. For LLMs, structural pruning focuses on decreasing the
number of Transformer layers, decreasing the hidden size, or removing
attention heads, and requires structural understanding of the model.
Specifically, Michel et al. (2019) and Voita et al. (2019) demonstrated that
the attention heads vary in terms of importance, and that for some tasks
nearly 40% of heads can be removed without major impact on accuracy. Other
work has focused on showing that many of the layers can be removed Sridhar and
Sarah (2020), on the order in which layers are removed Sajjad et al. (2020),
and on methods for improving model robustness to layer dropping Fan et al.
(2019). Models like BORT de Wynter and Perry (2020) combine structured pruning
with an optimization approach to produce smaller models. In some of our
experiments, we will apply standard “direct” layer dropping in conjunction
with unstructured and semi-structured pruning.

Semi-Structured Pruning is an intermediate approach, by which smaller groups,
e.g. rectangular sets of weights Lagunas et al. (2021), are set to zero. This
approach has recently gained in popularity thanks to efficient computational
support. Here, we extend the second-order pruning formulation to such
groupings, and show results for a specific grouping supported by a CPU
inference engine.

Quantization represents weights and activations in lower precision Courbariaux
et al. (2016), and was used to obtain quantized language models such as Q8BERT
Zafrir et al. (2019). TernaryBERT Zhang et al. (2020) used a complex
distillation-based approach to obtain extremely low-bit representations,
whereas Fan et al. (2020) proposes a scheme which randomly quantizes groups of
weights during training, leading to more accurate compressed models. Finally,
Q-BERT Shen et al. (2020) employed approximate information about the Hessian
spectrum in order to choose the quantization bit-widths applied to each layer.
Follow-up work Yu et al. (2022) applied a similar approach to structured
pruning in the context of convolutional and languange models, using an
approximation of the Hessian trace to decide which layers should be pruned. We
note that this approach is quite different from the one we employ here, as we
use completely different inverse-Hessian approximations to perform pruning
decisions.

The focus of this paper is on _weight pruning_ , and on computational speedups
achievable on commodity CPUs. As such, the methods we investigate are rather
different from quantization approaches. Moreover, it is impossible to directly
compare to low-bitwidth quantized models in terms of inference speeds, as CPU
inference frameworks do not currently support such custom bitwidth formats. We
will therefore only make use of standard Quantization-Aware Training (QAT) to
8-bit weights, which is well-supported on Intel CPUs, and showcase the
resulting computational speedups in conjunction with layer-dropping and weight
pruning.

Knowledge Distillation (KD) Hinton et al. (2015) trains a smaller “student”
model against the outputs of a larger and more accurate “teacher” model,
specifically adding a loss component which minimizes the KL-divergence between
the two output distributions. KD uses a hardness parameter to control the
mixture of regular loss and distillation loss, and a temperature parameter to
control the softness of the probability distribution. Models trained with KD
reformulate their loss as a linear mixture of regular and distillation loss,
which is the approach we adopt in our setup too.

Approaches like DistillBERT Sanh et al. (2019), have the student approximate
the output of the teacher model, while TinyBERT Jiao et al. (2020), MobileBERT
Sun et al. (2020), and MiniLM Wang et al. (2020) have the student approximate
the teacher’s intermediate representations, obtaining better results at the
cost of higher complexity.

##  3 Optimal BERT Surgeon (O-BERT-S) Pruning

###  3.1 Generalized Second-Order Block Pruning

The pruning problem starts from a well-optimized dense model
$\mathbf{w}^{*}\in\mathbb{R}^{d}$, and aims to find a sparse version of
$\mathbf{w}^{*}$, where many of the weights are set to zero, and the remaining
weights may be updated accordingly in order to preserve the loss. It is common
for this process to occur gradually, i.e. by progressively removing the
weights. A classic approach LeCun et al. (1989); Hassibi and Stork (1993) for
“optimal” pruning of weights from $\mathbf{w}^{*}$ at a step is to expand the
loss function $\mathcal{L}$ locally around $\mathbf{w}^{*}$ with respect to a
sparse 0/1 weight mask $\mathbf{M}$. If we denote by
$\mathbf{w}_{M}=(\mathbf{M}\odot\mathbf{w}^{*})$, the model resulting from the
Hadamard (element-wise) product between $\mathbf{M}\in\\{0,1\\}^{d}$ and
$\mathbf{w}^{*}$, we can use the Taylor expansion at $\mathbf{w}_{M}$ to
obtain:

$\displaystyle\mathcal{L}(\mathbf{w}_{M})\simeq\mathcal{L}(\mathbf{w}^{*})+(\mathbf{w}_{M}-\mathbf{w}^{*})^{\top}\nabla\mathcal{L}(\mathbf{w}^{*})$
$\displaystyle+\frac{1}{2}(\mathbf{w}_{M}-\mathbf{w}^{*})^{\top}\mathbf{H_{\mathcal{L}}}(\mathbf{w}^{*})(\mathbf{w}_{M}-\mathbf{w}^{*}).$

Given that $\mathbf{w}^{*}$ is well-optimized, it is reasonable in practice to
assume that $\nabla\mathcal{L}(\mathbf{w}^{*})\approx\mathbf{0}$ Singh and
Alistarh (2020). Then, the change in loss incurred by pruning a subset of
weights can be expressed as

$\delta\mathcal{L}(\delta\mathbf{w})\simeq\frac{1}{2}\delta\mathbf{w}^{\top}\mathbf{H_{\mathcal{L}}}(\mathbf{w}^{*})\delta\mathbf{w}$
(1)

where
$\delta\mathcal{L}(\delta\mathbf{w})\coloneqq\mathcal{L}(\mathbf{w}_{M})-\mathcal{L}(\mathbf{w}^{*})$
and $\delta\mathbf{w}\coloneqq\mathbf{w}_{M}-\mathbf{w}^{*}$. In this context,
a popular way of approximating the model’s Hessian at $\mathbf{w}^{*}$ is via
a dampened empirical Fisher Information Matrix Hassibi and Stork (1993):

$\mathbf{H_{\mathcal{L}}}(\mathbf{w})\simeq\widehat{\mathbf{F}}(\mathbf{w})=\lambda\mathbf{I}_{d}+\frac{1}{m}\sum_{i=1}^{m}\nabla\mathcal{L}_{i}(\mathbf{w})\nabla\mathcal{L}^{\top}_{i}(\mathbf{w}),$
(2)

where $\lambda\geq 0$ is a small dampening constant,
$\mathbf{I}_{d}\in\mathbb{R}^{d\times d}$ identity matrix and $m$ is the
number of gradient outer products used to approximate the Hessian. Given the
positive-definiteness of (2), the quadratic form (1) is always nonnegative
which is why we will refer to $\delta\mathcal{L}(\delta\mathbf{w})$ as a loss
increase incurred by pruning.

Returning to our pruning problem, assume we wish to identify a block of
weights $Q$ of a given shape whose removal by “zero-masking” would incur
minimum increase in loss. This leads to the following constrained optimization
problem:

$\begin{split}&\min_{\delta\mathbf{w}}\quad\frac{1}{2}\delta\mathbf{w}^{\top}\widehat{\mathbf{F}}(\mathbf{w}^{*})\delta\mathbf{w}\\\
&\phantom{x}\mathrm{s.t.}\quad\mathbf{e}_{k}^{\top}\delta\mathbf{w}+w_{k}=0,\quad\forall
k\in\textrm{Q}\end{split}$ (3)

where $\mathbf{e}_{k}\in\mathbb{R}^{d}$ stands for the $k$-th canonical basis
vector. This derivation is known Hassibi and Stork (1993); however, the
optimal pruning update has only been derived for pruning individual weights.

Here, we will provide a generalized solution, which applies to general $Q$.
First, for convenience, we express the system of $|\textrm{Q}|$ equality
constraints in matrix-equation form as
$\mathbf{E}_{\textrm{Q}}\delta\mathbf{w}+\mathbf{E}_{\textrm{Q}}\mathbf{w}^{*}=\mathbf{0},$
where $\mathbf{E}_{\textrm{Q}}\in\mathbb{R}^{|\textrm{Q}|\times d}$ is a
matrix composed of the corresponding canonical basis vectors
$\mathbf{e}_{k}\,(\forall k\in\textrm{Q})$ arranged in rows. This optimization
problem can be solved with the method of Lagrange multipliers. Specifically,
we wish to find stationary points of the Lagrangian
$\mathcal{L}(\delta\mathbf{w},\boldsymbol{\lambda})$, where
$\boldsymbol{\lambda}\in\mathbb{R}^{|\textrm{Q}|}$ denotes a vector of
Lagrange multipliers. Solving the system of equations
$\frac{\partial\mathcal{L}(\delta\mathbf{w},\boldsymbol{\lambda})}{\partial\delta\mathbf{w}}=\mathbf{0}$
and
$\frac{\partial\mathcal{L}(\delta\mathbf{w},\boldsymbol{\lambda})}{\partial\boldsymbol{\lambda}}=\mathbf{0}$
yields the following expression for the optimal weight update:

$\delta\mathbf{w}^{*}=-\widehat{\mathbf{F}}^{-1}(\mathbf{w}^{*})\mathbf{E}_{\textrm{Q}}^{\top}\left(\mathbf{E}_{\textrm{Q}}\widehat{\mathbf{F}}^{-1}(\mathbf{w}^{*})\mathbf{E}_{\textrm{Q}}^{\top}\right)^{-1}\mathbf{E}_{\textrm{Q}}\mathbf{w}^{*}.$

Now, the corresponding loss increase incurred by the optimal weight update
$\delta\mathbf{w}^{*}$, which prunes a set of weights $Q$ and updates the
remaining weights, can be expressed as the saliency score of the weights $Q$,
which we denote by:

$\rho_{\textrm{Q}}=\frac{1}{2}\left(\mathbf{E}_{\textrm{Q}}\mathbf{w}^{*}\right)^{\top}\left(\mathbf{E}_{\textrm{Q}}\widehat{\mathbf{F}}^{-1}(\mathbf{w}^{*})\mathbf{E}_{\textrm{Q}}^{\top}\right)^{-1}\mathbf{E}_{\textrm{Q}}\mathbf{w}^{*}.$

We will use this saliency/importance score to rank sets of weights for
pruning. As a sanity check, if we prune a single weight $w_{j}$ at a time, our
derivations will yield the standard formulas of Hassibi and Stork (1993);
Singh and Alistarh (2020).

###  3.2 An Efficient Implementation

Directly implementing and running the previously described second-order
pruning approach for large language models, where number of weights to prune
$\mathbf{w}\in\mathbb{R}^{d}$ is very large, is infeasible. In particular,
this is due to the dependence on the inverse of the dampened empirical Fisher
information matrix
$\widehat{\mathbf{F}}^{-1}(\mathbf{w})\in\mathbb{R}^{d\times d}$, appearing in
formulations of the saliency score and of the optimal weight update. We now
describe how to circumvent these issues.

####  3.2.1 Pruning the optimal set of weights

Assume a gradual pruning setup, in which at each pruning step we wish to prune
a model to a target sparsity $s\in[0,1]$, effectively zeroing out $s\times d$
weights, in groups of size $|Q|$. Typically $s\times d\gg|Q|$, meaning that we
want to remove multiple groups of weights of size $|Q|$ at the same time.
Finding the optimal set of $\frac{s\times d}{|Q|}$ sets $Q$ is an intractable
combinatorial problem, due to all the correlations between groups of weights.
Specifically, the number of possible combinations is given by the binomial
coefficient $\binom{n}{k}$, where $n=\frac{d}{|Q|}$ and $k=\frac{s\times
d}{|Q|}$. This problem can be alleviated by ignoring correlations between
different groups of weights $Q$, and solving only for correlations between the
weights within the same group. In practice, this boils down to evaluating the
saliency score $\rho_{\textrm{Q}}$ for each group $Q$, and pruning the
$\frac{s\times d}{|Q|}$ groups with the lowest score. As pruning many weights
in the same step can make the Taylor approximation of the loss function less
accurate, one can consider pruning with multiple smaller sub-steps with
recomputations of the Hessian approximation in-between them (without
intermediate finetuning). While this can further improve the quality of the
pruning step Frantar et al. (2021), we do not implement this additional
optimization since the competing methods do not utilize recomputations.

####  3.2.2 Inverse empirical Fisher computation

The key component of the above procedure in terms of space and time complexity
is computing products with the inverse empirical Fisher. A simple direct
approach would be to perform a block-wise diagonal approximation of this
matrix (which we detail next), and perform direct block inversion. However, we
found experimentally that this approach is too expensive in terms of time, and
quite sensitive to the choice of dampening parameter $\lambda$. As an
alternative, we rely on the fact that the matrix we wish to invert is a sum of
rank-1 matrices, and employ the Woodbury/Sherman-Morrison (WSM) inversion
formula. Specifically, given a sum $(\mathbf{A}+\mathbf{u}\mathbf{v}^{\top})$
between an invertible matrix $\mathbf{A}$ and an outer product of vectors
$\mathbf{u}$ and $\mathbf{v}$ with compatible dimensions, the inverse
$(\mathbf{A}+\mathbf{u}\mathbf{v}^{\top})^{-1}$ can be exactly calculated as
$\mathbf{A}^{-1}-\frac{\mathbf{A}^{-1}\mathbf{u}\mathbf{v}^{\top}\mathbf{A}^{-1}}{1+\mathbf{v}^{\top}\mathbf{A}^{-1}\mathbf{u}}$.
Placing the expression of the empirical Fisher in the WSM formula, we obtain
the following recursive formulation, where $m$ is the number of gradients
employed in the approximation:

$\displaystyle\widehat{\mathbf{F}}^{-1}(\mathbf{w})=\widehat{\mathbf{F}}^{-1}_{m}(\mathbf{w})$
$\displaystyle=$
$\displaystyle\left(\widehat{\mathbf{F}}_{m-1}(\mathbf{w})+\frac{1}{m}\nabla\mathcal{L}_{m}(\mathbf{w})\nabla\mathcal{L}_{m}^{\top}(\mathbf{w})\right)^{-1}.$

Unrolling the recursion with
$\widehat{\mathbf{F}}^{-1}_{0}(\mathbf{w})=\frac{1}{\lambda}\mathbf{I}_{d}$,
we can obtain an iterative formula to exactly calculate the inverse of the
empirical Fisher matrix as

$\displaystyle\widehat{\mathbf{F}}^{-1}(\mathbf{w})=\widehat{\mathbf{F}}^{-1}_{m}(\mathbf{w})=$
$\displaystyle\frac{1}{\lambda}\mathbf{I}_{d}-\sum_{i=1}^{m}\frac{\left(\widehat{\mathbf{F}}^{-1}_{i-1}(\mathbf{w})\nabla\mathcal{L}_{i}(\mathbf{w})\right)\left(\widehat{\mathbf{F}}^{-1}_{i-1}(\mathbf{w})\nabla\mathcal{L}_{i}(\mathbf{w})\right)^{\top}}{m+\nabla\mathcal{L}_{i}^{\top}(\mathbf{w})\widehat{\mathbf{F}}^{-1}_{i-1}(\mathbf{w})\nabla\mathcal{L}_{i}(\mathbf{w})}.$

The iterative formulation enjoys a number of computational advantages over the
direct implementation. The most notable ones are 1) avoiding explicit calls to
the expensive and dampening-sensitive matrix inversion operations, and 2)
allowing successive updates of the inverse as new gradients are computed,
never needing to store all $m$ gradients of size $d$ and thus significantly
reducing memory requirements.

###  3.3 Memory and run-time complexity

Computing and storing the inverse empirical Fisher
$\widehat{\mathbf{F}}^{-1}(\mathbf{w})\in\mathbb{R}^{d\times d}$ is
prohibitively expensive for modern LLMs, which have hundreds of millions of
parameters, due to the quadratic complexity on the number of weights $d$.
However, Singh and Alistarh (2020) have shown that a diagonal block-wise
approximation of the empirical Fisher matrix can be very accurate for pruning
of convolutional neural networks. We adapt the same approach here, in the
context of LLMs. Thus, for blocks of width $B$ along the main diagonal, memory
requirements for the computation of the inverse Fisher matrix are reduced from
the quadratic $\mathcal{O}(d^{2})$ to a linear $\mathcal{O}(Bd)$ dependence on
the number of weights $d$. At the same time, run-time complexity relaxes from
$\mathcal{O}(md^{2})$ to $\mathcal{O}(mBd)$. As we will show, this computation
can be efficiently and accurately performed for moderate values of $m$ and
$B$.

Another alternative we investigated was the matrix-free approach of Frantar et
al. (2021), which does not require a block-wise approximation and has
complexity $\Theta(dm)$. However, our investigation showed that this approach
required high values of $m$ to be accurate, which leads to high memory cost in
the case of BERT models.

###  3.4 Efficient and scalable implementation

On the practical side, we have identified general hyper-parameters $B=50$ for
the block size, and $m=1024$ for the number of gradients which produce state-
of-the-art pruning results for all analyzed BERT models. Moreover, for these
parameter values, the block-wise approximation of
$\widehat{\mathbf{F}}^{-1}(\mathbf{w})$ can be implemented very efficiently on
modern accelerators. Specifically, we take advantage of the fact that such
hardware favors batched matrix operations, and that the blocks of size
$B\times B$ in $\widehat{\mathbf{F}}^{-1}(\mathbf{w})$ are independent, making
it possible to update the inverse Fisher estimate with new gradients in one-
shot. With $N_{B}=\frac{d}{B}$ we refer to the total number of blocks, i.e.
the batch-dimension.

The procedure works as follows. First, we compute batched matrix-vector
products
$\widehat{\mathbf{F}}^{-1}_{i-1}(\mathbf{w})\nabla\mathcal{L}_{i}(\mathbf{w})\in\mathbb{R}^{N_{B}\times
B}$ and scalar denominators

$m+\nabla\mathcal{L}^{\top}_{i}(\mathbf{w})\widehat{\mathbf{F}}^{-1}_{i-1}(\mathbf{w})\nabla\mathcal{L}_{i}(\mathbf{w})\in\mathbb{R}^{N_{B}}.$

Then, we update the inverse Fisher for each block by computing the scalar-
scaled outer products

$\left(\widehat{\mathbf{F}}^{-1}_{i-1}(\mathbf{w})\nabla\mathcal{L}_{i}(\mathbf{w})\right)\left(\widehat{\mathbf{F}}^{-1}_{i-1}(\mathbf{w})\nabla\mathcal{L}_{i}(\mathbf{w})\right)^{\top},$

which are in $\mathbb{R}^{N_{B}\times B\times B}.$

In practice, for the 12-layer BERT-base model with $d=85M$ encoder weights to
prune and block size $B=50$, the $\mathcal{O}(Bd)$ memory requirement
translates to approximately 17GB, which can be easily kept on RTX 3090 card
with 24GB of memory. While this amount of memory is available on high-
performance GPUs, it is also straightforward to split the $N_{B}\times B\times
B$ tensor along the batch-dimension $N_{B}$ and utilize additional GPUs or
even memory swapping with CPU memory.

##  4 Experimental Validation

###  4.1 Compound Compression Setup

Figure 2: Overview of the compound compression approach when applied
_downstream_.

The general compression setup we consider, illustrated in Figure 2, compounds
multiple compression methods, applied iteratively, in order to progressively
remove parameter redundancies, at increasingly low levels of granularity. In
the basic setting, we are given a large “upstream” model pretrained on a large
initial task, usually some form of language modeling, and wish to obtain a
compressed “downstream” model, specialized to a finetuning task. Due to space
constraints, we mainly focus the presentation on the case where compression is
only applied on “downstream” applications, but we also present results when
compression is applied to the “upstream” pretrained model.

Downstream Compound Compression (DCC) is a combination of existing techniques
during finetuning, as we employ structured compression, distillation,
unstructured pruning, and quantization. A critical component, however, is the
_order_ in which the compression techniques are applied: given an original
large model, we first apply _structured compression_ to remove _logical
structures_ , e.g. layers and their intermediate representations, from the
model. These structurally compressed models can then be pretrained as language
modeling tasks. In conjunction, we leverage information from the uncompressed
model by using a _knowledge distillation_ loss to train the smaller model,
with respect to the teacher’s output. Both pretraining and distillation have
been shown to learn different linguistic aspect helpful for downstream tasks
Turc et al. (2019).

The second step aims to shed the remaining surplus parameters by performing
_unstructured compression_ , via iterative second-order or magnitude pruning,
followed by an additional finetuning period to recover accuracy. This stage
brings some of the largest compression gains, as well as significant
improvements versus prior work. For example, we will be able to compress
models by up to 10x with minimal accuracy loss, and more than double the
compression ratio for the same accuracy relative to prior work Sanh et al.
(2020).

The third and last step applies quantization on top of the resulting sparse
model, via quantization-aware training (QAT). Clearly, some of the above
compression steps are optional, depending on the inference infrastructure:
e.g., one can skip quantization if there is no hardware support for lower-
precision operations.

We note that, while prior work has also explored _pairs_ of compression
methods in conjunction, e.g. Sanh et al. (2020), we believe we are the first
to examine compounding _all_ methods in a single pipeline.

This framework is naturally extensible: we also present experiments on
Upstream Compound Compression (UCC), which adapts the above scheme to the case
where we wish to already obtain a compressed model on the large-scale
pretraining setup. UCC is able to remove the need for parameter tuning on the
downstream dataset; however, we observe experimentally that the success of UCC
relative to DCC varies by task and compression target.

Task |  Dense BERT | Sparsity |  GMP (MvP) | Soft MvP |  GMP (ours) |  O-BERT-S (ours) |  Lottery Ticket |  GMP (ours) |  O-BERT-S (ours)
---|---|---|---|---|---|---|---|---|---
| Epochs |  |  10 Epochs Pruning + Finetuning |  30 Epochs Pruning + Finetuning
SQuAD | 88.54 |  80% 90% 97% |  - 80.10 59.60 |  - 84.90 82.30 |  - 85.74 79.66 |  - 87.98 84.65 |  86.54∗ 68.00∗ -∗ |  88.64 87.44 83.82 |  89.04 88.31 85.98
MNLI | 84.54 |  80% 90% 97% |  - 78.30 69.40 |  - 81.20 79.50 |  - 81.83 79.34 |  - 83.20 81.00 |  82.60∗ 75.00∗ -∗ |  83.64 82.79 79.95 |  84.32 83.79 81.77
QQP | 91.06 |  80% 90% 97% |  - 79.80 72.40 |  - 90.20 89.10 |  - 90.77 89.80 |  - 90.89 90.23 |  90.30∗ 90.00∗ -∗ |  91.47 91.19 90.48 |  91.57 91.35 90.87

Table 1: Performance metrics of the 12-layer BERT-base-uncased model gradually
pruned over 10 and 30 epochs at a corresponding downstream task with
distillation loss from the Dense BERT model. We report mean F1 score for SQuAD
and mean accuracy for MNLI and QQP over three runs for 30 epochs and one run
for 10 epochs for the sole purpose of a fair comparison with Sanh et al.
(2020). For standard deviations and additional metrics please see Tables 7 and
9 in the Appendix. (∗ approximate results as the exact numbers are not
available.)

Our experimental validation first examines the accuracy-compression trade-off
for unstructured compression, to establish the viability of second-order
pruning for BERT. Then, we integrate these gains in our compounding
compression framework, and then finally evaluate the impact of compression
both in terms of model size and inference speedups.

Datasets and Baselines. We perform evaluation on a set of downstream English
tasks, using well-established training hyper-parameters. All of our
experiments are built using a modified version of a popular Transformer
library Wolf et al. (2020) with a focus on BERT-base-uncased, one of the most
commonly used LLMs, composed of 12 transformer encoder layers with
approximately 110M parameters. Following the community standards, we prune
weights of all encoder layers ($\sim$85M parameters) and report sparsities
w.r.t. this number. For comparisons, we follow the approach used to illustrate
the current state-of-the-art approaches Sanh et al. (2020); Zafrir et al.
(2021) for downstream and upstream pruning, respectively.

Specifically, our experiments evaluate performance on a variety of publicly
available “downstream” English-language tasks commonly used to evaluate the
impact of model compression: question answering (SQuAD v1.1 Rajpurkar et al.
(2016)), sentence classification (Quora Duplicate Query Dataset (QQP) Guo and
Tian (2017)) and natural language inference (MNLI Williams et al. (2018)).
Training hyper-parameters, schedules, and other details can be found in the
Appendix. Experimental data will be made public via Weights and Biases Biewald
(2020) along with compression “recipes”, models, and our full implementation.

###  4.2 Improved BERT Unstructured Pruning

We first revisit the accuracy-compression trade-off induced by unstructured
pruning, and show that it can be improved upon significantly.

Goals and Setup. We compare existing approaches, notably MvP Sanh et al.
(2020) downstream, Prune OFA Zafrir et al. (2021) upstream, and Lottery Ticket
Chen et al. (2020) both downstream and upstream, against gradual pruning,
using either magnitude (GMP) or Optimal BERT Surgeon (O-BERT-S). For a fair
comparison with MvP, we consider the 10-epoch pruning and finetuning regime
used to obtain the best results in Sanh et al. (2020). We optimized the
learning rate schedule and the pruning frequency for both magnitude (GMP) and
second-order (O-BERT-S) pruning. Specifically, we perform 2 epochs of
finetuning, followed by 6 epochs of pruning, and 2 further epochs of
finetuning of the compressed model. We prune uniformly per each layer with
GMP, and globally across all layers with O-BERT-S. (Global magnitude pruning
yielded worse results.) The full hyper-parameters are described in Appendix
A.1.1. The results are given in Table 1, where GMP (MvP) is the baseline used
by Sanh et al. (2020), and GMP (ours) presents our results.

Movement Pruning (MvP) Comparison. We observe that our variant of GMP
significantly outperforms the GMP baseline of Sanh et al. (2020), and can in
fact achieve competitive performance with MvP itself. Further, we observe that
Optimal BERT Surgeon outperforms other methods by a significant margin (> 2%
absolute F1 score at the same sparsity): remarkably, O-BERT-S pruned to 97%
sparsity has similar accuracies to MvP at 90% sparsity, which has roughly 3x
more weights. This suggests that second-order information can be quite useful
for pruning and recovery. Figure 4 presents the accuracy drops following each
pruning step for GMP versus O-BERT-S: notice the large gaps in favor of the
second-order method, which lead to sustained performance.

Extended Finetuning. Next, we examine the accuracy effects of extending the
pruning schedule to 30 epochs, matching the setup used for finetuning BERT
Lottery Tickets (LTH) Chen et al. (2020). We use the same schedule structure
as before, but scaled to 30 epochs instead of 10. (Please see Appendix A.1.1
for full hyper-parameters.) Table 1 (right) presents results for Lottery
Tickets, our variant of GMP, and O-BERT-S pruning.

The results show a clear accuracy difference between O-BERT-S and GMP, on the
one side, and Lottery Tickets on the other, especially at higher sparsities.
This difference is justified since the LTH approach attempts to mainly
transfer _network connectivity_ , whereas the other approaches can also
benefit from the weight values.

Finally, we examined the impact of 30-epoch finetuning with soft MvP on SQuAD,
targeting 90% sparsity (not shown in the Table). We obtained an (F1, EM)
combination of $(87.42,79.83)$ for MvP. The F1 gap in favor of O-BERT-S is
lower than at 10 epochs, suggesting that additional finetuning helps all
methods; yet, it is far from negligible.

###  4.3 Downstream Compound Compression

| Unstructured weight pruning | 4-block pruning | QAT
---|---|---|---
Layers | Sparsity |  GMP (ours) |  O-BERT-S (ours) |  GMP (ours) |  O-BERT-S (ours) |  GMP (ours) |  O-BERT-S (ours)
12 |  0% 80% 90% |  89.48 88.64 87.44 |  89.48 89.04 88.31 |  89.48 87.85 85.95 |  89.48 88.57 87.57 |  89.06 87.09 85.45 |  89.06 87.89 86.68
6 |  0% 80% 90% |  88.32 87.30 85.56 |  88.32 88.20 86.78 |  88.32 85.54 82.77 |  88.32 87.00 85.34 |  87.94 84.03 82.51 |  87.94 86.10 84.59
3 |  0% 80% 90% |  84.66 83.40 81.09 |  84.66 84.08 82.50 |  84.66 80.61 77.38 |  84.66 82.79 80.69 |  84.25 80.04 77.65 |  84.25 82.04 79.66
Table 2: F1 performance metrics of the 3-, 6-, and 12-layer BERT-base-uncased
models compound-compressed over 30 epochs on SQuAD v1.1. For the corresponding
exact-match (EM) results, please see Table 10 in the Appendix.

The above results suggest that our approach can yield significant improvements
when applied to unstructured pruning. However, high-performance inference
usually applies other types of compression, such as block (semi-structured)
pruning, layer dropping, and quantization. We now compound these compression
approaches, following the pattern from Section 4.1. For this step, we apply
layer dropping to 6 and 3 layers, followed by either GMP- or O-BERT-S-based
pruning, for a total of 30 epochs. We investigate both _unstructured pruning_
and _4-block pruning_ , in which contiguous blocks of 4 weights are either set
to zero or kept dense, for which we use the implementation described in
Section 3. Both pruning types can be leveraged for computational speedups on
CPUs NeuralMagic (2021), but 4-block pruning can be used in addition with INT8
quantization, providing additional performance gains. For this reason, we also
perform quantization-aware training (QAT) Jacob et al. (2018) on top of the
4-block pruned models. (See Appendix A.1.3 for its description.) The results
are given in Table 2, where we also report the accuracy of the dense BERT
models (0% sparsity) distilled from the 12-layer teacher in the same setup.

The results indicate that compression methods can be combined without model
collapse, although the loss increases do compound. We observe that Optimal
BERT Surgeon consistently outperforms (well-tuned) GMP, and that the accuracy
gaps are larger in the case of 4-block pruning, especially for smaller models.
The fact that the layer-dropped models are also highly compressible suggests
that structured and fine-grained (unstructured) compression are complementary.
We find it remarkable that our 6-layer unstructured O-BERT-S-pruned model is
competitive with the 12-layer MvP-pruned model when both are pruned to 90%.

Speedup-vs-Accuracy. We now evaluate the trade-off between speed and accuracy
on the well-established SQuAD v1.1 CPU-inference benchmark. As is common, we
use a sequence length of 128 and a batch size of 32. Our evaluation metrics
are number of items per second and the size of the model in MB, after standard
gzip compression. Figure 4 depicts relative accuracy versus magnitude of
improvement in speed/model size.

We emphasize that, as baseline for 100% recovery, we follow the community-
standard, e.g. Sanh et al. (2020) and adopt the dense 12-layer BERT-base model
with 88.54 F1 score, which is also the reason why some F1-recall values are
above 100%. As can be seen from the improved dense baselines in Table 2, this
model can also benefit from our 30-epoch pruning and finetuning framework. For
inference speed, the baseline is dense inference on DeepSparse, which matches
the industry-standard ONNXRuntime inference engine.

Results suggest a roughly-linear trade-off between compression and accuracy
loss, with a compression jump around 1% accuracy drop, due to quantization
being applied. Specifically, we observe 8.4x higher inference throughput at
less than 1% relative accuracy loss to the dense BERT-base model, 10x speedup
at < 2%, 15x speedup at < 3% drop, and 29x speedup at < 7.5% accuracy loss.
This demonstrates how compounding compression can be used to optimize a single
backbone LLM to various latency requirements. Detailed results are in Appendix
Table 11.

Task |  Dense BERT | Sparsity | Lottery Ticket | Prune OFA |  O-BERT-S (ours)
---|---|---|---|---|---
SQuAD F1, EM | 88.54, 81.42 |  90% 97% |  68.00∗, - - |  87.25, 79.83 - |  88.42, 81.31 84.39, 76.36
MNLI m, mm | 84.54, 85.06 |  90% 97% |  75.00∗, - - |  81.45, 82.43 - |  82.29, 83.40 78.85, 79.80
QQP acc, F1 | 91.06, 88.00 |  90% 97% |  90.00, 86.90 - |  90.93, 87.72 - |  90.83, 87.65 89.76, 86.36
Table 3: Transfer learning performance of sparse pre-trained models obtained
by applying our upstream pruning recipes, described in details in Appendix
A.1.2. We report mean results from three runs with hyper-parameters described
in Table 6. For standard deviations please see Table 9 in the Appendix. (∗
indicates approximate results as the raw data numbers are not available.)

Figure 3: F1 score on SQuAD v1.1 during unstructured pruning with GMP and
Optimal BERT Surgeon for 90% sparsity. GMP is affected by higher accuracy loss
and recovers less.

Figure 4: F1 recall of uncompressed community-adopted 12-layer BERT baseline
on SQuAD v1.1 relative to improvements in model inference speed and size.

###  4.4 Upstream Compound Compresssion

Our results so far applied compression on the “downstream” finetuning task. An
appealing alternative is to perform compression directly “upstream”, on the
semi-supervised pre-training task Zafrir et al. (2021).

For this, we begin by preparing a dense upstream teacher using the RoBERTa
procedure Liu et al. (2019), detailed in Appendix A.1.2. Then, following Prune
OFA Zafrir et al. (2021) we scale our downstream pruning recipes to the
upstream task, to create sparse pre-trained models. To evaluate the resulting
sparse models, we finetune the unpruned weights on three downstream tasks
(SQuAD v1.1, QQP, MNLI). We compare against Prune OFA Zafrir et al. (2019), as
well as against the LTH approach Chen et al. (2020).

The results in Table 3 show that upstream Optimal BERT Surgeon pruning
outperforms state-of-the-art methods by significant margins, except for the
QQP dataset, where Prune OFA and O-BERT-S results are within standard
deviation. Moreover, in contrast to Zafrir et al. (2021), which performed
extensive hyper-parameter tuning, our transfer learning recipe is very basic:
8-epochs on all downstream tasks, coupled with a linearly decaying learning
rate of 1.5e-4 (see Table 6). This suggests that sparse pre-trained models
found by O-BERT-S constitute a strong starting point for sparse transfer. Yet,
downstream transfer can yield more accurate models, matching the intuition
that additional compression gains can be obtained by specializing to the
transfer task.

##  5 Discussion

We presented the first approximate second-order pruning method to apply to
BERT scale, and investigated its practicality in the context of compounding
compression techniques. Our results significantly improve upon the known
compression-accuracy trade-offs for BERT models, and show that we can leverage
the resulting models for significant practical speedups on CPU inference. A
consequence of our results is that techniques such as layer dropping, KD,
second-order pruning, and quantization are complementary, and can therefore be
applied in conjunction.

We emphasize that our compression recipes are simple and generic, which should
help practical adoption and reproducibility, but also allow for future
improvements. For example, the transfer learning results in Table 3 are
obtained with the same set of hyper-parameters across all three datasets, and
in the 8-epochs setup for a fair comparison against other methods, leaving a
lot of space for additional improvements by extending the fine-tuning phase
and by performing a hyper-parameter search for each dataset independently. The
unstructured and semi-structured (4-block) pruning results in Tables 1 and 2
are obtained with the same recipes and minimal hyper-parameter modifications
across all models and datasets.

In terms of the computational cost of our Optimal BERT Surgeon second-order
pruning method for the 12-layer BERT-base model with 85M of prunable encoder
weights, our implementation updates inverse Fisher estimate with a new
gradient instantaneously and can run asynchronously while the next gradient is
being fetched. Extracting the diagonal, computing inverse Fisher vector
products, saliency scores and optimal weight updates for the dense model takes
less than one second when the model is dense during the first pruning step,
and then drops to almost zero as the model becomes more sparse. All these
computations consume slightly more than 17GB of GPU memory, making the
algorithm runnable on a single commodity RTX 3090 card, which has 24GB of
memory. With this parametrization, our O-BERT-S gradual pruning runs had the
same end-to-end running time as our GMP runs.

Finally, it is worth emphasizing that our parametrization of the standard
gradual magnitude pruner produced drastically better results (in some cases
even up to 20% absolute F1 score improvement) relative to magnitude-based
pruning results currently available in literature. This suggests that this
approach can be competitive when tuned carefully.

In terms of future work, we aim to investigate distillation of intermediate
model representations Wang et al. (2020), as well as the role of compound
compression applied on much larger language models, including generative ones.

##  6 Broader Impact

Our work is part of the general trend of producing inference efficient models
which approximate performance of their larger bases. By and large, this work
should help increase model efficiency, thereby reducing computational and
ultimately monetary cost of executing such models. Moreover, it could allow
models to be used by those who do not have access to expensive specialized
computing clusters: for instance, our main experimental speedup results are
aimed at widely-available CPUs.

## References

  * (1) Using deepspeed and megatron to train megatron-turing nlg 530b, the world’s largest and most powerful generative language model.  https://bit.ly/3DlbPIF/.  Accessed: 2021-11-09.
  * Biewald (2020) Lukas Biewald. 2020.  Experiment tracking with weights and biases.  Software available from wandb.com.
  * Chen et al. (2020) Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. 2020.  The lottery ticket hypothesis for pre-trained bert networks.
  * Courbariaux et al. (2016) Matthieu Courbariaux, Itay Hubara, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio. 2016.  Binarized neural networks: Training deep neural networks with weights and activations constrained to +1 or -1.  _arXiv: Learning_.
  * de Wynter and Perry (2020) Adrian de Wynter and Daniel J. Perry. 2020.  Optimal subarchitecture extraction for BERT.  _CoRR_ , abs/2010.10499.
  * Devlin et al. (2019) J. Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In _NAACL_.
  * Fan et al. (2019) Angela Fan, Edouard Grave, and Armand Joulin. 2019.  Reducing transformer depth on demand with structured dropout.  _arXiv preprint arXiv:1909.11556_.
  * Fan et al. (2020) Angela Fan, Pierre Stock, Benjamin Graham, Edouard Grave, Rémi Gribonval, Herve Jegou, and Armand Joulin. 2020.  Training with quantization noise for extreme model compression.  _arXiv preprint arXiv:2004.07320_.
  * Frankle and Carbin (2019) Jonathan Frankle and Michael Carbin. 2019.  The lottery ticket hypothesis: Finding sparse, trainable neural networks.  _arXiv: Learning_.
  * Frantar et al. (2021) Elias Frantar, Eldar Kurtic, and Dan Alistarh. 2021.  M-fac: Efficient matrix-free approximations of second-order information.  _Advances in Neural Information Processing Systems_ , 34.
  * Gale et al. (2019) Trevor Gale, Erich Elsen, and Sara Hooker. 2019.  The state of sparsity in deep neural networks.  _ArXiv_ , abs/1902.09574.
  * Guo and Tian (2017) Lei Guo and H. Tian. 2017.  Duplicate quora questions detection.
  * Han et al. (2015) Song Han, Huizi Mao, and William J. Dally. 2015.  A deep neural network compression pipeline: Pruning, quantization, huffman encoding.
  * Hassibi and Stork (1993) Babak Hassibi and David G Stork. 1993.  _Second order derivatives for network pruning: Optimal brain surgeon_.  Morgan Kaufmann.
  * Hinton et al. (2015) Geoffrey E. Hinton, Oriol Vinyals, and J. Dean. 2015.  Distilling the knowledge in a neural network.  _ArXiv_ , abs/1503.02531.
  * Jacob et al. (2018) Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. 2018.  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.
  * Jiao et al. (2020) Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020.  Tinybert: Distilling bert for natural language understanding.  _ArXiv_ , abs/1909.10351.
  * Kovaleva et al. (2021) Olga Kovaleva, Saurabh Kulshreshtha, Anna Rogers, and Anna Rumshisky. 2021.  BERT busters: Outlier layernorm dimensions that disrupt BERT.  _CoRR_ , abs/2105.06990.
  * Lagunas et al. (2021) François Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush. 2021.  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. Association for Computational Linguistics.
  * LeCun et al. (1989) Yann LeCun, John S. Denker, and Sara A. Solla. 1989.  Optimal brain damage.  In _NIPS_.
  * Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019.  Roberta: A robustly optimized bert pretraining approach.  _ArXiv_ , abs/1907.11692.
  * Michel et al. (2019) Paul Michel, Omer Levy, and Graham Neubig. 2019.  Are sixteen heads really better than one?  In _NeurIPS_.
  * NeuralMagic (2021) NeuralMagic. 2021.  Deep sparse: A fast cpu inference engine.
  * Radford et al. (2019) Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019.  Language models are unsupervised multitask learners.
  * Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016.  Squad: 100,000+ questions for machine comprehension of text.  In _EMNLP_.
  * Sajjad et al. (2020) Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. 2020.  Poor man’s BERT: smaller and faster transformer models.  _CoRR_ , abs/2004.03844.
  * Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  _ArXiv_ , abs/1910.01108.
  * Sanh et al. (2020) Victor Sanh, Thomas Wolf, and Alexander M. Rush. 2020.  Movement pruning: Adaptive sparsity by fine-tuning.  _ArXiv_ , abs/2005.07683.
  * Shen et al. (2020) Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. 2020.  Q-bert: Hessian based ultra low precision quantization of bert.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 8815–8821.
  * Singh and Alistarh (2020) Sidak Pal Singh and Dan Alistarh. 2020.  Woodfisher: Efficient second-order approximation for neural network compression.  _Advances in Neural Information Processing Systems_ , 33.
  * Sridhar and Sarah (2020) Sharath Nittur Sridhar and Anthony Sarah. 2020.  Undivided attention: Are intermediate layers necessary for bert?  _ArXiv_ , abs/2012.11881.
  * Sun et al. (2020) Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. 2020\.  Mobilebert: a compact task-agnostic bert for resource-limited devices.  In _ACL_.
  * Turc et al. (2019) Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019.  Well-read students learn better: The impact of student initialization on knowledge distillation.  _ArXiv_ , abs/1908.08962.
  * Vaswani et al. (2017) Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _ArXiv_ , abs/1706.03762.
  * Voita et al. (2019) Elena Voita, David Talbot, F. Moiseev, Rico Sennrich, and Ivan Titov. 2019.  Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned.  In _ACL_.
  * Wang et al. (2020) Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. 2020.  Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers.  _ArXiv_ , abs/2002.10957.
  * Williams et al. (2018) Adina Williams, Nikita Nangia, and Samuel Bowman. 2018.  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.
  * Wolf et al. (2020) 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. 2020.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pages 38–45, Online. Association for Computational Linguistics.
  * Xu et al. (2021) Dongkuan Xu, Ian En-Hsu Yen, Jinxi Zhao, and Zhibin Xiao. 2021.  Rethinking network pruning – under the pre-train and fine-tune paradigm.  In _NAACL_.
  * Yu et al. (2022) Shixing Yu, Zhewei Yao, Amir Gholami, Zhen Dong, Sehoon Kim, Michael W Mahoney, and Kurt Keutzer. 2022.  Hessian-aware pruning and optimal neural implant.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , pages 3880–3891.
  * Zafrir et al. (2019) Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. 2019.  Q8bert: Quantized 8bit bert.  _2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing - NeurIPS Edition (EMC2-NIPS)_ , pages 36–39.
  * Zafrir et al. (2021) Ofir Zafrir, Ariel Larey, Guy Boudoukh, Haihao Shen, and Moshe Wasserblat. 2021\.  Prune once for all: Sparse pre-trained language models.
  * Zhang et al. (2020) Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. 2020\.  Ternarybert: Distillation-aware ultra-low bit bert.  _arXiv preprint arXiv:2009.12812_.
  * Zhu and Gupta (2018) M. Zhu and Suyog Gupta. 2018.  To prune, or not to prune: exploring the efficacy of pruning for model compression.  _ArXiv_ , abs/1710.01878.

##  Appendix A Appendix

###  A.1 Our compression recipes

####  A.1.1 Downstream pruning

In Table 4 we describe in detail our pruning recipes and hyper-parameters for
all downstream results presented in Tables 1 and 2. For easier comprehension,
we also visualize learning rate schedules in Figures 6 and 8, and sparsity
schedules in Figures 6 and 8. We prepare our 3- and 6-layer models for
downstream pruning runs in two stages: layer dropping and retraining phase. We
drop layers from our pre-trained TeacherMLM model (more details on it in
Appendix A.1.2). After dropping, we retrain the remaining layers (following
insights from Turc et al. (2019)) in the same setup used to prepare the
TeacherMLM model, with addition of the distillation loss from it. KD hyper-
parameters are the same as the ones we used in our upstream pruning runs, i.e.
(hardness, temperature) = ($0.5$, $2.0$).

| 10 Epochs | 30 Epochs
---|---|---
Batch size |  xxxxxxxxxxx 16 for SQuAD, 32 for GLUE
Learning rate (initial, final) |  (8e-5, 3e-5) for SQuAD, (8e-5, 2e-5) for GLUE |  (8e-5, 8e-6) for SQuAD, (5e-5, 5e-6) for GLUE
Learning rate schedule | linear decay with rewinds
Learning rate rewinds |  one, start at epoch=8 |  periodic every 4 epochs, start at epoch=2
Knowledge Distillation (hardness, temperature) | (1.0, 2.0)
Student model |  xxx 12-layer: bert-base-uncased 6-layer: finetuned with KD on upstream 3-layer: finetuned with KD on upstream
Teacher model | Dense BERT
Prune start epoch | 2
Prune end epoch | 8 | 26
Pruning frequency |  100x per epoch for GMP 102x per epoch for O-BERT-S |  100x per epoch for GMP every 4 epochs for O-BERT-S
Initial sparsity pruned in one-shot, at prune start epoch with GMP |  xxx 12-layer: 70% unstructured, 30% 4-blocked 6-layer: 30% unstructured and 4-blocked 3-layer: 30% unstructured and 4-blocked
Initial sparsity pruned in one-shot, at prune start epoch with O-BERT-S |  xxx 12-layer: 70% unstructured and 4-blocked 6-layer: 30% unstructured and 4-blocked 3-layer: 30% unstructured and 4-blocked
Pruning type |  xxxxxxxx GMP: uniform per layer O-BERT-S: global
O-BERT-S parameters |  xxxxxxxxx Number of gradients $m=1024$ Block size $B=50$ Dampening $\lambda=10^{-7}$
Table 4: Downstream pruning recipes and hyper-parameters used to obtain
results presented in Tables 1 and 2.

Figure 5: Visualized learning rate schedule for 10-epoch runs.

Figure 6: Visualized sparsity schedule for 10-epoch runs with initial sparsity
of 70% and target sparsity of 90%, following the cubic interpolation Zhu and
Gupta (2018).

Figure 7: Visualized learning rate schedule for 30-epoch runs.

Figure 8: Visualized sparsity schedule for 30-epoch runs with initial sparsity
of 70% and target sparsity of 90%, following the cubic interpolation Zhu and
Gupta (2018).

####  A.1.2 Upstream pruning

Motivated by outstanding improvements in results when pruning and finetuning
models with knowledge distillation from a dense teacher, we apply the same
recipe at the upstream task as well. We prepare a teacher for upstream pruning
by following insights from Liu et al. (2019). More concretely we start with
the bert-base-uncased model, adopt pre-training on two datasets (BookCorpus &
English Wikipedia) with focus on the masked language modeling task (MLM) for
10-epochs with batch size 256 and learning rate linearly decaying to zero from
the initial value of 1e-4. We refer to this teacher as TeacherMLM. In Table 6
we describe in detail our upstream pruning recipe. As can be seen, our pruning
recipe is just a downscaled version of our 30-epoch downstream-pruning recipe
to 3-epochs. An important update in our downstream recipe is a decrease of the
knowledge distillation hardness from $1.0$ to $0.5$. In our initial
experiments we have noticed that pruning and finetuning at upstream only with
distillation loss (i.e. hardness $1.0$) doesn’t perform well. We conjecture
that in contrast to small scale downstream datasets, pruning and finetuning at
large upstream datasets over a short span of 3-epochs might be a very
difficult task.

| 3 Epochs
---|---
Datasets | BookCorpus & English Wikipedia
Batch size | 256
Initial learning rate | 5e-4
Learning rate schedule | linear decay with rewinds
Learning rate rewinds | periodic every 0.5 epochs
Max sequence length | 512
Weight decay | 0.01
Knowledge Distillation (hardness, temperature) | (0.5, 2.0)
Student model | TeacherMLM
Teacher model | TeacherMLM
Pruning frequency |  1000x per epoch for GMP 1002x per epoch for O-BERT-S
Table 5: Upstream pruning recipes and hyper-parameters. Transfer learning
results with sparse models obtained with this recipe are presented in Table 3.

| 8 Epochs
---|---
Initial learning rate | 1.5e-4
Learning rate schedule | linear decay to zero
Batch size |  12 for SQuAD, 32 for GLUE
Knowledge Distillation (hardness, temperature) | (1.0, 2.0)
Teacher model | Dense BERT
Table 6: Transfer learning hyper-parameters used to finetune upstream-pruned
sparse models at all downstream tasks (SQuAD and GLUE). This recipe is used to
obtain results presented in Table 3.

####  A.1.3 Downstream quantization

We perform QAT on top of dense and 4-block pruned models on SQuAD v1.1 as
shown in Table 2. We quantize to 8 bits the embedding matrices, linear modules
of all encoder units which includes matrices in their attention and feed
forward layers, and the linear module of the output layer. Weights that were
pruned are kept constant during quantization. Non-linear operations within the
Softmax, LayerNorm and GeLU are not quantized.

For each dense and block pruned model in Table 2, we perform a total of ten
epochs training where the quantization observers are active for the first five
and the remaining is finetuning. We do hyper-parameter search over the
learning rates of 1e-4, 8e-5, 5e-5, 3e-5 and the distillation hardness of 0.9
and 1.0. We then pick the model with the best F1 score.

###  A.2 Additional performance metrics

Due to the space constraints, in the paper we report F1 score for SQuAD v1.1,
matched accuracy for MNLI, and accuracy for QQP dataset. As all of our hyper-
parameters for MNLI and QQP are exactly the same, we refer to these two
datasets as GLUE. We believe these hyper-parameters will transfer to all other
datasets that are part of the GLUE benchmark. In Table 7 we report the
additional metrics too: exact match (EM) for SQuAD v1.1, mismatched accuracy
for MNLI, and F1 score for QQP dataset. Tables 9 and 9 present standard
deviations of the corresponding results in Tables 1, 3 and 7. Finally, Table
10 presents the exact-match metric for the corresponding results in Table 2.

Task |  Dense BERT | Sparsity |  GMP (MvP) | Soft MvP |  GMP (ours) |  O-BERT-S (ours) |  GMP (ours) |  O-BERT-S (ours)
---|---|---|---|---|---|---|---|---
| Epochs |  |  10 Epochs Pruning + Finetuning |  30 Epochs Pruning + Finetuning
SQuAD | 81.22 |  80% 90% 97% |  - 70.20 45.50 |  - 76.60 72.70 |  - 77.75 69.68 |  - 80.76 76.14 |  81.54 79.90 75.17 |  82.08 81.12 78.11
MNLI | 85.06 |  80% 90% 97% |  - 79.30 70.60 |  - 81.80 80.10 |  - 82.54 79.38 |  - 83.58 80.67 |  84.42 83.50 80.38 |  84.91 84.35 82.01
QQP | 88.00 |  80% 90% 97% |  - 65.00 57.80 |  - 86.80 85.50 |  - 87.56 86.34 |  - 87.69 87.05 |  88.48 88.11 87.15 |  88.63 88.30 87.66
Table 7: Performance metrics of the 12-layer BERT-base-uncased model gradually
pruned over 10 and 30 epochs at a corresponding downstream task with
distillation loss from the Dense BERT model. We report mean exact match for
SQuAD, mean mismatched accuracy for MNLI and mean F1 score for QQP over three
runs for 30 epochs and one run for 10 epochs for the sole purpose of a fair
comparison with Sanh et al. (2020).

Task | Sparsity |  GMP (ours) |  O-BERT-S (ours)
---|---|---|---
| Epochs | 30 Epochs
SQuAD F1, EM |  80% 90% 97% |  0.08, 0.02 0.07, 0.03 0.17, 0.26 |  0.11, 0.03 0.13, 0.13 0.11, 0.17
MNLI m, mm |  80% 90% 97% |  0.26, 0.15 0.13, 0.33 0.20, 0.01 |  0.14, 0.13 0.05, 0.04 0.35, 0.22
QQP acc, F1 |  80% 90% 97% |  0.05, 0.05 0.04, 0.04 0.11, 0.12 |  0.08, 0.08 0.04, 0.06 0.05, 0.08
Table 8: Standard deviations of three runs for results presented in Tables 1
and 7.

Task | Sparsity |  O-BERT-S (ours)
---|---|---
SQuAD F1, EM |  90% 97% |  0.08, 0.16 0.30, 0.30
MNLI m, mm |  90% 97% |  0.13, 0.01 0.28, 0.09
QQP acc, F1 |  90% 97% |  0.11, 0.14 0.08, 0.10
Table 9: Standard deviations of three runs for results presented in Table 3.

| Unstructured weight pruning | 4-block pruning | QAT
---|---|---|---
Layers | Sparsity |  GMP (ours) |  O-BERT-S (ours) |  GMP (ours) |  O-BERT-S (ours) |  GMP (ours) |  O-BERT-S (ours)
12 |  0% 80% 90% |  82.71 81.54 79.90 |  82.71 82.08 81.12 |  82.71 80.51 77.99 |  82.71 81.46 80.14 |  81.99 79.44 77.17 |  81.99 80.57 78.84
6 |  0% 80% 90% |  81.17 80.18 77.98 |  81.17 81.15 79.16 |  81.17 77.75 74.18 |  81.17 79.55 77.65 |  80.85 75.57 73.71 |  80.85 78.27 76.56
3 |  0% 80% 90% |  76.62 74.82 72.07 |  76.62 75.62 73.61 |  76.62 71.53 67.96 |  76.62 74.07 71.36 |  76.06 70.74 67.84 |  76.06 72.70 70.0
Table 10: Exact-match (EM) performance metric of the 3-, 6-, and 12-layer
BERT-base-uncased models compound-compressed over 30 epochs on SQuAD v1.1.

###  A.3 Inference speedups and compression ratios of compressed models

Details on the results shown in Figure 4 are drawn from Table 11. As shown in
the results, not all compound compressed models yield improvements in
inference or compression relative to retained model performance but those that
do allow for massive improvements.

Layers |  Sparsity (%) |  Compression Method |  F1 score |  F1 recall (%) |  Throughput (items per sec.) |  DeepSparse Speedup |  Model size (gzip MB) |  Compression Ratio
---|---|---|---|---|---|---|---|---
12 | 0 | none | 88.54 | 100.00 |  1165.81 |  11.00 | 384.7 |  11.00
12 | 80 | unstructured | 89.04 | 100.56 |  1222.66 |  13.38 | 173.1 |  12.22
12 | 90 | unstructured | 88.31 |  199.74 |  1292.40 |  14.44 | 140.1 |  12.75
12 | 80 | 4-block+QAT | 87.89 |  199.26 |  1552.22 |  18.39 |  137.8 | 10.18
6 | 80 | unstructured | 88.20 |  199.62 |  1419.68 |  16.38 | 128.3 |  13.00
6 | 90 | unstructured | 86.78 |  198.01 |  1663.02 | 10.07 | 111.8 |  13.44
6 | 80 | 4-block+QAT | 86.10 |  197.24 |  1989.54 | 15.04 |  126.2 | 14.70
3 | 80 | unstructured | 84.08 |  194.96 |  1737.62 | 11.21 | 105.9 |  13.63
3 | 90 | unstructured | 82.50 |  193.18 |  1974.00 | 14.80 |  197.7 |  13.94
3 | 80 | 4-block+QAT | 82.04 |  192.66 | 1892.27 | 28.75 |  120.3 | 18.92
Table 11: Compression effects on model size and inference speed, evaluated at
batch size 32 with sequence length 128 on SQuAD v1.1 dataset.