# VTrans: Accelerating Transformer Compression with Variational Information
Bottleneck based Pruning

Oshin Dutta 1 Ritvik Gupta 2 Sumeet Agarwal 1
1 Indian Institute of Technology, Delhi
2 Carnegie Mellon University
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

In recent years, there has been a growing emphasis on compressing large pre-
trained transformer models for resource-constrained devices. However,
traditional pruning methods often leave the embedding layer untouched, leading
to model over-parameterization. Additionally, they require extensive
compression time with large datasets to maintain performance in pruned models.
To address these challenges, we propose VTrans, an iterative pruning framework
guided by the Variational Information Bottleneck (VIB) principle. Our method
compresses all structural components, including embeddings, attention heads,
and layers using VIB-trained masks. This approach retains only essential
weights in each layer, ensuring compliance with specified model size or
computational constraints. Notably, our method achieves upto 70% more
compression than prior state-of-the-art approaches, both task-agnostic and
task-specific. We further propose faster variants of our method: Fast-VTrans
utilizing only 3% of the data and Faster-VTrans, a time efficient alternative
that involves exclusive finetuning of VIB masks, accelerating compression by
upto 25 times with minimal performance loss compared to previous methods.
Extensive experiments on BERT, ROBERTa, and GPT-2 models substantiate the
efficacy of our method. Moreover, our method demonstrates scalability in
compressing large models such as LLaMA-2-7B, achieving superior performance
compared to previous pruning methods. Additionally, we use attention-based
probing to qualitatively assess model redundancy and interpret the efficiency
of our approach. Notably, our method considers heads with high attention to
special and current tokens in un-pruned model as foremost candidates for
pruning while retained heads are observed to attend more to task-critical
keywords.

##  1 Introduction

Since their inception, Transformers (Vaswani et al., 2017) have fundamentally
transformed the NLP field, offering pre-trained self-supervised models
adaptable to specific downstream tasks (Wang et al., 2018). However, the surge
in popularity and scale of Transformer models has amplified deployment
challenges on resource-constrained devices, attributed to elevated latency and
substantial storage demands (Raffel et al., 2020; Brown et al., 2020).

This has motivated extensive work in transformer pruning focusing on various
components such as layers, heads, blocks within weight matrices, and hidden
states. However, embeddings, which account for over 22% of total model
parameters, are often overlooked (Xia et al., 2022; Hou et al., 2020) due to
challenges in maintaining consistency with skip connection. Moreover, existing
methods often use magnitude-based (Han et al., 2015) or simply sparsity-based
pruning (Xia et al., 2022), which overlooks the importance of weights for the
given task.

An alternative approach to achieving compact yet high-performing models is
knowledge distillation (Hinton et al., 2015), where insights from a larger
teacher model are transferred to a carefully designed smaller student model.
However, task-agnostic knowledge distillation can incur a prohibitively high
computational cost (Jiao et al., 2019). Quantization is another effective
approach that achieves high compression and speedup with minimal loss in
performance by leveraging techniques such as quantization-aware training (Bai
et al., 2021; Kim et al., 2021; Zafrir et al., 2019).

Figure 1: (a) Our primary method - VTrans involves training all pre-trained
model parameters and VIB masks on the entire dataset during pruning, but
during finetuning, only the unmasked important weights are updated. (b) Fast-
VTrans utilizes a subset of data for both pruning and finetuning. (c) Faster-
VTrans, the fastest among all, trains VIB masks, Add Norm layer, and model
bias parameters during pruning and finetuning, using only a subset of the
entire dataset. (d) The masked and redundant units are removed after the
finetuning phase, resulting in the dense compressed model.

To address the challenges, we propose a novel transformer pruning method based
on the Variational Information Bottleneck (VIB) principle. It effectively
removes redundant elements while preserving information flow. Additionally, we
leverage knowledge distillation to achieve higher pruning ratios without
sacrificing performance.

Our contributions are:

  * •

We propose a structured pruning framework for transformers that is evaluated
in both task-specific and task-agnostic contexts, while adhering to user-
defined constraints on either model parameters or FLOPs.

  * •

Unlike prior methods, we extend compression to embedding states, alongside
other structural components like layers, attention heads, and feedforward
networks (FFN) enabling higher compression levels.

  * •

We propose two alternative - Fast and Faster approaches (Figure 1) to the main
compression method which are both time and resource efficient with minimal
performance degradation compared to previous approaches.

  * •

We evaluate the proposed methods on GLUE and SQuAD tasks while compressing
BERT, ROBERTa and GPT-2 pre-trained models with superior performance to
previous SOTA (Figure 2).

  * •

Furthermore, we establish the scalability of our method by pruning and
evaluating LLaMA-2 with 7 billion parameters on WikiText-2 dataset (Table 1).

##  2 Related Work

Pruning. It involves removing redundant model parameters for substantial model
compression with minimal performance loss (Xia et al., 2022; Zhu & Gupta,
2017; Zafrir et al., 2021; Renda et al., 2020; Liang et al., 2023b). Recent
techniques for transformer model pruning employ structured approaches,
targeting specific components like layers (Fan et al., 2020; Sajjad et al.,
2020), heads (Michel et al., 2019; Voita et al., 2019; Liang et al., 2021),
intermediate dimensions (McCarley et al., 2021; Wang et al., 2020c) and blocks
of weight matrices (Lagunas et al., 2021) or multiple components jointly (Xia
et al., 2022; Sun et al., 2023). Low-rank approximation has also been combined
with pruning (Li et al., 2023) to further compression. But these prior methods
do not prune embedding parameters which often form more than 20% of the total
parameters. Methods utilizing only forward pass during pruning (Sun et al.,
2023) often involve semi-structured pruning that cannot be made dense post-
pruning, leading to slower inference speeds. Certain pruning methods focus
only on task-agnostic Liang et al. (2023a) and some on task-specific Nasery et
al. (2023); Yang et al. (2022).

Figure 2: Our method surpasses previous techniques in compressing BERT-base,
ROBERTa-base, and GPT-2-small models. Our faster variants have reduced
compression timeframes. All models have 28M parameters being compressed from
their respective teachers.

Knowledge distillation. Another compression method involves training smaller
(student) model from a larger (teacher) model. The technique has found
applications in both, task-specific (Tang et al., 2019; Turc et al., 2019;
Aguilar et al., 2020) as well as task-agnostic domains (Sanh et al., 2020a;
Khanuja et al., 2021; Chen et al., 2021). In addition, recent advancements
(Sun et al., 2019; 2020; Hou et al., 2020; Romero et al., 2015) have extended
its application by enabling the incorporation of information from intermediate
layers into the student model’s training. Moreover, studies (Ma et al., 2023;
Sanh et al., 2020b) demonstrate the effectiveness of combining pruning with
knowledge distillation. However, these methods often entail extensive training
time.

Model | Size | PPL | Inference
---|---|---|---
Speedup
LLaMA-2 | 7B | 5.11 | 1$\times$
Wanda 2:4 | 3B | 10.52 | 1.14$\times$
Bonsai | 3B | 19.47 | 1.58$\times$
LLM-pruner | 3B | 19.24 | 1.5$\times$
Faster-VTrans | 3B | 11.8 | 1.82$\times$

Table 1: Pruned models evaluated on Wikitext-2. Our method outperforms
structured pruning (Bonsai and LLM-pruner) and achieves faster inference than
semi-structured pruning (Wanda 2:4).

Variational Information Bottleneck. It approximates (Alemi et al., 2016) the
information bottleneck principle (Slonim & Tishby, 1999; Tishby & Zaslavsky,
2015), focusing on extracting relevant information from input variables for
output variables.It aims to maximize mutual information between intermediate
layers and outputs while minimizing inter-layer mutual information to
eliminate redundancy in information. It has been successfully applied to
remove neurons in CNN, linear architectures (Dai et al., 2018) and to RNN
models (Srivastava et al., 2021). Henderson & Fehr (2022) formulate a non-
parameteric variational autoencoder with VIB for transformers, but it is not
aimed at pruning weights. In our work, we apply the principle for transformer
compression.

##  3 Background

###  3.1 Basic Transformer Architecture

Transformer network $f(\cdot;\theta_{s})$ parameterised by $\theta_{s}$ used
for NLP tasks typically consist of an embedding module and $L$ layers with
each layer containing Multi-Head Attention (MHA) module and two Feed Forward
Networks (FFN). We denote the input to a transformer network as
$\boldsymbol{x}\in\mathbb{R}^{d}$ and the associated label as
$\boldsymbol{y}\in\mathcal{Y}$. Further, we represent the embedding layer
hidden states obtained after the addition of positional and input embeddings
as $\boldsymbol{m}$, output of MHA activations as
$\left\\{\boldsymbol{a}_{i}\right\\}_{i=1}^{L}$ and FFN layer embedding as
$\left\\{\boldsymbol{h}_{i}\right\\}_{i=1}^{L}$.

###  3.2 Compressed Representations with Variational Information Bottleneck

Tishby & Zaslavsky (2015) and Dai et al. (2018) conceptualised successive
layer representations in a deep neural network to form a Markov chain while
treating the input as a stochastic variable. In the context of transformers,
we extend it to the successive output representations from the embedding
layer, Multi-Head Attention (MHA) layers, and Feed-Forward Network (FFN)
layers. Our objective is to obtain condensed intermediate representations or
activations $\tilde{\boldsymbol{k}_{i}}$ after each module (embedding layer,
MHA, FFN) in the transformer architecture while preserving essential
information in the predicted output $\boldsymbol{\tilde{y}}$. Unlike the
compressed FFN representation as formulated in Dai et al. (2018),
$\boldsymbol{k_{i}}\in\mathbb{R}^{n\times seq\times d}$,
$\tilde{\boldsymbol{k_{i}}}\in\mathbb{R}^{n\times seq\times d^{{}^{\prime}}}$
are each sets of vectors with $n$ input examples and $seq$ tokens and
$d^{{}^{\prime}}$ represents the compressed $d$ dimension. As done by Dai et
al. (2018), we frame the optimization problem as:
$\mathcal{L}_{i}=\beta_{i}I(\tilde{\boldsymbol{k}_{i}},\boldsymbol{k}_{i-1})-I(\tilde{\boldsymbol{k}_{i}},\boldsymbol{Y})$.
Here, $I(\cdot)$ denotes mutual information between two random variables and
$\beta_{i}$ is a hyper-parameter controlling the trade-off between compression
and prediction accuracy. To simplify notation, we denote the compressed
version $\tilde{\boldsymbol{k}}$ as $\boldsymbol{k}$.

