Few-shot Task-agnostic Neural Architecture Search for Distilling Large Language Models
Dey
Abstract
Traditional knowledge distillation (KD) methods manually design student architectures to compress large models given pre-speciﬁed computational cost.
This requires several trials to ﬁnd 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-speciﬁc 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 ﬁne-tuning on any downstream task.
Experiments on GLUE benchmark demonstrate AutoDistil to outperform state-of-the-art KD and NAS methods with upto 41 x reduction in computational cost.
Code and models are available at aka.ms/autodistil.
Introduction
While large pre-trained language models (e.g., BERT [1], GPT-3 [2]) are effective, their huge size poses signiﬁcant challenges for downstream applications in terms of energy consumption and cost of inference [3] limiting their usage in on the edge scenarios and under constrained computational inference budgets.
Knowledge distillation [4, 5, 6, 7] has shown strong results in compressing pretrained language models into small student models.
However, these works require pre-speciﬁcation of the student architecture and computational cost (e.g., number of parameters, FLOPs) for distillation.
This poses two signiﬁcant challenges: (i) it requires several trials to come up with viable architectures as they are hand-engineered and to deﬁne several hyper-parameters (e.g., number of layers and attention heads, hidden dimension, etc.)
; (ii) one has to re-run distillation with any change in speciﬁcation for the student architecture or computational cost for using it in a target environment.
Neural Architecture Search (NAS) [8, 9, 10, 11] provides a natural solution to automatically search through a large space of candidate models.
The dominant NAS paradigm consists of two main steps: (a) Training a Super model combining all possible architectures into a single graph and jointly training them via weight-sharing; (b) Searching for optimal architecture from Super model with best accuracy on a downstream task, satisfying user-speciﬁed latency constraint for target device.
Parallel to above computer vision (CV) works, NAS has shown strong results in recent works like DynaBERT [12],
AutoTinyBERT [13] and NAS-BERT [14] for natural language understanding (NLU).
Drawbacks of existing NAS methods.
[D1: Co-adaptation in weight-sharing]
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 conﬂicts in weight-sharing where bigger student models converge faster in contrast to smaller ones converging slower [16, 11].
[D2: Multi-stage training]
A single SuperLM may not have sufﬁcient capacity to encode a large search space.
As a result, these works use multi-stage training process, where they ﬁrst conduct
NAS to identify candidate students and then perform further pre-training [13] and knowledge distillation [14] of the candidates.
[D3: Task-speciﬁc training] NAS works in CV domain (e.g AutoFormer [17], Once-for-all [10],
One-Shot NAS [11, 18]) leverage hard class labels from a task (e.g., image classiﬁcation) or soft labels from ImageNet pre-trained models (e.g., MobileNet [7], RegNet [19]) for task-speciﬁc optimization with accuracy as evaluation metric.
Different from CV domain, NLU tasks have different objectives and evaluation metrics for classiﬁcation (e.g., MNLI), regression (e.g., STS-B) and correlation (e.g., CoLA).
This makes it challenging to adapt existing NAS works to the NLU domain in a task-agnostic setting.
Recent NAS works in the NLU domain are not fully task-agnostic.
For instance,
DynaBERT [12] accesses both task labels for knowledge distillation and task development set for network rewiring.
NAS-BERT [14] performs two-stage knowledge distillation with pre-training and ﬁne-tuning of the candidates.
While AutoTinyBERT [13] 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 a detailed discussion on the differences between AutoDistil and existing KD and NAS methods based on ﬁne-grained search space, different training strategies and amortized training cost in Table 1.
More comparison details are summarized in Section A of Appendix.
Contributions.
We address above challenges with fully task-agnostic few-shot NAS consisting of three steps.
(S1) Search space design.
We partition the Transformer search space into K sub-spaces considering important architectural hyper-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
[Table]
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 (e.g., BERT) to the student and use weight-sharing to train the SuperLM.
(S3)
Lightweight optimal student search.
We obtain optimal student(s) directly from well-trained SuperLM(s) without any re-training that can be simply ﬁne-tuned on downstream tasks.
Our contributions over existing NAS works can be summarized as:
• In contrast to prior works (e
g., DynaBERT, AutoTinyBERT, NAS-BERT), we do a single-stage training combining NAS and distillation with no further pre-training or augmentation and demonstrate superior performance of the NAS process itself with signiﬁcantly reduced training cost.
Obtained subnetworks are simply ﬁne-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 conﬂicts in SuperNet training compared to prior One-shot NAS works in NLU (e.g., DynaBERT, AutoTinyBERT, NAS-BERT).
AutoFormer in the CV domain is an exception to 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 3 x 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 selfattention (MHA) and feed-forward network (FFN).
Transformer layers are stacked to encode contextual information for input tokens as: X l = Transformer l ( X l  1 ) , l 2 [1 , L ] where L is the number of Transformer layers, X l 2 R s ⇤ d hid , s is the sentence length, and d hid is the hidden dimension.
In the following, we omit the layer indices for simplicity.
Multi-Head Self-Attention (MHA).
Given previous Transformer layer’s output X , MHA computes:
Attention
[Equation]
[Equation] [Equation] where W Qh , W Kh , W Vh 2 R d hid ⇤ d head , W O 2 R d hid ⇤ d hid are linear transformations.
Q h , K h , V h 2 R s ⇤ d head are called queries, keys, and values, respectively.
H is the number of heads.
head h =
Attention( Q h , K h , V h ) denotes the h -th attention head.
Concat is the concatenating operation.
d head =
d hid /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 a ReLU activation as:
[Equation]
Search Engine
TaskAgnostic TaskProxy ! ! ∗
Performance Cost # !
K -shot Learning for SuperLM
SubnetworkSearch
K Search Partitions …
Transformer Transformer
Transformer Transformer Transformer
Transformer T x 1 x 2 x…
Transformer
Transformer
Transformerx 1 x 2 x
Transformer S
Q
K
V
Distill Q
K
V
Attention
W Q
W K
W V
W o
W W
Sampled Subnetworks
Pre-trained Teacher
Adaptive Transformer Block ofSampled Subnetworks
Multi-head Self-attention Relation Distillation where W 1 2 R d hid ⇤ d f , W 2 2 R d f ⇤ d hid , b 1 2 R d f , and b 2 2 R d hid .
In addition, there are residual connection and layer normalization on top of MHA and FFN (denoted by  in Figure 2), which are formulated as LayerNorm(x + MHA(x)) and LayerNorm(x + FFN(x)), respectively.
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 to include: (1) Feed-forward network (FFN) dimension - we encode this by the MLP (multi-layer perceptron) ratio deﬁned as r =
d f d hid , with d f and d hid representing the intermediate dimension of the FFN and hidden dimension respectively; (
2) Number of layers ( L ) to capture the network depth; (3) Hidden dimension ( d hid ) to encode input; (4) Attention heads ( H ) for multi-head self-attention
All of the above factors are important for model capacity and have a signiﬁcant impact on 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 is one of the most computation intensive components [23].
Therefore, we 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, we make the number of attention heads searchable as well.
Inductive bias.
Prior works on CNNs [26] and Transformers [27] demonstrate that thinner and deeper neural networks with improved representation capacity perform better than wider and shallower ones.
We incorporate this as an inductive bias to decide the number of layers to consider for the students in each of 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 a 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 [13].
Search space partition.
Existing works [13, 14] train a single large SuperLM 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 [11].
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 conﬁgurations 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 subnetwork 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 to a 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 conﬂicts in weight-sharing.
[Table]
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.
We extract student subnetworks from the SuperLM by a simple truncation strategy like bottom-left extraction.
In particular, given a speciﬁc architecture ↵ = { l, d hid , r, h } , (i) we ﬁrst extract alternate l Transformer layers from the SuperLM; (ii) then extract bottom-left sub-matrices in terms of d hid and r from the original matrices that represent the hidden dimension and the MLP ratio respectively; (iii) ﬁnally, for the attention heads, we extract the leftmost h heads and retain the dimension of each head as the SuperLM.
This strategy is used during sampling sub-networks for SuperLM training via weight sharing after which the sub-network weights are 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 Task-agnostic SuperLM Training
We illustrate SuperLM training process in Algorithm 1.
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 SuperLM 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 SuperLM.
During sampling, we employ Sandwich rule [28], also used in BigNAS [11], that samples the smallest subnetwork, the largest subnetwork and M random ones for updating SuperLM.
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 [4] for task-agnostic training.
To this end, we employ multi-head 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
A 1 , A 2 , A 3 to denote the queries, keys and values of multiple relation heads of teacher model, and B 1 , B 2 , B 3 respectively for a sampled subnetwork.
Mean squared error ( MSE( · ) ) between multi-head self-attention relation of teacher and sampled subnetwork is used as distillation objective:
[Equation]
where R Ti = softmax( A i A > i / p d k ) , R Si = softmax( B i B > i / p d k ) , H is the number of attention heads; R Ti represents the teacher’s Q  Q , K  K , or V  V relation; R Si represents the same for student.
R Tik is the relation information based on one attention head, and d k is the attention head size.
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 ﬁrst concatenate the selfattention 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.
Input: Partitioned K [Table] A k ; initialized K [Table]
- perLMs S k on A k ; [Table] [Table]
[Table] ; sampling steps M Output: Trained SuperLMs { S k } for k = 1 to K do for i = 1 to E do Get a [Table]
of data from D for batch in D
do Clear gradients in [Table] S k for m = 1 to M do Randomly sample a [Table]
s from S k Calculate self-attention distil.
loss between subnetwork s and [Table] T with Eqn.
Accumulate gradients end for Update S k with the accumulated gradients end for end for end for
The SuperLM for sub-space
A k is trained as:
[Equation]
where, K is the number of sub-space partitions; W are the weights of the SuperLM;
W ↵ are the weights in W speciﬁed by the architecture ↵ ; U are the weights of the teacher model including the self-attention module used for distillation; D train is the training data set, and L ( · ) is the self-attention loss function from Eqn.
3.3 Lightweight Optimal Student Search
We outline two search strategies for selecting the optimal student subnetwork.
Task-agnostic search.
We adopt this to be our primary strategy to compare against all baselines since it does not access any task label information.
We compute the task-agnostic self-attention distillation loss for all student subnetworks using Eqn.
(4) on a heldout validation set from the unlabeled training corpus.
The student subnetworks are directly obtained by bottom-left extraction from the well-trained SuperLM (outlined in
Section 3.1).
This process is 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.
[Equation]
where W ⇤ ↵ is the weights of architecture ↵ obtained from W ⇤ A k , D val is the validation data set, L is the self-attention distillation loss, and g ( · ) is a function to calculate the computational cost (e.g., #FLOPs, #parameters) of the subnetwork subject to a given user-speciﬁed resource constraint c .
Task-proxy search.
We compare our task-agnostic search against another strategy that considers a proxy task (e.g., MNLI [29]) with label information to ﬁne-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) benchmark [30].
We compare our method with the baseline methods on two single-sentence classiﬁcation tasks (CoLA [31], SST-2 [32]), two similarity and paraphrase tasks (MRPC [33], QQP [34]), and three inference tasks (MNLI [29], QNLI [35], RTE [36, 37, 38, 39]) 3 .
We report accuracy for MNLI,
QNLI, QQP, SST-
Baselines.
We compare against several task-agnostic methods 4 generating compressed models from
BERT base teacher, using (i) knowledge distillation like BERT SMALL [40], Truncated BERT [29],
DistilBERT [5], TinyBERT [6], MINILM [4]; as well as those based on Neural Architecture Search, like AutoTinyBERT [13], DynaBERT [12], and NAS-BERT [14].
AutoDistil conﬁguration.
We use uncased BERT BASE as the teacher consisting of 12 Transformer layers, 12 attention heads; with the hidden dimension and MLP ratio being 768 and 4, respectively.
It consists of 109 M parameters with 11 .
2 G FLOPs.
We use English Wikipedia and
BookCorpus data for SuperLM training with WordPiece tokenization.
We use a batch size ofand 4 e - 5 as the peak learning rate for 10 epochs.
The maximum sequence length is set to 128 .
The coefﬁcients in distillation objective (Eqn.
(4)),  1 ,  2 , and  3 , 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 BASE .
Other hyper-parameter settings are shown in Appendix.
We use 16 V 100 GPUs to train the SuperLM with 336 GPU-hours as the training cost.
4.1 Finding the Optimal Compressed Models
AutoDistil Agnostic is obtained by fully task-agnostic training and task-agnostic search without using any task label information.
We set a constraint in Eqn.
(6) 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.
AutoDistil Proxy uses MNLI [29] as a 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, we ﬁne-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
AutoDistil Proxy B , AutoDistil Proxy S and AutoDistil Proxy T obtained from the corresponding sub-spaces of SuperLM Base , SuperLM Small and SuperLM Tiny , respectively.
Notably all students are obtained from the AutoDistil SuperLM still trained in a fully task-agnostic 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
BERT BASE in Table 3 with respect to the following measures: computational cost in the form of (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 AutoDistil Agnostic generated via our task-agnostic SuperLM training leads to upto 3 x reduction in FLOPs over state-of-the-art distilled models (e.g., MINILM [4], TinyBERT [6], DistilBERT [5]) that are hand-engineered while matching the overall task performance.
The most aggressive compressed version corresponding to
AutoDistil Proxy T obtains a massive 41 x reduction in FLOPs over BERT BASE 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 classiﬁcation.
This depicts an interesting impact of massive model compression on varying task types.
4.1.2 Comparison with Neural Architecture Search Baselines
[Table]
[Table]
AutoTiny- BERT does not report the cost of SuperNet training - typically the most expensive step.
Further Training refers to additional pre-training applied to NAS-generated candidates.
We report the performance of several NASgenerated student models of comparable FLOPs and parameters from corresponding papers in
Table 3.
AutoDistil outperforms all competing methods on aggregate for all sizes; except for small-sized model; where it has marginally lower performance (0.1 points on avg) compared to AutoTinyBERT.
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 use multi-stage training; where additional pre-training and distillation is applied to NAS-generated candidates.
To better understand the impact of single-stage vs. multi-stage methods on the training cost,
we
[Table]
compare the overall cost of NAS for AutoDistil and that reported in AutoTinyBERT 5 for the small model segment in Table 4.
AutoDistil is much cheaper due to its single-stage training protocol; where no additional pre-training or distillation is needed.
It is worth noting that the overall SuperNet training cost of AutoDistil (the most expensive component of NAS) is less or comparable to the additional training cost of re-training candidate models for AutoTinyBERT.
Note that AutoTinyBERT does not report their SuperNet training cost.
Additionally, AutoDistil has a much faster search mechanism due to (1) inductive biases built into the search space deﬁnition 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 3 for the MNLI task.
The blue points represent all the subnetworks extracted from
AutoDistil and red points denote the optimal ones, all ﬁne-tuned on the MNLI task.
We observe the optimal AutoDistil models to outperform several competing methods.
4.1.3 Task-agnostic Training Strategies
[Table]
We study different task-agnostic strategies for SuperLM training in AutoDistil .
Speciﬁcally, we compare three strategies in Table 5.
(i)
We replacing the KD loss in Eqn.
(4) with masked language modeling (MLM) loss [1] to calculate gradients which is the most widely used task-agnostic pre-training and distillation strategy.
(ii) KD att
+Cont further continues training the searched compressed models on the large language corpus.
(iii) KD att is the strategy adopted in AutoDistil for self-attention distillation.
We evaluate subnetworks with the same architecture ( 6 layers,hidden, 12 heads, MLP ratio 4 ) from the trained SuperLM.
We ﬁne-tune subnetworks on MRPC,
RTE, and MNLI tasks, and report f1, 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 K labels each) for ﬁne-tuning.
This allows us to evaluate transferability of the compressed models trained on unlabeled general domain (e.g., Wikipedia) and ﬁne-tuned on downstream tasks with limited labels.
MNLI has a larger number of
[Figure]
( 400 K ).
Prior work [41] has shown MNLI to transfer well to other tasks in GLUE.
First, we observe self-attention distillation to perform better than MLM, for SuperLM training.
Second, we 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]
For few-shot NAS, we choose K = 3 (i.e. 3 subspaces) for following reasons: (1) 3 sub-spaces correspond to base, small and tiny model sizes; as used in prior work in CV; e.g. AutoFormer [17], (2) searching over different values of K is resource-extensive since it requires training K SuperLMs for each choice of K , (3) As K increases, search process becomes 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 a single space ( K = 1 ) to 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 6.
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 ﬁne-tune subnetworks on RTE and MRPC tasks, and report accuracy and f1 respectively.
We observe fewer subnetworks in a single search space for one-shot NAS result in better performance.
This results from optimization interference and gradient conﬂicts 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 beneﬁt of few-shot NAS for language model distillation .
4.1.5 Comparing Search Strategies and Optimal Architectures
From Table 3, we observe that the student models AutoDistil Agnostic and AutoDistil Proxy S obtained from SuperLM small by task-agnostic and task-proxy search strategies respectively obtain a similar trade-off between performance and cost.
The task-proxy search results in a minor performance gain 0 .
3 over the fully task-agnostic search mechanism.
Table 7 shows the conﬁguration of searched optimal architectures from AutoDistil with corresponding computational cost.
For reference, we also show the architecture of the teacher BERT BASE and a state-of-the-art distilled model MINILM [4] that are hand-engineered.
We observe that the obtained architectural hyper-parameters are quite
[Table]
non-standard and difﬁcult to by trial and
We also observe that optimal compressed models have thin-and-deep structure consistent with ﬁndings that thinner and deeper models perform better [26] than wider and shallower ones.
While we use this as an inductive bias for sub-space partitioning, our search space (Table 2) also contains diverse subnetworks with different depth and width.
Non-maximal MLP ratio and attention heads for optimal compression indicate that self-attention and feed-forward layers of Transformers are overparameterized [21, 22].
Related Work
Task-speciﬁc knowledge distillation.
Knowledge distillation (KD)
[42] is a widely used technique for model compression, which transfers knowledge from a large teacher to a smaller student.
Taskspeciﬁc KD aims to generate smaller students by using downstream task label information.
Typical task-speciﬁc KD works include BERT-PKD [43], BERT SMALL [40], TinyBERT [6], DynaBERT [12], and SparseBERT [44].
While task-speciﬁc KD often achieves good task performance, a typical drawback is that it is resource-consuming to run KD for each and every task, and also not scalable.
Task-agnostic knowledge distillation.
In contrast to task-speciﬁc 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 ﬁne-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 task-agnostic KD works include DistilBERT [5], MobileBERT [7], and MiniLM [4].
MobileBERT assumes that students have the same number of layers as the teacher for layer-by-layer distillation.
MiniLM 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 be repeated for a new student with a different cost.
In contrast, we develop techniques to automatically design and distill several student models with variable cost using NAS.
Neural Architecture Search.
While NAS has been extensively studied in computer vision [8, 9, 10, 11], there has been relatively less exploration in natural language processing.
Evolved
Transformer [45] and HAT [46] search for efﬁcient sub-networks from the Transformer architecture for machine translation tasks.
Some recent approaches closest to our method include, DynaBERT [12],
AutoTinyBERT [13] and NAS-BERT [14].
DynaBERT performs task-speciﬁc distillation.
NASBERT performs two-stage knowledge distillation with pre-training and ﬁne-tuning of candidates.
Similar to above approaches, AutoTinyBERT also employs one-shot NAS with a single large search space containing millions of subnetworks that result in co-adaption and weight-sharing challenges for SuperLM training.
Further it also uses a multi-stage training protocol for further pre-training and distillation of the NAS-generated candidates.
In contrast, AutoDistil employs few-shot NAS with a compact search space design with a single-stage task-agnostic training protocol.
This further allows us to do a lightweight search for the optimal student without re-training.
[16] studies few-shot task-speciﬁc NAS in the CV domain with CNN architectures, and different design and search spaces from AutoDistil .
More comparisons between them can be found in the Appendix (Section A.2).
Conclusion
We develop a few-shot task-agnostic NAS method, namely AutoDistil 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 ﬁne-tuning on tasks.
AutoDistil generates students with 3 x less computational cost (FLOPs) than state-of-the-art task-agnostic KD methods while obtaining a similar task performance in GLUE.