PASS: Pruning Attention Heads with Almost-sure Sparsity Targets
Authors: Dujian Ding
Ganesh Jawahar
Laks V. S. Lakshmanan
Venue: Accepted by TMLR ()

Abstract: Transformer models have been widely used to obtain high accuracy values in multiple fields including natural language processing (NLP), computer vision, and more. This superior performance typically comes at the expense of substantial computational overhead. Multi-head attention is the key factor in the success of Transformer models that has been found to be computationally expensive. Significant research effort has been devoted to improving attention compute efficiency by pruning redundant attention heads. A widely adopted paradigm is to jointly learn a set of gate variables and apply thresholds on gate values to prune heads. Previous work shows a high level of sensitivity to threshold tuning which can limit subnetwork performance and prevent them from wider adoption in practice. We propose the notion of almost-sure sparsity to overcome this limitation and develop a generic framework for Pruning with Almost-Sure Sparsity (PASS) targets over attention heads. To further boost efficiency, we design a novel technique, concentrator, based on which we develop PASSCONC (PASS with CONCentrator). We also present a simple-yet-effective strategy to further improve subnetwork performance by clipping and selectively reopening learned gates. We investigate PASS and PASSCONC on two widely studied architectures: encoder-decoder (ED) Transformer and encoder-only Transformer (e.g., BERT-base). Experiments on IWSLT14 German-to-English translation and GLUE benchmark tasks demonstrate that our approaches outperform the SOTA by achieving up to 1.33 higher BLEU scores, 1.44% higher accuracy, and 60% higher attention speedups.

Published in Transactions on Machine Learning Research (09/2024)
PASS: Pruning Attention Heads with Almost-sure Sparsity Targets
Dujian Ding
dujion dittinn.ding@gmma... com
Department of Computer Science the University of British Columbia
Ganesh Jawahar
ganeshjuhr @gmaail com
Google DeepMind
Laks V.S. Lakshmanan
laks@ccs ubc.
Department of Computer Science the University of British Columbia
Reviewed on OpenReview: https: // openreview. net/ forum? id= SQduStTKGL
Abstract
Transformer models have been widely used to obtain high accuracy values in multiple fields including natural language processing (NLP), computer vision, and more. This superior performance lypically comes at the expense substantial computational overhead. Multi- head attention is the key factor in the success of Transformer models that has been found to be computationally expensive. Significant rescarch effort has been devoted to improving attention compute efficiency by pruning redundant attention heads. A widely adopted paradigm is to jointly learn set of gate variables and apply thresholds on gate values to prune heads. Previous work shows high level of sensitivity to threshold tuning which can limit subnetwork performance and prevent them from wider adoption in practice. We propose the notion of almosl-sure sparsity overcome this limitation and develop a generic framework for Pruning with A lmmost-Sure Sparsity (PASS) targets over attention heads. To further boost efficiency, design novel technique, concentrator, based on which develop PASSCONC (PASS with CONCentrator). We also present simple-yet-effective strategy to further improve subnetwork performance by clipping and selectively reopening learned gates. We investigate PASS and PASSCONC on two widcly studied architectures: encoder- decoder (ED) Transformer and encoder-only Transformer (e.g. BERT-baase.. Experiments on IWSLT14 German-to-Enun English translation and GLUE benchmark tasks demonstrate that our approachess outperform the SOTA by achieving up to 1.33 higher BLEU scores, 1.44% higher accuracy, and 60% higher allention speedups.
1 Introduction
Transformer models (Vaswani al. 2017) have become lead force the study of natural language pro-
cessing (NLP), computer vision, information retrieval, and other domains (Hua et al. 2024 2023 Asai al., Darcet et al. 2024 Ding et al.) 2024). As Transformers grow deeper and larger, however, their ap- plication on longer contexts remains challenging because attention computation, which is at the heart of Transformer architectures, is of quadratic time and memory complexity with respect to the input length Dao et al. 2022 For example, Wang et al. (2220a) observed that attention computation typically accounts for over 50% end-to-end latency of GPT-2 model on multiple hardware platforms.
2024
Significant research efforts have been devoted improving attention computation efficiency from two or- thogonal perspectives: reducing allention complexity and pruming attention heads. As successful attempt

---