To make the problem tractable, we invoke the variational upper bound as done
by Alemi et al. (2016),

$\tilde{\mathcal{L}_{i}}=\mathbb{E}_{\boldsymbol{X},\boldsymbol{Y},\boldsymbol{k}_{i},\boldsymbol{k}_{i-1}}\bigg{[}\beta_{i}\mathbb{D}_{KL}\Big{[}p\big{(}{\boldsymbol{k}_{i}}\mid\boldsymbol{k}_{i-1}\big{)}\|q\big{(}{\boldsymbol{k}_{i}}\big{)}\Big{]}-\log
q\big{(}{\boldsymbol{y}\mid{\boldsymbol{k}_{L}}}\big{)}\bigg{]}\geq{\mathcal{L}_{i}}$
(1)

Here, $q(\boldsymbol{k}_{i})$ and $q(\boldsymbol{y}\mid\boldsymbol{k}_{L})$
approximate $p(\boldsymbol{k}_{i})$ and
$p(\boldsymbol{y}\mid\boldsymbol{k}_{L})$, respectively. We get the compressed
representations by performing a dot product of the transformer activations
with a random set of vectors $\boldsymbol{z_{i}}\in\mathbb{R}^{n\times
seq\times d}$, where
$\boldsymbol{z_{i}}=\boldsymbol{\mu_{i}}+\boldsymbol{\epsilon_{i}}\odot\boldsymbol{\sigma_{i}}$.
$\boldsymbol{\mu_{i}}\in\mathbb{R}^{1\times
d},\boldsymbol{\sigma_{i}}\in\mathbb{R}^{1\times d}$ are learnable parameters,
while $\epsilon_{i}\in\mathbb{R}^{n\times seq\times d}$ is sampled from non-
parameterized $\mathcal{N}(\mathbf{0},\boldsymbol{I})$, thus broadcasting
$\boldsymbol{\mu_{i}}$ and $\boldsymbol{\sigma_{i}}$.

$\begin{array}[]{cc}\boldsymbol{k}_{i}=\boldsymbol{z}_{i}\odot
f_{i}(\boldsymbol{k}_{i-1})&;\boldsymbol{z_{i}}\in\mathbb{R}^{n\times
seq\times d}\end{array}$ (2)

$f(\cdot)$ denotes a standard sub-layer of the transformers like the value
head projection or the up projection sub-layer. Following Dai et al. (2018),
we derive the final VIB objective function as,

$\tilde{\mathcal{L}}_{\text{VIB}}=\sum_{i}^{L}\beta_{i}\sum_{j=1}^{r_{i}}\log\left(1+\frac{\mu_{i,j}^{2}}{\sigma_{i,j}^{2}}\right)-\mathbb{E}_{{\mathbf{X},\mathbf{Y},\mathbf{k}}}\left[\log
q\left(\boldsymbol{y}\mid\boldsymbol{k}_{L}\right)\right]$ (3)

where $r_{i}$ denotes the total hidden states or neurons within layer $i$. The
expectation is approximated with the task classification layer. Unlike the
application of VIB on non-transformer-based networks as done by Dai et al.
(2018), compressing one token’s representation in transformer-based models can
inadvertently affect the representations of other tokens due to their inter-
dependencies. However, by sampling a random vector with the shape
$(samples,sequence\\_length,dimensions)$, each token’s representation can be
uniquely adjusted by its own random vector. This approach potentially
preserves more detailed context and dependencies, maintaining the integrity of
the token-level information.

##  4 Method

Our approach involves two phases: pruning and finetuning.

During pruning, we systematically prune the student model, initialized from
the finetuned teacher, using VIB-based pruning while distilling knowledge from
the teacher model. In the finetuning phase, we optimize the remaining model
parameters and VIB masks to achieve the final compressed model with optimal
performance. We elaborate on the primary method in section 4.1, and discuss
the Fast and Faster variants in sections 4.3.

###  4.1 VIB-based Pruning

Following Equation 2, we define compressed representations for each structural
module in a transformer by multiplying the output of the module element-wise
to a random vector. The compressed embedding representation is thus defined as
$\boldsymbol{m}=\boldsymbol{z}_{m}\odot\text{Emb}(X,W)$, where
$\boldsymbol{z}_{m}$ is the random vector multiplied element-wise to the
output of the embedding layer. $W$ denote the weight matrices of the embedding
layer. Similarly, for layer $i$ the compressed MHA representation is defined
as
$\boldsymbol{a}_{i}=\boldsymbol{z}_{a_{i}}\odot[\operatorname{Att}_{1},\operatorname{Att}_{2},..\operatorname{Att}_{J}]$
, compressed FFN representation as
$\boldsymbol{h}_{i}=[\operatorname{gelu}\left(XW_{iU}\right)\operatorname{diag}(z_{inter,i})\cdot
W_{iD}]\odot\boldsymbol{z}_{out,i}$, where $W_{iU}\in\mathbb{R}^{r_{i}\times
r_{f}}\text{ and }W_{iD}\in\mathbb{R}^{r_{f}\times r}$ are intermediate up-
projecting and final down-projecting layer matrices respectively.
Layer pruning. We incorporate the ability to prune whole MHA and FFN layers by
introducing additional random vector $z_{layer}$. This enables effective
pruning in scenarios with high sparsity requirements. Thus, the MHA and FFN
representation for layer $i$ are defined as
$MHA_{{i}}=z_{layer,i}.(\boldsymbol{a}_{i}.W_{i}^{O});\;\;\;\;FFN_{i}=z_{layer,i}\cdot\boldsymbol{h}_{i}$,
where $W_{i}^{O}\in\mathbb{R}^{r_{i}}$ denotes the output attention matrix.

Implementation of masks for pruning. The $\mu_{i,j}$ and $\sigma_{i,j}$
parameters of the $z_{i,j}$ vectors of layer $i$ and neuron $j$ are initially
sampled from normal distributions and iteratively refined during pruning. As
shown by Dai et al. (2018), at the minima of the Equation 2, whenever
$\alpha_{i,j}=\mu_{i,j}^{2}/\sigma_{i,j}^{2}=0$, then that neuron $j$ of layer
$i$ is redundant. In the inference phase, these $z_{i,j}$ vectors are
converted to hard masks $z_{i,j}^{mask}$ using a thresholding operation on
$\alpha_{i,j}$ for pruning neurons as defined by,

$\displaystyle z_{i,j}^{mask}$ $\displaystyle=\begin{cases}0&\text{if
}\log\frac{\mu_{i,j}^{2}}{\sigma_{i,j}^{2}}\leq\tau\\\ 1&\text{if
}\log\frac{\mu_{i,j}^{2}}{\sigma_{i,j}^{2}}>\tau\\\ \end{cases}$ (4)

Sparsity Control. Using $z_{i,j}^{mask}$ hard masks, we calculate model
sparsity as the ratio of pruned parameters to the initial count. For FLOPs
constraints, sparsity represents the reduction ratio in FLOPs to the initial
value. We adopt an approach akin to Xia et al. (2022); Dutta et al. (2023) to
incorporate a Lagrangian term which enhances stability and provides finer
control over pruning by enforcing an equality constraint $s_{e}=t$ and
introducing a violation penalty as
$\mathcal{L}_{s}=\lambda_{1}\cdot(s_{e}-t)+\lambda_{2}\cdot(s_{e}-t)^{2}$,
where $s_{e}$ is the expected model sparsity, $t$ is the target sparsity and
$\lambda_{1},\lambda_{2}\in\mathbb{R}$ are Lagrangian multipliers jointly
updated during the pruning process. We compute $s_{e}$ by employing the
sigmoid differentiable form of $z_{i,j}^{mask}$.

Model | Pars | SST-2 | QNLI | MNLI | QQP | CoLA | RTE | STS-B | MRPC | Avg | Prune Time
---|---|---|---|---|---|---|---|---|---|---|---
| (Mil) | (Acc) | (Acc) | (Acc) | (Acc) | (M. Corr) | (Acc) | (P. Corr) | (Acc) | Score | (GPU hrs)
BERTbase(teacher) | 109 | 93.1 | 91.5 | 84.5 | 91.1 | 61.2 | 79.8 | 89.3 | 88.3 | 84.8 | -
DynaBERT (Hou et al., 2020) | 34 | 89.4 | 83.9 | 76.3 | 88.8 | 15.8 | 66.4 | 85.4 | 84.6 | 73.8 |  $\sim$ 15
BlockPrune (Lagunas et al., 2021) | 28 | 89.3 | 84.7 | 78.8 | 88.8 | 15.8 | 66.4 | 85.4 | 81.9 | 73.9 |  $\sim$20
CoFi (Xia et al., 2022) | 28 | 90.6 | 86.1 | 80.6 | 90.1 | 35.6 | 64.7 | 83.1 | 82.6 | 76.7 |  $\sim$33
PostPrune (Kwon et al., 2022) | 28 | 75.7 | 63.8 | 45.5 | 67.5 | 32.4 | 58.6 | 80.2 | 78.7 | 62.8 |  $\sim$0.08
VTrans (ours) | 28∗ | 92.2 | 89.7 | 83.8 | 90.5 | 43.3 | 73.4 | 87.9 | 87.0 | 81.0 |  $\sim$20
Fast-VTrans (ours) | 28∗ | 90.8 | 85.8 | 79.3 | 89.2 | 37.6 | 68.0 | 87.9 | 83.8 | 77.8 |  $\sim$ 2
Faster-VTrans (ours) | 28∗ | 82.3 | 81.8 | 72.1 | 81.4 | 36.3 | 67.6 | 86.8 | 82.5 | 73.9 |  $\sim$ 0.7
RoBERTabase(teacher) | 124 | 95.3 | 93.2 | 87.7 | 73.8 | 62.0 | 78.7 | 91.2 | 90.1 | 84.0 | -
MINILMv26 (Wang et al., 2020b) | 82 | 93.5 | 91.6 | 84.3 | 72.8 | 57.8 | 72.1 | 87.5 | 88.2 | 81.0 |  $\sim$300
FeatureCorr (Huang et al., 2023) | 82 | 93.8 | 92.0 | 85.6 | 73.5 | 60.3 | 72.5 | 88.3 | 89.6 | 82.0 | -
VTrans (ours) | 82∗ | 94.8 | 93.0 | 87.3 | 73.6 | 61.0 | 73.1 | 90.9 | 90.1 | 83.0 |  $\sim$25
Fast-VTrans (ours) | 82∗ | 94.1 | 91.9 | 85.4 | 72.9 | 60.4 | 72.7 | 89.6 | 89.8 | 82.1 |  $\sim$ 2.8
Faster-VTrans (ours) | 82∗ | 93.4 | 90.6 | 84.8 | 72.3 | 59.2 | 72.6 | 87.9 | 88.5 | 81.2 |  $\sim$ 1
GPT2small(teacher) | 125 | 92.7 | 88.5 | 82.3 | 89.7 | 47.4 | 74.4 | 88.9 | 85.8 | 81.2 | -
DistilGPT2 (Sanh et al., 2020a) | 82 | 90.7 | 87.9 | 81.6 | 66.8 | 39.4 | 68.3 | 79.6 | 87.9 | 75.3 |  $\sim$ 300
FeatureCorr (Huang et al., 2023) | 82 | 92.0 | 88.5 | 83.4 | 70.6 | 42.3 | 70.2 | 81.6 | 88.4 | 77.1 | -
VTrans (ours) | 82∗ | 92.4 | 88.8 | 82.8 | 89.5 | 45.4 | 72.9 | 88.1 | 86.0 | 80.7 |  $\sim$25
Fast-VTrans (ours) | 82∗ | 92.0 | 88.4 | 82.8 | 88.7 | 43.5 | 71.4 | 86.5 | 86.0 | 80.0 |  $\sim$2.8
Faster-VTrans (ours) | 82∗ | 91.9 | 87.3 | 80.1 | 86.2 | 42.1 | 70.8 | 86.5 | 86.0 | 78.9 |  $\sim$1

