Few-shot Task-agnostic Neural Architecture Search for Distilling Large Language Models
Authors: Dongkuan Xu
Subhabrata Mukherjee
Xiaodong Liu
Debadeepta Dey
Wenhui Wang
Xiang Zhang
Ahmed Hassan Awadallah
Jianfeng Gao
Venue: NeurIPS 2022 Accept (2022)
Keywords: Pre-trained Language Models
Knowledge Distillation
Neural Architecture Search

Abstract: Traditional knowledge distillation (KD) methods manually design student architectures to compress large models given pre-specified computational cost. This requires several trials to find viable students, and repeating the process with change in computational budget. We use Neural Architecture Search (NAS) to automatically distill several compressed students with variable cost from a large model. Existing NAS methods train a single SuperLM consisting of millions of subnetworks with weight-sharing, resulting in interference between subnetworks of different sizes. Additionally, many of these works are task-specific requiring task labels for SuperLM training. Our framework AutoDistil addresses above challenges with the following steps: (a) Incorporates inductive bias and heuristics to partition Transformer search space into K compact sub-spaces (e.g., K=3 can generate typical student sizes of base, small and tiny); (b) Trains one SuperLM for each sub-space using task-agnostic objective (e.g., self-attention distillation) with weight-sharing of students; (c) Lightweight search for the optimal student without re-training. Task-agnostic training and search allow students to be reused for fine-tuning on any downstream task. Experiments on GLUE benchmark demonstrate AutoDistil to outperform state-of-the-art KD and NAS methods with upto 3x reduction in computational cost and negligible loss in task performance. Code and model checkpoints are available at https://github.com/microsoft/autodistil.