Published in Transactions on Machine Learning Research (09/2024)
Gate E [0,1]
90
Prev. Prob. Pruning SOTA
# Pruned Heads
60
Our Method Unpruned
Accuracy (%)
88-
attn head1
50
a7-
40
86
attn heady
Prev. Prob. Pruning SOTA
Our Method
85
Sparsity Target
84-
0.2
0.3
0.4
0.5
066
0.7
0.8
100
120
140
160
180
Gate zn E [0,1]
Gate Thresholds
Speedup (%)
(a) Pruning attention heads w/ gates.
(b) Threshold sensitivity.
(c) Pruming results.
Figure 1: (a) Pruning attention heads by learning a set of gate variables. Gates take values from [0, 1] and apply to attention heads before summation. Heads of low gate values are more likely to be pruned. (b) With a Transformer model of 72 heads and a sparsity target aiming to prune 56 heads, previous probabilistic pruning SOTA (Xia et al. 2022 is highly sensitive to gate threshold tuning (a 0.5 threshold only prunes 49 heads) while our approach consistently achieves the desired sparsity target (c.g.. pruning 56 heads) in a chreshold-independeen manner. (c) On GLUE benchmark tasks (Wang et al., 2018), our approach achieves up to 1.4% higher average accuracy and 60% higher speedups than previous probabilistic pruming SOTA. Points in the plot represent the accuracy-speoduup tradc-offs achieved by different approaches under various sparsily targets. Unpruned model performance is also included for reference purpose.
at reducing attention complexity, sparse attention (Roy et al. 2021 Tay et al. 2020 Child et al., 2019) focuses on sparsifying the attention distribution over tokens for each head to improve efficiency. Linformer (Wang et al. 2020b) reduces the attention compute complexity from O(N²) to O(N) with low-rank matrix approximation. FlashAttention (Dao et al. 2022 Dao 2023) focuses on memory access efficiency in attention computation and achieves linear memory complexity by exploiting the asymmetric GPU memory hierarchy and minimizing unnecessary data transfers. However, implementing these approaches introduces extra chal- lenges, particularly in rewriting the attention arithmetic operations as well as the underlying CUDA kernels to improve hardware utilization, which prevents them from wider adoption in practice.
The second line of work focuses on praning attention heads (Voita et al., 2019 Li et al. 2021 2023 to achieve significant inference specdups without changing the arithmetic operations of attention modules and therefore can be applied lo majority of Transformer models with only minor training/ /fine- fine- tuning configuration changes. These works utilize the fact that properly trained Transformers are highly over-paramcterized, and study how to extract efficient subnetworks by removing redundant heads without significant, performance drops. widely adopted paradigm is to jointly learn a set of trainable gale variables for each attention head, as shown Figure Ta] At test. time, attention heads associated with low gate values are pruned subject to predefined thresholds or sparsity targets. Typically, Li et al. (2021) achieves user- specified attention head sparsity by iteratively applying the Gumbel-softmax trick (Gumbee
Xia et al.,
1954 to select
the top-K K most important attention heads for given sparsity targels. At each training iteralion. however, only selected attention heads get updated by the training optimizer (e.g., Adam (Kingma & Ba 2015)) which prevents models from being trained with more heads in early training stages and limits the final subnetwork performance. VVita et al. (2019) and Xia et al. (2022) overcome this limitation by allowing all heads to participate the training process and achieve sparsity a probabilistic manner. Specifically, (Voia et al., (2019) learns the probability distribution for gate values and sparsifies the models by regularizing the gate closing probability (the likelihood that gate variables equal 0). Xia et al. (2022) follows similar probabilistic forumulation and achieves the target sparsity in expectation by using the lagrangian multiplier (Wang et al. 2020ch to explicitly enforce sparsity constraints on attention heads. At test time, both Voita et al. (2019) and Xia et. al. (2022) use a threshold (e.g., 0.5) on the gale closing probabilities to determine if an attention head can be pruned confidently. In practice, however, these probabilistic pruning approaches suffer from two limitations. Firstly, setting correct thresholds challenging and mistakenly chosen threshold can lead failed sparsity targets, as illustrated in Figure [b] Secondly, thresholding gate variables at test time may bring significant architecture changes to Transformer models and lead to decayed subnetwork performance.
2

---

Published in Transactions on Machine Learning Research (09/2024)
observe that previous probabilistic pruning SOTA (Xia et al. 2022 may lead to drastic performance decay even with correctly calibrated gate thresholds, while our approach achieves up to 1.4% higher average accuracy and 60% higher speedups (see Figure Icp.
On GLUE benchmark tasks (Wang et al. 2018),
Motivated by these limitations, in this work we propose a novel pruning approach, Pruning with Almost- Sure Sparsity targets (PASS), to achieve good subnetwork performance by involving all heads during model training and prune models to the desired sparsity levels in a hureshold-independeed et al. (2019) and Xia et al. (022)) PASS takes the probabilistic pruning formulation but explicitly enforces all gate variables to diverge as the pruning process proceeds. The most important observation is that, if the closing probabilitios for all gate variables approach either 0 or 1 with respect to the sparsity targets after training, we can easily extract the desired sparse subnetworks by introducing no further changes in attention compute at test time and therefore retain the subnetwork performance. We use the weal-sstablished notion "almost surely" in probability theory (Jacod Protter 2003) and propose the notion of almost-sure sparsity (see Section 3.1. We say gate can be closed almost surely if the corresponding gate closing probability equals 1. In PASS, we express the sparsily largets in terms of the almost sure sparsity and extract the desired subnetworks by enforcing the sparsity constraints through regularization techniques in model training process (sce Scction 3.2) To push the envelope on inference efficiency, we propose novel techmique, concentrator, based on which we develop PASSCONC (PASS with CONCentrator), as discussed in Section 3.3 Observing the gradient vanishing problem in the gate variable training process, we present a simfle-vet-fffectivee strategy to clip and selectively reopen learned gates which leads to improved subnetwork performance (see Section 3.4.. We evaluate our methods with encoder-decoder (ED) Transformer models and BERT models on IWSLT German-to-Fanglshh translation (Cettolo et al., 2014) and GLUE benchmark tasks (Wang et al., 2018). We explore the Pareto front between model performance and inference efficiency
manner. Similar to VVoita
for surnetworks identified by PASS, PASSCONC, and recent work (Li et al.. 2021 Xia et al. 2022: Voita et al.] 2019). Experiments show that PASS and PASSCONC outperform all baselines across a majority of experiment settings, by identifying subnetworks of higher speedups and better model performance (see Section 4. For example, on GLUE benchmark tasks, PASSCONC achieves a 185.2% attention speedup on average, which is 60% higher than all baselines, while providing even higher accuracy. This observation suggests that PASS and PASSCONC are capable of identifying subnetworks with high model capability and can be applicd to resourco-limited applications to achieve good performanco-cffciineyy trade-offs.
In this work, we make the following contributions.
We propose a novel notion of almost-sure sparsity and develop an effective model pruning framework PASS to prune models to specified almost-sure sparsity levels on attention heads.
2. We propose a novel technique, concentrator, to further push the envelope on model inference effi- ciency and develop PASSCONC.
We present a simple-yet-cffective strategy further improve subnetwork performance by clipping and selectively reopening learned gates.
4. We evaluate PASS and PASSCONC on ED Transformer and BERT models with well established NLP tasks. Experiments show that PASS and PASSCONC outperform baselines by obtaining significant efficiency improvements and better performance--ffecenccy trade-offs.
2 Preliminaries
A frequently encounterod task in machine learning is to find the model that minimizes the negative log- likelihood of an observed dataset, which can be formulated as follows,
0r = arg min log P(D(0)
(1)
where D is an observed dataset and 0 = (0,, ... p } stands for the parameters of a parameterized model
(e.g., a neural network). In real-world applications, we typically have model sparsity constraints to prevent
high inference latency reduce memory footprints (Gupta & Agrawal 2022). A recent line of work (Louizos
3

---

Published in Transactions on Machine Learning Research (09/2024)
et al. 2017 Voita et al. 2019) pursues this goal by training gate variables, == == [2,,"2"" 2 jointly with parameters, 0. Each E has support [0, 1]. The objective function Eq. 1 can be re-parameterized as,
arg mim log P(D @ x)
where O indicates component-win multiplication between network parameters and the gate variables z. Typically, is a latent variable following the posterior distribution D)z which reflects the user-defined sparsity constraints. The probabilistic pruning approaches (Voita et al., 2019 Xia et al. 2022 aim to optimize the expected likelihood over the posterior distribution of the gate variables 2,
0 = arg min - log Ep(n Oz)
(3)
is mathematically intractable when the posterior p(z[D) is priori unknown. As an attempt to tackle such intractability, we can first derive the evidence lower bound of the log-likelihood in Eq. 3 which is widely used technique in previous variational inference work (Vahdat eet al.) 2018a b). Since we are interested in minimizing the negative log-likelihood it gives us an upper bound for the objective in Eq. 3
The objective function described by Eq. 3
log pl Dj (P(D)0 0z)] S - [log P(D @ x)] + (L( () ()
(1)
where q(z; L) is an approzimate posterior distribution parameterized by Φb = f0r f01.02.... 02. 010 Detailed derivation can be found in Appendix A.1 Minimizing this upper bound with respect to ((z; b) results in q(z; 2) D(z]D) and turns the inequality into an equality (Beal) 2003). By denoting this upper bound C(0, 2) we can then formulate the learning problem as,
C(0, 4) E log P(D @ 2) + KLi((z; ((ll())))lll(()))
0-,4" " = arg min ((0 ()
0,0
We aim jointly learn the optimal network parameters 0r and the distribution of gate variables, 4", by minimizing the upper bound C(0, 4)
The foregoing analysis gives a generic framework to enforce sparsity over neural models which is agnostic to the underlying network structures. To prune attention heads, all we need is to assign cach head a gate variable and solve Eq. 5 with z = {z1, #2, ] where H is set of all attention heads (see Figure la
3 Methodology
3.1 Almost-sure Sparsity
The KL-divergence term in Eq. 5 is mathematically intractable when the true posterior p(z|D) is unknown.
A line of work (Voita et al., 2019 Xia et al.) 2022) attempts to tackle this intractability by replacing the surrogates. A widely used surrogate (Voita et al., 2019) is > Pr[z # 0), which can be seen as a special case of the KT-divergence term that assumes a constant (2())) ratio log D)= A. Though this surrogate circumvents the intractability issue, it is often challenging identify the right& for a given sparsity target 8 (Li et al. 2021) Other work (Xia et al. 2022) utilizes surrogates in the form of Lagrangian Multipliers (Wang et al. 2020c) to enforce sparsity in expectation for a given target. Though this approach able to achieve target sparsities in probabilistic manner, its performance highly relies on the gate thresholds and may lead to limited subnetwork performance, illustrated in Figures Ib and Ic
KL-divergence term with iestribttioninddpenddd
In light of the limitations of previous work, we introduce the notion of almost-sure sparsity and propose a novel surrogate which allows us to learn empirically good approximate posteriors as well as discover subnetworks with desired target sparsities almost surely. The intuition behind the almosl-sure sparsity is straightforward. Note that a model has sparsity s provided a fraction s of the gates are closed in the network. From a probabilistic perspoctive, it is natural ask a subnetwork to be "confident" about which gates should
The posterior distribution p also depends the models but ignore it here since does change the inequality.
4

---

Published in Transactions on Machine Learning Research (09/2024)
be closed. In other words, gates should be closed with high probability. Mathematically, event is said to happen almost surely, if it happens with probability 1 (Jacod & Protter 2003 Formally, we define almost-sure sparsity as follows.
Definition 1 Caaiiiion Allaaaaaiai (Almost-sure Sparsity) Given 8 E (0, 1), gate variables z have almost-sure sparsity s if 3zloose: Zopen - 2, such that Pr[z; = 0] = 1, 4/2 E Zelose and Pr[z, = 1] = 1, V2z E Zopen: where Zelose N Zopen Zetase Zopen = 2, and |ictose = s/2/-
We argue that the almost-sure sparsity is better aligned with the sparsity notion we need in static subnetworks and enables the subnetwork discovery with desired sparsity targets. Next, we present learning objectives designed to achieve almost-sure sparsity targets specified by users.
3.2 Learning Objective with Almost-sure Sparsity
We aim to learn good approximate posterior g(2; b) with desired almost-sure sparsity. In this paper, we adopt the Hard Concrete distribution (Louizos et al. 2018) the basic form of the approximate posterior g(z; ②), given its contimuous-discrete nature and its wide application in model pruning (Voita et al. 2019 Xia et al. 2022
Hard Concrete distribution has its support over the closod interval [0, 1] and non-zero probability mass at and 1. Hard Concrete dis- tribution is derived by stretching and collapsing the Concrete distri- bution Madddisn et al.,
Roaso ( (0r, $2). 0.5)
Rpass (0
{01, b, 0.5)
1.2
0
2016), as
illustrated in Figure 3a We in- troduce derivation details in Ap- pendix A.2 For each gate 2 E [0,1] following Hard Concrete distribu- tion, the corresponding probability mass at 0 and with respect to g(2;;)) are given as q(4 = =0;
1.2
-10 -10
0.8
-10
0.8
0
10
-10
10
Φ1
Φ
Figure 2: Values of Roase and Rpass with Φ (01,02) and 0.5.
sigmoid (3 log (7) - -.). q(2, = 1:0,) = sigmoid (0. - B log (-)) For simplicity of notation, we de- note (d.))) = ((= 0: (b), the gate closing probability, and qi (0:) == a(2; 1; the gate opening probability. Due the monotonicity of the sigmoid function, when di increases, qı (óx) increases and (o( decreases, and gate is more likely to open. We further define (6, 1- (0,)) (d:) as the probability for 24 being non-binary. We use 8 = 0.33, 7 -0.1, and Ç = 1.1 by default, following previous work (Voita ct al. 2019). Clcarly, the closing and opening probability of each 2; E z are differentiable functions of Φ; E &. as shown in Figure 3b) By jointly learning & with the network parameters, we are able to almost-surely close (resp. open) gates 2 E z by contimuuusy increasing (resp. decreasing) the values of Φi E & using gradient-descent optimizers (c.g., Adam (Kingma Ba 2015)). At each training iteration, gates are sampled w.r.t. the learnt distribution and then applied attention heads to achieve pruning.
At the end of pruning, we want q(z; b) to achieve almost-sure sparsity for given target s. Our strategy is to design a learning objective that meets the desired almost-sure sparsity at its optimum, and optimize along with model training. It is worth pointing out that there exists a family of learning objectives satisfying this criterion. However, not all of them can be easily optimized their minimum, especially by gradient descent optimizers (Kingma & Ba 2015 For example, one may propose to minimize the following objective.
[6]
10
Roase ( 8) = Σ (mb( +
s|0|
Σ
(x))
(6)
It can be casily seen that Rease takes on its minimum value when achicving almost-sure sparsity s. However, there exist local optima that may prevent gradient descent optimizers from converging to the global optimum. To illustrate this, for simplicity, we visualize the values of Raase in 2-gates setting {21, 22) in Figure 2
5

---

Published in Transactions on Machine Learning Research (09/2024)
Gate Close&Open Prob.
Gate Prob. Gradients
equi0ilee aq1 (0:l/0p
PDF Prizi ( 0)
Stretched Concrete Concrete
9:(h)
0.10
0.2
0.05
0.0
0.00
0.2
0.4
0.6
0.8
-10
-5
5
10
-10
Zi
(a) Hard concrete distribution.
(b) Closing and opening probabilities.
(c) Gate probability gradients.
Figure (a) Hard concrete distribution derived by streching-ande and-collapsing a Concrete distribution. (b) Closing and opening probability of gating variables are differentiable functions of Φ, (c) The gradient values of both gale closing and opening probabilities quickly approach 0 as o; increases or decreases. Hard concrete distribution is parameterized with B = 0.33, -01,, Ç 1.1, following previous work (Voita et al. 2019
With 2 gates and a sparsity target 0.5, we want one gate to be almost-surely closed and the other gate almost-surely opened. In Figure 2 such global optima correspond to the top-left and bottom-right corner where one of Φi or d2 takes on high value and the other takes on a low value. However, can be clearly obscrved that there exist a local optimum both gates are open with high probability. In other words, with Rease, if both Φ1 and Φ2 happen to take positive values due to noise from the training process or bad initialization, the gradient descent direction will increase the probability for both gates to be open and fail to meet the sparsity target 0.5 by delivering overly dense models. Under weak condition we can prove that the gradient descent direction of Roase
the top-right region which corresponds to the situation where
always leads to higher opening probability for gate 2, øf 2 log( -I-v---a-aall-------- where g(a) 2c = logll( ) The proof is presented in Appendix A.3
[(a)
In light of the limitation of Rbase, we propose the following learning objective,
8 Rpasa (2, 8) = Σ (ms(ds) + s/0 Σ (o(+)) + (1 - s)10 Σ (or)
10
10
Rposs does not suffer from the local optimality issue that Rease does, as shown in Figure 2 In fact, show that minimizing Rpass always generates neither over-sparse nor over-dense subnetworks. In order to show this formally, we define the expected number of opened gates ( E!e q (1k)) as the expected densily, and the expected number of closed gates 1 (o(())) as the espected sparsily. We have the following lemma
Lemma 1 Minimizing Rease always generates sparse subnetworks whose expected sporsities no than 8 and expected densities are no more than 1 for any given sparsily target E (0,1).
Proof can be found in Appendix A.4. By substituting the KL--divereence term in Eq. with Rpass, we obtain the PASS optimization objective where À is the regularization coefficient.
Lpass (0, 3) - P(D @ x)) + XR (
0 0pass, Dpass arg
arg.mmn min Lpass (9, b))
(8)
6,6
3.3 Concentrator
To further improve model inference efficiency, we propose the use of concentrator. Wang et al. (2020a) observed that the auxiliary operations in multi-head attention computation (e.g., reshaping and transposing
We assume the Hard Concrete distribution is equally stretched in both directions, which gives 7 =

---

Published in Transactions on Machine Learning Research (09/2024)
matrices, heads splitting, and concatenation) account for 73% of the overall latency in attention layers. The run-time overhead can hardly be avoided as long as there exist unprumed heads in the attention layers. Consider subnetworks of the same attention head sparsity. Intuitively, if the unpruned attention heads are inclined to concentrate among a few layers, the other layers can be entirely skipped, saving the run-time overhead and improving inference efficiency. Given this, we propose the concenlralor to encourage the unpruned attention heads to be concentrated on as few layers as possible.
Given Transformer-based model of L layers and H heads per layer, the concentrator is defined as Reone ( El (1 11-IT' (1-IT 2o (o(()))): where ΦI, is the distribution paramcter for the h-th gate variable on the I-th
layer. Notice that 1 [ go((r,,) indicates if the I-th layer can be entirely skipped: takes on value 0 only if all heads of the layer have closing probability 1. Reone is summation of the layer-wise indicators over all layers and has regularization effects by penalizing the levels of unconcentration. We introduce de control the concentrator effects and obtain the following optimization objective for PASSCONC, i.e., PASS with concentrator.
passsonn (0, () = -By-aaa P(D O 2)] + (Rpaes00.)) 5) + de Reonc
0 passcoone $ passcone arg min Lpasscooe (0, ()
(9)
8,0
3.4 Clipping and Reopening
In practice, with proper training settings, the proposed approach can discover subnetworks with the desired sparsities and high accuracy. Note that we approach almost sure sparsity by increasing or decreasing d; E & with gradient-descent optimizers. However, as é increase or decrease, their gradients quickly converge to 0 as illustrated in Figure 3c| Consequently, gates closed (resp. opened) with high probability in early training stage are unlikely to be self-adaptively re-opened (resp. closed) in later training iterations by gradient- descent optimizers, which may lead to sub-optimal pruning results. We propose to resolve this issue with a clipping and selective reopening strategy. The idea of clipping has been widely used in training deep learning models to avoid gradient exploding and vanishing (Zhang et al. 2019 Koloskova et al.) 2023). In this same spirit, we clip Φ; to predefined ranges to alleviate the aforementioned issucs caused by small gradients. In our implementation,, we empirically clip all éx's 's to the range [-5, 5] to avoid the vanishing of gradients to excessively small values (see Figure Bcb. Randomness has been widely observed be helpful in neural network training (Bottoul 2010). To further incentivize training dynamics,
propose to randomly reopen closed gates with respect to the gate quality. There is line of work on how measure gate qualities (Michel eet al. 2019 Voita et al. 2019 Ding et al. (2017), among which we choose head confidence (Voita et al. 2019) in our implementation because it has been found to be an informative notion (Behnke & Heafield 20200 and requires little to no additional computation. The confidence of a head is the average maximum attention weights of tokens in a set of sentences (Voita et al. 2019). We normalize confidence scores for each attention head and reopen almost-surely closed gates with a probability equal to the normalizod scores.
Notably, the clipping and reopening strategy is designed to be applied jointly. Without clipping to a proper range, opened gate may remain fully open due to the nearly zero gradients and finally result in overly- dense subnetworks as the reopening strategy progresses. Similarly, without the randomly reopening strategy, the subnetwork may be trapped by certain local optimums and lead to suboptimal performance even with the clipping strategy.
4 Evaluation
4.1 Evaluation Setup
4.1.1 Model, Data, and Metrics
Model: We investigate two widely-studied Transformer models: encoder-decoder (ED) Transformer (Vaswani et al. (017)) and the encoder-only Transformer, BERT (Devlin et al.) 2018). We use the FAIRSEQ toolkit OOt et al. 2019 to implement a 6-layer ED Transformer with 72 heads in total, and the HUGGING
3To reopen an almost-surely closed gate 2i, we manually decrease closing probability by increasing di.
7

---

Published in Transactions on Machine Learning Research (09/2024)
FACE codebase (Wolf et al., 2020) to implement 12-layer BERT-base with 144 heads in total. We believe our setting has covered the most important attention head variants, especially that the ED Transformer includes both encoder (self-attention) and decoder (self- and cross-attention) modules, and therefore suffices a comprehensive evaluation setup.
Datasets: Following previous work [Li et al. 2021p, the ED Transformer model is trained and evaluated on the IWSLT11 German-to-Englisi translation dataset (Cettolo et al. 2014). The BERT-base model is fine- tuned and evaluated on 4 benchmark NLP tasks from the GLUE benchmark (Wang al.) 2018) inchuding the Multi-Genn Natural Language Inference (MNLI) dataset (Williams et al. 2018), the bquestin-ansseriin NLI (QNLI) dataset (Rajpurkar et al. 2016), the Quora Question Pairs (QQP) datasct (Sharma et al.) 2019). and the Stanford Sentiment Treebank (SST-2) dataset (Socher et al. 2013). We choose these 4 datasets of the largest training splits from the GLUE benchmark to ensure that gate variables can be sufficiently trained to convergence.
Metrics: We use BLEU score (Papineni al. 2002) as the metric to measure model performance on the translation task following previous work (Li et al. 2021 Michel et al. 2019 and use accuracy as the metric the GLUE benchmark tasks following Wang et al. (2018). In addition, are also interested in the efficiency improvements achieved by PASS and PASSCONC. We use wall clock time to measure the efficiency w.r.t. latency.
4.1.2 Baselines
We consider three strong baselines that prune attention heads to a specified sparsity level.
Differentiable Subset Pruning (DSP). DSP {Li et al. 2021 applies the Gumbel-softmax trick (Gumbel,
1954) to select the top-K attention heads for a given sparsity target. DSP learns a K-hot vector Oh by
iteratively applying Gumbel-softmax K times, where gh = EK of Ex= En= /xplr Th = h k-1
cmp(r\/7)
(0 )) and r = wh + nu- Wh denotes trainable parameter indicating head importance, n Gumbel(0, 1) is Gumbel noise, and 7 is a hyper-parameter that controls the annealing temperature.
Lagrangian Multiplier (LAG). A recent line of work (Xia et al.) 2022 Wang et al., 2020c) employs Lagrangian Multiplier (Wang al. 2020c) to enforce sparsily in expectation. Given sparsity target 8; LAG trains models along with the regularization term Riag- (1 - s) + 22 s)², where is the expected sparsity. A and X2 are trainable parameters and will be optimized jointly in training.
Voita et al. (2019) (Voita). Voita et al. (2019) prunes attention heads by applying the stochastic approx-
imation to Lo regularization (Louizos et al. 2018 to gate closing probabilities. Voita et al. (2019) achieves
pruning by jointly training models with the following regularization term Rvoita (b)
Ao Σ h1 qo (0(0)))
where Ao can be used to indirectly control the achieved sparsities. Following the previous work (Li et al. 2021), we use grid search to find the correct X, values for each sparsity setting.
As observed in Li et al. 2021), be is in general hard to tune and there could exist certain sparsity targets that cannot be achieved by tuning Aø, which makes Voita's approach not a good fit if the user demands a sparse Transformer model of customized number of attention heads. In our work, we propose enforce the spansity constraints explicitly using a novel sparsifier (see Section 3.2 and further push the envelope of efficiency by developing the concentrator (see Section 3.3). The design of our sparsifier, concentrator, as well as the clipping and reopening strategy has not been proposed in previous attention pruning literature and contributes to the novelty of our work.
We adopt the same implementation of DSP and Voita as in Li et al. 2021 We implement LAG by training models with the regularization term Ring = (s (2 ()-s)) 5)², where the user-specified sparsity target, \ and A are trainable parameters, and å the expected sparsity over all attention heads. Following Louizos ct al. (2017),
cstimate the value of cach gate as 2; max max sigmoid (e:)(5 77 7)}}, for
di E 4. The expected sparsity 3 over attention heads H. is computed as = [M| E 1-2.
*Observing the emergence of large language models (LLM) Chang et al. 2024 that are typically decoder-only Transformer models, it is intriguing to evaluate the scalability of our approaches on large-scale LLMs and we leave it in our future work. httpps://gthuu..oom// httpe//www.ooooooooobaoooc/eaaee
8

---

Published in Transactions on Machine Learning Research (09/2024)
Table 1: Subnetwork performance on IWSLT14 En translation (left) and GLUE benchmark tasks (right).
BLEU
Accuracy (%)
K
PASS
PASSCONC
DSP
LAG
Voita
PASS
PASSCONC
DSP
LAG
Voita
16
32.73
32.70
31.40
30.91
27.55
86.27
85.25
84.47
83.84
84.83
32
33.45
33.48
33.42
32.66
32.80
87.59
86.47
88.36
86.99
87.15
48
33.89
33.91
34.00
33.12
32.97
88.65
88.30
88.52
88.02
88.02
64
34.01
34.05
33.89
33.02
33.20
88.72
88.62
88.81
88.40
84.20
4.1.3
Protocols
We express sparsity targets over attention heads H. interchangeably 8 E (0, and as integer K where K = [)H the number of unprumed heads. Unless stated otherwise, for a given sparsity target K, we evaluate all medh by selecting the top- K most important heads w.r.t. the corresponding ranking metries, i.., the gate opening probabilities for PASS, PASSCONC, Voita, and LAG, and the head importance score wh for DSP. Detailed hyper-parameter settings are Appendix B We test all methods on both architectures with target tasks (30 training epochs for ED Transformer; fine-tuning epochs for BERT-base as in Li et al. (2021) )- All experiments are conducted on high performance compute cluster equipped with NVIDIA P100 GPUs (each with 12GB GPU RAM). Codebase is available on https: /gtpggtttttoo////////m Dww.oom/iiiiiiiiiinnini. PASS.
We evaluate the performance of PASS and PASSCONC in Section 4.2 demonstrate the speedups brought by PASSCONÇ in Section 4.3 validate the eveectiveness of rancentrator as well as the clipping and reopening
strategy in Section 4.4 present the patterns of pruned attention heads in Section 4.5) and show more performance results in Appendix
4.2 PASS and PASSCONC Improve Model Performance
We investigate the model performance of subnetworks identified by PASS, PASSCONC and all baselines under various sparsity constraints. We compare all five methods on both ED Transformer and BERT-base models. The results are summarized in Table [ Detailed performance results on each GLUE benchmark lask are Appendix C
On IWSLT14 German-too-nnllie translation task, PASS and PASSCONC outperform all 3 baselines in majority of sparsity settings. When K = 16, both PASS and PASSCONC achieve BLEU scores of 32.7, which is 1.3 higher than DSP, 1.8 higher than LAG, and 5.2 higher than Voita. On the GLUE benchmark lasks, we observe a similar trend in high sparsity situations. When K = 16, PASS and PASSCONC achieve average model accuracy of 86.27% and 85.25% respectivcly, while DSP drops to 84.47%, LAG drops to 83.84%, and Voita drops to 84.83%. When sparsity largets are low, PASS able to match or outperform all 3 baselines, while PASSCONC can be outperformed by the strongest baseline DSP while still being comparable to the remaining two.
One interesting observation is that Voita delivers surprisingly low accuracy in low sparsity settings (e.g.. K = 64) with GLUE benchmark tasks. The degraded performance can be attributed to its intrinsic sensitivity to the choice of Åv, which is used to indirectly control sparsity targets. Li al. (2021) observed that a small increase in A (e.g., 0.0009 > 0.0014) may lead to drastic change of achieved sparsity (e.g., the number of unpruned heads decreases from 30 to 11), which suggests that Voita. is inadequate when users require subnctworks of pro-defined number of attention hcads.
4.3 PASSCONC Improves Model Efficiency
We evaluate the attention speedups for subnetworks identified under various sparsity constraints, at inference time. We report the inference speedups in comparison to the unprumed model. The results are summarized in Figure and Table Detailed results on each GLUE benchmark task can found in Appendix C
On the GLUE benchmark tasks with BERT-base models, PASSCONC outperforms all basclines across a majority of sparsity constraints with great efficiency improvements and comparable or better accuracy (see

---

Published in Transactions on Machine Learning Research (09/2024)
Table Attention speedups on 1WSLT14 De-En translation (left) and GLUE benchmark tasks (right).
Speedup (%)
Speedup (%)
K
PASS
PASSCONC
DSP
LAG
Voita
PASS
PASSCONC
DSP
LAG
Voita
16
144.3
162.8
141.1
141.1
142.7
114.4
185.2
123.1
120.4
126.1
32
115.5
118.7
118.7
110.4
117.6
107.1
135.6
107.3
105.4
105.3
48
101.6
104.1
105.8
102.4
102.4
103.1
109.3
102.7
102.9
103.8
64
100.8
104.1
100.0
100.0
100.0
103.2
106.4
102.9
103.0
103.0
K=64
K=48
90
K=32
LAG
K=16
K=64
Voita
Accuracy (%)
K=48
DSP
PASS
BLEU
LAG
87
K=32
PASSCONC
Voita
Unpruned
DSP
86
PASS
PASSCONC
85
K=16
28
Unpruned
100
1i0
120
130
140
150
160
100
120
140
160
180
Speedup (%)
Speedup (%)
(a) IWSLT14 De-En translation.
(b) GLUE benchmark tasks.
Figure 4: Subnetwork performance v.s. attention speedups on IWSLT14 De-En translation task and GLUE benchmark tasks. We add dashed lines (--) to connect data points with the same sparsity constraints for better comparison. The unpruned model performance is included for reference purpose.
Figure 4 When = 16, PASSCONG achieves 185.2% speedup, which 60% higher than all baselines, and am average accuracy 85.25% that is also higher than DSP, LAG, and Voita. PASS has better accuracy but lower speodup. As the sparsity targets decrease (i.c, as K increases), the speedups achieved by all methods in general goes down but PASSCONC always dominates the competition in terms of elficiency, at the price of a relatively small drop in performance. On IWSLT11 PormanntooFngliian task with ED Transformer model, PASSCONC outperforms all baseline methods in almost all sparsity settings (see Table 2 When = 16, PASSCONC achieves a 162.8% speedup, which is more than 20% higher than all baselines, with at least higher BLEU scores.
4.4 Ablation Study
Previous analysis of PASS and PASSCONC in Section 4.3 demonstrates the significant efficiency improve- ments brought about by the concentrator (see Section 3.3 this section, we validate that (1) by introducing the concentrator loss, PASSCONC is able to prune more attention layers entirely (Figure 5a), and (2) the clipping and reopening strategy is necessary for PASSCONC to obtain significant efficiency improvements (Figure 5bh. We report results using BERT-base models and the GLUE benchmark tasks. In Figure 5a we observe that PASSCONC with concentrator is able to prune up to 7 out of 12 attention layers with high sparsily largets, which is 3 layers more than not using concentrator and therefore validates its effectiveness in pruning redundant attention layers entirely to bring significant speedups. In Figure [5b] we observe that without the clipping and reopening strategy, the speedups achieved by PASSCONC can reduce by up 70%! This observation demonstrates the necessity of dynamically ro-activating closed gates to help the model converge cost-effective regions, desired by the concentrator.
10

---

Published in Transactions on Machine Learning Research (09/2024)
Concentrator
180
Clip&Reopen w/o Clip&Reopen
# Pruned Layers
w/o Concentrator
Speedup (%)
160
4-
140
2
120
0
100
20
30 40
50
60
20
30
40
50
60
# Unpruned Heads
# Unpruned Heads
(a) Concentrator
(b) Clipping and Reopening
Figure 5: (a) Concentrator is critical for PASSCONC to prune more attention layers entirely. (b) The clipping and reopening strategy is necessary for PASSCONC to obtain significant efficiency improvements. All results are reported with BERT-base models on the GLUE benchmark tasks.
To disentangle the benefit of the clipping and reopening strategy from the proposed sparsity objective, implement the Voita approach with the clipping and reopening strategy and compare it with PASS and PASSCONC on the SST-2 dataset from the GLUE benchmark. Results are summarized in Figure G By clipping and reopening gates as Voita prunes attention heads, the accuracy of the resulting subnetworks drastically drops 50.9% when K = 16 while PASS and PASSCONC achieve accuracy at least 89.3%. One possible reason that, the clipping and reopening strategy reopens fully closed gates along the pruning process which tends to continue reducing the subnetwork sparsity levels, Without our proposed sparsity objective, Voita fails to converge to the desired sparsity levels and lcads to decayed performance when extract the sparse subnetworks according to the user-specified sparsity targets.
4.5 Distribution of Heads
We visualize the distribution of unprumed hcads bet- ter understand the dynamics of the pruning process (see Figure 7 We report the frequency of heads unpruned by PASSCONG, CONC, averaged over all datasets (IWSLT14 De-En translation task for BD-Transformer and GLUE bench- mark tasks for BERT-base) and sparsity settings (K=16, 32, 48, 64). In each layer, the heads are sorted by the unpruned frequency for clearer presentation. We also port the average number of unpruned heads per layer for both eD-TTrasfformer and BERT- base to provide a
K=84
SST-2
90
Accuracy (%)
K-32
K=16
80
Voita
70
PASSCONC
PASS
60
Voita (Clip&Reopen)
quantitative comparison (see Tables and 4 For ED
Unpruned
50
Transformer (Figure 7a and table 3 we observe that the encoder self-attention heads are typically less important than cross-attention and decoder self-attention heads test time, and therefore pruned more frequently (the average number of unpruned encoder self-attention heads are at least 5 fewer than the other two attention types). This observation suggests the high level of redundancy in encoder attention heads from properly trained Enc-Dec Transformers, which conforms with the observations
100
110
120
130
160
Speedup (%)
Figure 6: Voita augmented with clipping and reopening strategy leads to decayed subnetwork performance.
Li et al. (2021) and Voita al. (2019). For BERT-base (Figure 7b and table we observe that the middle layers (Layer 2,3,4,6) likely to contain unpruned heads, in comparison to the bottom (Layer and
11

---

Published in Transactions on Machine Learning Research (09/2024)
1.0
1-
1
0.8
2
0.8
3
0.6
4
Layers
0.6
Layers
5
6
0.4
7
0.4
9
0.2
0.2
10
11
12
0.0
0.0
Enc Self
Cross
Dec Self
2
3
5
7 8
101112
Heads
Heads
(a) TTD-rrasffomer
(b) BERT-base
Figure 7: Distribution of unpruned heads across layers. Different colors represent the unpruned frequency for each head across all datasets (IWSLT14 De-En translation task for SD-Transformer and GLUE benchmark tasks for BERT-base) and sparsity settings (K=16, 32, 48, 64).
Table 3:
The average number of unpruned heads per layer in ED- -0--srasformer on IWSLT14 De-En translation task across sparsity settings (K=16, 32, 48, 64).
Layer
2
3
4
5
SUM
Enc Sclf
1.67
1.00
1.67
2.17
0.92
2.17
9.58
Cross
1.33
2.00
2.58
2.92
2.50
3.25
14.58
Dec Self
2.58
2.83
3.00
2.58
3.00
1.83
15.83
top (Layer 7-12) layers, which is also aligned with the observation in previous work (Prasanna et al. 2020b Sajjad et al. 2023).
4.6 Overhead Analysis
To realize the efficiency improvements achieved by attention head pruning, we index the unpruned attention heads and only perform matmul the indexed tensors during inference. In Sections 4.2 to 4.4 implement all approaches with indexing and evaluate the speedups w.r.t. the unpruned model using indexing to better understand the efficiency improvements solely caused by different attention sparsities. In this subsection, we investigate the overhead of indexing by comparing our methods with indexing to the umpruned model without indexing (matmul is performed on the original unpruned matrix), using a 72-hcad Encoder-Decoder Transformer on Intel Broadwell CPUs @ 2.2GHz, and NVIDIA P100 Pascal GPU. Comparison results are summarized in Table 5] where we report the average latency achieved by subnetworks of K unpruned hcads CPU and GPU devices scparately.
Table 4: The average number of unpruned heads per layer in BERT- on GLUE benchmark tasks across sparsity settings (K-16, 32, 48, 64).
Layer
2
4
5
6
7
9
10
11
12
Enc Self 3.38
3.81
4.00
4.44
3.19
4.00
2.88
3.25
2.63
3.25
2.13
3.06
12

---

Published in Transactions on Machine Learning Research (09/2024)
Table 5: The average inference latency for subnetworks extracted by indexing unpruned heads, compared to the full model latency without indexing. On CPU-only devices, the subnetworks can achieve speed-ups up to 48 heads umpruned out of 72.
K GPU (1e-4 s) CPU (1c4 s)
4
16
32
48
64
No Indexing
1.6
1.6
3.2
4.6
8
9.1
10.8
11
4.3
16.9
16.2
27.5
44.3
68.1
82.2
95.5
102.9
96
Intuitively, on CPU-only devices where matmul operation is dominantly time- time-consum,ng, the overhead of indexing could be negligible and the subnetworks can achieve speed-ups up to 48 heads unpruned out of 72. In contrast, on devices specialized for in-parallel matmuu operations such as GPUs, the indexing approach may cause non-negligible overheads and head pruning achieves speed-ups only with high-sparsity targets, such when less than 8 heads are retained from the subnetworks, as shown in Table 5
It is worth noting that, indexing in our work is used to reflect the attention sparsities and is only necessary when we assume no specialized software supports. With software supports such as high-performance sparse matrix multiplication libraries the indexing operation is no longer necessary and can be completely removed to avoid overheads.
5 Related Work
Unstructured pruning has been well studied the literature (Gupta Agrawal 2022) and dates back to Op- timal Brain Damage (LeCun et al., 1989). Unstructured pruming prunes individual parameters and identifies subnetworks of high sparsity. However, unstructured pruning hardly achieves practical efficiency improve- ments without specialized software and hardware support (Xia et al. 2022). In contrast, struclured pruning prunes groups of parameters within certain structures (e.g., channels and attention heads). Structured prun- ing has been widely explored in computer vision tasks (He et al. 2017) and has started to attract research interest in the NLP community. Research efforts have been dovoted to designing pruning strategics at both coarse- and fine-grained levels (Xia et al., 2022 Prasanna et. al.] (020aa) over structures like feed-forward layers and attention heads. Previous work on attention head pruning (Li et al.) 2021; Michel al.) 22019) Voita et al., 2019) either highly relies threshold tuning or enforces hard structural constraints model structure in the early training stages, which could lead to limited subnetwork performance. We focus structured pruning and propose the notion of almost-swre sparsity to overcome the above limitations.
In addition to pruning, many other techniques have been developed to obtain inference efficiency for deep learning models. Other than sparsity over the number attention heads, a line of work (Roy et al. 2021 Child et al.] 2019) focuses on sparsifying the attention distribution over tokens for each head to improve efficiency and head diversity. Recently, Correia et al. (2019) and Trrevioo et al. (2021) propose to adaptively sparsify the attention distribution by enforcing low attention scores to be exactly 0 through a-entmax (Peters)
et al. 2019) Linformer (Wang et al. 2020b) develops efficient self-attention design of linear complexity by using low-rank matrix approximation. FlashAttention (Dao et al.) 2022 Dao 2023 focuses on I/O access in atlention computation and achieves linear memory complexity by leveraging the hardware memory hierarchy and minimizing unnecessary data transfers. Other techniques include quantization (Shen et al. 2020). knowledge distillation (Hinton et al., 2015), paramcter sharing (Ling et al.]
2015). tensor dccomposition
(Oseledet) 2011). neural architecture search (Zhang et al. 2024 Zheng et al. 2023 and more. We refer
interested readers to (Menghani 2023 Gupta & Agrawal survey.
2022
Treviso al.
2022 for comprehensive
6 Discussion and Conclusion
We propose a novel notion of almost-sure sparsity, develop a generic framework for Pruning with Almost- Sure Sparsity (PASS) targets, and demonstrate its pruning capacily with attention heads. To further push
ahttes//daoossniidda.ccm////// /
13

---

Published in Transactions on Machine Learning Research (09/2024)
the envelope on inference efficiency, propose novel technique, concentrator, based on which we develop PASSCONC (PASS with CONCentrator).. We also present simple-yell-efecttie strategy to improve subnetwork performance by clipping and selectively reopening learned gates. We investigate PASS and PASSCONC on two widely studied Transformer models: encoder-docoder (ED) Transformer and encoder- only Transformer, BERT-base. Experiments on IWSLT14 German-too-nngissh translation (Cettolo et al. 2014) and GLUE benchmark tasks (Wang al., 2018) demonstrate that PASS and PASSCONG outperform the SOTA methods by identifying subnetworks of up to 1.33 higher BLEU scores, 1.44% higher accuracy, and 60% higher speedups, at the same sparsity levels.
We conclude that PASS and PASSCONC can be used to identify high performance subnetworks and help address the challenge of deploying Transformer models in resouroe-limited applications. We identify several important extensions for future work: (1) More pruning structures and metrics. We would like to explore the possibility of extending the proposed framework to multiple model structures (e.g., ford-forward layers) and prune for meeting other target footprint metrics such as lateney and memory which are more accessíble to users in real-world applications, in addition to sparsity. (2) Combination with other tech- niques. Since both PASS and PASSCONC are agnostic to the underlying self-attention implementation, it is intriguing to investigate the compound efficiency improvements achieved by combining our approaches with other efficiency improvement techniques such as flash attention (Dao et al. 2022 Dao 2023) and quan- tization (Shen et al., 2020; Frantar et al. 2022). For example, a naive combination strategy is to first apply our attention pruning techniques to reduce the redundant heads and then leverage quantization or flash attention techniques to accelerate the computation of those umpruned heads. In the future, we would like to explore other non-trivial combination strategies to further push the envelope efficiency gains. (3) Large Language Models (LLMs) pruning. LLMs are typically large-scale decoder-only Transformer models (Chang et al. 2024). Recent work (Xia et al. 2023) shows that LLMs possess high level of parameter redundancy which can benefit from model pruning to improve inference efficiency while maintaining high performance. Though we have demonstrated the effectiveness of our approaches on both encoder and de- coder modules in standard Transformer models, it is still intriguing to test the scalability of our approaches on large scale Transformer models such as L.L.s. On the one hand, is worth noting that our sparsity convergence analysis agnostic to the model size and attention types (see Lemma [ For decoder-only LLMs, we believe our approach will be able prune subnetworks to desired sparsity levels. On the other hand, we acknowledge that currently the intrinsic mechanism behind the impressive emergent abilities of LLMs (Wei et al. 2022) remains unclear. How to effectively and efficiently prune attention heads from LLMs brings new challenges and opportunities which we plan to investigate in our future work.
References
Akari Asai, Zeqiu WU, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-RAG: Learning to re- trieve, generate, and critique through self-reflection. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreriii www.aaarrrreeeeeeeoouuuuuiaeeeee net/fforurrdd-hhyyy5oorvv
Matthew Beal. Variational algorithms for approximate bayesian inference /. PhD thesis, 01 2003.
Maximiliana Behnke and Kenneth Heafield. Losing heads the lottery: Pruming transformer attention in neural machine translation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 2664 2664-2674, Online, November 2020. Association for Computational Linguistics. doi: 22.1663/vv12020.0. eennn-mmmalppaae main. 211. URL https://aclantth// wrg.hholooyyorg/222000 emn.p-mmmmnp-mm....
211
Léon Bottou. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT2010: 19th International Conference on Computational StatisticsParis France, August 22-27, 2010 Keynote, Invited and Contributed Papers, pp. 177-186 Springer, 2010.
Mauro Ccttolo, Jan Nichues, Scbastian Stüker, Luisa Bentivogli, and Marcello Federico. Report on the 11th iwslt evaluation campaign, iwslt 2014. In Procoedings of the International Workshop on Spoken Language Translation, Hanoi, Vietnam, volume 57, 2014.
14

---

Published in Transactions on Machine Learning Research (09/2024)
Yupeng Chang, Xu Wang, Jindong Wang, Yuan Wu, Linyi Yang, Kaijie Zhu, Hao Chen, Xiaoyuan Yi, Cunxiang Wang, Yidong Wang, et al. A survey on evaluation of large language models. ACM Transactions on Intelligent Systems and Technology, 15(3):1 45, 2024.
Rewon Child, Scott Gray, Alec Radford, and Nya Sutskever. Generating long sequences with sparse trans- formers. arXiv preprint arXiv:1904.10509 2019.
Gonçalo M Corrcia, Vlad Niculac, and André FT Martins. Adaptively sparse transformers. arXiv preprint MrXiv:1909.00015 2019.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint MrXiv2337.00661,, 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memoor-eficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024. URL htttps://openrevie:
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirec- tional transformers for language understanding. CoRR, abs/1810.04505, 2018. URL http://arxiv.org
aba/1810. 04805
Dujian Ding, Ankur Mallick, Chi Wang, Robert Sim, Subhabrata Mukherjee, Victor Rühle, Laks V. S. Takshmanan, and Ahmed Hassan Awadallah. Hybrid LLM: Cost-efficient and quality-aware query routing. The Twelfth International Conference on Learning Representations, 2024. URL https://openrevieu..
net/forum?id=02E3mull 2id=02E3mUequi
Yanzhuo Ding, Yang Liu, Huanbo Luan, and Maosong Sun. Visualizing and understanding neural machine translation. In Proceedings of the 55th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1150-1159, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10.18653 N..18853/11/P11-1100 URL https://aclaattholooe ortp:p7711o.0////1111100
Elias Frantar, Salch Ashkboos, Torsten Hoefter, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323 2022.
Emil Julius Gumbel. Statistical theory of extreme values and some practical applications : A series of lectures. 1954.
Manish Gupta and Puneet Agrawal. Compression of deep learning models for text: A survey, ACM Trans. Knowl. Discov. Data, 16(4), jan 2022. ISSN 1556-4681. doi: 10.1145/3487045, URL https://doi.orr/10
1145/3487045
Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. Proceedings of the IEEE international conference on computer vision, pp. 1389-1388 1999-197,490 2017.
Geoffrey Hinton, Oriol Vinyals, Jelf Dean, et al. Distilling the knowledge in neural network. arXiv preprint MtXiv:1503.02531 2(7), 2015.
Wenyue Hua, Yingqiang Ge, Shuyuan Xu, Jianchao Ji, and Yongfeng Zhang. Up5: Unbiased foundation model for fairness-aware recommendation. arXiv preprint 12X0 2023.
Wenyue Hua, Jiang Guo, Mingwen Dong, Henghui Zhu, Patrick Ng, and Zhiguo Wang. Propagation and pitsalls: Reasoning-basee based assessient of knowledge editing through counterfactual tasks. arXiv preprint arXiv:2401.17585, 2024.
Jean Jacod and Philip E. Protter. Probability essentials. Springer, 2003.
15

---

Published in Transactions on Machine Learning Research (09/2024)
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, TCLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. CRL http:///xxiv.oorraa////////
Anastasia Koloskova, Hadrien Hendrikx, and Sebastian U Stich. Revisiting gradient clipping: Stochastic bias and tight convergence guarantoes. arXiv preprint arXiv:2305.0158; 2023.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
Jiaoda Li, Ryan Cotterell, and Mrinmaya Sachan. Differentiable subset pruning of transformer heads, 2021. .46557
URL attps://arxiv..t
Wang Ling, Chris Dyer, Alan W Black, Isabel Trancoso, Ramón Fermandez, Silvio Amir, Luís Marujo, and Tiago Luís. Finding ionction in form: Compositional character models for open vocabulary word representation. In Procoedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 1520 1530, Lisbon, Portugal, September 2015. Association for Computational Linguis- tics. doi: (L18833/v1/DDLL 11761--1176. URL h www.at/oooooooogg/6111111 org/D15-1176
Christos Louizos, Max Welling, and Diederik P. Kingma. Learning sparse neural networks through to regu- larization, 2017. URL https://arxiv. org/ abtp//1/a/11...../ 013122
Christos Louizos, Max Welling, and Diederik Kingma, Learning sparse neural networks through to regu- larization, 2018.
Chris J. Maddison, Andriy Mnih, and Yee whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. CoRR, abs/1611.00712 2016. URL
Gaurav Menghani. Efficient deep learning: A survey on making deep learning models smaller, faster, and better. ACM Computing Surveys, 55(12):1-37, 2023.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? CoRR, abs/1905.10650, 2019. URL httpp://aazxvv.r/abbb////ss
Ivan V Osclodets. Tensor-train decomposition. SIAM Journal on Scientific Computing, 33(3):2295-25 23177, 2011.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling, arXiv preprint arXiv:1904.01038, 2019.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pp. 311-318, 2002.
Ben Peters, Vlad Niculac, and André FT Martins. Sparse sequence-to-soqqeeec models. arXiv preprint arXiv:1905.05702, 2019.
Sai Prasanna, Anna Rogers, and Anna Rumshisky. When BERT Plays the Lottery, All Tickets Are Winning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 3208 3229, Online, November 2020a. Association for Computational Linguisties. doi: 18.16633vv1 2020.emnlp-main.c main.259. URL
wttp://..gg/2//.//1.mm.. emnlp-main. 259
Sai Prasanna, Anna Rogers, and Anna Rumshisky. When bert plays the lottery, all tickets are winning. arXiv preprint arXiv:2005.00661 2020b.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint wwXiv:1606.05250 2016.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-basod, sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53 68, 2021.
16

---

Published in Transactions on Machine Learning Research (09/2024)
Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. On the effect of dropping layers of pre- trained transformer models. Computer Speech & Language, 77:101429, 2023.
Lakshay Sharma, Laura Graesser, Nikita Nangia, and Utku Evei. Natural language understanding with the quora question pairs dataset. arXiv preprint wtXiv:1907.0104. 2019.
Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Q-bert: Hessian based ultra low precision quantization of bert. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 0815-8821, 2020.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y 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, 2013.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference Machine Learning, pp.
PMLR, 2020.
Marcos Treviso, António Góis, Patrick Fernandes, Erick Fonseca, and André FT Martins. Predicting atten- tion sparsity in transformers. arXiv preprint arXiv:2109.1211 12188, 2021.
Marcos Treviso, Tianchu Ji, Ji-Ung Lee, Betty van Aken, Qingqing Cao, Manuel R. Ciosici, Michael Hassid, Kenneth Heafield, Sara Hooker, Pedro H. Martins, André F. T. Martins, Peter Milder, Colin Raffel, Edwin Simpson, Noam Slonim, Niranjan Balasubramanian, Leon Derczynski, and Roy Schwartz. Efficient methods for natural language processing: A survey, 2022. URL https://arziv. www..ogg///a........ooo hogg/ab//2209.... 00099
Arash Vahdat, Evgeny Andriyash, and William Macready. Dvae#: Discrete variational autoencoders with relaxed boltzmann priors. Advances in Neural Information Processing Systems, 31, 2018a.
Arash Vahdat, William Macready, Zhengbing Bian, Amir Khoshaman, and Evgeny Andriyash. Dvae++: Dis- crete variational autoencoders with overlapping transformations. In International conference on machine learning, pp. 5035 5044. PMLR, 2018b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR,
2017. URL http://arxiv.om/ org/
abs/1706. 03766
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. CoRR, hass/1905.0941, 2019.
URL http:://rxix... http::/arxtv.orr/aab////0 absp///s//111...... 09418
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A
multi-task benchmark and analysis platform for natural language understanding. CoRR, abs/ 1bs/1180.00461,
2018. URL http://arxivv www//bs/////eeeoooeo/../eeeeeeet aas////ss////// 1804. 07461
Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade token and head pruning, 2020a. URL https://arxiv.. wrgg/be//// .98552
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with lincar complexity. arXiv preprint arXiv:2006.04777 2020b.
Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Proccedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 6151-6162, 6162, Online, November 2020c. Association for Computational Linguistics. doi: 10. 1ww51//v1/22200ooo.m wttp2/00000000mmmmeem main. 496. URL https: org//020. emnlp-main.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Scbastian Borgcaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. Transactions Machine Learning Research, 2022.
17

---

Published in Transactions on Machine Learning Research (09/2024)
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 Linguisties: Human Language Technologies, Volume 1 (Long Papers),
pp. 1112-1122. Association for Computational Linguistics, 2018. URL N18-1101
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of-the-arg natural language processing. In Procoodings of the 2020 conference on empirical methods in natural languago processing: system demonstrations, pp. 38 45, 2020.
Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models, 2022. URL
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared Ilama: Accelcrating language model pre-training via structured pruning. arXiv preprini arXiv:2310.06694 2023.
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: theoretical justification for adaptivity. arXiv preprint arXiv:1905. 11881, 2019.
Shaokun Zhang, Xiawu Zheng, Guilin Li, Chenyi Yang, Yuchao Li, Yan Wang, Fei Chao, Mengdi Wang, Shen Li, and Rongrong Ji. You only compress once: Towards effective and elastic bert compression via exploit-cxplore stochastic nature gradient. Neurocomputing, 599:128140, 2024.
Xiawu Zheng, Chenyi Yang, Shaokun Zhang, Yan Wang, Baochang Zhang, Yongjian Wu, Yunsheng Wu, Ling Shao, and Rongrong Ji. Ddpnas: Efficient neural architecture search via dynamic distribution pruning. International Journal of Computer Vision, 131(5))1234 1249, 2023.
A Proofs and Analysis
A.1 Derivation of the Upper Bound in Eq. 4
Here we present how derive the upper bound in Eq. 4
loog Ep(
[P [P @ z)]
1 log
P(D)00 2))
(z)
= log
p(z) D)P(D
(z)
(10)
D)
= log (x)
P(D O
g0 (z)
p(z) D)
< Ez) log
P @ 2)
40 (z)
= (x (z) [log P(DI0 3 z)] + KL (ox (z) 1| D))
A.2 Hard Concrete Distribution (Louizos et al. 2018
Hard Concrete distribution is derived from the Concrete distribution (Maddison et al. 2016 The cumulative
distribution function of a binary Conrete random variable zf E (0, 1) is as lollows:
(11) constant. We can obtain the Hard Concrete distribution by first stretching the support of the Concrete distribution from (0, 1) to (v, {), where 7 < 0, Ç 1, and then collapsing the probability mass on (7, 0] (resp. [1, ()) to the endpoint (resp. 1).
Q.(2f:8,04) = sigmoid ((log 2 log(1 - = 3 +x)
where sigmoid(x) = I+e-z is the sigmoid function and 0 < B <
18

---

Published in Transactions on Machine Learning Research (09/2024)
Equivalently, for each Concrete random variable 2f E (0. 1), can obtain a Hard Concrete random variable
2; = -min(1, max 25 7) + 7). For each 2i, the closing probability g(z; = 0; Pr(2f
S إير
Qa( B, Φt), and the opening probability ((= = 1; = Pr(=f 2 ( = 1 - Qa 2o( B, 0.). By plugging
in Eq. we have the gate closing and opening probabilities.
A.3 Rbose Local Optimum Analysis
We show that the gradient descent direction of Rbase always leads to a higher opening probability for gate z; if Φb; 2 log( log(= -l-----gloaa---) ) where g(a) = 2e -e e " a = 3 log( -2) We assume that the Hard Concrete distribution is equally stretched in both directions, which gives + 1 and has been widely adopted in previous work (Voita et al. 2019 Xia et al. 2022).
[(a)
focus on the high sparsity target situation (s|0) > E (o())) and ignore the sparsily constraint 8 in following compulation because it does not change the equation. We first compute the derivative of Roase (Φ) with respect to di,
Without loss of generality,
ÜBbase () X (2sig(-0: + a)sig(di - a)
Ib:
(12)
-sig( 0, s)si((> aa))
where x > a= Blog Bl (-7) and 3 0, 1 In order to compute the domain for :Rams:
0,
00:
we can equally solve the following inequalities,
aBsase (
00
0o:
2sig( 0i a)sig(0; - a) S
(13)
sig( 0; a)sig(dr aa)
sig( 6; + a)
sig(0; + a)
<
sig( 0. - a) sig(0: a)
where the last step is due to the non-negativity of the sigmoid function. By denoting fa(Ø)) can simplify the incquality to be 2fa(- I f(()). The derivative of fa())) with respect to di is as follows,
abj we
sig(dsta)
(or)
sig(o, + a) sig(Ø. a)
(sig(-ò; - -a) --ii(---r + +a)
(14)
0Φ.
Because = 9 log ( = Blog (5²) due to the fact that the sigmoid function monotonically increasing. As a result, the derivative of fox) with respect to oi is always positive (ff((0)) > 0) and therefore fa(di) is monotonically increasing while (=)- monotonically decreases as function of di.
< 0, we have - -0. +a and hence sig(-o;-a) sig(-0, 00: +a) +a) -0,
Assuming there is of having the equality 2fa(-o() = fa(d). It is easy to show that for di 2 %, the incquality 2fa (-Φ+) S fa())) holds true due to the monotonicity of fa())) and fa(-Øs), and hence we have @Raave (0) < 0 (the gradient descent direction). Notably, it indicates that, for di 2 6/, any gradient-descent optimizer (e.g., Adam (Kingma & Bal (2015)) will continue increasing the value of di which opens more gates and leads to undesired low sparsity, as dicussed in Section 3.2 We derive the value of oi by solving
do:
19

---

Published in Transactions on Machine Learning Research (09/2024)
equation 2fa(-01 fa(d)) as follows,
2 -1) fa(o))
sig( - + +o)
sig(o; + a)
sig 0f a)
siglo a)
+ etaa
+ e iita
(15)
+ 2 + e0
1 e -M-a
+ (2c" ") e
-2 + (e" 22 2)=
+ (2e" - e e e 20 e = -0 + E" a
which is quadratic equation with respect to ex: and gives us of log( 1 ala))(aaa))))))
where g(a)
s(a)
2e -ea
A.4 Proof of Lemma
We include the proof of Lemma Mas follows.
Proof: We prove this lemma by showing that, in every possible case, the gradient descent direction of the objective Rpass always leads to subregion in the scarch space where the expected sparsity is no more than 8 and expected density is no more than - 8. Recall that we define erpected density as 1 |o| E q (6i), and the expected sparsity as 1 Σ (o())). For the simplicity of proof language, use Ea to denote the expected density and E, to denote the expected sparsity. Intuitively, as the expected density increases (resp. decreases), the expected sparsity will monotonically decrease (resp. increase). One important observation is that, the sum of expected density and sparsity is always less than 1, due to the fact that qı (Øx) + go( 1 1 (0)n and (mb(())) 0.
[o|
Given a sparsity target depending on the values of expected density and sparsity, there are three possible situations: (i) the easiest case is that the model is already neither over-sparse nor over-dense (i.c., Ea S 1 11- and E. < s). In this case, can rewrite Rpass (2, 2> i- Ind (o.). Minimizing R s) amounts to minimizing (no( which polarizes the gates to achieve the required almost-sure sparsity, until either Eg = - 8, E, < or Eg < 1- 8, E, = 8 both of which satisfy Lemma [ (ii) high expected density and low expected sparsity (i.c., Ed 21 obtain Rpase(0,
-1- 21-a and Eg S s), in which case we can rewrite Eq. 7 to Σ (n.( +|01 (0(
Σ 1= (1.)).. Because gno(ó.) + t1 (0.)
B 1 1-9u (o()),, we can further simplify the equation to be Rpase ) = 2 " i=1 (s - (o-())))) Clearly, minimizing Rpass ( amounts to increasing (o(b)) that leads to lower expected density and higher expected sparsity until Ed = 8, E, < 8, which satisfies Lomma [ (iii) low expected density and high expected sparsity (i.e., Es < 1 and E, 2 s). Similarly, we can rewrite Rpaee (9,))((.s)
2 | (1 - (o:)). Minimizing Rpass(0, ()) amounts to increasing & ((s)) that brings down expected sparsity as well as increases expected density until Ea < 1 8, E. = 8, which satisfies Lemin [ It is impossible to have both high expected density and expected sparsity due E, + E, < We have shown that Lemma holds in all possible cases.
B Hyper-parameter Settings
We adopt an exponential escalation strategy to increase the value of regularization coefficient. X.
2 Absee
X #ftrain its/ stop
(16)
where Atase and A are hyper-paramnters. We choose #n__seee_ step as 1, 000 in all experiments. We use the same inverse square root learning rate schedules for model training as in Li al. (2021). Values of corresponding hyper-pane ppper-parrmeters are summarized in Table
As for PASSCONC, X. is defined as the minimal ratio between sparsifier gradients and concentrator gra- dients of Φi 6 for all heads to ensure that the concentrator is not dominated by the sparsifier while the
20

---

Published in Transactions on Machine Learning Research (09/2024)
Table 6: Hyper-parameters
ED Transformer
BERT-bas
Abase
le-5
A
1000
learning rate for +
0.2
0.5
Table 7: Subnetwork performance at different sparsity levels, on MNLI and QQP.
Accuracy (MNLI) (%)
(MNLI)
Accuracy (ocuracy(QQP)) (%)
PASS
PASSCONC
DSP
LAG
Voita
PASS
PASSCONC
DSP
LAG
Voita
16
79.38
78.38
80.69
78.19
80.14
89.18
88.50
89.25
86.87
85.53
32
80.73
79.63
82.64
81.33
80.30
89.94
89.36
90.09
88.99
89.99
82.59
82.51
82.92
82.80
82.03
90.57
90.40
90.16
89.67
90.20
64
83.23
83.06
83.55
83.74
75.70
90.80
90.74
90.28
90.02
87.21
concentration effect is modest to avoid damaging model training quality.
= A-A.min min { RoRss
&Ronne
| E
(17)
abs
An is adaptively updated in the middle of training, and set to O in both the early training phase and the last few training iterations to help improve model performance and achieve sparsity convergence. Specifically, de is set to 0 during the first 20, 000 itcrations and the last 7, 7,000 iterations when training ED Transformer models. For BERT-base models, de is set except for ilerations between 2, 000 and 5,000.
As to the clipping range, there trade-off between choosing a larger range (e.g., [-0, 10]) and a smaller one (c.g., [-1, 1)). In general, larger range allows us to better achieve subnetworks of desired almost-sure sparsities but the gradients will vanish to excessively small values (see Figure 3c| and make the gradient- descent optimization extremely slow. On the other hand, small range allows efficient gradient-descent optimization but the corresponding gates may be neither almost surely closed nor opened (see Figure 3b) and therefore fail the almost-sure sparsity guarantees. For example, for a large range [-10, 10), the smallest gradient approaches 2e-5 (see Figure 3cb which leads to slow optimization rates, while for a small range (1, 1], the maximal gate open/close probability only 55% (sce Figure 3b) that neither almost surely opened nor closed. We
held-out set to empirically select. [-5, 5] as the clipping ranges because it gives maximal gate open/close probability that is roughly 99% while the smallest gradient is above 0.003, which is practically sufficient for efficient optimization according to our experiments.
C More Performance Results
In this section, provide detailed performance and speedup results each of the 4 GLUE benchmark tasks: MNLI, QQP, QNLI, and SST-2 (see Section 4.1.1.. Performance results are summarized in Tables 7 and 8 which resemble our analysis in Section 4.2 In general, PASS and PASSCONC are able to match or outperform all 3 basclincs across a majority of experiment scttings. Speedups achieved by different approaches are presented Tables g and 10 Notably, PASSCONC outperforms all baselines almost all cases which demonstrates its effectiveness in delivering efficiency improvements. Lastly, we report the Pareto front between model performance and speedups for subnetworks identified by different methods (sce Figure 8 where PASSCONC dominates all other methods in most cases by achieving better performance--ffccinneyy trade-offs, similar to our observation in Section 4.3
21

---

Published in Transactions on Machine Learning Research (09/2024)
MNLI
QQP
91
LAG
LAG
84
Accuracy (%)
Voita
Accuracy (%)
90
Voita
DSP
DSP
82
89
PASS
PASS
80
PASSCONC
88
PASSCONC
Unpruned
Unpruned
78
87
86
76
100
150
200
250
100
120
140
160
Speedup (%)
Speedup (%)
QNLI
SST-2
92
LAG
Accuracy (%)
90
Voita
Accuracy (%)
DSP
90
LAG
88
PASS
Voita
PASSCONC
DSP
86
Unpruned
86
PASS
PASSCONC
84
84
Unpruned
100
110
120
130
140
100
110
120
130
140
150
160
Speedup (%)
Speedup (%)
Figure 8: Atlention speedups v.s. subnetwork performance on MNLI, QQP, QNLI, and SST-2, The unprumed model performance is included for reference purpose.
Table Subnetwork performance different sparsity levels, on QNLI and SST-2.
AcCuracy(QNLI) (%)
Accuracy (SST-22 (%)
PASS
PASSCONC
DSP
LAG
Voita
PASS
PASSCONC
DSP
LAG
Voita
86.03
84.79
84.59
83.49
84.90
90.48
89.33
83.37
86.81
88.76
88.43
85.83
88.38
87.50
87.59
91.28
91.06
92.32
90.14
90.71
48
89.13
88.80
89.58
89.13
88.91
92.32
91.51
91.40
90.48
90.94
64
89.57
89.51
89.68
89.69
82.61
91.28
91.17
91.74
90.14
91.28
Table 9; Attention speedups at different sparsity levels, MNLI and QQP.
Speedup(MNLI) (%)
Speedup QQP) (%)
K
PASS
PASSCONC
DSP
LAG
Voita
PASS
PASSCONC
DSP
I.AG
Voita
16
124.2
262.7
112.5
113.2
132.9
108.2
171.2
112.7
133.1
125.4
32
117.7
141.3
106.7
103.5
102.8
103.4
138.8
112.3
107.8
107.8
48
103.2
107.4
102.5
102.5
106.7
103.4
118.9
102.7
102.4
103.0
64
103.5
107.8
102.8
102.5
104.1
103.4
107.5
102.7
103.4
102.1
22

---

Published in Transactions on Machine Learning Research (09/2024)
Table 10: Attention speedups at different sparsity levels, on QNLI and SST-2.
Speedup(QNLI) (%)
Sppedup(SSTT22 (%)
K
PASS
PASSCONC
DSP
LAG
Voita
PASS
PASSCONC
DSP
LAG
Voita
16
113.8
143.7
118.3
127.5
138.8
111.5
163.3
148.8
107.9
107.2
32
105.1
125.6
103.1
108.2
107.9
102.2
136.7
106.9
102.2
102.8
48
103.8
103.5
103.1
103.5
102.8
101.9
107.2
102.5
103.1
102.5
64
103.8
107.5
103.1
103.1
103.1
102.2
102.8
102.8
102.8
102.8
23


Reviews:
[{"review_id": "WAWWNwtvG6", "reviewer": "Reviewer_wny5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- inference speed-up of transformers is a topic of clear interest to the audience of this journal\n- the paper is clearly written, well structured, and the appendices offer meaningful support\n- good theoretical treatment and justification of the proposed learning objective with almost-sure sparsity\n- the custom learning objective for PASS, and the additional \"concentrator\" and \"clipping and reopening\" strategies appear to be novel\n- the best combination of strategies results in remarkable speedups with limited accuracy degradation\n\nWeaknesses\n- clipping and gate reopening (quite effective in improving PASSCONC speed-ups) are always applied concurrently. What is the individual effect of these two distinct strategies? Have these been applied only to PASSCONC or to PASS as well?\n- regarding clipping, $\\phi_i$ is bound to a range of [-5,5], selected empirically. What is the impact of changing this threshold? Are the results sensitive to this arbitrary selection?\n- in an effort to highlight which heads are most commonly pruned, Fig 6 averages the pruning frequencies across multiple sparsity settings and, for BERT, also across tasks; the averages are then sorted; however, without additional information on run-to-run variations, it's unclear what conclusions can be drawn from it; are the same heads consistently pruned as the sparsity target is raised? are they consistently pruned from one task to another? \n- in addition, the text related to Fig 6b states that the heads in the middle layers (2-7) of BERT are more important and consequently pruned less frequently; this qualitative observation can't be easily derived from the figure itself; overall, I think a better visualization for Fig 6a,b or the use of a more concrete metrics for \"overall head pruning\" is warranted\n- results are reported on just 2 models and limited tasks; it's unclear whether the proposed techniques are scalable to larger, more relevant models (admittedly, this limitation is mentioned by the authors as future work)\n\nOther comments\n- it should be specified that the model used is BERT-base; the manuscript does state it is a 12-layer BERT model and provides a citation to the original paper, but it's still worth using the common name for this model to facilitate the reader's understanding", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a methodology (Pruning with Almost Sure Sparsity, PASS) for pruning attention heads in transformer models to achieve desired sparsity targets. This method relies on training head gates with an objective that ensures a binary open/close outcome for each gate. In addition, a concentrator is proposed which modifies the objective to encourage the unpruned heads to be concentrated in few layers as possible, so that entire layers can be pruned out (PASSCONC). A strategy for clipping and reopening gates is also introduced. Results on two transformer models show remarkable inference speedups with limited accuracy degradation, typically outperforming competing methods for attention head pruning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See Weaknesses in previous section.", "broader_impact_concerns": "Not address in the manuscript. No concern on my end.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "G50LbQLnN5", "reviewer": "Reviewer_yuW5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. sufficient experimental data to demonstrate the benefits over prior works.\n2. Benefit of Concentrator is supported by Fig 5a. Benefit of clipping/reopening is supported by Fig 5b.\n\nWeaknesses:\n1. The reasoning of using an asymmetric design of probability function, i.e. q(). \nCompared to previous works, the choice of probability function is unclear. Intuitively, one would choose a probability function that enables \"gate_close_probability = 1.0 - gate_open_probability\" which naturally would result in a symmetric q(phi) about phi=0, such as used in Louizos 2018 Fig 2b. However, the author chose to use 2 separate probability functions for open and close and they are not complementary, as shown in Fig 2b of this manuscript. For example, when phi=0, q_open and q_close are both ~0.3, which further requires the introduction of an extra variable q_nb = 1 - q_open - q_close. Because of this, when trying to encourage q_0 to converge to sparsity s in Eq 6, q_1 will not automatically satisfy (1-s), hence the need of introducing an extra term in Eq 7 to resolve this. If the open and close probability were complementary for a given phi, as in Louizos, the objective function could be simplified greatly. Author may want elaborate on this seemingly unnatural choice. \n\n2. Insight on the source of speed-up. \nDepending on actual code implementation, head pruning could be translated into different computation schemes. One example (as in HuggingFace implementation) is that the pruned Linear layer will store a smaller weight tensor. And together with an additional indexing/slicing step on the incoming activation tensor, it can be used together with the pruned weights in the matmul. One important consideration here is that all heads in a layer, no matter how many there are, are computed at the same time. Splitting by heads and computing them one after another then concatenate would be very inefficient. In this scenario, speed-up would derive from performing matmul on smaller tensors, but the overhead of indexing activations is not negligible. It is unclear how authors implement their codes, but a brief discussion or profiling about the relative saving in compute vs additional overhead from indexing could give the readers a better idea how to benefit from head pruning schemes.\n\nAdditionally, according to Table 2 and Fig 5a, the main speed-up of this work comes from layer-pruning, not head-pruning, indicating the saving in computation is comparable with the overhead even when more than 50% of the heads were pruned (ie, BERT, keeping 64 out of 144 heads). But layer-pruning would be less preferrable as it will likely result in much more discernable accuracy degradation compared to fine-grained structured pruning at the same sparsity. (For example, https://developer.nvidia.com/blog/accelerating-inference-with-sparsity-using-ampere-and-tensorrt/). That being said, the reviewer would suggest to add some discussions, and experimental data if possible, to consider the memory-bound cases in addition to compute-bound cases. Especially for LLM, faster or reduced computation may not benefit as much because most of the time were spent on transferring weights and large tensors. If one considers pruning as a compression technique, under memory-bound situation, it might exhibit a much greater speed-up. A good example would be GPTQ in the quantization field.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a method for attention head pruning which incorporates user-specified sparsity target into objective function so that the desired sparsity can be achieved without additional hyper-parameter tuning. Furthermore, this paper introduced a second term into the objective function that encourages the unpruned heads to be localized in fewer layers. In another word, there will be a better chance that all heads in a certain layer, and this layer itself, could be pruned, and therefore, the computation of this entire layer could be skipped. To address potential training challenges, gradient clipping and gate reopening are employed. Finally, accuracy and speed-up are demonstrated on encoder-decoder Transformer and BERT, showing improvement over prior works.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "please see weakness above.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "J23Nz0OhCJ", "reviewer": "Reviewer_j5gS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengthes:\n1. The paper introduces the concept of \"almost-sure sparsity\" and leverages it effectively for attention head pruning.\n\n2. PASS and PASSCONC consistently outperform baseline methods across various sparsity levels and datasets, achieving both high accuracy and significant speedups.\n\n3. The paper provides a clear mathematical formulation and justification for the proposed methods. \n\nWeaknesses:\n\n1. Although the evaluations are thorough for the presented tasks, the paper does not investigate the scalability of the methods to very large models or other language models, such decoder-based large language models (LLMs).\n\n2. The success of the methods appears to depend on careful tuning of hyperparameters, such as the regularization coefficients and clipping ranges. The paper could benefit from providing more discussion or guidelines on selecting these hyperparameters for different tasks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces PASS and PASSCONC, two innovative methods for pruning attention heads in Transformer models. These methods leverage the concept of \"almost-sure sparsity\" to address the limitations of previous probabilistic pruning techniques, achieving desired sparsity levels with minimal performance degradation. Experiments on translation and GLUE benchmark tasks demonstrate that these methods outperform existing pruning techniques.\n\nContributions:\n1. The authors propose a new concept termed \"almost-sure sparsity,\" where the probability of a gate being open or closed approaches certainty. This ensures that the pruned model maintains high performance by avoiding the common threshold tuning problems associated with other pruning methods.\n\n2. PASS: The paper introduces the Pruning with Almost-Sure Sparsity (PASS) framework, which enables effective pruning of attention heads during Transformer model training. This method involves learning a set of gate variables for each attention head, determining their likelihood of being pruned.\n\n3.PASSCONC: To further enhance efficiency, the authors present a novel technique called \"concentrator\" within the PASS framework, named PASSCONC. This technique concentrates active attention heads in fewer layers, potentially reducing computational overhead significantly.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "To strengthen the work:\n\n1. Include a more detailed discussion or guidelines on selecting hyperparameters, such as regularization coefficients and clipping ranges, for different tasks and models. This could involve providing heuristic methods, empirical results, or a systematic approach to tuning these parameters to achieve optimal performance.\n\n2. The paper could benefit from a more thorough discussion on the applicability of these methods to other domains or more extensive models, for example decoder-only LLMs. This could include theoretical considerations or preliminary empirical results on diverse model architectures and tasks, highlighting the versatility and potential impact of the proposed methods.\n\n3. A more detailed comparative analysis with other pruning methods, beyond just attention head pruning, could provide a clearer picture of the strengths and potential limitations of the proposed methods. This analysis could include comparisons with structured pruning, unstructured pruning, and other state-of-the-art techniques, supported by quantitative metrics and qualitative insights.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a method for attention head pruning which consists of training head gates variables with an objective that ensures the network has expected sparsity matching the desired target sparsity. The authors also propose another variant which introduces another term in the objective which encourages the unpruned heads to be concentrated in few layers, to enable pruning entire attention layers. They also employ a clipping and gate reopening strategy to address training challenges. Experiments on two types of transformer models (encoder-decoder and encoder-only) on a translation task and GLUE tasks show improvements over baselines.

Strengths:
- Paper is clearly written
- The proposed method PASS is motivated theoretically
- Experiments show that the proposed methods outperform baselines
- Ablation studies show the benefit of the additional concentrator loss and of the clip/reopening strategy on PASSCONC. 

Weaknesses:
- No evaluation on large models or decoder-only models like LLMs.
- Proposed methods require tuning multiple hyper-parameters.
- Compute saving from pruning can be offset by indexing overhead on GPUs at medium to low sparsity.
- Missing implementation details for baselines and code not provided for now.
- It is unclear if the superior performance of PASS is due to the chosen sparsity objective or the clipping/reopening strategy or both. Need more ablation experiments to evaluate that.
- Important baseline missing. 

The overall recommendations from the reviewers were mixed, with two reviewers leaning to accept and one leaning to reject. I summarized above the strengths and weaknesses highlighted by reviewers, as well as my own assessment. The authors revised the paper to address some of the reviewers comments but not all.  I recommend to "accept with minor revision". I request the following revisions to be made:
- Compare with the method of  [Voita et al., 2019] with clipping/reopening applied to it, to disentangle the benefit of clipping/reopening from the proposed sparsity objective.
- Compare with the CoFi method of [Xia et al, 2022], adapted to prune only heads and attention layers (so without pruning FFN layers). This method outperforms for example the method of [Wang et al, 2020c], which is the most recent baseline considered in the experiments. Alternatively, explain why you did not compare with this method.
- Release code to allow reproducing all empirical results.
- Include the discussion on the trade-off between saving in compute from pruning heads vs additional overhead from indexing and the table shown in the rebuttal (from your response to Reviewer yuW5). Clarify the apparent mismatch between the results in the paper showing speedups on a GPU with up to K=64 unpruned heads, and the ones in the new table showing that speedups on a GPU are only achieved up to k=4 unpruned heads!
- Include the discussion on how you chose the clipping range and on the importance on applying clipping and gate reopening jointly (from your response to Reviewer wny5).
- Provide more details on how you implemented the baselines.
Decision: Accept (poster)