Table 2:  The performance of models on GLUE dataset. ∗ indicates the average
number of parameters across all datasets. Average variance in performance of
our models is about $\pm 0.4\%$ obtained over five random initial seeds.

###  4.2 Knowledge Distillation

We distill knowledge from the original full-sized Transformer to the student
model during pruning by minimizing the cross entropy between their output
probability distributions $\boldsymbol{p}_{s}$ and $\boldsymbol{p}_{t}$ as
$\mathcal{L}_{\text{pred}}=D_{\mathrm{KL}}\left(\boldsymbol{p}_{s}\|\boldsymbol{p}_{t}\right)$.
We also adopt a layer-wise distillation approach inspired by Xia et al.
(2022). This method aligns the teacher’s layer with the evolving student’s,
defining an intermediate distillation loss function as
$\mathcal{L}_{\text{layer}}=\sum_{i\in\mathcal{T}}\operatorname{MSE}\left(W_{\text{layer}}\boldsymbol{H}_{s}^{m(i)},\boldsymbol{H}_{t}^{i}\right)$
where $W_{\text{layer }}$ is a transformation matrix initialized to identity.
$\boldsymbol{H}_{s}^{m(i)}$ and $\boldsymbol{H}_{t}^{i}$ represent the
student’s and teacher’s hidden layer representations, respectively, for layer
distillation, with $m(\cdot)$ mapping teacher layer representations to the
closest student layer as $m(i)=\underset{j:\text{unpruned
layers}}{\arg\min}\operatorname{MSE}\left(W_{\text{layer
}}\boldsymbol{H}_{s}^{j},\boldsymbol{H}_{t}^{i}\right)$. Hence, the
distillation loss encompasses both the layer-wise loss and the logit-based
cross entropy loss with $\eta$ controlling their relative influence:
$\mathcal{L}_{\text{distil
}}=\eta\mathcal{L}_{\text{pred}}+(1-\eta)\mathcal{L}_{\text{layer }}$. The
comprehensive loss function to be minimised during pruning is
$\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{distil
}}+\tilde{\mathcal{L}}_{\text{VIB}}+\mathcal{L}_{s}$. The steps of the our
method is reiterated in Algorithm 1 in the Appendix.

###  4.3 Faster variants of our method

Faster compression with Data Subset. Leveraging the valuable features and
patterns learned during the initial training phase, our Fast-VTrans approach
involves training VIB-masks and refining un-masked pre-trained model
parameters with a significantly reduced data requirement as shown in figure
1(b). Intuitively, this process is expected to facilitate adaptation to lower
model parameters without overfitting. To achieve this, we employ a 3% data
subset for pruning and fine-tuning, balancing efficiency and adaptability.

VIB mask training. Motivated by the recent advancements in post-training
pruning methodologies (Kwon et al., 2022), we propose the fastest variant,
Faster-VTrans that freezes all pre-trained model parameters(Figure 1(c)).
During pruning, only VIB masks are trained to target irrelevant elements.
Additionally, all norm and bias terms (where available) are fine-tuned to
compensate for information loss. This approach, using only 3% of the data,
significantly reduces compression time while maintaining performance, making
it ideal for scenarios with limited time or resources.

Figure 3: Pruning BERT-base (teacher model) on GLUE and SQuAD tasks: Each
point denotes the mean performance of pruned models averaged over five random
trials with different seeds with average variance of $\pm 0.4\%$ across all
models. The dashed line represents a 1% performance drop from the teacher
model.

##  5 Experiments

###  5.1 Implementation Details

Datasets. We assess the efficacy of our proposed method across a comprehensive
set of linguistic understanding tasks - the General Language Understanding
Evaluation (GLUE) benchmark (Wang et al., 2018), Stanford Question Answering
Dataset (SQuAD) version 1.1 (Rajpurkar et al., 2016) and WikiText-2 (Merity et
al., 2016).

Baseline Models. We compare our method variants with several previous
techniques: DynaBERT (Hou et al., 2020), Block Pruning (Lagunas et al., 2021),
PostPrune (Kwon et al., 2022), CoFi (Xia et al., 2022), FeatureCorr (Huang et
al., 2023), DistilBERT (Sanh et al., 2020a), TinyBERT-GD (Jiao et al., 2019),
MiniLM (Wang et al., 2020b), and HomoBERT (Liang et al., 2023a). We also
compare our pruned LLaMA-2 models with previous semi-structured pruning
techniques- SparseGPT (Frantar & Alistarh, 2023), Wanda (Sun et al., 2023) and
structured pruning techniques- LLM-pruner (Ma et al., 2023), Bonsai (Dery et
al., 2024). Details are deferred to Appendix A.2.

Training Details. The primary method utilizes entire datasets for pruning and
finetuning, while the Fast and Faster variants operate on a randomly sampled
subset, approximately 3% of the data. For Fast and Faster variants, 8000
samples are used for pruning and finetuning on large datasets and 2000 on
smaller ones of GLUE and SQuAD. Finetuning uses 16000 samples for large
datasets and the entire dataset for smaller ones. We conduct five runs with
random seeds for each sparsity constraint, reporting average performance.
Experiments conducted to come to the final hyper-parameter settings are
provided in Appendix A.5.

###  5.2 Evaluation

Performance comparison and Compression speedup. In Table 2, we compare pruned
models derived from our primary method and its faster variants (Fast and
Faster) with alternative distillation and pruning approaches. Our process
starts from three teachers - BERTbase, ROBERTa, and GPT-2 small. To ensure a
fair comparison of performance on GLUE datasets and compression time (measured
on the largest dataset MNLI) with previous methods that share the same teacher
model, we categorize the obtained models into three sections. In the uppermost
group, we compare models with 28 million parameters and 75% sparsity to match
previous smallest models (Xia et al., 2022). Our primary method outperforms
prior techniques with similar order of compression time. Fast-VTrans shows a
16x acceleration compared to CoFi, and Faster-VTrans is at least 20x faster
than DynaBERT and BlockPrune, with similar performance. Although PostPrune
(Kwon et al., 2022) takes lower compression time, our method yields stable
performance across different seed initialization (Appendix B.2). In the middle
group, our 66% sparse model by VTrans, outperforms both MINILMv2 and
FeatureCorr, which utilize six-layer student models, with a 1% higher GLUE
score. Additionally, the faster variants demonstrate comparable performance
with substantially reduced compression time. Finally, in the bottom group,
when applied to the decoder-based GPT-2-small model, our compression technique
demonstrates superior performance over previous approaches with minimal
compression time. Furthermore, our pruned model compressed from a larger
model- uncased BERT-large- outperforms DistilBERT and MINILMv2 on two GLUE
tasks while achieving comparable results on the remaining tasks (Appendix
B.3).

Model | Parameters | SST-2 | QQP | MRPC | MNLI
---|---|---|---|---|---
| (Million) | (Acc) | (Acc) | (Acc) | (Acc)
BERTbase | 109 | 93.1 | 91.1 | 88.3 | 84.5
TinyBERT4 | 14.5 | 89\. 7 | 90.0 | 81.4 | 80.4
MiniLM3 | 17.0 | 89.3 | 88.8 | 81.9 | 78.8
HomoBERT | 14.1 | 90.1 | 89.9 | 87.3 | 81.2
VTrans w GD | 14.5 | 91.8 | 90.3 | 87.8 | 80.6

Table 3: The performance of models pruned using task-agnostic general
distillation (GD) on GLUE: The pruned models are fine-tuned on downstream
tasks. TinyBERT results are reported without data augmentation for fair
comparison.

Model Size and Performance comparison. In Figure 3, we compare the performance
of the models obtained with our task-specific approaches with previous methods
across various sparsity levels in terms of relative model parameters. We note
that our primary approach is able to achieve about 70 to 80% compressed models
with less than 1% loss in accuracy from the un-pruned teacher model over all
the tasks. It achieves about 20 to 70% higher sparsity than other pruning
methods like CoFi, BlockPruning and DynaBERT as it can prune embedding states
along with intermediate layers, heads and hidden units unlike others. The
faster variants of our method retain performance comparable to previous
approaches but with 10 to 20x speedup in compression as seen in Table 2.