Few-shot Task-agnosti Neural Architecture Search for Distilling Large Language Models
Dongkuan Xu* NC State University dxu27@ncsu.c edu
Subhabrata Mukherjee Microsoft Research
submukhe@miczooofe com
Xiaodong Liu Microsoft Research
Debadeepta Dey Microsoft Research
Wenhui Wang Microsoft Research
Xiang Zhang Penn State University
Ahmed Hassan Awadallah Microsoft Research
Jianfeng Gao Microsoft Research
Abstract
Traditional knowledge distillation (KD) methods manually design student archi- tectures to compress large models given pre-specified computational cost. This re- quires several trials find viable students, and repeating the process with change computational budget. We use Neural Architecture Search (NAS) to automatically distill several compressed students with variable cost from large model. Existing NAS methods train single SuperLM consisting of millions of subnetworks with weight-sharing, resulting in interference between subnetworks of different sizes. Additionally, many of these works are task-specific requiring task labels for Su- perLM training. Our framework AutoDist addresses above challenges with the following steps: (a) Incorporates inductive bias and heuristics to partition Trans- former search space into compact sub-spaces (e.g., K=3 can generate typical student sizes of base, small and tiny); (b) Trains one SuperLMM for each sub-space using task-agnostic objective (e.g., self-attention distillation) with weight-sharing of students; (c) Lightweight search for the optimal student without re-training. Task-agnostic training and search allow students to be reused for fine-tuning on any downstream task. Experiments GLUE benchmark demonstrate AutoDiste to outperform state-of-the of-the-artttttt- KD and NAS methods with upto 41x reduction in computational cost. Code and models are available at alfak..s/auuodisti
Introduction
While large pre-trained language models (e.g., BERT I. GPT-3 (2) effective, their huge size poses significant challenges for downstream applications terms of energy consumption and cost of inference
[3] limiting their usage in the edge scenarios and under constrained computational has shown strong results in compressing pre- trained language models into small student models. However, these works require pre-specification of the student architecture and computational cost (e.g., number of parameters, FLOPs) for distillation. This poses significant challenges: (i) it requires several trials to come up with viable architectures as they are hand-engineered and to define several hyper-parametent (e.g.. number of layers and attention heads, hidden dimension, etc.); (ii) one has to re-run distillation with any change in specification for the student architecture or computational cost for using it in target environment.
inference budgets. Knowledge distillation 4 [5]
6. /
'Work done while interning Microsoft Research. CCorresoonding author and project lead.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

84
82
80
Accuracy
78
DySaBERT
Autolistlayes
AutoTinyEENT-ADS
76
AutoDistlleo
AutoDisti
AdtoTinyENTTKD- -222-
AutoTinyGENT-CG
74
Autoditt
BERTE PKD
AutoTinyBERT.OO.
NAS BEÄTu
NAS-BERTy
NAS-BERT,a
72
ThyyBRRTS MobileBERT
MinIMM
NAS-BERT,
0.0
0.2
0.4
0.6
0.8
1.0
1.2
#FLOPS
1e10
Figure 1: AutoDistil uses few-shot task-agnostic NAS distill several compressed students with variable #FLOPs (X-axis) from K=3 SuperLMs (corresponding to each point cloud) trained on K sub-spaces of Transformer search space. Each student (blue dot) extracted from the SuperLM is fine-tuned on MNLI with accuracy on y-axis. The best student from each SuperLM is marked in red. Given any state-of-the-art distilled model, AutoDist itiisill generates better candidate with less #FLOPs and improved task performance from corresponding search space. More discussions on trade-off and selection strategies can be found in the Appendix (Sections G
provides a natural solution to automatically search through a large space of candidate models. The dominant NAS paradigm consists two main steps: (a) Training Super model combining all possible architectures into single graph and jointly training them via weight-sharing; (b) Searching for optimal architecture from Super model with best accuracy on downstream task, satisfying user-specified latency constraint for target device. Parallel to above computer vision (CV) works, NAS has shown strong results in recent works like DynaBERT [12] for natural language understanding (NLU).
Neural Architecture Search (NAS) g
AutoTinyBERT [K] and NAS-BERT 14
Drawbacks of existing NAS methods.
[D1: Co-adaptation in weight-sharingl Above works train one single large Super Language Model (SuperLM) consisting of millions of diverse student architectures. This results in some undesirable effects of co-adaptation [15] like conflicts in weight-sharing where bigger student models converge faster in contrast to smaller ones converging slower [6 [
[D2: Multi-stage training] A single SuperLM may not have sufficient capacity encode large search space. As a result, these works use multi-stage training process, where they first conduct NAS identify candidate students and then perform further pre-training tion 14 of the candidates.
K and knowledge distilla-
[D3: Task-specific training] NAS works CV domain (e.g., AutoFormer 17. Once-for-all [[O] One-Shot NAS [II I&) leverage hard class labels from a task (c.g.. image classification) or soft labels from ImageNet pre- pre-trained models (c.g.. MobileNet 77.. RegNet (I9]) for task-specific opti- mization with accuracy as evaluation metric. Different from CV domain, NLU tasks have different objectives and evaluation metries for classification (e.g., MNLI), regression (e.g., STS-B) and corre- lation (e.g., CoLA). This makes it challenging to adapt existing NAS works to the NLU domain in task-agnostic setting. Recent NAS works in the NLU domain are not fully task agnossic For instance, DynaBERT [12 accesses both task labels for knowledge distillation and task development set for network rewiring. NAS-BERT I4] performs two-stage knowledge distillation with pre-training and fine-tuning of the candidates. While Auto'TinyBERY [K]] also explores task-agnostic training, we demonstrate better performance from few-shot NAS and much cheaper cost from single stage training without additional pre-training and distillation. We present detailed discussion the differences between Aut oDistil and existing KD and NAS methods based on fine-grained search space, different training strategies and amortized training cost in Table I More comparison details are summarized in Section A Appendix.
Contributions. We address above challenges with fully task-agnoster few-shot NAS consisting of three steps. (S1) Search space design. We partition the Transformer search space into sub-spaces considering important architectural hypee-parameters like the network depth, width and attention heads. We further leverage inductive bias and heuristics to limit the number of student architectures in each sub-space. (S2) Fully task-agnostic SuperLM training. We train K SuperLM overall, one

---

Autstistil Table Comparing AutsDistil with existing KD and NAS methods on aspects as task aasssoiosttcc training; generating students with variable compression cost, single-stage training without additional adaptation; Su training with compact search space to mitigate interference (P denotes partial).
Method
Task-
Variable Compression
NAS
agnostic
Single Stage
SuperLM
Compact Search
Training
BERT-PKD
X
SparseBERT
X
DistilBERT
X
TinyBERT
X
N/A
MOBILEBERT
X
MINILM
X
DYnaBERT
One-shot
NAS-BERT
P
One-shot
X
AutoTTnyBERT
P
One-shot
X
Few-shot
for every sub-space. This allows each SuperLM more capacity to encode a sub-space as opposed to a single large one. We train each SuperLM with a fully task-agnostic objective (without accessing any task labels) like deep self-attention distillation, where we transfer knowledge from the self-attention module (including keys, queries and values) of a pre-trained teacher (c.g.. BERT) to the student and use
to train the SuperLM. (S3) Lightweight optimal student search. We obtain optimal student(s) directly from well-trained (uperLM(s) without any re-training that can be simply fine-tuned on downstream tasks. Our contributions over existing NAS works can be summarized as:
In contrast prior works (e.g., DynaBERT, AutoTinyBERT, NAS-BERT), do single-stage training combining NAS and distillation with no further pre-training or augmentation and demon- strate superior performance of the NAS process itself with significantly reduced training cost. Obtained subnetworks are simply fine-tuned on downstream tasks.
Fully task-agnostic training with subnetwork attention state alignment for self-attention relation distillation and search in contrast to prior works in NLU (e.g.. DynaBERT, NAS-BERT) and CV (e.g., AutoFormer, BigNAS, Once-For-All,
Few-shot NAS to mitigate gradient conflicts in SuperNet training compared to prior One-shot NAS works in NLU (e.g., DynaBERT, BEERR,,, NAS-BERT). AutoFormer in the CV domain is an exception this point which also uses few-shot NAS but accesses task labels during training.
Strong results over all the above NAS and distillation works in NLU with 3x additional compression over best performing distillation technique with negligible performance drop.
Background
We present an overview of Transformers 20 especially its two main sub-layers, multi-head self- attention (MHA) and ferw-forward network (FFN). Transformer layers are stacked encode contex- tual information for input tokens as: X' Transformer, (X' (X'-1), I [I, L] where L is the number of Transformer layers, X' E R sedaad s is the sentence length, and daid is the hidden dimension. In the following, we omit the layer indices for simplicity. Multi- Self-Attention (MHA), Given previous Transformer layer's output X. MHA computes:
Attention Q+. Ki, Vn. softmax( Q.K...
Vh:: KR, V, XW XWX XWX.
(1)
Vdhead
MHA(X) = Concat Chead, ... headh W
(2)
where WO W.K. K. WY E Rddiaadnead WO E adiaa are linear transformations. Q., K,, V E Retdhead are called queries, keys, and values, respectively. H is the number of heads. head, Attention Q Kr,, Vv) denotes the h-th attention head. Concat is the concatenating operation. dhead = drea/ H is the dimension of each head.
Feed-Forward Network (FFN). Each Transformer layer contains an FNN sub-layer, which is stacked on the MHA. FFN consists of two linear transformations with ReLU activation as:
FFN(x) max(0, x. + +b,)
W + b2.
(3)
3

---

Transformer Sumpled buknetwwrks
Multi Selfaatendion Balation Dissllation
Traclumer
Performnce
Search Engine
w
Trardorners
Tranafomer's
w
Trarsfomeer
Trandorner's
Trats former
Transfornet
Attention
Sampled Subnetworks
Rg
KSearch Partitions
Learning SupertM
Figure 2: Overview of AutoDistil considers K partitions of the Transformer architecture subspace to train one SuperLM for each partition with weight- sharing of the constituent subnetworks trained via task-agnostic deep self-attention distillation. Optimal compressed subnetworks can be easily extracted from the SuperLMs without additional training or distillation.
where W Rodnaady 2 E Rdd+ddid, by E R2,, and by E Riddid, In addition, there are residual connection and layer normalization top of MHA and FFN (denoted by + in Figure 2 which are formulated as LayerNorm(x + MHA(x)) and LayerNorm(x + FFN(x)), respectively.
3 Few-shot Task-agnostic NAS
Given a large pre-trained language model (e.g., BERT) as teacher, AutoDistil distills several compressed models with variable computational cost with the following major components.
3.1 Search Space Design
Searchable transformer components. From Transformers overview (Section 2 and our framework (Figure 2) we observe four important hyper-parameters for the Transformer blocks include: (1) Ferdarorward network (FFN) dimension we encode this by the MLP (multi-layer perceptron) ratio defined
dy Thid with d and dhad representing the intermediate dimension of the FFN and hidden dimension respectively; (2) Number of layers (L) capture the network depth; (3) Hidden dimension (dhid) to encode input; (4) Attention heads (H) for multi-head self-attention.
All of the above factors are important for model capacity and have significant impact the model size and computational cost. For instance, different layers have different feature representation capabilities. Recent works show that Transformer models are overparameterized 21 22 such as the feed-forward layer (FFN), which one of the most computation intensive components [ Therefore, search for optimal MLP ratio and hidden dimension that reduce computational cost resulting from FFN layers. Furthermore, studies [24. 25 show that attention heads can be redundant when they learn to encode similar relationships for each word. Thus, make the number of attention heads searchable as well.
Inductive bias. Prior works on CNNs 26 and Transformers 2 demonstrate that thinner and deeper neural networks with improved representation capacity perform better than wider and shallower ones. We incorporate this inductive bias to decide the number of layers to consider for the students in each our K sub-spaces (base, small, tiny), where we prefer deeper students in terms of the number of layers. Furthermore, we constrain all the Transformer layers in given student model to share identical and homogeneous structures, i.e., the same number of attention heads, hidden dimension, etc. This not only reduces the size of the search space, it is also more friendly to hardware and software frameworks [3
Search space partition. Existing works 11 14 train single large SuperI.M containing millions of student architectures by weight-sharing. This leads to performance degradation due to optimization interference and convergence of subnetworks with very different sizes |Л]. To mitigate such

---

interference, we employ a few-shot learning strategy 17 16 as follows. We partition the whole Transformer search space into K sub-spaces such that each sub-space covers different sizes of student models given by the number of parameters. For instance, K = 3 can cover typical student sizes, namely base, small and tiny versions. Table 2 shows the parameter ranges for the K sub-spaces, along with the student configurations contained in each. Extensive details on the search space design can be found in the Appendix (Section D)
We now encode each sub-space into a SuperLM, where each student model in the space is a subnet- work of the SuperLM. Furthermore, all the student subnetworks share the weights of their common dimensions, with the SuperLM being the largest one in the search space. Considering K independent SuperLMs, each one now has more capacity to encode a sub-space, in contrast limited capacity single SuperLM as in prior works. Furthermore, our choices for the heuristic partition and inductive bias result in less number of student models of comparable size in each sub-space which alleviates conflicts in weight-sharing.
We extract student subnet- works from the SuperLM by a simple truncation strat- egy like bottom-left extrac- tion. In particular, given a specific architecture a dnidd r. h), (i) we first extract alternate 1 Trans- former layers from the Su- perLM; (ii) then extract bottom-left sub-matrices in terms of daid and from the original matrices that repre- sent the hidden dimension and the MLP ratio respec-
SuperLMMsmall MumerLMrmy
SuperLMBase
BERT
#Subnets #Layers
256
256
256
N/A
(4,7, 7, 1)
(9, 12,
(9, 12, 1)
12
#Hid dim
(128, 224, 32)
(256, 352, 32)
(544, 640, 32)
768
MIP Ratio
(2.0, 3.5, 0.5)
(2.5, 4.0, 0.5)
(2.5, 4.0, 0.5)
4.0
#Heads
(7, 10, 1)
(7, 10, 1)
(9, 12, 1)
12
#FLOPs
40-367.M
0.5-2.1G
2.1-7.9G' 39-79M
11.2G
#Params
4-10M
12-28M
109M
Table 2: The search space of AutoDi st with K=3 partitions, each consisting of 256 subnets with variable computational cost. We train one SuperLM with weight-sharing for each partition with child models sharing transformer blocks. Each tuple represents the lowest value, highest value, and steps for each factor.
tively; (iii) finally, for the attention heads, we extract the leftmost h heads and retain the dimension of each head as the MuperLM This strategy is used during sampling sub-networks for SuperLM training via weight sharing after which the sub-network weights updated via self-attention relation distillation. There can be better strategies to extract sub-networks (e.g.. ordering heads by importance) that we defer to future work. We did explore some strategies for selecting layers (alternate vs. top vs. bottom) with results in Section B of Appendix.
3.2 aano-ig SuperLM Training
We illustrate SuperLM training process in Algorithm I Given a large pre-trained language model (e.g.. BERT) as teacher, we initialize the SuperLM with the weights of teacher. In each step of SuperL training. we randomly sample several student subnetworks from the search space; apply knowledge distillation between sampled subnetworks and the teacher to accumulate gradients; and then update the SuperL.M. During sampling, we employ Sandwich rule (28]. also used in BigNAS [. that samples the smallest subnetwork, the largest subnetwork and M random ones for updating SuperI.M. The motivation is to improve the performance of all subnetworks by increasing the performance lower bound (smallest subnetwork) and upper bound (largest one) across all subnetworks.
We leverage deep self-attention distillation for task-agnostic training. To this end, we employ multi-head self self-attention relation distillation to align the attention distributions as well as scaled dot-product of keys, queries and values of the teacher and sampled student subnetworks. Consider A1, Ag. As to denote the queries, keys and values of multiple relation heads of teacher model, and B,, B2, B3 respectively for sampled subnetwork. Mean squared error (MSE(.) between multi-head self-attention relation of teacher and sampled subnetwork is used as distillation objective:
" MSE(R. R
1
L BeCi
L.
H
where R.- softmax(A,A) / de). R (B, /Vde. H is the number of attention heads; R.' represents the teacher's Q K. K, V - V relation; R represents the same for student. R.T is the relation information based attention head, and dy is the attention head size.
5

---

Relation knowledge distillation avoids the introduction of additional parameters to transform the student's representations with different dimensions to align to that of the teacher. For the teacher model and subnetworks with different number of attention heads, we first concatenate the self- attention vectors of different attention heads of the subnetwork and then split them according to the number of relation heads of the teacher model. Then, we align their queries with the same number of relation heads for distillation. In addition, we only transfer the self-attention knowledge from the last layer of the teacher model to the last layer of the student model.
The SuperLM for sub- subcepace Ag trained as: U;Drraiss Derain)) )}
Algorithm Few-shot Task-agnostic Knowledge Distillation with AutoDistil.
W = argminw A U;
(5)
Input: Partitioned K sub-spaces Au; initialized K Su- perLMs St on Au; pre-trained teacher T; unlabeled data D: training epochs E; sampling steps M
where, K' is the number of sub-space parti- tions; W are the weights of the SuperLM; Wa the weights in W specified by the chitecture a; U are the weights of the teacher model including the self-attention module used for distillation; Drrain is the training data set, and L(-) is the self-attention loss function from Eqn.
Output: Trained MuperL {Su}
for k = to K do
for i = 1 to do
Get batch of data from for batch
do
Clear gradients in SuperL S₁
for m = 1 to M do
Randomly sample a subnetwork from Ss
Calculate self-attention distil. loss between subnetwork and teacher T with Eqn. Accumulate gradients
3.3 Lightweight Optimal Student Search
We outline two search strategies for selecting the optimal student subnetwork.
end for Update S& with the accumulated gradients end for end for end for Task-agnostic search. We adopt this to be our primary strategy to compare against all base- lines since does not access any task label
information. We compute the task-agnostic self-attention distillation loss for all student subnetworks using Eqn. on a heldout validation set from the unlabeled training corpus. The student subnet- works are directly obtained by bottom-left extraction from the well-trained SuperLM (outlined in Section B.T. This process lightweight since it does not require any training or adaptation of the student and number of subnetworks is limited. The optimal student is given by the subnetwork with least validation loss subject to following constraint.
a'A = ar argminaeAu...... ; Deat), s.t. g(a) <
(6)
where W is the weights of architecture obtained from W Duai is the validation data set, L the self-attention distillation loss, and 9(-) is a function calculate the computational cost (e.g.. #FLOPs, #parameters) of the subnetwork subject to given user-specified resource constraint c.
Task-proxy proxy search. We compare our task-agnostic search against another strategy that considers a proxy task (c.g., MNLI (9) with label information to fine-tune the 256 candidate subnetworks in each sub-space. The optimal student in each sub-space is given by the one with the best downstream task performance (e.g.. accuracy). Note that, for this strategy, the proxy task is used only during search while the NAS training is still fully task-agnostic.
Experiments
Datasets. We conduct experiments on General Language Understanding Evaluation (GLUE) bench- mark 330 We compare our method with the baseline methods on two single-sentence classification
tasks (CoLA GT. SST-2
(32). two similarity and paraphrase tasks (MRPC 33 QQP
B4 ). and
three inference tasks (MNLI 29 QNLI 35), RTE
C6
37
38
39
We report accuracy for MNLI,
QNL.I, QQP, SST-2, RTE, report fl for MRPC, and report Matthew correlation for CoI.A. El
Baselines. We compare against several task-agnostic method generating compressed models from BERT BaRTbase base teacher, using (i) knowledge distillation like SMALL [40, Truncated BERT 29)
DistilBERT [5). Tiny BERT G. MINILM 4: as well as those based on Neural Architecture Search,
like AuEoTTyBERRT 13. DERT 12 and NAS-BERT [14]
We ignore STS-B for fair comparison with our strongest KD baseline MINILM 4 that do not report it.
For a fair comparison, we do not include MoR that BERT as teacher.

---

AutoDistil configuration. We use uncased BART as the teacher consisting of 12 Trans- former layers, 12 attention heads; with the hidden dimension and MLP ratio being 768 and respectively. It consists of 109M parameters with 11.2G FLOPs. We use English Wikipedia and BookCorpus data for LuperLM training with WordPiece tokenization. We use a batch size of 128 and 4e-5 as the peak learning rate for 10 epochs. The maximum sequence length is set to 128. The coefficients in distillation objective (Eqn. 4 31. 32. and 33, are all set to 1, We distill the self-attention knowledge of the last layer to train the SuperLM Both the teacher and SuperLM are initialized with pre-trained BERT Other hyper-paramett settings are shown in Appendix. We use 16 V100 GPUs train the SuperLM with 336 GPU-hours as the training cost.
4.1 Finding the Optimal Compressed Models
AutoDistil Aguosstti using any task label information. We set constraint in Eqn. O such that the #FLOPs of the optimal compressed model is atleast 50% less than the teacher model. We rank all the subnetworks contained in all the partitions of the trained SuperLM by their self-attention distillation loss on the heldout validation set, and select the one that meets the constraint with the minimum loss.
Aguostic is obtained by fully task-agnostic training and task aaskstgnostic search without
iittDiistilpree uses MNLI 29 as proxy to estimate downstream task performance of different subnetworks. Prior work 41 has demonstrated performance improvements in MNLI to be correlated to other GLUE tasks. To this end, fine-tune all the 256 subnetworks in each partition of the trained superLMs, and select corresponding subnetworks with the best trade-off (more discussions on the optimal trade-off can be found in the Appendix (Section G between task performance (accuracy) and computational cost (#FLOPs). This results in K=3 optimal students, corresponding to AuteDist P oD i i Proxys and AutoDi sttpissiiin i obtained from the corresponding sub-spaces of SuperL SuperLMsmall and SuperL respectively. Notably all students are obtained from the AutoDisti SuperLM still trained in a fully task agnkstgn fashion.
4.1.1 Comparison with Traditional Knowledge Distillation Baselines
We compare AutoDistil against state-of-the-art KD models distilled from the same teacher BERTBASE Table D with respect to the following measures: computational cost in the form (i) FLOPs and (ii) parameters, along with (iii) improvement in the average task performance aggregated over all the GLUE tasks. We observe that the compressed model outoDist Agnostic generated via task-agnostic SuperLM training leads to upto 3% reduction in FLOPs over state-of-the-ar distilled models (c.g., MINILM [4) TinyBERT (,, DistilBERT (5)) that are hand-cnginne -eginneeed while matching the overall task performance. The most aggressive compressed version corresponding to AutoDi
obtains a massive 41x reduction in FLOPS over BERTBASE while incurring 5 point accuracy drop in GLUE (excluding CoLA) and 10 point drop (including CoLA). Notably CoLA is a syntactic task in contrast to other semantic tasks in the benchmark like natural language inference, paraphrase detection and sentiment classification. This depicts an interesting impact of massive model compression on varying task types.
4.1.2 Comparison with Neural Architecture Search Baselines
We report the performance of several NAS- generated student models of comparable FLOPs and parameters from corresponding papers in T AutoDistil outperforms all compet- ing methods on aggregate for all sizes; except for small-sized model; where it has marginally lower performance (0.1 points on avg) compared to AnyoTTTyBERT. It is worthwhile to note that computational cost of training process is another important dimension for comparing methods. This is especially important when comparing to NAS methods that multi-stage training: where additional traitraining and distillation is applied to NAS-generater candidates.
Cost (GPU
AutoTiny AutoTiny
Auto Distil
hours)
BERT
BERT-Fast
SuperNet Training
NR
NR
336
Search
150
12
<1
Further Training
870
290
0
Table 4: Training cost (V100 GPU hours) com- parison for generating students of similar FLOPs. N Tuto BERT does not report the cost of Su- perNet training typically the most expensive step. Further Training refers to additional pre-training applied to NAS -enegeree candidates.
To better understand the impact of single-stage vs. multi-stage methods on the training cost, we
7

---

Table Performance comparison between students from traditional task-agnostic distillation; multi- stage one-shot NAS with additional pre-training, distillation; and single-stage few-shot Aut Our results are averaged over runs with baselines reported from corresponding papers.
Model
#FLOPs
#Para
MNLI-m QNLI QQP SST-2 CoLA MRPC RTE
Average
(Metric)
(G)
(M)
(Acc)
(Ace) (Ace) (Ace) (Mce) (Ace) (Ace)
BERT BASE (teacher)
11.2
109
84.5
91.7
91.3
93.2
58,9
87.3
68.6
82.2
Base-sized Models from Task-agma agnostic KD Methods and Autobistil
SMALL
SMALL
40
5.66
66.5
81.8
89.8
90.6
91.2
53.5
84.9
67.9
80.0
Truncated BERT
29
5.66
66.5
81.2
87.9
90.4
90.8
41.4
82.7
65.5
77.1
DistilBRRT 5
5.66
66.5
82.2
89.2
88.5
91.3
51.3
87.5
59.9
78.6
TinyBERT 6
5.66
66.5
83.5
90.5
90.6
91.6
42.8
88.4
72.2
79.9
MINILM
5.66
66.5
84.0
91.0
91.0
92.0
49.2
88.4
71.5
81.0
AutoDisti 1 Proxy's
4.40
50.1
83.8
90.8
91.1
91.1
55.0
88.8
71.9
81.7
Small-sized Models from Multi- One-shot NAS Methods and AutoDistii
13
1.69
30.0
82.3
89.7
89.9
91.4
47.3
88.5
71.1
80.0
DynaBERT 12
1.81
37.7
82.3
88.5
90.4
92.0
43.7
81.4
63.2
77.4
NAS-BERT 14
2.30
10.0
76.4
86.3
88.5
88.6
34.0
79.1
66.6
74.2
AutoDistio 1
2.02
26.1
83.2
90.0
90.6
90.1
48.3
88.3
69.4
79.9
AutoDi ittiiistiii 1 Agnoatic
2.13
26.8
82.8
89.9
90.8
90.6
47.1
87.3
69.0
79.6
Tiny-sized Models from Multi-stage One-shot NAS Methods and AutoDi sttii
BEERT-NYBERT----SS [3
0.30
10.1
76.0
85.5
86.9
86.8
20.4
81.4
64.9
71.7
NAS-BERTs 4
0.87
5.00
74.4
84.9
85.8
87.3
19.8
79.6
66.6
71.2
AutoDisti iloxyy Proxyr
0.27
6.88
79.0
86.4
89.1
85.9
24.8
78.5
64.3
72.6
compare the overall cost of NAS for AutoDistil and that reported in BEERT for the small model segment in Table td AutoDistil is much cheaper due to its single-stage training protocol; where no additional pre-training distillation is needed. It worth noting that the overall SuperNet training cost of Aut ADistiittin (the most expensive component of NAS) is less or comparable to the additional training cost of re-training candidate models for Tiny BERT. Note that Auto BEERRTT does not report their SuperNet training cost. Additionally, AutoDi i still has much faster search mechanism due (1) inductive biases built into the search space definition to limit the number of student architectures and (2) task-agnostic search that only requires computing self-attention validation loss without the need for any training.
Finally, we show the pareto frontier of student subnetworks generated by several KD and NAS methods in Figure for the MNLI task. The blue points represent all the subnetworks extracted from AutoDis st il and red points denote the optimal ones, all fine-tuned on the MNLI task. We observe the optimal Aut oDi it models outperform several competing methods.
4.1.3 Task-agnostic Training Strategies
We study different task task-agnostic strategies for SuperLM train- ing in AutoDis st itl Specifically, compare three strate- gies in Table 5) (i) We replacing the KD loss in Eqn. 4 with masked language modeling (MLM) loss [ to calculate gradi- ents which is the most widely used task-agnostic pre-training and distillation strategy. (ii) KD..+Coont further continues training the searched compressed models on the large language
Table 5: Comparing task-agnostic SuperLM training strategies.
Strategy
MRPC
RTE
MNLI
MLM
89.4
68.2
82.2
KDau+Comm
91.0
71.8
83.5
KDatt
91.2
71.5
83.2
corpus. (iii) KDatt is the strategy adopted in AutoDistil for self-attention distillation. We evaluate subnetworks with the same architecture (6 layers, 768
hidden, 12 heads, MLP ratio from the trained SuperLM. We fine-tune subnetworks on MRPC, RTE, and MNLI tasks, and report fll, accuracy, and accuracy, respectively. MRPC and RTE are paraphrase detection and natural language inference tasks, respectively. They are low-resource tasks with limited training examples (3 labels each) for fine-tuning, This allows us to evaluate transferability of the compressed models trained on unlabeled general domain (e.g.. Wikipedia) and fine-tuned on downstream tasks with limited labels. MNLI has a larger number of examples
Other NAS methods either use different hardware for training or do report the cost.

---

83.5
$3.0
Accuracy
$2.5
62.6
81.3
81.0-
61.0
0.4
0.6
0.8
10
0.0
0.5
3.0
1.5
2.0
#FLOPs
#11OPs
#FLOPs
(a) Acc vs #FLOPs (SuperLMRaaol (SuperLM (mall)..
Acc
VS
#FLOPs (c) Acc vs #FLOPs (SuperLMrrny)
st
83.0
Accuaty
Accuato
02.5
81.5
$10
81.0
2.5
2.2
2.5
3.0
3.5
0.4
0.5
0.8
1.5
2.4
iPam
PPara
ePara
(d) Acc vs. #Para (SuperLMR.ae).
(e) Acc vs #Para (SmaarL)Msmat).. (1) Acc vs #ara (SuperLMTrin)))
Figure 3: Computational cost task (MNLI) performance trade-off for all 256 subnetworks contained in each of K SuperLMs (base, small and tiny).
( show the trade-off between accuracy (Y-axis) and #FLOPs (X-axis), and 3(d) 3 show the trade-off between accuracy (Y-axis)
((c)
and #Para (X-axis) with optimal compressed AutoDistil student for each SuperI. red.
(400K). Prior work 41 has shown MNLI to transfer well to other tasks in GLUE. First, we observe self-attentio distillation to perform better than MLM, for SuperLM training. Second, observe limited performance gains with continued training of the optimal subnetworks from NAS as done in existing works demonstrating the effectiveness of our single-stage training protocol.
4.1.4 One-shot vs. Few-shot NAS with Varying K
Table 6: Search space design strategies.
For few-shot NAS, we choose K=3 (i.e. 3 sub- spaces) for following reasons: (1) sub-spaces corre- spond to base, small and tiny model sizes; as used in prior work in CV; e.g. AutoFormer
Search Space Size (#subnetworkks) One-shot (K Few-shot (K 3)
Task
17 (2) search- ing over different values of K is resource-extensivg since it requires training K SuperLMs for each choice of K, (3) As K increases, search process becomes
27
864
11232
256*3
MRPC
88.2
87.5
85.1
91.2
RTE
67.2
64.5
62.8
71.8
similar to the undesirable brute-force discrete search that trains all models in search space individually.
To understand the effect of few-shot vs. one-shot NAS, we compare the performance of single space (K 1) multiple sub-spaces (K 3). We extract subnetworks with same architecture (6 layers, 768 hidden, 12 heads, MLP ratio 4) from trained SuperLMs for each strategy for evaluation with results in Table For one-shot NAS, we consider a single search space containing varying number of subnetworks (e.g., 27, 864, 11232). Few-shot NAS contains 256 subnetworks in each partition. We fine-tune subnetworks on RTE and MRPC tasks, and report accuracy and fl respectively. We observe fewer subnetworks in single search space for one-shot NAS result in better performance. This results from optimization interference and gradient conflicts as the number and size of subnetworks increase in the space. Finally, our design strategy performs the best while containing lesser number of subnetworks demonstrating the benefit of few-shot NAS for language model distillation.
4.1.5 Comparing Search Strategies and Optimal Architectures
From Table we observe that the student models AutoDi st Agnostic and AutcDi it Proxys obtained from SuperI by task-agnostic and task-proxy search strategies respectively obtain a similar trade-off between performance and cost. The task- search results in minor performance gain 0.3 over the fully task-agnostic search mechanism. Table / shows the configuration of searched optimal architectures from Aut oDistil with corresponding computational cost. For reference, we also show the architecture of the teacher BERT BASE and a state-of-the-ar distilled model MINILM ! that are -andnenngeeeree. We observe that the obtained architectural
are quite

---

non-standard and difficult to obtain by trial and error considering the large space of Transformer architectures. We also observe that optimal com- pressed models have thin-and-deep structure consistent with findings that thinner and deeper models perform better 126] than wider and shallower ones. While use this an induc- tive bias for sub-space partitioning, our search space (Table 2 also con-
Table 7: Architecture comparison between the optimal com-
pressed students searched by AutoDistil with state-of- the-art hand-engineered students distilled from BERT BASE-
Model
#Layers #Hid Ratio #HHads #FLOPs #Para
BERTBASE MINILM
12
768
12
11.2G
109M
6
768
4
6
5.66G
66.5M
Autooiis Agnostic
11
352
4
10
2.13G 26,8M
AutoDis Proxys
12
544
3
9
4.40G
50.1M
AutoDis Proxys
11
352
4
8
2.02G
26.1M
AutoDis Proxyt
160
3.5
10
0.27G 6.88M
tains diverse subnetworks with differ-
ent depth and width. Non-maximal MIPP ratio and attention heads for optimal compression indicate that self-attention and feed- forward layers of Transformers are overparameterized [21
5 Related Work
Tpsk-specific knowledge distillation. Knowledge distillation (KD)
42
a widely used technique
for model compression, which transfers knowledge from a large teacher to smaller student. Task- specific KD aims to generate smaller students by using downstream task label information. Typical task-specific KD works include BERT-PKD 43 BERTSMALL 40.. TinyBERT (6. DynaBERT [12). and SparseBERT 4 While task-specific KD often achieves good task performance, a typical drawback is that it is resource-consumingg to run KD for each and every task, and also not scalable. Task-agnostic knowledge distillation. In contrast to task-specific KD, we explore task-agnostic KD that does not use any task label information. The distilled task-agnostic models can be re-used by simply fine-tuning on downstream tasks. Task-agnostic KD leverages knowledge from soft target probabilities, hidden states, layer mappings and self-attention distributions of teacher to train student models. Typical aask-agi KD works include DistilBERT [5], MobileBERT 77 and MiniLM [4. MobileBERT assumes that students have the same number of layers as the teacher for layer-by- by-ley distillation. Minil M transfers self-attention knowledge from the last layer of the teacher to that of the student. These works rely on hand-designed architecture for the student models for KD that requires several trials, and needs to he repeated for a new student with a different cost. In contrast, we develop techniques automatically design and distill several student models with variable cost using NAS. Neural Architecture Search. While NAS has been extensively studied in computer vision X g 10 M. there has been relatively less exploration in natural language processing. Evolved Transformer 45 and HAT [Δ] search for efficient sub-networks from the Transformer architecture for machine translation tasks. Some recent approaches closest to our method include, DynaBERT (I2). AutoTinyBERT [B] and NAS-BERS [1]. DynaBERT performs task task-specific distillation. NAS- BERT performs two-stage knowledge distillation with pre-training and fine-tuning of candidates. Similar to above approaches, Auto' Auto'Tiny BnRyBERT' also employs one-shot NAS with a single large search space containing millions of subnetworks that result co-adaption and weight-sharing challenges for Su training. Further it also uses a multi-stage training protocol for further pre-training and distillation of the NAS-generated candidates. In contrast, Aut oDiisssss still employs few-shot NAS with compact search space design with single-stage task-agnostic training protocol. This further allows us to do a lightweight search for the optimal student without re-training. [6 studies few-shot task-specific NAS in the CV domain with CNN architectures, and different design and search spaces from ADtstii More comparisons between them can be found in the Appendix (Section A.2)
6 Conclusion
We develop few-shot task-agnostic NAS method, namely AutoDistii to distil large language models into compressed students with variable computational cost. To address the co-adaption and weight-sharing challenges for SuperLM training. we partition the Transformer search space into K compact sub-spaces covering important architectural components like its depth, width, and attention heads. We leverage self-attention distillation for fully task- agnostic SuperLM training and lightweight optimal search without any re-training. Obtained students can be re-used by simply fine-tuning on tasks. Aut AutoDistil generates students with 3r less computational cost (FLOPs) than tfte-of-the-art task-agnostic KD methods while obtaining similar task performance in GLUE.
10

---

References
[1] Jacob Devlin. Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirec- tional transformers for language understanding. In NAACL, pages 4171-4186, Minncapolis, Minnesota, June 2019. Association for Computational Linguistics.
[2] Tom Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah. Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Fric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Bemer, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[3] Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in NI.P. In ACI, pages 3645-3650, Florence, Italy, July 2019. Association for Computational Linguistics.
[4] Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep sclf-attention distillation for task-agnostic compression of pre-trained transformers. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 5776-5788. Curran Associates, Inc., 2020.
[5] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv: 108, 2019.
[6] Xiaoqi Jiao, Yichun Yin, Lifeng Shang. Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings, pages 4163 4174, 2020.
[7] Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobilebert: compact task-agnostic bert for resource-limited devices. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2158-2170, 2020.
[8] Hieu Pham, Melody Guan, Barret Zoph. Quoc Le, and Jeff Dean. Efficient neural architecture search via parameters sharing, In International Conference on Muchine Learning, pages 4095 4194-4104. PMLR, 2018.
[9] Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, and Quoe V
Le. Mnasnet: Platform-aware neural architecture search mobile. Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 2820-2828, 2019.
[10] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once for all: Train one network and specialize it for efficient deployment. In International Conference on Learning Representations, 2020.
[11] Jiahui Yu, Pengchong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang, and Quoc Le. Bignas: Scaling up neural architecture search with big single-stage models. In European Conference Computer Vision, pages 702-717. Springer, 2020.
[12] Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang. Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 9782 9793, Curran Associates, Inc., 2020.
[13] Yichun Yin, Cheng Chen, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. BERRT: Automatic hyper-parameter optimization for efficient pre-trained language models. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the I 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 5146-5157, 5157. Association for Computational Linguistics, August 2021.
[14] Jin Xu, Xu Tan, Rengian Luo, Kaitao Song, Jian Li, Tao Qin, and Tie-Yan Liu. NAS-BERT: task-agnostic and adaptive-size BERT compression with neural architecture search. In Feida Zhu, Beng Chin Ooi, and Chunyan Miao, editors, KDD '11: The 27th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Virtual Event, Singapore, August 14-18, 2021, pages 1933-1943. ACM, 2021.
[15] Gabriel Bender, Jan Kindermans, Barret Zoph, Vijay Vasudevan, and Quoc Le. Understanding and simplifying one-shot architecture search. In International Conference on Machine Learning, pages 550-559. PMLR, 2018.
11

---

[16] Yiyang Zhao, Linnan Wang, Yuandong Tian, Rodrigo Fonseca, and Tian Guo. Few-shot neural architecture search. In International Conference Machine Learning, pages 12707-12718. PMLR, 2021.
[17] Minghao Chen, Houwen Peng, Jianlong Fu, and Haibin Ling. Autoformer: Searching transformers for visual recognition. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 12270-12280, 2021.
[18] Tianyi Chen, Bo Ji, Tianyu Ding, Biyi Fang, Guanyi Wang. Zhihui Zhu, Luming Liang, Yixin Shi, Sheng Yi, and Xiao Tu, Only train once: A one-shot neural network training and pruning framework. Advances in Neural Information Processing Systems, 34, 2021.
[19] Jing Xu, Yu Pan, Xinglin Pan, Steven Hoi, Zhang Yi. and Zenglin Xu. Regnet: Self-regulated network for image classification. IEEE Trunsactions on Neural Networks and Learning Systems, 2022.
[20] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
[21] Paul Michel, Omer Lovy, and Graham Neubig. Are sixteen heads really better than onc? In NeurlPS, pages 14014-14024, 2019.
[22] 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 pruned. arXiv preprint KKiv:1905,09411 2019.
[23] Prakhar Ganesh, Yao Chen, Xin Lou, Mohammad Ali Khan, Yin Yang, Deming Chen, Marianne Winslett, Hassan Sajjad, and Preslav Nakov. Compressing large-scale transtormer-based models: A case study on bert. arXiv preprint arXiv:2002. 19855, 2020.
[24] Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really hetter than one? In H. Wallach, H. Larochelle, A. Beygelzimer, F. l'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[25] 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. Proceedings the 57th Annual Meeting of the Association for Computational Linguistics, pages 5797-5808, Florence, Italy, July 2019. Association for Computational Linguistics.
[26] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitett:: Hints for thin deep ncts. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[27] Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez. Train big, then compress: Rethinking model size for efficient training and inference transformers. In ICML, pages 5958-5968. PMLR, 2020.
[28] Jiahui Yu and Thomas Huang. Universally slimmable networks and improved training techniques. In Proceedings of the IEEF/CVE international conference on computer vision, pages 1803-1811, 1, 2019.
[29] Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long Papers), pages 1112-1122, New Orleans, Louisiana, June 2018 Association for Computational Linguistics.
[30] Alex Wang. Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy. and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics.
[31] Alex Warstadt, Amanpreet Singh, and Samuel R. Bowman. Neural network acceptability judgments, 2018.
[32] Richard Socher et al. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 1631-1642, Seattle, Washington, USA, October 2013, Association for Computational Linguistics.
[33] William B Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005), 2005.
12

---

[34] Zihan Chen, Hongbo Zhang, Xiaoji Zhang, and Leqi Zhao. Quora question pairs. URL https:///ww. kaggle. comlclquora-quesiion-nuie 2018.
[35] Pranav Rajpurkar, Jian Zhang. Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for 2016.
machine comprehension of text. arXiv preprint
[36] Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Machine Learning Challenges Workshop, pages 177-190. Springer, 2005.
[37] R Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Szpektor. The second pascal recognising textual entailment challenge. In Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment, 2006.
[38] Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and William Dolan. The third pascal recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing, pages 1 1-9, 2007.
[39] Luisa Bentivogli, Peter Clark, Ido Dagan, and Danilo Giampiccolo. The fifth pascal recognizing textual entailment challenge. In TAC. 2009.
[40] Tulia Ture, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint (rXXv:1908.08960 2019.
[41] Tianlong Chen, Jonathan Franklc, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis pre-trained bert networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 15834-15846. Curran Associates, Inc., 2020.
[42] Geoffrey Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in neural network. In NIPS Deep Learning and Representation Learning Workshop. 2015.
[43] Siqi Sun, Yu Cheng, The Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. In Proceedings of the 2019 Conference on Empirical Methods in Natural Longuage Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNIPP) pages 4323-433 4332,, 2019.
[44] Dongkuan Xu, Ian EH Yen, Jinxi Zhao, and Zhibin Xiao. Rethinking network pruning- under the pre-train and fine- tune paradigm. In Proceedings of the Human Language Technology Conference of the NAACL, 2021.
[45] David So. Quoc Le. and Chen Liang. The evolved transformer. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Larning Research, pages 5877-5886. PMLR, 09-15 Jun 2019.
[46] Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai. Ligeng Zhu, Chuang Gan, and Song Han. Hat: awree transformers for efficient natural language processing. In Annual Conference of the Association for Computational Linguistics, 2020.
[47] Hassan Sajjad, Fahim Dalvi, Nadir Durrani, and Preslav Nakov. On the effect of dropping layers of wre-trained transformer models. arXiv preprint arXiv:2004.00 038444, 2020.
[48] John Wieting and Kevin Gimpel. ParaNMT-5ON Pushing the limits of paraphrastic sentence embeddings with millions of machine translations. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguisties (Volume I: Long Papers), pages 151-662, Melbourne, Australia, July 2018. Association for Computational Linguistics.
[49] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story- -tiry-like visual explanations by watching movies and reading books. In Proceedingy of the IEEE international conference on computer vision, pages 19-27, 2015.
[50] Jianping Gou, Baosheng Yu, Stephen Maybank, and Dacheng Tao. Knowledge distillation: A survey. International Journal of Computer Vision, 1899111188911819 2021.
[51] Ting Chen, Simon Komblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML. pages 1597-1607. PMLR, 2020.
13


Reviews:
[{"review_id": "fsji_N7FpBx", "reviewer": "Reviewer_1T87", "summary": "This work partitions the search space of neural architecture search into K sub-spaces, trains a SuperLM for each partition with task-agnostic self-attention distillation, and finds optimal students without additional training.", "questions": "How you define the search space in Table 1 is not straightforward, though the authors mentioned it is based on inductive bias and heuristic. Could you elaborate on the heuristic part? The prior work mentioned in Line 129 is not evaluated by Transformer architectures, so I doubt inductive bias from this work is valid to apply.\n\nSuperLM in a smaller region is also initialized with pre-trained BERT-Base, so only the leftmost layers, dimensions, and heads are used. However, dimensions and heads are order-independent, meaning that the performance might change after re-ordering.\n\nSection 4.1.3 and Section 4.1.4 are only tested on MRPC and RTE datasets. I am suspicious they are enough or good representatives.\n\nAutoDistil-proxy only uses MNLI as a source task. I am curious how the transferability of optimal architectures changes between source and target tasks?", "limitations": "I wonder whether the few-shot task-agnostic NAS method is applicable to other domains, including computer vision.\n", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed methods outperform traditional knowledge distillation baseline and other NAS baselines. The idea of partitioning search space by using multiple supernets is intuitive and practically useful. Other implementations are borrowed from the existing methods like BigNAS and MiniLM. The paper provides adequate ablation study results to confirm the effectiveness of each component (e.g., training strategies, few-shot NAS, search strategy).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Q2LJT-96f8I", "reviewer": "Reviewer_ov4n", "summary": "The paper preposes AutoDistill, a recipe for distilling large language models into smaller sizes when a budget is given. \n\nAutoDistill has three major steps: \n* designing the search space with 4 hyperparameters, and partition the large search space into three smaller ones (K=3; tiny, small, base).\n* training a SuperLM in each partition, with a task-agnostic objective similar to MiniLM (Wang et al., 2020)\n* obtain all possible student architectures by pruning the SuperLM, and find the best one by enumerating the model and computing the loss on either validation data or task-specific data\n\nAutoDistill achieves superior performance on GLUE benchmark when compared to baselines such as (a) fine-tuning, (b) task-agnostic distillation baselines: DistilBERT, TinyBERT, and (c) NAS baselines: NAS-BERT. Meanwhile AutoDistill is more computationally efficient as demonstrated in Table 3.\n", "questions": "(Disclaimer: I know little about NAS or one-shot/few-shot NAS. My ratings are based on my current understanding and I hope the authors can clarify the following for me)\n\n* Does \"few-shot\" in the method mainly mean splitting the search space into K=3 subspaces?\n* What is the major difference between AutoDistill and Reference [16] (Zhao et al., 2021)? \n* The motivation of few-shot NAS is to alleviate conflicts in weight-sharing (Line 161) and mitigate gradient conflicts (Line 90). As this is one contribution highlighted by the authors, I believe some ablation analysis or discussion on not using the 3 subspaces design should be included.\n* In figure 1 and in Line 251, how do you quantify \"best trade-off\"? What is the metric used to determine this?\n* Following on the previous question, in figure 1, since MNLI performance is the criteria and also the y-axis, why is the selected models (red) not always the best performing models?\n\nMinor Questions:\n* How large is the validation set used in Task-agnostic search (Line 206)? It is quite surprising to me that evaluating the 256 candidate subnetworks on this set (or on MNLI) takes less than one hour.\n* Line 300. Not clear about the MLM baseline. Does it mean the SuperLM is directly trained on MLM objective? Is the \"random sample a subnetwork\" step still used?", "limitations": "Limitations are discussed in Appendix E.", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n* AutoDistill can be practically useful. It achieves 3x compression over other distillation methods. The resulting models can be re-used by practitioners and researchers.\n* Solid experiment and clear visualization. \n* The authors have uploaded code with running instructions to ensure reproducibility.\n\nWeakness:\n* The novelty of AutoDistill is still unclear. Components of AutoDistill seem to be adapted from existing methods.\n* Writing is sometimes confusing and not coherent. See questions below.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "MedPeHY2XuZ", "reviewer": "Reviewer_swEZ", "summary": "This paper proposes an effective method AutoDistill to compress several student models with various computation costs at once. By combining few-shot NAS and task-agnostic knowledge distillation, AutoDistill alleviates optimization interference and speed-up the training process, enabling single-stage training without further pre-training and distillation. Experimental results demonstrate that AutoDistill outperforms strong NAS and knowledge distillation baselines.", "questions": "1. How transferable is the optimal student architecture derived from the proposed AutoDistill method across teachers and knowledge distillation objectives? \n2. Given some optimal student architecture derived by NAS, how would traditional task-agnostic knowledge distillation method (e.g., MiniLMv2) perform?", "limitations": "The authors address the limitations adequately.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper is clearly written and well structured. The authors first demonstrate the weaknesses of existing knowledge distillation and NAS-based language model compression methods, then they introduce their method clearly, making it easy for readers to follow.\n2. The proposed method is reasonable and technically sound. By adopting the idea of few-shot NAS, the problem of optimization interference is alleviated. Moreover, by fully task-agnostic knowledge distillation and search, AutoDistill significantly reduces training costs.\n3. The authors conduct comprehensive experiments and the results verify the superiority of AutoDistill over strong baseline methods.\n\nWeaknesses:\n1. Although the proposed method is effective, the major two components of the method (i.e., few-shot NAS and task-agnostic knowledge distillation) are not novel. When compared with DynaBERT, it seems that AutoDistill is its task-agnostic version with a larger searching space and a few other modifications.\n2. When compared with existing NAS-based language model compression method (e.g., AutoTinyBERT), AutoDistill enjoys much less computation cost in training and searching. However, when compared with traditional task-agnostic knowledge distillation method (e.g., MiniLMv2), the advantage of training speed is gone.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept