Structural Pruning of Pre-trained Language Models via Neural Architecture Search
Authors: Aaron Klein
Jacek Golebiowski
Xingchen Ma
Valerio Perrone
Cedric Archambeau
Venue: Accepted by TMLR ()

Abstract: Pre-trained language models (PLM), for example BERT or RoBERTa, mark the state-of-the-art for natural language understanding task when fine-tuned on labeled data.  However, their large size poses challenges in deploying them for inference in real-world applications, due to significant GPU memory requirements and high inference latency.
This paper explores neural architecture search (NAS) for structural pruning to find sub-parts of the fine-tuned network that optimally trade-off efficiency, for example in terms of model size or latency, and generalization performance. We also show how we can utilize  more recently developed two-stage weight-sharing NAS approaches in this setting to accelerate the search process. Unlike traditional pruning methods with fixed thresholds, we propose to adopt a multi-objective approach that identifies the Pareto optimal set of sub-networks, allowing for a more flexible and automated compression process.

Published in Transactions on Machine Learning Research (09/2024)
Structural Pruning of Pre-trained Language Models via Neu- ral Architecture Search
Aaron Klein Amazon Web Services
kleinoro Batasrro@mmzzoo.com
Jacek Golebiowski Amazon Web Services
jacekgo
Xingchen Ma Amazon Web Services
zgchenmaGomczzon.com
Valerio Perrone
vperrone@amaazon.com
Amazon Web Services
Cedric Archambeau Helsing
cedric. Sabree.rcchmmbauu@hhllinngg
Reviewed on OpenReview: https: openreiiew. net/ forum? id- XIXBX
Abstract
Pre-trained language models (PLM), for example BERT or RoBERTa, mark the state-of-the- art for natural language understanding task when fino-tuned labeled data. However, their large size poses challenges in deploying them for inference in real-world applications, due to significant GPU memory requirements and high inference latency. This paper explores neural architecture search (NAS) for structural pruming to find sub-parts of the fine-tuned network that optimally trade-off efficiency, for example terms of model size or latency, and generalization performance. We also show how we can utilize more recently developed two- stage weight-sharing NAS approaches in this sctting to accelerate the scarch process. Unlike traditional pruning methods with fixed thresholds, we propose to adopt multi-objective approach that identifies the Pareto optimal set of sub-networks, allowing for more flexible and automated compression process.
1 Introduction
Pre-trained language models (PLMs) such as BERT DDevlin et al. 2019 widely used for natural language understanding (NLU)) tasks when large amount of labelled data is available for fine-tuning. However, deploying PLMs for inference can be challenging due to their large parameter count. They demand significant GPU memory and exhibit high inference latency, making them impractical for many real-world applications, for example when used in an end-point for web service or deployed on an embedded systems. Recent work (Blalock et al.) 2020 Kwon et al. 2022: Michel al.] 2019 Sajjad et al., 2022) demonstrated that in many cases only subset of the pre-trained model's parameters significantly contributes to the downstream task performance. This allows for compressing the model by pruning parts of the network while minimizing performance deterioration.
RoBERTa (Liu et al. 2019bb are
Unstructured pruning (Blllock et. al. 2020) computes a score for each weight in the network, such as the weight's 's magnitude, and removes weights with scores below a predetermined threshold. This approach of ten achieves high pruning rates with minimal performance degradation, but also leads sparse weight matrices, which are not well-supported by commonly used machine learning frameworks. Structured prun-
ing (Michel al. 2019: Sajjad et al. 2022 removes larger components of the network, such as layers

---

Published in Transactions on Machine Learning Research (09/2024)
FER Mark
MFA Taye
IT5 Line
validation error
MEA Agg
LIBA 3ince
number of parameters
(a) Extracting sub-networks from pre-trained network.
(b) Pareto front of sub-networks.
Figure 1: Illustration of our approach. a) We fine-tune the pre-trained architecture by updating only sub- networks, which we select. by placing a binary mask over heads and units in each MHA and FFN layer. b) Afterwards, we run a multi-objective search to select the optimal set of set-rksworks that balance parameter count and validation error.
heads. Although it lypically does not achieve the same pruning rates as unstructured pruning, it only prunes entire columns/rows the weight matrix, making it compatible with popular deep learning frameworks and hardware.
Neural architecture search Zooph & Le 2017 Real et al.
2017
Bergstra et al.
2013) (NAS) finds more
resource efficient neural network architectures in data-driven way by casting it optimization problem. To reduce the computational burden of vanilla NAS, which needs to train and validate multiple architectures.
weight-sharing--esse baaed neural architecture search (Pham et al. 2018 Liu et. al. 2019b Elsken et al. 2018) first trains a single large network, called the super-nebwork, and then searches for sub-networks within the super- neterr.etwork.
We propose to use NAS for structural pruning of pre-trained networks, to find sub-networks that sustain performance of the pre-trained network after fine-tuning (see Figure 1 for illustration). Most structural pruning approaches prune the networks based predefined threshold on the pruning ratio. In scenarios where there is no strict constraint on model size, it can be challenging to define such a fixed threshold in advance. NAS offers a distinet advantage over other pruning strategies by enabling a multi-objeclive approach to identify the Pareto optimal set of sub-networks, which captures the nonlinear relationship between model size and performance instead just obtaining a single solution. This allows us to automale the compression process and to select the best model that meets our requirements post-hoc after observing the non-linear Pareto front, instead of running the pruning process multiple rounds to find the right threshold parameter.
While there is considerable literature on improving the elficiency of PLMs, to the best of our knowledge there is no work yet that explored the potential of NAS for pruning fine-tuned PLMs. Our contributions are the following:
We discuss the intricate relationship between NAS and structural pruning and present NAS ap- proach that compresses PLMs for inference after fine-tuning downstream tasks, while minimizing performance deterioration. Our focus lies not in proposing a novel NAS method per but rather in offering a practical use-case for NAS in the context of PLM that works competitively to structural pruning methods from the literature.
We propose four different search spaces with varying complexity to prune components of transformer based PLM and discuss how their definition affect the structure of sub-networks. Two of these search
spaces are typically used by existing structural pruning approaches (see Section 4.2). While one of these commonly used scarch spaces exhibits the highest degree of freedom, we show in Section [.1.1 that search space with lower complexity can be more efficient to explore and eventually lead better performance within reasonable budget.
2

---

Published in Transactions on Machine Learning Research (09/2024)
We contribute benchmarking suite for multi-objective NAS. We also show how can apply recently proposed weight-sharing based NAS methods in this setting. Based our benchmarking suite, we perform a thorough ablation study of standard and weight-sharing based NAS. In the long run we anticipate that our work will drive the development of future NAS methods.
We present an overview of related work in Section 2 and describe our methodology in Section 3 Section 4 provides an empirical comparison of our proposed approach with other structural pruming methods from the literature, along with an in-depth ablation study. Code is available at httpss//github.com/ /whittle- org/ plm pruning..n
2 Related Work
Neural Architecture Search (NAS) (see Elsken et al. (2018) for an overview) automates the design of neural network architectures to maximize generalization performance and efficiency (e.g,, in terms of latency, model size or memory consumption). The limiting factor of conventional NAS is the computational burden of the search, which requires multiple rounds of training and validating neural network architectures (Zoph & Le 2017 Real et al., 2017). To mitigate this cost, various approaches have been proposed to accelerate the search process. For example, some of these methods early terminate the training process for poorly performing configurations (Li et al., 2018) or extrapolating learning curves (White et al.) 2021b). Weight-sharing NAS (Pham et al. 2018: Liu et al. 2019a) addresses the cost issue by training a single super-network consisting of all archilectures in the search space, such that each path represent a unique architecture. Initially, Liu et al. (2019a) framed this as a bi-level optimization problem, where the inner objective involves the optimization of the network weights, and the outer objective the selection of the architecture. After training the super-network, the best architecture is selected based on the shared weights and then re-trained from scratch. However, several papers (Li & Talwalkar 2020 Yang et al., 2020) reported that this formulation heavily relies on the search space and does not yield better results than just randomly sampling architectures. To address this limitation, Yu et al. (2020) proposed a two-stage NAS process. In the first stage, the super-nel work work is trained by updating individual sub-networks in each iteration, instead of updating the entire super-network. After training, the final model is selected by performing gradient-free optimization based on the shared weights of the super-network, without any further training. Concurrently, Cai et al. (2020) applies a similar approach for convolutional neural networks in the multi-objectivee setting by first training single super-network and then searching for sub-networks to minimize latency on some target devices. Related to work is also the work by Xu et al. (2021). which searchos for more efficient BERT archilectures during the pre-training phase.
Structural Pruning involves removing parts of a trained neural network, such as heads (Michel et al. 2019) or entire layers (Sajjad et al. 2022). to reduce the overall number of parameters while preserving performance. Individual components are pruned based on a specific scoring function, using a manually defined threshold. For transformer-he based architectures, Michel al. (2019) observed that a significant number of heads, up to single head in a multi-head attention layer, can be deleted after fine-tuning without causing a significant loss in performance. Voita et al. (2019) proposed LO regularization
means prune individual heads in multi-head attention layer. [won et al. (2022) prunes individual heads and units in the fully-connected layers after fine-tuning according to the Fisher information matrix. Saijad et al. (2022) demonstrated that it is even possible to remove entire layers of a pre-trained network prior to fine-tuning, with minimal impact on performance. In comparison to our data-driven approach, [Sajaad et al. (2022) suggested using predefined heuristics (c.g., delcting top / odd / even layers) to determine layers to prunc. However, as shown in our experiments, the appropriate architecture depends on the specific task, and more data-driven methods are necessary to accurately identify the best layers to prune.
Distillation (Hinton et al. 2015) trains a smaller student model to mimic the predictions of pre-trained teacher model. For instance, Sanh el al. (2020) used this approach to distill a pre-trained BERT model (De- vlin et al.] 2019) into a smaller model for fine-tuning. Jiao et al. (2019) proposed a knowledge distillation approach specifically for transformer-based models, which first distills from pro-trained teacher into a smaller model and then performs task-specific distillation in a second step based on a task augmented dataset. Related to our method is also AdaBERT (Chen et. al. 2020 which trains task-specific convolu-
3

---

Published in Transactions on Machine Learning Research (09/2024)
tional neural networks based on differentiable NAS such as BERT.
Liu et al. 2019a by distilling the knowledge of a PML
Unlike pruning-based methods, distillation allows for complete architectural changes beyond merely drop- ping individual components. However, from a practical standpoint, determining the optimal structure and capacity of the student, network needed match the performance of the teacher network also amounts to a hyperparameter and neural architecture search problem. Additionally, training a student network requires significant amount of computational resources. For example, the model by |anh et al. (2020) was trained for around 90 hours on 8 16GB V100 GPUs. This cost be amortized by fine-tuning the student model to solve many different tasks, but, depending on the downstream tasks, it potentially requires a substantial amount of iterations which is not always desirable for practitioners who aim to solve a single specific task. This is especially important in the multi-objective setting where many networks need to be distilled to cover the full size/aceuracy Pareto front.
Quantization (Dettmers et al. 2022 Dettmers & Zettlemoyer 2023 reduces the precision of model pa- rameters from footingg-poin numbers to lower bit representations (e.g., 8-bit integers). The main advantage of quantization is the reduction in memory footprint. However, necessarily lead to faster latency. Quantization is independent of our NAS approach and can be employed on the pruned network to further decrease memory usage.
show in the Appendix F this does not
3 Structural Pruning via Neural Architecture Search
We first provide multi-objective problem definition for structural pruming of fine-tuned PLMs via neural architecture search. Afterwards, we describe how we can apply weight-sharing based NAS. At the end, we present four search spaces to prune transformer-based architectures, which exhibit different degree of pruning.
3.1 Multi-Objective Sub- Network Selection
We consider pre-trained transformer model based encoder-only architecture, such for example BERT VVaswani et al. 2017), with L non-embedding layers, each composed of multi-head attention (MHA) layer followed by a fully connected feed forward (FFN) layer. However, all methods presented here can also be applied to decoder or encoder-decoder based architectures. Given an input sequence E Ruxdmaded where n represents the sequence length and dimodet the size of the token embedding, the MHA layer defined by: MHA(X) = Σ. Att(w (t) W TRe Hddddmode! are weight matrices. ALt(.) is a dot product attention head (Bahdanau et al. 2015) and H is the number of heads. The output is then computed by XMHA LN + MHA(X)) )), where LN denotes layer normalization (Ba et al. 2016). The FFN layer is defined by FFN(X) W1a with Wo E RUddmodd and W1 E R Rdmodal XU where U denotes the intermediate size and a()) is a non-linear activation function. Also here we use a residual connection to compute the final output: #PEN = LN(X + FF)(
Q K w/(") w X) where W ((()
K ((()) E IR Rddmodetxd and w E
We define a binary mask Mirad E [0,1) 1x If for each head in the multi-heat attention layer and a binary 10,1) for each neuron in the fully-connected layers. The output of the I-th MHA layer and FFN layer is computed by (X = -. Mheadli, Au)-)) and FFM(X) = W1 o Mneuron!!,; (W, respectively, where o denotes element-wise multiplication.
mask Mnsuron
Now, let's define scarch space 0 E ② that contains finite set of configurations to define possible sub- networks sliced from the pre-trained network. We define function CREATEMASK that maps from a configu- ration 0 + Mhead, M to binary masks. Let's denote the function fo : θ - R as the validation error of the sub-network defined by configuration 0 after fine-tuning on some downstream task. To compute the validation score induced by we place corresponding masks Mhead: Mneuron over the network. Additionally, we define the total number of trainable parameter f1 : θ -+ N of the sub-network. Our goal to solve the following multi-objective optimisation problem:
(1)
4

---

Published in Transactions on Machine Learning Research (09/2024)
In the multi-objective setting, there is no single 0. E θ that simultaneously optimizes all M objectives. Let's define 0 > O' iff J:(0) S Ii("",, E [M] and 7i E [k] [i(0) < fi(0"). We aim to find the Parelo Set: Py =0 01 E e 0' > 0} of points that dominate all other points in the search space in at least one objective.
To solve this optimization problem, we can utilize standard multi-objective search methods that are com- monly used for NAS, such as random search. Here each function evaluation consists of fine-tuning a sub- network 0 initialized with the pre-trained weights instead of random weights. We can also directly adopt more advanced strategies, such multi-fidelity NAS, for example MO-ASHA (Schmucker et al. 2021 to accelcrate the scarch process.
3.2 Weight-sharing based Neural Architecture Search
Following previous work (Yu al. 2020 Wang et al. 2021 our weight-sharing based NAS approach consists of two stages: the first stage is to treat the pre-trained model as super-netword and fine-tune it on the downstream task. We explore different super-network training strategies from the literature that update only parts of the network in each step, to avoid co-adaption of sub-networks. The second stage, utilizes multi-objective search strategies to approximate the Pareto-optimal set of sub-networks.
3.2.1 Super-Network Training
In the standard NAS setting, we would evaluate fo(0) by first fine-tuning the sub-network defined by 0 on the training data before evaluating on the validation data. The weights of the sub-nctwork are initialized based on the pre-trained weights. While more recent. NAS approaches (Li & Talwalkar, 2020 Klein et al. 2020) accelerate the search process by early stopping poorly performing sub-networks, this still amounts an optimization process that requires the compute of multiple independent fine-tuning runs.
The idea of two-stage weight-sharing-Be based NAS Yu et al., 2020 is to train single-set of shared weights, dubbed super-network, that contains all possible networks in the search space. After training the super- networks, evaluating fo) only requires a single pass over the validation data.
We consider the pre-trained network as super-network network with shared weights that contains all possible sub- networks E ④. To avoid that sub-networks co-adapt and still work outside the super-network, previous work (Yu et al. 2020 Wang al. 2021) suggested to update only a subset of sub-networks in each stochastic gradient descent step, instead of updating all weights jointly. We adapt this strategy and sample sub-networks according to the sandwich rule (Yu et al. 2020 Wang et al. 2002]) in each update step, which always updates the smallest, the largest and k random sub-networks. The smallest and largest sub-network correspond to the lower and upper bound of θ, respectively. For all search spaces θ define below, the upper bound is equal to full network architecture, i.e, the super-network and the lower bound removes all layer except the embedding and classification layer.
Additionally, we use in-place knowledge distillation Yu et. al. 2019) which accelerate the training process of sub-networks. Given the logits supernet (z: of the super-network, which obtain for free with the sandwich rule, and the logits of a sub-network To( the loss function to obtain gradients for the sub-networks follows the idea of knowledge distillation:
LKD = LCE + DKL (T
(Tsupernet T
o
T (T))
(2)
where (.I(.) denotes the Lulbback-Leibler divergence between the logits of the super-network and the sub- network, T temperature parameter, (-) the softmax function and Lce is the cross-entropy loss of the training data.
3.2.2 Sub-network selection
After training the super-network, we compute the validation error fo( by applying Mhead and Mneuron to the shared weights and performing a single pass over the validation data. This substantially reduces the computational cost involved in the multi-objective problem stated in Equation 1 To solve this problem,
5

---

Published in Transactions on Machine Learning Research (09/2024)
essentially use the same multi-objective approaches as for standard NAS (see Appendix D except for multi-fidelity approaches such as MO-ASHA. Multi-fidelity methods early stop the training process by evaluating the model at intermediate time steps, called rung levels. Since we do not perform any additional training steps anymore and just evaluate the final sub-network, these rung levels are not defined.
3.3 Search Space
The search space @ defines sub-networks of the pre-trained network architecture. An expressive θ allows for fine-grained pruning but might also become infeasible to explore. We propose the following search spaces that exhibit different levels of complexity. For each search space we provide pseudo code to define the CREATEMASK function in Appendix B
We define the number of heads H = [0, H]. the number of units U = [0, U] and the total number of layers L = [0, L], such that θ = H x 2 C. Compared to the other search spaces, the dimensionality of this search space remains fixed with different model sizes, and only its upper bound i.e. (H,U, U, L) increases. For each layer, we always keep the first h H heads and u E U units, respectively, to enforce that CREATEMASK is bijective mapping (see Appendix B
SMALL:
LAYER: Inspired by Sajjad et al. 20222 we prune individual attention and fully-connected layers instead of single heads and neurons. We define a search space e = {00 1). that contains one binary hyperparameter for each layer that determines if the corresponding layer is removed.
MEDIUM:
Based on the previous search space, we allow for a flexible mumber of heads / units per layer. For each layer I E [0, L], we define Hi [0, H] and Ui [0, U7, such that the final scarch space is θ = Ho x Uo Hr. x Ur. As for the SMALL search space we also keep the first heads and units in each layer.
For each head and neuron in the fully-connected layer we define a single binary θ = 10 which is combined to form the search space θ Oo x ... x + L(H+1)) This is the most expressive search space, but also grows quickly with the model size. The search space is also commonly used by other structural pruning approaches (Kwon et al., 2022). It. might not be very useful in practice, because cannot easily remove single tows/columns of the weight matrix with most transformer implementations and hence it will not necessarily reduce the inference latency. However, it provides us a reference in terms of predictive performances that can be retained under certain pruming ratio.
LARGE:
Each search space induces a different pattern for Miead and Mneuron that we place over the super-network to select sub-networks (sce Figure 2 for some examples). To see how this effects the distribution over param- eter count and hence the sampling during the super-network training, we sample N = 500 configurations {oo.... ..., O uniformly random and compute the number of trainable parameters (f. fi (0x) for all four scarch spaces (sce Figure 3). The SMALL scarch space is somewhat biased to smaller networks. The MEDIUM search space, even though more expressive, is highly biased towards mid-size networks, since average half of the heads neurons are masked out. For the two binary search spaces LAYER and T.ARGE, we can achieve uniform distribution over the number of parameters, by using the following sampling pro- cess. We first sample an integer ~ (0, K), where & = L for the LAYER search space, and k = L(H + I) for the L. search space. Afterwards, we randomly select entries of the binary vector O E and set them 1.
4 Experiments
We evaluate different lypes of NAS for structural pruning on eight lext classification tasks, including textual entailment, sentiment analysis and multiple-choice question / anssering. We provide a detailed description of cach task Appendix [ All tasks come with a predefined training and evaluation set with labels and hold-out test set without labels. We split the training set into a training and validation set (70%/30% split) and use the evaluation set test set. We fine-tune every network, sub-network or super-network, for

---

Published in Transactions on Machine Learning Research (09/2024)
SMALL
MEDIUM
LARGE
layer
layer
oyer
head
head
head
head
Figure 2: Examples of head masks Mhead sampled uniformly at random from different search spaces. Dark color indicates that the corresponding head is masked. The same pattern can be observed for M Mnearon
SMALL
MEDIUM
LAYER
332
relative number parameters
relative number of parameters
relative number parameters
relattive number parameters
Figure 3: Distribution of the parameter count (0) for uniformly sampled O ~v O
epochs on single GPU. For all multi-objective scarch methods, we use Syne Tune (Salinas et al. 2022 a single GPU instance. We use BERT-base (Devlin et al. 2019) (cased) and RoBERTa-bass (Liu et al. 2019b) pre-trained network, which consists of L = 12 layers, I 3072 units and H 12 heads (other hyperparaneters are described in Appendix A.. While arguably rather small for today's standards, they still achieve competitive performance on these benchmarks and allow for more thorough evaluation. We also present a comparison to quantization in Appendix
4.1 Benchmarking Neural Architecture Search
We now present an evaluation of different multi-objective NAS approaches on our benchmarking suite. To quantify the performance of Pareto set Pr, we compute for each Pareto set the Hypervolume (Zitzler al. 2003) and report the regret, i the difference to the best possible Hypervol- ume averaged across all repetitions. Given a reference point E M, the Hypervolume HV(p)rr) A(UyEP, [y, rJ) is de- fined as the M-th dimensional Lebesgue measure between the Pareto set Py and the reference point r, where [y. r] rep- resents the hyper-roctangle between y and r (sce Figure for an example).
1/0
abjoctiove
500
objective fo
Figure 4: Example to compute the Hypervol- ume HV(P,(r), corresponding to the sum of the rectangles, across a reference point r and a sot of points Py = (yo. (yo.Y1. V2.
To compute the Hypervolume, first normalize each ob- jective based on all observed values across all methods and repetitions via Quantile normalization. This results in uni- form distribution between 1], and we use r (2 as reference point, which means the highest possible Hypervol- ume would be We evaluate each method with 10 different seeds for the random number generation.

---

Published in Transactions on Machine Learning Research (09/2024)
4.1.1 Search Space
First, we compare the search spaces definitions from Section 3.3 using weight-sharing based NAS. We fine- tune the super-network as described in Section 3.2 and sample 100 sub-networks uniformly at random to compute the hypervolume.
Conclusions: Within this budget (see Figure 5 the SMALL search space achieves the best performance across all datasets, except for COLA. Interestingly, even though the MEDIUM search space allows for a more fine-grained per layer pruning, leads to worse results. We altribule this to the non-uniform distribution of parameter count as described in Section 3.3. The LAYER search space often out-performs the MEDIUM and LARGE scarch space, but, except for COLA, leads to Pareto sets that under-perform compared to the SMALL search space. The LARGE search space, which is a superset of the other search spaces, seems infeasible to explore with random sampling over so few observations. We use the SMALL search space for the remaining experiments.
4.1.2 Standard Neural Architecture Search
We compare the following multi-objective search methods to tackle the NAS problem described in Scction 3 where each sub-network is fine-tuned in isolation. We provide a more detailed description of each method in Appendix D A simple multi-objectiive local search (LS) described in Appendix D This is inspired by the work of White al. (2021a), which showed that local search often performs competitively on NAS problems. Random search (RS) (Bergstra & Bengio 2012) samples architectures uniformly at random from the search space. A multi-objective version of the regularized evolulion (REA) algorithm (Real el al. 2019 frequently used in the NAS literature. Compared to the original singe-objective algorithm, we sort elements in the population via non-dominated sorting. Expected Hypervolume Improvement (EHVI) (Daulton et al. 2020 is a multi-objective Bayesian optimization strategy that samples candidate points using a Gaussian process model of the objective function. Lastly, we include MO-ASHA (Schnucker al.] 2021 multi- objective version of asynchronous successive halwing (Li & Talwalkar 2020; Jamieson & Talwalkar 2016) that terminates the training process of poorly performing candidates early to accelerate the overall optimization process. While MO-ASHA could potentially be combined with a model-based approach, commonly done for singlo-objeetive optimization (Falkner et al.} 2018 Klein et al., 2020). here we followed the original algorithrn and sample candidate uniformly at random from the search space.
Following common experimental practice from the HPO literature, we aggregate results by computing the average ranks of cach methods across repetitions, datasets and time steps. Following Feurer et al. (2015). we sample 1000 bootstrap samples across all repetitions and lasks, compute the rank of each method and average across all samples. Results are shown in Figure 6a. We shows results for each individual task in Appendix D
Somewhat surprisingly RS is a strong baseline on these benchmarks, outperforming more sophisticated approaches such EHVI or MO-REA. Fine-tuning these models is often unstable (Mosbach) eet al. 2021) especially on smaller datasets, resulting in high observation noise. For the RoBERTaa-ase model, LS often performs competitively to RS given a sufficient large budget. MO-ASHA quickly stops the evaluation of poorly performing sub-networks and hence outperforms RS on average. However, on small datasets such as RTE, fine-tuning is faster than the non- con-domintted sorting of MO-ASHA, such that converges slower than RS (see Appendix E
Conclusions:
4.1.3 Weight-sharing based Neural Architecture Search
Next, we evaluate different techniques for two-stage NAS for this problem. We distinguish between, fine- tuning the super-network and multi-objective scarch.
Super-network fino-tuning: First, we compare the following strategies to fino-tune the super-network from the literature. To compare these methods, after fine-tuning the super-network, we sample 100 sub- networks uniformly at random from the SMALL search space to estimate the Pareto set and report here the
8

---

Published in Transactions on Machine Learning Research (09/2024)
KTB
STSR
Seperodlenn
Viikem
1dy
Villem
Tary
Mi
Tayy
Tart
Madees
Taig
spar
sporn
spar
1KDB
QNLI
Apperolunee
Stedie
Hedin
Seal
Miabes woonch squee
teal
sasrch ap00n
March spar
(a) BERT-base-com
XBRO
STAD
C0CA
Iogeroolaro
Egps
Saal
10/t
Viifee Seal seanch qpace
Lxyr
Hillen nearch spar
Lent
Mickas
mnach spare
wwech spare
IMDB
8812
SWAG
QNL
typersatnee
begeroolene
soardh spare
sestcc SUCAE
seurc space
scnuch
(b) RoBERTa-base
Figure 5: Comparison of the four different search spaces using weight-sharing based NAS. We sample 100 random sub-networks uniformly at random using the fine-tuned weights of the super-network. The SMALL search space dominates the other search spaces except for the COLA dataset. While SMALL is a subset of MEDIUM and LARGE, these spaces are too high-dimensional to explored with a sensible compute budget. First two rows show results for BERT-base-caser and last two rows for RoBERTa-base.
Hypervolume. We repeat this process 10 times with a different seed for the random number generation. For each repetition we use the exact same set of random sub-networks for all super-network network training strategies.
standard: Which trains all weights of super-network in the standard fine- Ano-tuning setting
random: Samples single random sub-network each update steps
random-linear: Inspired by Bender et al. 2018 we either sample a random sub-network with probability or the full-network with probability of 1-p cach update step. Thereby, p lincarly increased from to after each update step over the course training.
sandwich: The super-network updated according the sandwich rule (Yu et al. 2020 Wang
ct al. 2021) described in Section 3.2 We set the number of random sub-networks each update
step h: = 2.
9

---

Published in Transactions on Machine Learning Research (09/2024)
KOBOYRANCE
time steps
time steps
line slopp
Gine staps
(a) Average ranks standard-NAS
(b) Average ranks weight-sharing based NAS
Figure 6: Average ranks of the different multi-objective search methods across all dataset for a) standard-NAS and b) weights-sharing based NAS for both BERT-based cased and RoBERTa- base models. Random search based approaches (RS, (O0-ASHA)eerfor
competitively the standard NAS setting. EHVI outperforms other methods earlier time steps which we attribute to its probabilistic model. Given a sufficient amount of time often finds well-performing Pareto fronts.often finds well-perofforming Pareto fronts.
kdl: Update k = random sub-networks using in-place knowledge distillation
Yu et al. 2020
Wang et al. 2021 according to Equation 2
full: Implements the training protocol described in Section 3.2 i. it combines the sandwich rule
with in-place knowledge distillation to update sub-networks.
Conclusions: Figure shows the Hypervohume across all task for BERT-base and RoBERTa-base base, re- spectively. Standard fine-tuning and just randomly sampling sub-network leads to poorly performing Pareto sets compared to other methods. The only exception is the COLA dataset, where standard fine-tuning some- times works best. However, also observe high variation across the probability of sampling random sub-networks improves to just random sampling. Better results are achieved by using the sandwich rule or knowledge distillation. Thereby, combining both slightly improves results further.
this datasets. Linearly increasing
Multi-Objective Search Lastly, we compare in Figure search methods as for standard-NAS. We follow the same process as described Section 4.1.2 compute averange ranks. We do not include MO-ASHA in this setting, since each function evaluation only consists of validating sub-network based on the shared weights of the super-network after finc-tuning and hence does not allow for multi-fidelity approach. Optimization trajectories for all dalasets are in Appendix E.
6b average ranks of the same multi-objective
Conclusions: As for standard NAS, RS is a surprisingly strong baseline. EHVI performs better at early workss
lime-sleps. We found using the shared weights of the super-nel works for evaluation results in a much smaller observation noise than fine-tuning each sub-network in isolation, which less deceiving for the probabilistic model of EHVI. Given enough time, LS starts outperforming RS and EHVI on RoBERTa-base model and performs competitively to EHVI on the BERT-base model.
4.2 Comparison to other Structural Pruning Approaches
We now present a comparison against other structural pruming approaches. For NAS we use the SMALL search space defined in Section 3.3 based ablation study in Section [41] Fach method had the same total amount of wall-clock time and compute, which include both finc-tuning and scarch. We compare the following methods:
Head-Pruning (HP) Michel et al. 2019) prunes heads greedily using proxy score for the impor- tance of each head for final performance based the gradient of the loss function.
Retraining Free Pruning (RFP) Kwon et al. 2022) uses a three-phased pruning strategy that, based threshold a, prunes individual heads in the MHA layer and units in the FFN layer.
10

---

Published in Transactions on Machine Learning Research (09/2024)
7TR
CDLA
Sepprrvtiinn
Togger Togeesolome
bag arpe-ratoedd wonnin
espronntonr
aupertatook
hop hoperootine womemne
Inypersliume 2
maatie masa Bay mussen apacod
Maadicd makio leae mudioo
esperseeook toining
yoataa mudic
agpe-sstordi strengy
netiica upliceks
dasaed rasie
zipenstork tning siralagy
reorso Iglandd
traningg
(a) BERT BERT-baaeed base-cased
2T1
MBPC
STSD
COLA
Inpernaluan
Experoalan
eppeasionk Iraning
sepreaattnek training
IM38
8T1
SWAG
QNL
Iepersdlane
madeen: Bone mochio
mphieks
mesees repaiaa
(b) RoBERTa-buse
Figure 7: Comparison of different strategies to fine-tune the super-network. First two rows show BERT-bast and last two rows show RoBERTa-bee base.
The first phase computes a binary mask for heads and units by computing the diagonal Fisher information matrix. The matrix is then rearranged by a block-approximatel Fisher information matrix. In the last step, the masked is further tuned by minimizing the layer-wise reconstruction error. This method operates in the LARGE search space described in Section 3.3 We run RFP with different values for E (0.1, 0.2.0........03.. to obtain Pareto set of architectures.
Layer Dropping (LD): Following Sajjad et al. (2022) we first remove the top n E 1, L I layers and fino-tune the remaining layers directly on the downstream task. To obtain Parcto set of N points, we fine-tune N models with different amount. of layers removed. This method serves simple heuristic to explore the LAYER search space.
Standard NAS (S-NAS): for the previous experiments, we used standard NAS using random search where each sub-network is initialized with the pre-trained weights and then fino-tuned inde- pendently.
11

---

Published in Transactions on Machine Learning Research (09/2024)
COLA
010
1PN
pruned
number paramatara
A2N
purarmterrrs pruned
SWAN
ONL
number of corammees prined
sumber paarmmttees pured
number paramrtees prined
number ofparmeters prined
COLA
C4B1
Clat
TUNDET Carmmnetes pruned
sumbbrr parameters pruved
TUMDER of peprrmeeer punee
nuMDEr of caramaters praned
MD3
SST2
3/50
2/%
Cn A/ 25h 22 22% 72%
20% 755
Figure 8: Loss in test performance (the higher the better) versus the parameter count relative to the un- pruned model on all 8 text classification datasets. First two rows show BERT-base and last two rows show RoBERTa-base.
Weight-sharing NAS (WS-NAS): Follows the two-stage weight sharing based NAS approach out- lined in Section 3.2 We use sandwich rule and in-place knowledge distillation to train the super- network and EHVI to search for the Pareto optimal set of networs......
To compare results, we normalize the number of parameters to [0.1] and bin results based on different thresholds / E (0.2)...0.9). Note that roughly 20% of the parameters of BERT-base / RoBERTa-base are included in the embedding and classification head, and hence cannot be trivially pruned without changing the embedding dimension or the number of classes. For each bin, we report the best performance of the solution with < parameters. We discuss the relationship between parameter count and model inference time in
Figure & shows the parameter count (horizontal axis) and the test error (vertical axis) relative to the unpruned network for all datascts. For reference, we indicate 95% and 90% relative performance to the unpruned network as well as the original performance by dashed lines. We sort dataset by their size from left to right.
First, both S-NAS and WS-NAS achieve competitive performance to structural pruning meth- ods. Especially for higher praning ratios, both S-NAS and WS-NAS outperform RFP and HP even though they operate in the LARGE search space, which much higher dimensional than the SMALL search space used for NAS. This indicates that these methods cannot rigorously handle such high dimensional spaces.
Conclusion:
12

---

Published in Transactions on Machine Learning Research (09/2024)
Second, NAS methods seems to perform better for larger datasets. Given sufficient amount of budget, simple L.D, which operates in the LAYER search space, also achieves competitive results. This is in-line with our results in Section 4.1.1 showing that the LAYER search space can still provide sensible results compared to more expressive search spaces.
WS-NAS substantially reduces the overall runtime by fine-tuning only a single super-network, which enables it achieve much better performance with limited amount of compute compared S-NAS, However, we expect S-NAS to outperform WS-NAS if we scale the total compute, since each sub-network if s fine-tuned in isolation and hence able to optimally adapt to the training data.
SST2
ONEI
000
erpraned nolivrt
0.000
O000
anotured newer.
1V% NGS
S-MG
1.4
5-M66
64-
LO
34-
tost arror
028
00066000000
06 07
6a
08
07
66
09
parameters
parameters
parameters
Figure 9: Pareto fronts of single runs for each method for SST2, IMDB and QNLI dataset to balance paramcter-count and test-crror (the lower the better).
For qualitative comparison, we show the results for single run on three datasets in Figure 9 example runs are shown in Appendix
right. More
4.3 Conclusions
We propose NAS for structural pruning of fine-tuned PLMs. By utilising multi-objective approach, find the Pareto optimal set of sub-networks that balance between model size and validation error. Returning a Pareto set of sub-networks allows practitioners to select the optimal network without running the pruning process multiple times with different thresholds. We also provide an in-depth analysis of recently developed two-stage weight-sharing approaches this setting, which require only a single fine-tuning run of the PLM.
Future work could explore the instruction tuning (Wei et al. 2022) setting, where the final model is evaluated in few-shot setting. Our approach samples sub-networks uniformly at random, which allocates the same amount of update steps to all sub-networks on average. Future work could explore more complex sampling distribution biased towards sub-nctworks closer to the Pareto set.
References
J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXto:1607.06.. [stat.ML), 2016.
D. Bahdanau, K. Cho, and Y. Bengio. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations (ICLR'15). 2015.
G. Bender, P. Kindermans, B. Zoph, V. Vasudevan, and Q. Le. Understanding and simplifying one-shot architecture search. In Proceedings of the 35th International Conference on Machine Learning (CCML'18). 2018.
J. Bergstra and Y. Bengio. Random search for Research (JMLR-12) 2012.
optimization. Journal of Machine Learning
J. Bergstra, D. Yamins, and D. Cox. Making a science of model search: Hyperparameter optimization hundreds of dimensions for vision architectures. In Proceedings of the 30th International Conference on Machine Learning (TCML 13), 2013.
13

---

Published in Transactions on Machine Learning Research (09/2024)
D. Blalock, J. J. G. Ortiz, J. Frankle, and J. Guttag. What is the state of neural network pruning? ar arXiv:2003.0000 03033 fcs. L.J, 2020.
Cai, C. Gan, T. Wang, Z. Zhang, and S. Han. Once-for-all: Train one network and specialize for efficient deployment. In International Conference on Learning Representations (ICLR'20) 2020.
D. Chen, Y. Li, M. Qiu, Z. Wang, B. Li, B. Ding, H. Deng, J. Huang, W. Lin, and J. Zhou. Adabert: Task-adaptive bert compression with differentiable neural architecture search. In Proceedings of the 29th International Joint Conference on Artificial Intelligence (IJCAI'20), 2020.
S. Daulton, M. Balandat, and E. Bakshy. Differentiable expected hypervolume improvement for parallel multi-objective bayesian optimization. In Proceedings of the 94th International Conference on Advances in Neural Information Processing Systems (NeuRIPS 2020.
T. Dettmers and L. Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws. arXim:2212.0979) [es.LG], 2023.
T. Dettmers, M. Lewis, Y. Belkada, and L. Zetteemoyer. Llm.int8(): 8-bit matrix multiplication for trans- formers at scale. In Proceedings of the 36th International Conference on Advances in Neural Information Processing Systems (NeuRIPS' 22). 2022.
J. Devlin, M. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapler of the Association for Computational Linguistics: Human Language Technologies, 2019.
T. Flssen, J. H. Metzen, and F. Hutter. Neural architecture search: A survey. orXic:180800597 [sLal.MT). 2018.
S. Falkner, A. Klein, and F. Hutter. BOHB: Robust and efficient hyperparameter optimization at scale, In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), 2018.
M. Feurer, T. Springenberg, and F. Hutter. Initializing bayesian hyperparameter optimization via meta- learning. In Proceedings of the 29th National Conference on Artificial Intelligence (AAAT'15), 2015.
G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in neural network. 1503.0150..0200 MLJ, 2015.
K. Jamicson and A. Talwalkar. Non-stochastic best arm identification and hyperparameter optimization. In Proceedings of the 17th International Conference on Artificial Intelligence and Statistics (AISTATS16), 2016.
X. Jiao, Y. Yin, L. Shang, X. Jiang, X. Chen, L. Li, F. Wang, and Q. Liu. Tinybert: Distilling bert for natural language understanding. arXiv: 10351 CL.. 2019.
A. Klein, L. C. Tiao, T. Lienart, C. Archambeau, and M. Seeger. Model-based asynchronous hyperparameter optimization. arxiv:203. 10865 [cs.LG], 2020.
W. Kwon, S. Kim, M. W. Mahoney, J. Hassoun, K. Keutzer, and A. Gholami. A fast post-training pruning framework for transformers. arXiv:2204.0000 .96556 fes. (cs) 2022.
L. Li and A. Talwalkar. Random search and reprodueibility for neural architecture search. In Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, 2020.
L. Li, K. Jamieson, G. DeSalvo, A. Rostamizadeh, and A. Talwalkar. Hyperband: Bandit-based configuration evaluation for hyperparameter optimization. In International Conference on Learning Representations (7),R'17, 2017.
L. Li, K. Jamicson, A. Rostamizadeh, K. Gonina, M. Hardt, B. Rocht, and A. Talwalkar. Massively parallel hyperparameter tuning. arXiv:1810. 810.05994 [cs. 2018.
14

---

Published in Transactions on Machine Learning Research (09/2024)
H. Liu, K. Simonyan, and Y. Yang. DARTS: Differentiable architecture search. In International Conference on Learning Represenlations (ICLR'19) 2019a.
Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXi:: 1907. 1 11692 [cs. CL,, 2019b.
P. Michel, O. Levy, and G. Neubig. Are sixteen heads really better than onc? In Proceedings of the 32th International Conference on Advances in Neural Information Processing Systems (NeurIPS'19) 2019.
M. Mosbach, M. Andriushchenko, and D. Klakow. On the stability of finc-tuning bert: Misconcoptions, explanations, and strong baselines. In International Conference on Learning Representations (ICLR' 21), 2021.
H. Pham, M. Guan, B. Zoph, Q. Le, and J. Dean. Efficient neural architecture search via parameters sharing. In Proceedings of the 35th International Conference on Machine Learning (ICML (18), 2018.
E. Real, S. Moore, A. Selle, S. Saxena, Y. L. Sucmatsu, J. Tan, Q. V. Le, and A. Kurakin. Largo-scale evolution of image classifiers. In Proceedinggs of the 34th International Conference on Machine Learning (ICML '17), 2017.
E. Real, A. Aggarwal, Y. Huang, and Q V. Le. Regularized Evolution for Image Classifier Architecture Scarch. In Proceedings of the Conference on Artificial Intelligence (AAAI119), 2019.
Sajjad, F. Dalvi, Durrani, and P. Nakov. On the effect of dropping layers of pre-trained transformer models. arXiv:2004.088 fes. CLJ, 2022.
D. Salinas, M. Seeger, A. Klein, V. Perrone, M. Wistuba, and C. Archambeau. Syne tune: A library for large scale hyperparameter tuning and reproducible research. In First Conference on Automated Machine Learning (Main Track), 2022.
V. Sanh, L. Debut, Chaumond, and T. Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv:1910. 0000000.011111111111110011111111 fcs. CL), 2020.
R. Schmnuckrr, M. Donini, M. B. Zafar, D. Salinas, and C. Archambeau. Multi-objective asynchronous successive halving, orXiv:2106. 12639 [stat. 2021.
Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. Gomez, L. Kaiser, and I. Polosukhin, Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS'17), 2017.
E. Voita, D. Talbot, F. Moiseev, R. Semnrich, and 1. Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
D. Wang, M. Li, C. Gong, and V. Chandra. AtlentiveNAS: Improving Neural Architecture Search via Attentive Sampling. arXic:2011.0001 [cs. CV3, 2021.
J. Wei, M. Bosma, V. Zhao, K. Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q, V Le. Finetuned language models are zero-shot learners. In International Conference on Learning Representations (ICLR 222, 2022.
C. White, Nolen, and Y. Savani. Exploring the loss landscape in neural architecture search. In Proceedings of the 37th conference on Uncertointy in Artificial Intelligence (UAI22)) 2021a.
C. White, A. Zela, R. Ru, Y. Liu, and F. Hutter. How powerful are performance predictors in neural archi- tecture search? In Proceedings of the 35th International Conference on Aduances in Neural Information Processing Systems (NeuRIPS'21) 21), 2021b.
J. Xu, X. Tan, R. Luo, K. Song, Li, T. Qin, and T. Liu. NAS-BERT: task-agnostic and adaptive-size bert compression with neural architecture search. In Proceedings of the 27th ACM SIGKDD International Conference Knowledge Discovery & Data Mining (KDD'21) 21), 2021.
15

---

Published in Transactions on Machine Learning Research (09/2024)
A. Yang, P. M. Esperança, and F. M. Carlucci. NAS valuation is frustratingly hard. In International Conference Learning Representations (ICLR 20), 2020.
J. Yu, L. Yang, N. Xu, J. Yang, and T. Huang. Slimmable neural networks. In International Conference Learning Representations (ICLR'19), 2019.
J. Yu, P. Jin, H. Liu, G. Bender, J. Kindermans, M. Tan, T. Huang, X. Song, R. Pang, and Q Le. BigNAS: Scaling Up Neural Architecture Search with Big Single-Stage Models. In The European Conference Computer Vision (ECCV'20), 2020.
E. Zitzler, L. Thicle, M. Laumanns, C. M. Fonseca, and V. G. Da Fonseca. Performance assessment of multiobjective optimizers: An analysis and review. IEEE Transactions evolutionary computation, 2003.
B. Zoph and Q V. Le. Neural architecture search with reinforcement learning. In International Conference on Learning Representations (ICLR'17) 2017.
A Hyperparameters
Table shows the hyperparameters for fine-tuning the super-network. We largely follow default hyperpa- rameters recommended by the HuggingFace transformers library For all multi-objective search method, we follow the default hyperparameter of (Syne Tune
Hyperparameter
Value
Learning Rate
0.00002
Number of random sub-networks
Temperature
Batch Size
B
Masking
Algorithm and 4 show pscudo code for the LAYER, SMALL, MEDIUM and LARGE search space, respectively. Note that, 1 indicates a vector of ones. For matrix M, we write M[;;; N] to denote the first. N columns for all rows and, vice versa, M :] for the first N rows.
input sub-network configuration E 1
output: Mheadd M Mneuron
Mhead f [0]L*H,
Maeuron + Lx],
for 1= =0,.... O,...,
do
M Mhead(,; - A1];
Maearon |1 0[l];
end
A Algorithm 1: CREATEMASK function for LAYER search space
16

---

Published in Transactions on Machine Learning Research (09/2024)
input sub-network configuration 0 No .....H U1
output: Mhead, Mecroo
Mhead f [0]Lx8;
Maeuron [0]Lx]
for / == 0,..., do
-0
number of heads in layer */ number of units in layer */
u = 0
M Mhadd [1; hi 1;
M Mnearon! [L, u| 1;
end
Algorithm 2: CREATEMASK function for MEDIUM search space
input- sub-network configuration E H x x C
output: Mhead, Mneuron
h 0
/+ number of heads */
u 0[1
/+ number of units */
- 0[2];
number of layers */
Mhead
[0]L*H:
Mneuron LXT [x]
Mhead|: 1,: h1 1 1;
Mneuron |: 1;: u] + 1;
Algorithm 3: CREATEMASK function for SMALL search space
input sub-nnttoork configuration A (0, 1).1 ()
output: Mhead, Mnearon
Mhend &- [ H]:
Mneuron + 0 H :]:
Algorithm 4: CREATEMASK function for LARGE search space
C
Datasets
We use the following dataset test classification datasots. All dataset are classification tasks, except for STSB, which is regression dataset.
The Recognizing Textual Entailment (RTE) dataset aims to identify the textual entailment of two sentences.
The Microsoft Research Paraphrase Corpus (MRPC) dataset consists of sentence pairs extracted from online news sources. The task is to predicts these pairs are semantically equivalent to cach other.
The Semantic Textual Similarity Benchmark (STSB) consists sentences pairs that are scored between and based on their similarity.
The Corpus of Linguistics Acceptability (COLA) dataset contains English sentences that are labeled as grammatically correct or not.
The TMDB dataset for sentiment classification (positive negative) of movie reviews.
The Stanford Sentiment Treebank (SST2) datasets classifies the positive / negative sentiment of sentences extracted from movie reviews.
Situations With Adversarial Generations (SWAG) dataset for multiplo-choice question answering.
QNLI is a modified version of the Stanford Question Answering Datasct which collection of
question / answer pairs where question are written by human annotators and answers are extracted from Wikipedia. The task is to predict whether the answers is correct.
17

---

Published in Transactions on Machine Learning Research (09/2024)
D Additional Details NAS Search
In this section, we present more details about all multi-objective search strategies that we used for both the standard NAS scenario, where each sub-network is fine-tuned from the pre-trained weights, and the weight-sharing scenario, where sub-networks are evaluated using the shared weights of the super- network.
Random Scarch: We follow the standard random search approach (Bergstra & Bengio 2012
where for a fixed number of iterations T, we sample in each iteration / E {o,..., T} a random sub-networks uniformly from the search space 0, ~ u()).
MO-Regularized Evolution: We adapt the original algorithm proposed by (Real et al. 2019) which maintains a population of architectures and, in each iteration, removes the oldest architecture from the population. To sample a new configuration, we follow Real et al. (2019) and first sample a set of random configurations and mutate the configuration from the set with the lowest rank. Instead of just using the validation performance, we rank each element in the population via non-dominated sorting. We mutate architecture 0 by first sampling random dimension of the search space d ~ (0, (0) and than sample new value for this dimension O[d] = U(Oa). This follows the same process as for our multi-objective local search described below.
Expected Hypervolume Improvement: Bayesian optimization optimizes single function f - E, where in each iteration t we select the most promising candidate in the input space 0, E arg max (.) (0) according to some acquisition function a θ + RR The idea of this ac- quisition function is to trade-off exploration and exploitation, based on a probabilistic model of the objective function D)), trained on some observed data points D. (0) 30). (0., ...
y: ), where y (f) To adapt Bayesian optimization to the multi-objective setting, we follow Daulton al.) (2020) and use single Gaussian process p(JID) for all objective functions ʃ(0) [fo (), .... (0)). Here, the acquisition function a0) (H the hypervolume HVI(y) = HV(P, yy.r) HV based on the current Pareto front Py and some reference point
computes the expected improvement of
Multi-Objectivv ASHA: Given a halving constant 77, minimum I'in and maximum T'aa number of epochs for fine-tuning, successive halving defines set of rungs R = (r.hiilk lk O..., K} where for simplicity we assume that Tmaz = ,, Starting from a set of randomly sampled configurations Tpms
C = (60 O,), successive halving evaluates all configuration on the first rung level ro and promotes the top 7 configuration for the next rung while discarding the others. This process is iterated until we reach the maximum rung level F% = Tmaz- We follow common practice (Li et al. 2017 Falkner et al., 22018) and run multiple rounds of successive halving until we hit maximum budget (defined in wall-clock time). Asynchronous successive halving (Li et al. 2018) standard successive halving to distributed asynchronous case, which requires some changes in the decision making routing. To cater for the multi-objective setting, we use multi-objective ASHA (Schmucker et al.) 2021) which uses non-dominated sorting instead of just the validation accuracy to rank configurations on each rung-level.
Multi-objective Local Search: Previous work (White el al. 2021aa) has demonstrated that simple local search often performs competitively compared to more advanced NAS methods. We propose a straightforward multi-objective local scarch approach. Starting from the current Parcto front P,- which is initialized by some starting point, we randomly sample an element 0, ~ Py and then generate random neighbor point by permuting single random entry of 0,.
18

---

Published in Transactions on Machine Learning Research (09/2024)
STWW
COLA
$0000
Taxt eroe
000
30
64
64 ONU
paramttt
5372
SIAG
000
0.00g
Teee arce
0000
500
04
Figure 10: Single Pareto fronts random run for each method on all 8 GLUE datasets.
input Search space θ, number of iteration T, starting point O start
output: Pareto front P
/* evaluate starting point
Po (ttart):
Ystart = (Ostaar)) (Ostar))
Y {
/* main loop
*/
for = I....,T do
/* sample random element from the population
*/
0, ~ ((P-11):
/* mutate
*/
d u(0, (0,)):
sample random dimension // sample a new value from the search space
0 + (O:)
0 u(a)):
/* evaluate
(()):
Y + YU
/* update population
S(Y) = (y' EY: fv" Y y > v.v/ + y"} = 0}:
// Pareto front
P (0: y() S(Y)):
end
Algorithm 5: Local Search
E
Additional Results
In this section we present additional results from the experiments described in Section 4 Figure 10 shows the Pareto front of randomly chosen runs for each dataset for BERT-base, sorted by the training dataset size. RTE and MRCP are small datasets compared to other datasets, leading often to an unstable fine-tuning process.
Figure and 12 shows the optimization trajectories for each methods using standard NAS and weight- sharing based NAS, respectively. Solid lines indicate the mean and shaded area the standard devication.
19

---

Published in Transactions on Machine Learning Research (09/2024)
ST5B
COLA
RTE
regret Toy
120
File
W0
266
6i
189
200
runtime
runtime
runtime
runtime
IMDB
SST2
SWAG
QNU
hypereolume
hngperveture
0:4
rearet
regret
041
runtime
ninfimme
runtime
runtime
RTE
MRFC
STSB
rogret hypervalume
hryper
rogret hop
regrett
runtime IMDS
runtime
runtima
runtime
SST2
SWAG
QNU
114
reg rogret hopee hypervolume
0:2
regret
Dolr nuntime
140
nuntime
runtime
runtime
Figure 11: Hypervolume of standard NAS methods (last two rows).
(first two rows) and RoBERTa-base
F
Quantization
Quantization (Dettmers et al. 2022 Dettmers & Zettlemoyer 2023 is a powerful technique that significantly
reduces the memory footprint of neural networks. However, its impact on latency is not immediate, especially when dealing with batch sizes that can not fit into the cache of the device Dettmers & Zettlemoyer (2023) With our flexible NAS framework we can simply replace objectives and directly optimize latency on the target device instead of parameter count.
Figure 13 left shows the Pareto set obtained with our NAS approach, where we optimize latency instead of parameter count on the COLA dataset across different GPU types. Additionally, we evaluate the perfor- mance of the unpruned super-nctwork with 8-bit (Dettmers et al., 2022) and 4-bit (Dettmers & Zettlemoyor 2023) quantization. While quantization substantially reduces the memory footprint. (Figure 13 right), it ac- tually leads to worse latency. While quantization introduces small overhead due to the additional rounding steps, the latency could potentially be reduced by optimizing the low-level CUDA implementation. Some- what surprisingly using int-8bit quantization leads to high performance drop on some hardware. NAS effectively reduces the sizes of weight matrices, leading to reduced GPU computation and, thus, is less hardware depend.
We can also apply quantization to sub-networks, making orthogonal to our NAS methodology and offering further improvements to the memory footprint. Overall, these findings shed light on the trade-offs wetween
20

---

Published in Transactions on Machine Learning Research (09/2024)
S15B
DOLA
D00022 leyer vouorrttuu
1nppect Teypet
pusiiin (ancedds)
Furiin (sonndd)))
stailine (scordda)
SWAG
ONLI
[axsesds
(soond)))
(somnth))
cuallin (acorda)
MEBPC
STAB
COLA
(oconds)
(coond))
(coonds)
IMUD
SWAG
QNLI
26
Figure 12: Hypervolume of search methods for weight-sharing based NAS on BERT- (first two rows) and RoBERTa-base (last two rows).
COLA
COLA
1.0
1.0
0.6 tbtt
0.9
0.6
test error
89
0est
8
0.6
0.6
0.5
O 00
05-
V100
PoPo
4L6
memory footprint
latency (milli-seconds)
Figure 13: Test error versus memory footprint (left) and latency (right) on 3 different GPU types for the Pareto front found by our NAS strategy and the un un-pruned network with 8bit and 4bit quantization.
memory footprint reduction and latency optimization. We leave it to future work to explore the connection between NAS and quantization.
21


Reviews:
[{"review_id": "UDT4k1dekY", "reviewer": "Reviewer_TYiE", "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) relevant approach to the structural pruning procedure\n2) extensive experiments that confirm the efficiency of the introduced approach\n3) clear description of the problem statement and suggested algorithm\n\nWeaknesses:\n1) the main text and plots should be carefully proofread to sync them and use the same notation\n2) only two not large-scale models are considered in experiments\n3) absence of the runtime comparison for the considered approaches", "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 presented paper considers the problem of structural pruning of the pre-trained language model from the neural architecture search perspective. The authors introduce four search spaces that define subnetworks. They state the corresponding bi-objective optimization problem, where generalization error and latency are used as evaluation criteria. The experiments consist of two parts: comparing different NAS methods and comparing the best NAS method with alternative structural pruning. The BERT-base and RoBERTa-base models are used for evaluation purposes and multiple text classification tasks. The experiments show that the NAS-based approach to structured pruning of the considered models is competitive or even better scalable than alternative approaches.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Major requests:\n1) in the contribution list, the authors mention four search spaces but then map the existing structural pruning only to two of them. Please comment on why there are only two. This contribution point induces more intrigue than makes your results clear.\n2) please add a caption to the two subplots with average ranks vs time steps\n3) what are the reported ranks in Figure 5b? No definition in the paragraph \"Multi-Objective Search\"\n4) the runtime comparison of the discussed methods is missing in the text. That is an essential factor in selecting the method for structural pruning. Please add such a comparison in section 4 and discuss the observed trade-off (if any) between the Pareto optimality of the obtained subnetworks and the computational complexity of the tested methods.\n5) please add the source code link to reproduce your results \n\n \nMinor requests:\n1) \"They demand significant GPU memory and exhibit high inference latency, making them impractical\nfor many real-world applications, for example when used in an end-point for a web service or deployed on an\nembedded systems\" - make consistent usage of article \"an\" and the following adjective+noun \"embedded systems\". The text has other typos; please proofread it carefully.\n2) the notation $\\circ$ in applying the binary mask is not defined\n3) please make consistent y-label and text description for Figure 7 (\"test performance\" vs \"test error\"), or add explicit clarification that the smaller, the better\n4) \"Especially for higher pruning ratios, NAS outperforms ...\" - what NAS? S-NAS or WS-NAS?", "broader_impact_concerns": "No concerns about the ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "nS2K0wh4ay", "reviewer": "Reviewer_cEWC", "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- The paper is well written and serves as a gentle introduction to NAS approaches in a relevant context (pruning PTL models).\n- The main multistage approach is well presented (additional questions presented in the request changes section).\n- The set of experiments is quite thorough. Eight benchmarks are considered on two different PTL models. More than ten approaches are compared.\n- The overall approach of the authors can serve as a benchmark for future research on NAS techniques. Benchmark papers can be extremely relevant in the field for further progress.\n- Tackling pruning via a Pareto frontier can offer simple selection tools for end-users. This point could even be more highlighted. For example, the experiments in the Appendix on quantization show alternative multi-objective definitions that would be relevant for users.\n\nWeaknesses:\n- The exhaustive nature of the paper is remarkable. However, currently, take-away messages are not fully clear-cut. This is mostly a question of presentation. The current material is enough to raise the interest of peers.\n- As the main contribution consists in a benchmark, the metrics should be fully defined in the paper rather than using some references.", "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 considers pruning pre-trained language (PTL) models for downstream tasks requiring low latency and/or low memory footprint using various neural architecture search (NAS) techniques. As a byproduct, the paper stands as a benchmark for numerous NAS methods.\nThe paper starts with a comprehensive overview of the challenges raised by downstream applications of PTL and different NAS techniques in the first two sections. The third section lays down the NAS approach for structural pruning of PTL models using masks on multi-head attentions as well as fully connected layers. Several search spaces are presented with a quick overview of their relevance. A complete two-stage search procedure is detailed. Experiments delve into the relevance of NAS approaches for PTL models pruning. A first set of experiments demonstrate that a small search space is the most efficient in most benchmarks. Then a set of experiments study the performance of various NAS searches on this small search space. The authors present then comparisons of multistage NAS searches that also serve as an ablations study of the main approach presented earlier. Additional comparisons with quantization techniques are present in the Appendix.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Questions:\n- Why using the *first* units/heads? Since one can simply permute the columns of the next matrix multiplying the inputs, it seems that this choice is arbitrary. The argument of about the problematic concentration of the MEDIUM search space is great. But why not adding a permutation of the units/heads?\n- An important point raised by the authors is the actual hardware constraints imposed on such models. I suppose one such constraint is compatible sharding. It would be great to expand on that since this can offer a principled approach to efficient pruning that satisfies the constraints (if I understand well the LARGE search space does not satisfy such constraints). \n\nSuggested changes:\n- Give a complete mathematical definition of the hypervolume as it is the main metric used in the benchmark. \n- Maybe give more mathematical details on the NAS search that appears to perform the best in the main text and more details on all approaches in the Appendix to ensure that the paper is self-contained. \n- Consider making paragraphs with take-aways from the experiments in bold, and then add additional details. \n- Maybe consider making a streamlit link to ease the comparisons between approaches. A plot with more than three curves is generally hard to parse. (This is just a detail, feel free to ignore).\n- Appendix F is quite interesting. It shows in particular the flexibility of the multi-objective approach to tackle issues beyond mere memory footprint. In fact latency is a leitmotif in the introduction that does not appear anywhere else. So I would suggest adding this section in th main text to strengthen the approach and maybe consolidate that experiment by adding for example Pareot frontiers and experiments on all benchmarks. \n\nDetails:\n- On Fig. 1 MHA and FFN are acronyms which are defined much later, consider using the full names here.\n- PTL (page 3) not defined.\n- \"For example, Sanh et al (2020) was trained...\" (page 3) -> the author was not trained for 90 hours hopefully. I should read \"the model of Sanh et al...\"\n- Maybe fully detail the Attention layer page 4. The paper has the potential to not only serve as a benchmark but also a gentle introduction to such pruning techniques, so it could be worth going all in on the the mathematical details.\n- The notation U in the definition of $M_{neuron}$ is not given (it is way later but not at the time of the definition page 4). \n- The notation $\\circ$ in the application of the masks (page 4 second paragraph of Section 3.1) is not defined and also uncommon. A $\\cdot$ may be more appropriate.\n- The dimensions of the matrix multiplications in the definition of FFN (page 4) are ill-posed (you may use $X^\\top$). Then when using $M_{neuron}$ it could be worth fulyl detailing mathematically what the $\\circ$ means in this context (namely multiplications of rows/columns, depending on your choice of definition of FFN).\n- Section 3.2.1 second sentence: \"initialize\" -> \"initialized\".\n- Eq. 2 make parentheses larger for terms inside $\\sigma$.\n- Mention that $\\mathcal{L}_{CE}$ is defined on some data (and maybe detail which one in practice).\n- Last sentence of Section 3.2.2 not clear for non-knowledgeable readers.\n- What do you mean by the dimensionality of the search space and its upper bound? (Section 3.3, SMALL search space second sentence).\n- Put in italic/bold the NAS searches presented and used in Section 4.1.2.\n- Add a suptitle on the first and third rows of e.g. Fig. 4 that says \"BERT\" or \"RoBERTa\" for ease of reading.\n- In the last sentence of the paragraph Multi-Objective Search (page 9) there is a verb missing.\n- In Appendix C, maybe give references to the different datasets used.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "s121dI2KxM", "reviewer": "Reviewer_NdN4", "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\n1. The technology foundation of this paper is reasonable and sound. Model pruning and NAS share the same target: keeping only a part of the model while not harming the performance. Therefore, it is natural to apply the NAS technique to the model pruning problem.\n\n2. The evaluation of this paper is very solid. This paper evaluates eight different datasets and two different language models with different pruning algorithms and settings. The substantial experiment result makes the conclusion quite convincing.\n\n## Weaknesses:\n\n1. This paper does not prove that the NAS-based solution has an obvious advantage over traditional solutions. According to the evaluation, traditional solutions have higher accuracy with small or medium pruning sizes. Though NAS solutions have higher accuracy with large pruning sizes, the accuracy loss could be 50% or more. To the reviewer's understanding, such accuracy loss is unacceptable. \n\n2. The major metric Hypervolume is not well described. The description in Sec 4.1 is not clear, and it is hard to understand the meaning of the hypervolume and how we calculate it.", "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": "Though language models are very powerful in solving many natural language tasks, deploying them could be challenging because they require huge GPU memory and usually have high inference latency. This paper explores the model pruning to find a sub-part of the neural network, which is a trade-off between the resource consumption and the model capability. \n\nThe key idea of this paper is to use the Neural Architecture Search (NAS) technique to solve the model pruning problem. It first pre-trains the super network. Then, it generates many sub-networks by using different masks and fine-tunes those sub-networks. Finally, it selects those sub-networks that balance the model size and predict accuracy. \n\nIn summary, this paper makes three main contributions:\n\n1. This paper proposes a new use case of NAS that can help solve the model pruning problem.\n2. This paper proposes four search spaces of pruning language models that cover existing structural pruning strategies.\n3. This paper proposes a new benchmark for multi-objective NAS. It evaluates different NAS algorithms on eight classification tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The author should add some experiments or data to prove that the NAS-based solutions do have some unique advantages. For example, the NAS-based solutions have higher and acceptable model accuracy with the same pruning size, or they can do pruning much faster than traditional solutions.\n\n2. The author should modify Sec 4.1 to define the hypervolume clearly.", "broader_impact_concerns": "No ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper explores neural architecture search (NAS) for structural pruning.
The goal of structural pruning is to identify sub-parts of a fine-tuned network that optimally trade off efficiency, such as model size and generalization performance. A key contribution of this paper is that it does not propose a pruning method with a fixed threshold. Instead, it advocates for a multi-objective approach that identifies a Pareto optimal set of sub-networks using the hypervolume metric. This allows for a flexible and automated compression process.

The reviewers commended the thorough experimental verification. Eight benchmarks are considered on two different pre-trained language models, and more than ten approaches are compared. The authors' overall approach can serve as a benchmark for future research on NAS techniques.

The paper is well-written and briefly introduces NAS approaches in the context of pruning pre-trained language models.

The reviewers find that the paper's exposition is extremely clear, the connection between NAS and pruning is well explained and explored, and the general approach (Pareto vs. focus on a single technique) could be of interest to end-users. Therefore, we propose this paper for a survey certification.

When preparing the final version of this manuscript, please consider the reviewers' suggestions for making this paper even more self-contained.
Decision: Accept (poster)