Task-agnostic performance. To evaluate our primary method on task-agnostic
setting, we use general distillation during pruning and obtained a compressed
model of similar size as obtained by previous approaches (Liang et al., 2023a;
Jiao et al., 2019). We finetune the model on each of the tasks in GLUE and
report results of four of the tasks in Table 3. Our model performs better on
three of the four tasks than other models obtained with previous task-agnostic
approaches such as TinyBERT, MiniLM and HomoBERT.

| QNLI (60%) | MNLI (60%) | QNLI (95%) | MNLI (95%) | SQuAD (60%) | SQuAD (95%)
---|---|---|---|---|---|---
| Speedup | Acc | Speedup | Acc | Speedup | Acc | Speedup | Acc | Speedup | Acc | Speedup | Acc
Our Method | 1.5$\times$ | 91.5 | 1.5$\times$ | 84.3 | 9$\times$ | 81.9 | 9$\times$ | 74.2 | 1.4$\times$ | 87.0 | 7$\times$ | 79.6
-Embeddings | 1.7$\times$ | 90.2 | 1.5$\times$ | 83.9 | 9.2$\times$ | 80.6 | 9.4$\times$ | 73.5 | 1.4$\times$ | 84.3 | 7.5$\times$ | 77.6
-Layers | 1.5$\times$ | 91.5 | 1.5$\times$ | 84.2 | 7$\times$ | 81.9 | 7$\times$ | 73.8 | 1.4$\times$ | 85.8 | 5$\times$ | 78.0
-Embeddings & Layers | 1.5$\times$ | 90.2 | 1.5$\times$ | 83.9 | 6$\times$ | 80.4 | 6$\times$ | 72.6 | 1.4$\times$ | 84.5 | 5$\times$ | 75.8

Table 4: Investigating Pruning Strategies in BERT for varied sparsity
constraints in GLUE and SQuAD Tasks: Removing embedding masks results in a
performance decline. Removing layer masks reduces speedup, while removing both
increases pruning of hidden states and heads, ultimately leading to diminished
performance and reduced speedup in high sparsity models.

###  5.3 Feasibility of Compression with Faster Variants

Figure 4: Speeding-up Compression: Faster-VTrans achieves significant speed-up
over VTrans on GLUE tasks, with minimal performance loss for models 35%
smaller than GPT-2small.

Influence of data-subset on performance. To explore the relationship between
data sample size and compression performance, we conducted pruning experiments
across GLUE tasks. In our trials, we used 8000 samples for larger datasets
like MNLI, QQP, SST2, and QNLI, and 2000 samples for smaller datasets such as
MRPC, STSB, CoLA, and RTE. Interestingly, increasing sample size for larger
datasets in the faster variant led to a 1% performance boost, but at the
expense of compression times 2 to 4x longer (Appendix A.3). This suggests the
viability of compressing pre-trained models with limited datasets, while
maintaining crucial characteristics and preserving performance to a
significant degree.

Pruning Speedup in Decoder Models. In order to assess the compression speedup
for decoder models, we employ our primary method and the Faster variant to
compress GPT-2small. We compare the performance and speedup with both the
methods across four GLUE datasets. Notably, on large datasets such as MNLI and
SST-2, as depicted in Figure 4, there is a substantial speedup (ranging from
11 to 25x) in compression achieved with Faster-VTrans, while incurring less
than 3% accuracy degradation compared to our primary method.

Method | WikiText2↓ | BoolQ | PIQA | HellaSwag | WinoGrande | ARC-e | ARC-c | Average↑ | Prune
---|---|---|---|---|---|---|---|---|---
| PPL | Acc | Acc | Acc | Acc | Acc | Acc | Acc | Time(hrs)
Unpruned model | 5.68 | 76.5 | 79.8 | 76.1 | 70.1 | 72.8 | 47.6 | 70.48 | -
LLM-Pruner† | 16.41 | 60.28 | 69.31 | 47.06 | 53.43 | 45.96 | 29.18 | 45.95 | 1
LoRAPrune† | 11.60 | 61.88 | 71.53 | 47.86 | 55.01 | 45.13 | 31.62 | 52.17 | $>24$
Bonsai† | 10.92 | 67.22 | 60.85 | 43.09 | 61.64 | 54.92 | 26.28 | 52.33 | 40
Faster-VTrans† | 10.58 | 63.27 | 68.56 | 42.0 | 57.38 | 56.97 | 26.46 | 52.44 | 2

Table 5: Zero shot performance of the compressed LLaMA-7B. †Finetuned with
LoRA . Our method takes 10 to 20 times lower time to prune than other methods
with similar performance. Method with lower compression time have much lower
average performance.

###  5.4 Scaling-up to prune more than billion parameters

We apply our Faster-VTrans variant without distillation to remove redundant
units from the LLaMA-2 7 billion model. Table 1 displays our model’s
(1.6$\times$) higher inference speedup compared to Wanda (Sun et al., 2023), a
semi-structured pruning method yielding sparse models not suitable for most
hardware. Compared to LLM-Pruner (Ma et al., 2023) and Bonsai (Dery et al.,
2024), our model achieves (over 1.2$\times$) higher inference speedup with
significantly better performance (11.8 vs 19.24 perplexity). Pruning takes
only 2 hours on a single GPU, about half the time of Bonsai. Additionally,
low-rank weight adaptation (Hu et al., 2021) enables us to fine-tune all
modules within the same timeframe. Fine-tuning results are in Appendix B.1. We
also test our 50% pruned LLaMA on zero-shot reasoning tasks shown in Table 5.
Overall, our pruned model outperforms other techniques with significantly
lower prune time.

###  5.5 Ablation Study

Effect of pruning different structures. To investigate the effect of pruning
embedding states unlike other previous state-of-the-art approaches, we conduct
ablation studies. On removing the choice to prune the embedding states, as
seen in Table 4, the performance decreases across all tasks. This is
attributed to the the increased pruning of other model structures - the
layers, hidden states and the attention heads. On not pruning the layers, the
models with higher sparsity levels suffer most in terms of speedup while
performance degrades slightly. When both the embeddings and layers are not
pruned, the algorithm prunes more attention heads and hidden states leading to
lower performance of models at 60% sparsity levels with higher sparsity (95%)
levels affected most in terms of both speedup and performance.

| QNLI | MNLI | SQuAD
---|---|---|---
| (60%) | (95%) | (60%) | (95%) | (60%) | (95%)
Our Method | 91.5 | 81.9 | 84.3 | 74.2 | 87.0 | 79.6
-$\mathcal{L}_{\text{layer }}$ | 90.2 | 80.4 | 83.4 | 73.6 | 86.5 | 76.3
-$\mathcal{L}_{\text{pred }},\mathcal{L}_{\text{layer }}$ | 89.8 | 79.6 | 82.9 | 72.4 | 86.2 | 75.5

Table 6: Ablation Experiments with and without distillation losses at
different model sparsity levels.

Importance of distillation. We analyze the impact of the distillation
objectives in conjunction with our VIB-based pruning approach on pruned model
performance in Table 6. We note that ablating layer-wise knowledge
distillation from dynamically matched teacher layers to student layers results
in a performance drop of 1 to 3 points. Completely ablating knowledge
distillation leads to a performance drop of 2 to 4 points across all datasets.
This shows that distillation helps improve the pruning performance of our VIB-
based method by retaining performance of the student models during the
iterative pruning process.

Figure 5: (a) Comparative attention allocation towards special tokens in un-
pruned (top) and pruned (bottom) models. Dots represent token attention by
heads, with lines indicating mean attention by remaining heads. Pruned model
exhibits reduced special token attention (b) Eliminated heads example: those
highly attentive to CLS, SEP, current token, common articles like ’the’ and
displaying broad attention. Line thickness reflects attention weight towards
the token.

##  6 Qualitative Analysis

Computing Average Attention in Heads. We evaluate information preservation and
redundancy reduction in pruned models while maintaining performance by
analyzing attention heads in both pre-trained BERT and an equivalently
performing pruned model. On the SST-2 downstream task, we compare the un-
pruned BERT base model with a 50% smaller sized model. Our examination
includes computing average attention directed towards special tokens (SEP,
CLS, periods, and commas) as done by Clark et al. (2019). In Figure 5(a), we
observe significant attention allocation to special tokens in the un-pruned
model. Clark et al. (2019) speculate that attention to these tokens might
serve as ’no-ops’ when a head’s function is not applicable, making such heads
viable candidates for pruning. In the 50% pruned model, we confirm this
observation, noticing a reduction in attention heads in the latter layers
(10-12) and a drastic decrease in average attention towards these tokens in
the remaining heads. Additionally, we observe a shift in average attention in
each layer from the current token to the previous and next token in the pruned
model (see Appendix D.3).
Redundancy in Attention Heads. We visualize attention
heads111https://github.com/jessevig/bertviz to identify redundant and
preserved language aspects in pruned models. In Figure 5(b) displaying
attention weights, we observe that heads with high focus on CLS or SEP tokens
and excessive attention to the current token are often pruned. Similarly,
heads with broad attention or a focus on common articles like ’the’ or ’a’ are
removed. Moreover, we observe that our method aligns with the interpretation
by Burgess et al. (2018) suggesting that VIB encourages the acquisition of
more disentangled representations (see Appendix D.4).

##  7 Conclusion

Integrating the Variational Information Bottleneck principle, our structured
pruning method efficiently prunes redundant units across all structural levels
in transformers. This results in upto 70% higher compression compared to prior
methods, accompanied by minimal accuracy loss ($<$1%) on GLUE and SQuAD tasks,
along with up to 4x inference speedup. Unlike previous approaches, we address
both parameter and FLOPs constraints with our task-specific distillation
approach, requiring significantly lower training resources (upto 60x less)
than task-agnostic methods while achieving similar performance. Our efficient
compression variants, namely Fast and Faster-VTrans, deliver over 10x pruning-
speedup with performance comparable to previous state-of-the-art approaches.
Our evaluation encompasses various architectures, including decoder-based GPT
models. Demonstrating the effectiveness and scalability of our method, we
compress the large LLaMA-2 model by 50% with better performance (7 points) and
enhanced inference speed (1.2$\times$) than other structured pruning methods
with lower compression time. Our qualitative analysis supports the
competitiveness of our approach, with potential for further quantification in
future work.

## References

  * Aguilar et al. (2020) Gustavo Aguilar, Yuan Ling, Yu Zhang, Benjamin Yao, Xing Fan, and Chenlei Guo.  Knowledge distillation from internal representations.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pp. 7350–7357, 2020.
  * Alemi et al. (2016) Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy.  Deep variational information bottleneck.  _arXiv preprint arXiv:1612.00410_ , 2016.
  * Bai et al. (2021) Haoli Bai, Wei Zhang, Lu Hou, Lifeng Shang, Jing Jin, Xin Jiang, Qun Liu, Michael Lyu, and Irwin King.  Binarybert: Pushing the limit of bert quantization, 2021.
  * Bar-Haim et al. (2006) Roy Bar-Haim, Ido Dagan, Bill Dolan, Lisa Ferro, and Danilo Giampiccolo.  The second pascal recognising textual entailment challenge.  _Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment_ , 01 2006.
  * Bentivogli et al. (2009) Luisa Bentivogli, Bernardo Magnini, Ido Dagan, Hoa Trang Dang, and Danilo Giampiccolo.  The fifth PASCAL recognizing textual entailment challenge.  In _Proceedings of the Second Text Analysis Conference, TAC 2009, Gaithersburg, Maryland, USA, November 16-17, 2009_. NIST, 2009.  URL https://tac.nist.gov/publications/2009/additional.papers/RTE5_overview.proceedings.pdf.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Burgess et al. (2018) Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner.  Understanding disentangling in $\backslash beta$-vae.  _arXiv preprint arXiv:1804.03599_ , 2018.
  * Cer et al. (2017) Daniel Cer, Mona Diab, Eneko Agirre, Iñigo Lopez-Gazpio, and Lucia Specia.  SemEval-2017 task 1: Semantic textual similarity multilingual and crosslingual focused evaluation.  In _Proceedings of the 11th International Workshop on Semantic Evaluation (SemEval-2017)_ , pp. 1–14, Vancouver, Canada, August 2017. Association for Computational Linguistics.  doi: 10.18653/v1/S17-2001.  URL https://aclanthology.org/S17-2001.
  * Chen et al. (2021) Cheng Chen, Yichun Yin, Lifeng Shang, Zhi Wang, Xin Jiang, Xiao Chen, and Qun Liu.  Extract then distill: Efficient and effective task-agnostic bert distillation, 2021.
  * Clark et al. (2019) Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning.  What does BERT look at? an analysis of BERT’s attention.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pp. 276–286, Florence, Italy, August 2019. Association for Computational Linguistics.  doi: 10.18653/v1/W19-4828.  URL https://aclanthology.org/W19-4828.
  * Dagan et al. (2006) Ido Dagan, Oren Glickman, and Bernardo Magnini.  The pascal recognising textual entailment challenge.  In Joaquin Quiñonero-Candela, Ido Dagan, Bernardo Magnini, and Florence d’Alché Buc (eds.), _Machine Learning Challenges. Evaluating Predictive Uncertainty, Visual Object Classification, and Recognising Tectual Entailment_ , pp. 177–190, Berlin, Heidelberg, 2006. Springer Berlin Heidelberg.  ISBN 978-3-540-33428-6.
  * Dai et al. (2018) Bin Dai, Chen Zhu, Baining Guo, and David Wipf.  Compressing neural networks using the variational information bottleneck.  In _International Conference on Machine Learning_ , pp. 1135–1144. PMLR, 2018.
  * Dery et al. (2024) Lucio Dery, Steven Kolawole, Jean-Francois Kagey, Virginia Smith, Graham Neubig, and Ameet Talwalkar.  Everybody prune now: Structured pruning of llms with only forward passes.  _arXiv preprint arXiv:2402.05406_ , 2024.
  * Dolan & Brockett (2005) William B. Dolan and Chris Brockett.  Automatically constructing a corpus of sentential paraphrases.  In _Proceedings of the Third International Workshop on Paraphrasing (IWP2005)_ , 2005.  URL https://aclanthology.org/I05-5002.
  * Dutta et al. (2023) Oshin Dutta, Tanu Kanvar, and Sumeet Agarwal.  Search-time efficient device constraints-aware neural architecture search.  _arXiv preprint arXiv:2307.04443_ , 2023.
  * Fan et al. (2020) Angela Fan, Edouard Grave, and Armand Joulin.  Reducing transformer depth on demand with structured dropout.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=SylO2yStDr.
  * Frantar & Alistarh (2023) Elias Frantar and Dan Alistarh.  Sparsegpt: Massive language models can be accurately pruned in one-shot.  In _International Conference on Machine Learning_ , pp. 10323–10337. PMLR, 2023.
  * Giampiccolo et al. (2007) Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan.  The third PASCAL recognizing textual entailment challenge.  In _Proceedings of the ACL-PASCAL Workshop on Textual Entailment and Paraphrasing_ , pp. 1–9, Prague, June 2007. Association for Computational Linguistics.  URL https://aclanthology.org/W07-1401.
  * Han et al. (2015) Song Han, Jeff Pool, John Tran, and William Dally.  Learning both weights and connections for efficient neural network.  In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 28. Curran Associates, Inc., 2015.  URL https://proceedings.neurips.cc/paper_files/paper/2015/file/ae0eb3eed39d2bcef4622b2499a05fe6-Paper.pdf.
  * Henderson & Fehr (2022) James Henderson and Fabio Fehr.  A variational autoencoder for transformers with nonparametric variational information bottleneck.  _arXiv preprint arXiv:2207.13529_ , 2022.
  * Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network, 2015.
  * Hou et al. (2020) Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu.  Dynabert: Dynamic bert with adaptive width and depth.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 9782–9793. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/6f5216f8d89b086c18298e043bfe48ed-Paper.pdf.
  * Hu et al. (2021) E Hu, Y Shen, P Wallis, Z Allen-Zhu, Y Li, S Wang, L Wang, and W Chen.  Low-rank adaptation of large language models.  _arXiv_ , 2021.
  * Huang et al. (2023) Kun Huang, Xin Guo, and Meng Wang.  Towards efficient pre-trained language model via feature correlation distillation.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Jiao et al. (2019) Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu.  Tinybert: Distilling bert for natural language understanding.  _arXiv preprint arXiv:1909.10351_ , 2019.
  * Khanuja et al. (2021) Simran Khanuja, Melvin Johnson, and Partha Talukdar.  MergeDistill: Merging language models using pre-trained distillation.  In _Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021_ , pp. 2874–2887, Online, August 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.findings-acl.254.  URL https://aclanthology.org/2021.findings-acl.254.
  * Kim et al. (2021) Sehoon Kim, Amir Gholami, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer.  I-bert: Integer-only bert quantization.  _International Conference on Machine Learning (Accepted)_ , 2021.
  * Kruskal (1964) Joseph B Kruskal.  Multidimensional scaling by optimizing goodness of fit to a nonmetric hypothesis.  _Psychometrika_ , 29(1):1–27, 1964.
  * Kwon et al. (2022) Woosuk Kwon, Sehoon Kim, Michael W. Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gholami.  A Fast Post-Training Pruning Framework for Transformers, October 2022.  URL http://arxiv.org/abs/2204.09656.  arXiv:2204.09656 [cs].
  * Lagunas et al. (2021) François Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush.  Block pruning for faster transformers.  In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_ , pp. 10619–10629, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.emnlp-main.829.  URL https://aclanthology.org/2021.emnlp-main.829.
  * Li et al. (2023) Yixiao Li, Yifan Yu, Qingru Zhang, Chen Liang, Pengcheng He, Weizhu Chen, and Tuo Zhao.  Losparse: Structured compression of large language models based on low-rank and sparse approximation.  _arXiv preprint arXiv:2306.11222_ , 2023.
  * Liang et al. (2021) Chen Liang, Simiao Zuo, Minshuo Chen, Haoming Jiang, Xiaodong Liu, Pengcheng He, Tuo Zhao, and Weizhu Chen.  Super tickets in pre-trained language models: From model compression to improving generalization, 2021.
  * Liang et al. (2023a) Chen Liang, Haoming Jiang, Zheng Li, Xianfeng Tang, Bin Yin, and Tuo Zhao.  HomoDistil: Homotopic Task-Agnostic Distillation of Pre-trained Transformers, February 2023a.  URL http://arxiv.org/abs/2302.09632.  arXiv:2302.09632 [cs].
  * Liang et al. (2023b) Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, and Tuo Zhao.  Less is more: Task-aware layer-wise distillation for language model compression.  In _International Conference on Machine Learning_ , pp. 20852–20867. PMLR, 2023b.
  * Ma et al. (2023) Xinyin Ma, Gongfan Fang, and Xinchao Wang.  Llm-pruner: On the structural pruning of large language models.  _arXiv preprint arXiv:2305.11627_ , 2023.
  * McCarley et al. (2021) J. S. McCarley, Rishav Chakravarti, and Avirup Sil.  Structured pruning of a bert-based question answering model, 2021.
  * Merity et al. (2016) Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher.  Pointer sentinel mixture models.  _arXiv preprint arXiv:1609.07843_ , 2016.
  * Michel et al. (2019) Paul Michel, Omer Levy, and Graham Neubig.  Are sixteen heads really better than one?  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper_files/paper/2019/file/2c601ad9d2ff9bc8b282670cdd54f69f-Paper.pdf.
  * Nasery et al. (2023) Anshul Nasery, Hardik Shah, Arun Sai Suggala, and Prateek Jain.  End-to-end neural network compression via $\frac{\ell_{1}}{\ell_{2}}$ regularized latency surrogates.  _arXiv preprint arXiv:2306.05785_ , 2023.
  * Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551, 2020.
  * Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang.  Squad: 100,000+ questions for machine comprehension of text, 2016.
  * Renda et al. (2020) Alex Renda, Jonathan Frankle, and Michael Carbin.  Comparing rewinding and fine-tuning in neural network pruning, 2020.
  * Romero et al. (2015) Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio.  Fitnets: Hints for thin deep nets, 2015.
  * Sajjad et al. (2020) Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov.  Poor man’s bert: Smaller and faster transformer models.  _ArXiv_ , abs/2004.03844, 2020.
  * Sanh et al. (2020a) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter, 2020a.
  * Sanh et al. (2020b) Victor Sanh, Thomas Wolf, and Alexander Rush.  Movement pruning: Adaptive sparsity by fine-tuning.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 20378–20389. Curran Associates, Inc., 2020b.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/eae15aabaa768ae4a5993a8a4f4fa6e4-Paper.pdf.
  * Slonim & Tishby (1999) Noam Slonim and Naftali Tishby.  Agglomerative information bottleneck.  In S. Solla, T. Leen, and K. Müller (eds.), _Advances in Neural Information Processing Systems_ , volume 12. MIT Press, 1999.  URL https://proceedings.neurips.cc/paper_files/paper/1999/file/be3e9d3f7d70537357c67bb3f4086846-Paper.pdf.
  * Socher et al. (2013) Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts.  Recursive deep models for semantic compositionality over a sentiment treebank.  In _Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing_ , pp. 1631–1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics.  URL https://aclanthology.org/D13-1170.
  * Srivastava et al. (2021) Ayush Srivastava, Oshin Dutta, Jigyasa Gupta, Sumeet Agarwal, and Prathosh AP.  A variational information bottleneck based method to compress sequential networks for human action recognition.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision_ , pp. 2745–2754, 2021.
  * Sun et al. (2023) Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter.  A simple and effective pruning approach for large language models.  _arXiv preprint arXiv:2306.11695_ , 2023.
  * Sun et al. (2019) S. Sun, Yu Cheng, Zhe Gan, and Jingjing Liu.  Patient knowledge distillation for bert model compression.  In _Conference on Empirical Methods in Natural Language Processing_ , 2019.
  * Sun et al. (2020) Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou.  MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pp. 2158–2170, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.195.  URL https://aclanthology.org/2020.acl-main.195.
  * Tang et al. (2019) Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin.  Distilling task-specific knowledge from bert into simple neural networks, 2019.
  * Tishby & Zaslavsky (2015) Naftali Tishby and Noga Zaslavsky.  Deep learning and the information bottleneck principle.  In _2015 IEEE Information Theory Workshop (ITW)_ , pp. 1–5. IEEE, 2015.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample.  Llama: Open and efficient foundation language models, 2023.
  * Turc et al. (2019) Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Well-read students learn better: On the importance of pre-training compact models, 2019.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Voita et al. (2019) Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov.  Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned, June 2019.  URL http://arxiv.org/abs/1905.09418.  arXiv:1905.09418 [cs].
  * Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman.  GLUE: A multi-task benchmark and analysis platform for natural language understanding.  In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pp. 353–355, Brussels, Belgium, November 2018. Association for Computational Linguistics.  doi: 10.18653/v1/W18-5446.  URL https://aclanthology.org/W18-5446.
  * Wang et al. (2020a) Wenhui Wang, Hangbo Bao, Shaohan Huang, Li Dong, and Furu Wei.  Minilmv2: Multi-head self-attention relation distillation for compressing pretrained transformers.  _arXiv preprint arXiv:2012.15828_ , 2020a.
  * Wang et al. (2020b) Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou.  Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers.  _Advances in Neural Information Processing Systems_ , 33:5776–5788, 2020b.
  * Wang et al. (2020c) Ziheng Wang, Jeremy Wohlwend, and Tao Lei.  Structured Pruning of Large Language Models.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pp. 6151–6162, 2020c.  doi: 10.18653/v1/2020.emnlp-main.496.  URL http://arxiv.org/abs/1910.04732.  arXiv:1910.04732 [cs, stat].
  * Warstadt et al. (2019) Alex Warstadt, Amanpreet Singh, and Samuel R. Bowman.  Neural network acceptability judgments.  _Transactions of the Association for Computational Linguistics_ , 7:625–641, 2019.  doi: 10.1162/tacl˙a˙00290.  URL https://aclanthology.org/Q19-1040.
  * Williams et al. (2018) Adina Williams, Nikita Nangia, and Samuel Bowman.  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)_ , pp. 1112–1122. Association for Computational Linguistics, 2018.  URL http://aclweb.org/anthology/N18-1101.
  * Xia et al. (2022) Mengzhou Xia, Zexuan Zhong, and Danqi Chen.  Structured pruning learns compact and accurate models.  In _Association for Computational Linguistics (ACL)_ , 2022.
  * Yang et al. (2022) Nakyeong Yang, Yunah Jang, Hwanhee Lee, Seohyeong Jung, and Kyomin Jung.  Task-specific compression for multi-task language models using attribution-based pruning.  _arXiv preprint arXiv:2205.04157_ , 2022.
  * Zafrir et al. (2019) Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat.  Q8bert: Quantized 8bit bert.  In _2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing - NeurIPS Edition (EMC2-NIPS)_. IEEE, December 2019.  doi: 10.1109/emc2-nips53020.2019.00016.  URL http://dx.doi.org/10.1109/EMC2-NIPS53020.2019.00016.
  * Zafrir et al. (2021) Ofir Zafrir, Ariel Larey, Guy Boudoukh, Haihao Shen, and Moshe Wasserblat.  Prune once for all: Sparse pre-trained language models, 2021.
  * Zhu & Gupta (2017) Michael Zhu and Suyog Gupta.  To prune, or not to prune: exploring the efficacy of pruning for model compression, 2017.

Algorithm 1 Our primary method variant - Identifying redundant units with VIB-
based masks with user-defined parameters or FLOPs constraint

Input: Pruning metric (parameters or FLOPs), target metric $t$, Teacher Model

Initialize:VIB masks $z_{m},z_{a},z_{out},z_{layer}$

for $e={1,...,Epochs}$ do

Calculate VIB loss $\tilde{\mathcal{L}}_{\text{VIB}}$

Calculate expected sparsity $s_{e}$

Calculate sparsity loss $\mathcal{L}_{s}$

Total loss
$\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{distil}}+\tilde{\mathcal{L}}_{\text{VIB}}+\mathcal{L}_{s}$

Update model $f(\cdot;\theta_{s})$ weights $\theta_{s}$

Update $z$ masks

Update $\lambda_{1},\lambda_{2}$

end for

Convert $z$ masks to binary masks $\hat{z}_{i,j}^{mask}$

Fine-tune VIB mask-selected model weights by minimizing
$\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{distil}}+\tilde{\mathcal{L}}_{\text{VIB}}$

Prune VIB mask-selected units from model

$\hat{f}(\cdot;\theta_{s})\leftarrow f(\cdot;\theta_{s})$
$\odot\hat{z}_{i,j}^{mask}$

Output: Compressed model $\hat{f}(\cdot;\theta_{s})$

##  Appendix A Further Training details

###  A.1 Datasets

General Language Understanding Evaluation (GLUE) (Wang et al., 2018) is a
collection of nine natural language understanding tasks. The GLUE tasks
encompass various domains, including sentiment analysis (SST2, (Socher et al.,
2013)), natural language inference (MNLI, (Williams et al., 2018)), paraphrase
identification (QQP and QNLI), textual entailment (MRPC, (Dolan & Brockett,
2005)), linguistic acceptability (CoLA (Warstadt et al., 2019)), semantic
textual similarity (STS-B (Cer et al., 2017)), and recognizing textual
entailment (RTE (Dagan et al., 2006; Bar-Haim et al., 2006; Giampiccolo et
al., 2007; Bentivogli et al., 2009)).

Stanford Question Answering Dataset (SQuAD) (Rajpurkar et al., 2016) is a
reading comprehension dataset, consisting of questions posed on Wikipedia
articles, where the answer are a segment of the comprehension text.

WikiText-2 Created from Wikipedia articles Merity et al. (2016), this dataset
consists of 2 million tokens for training and more than 200,000 tokens in the
validation and test set.

Method | Pruning Setting | Pruning Granularity | Distillation Setting | Student Initialization
---|---|---|---|---
DistilBERT (Sanh et al., 2020a) | no pruning | - | TA | Un-trained weights
TinyBERT-GD (Jiao et al., 2019) | no pruning | - | TA | Un-trained weights
MiniLM (Wang et al., 2020b) | no pruning | - | TA | Un-trained weights
HomoDistil (Liang et al., 2023a) | prune while distill | Row, Column | TA | pre-trained weights
DynaBERT (Hou et al., 2020) | prune, then distill | FFN, Head | TS | Pruned, Pre-trained weights
Block Pruning (Lagunas et al., 2021) | prune, then distil | weight blocks | TS | Pruned, Pre-trained weights
Post-prune (Kwon et al., 2022) | prune | MHA, FFN | none | -
CoFi(Xia et al., 2022) | prune while distill | Head, FFN, Layers | TS | Pre-trained weights
Our method | prune while distill | Embeddings, Head, | TA/TS | Pre-trained/finetuned
|  | FFN, Layers |  | teacher weights
SparseGPT (Frantar & Alistarh, 2023) | prune | MHA, FFN | none | -
Wanda (Sun et al., 2023) | prune | MHA, FFN | none | -
LLM-Pruner (Ma et al., 2023) | prune | MHA, FFN | none | -
Bonsai (Dery et al., 2024) | prune | MHA, FFN | TS | Pre-trained compressed model
Our method (for LLM pruning) | prune | Embeddings, Head, | TS | Pre-trained
|  | FFN, Layers |  | compressed model

Table 7: Comparison of the baseline methods with our method in terms of
pruning and distillation techniques. TS - indicates task specific; TA - Task
agnostic distillation

###  A.2 Further details about the baseline models

In Table 7, we present a comprehensive comparative analysis of the various
baseline methods used for performance evaluation in Section 5.

DistilBERT (Sanh et al., 2020a) employs a vanilla distillation approach during
the pre-training phase by considering the distillation loss over the outputs
of the models. TinyBERT (Jiao et al., 2019) builds upon DistilBERT by
leveraging knowledge from intermediate Transformer layers. MiniLM (Wang et
al., 2020b) targets the discrepancy between queries-keys scaled dot product
and values-values scaled dot product in the final layer’s self-attention
module. Extending on MiniLM, MiniLMv2 (Wang et al., 2020a) emphasizes on the
emulation of attention head relations from the teacher model. Employing a
task-agnostic distillation approach, HomoDistil (Liang et al., 2023a) distils
knowledge into iteratively pruned students initialized to the teacher’s
weights. DynaBERT (Hou et al., 2020) offers flexibility to adjust the model
size and latency by selecting an adaptive width and depth. The training
involves an initial training of a width-adaptive model followed by enabling
adaptability in both width and depth, using distillation from the full-size
model. CoFi (Xia et al., 2022) adopts an iterative pruning strategy during the
finetuning phase, utilizing distillation from the larger task-specific teacher
models to the student. It performs simultaneously pruning of attention heads,
FFNs and whole layers. Movement Pruning (Sanh et al., 2020b) proposes a
deterministic first-order weight pruning method, that is effective in
finetuning regime of pre-trained models. BlockPruning Lagunas et al. (2021)
uses this method to prune blocks of weight matrices, allowing for better
optimizations on dense hardware. Post-Prune uses post-training pruning
framework (Aguilar et al., 2020) that proposes a Fisher-based technique to
train masks for identification of redundant neurons and selectively fine-tunes
only the masks to specific tasks.

We show the scaling ability of our pruning technique to prune large language
models like LLaMA-2 (Touvron et al., 2023). Comparison with previous
techniques includes SparseGPT (Frantar & Alistarh, 2023) which proposes a
second-order layer-wise pruning method that approximates closed form equations
thus being able to scale up pruning LLMs. Wanda (Sun et al., 2023) takes into
account the norm of weights and input activations for pruning weights in an
unstructured/structured manner. Bonsai (Dery et al., 2024) is a gradient-free
structured pruning method that estimates module importance perturbatively by
generating sub-models and evaluating their performances. LLM-pruner (Ma et
al., 2023) is a structured pruning method that uses gradient information to
prune large language models in a task-agnostic manner.

###  A.3 Influence of Dataset size on pruning and finetuning time

We prune with different number of samples from the datasets within GLUE with
our Fast-VTrans method as seen in Figure 6. For both the faster variants, we
use 8000 samples from large datasets (MNLI, QNLI, QQP, SST-2, SQuAD) and 2000
samples from the small datasets (MRPC, RTE, STSB, CoLA) during pruning.
Further increasing the samples leads to a 2x to 4x increase in prune time with
only about 1% boost in performance. During finetuning with the faster
variants, we use 16000 samples for large datasets and the whole small
datasets.

Figure 6: Data subset impact: Prune time rises exponentially with data
samples, while model performance stabilizes.

###  A.4 Influence of number of epochs

The number of epochs used to finetune the model is varied from 5 to 20 as seen
in Figure 7. For smaller GLUE tasks, we finetune the weights of the student
model with distillation for 20 epochs. For larger GLUE tasks, we finetune for
10 epochs with distillation. Increasing the number of epochs for larger
datasets (MNLI, SST-2) is seen to increase the time by 2x with less than 1%
increase in accuracy.

Figure 7: We finetune the models obtained with sparsity of 75% with different
number of epochs and evaluate the final performance.

###  A.5 Implementation Details of experiments with GLUE and SQuAD

We list the hyper-parameter setting for task-specific pruning and finetuning
with our method on GLUE and SQuAD tasks in Table 8. Experiments are performed
on a single NVIDIA V100 GPU (32GB). In task-specific setting, we finetune the
pre-trained BERT-base to particular tasks for 3 (large datasets) and 10 (small
datasets) epochs to obtain teachers as used in previously Hou et al. (2020).
BERT model previously finetuned on MNLI task is utilized to get teachers of
smaller GLUE datasets - RTE, MRPC and STS-B. These teacher yield high
performing models.

For VTrans and Fast-VTrans: We prune and finetune with our objective function
for 20 epochs to obtain smaller student models within the given sparsity
constraint. Epochs lower than 20 does not let the pruning algorithm converge
for higher sparsity levels(¿80%). For lower sparsity levels, the algorithm
converges within less than 10 epochs. For Faster-VTrans, we freeze the model
parameters while training only the VIB masks.

For our task-agnostic variant, we use similar hyperparameters used by Jiao et
al. (2019) and VIB parameters are kept as in Table 8. We use 8 NVIDIA V100 to
prune the BERT-base pre-trained model for 3 epochs. The pruned model is
finetuned for 20 epochs on small and 10 the large tasks of GLUE and SQuAD as
seen optimal in Figure 7.

Hyper-parameters | GLUE(small) | GLUE(large) | SQuAD
---|---|---|---
VIB learning rate | {1 to 2} $\times 10^{-2}$ | {1 to 2} $\times 10^{-2}$ | {1 to 2} $\times 10^{-3}$
Training Learning Rate | {2} $\times 10^{-5}$ | {2} $\times 10^{-5}$ | {3} $\times 10^{-5}$
Fine-tuning Learning Rate | {2 to 7} $\times 10^{-5}$ | {2 to 7} $\times 10^{-5}$ | {2 to 7} $\times 10^{-5}$
Batch Size | 32 | 32 | 16
Fine-tune Epochs | 20 | 10 | 3
Distillation $\eta$ | 0.5 | 0.5 | 0.5

Table 8: Hyper-parameter configurations for pruning and finetuning with our
method and its variants on GLUE and SQuAD datasets. Experiments are conducted
on NVIDIA V100 GPU.

VIB LR | Dataset size | block_size
---|---|---
5x$10^{-2}$ | 4000 | 512

((a))

weights LR | LoRA-rank | LoRA-$\alpha$ |  $\eta$ (distill Weight) | block_size
---|---|---|---|---
1x$10^{-4}$ | 128 | 4$\times$rank | 0.01 | 512

((b))

Table 9: Hyper-parameters for (a) pruning with Faster-VTrans and (b) for fine-
tuning compressed model on WikiText-2 Dataset

Figure 8: Training plots for the CoLA dataset under a 40% model size sparsity
constraint. Top left: illustrates the evolution of various loss objectives
over training iterations. Top right: the pruning trajectory stabilizes with
increasing iterations. Bottom left: depicts the trajectory of remaining
parameters in each layer. Bottom right: showcases different loss trajectories
during finetuning.

###  A.6 Implementation details with experiments on WikiText-2

Our pruning method Faster-VTrans essentially only introduces a single hyper-
parameter, the VIB learning rate, which we set at 5x$10^{-2}$. Increasing the
VIB learning rate (LR) reduces pruning time but may result in lower model
performance. Conversely, decreasing the VIB LR improves model performance but
requires more data samples to converge to the user-defined sparsity level,
thereby increasing pruning time. We prune with only 4000 data samples from the
train set to keep our prune time within 2 GPU hours. During finetuning, we use
LoRA to finetune the all modules except embeddings. The hyper-parameters are
listed in Table 9(b). All experiments with LLaMA are run on NVIDIA A6000 GPU
(48GB).

###  A.7 Resources required for pruning LLMs

The details about the number of data samples and maximum sequence length are
provided in Table 10 below. Although Bonsai is a forward-pass-only method, it
takes approximately 37 hours to prune 50% of the model parameters. This
duration is significantly longer than other forward-pass methods and our
method due to Bonsai’s exploration of sub-models and repeated evaluations.
Bonsai processes one sample with a sequence length of 4096 in a single pass,
whereas our method handles one sample with a maximum sequence length of 512
using a single NVIDIA A6000 (48GB) GPU. In contrast, LLM-Pruner requires two
GPUs of 80GB each or four 40GB GPUs, which is four times the GPU memory
required by our method.

Method | Wikitext2 PPL(w/o) | Max Memory | Pruning | Sequence | Pruning | Inference
---|---|---|---|---|---|---
| finetune | during pruning | Samples | length | Time | Speedup
LLM-pruner | 19.24 | 154 GB | 10 | 64 | 1 hr | 1.5x
Bonsai | 19.47 | 25 GB | 32 | 4096 |  $\leavevmode\nobreak\ 37$ hrs | 1.58x
Ours | 11.8 | 36 GB | 5000 | 512 |  $\leavevmode\nobreak\ 2$ hrs | 1.82x

Table 10: Comparison of details of pruning LLaMA-2-7B model using structured
pruning methods

###  A.8 Pruning/Training plots

We show the trajectory of different losses and layer parameters in our final
objective function during pruning and finetuning for a particular task in
Figure 8. Our sparsity constraint is seen to stabilize pruning after certain
epochs.

###  A.9 FLOPs and Speedup

With FLOPs constraint in our objective function instead of model parameters,
we calculate the sparsity as the reduction in FLOPs of the model to the
original FLOPs. We evaluate our model using the Pytorch Profiler
(https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) and
compare the inference speedup and the FLOPs (excluding the embeddings for
consistency with other methods) of the baseline models with a model obtained
using our method in Table 11. Although certain methods yield models that have
faster inference speed, our model can be compressed more in terms of number of
parameters than the others while retaining performance.

Model | Params (million) | Inference Speedup | # FLOPs (non-embedding)
---|---|---|---
BERT-base | 109 | 1.00 $\times$ | 1.00 $\times$
DistilBERT6 | 66 | 1.98 $\times$ | 0.50 $\times$
TinyBERT6\- GD | 66 | 1.98 $\times$ | 0.50 $\times$
MiniLMv16 | 66 | 1.98 $\times$ | 0.50 $\times$
HomoBERT-base | 65 | 1.30 $\times$ | 0.56 $\times$
BERT-small | 28.6 | 4.77 $\times$ | 0.15 $\times$
$\mathrm{CoFi}_{5\%}$ | 28.4 | 5.16 $\times$ | 0.05 $\times$
TinyBERT3×384-$\mathrm{GD}$ | 17.0 | 7.34 $\times$ | 0.06 $\times$
MiniLMv13 | 17.0 | 7.34 $\times$ | 0.06 $\times$
TinyBERT4×312-GD | 14.5 | 6.28 $\times$ | 0.07 $\times$
HomoBERT-xsmall | 15.6 | 2.51 $\times$ | 0.10 $\times$
HomoBERT-tiny | 14.1 | 2.55 $\times$ | 0.09 $\times$
Ours | 14.5 | 4 $\times$ | 0.07 $\times$

Table 11:  We analyze and compare the inference speedup and FLOPs count
(excluding embedding) for the baselines and our method, relative to BERT-base.

##  Appendix B Further comparison with other methods

###  B.1 Comparison on WikiText-2 after fine-tuning pruned LLMs

The performance of models after finetuning on WikiText-2 is seen in Table 12.
The performance of our model is slightly worse since we use the WikiText-2
dataset with 5000 data samples for fine-tuning. Whereas, Bonsai, Wanda use c4
(Raffel et al., 2020) dataset with over 15000 samples for finetuning, thus
requiring 3 times more fine-tuning time. Before fine-tuning, the model
produced by Wanda 2:4 achieves a speedup over the parent model (1.14×). While
the performance gap can be bridged by LoRA finetuning $(10.52\to 8.34)$, the
adapted semi-structured model experiences a drastic slowdown (0.75×), since
the learned low-rank matrices cannot be merged with the original sparsified
ones without reverting back to dense computation. SparseGPT already uses
weight updates in its algorithm.

Model | Prune | Size | Data | PPL | Inference | Prune Time
---|---|---|---|---|---|---
Method | samples | Speedup | (GPU hrs)
SparseGPT 2:4 | semi-structured | 3B | - | 10.17 | 1.24x | 0.5
Wanda 2:4 w FT | semi-structured | 3B | 15000 | 8.34 | 0.75x | 0.06
Bonsai w FT | structured | 3B | 15000 | 8.89 | 1.58 | 35
Faster-VTrans w FT | structured | 3B | 5000 | 9.47 | 1.82x | 2

Table 12: Performance of models pruned from LLaMA-2-7B evaluated on Wikitext-2
after (w) finetuning (FT).

###  B.2 Stability of the fastest Faster-VTrans variant.

We conduct a comparative assessment between the Faster-VTrans variant and an
analogous fast pruning technique based on Fisher information (Kwon et al.,
2022) across various GLUE tasks, as depicted in Figure 9. The models are
generated with different FLOP constraints, resulting in models with reduced
FLOPs. Notably, for extreme compression (exceeding 50% FLOPs reduction), our
method outperforms the alternative approach, yielding models with superior
performance. Furthermore, our approach demonstrates models with 2 to 10% lower
variance in performance over 10 different random seed initialization,
showcasing enhanced stability, as illustrated in the figure.

Figure 9: Comparison: Faster-VTrans vs. PostPrune (Kwon et al., 2022) on 3
GLUE tasks. Faster-VTrans achieves superior accuracy and stability on
compression from BERT-base

###  B.3 Pruning BERT-large

We also evaluated our primary approach for compressing large models such as
uncased BERT-large. Our models outperform DistilBERT and MINILMv2 on two GLUE
tasks while achieving comparable results on the remaining tasks as seen in
Table 13.

Model | Parameters | SST-2 | MNLI | RTE | MRPC
---|---|---|---|---|---
| (Million) | (Acc) | (Acc) | (Acc) | (Acc)
BERTlarge | 334 | 94.9 | 86.6 | 83.8 | 89.3
DistilBERT6 | 66 | 92.5 | 82.4 | 58.4 | 86.9
MINILMv26 | 66 | 92.4 | 83.4 | 71.3 | 88.6
VTrans | 66∗ | 92.5 | 83.1 | 73.6 | 89.1

Table 13:  Compression of BERT-large model: Comparing performance of
compressed models on GLUE dataset. ∗ indicates the average number of
parameters across all datasets.

###  B.4 Key differences from other iterative pruning methods

We provide a detailed discussion of the key differences of our method from
other pruning methods like HomoDistil (Liang et al., 2023a) and CoFi (Xia et
al., 2022):

Flexibility in pruning and distillation settings: In Table 7 presents a
comprehensive comparison of our method with other baseline techniques
concerning pruning and distillation.
Notably, our training method exhibits flexibility in working within both task-
agnostic and task-specific settings. While HOMODISTIL operates primarily in a
specific setting, our approach accommodates both. Moreover, our task-specific
setting demonstrates remarkable efficiency with significantly lower resource
requirements, utilizing just 1 GPU for tasks such as GLUE and SQuAD, as
opposed to the substantial resource demand of at least 8 GPUs and extensive
training time in the broader setting.

Performance and Resource Efficiency: Table 3 showcases that our task-specific
variant achieves either better or comparable performance in certain tasks
compared to HOMODISTIL, despite utilizing fewer GPUs. This efficiency is
crucial, demonstrating improved performance with reduced resource demands.

Flexibility in Pruning and Layer Adaptability: Our task-specific pruning
methodology, as depicted in Figure 1 identifies varying optimal numbers of
embedding states for different tasks. This adaptability is absent in
HOMODISTIL, which operates under a task-agnostic paradigm. Furthermore, our
embedding masks, FFN layer masks enable flexibility in pruning of layers
enabling higher pruning in certain layers - an aspect not accommodated by
HOMODISTIL, which maintains consistent layer sizes across all layers.

Metric Constraints : Our method has been tested using both parameters and
FLOPs constraints, recognizing that each constraint affects computational
operations differently. We emphasize that a lower number of layers is often
preferred for inference speedup, while hardware limitations, especially in
edge devices like Raspberry Pi, typically align with FLOPs constraints. In
contrast, HOMODISTIL manages local sparsity within each layer, focusing on
controlling total model parameters.

Explainability and VIB-based Approach: Through our VIB-based method, our
approach offers explainability by identifying and pruning redundant units such
as attention heads. This feature distinguishes our methodology from
HOMODISTIL, which does not explicitly provide explainability for its pruning
decisions.

##  Appendix C Further Ablation Experiments on GLUE dataset with BERT

###  C.1 How many heads are enough without losing accuracy?

On pruning just the attention heads in each layer of the transformer, we see
that most GLUE tasks need less than 40% of the attention heads to retain the
performance of the model. As seen in Figure 10, for task SST-2, 30% of the
heads seem to be enough to retain the performance of the large teacher model.

Figure 10:  When only attention heads are pruned, it is seen that for
downstream tasks like SST-2, only 30% of the total heads are sufficient to
retain performance.

###  C.2 Ablating Embedding states pruning

We compare the performance of our pruning method specifically with that of
CoFi to gauge the efficiency of our VIB-based pruning method with flexible
layer sizes in Table 14. For fairness, we ablate the Embedding states pruning
and just prune the other modules of the transformer. The teacher which is also
the initial student being pruned is the kept same for both the methods.

Tasks | CoFi | Our Method
---|---|---
SST-2 | 90.6 | 91.2
MRPC | 82.6 | 82.8
RTE | 64.7 | 65.3
QNLI | 86.1 | 86.9

Table 14: Results showing efficiency of VIB-based pruning coupled with
flexible layer sizes in our method. Both methods prune the same structural
components- all components except Embedding parameters of BERT base. All
models have about 5 million parameters (excluding embeddings)

###  C.3 Analysis of total prune time

The total compression time mentioned in Table 2 includes the pruning and
finetuning time. The breakdown of the same is given in Table 15.

| Prune | Fine-tune
---|---|---
GLUE-high | 23% | 78%
SQuAD | 20% | 80%

Table 15: Breakdown of prune time and finetuning time on pruning BERT-base

##  Appendix D Further Qualitative Analysis

###  D.1 Pruning pattern

We analyze the pruned model structure obtained with our primary compression
method. Figure 11 shows the remaining attention heads, intermediate dimensions
and hidden states in FFN layers at three different sparsity levels(60%, 70%
and 90%) averaged on five datasets within GLUE. From the left-most plot in
Figure 11, we note that the over-parameterized embedding layer can be
compressed such that 20% of the total embedding states are enough to retain
performance of the un-pruned model. We also observe from the other plots in
the figure that there is a significant decrease in the intermediate dimensions
at all sparsity levels. At higher sparsity levels (at 90%), the latter
intermediate and FFN layers are deemed redundant and eliminated. Most of the
pruning happens from the latter layers.

Figure 11: The average pruning pattern across different structural elements of
BERT for GLUE tasks. Left plot: Shows pruned models (dots) retain performance
at more than 80% lower embedding states.

###  D.2 More examples of redundant heads

Figure 12 shows more examples of redundant heads with a different input
sentence than shown in the main paper. Repetitive heads with attention to CLS
tokens, attention to own tokens and broad attention in deemed redundant and
can be pruned out from the teacher model.

###  D.3 Shift of attention to next or previous token

Figure 12 shows how average attention of input tokens in the compressed models
change from current token to next or previous token. Attending to current
token does not give the model any new information while focusing more on
previous or next token might yield better performance.

Figure 12: Another example of heads that are deemed redundant and pruned.
Repetitive Heads with tokens attending to the CLS token, tokens attending to
themselves and heads with broad attention are deemed redundant. Plots show the
average attention paid to next, current and previous token in (Top:) Un-pruned
model and (Bottom:) pruned model

###  D.4 Divergent Behavior Among Attention Heads.

We calculate the Jensen-Shannon Divergence $\sum_{\text{token }\in\text{ data
}}JS\left(\mathrm{H}_{i}(\text{token}),\mathrm{H}_{j}(\text{token})\right)$ to
measure the pairwise distances between attention head distributions over input
data. Using multidimensional scaling Kruskal (1964), we create a two-
dimensional representation, ensuring the resulting embeddings closely
approximate the Jensen-Shannon distances between the corresponding heads. We
then visualize the behavior among heads in Figure 13. In the un-pruned BERT
model clustering in layers show similar head behavior. Layers 9 to 12 and
layers 1 to 4 exhibit close grouping among themselves suggesting analogous
behavior in adjacent layers. Conversely, in the pruned model, heads from
layers 1 to 4 show a more dispersed distribution, indicating a broader range
of behaviors. In the later layers (8 to 11), heads appear sparse and widely
dispersed with layer 12 being entirely pruned. This validates our approach,
demonstrating its effectiveness in removing redundant heads within a layer and
entire layers with similar behavior. Moreover, our method aligns with the
interpretation by Burgess et al. (2018) suggesting that VIB encourages the
acquisition of more disentangled representations.

Figure 13: (a) Comparison between full BERT-base and 50% pruned model
attention heads in 2-D space. Distance between points reflect average Jensen-
Shannon divergences between head representations. (b) Heads in pruned model
exhibit sparser, divergent behavior compared to un-pruned model. Heads
retained in the 50% pruned model showing attention among input tokens from
SST-2 dataset.