Few-shot Task-agnostic Neural Architecture Search for Distilling Large Language Models
Anonymous Author(s)
AffiliationAddress
Abstract
Traditional knowledge distillation (KD) methods manually design student architectures to compress large models given pre-specified computational cost.
Thisrequires several trials to find viable students, and repeating the process with changein 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 ofdifferent sizes.
Additionally, many of these works are task-specific requiring tasklabels for SuperLM training.
Our framework AutoDistil addresses above challenges with the following steps: (a) Incorporates inductive bias and heuristics topartition Transformer search space into K compact sub-spaces (e.g., K = 3 cangenerate typical student sizes of base, small and tiny); (b) Trains one SuperLM foreach sub-space using task-agnostic objective (e.g., self-attention distillation) withweight-sharing of students; (c) Lightweight search for the optimal student withoutre-training.
Task-agnostic training and search allow students to be reused for finetuning on any downstream task.
Experiments on GLUE benchmark demonstrate
AutoDistil to outperform state-of-the-art KD and NAS methods with upto 3 xadditional reduction in computational cost and negligible loss in task performance.
Introduction
While large pre-trained language models (e.g., BERT [1], GPT-3 [2]) are effective, their huge sizeposes significant challenges for downstream applications in terms of energy consumption and costof inference [3] limiting their usage in on the edge scenarios and under constrained computationalinference 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-specification ofthe student architecture and computational cost (e.g., number of parameters, FLOPs) for distillation.
This poses two significant challenges: (i) it requires several trials to come up with viable architecturesas they are hand-engineered and to define several hyper-parameters (e.g., number of layers andattention heads, hidden dimension, etc.)
; (ii) one has to re-run distillation with any change inspecification 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 searchthrough 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 trainingthem via weight-sharing; (b) Searching for optimal architecture from Super model with best accuracyon a downstream task, satisfying user-specified latency constraint for target device.
Parallel to abovecomputer 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 workstrain one single large Super Language Model (SuperLM) consisting of millions of diverse student architectures.
This results in some undesirable effects of coadaptation [15] like conflicts in weight-sharing where bigger student models converge faster in contrast to smallerones converging slower [16, 11].[D2: Multi-stage training]
A single SuperLM may nothave sufficient capacity to encode a large search space.
As a result, these works use multi-stage training process,where they first conduct NAS to identify candidate students and then perform further pre-training [13] and knowledge distillation [14] of the candidates.[D3: Task-specific 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 classification) or soft labels from ImageNet pretrained models (e.g., MobileNet [7], RegNet [19]) fortask-specific optimization with accuracy as evaluation metric.
Different from CV domain, NLU tasks have different objectives and evaluation metrics for classification(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 taskagnostic.
For instance, DynaBERT [12] accesses both tasklabels for knowledge distillation and task development setfor network rewiring.
NAS-BERT [14] performs two-stage knowledge distillation with pre-trainingand fine-tuning of the candidates.
While AutoTinyBERT [13] also explores task-agnostic training, wedemonstrate better performance from few-shot NAS and much cheaper cost from single stage trainingwithout additional pre-training and distillation.
We present a detailed discussion on the differencesbetween AutoDistil and existing KD and NAS methods based on fine-grained search space,different training strategies and amortized training cost in Section A of Appendix.
Contributions.
We address above challenges with fully task-agnostic few-shot NAS consisting ofthree steps.
(S1) Search space design.
We partition the Transformer search space into K sub-spacesconsidering important architectural hyper-parameters like the network depth, width and attentionheads.
We further leverage inductive bias and heuristics to limit the number of student architecturesin each sub-space.
(S2) Fully task-agnostic SuperLM training.
We train K
SuperLM overall, onefor every sub-space.
This allows each SuperLM more capacity to encode a sub-space as opposed to asingle large one.
We train each SuperLM with a fully task-agnostic objective (without accessing anytask labels) like deep self-attention distillation, where we transfer knowledge from the self-attentionmodule (including keys, queries and values) of a pre-trained teacher (e.g., BERT) to the student anduse weight-sharing to train the SuperLM.
(S3)
Lightweight optimal student search.
We obtainoptimal student(s) directly from well-trained SuperLM(s) without any re-training that can be simplyfine-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-stagetraining combining NAS and distillation with no further pre-training or augmentation and demonstratesuperior performance of the NAS process itself with significantly reduced training cost.
Obtainedsubnetworks are simply fine-tuned on downstream tasks.• Fully task-agnostic training with subnetwork attention state alignment for self-attention relationdistillation 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 NASworks in NLU (e.g., DynaBERT, AutoTinyBERT, NAS-BERT).
AutoFormer in the CV domain is anexception 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 compressionover best performing distillation technique with negligible drop in task performance.
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
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 ∈ [1 , L ] where L is the numberof Transformer layers, X l ∈ R s ∗ d hid , s is the sentence length, and d hid is the hidden dimension.
Inthe following, we omit the layer indices for simplicity.
Multi-Head Self-Attention (MHA).
Given previous Transformer layer’s output X , MHA computes:
[Equation]
where W Qh , W Kh , W Vh ∈ R d hid ∗ d head , W O ∈ R d hid ∗ d hid are linear transformations.
Q h , K h , V h∈ 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 stackedon the MHA.
FFN consists of two linear transformations with a ReLU activation as:
[Equation]
where W 1 ∈ R d hid ∗ d f , W 2 ∈ R d f ∗ d hid , b 1 ∈ R d f , and b 2 ∈ R d hid .
In addition, there are residualconnection and layer normalization on top of MHA and FFN (denoted by ⊕ in Figure 2), which areformulated 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 severalcompressed 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 defined as r =
d f d hid , with d f and d hid representing the intermediate dimension of the FFN andhidden dimension respectively; (2) Number of layers ( L ) to capture the network depth; (3) Hiddendimension ( 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 significant impact on the modelsize and computational cost.
For instance, different layers have different feature representationcapabilities.
Recent works show that Transformer models are overparameterized [21, 22], such asthe 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 costresulting from FFN layers.
Furthermore, studies [24, 25] show that attention heads can be redundantwhen they learn to encode similar relationships for each word.
Thus, we make the number of attentionheads searchable as well.
Inductive bias.
Prior works on CNNs [26] and Transformers [27] demonstrate that thinner and deeperneural 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 ineach of our K sub-spaces (base, small, tiny), where we prefer deeper students in terms of the numberof layers.
Furthermore, we constrain all the Transformer layers in a given student model to shareidentical 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 andsoftware frameworks [13].
Search space partition.
Existing works [13, 14] train a single large SuperLM containing millions ofstudent architectures by weight-sharing.
This leads to performance degradation due to optimizationinterference and convergence of subnetworks with very different sizes [11].
To mitigate suchinterference, 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 studentmodels given by the number of parameters.
For instance, K = 3 can cover typical student sizes,namely base, small and tiny versions.
Table 1 shows the parameter ranges for the K sub-spaces,along with the student configurations contained in each.
Extensive details on the search space designcan be found in the Appendix (Section D).
[Table]
BERT
[Table]
MLP Ratio
[Table]
We now encode each subspace into a SuperLM,where each student modelin the space is a subnetworkof the SuperLM.
Furthermore, all the student subnetworks share the weights oftheir common dimensions,with the SuperLM being thelargest one in the searchspace.
Considering K independent SuperLMs, eachone now has more capacityto encode a sub-space, incontrast to a limited capacity single SuperLM as in prior works.
Furthermore, our choices for the heuristic partition andinductive bias result in less number of student models of comparable size in each sub-space whichalleviates conflicts in weight-sharing.
We extract student subnetworks from the SuperLM by a simple truncation strategy like bottom-leftextraction.
In particular, given a specific architecture α
= { l, d hid , r, h } , (i) we first extract alternatel Transformer layers from the SuperLM; (ii) then extract bottom-left sub-matrices in terms of d hidand r from the original matrices that represent the hidden dimension and the MLP ratio respectively;(iii) finally, for the attention heads, we extract the leftmost h heads and retain the dimension of eachhead as the SuperLM.
This strategy is used during sampling sub-networks for SuperLM training viaweight 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 wedefer 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 SuperLMtraining, we randomly sample several student subnetworks from the search space; apply knowledgedistillation between sampled subnetworks and the teacher to accumulate gradients; and then updatethe SuperLM.
During sampling, we employ Sandwich rule [28], also used in BigNAS [11], thatsamples 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 performancelower 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 employmulti-head self-attention relation distillation to align the attention distributions as well as scaleddot-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( · ) ) betweenmulti-head self-attention relation of teacher and sampled subnetwork is used as distillation objective:
[Equation]
where R Ti = softmax( A i
A ⊤ i / √ d k ) , R Si = softmax( B i B ⊤ i / √ d k ) , H is the number of attentionheads; R Ti represents the teacher’s Q − Q , K − K , or V − V relation; R Si represents the same forstudent.
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 thestudent’s representations with different dimensions to align to that of the teacher.
For the teachermodel and subnetworks with different number of attention heads, we first concatenate the selfattention vectors of different attention heads of the subnetwork and then split them according to thenumber of relation heads of the teacher model.
Then, we align their queries with the same number ofrelation heads for distillation.
In addition, we only transfer the self-attention knowledge from the lastlayer of the teacher model to the last layer of the student model.
Input: Partitioned K sub-spaces A k ; initialized K
Su- perLMs S k on A k ; pre-trained ; unlabeled epochs steps M
Output: Trained SuperLMs { S k } for k = 1 to K do for i = 1 to E do
Get a batch of data from D for batch in D do
Clear gradients in SuperLM S k for m = 1 to M do
Randomly sample a subnetwork s from S k Calculate self-attention distil.
loss between subnetwork s and teacher T with Eqn.
(4) 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 specified by the architecture α ; U are the weights of the teachermodel including the self-attention moduleused for distillation; D train is the training dataset, and L ( · ) is the self-attention loss functionfrom Eqn.
3.3 Lightweight Optimal Student Search
We outline two search strategies for selectingthe optimal student subnetwork.
Task-agnostic search.
We adopt this to be ourprimary strategy to compare against all baselines since it does not access any task labelinformation.
We compute the task-agnostic self-attention distillation loss for all student subnetworksusing 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 thestudent and number of subnetworks is limited.
The optimal student is given by the subnetwork withleast 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, Lis 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-specified resource constraint c .
Task-proxy search.
We compare our task-agnostic search against another strategy that considers aproxy task (e.g., MNLI [29]) with label information to fine-tune the 256 candidate subnetworks ineach sub-space.
The optimal student in each sub-space is given by the one with the best downstreamtask performance (e.g., accuracy).
Note that, for this strategy, the proxy task is used only duringsearch while the NAS training is still fully task-agnostic.
Experiments
Datasets.
We conduct experiments on the General Language Understanding Evaluation (GLUE)benchmark [30].
We compare our method with the baseline methods on two single-sentence classification 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]) 1 .
We report accuracy for
MNLI, QNLI, QQP, SST-2, RTE, report f1 for MRPC, and report Matthew’s correlation for CoLA.
Baselines.
We compare against several task-agnostic methods 2 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
configuration.
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 of 128and 4 e - 5 as the peak learning rate for 10 epochs.
The maximum sequence length is set to 128 .
The coefficients in distillation objective (Eqn.
(4)), β 1 , β 2 , and β 3 , are all set to 1 .
We distill theself-attention knowledge of the last layer to train the SuperLM.
Both the teacher and SuperLM areinitialized with pre-trained BERT BASE .
Other hyper-parameter settings are shown in Appendix.
Weuse 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 withoutusing any task label information.
We set a constraint in Eqn.
(6) such that the #FLOPs of the optimalcompressed model is atleast 50% less than the teacher model.
We rank all the subnetworks containedin all the partitions of the trained SuperLM by their self-attention distillation loss on the heldoutvalidation 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 differentsubnetworks.
Prior work [41] has demonstrated performance improvements in MNLI to be correlatedto other GLUE tasks.
To this end, we fine-tune all the 256 subnetworks in each partition of thetrained superLMs, and select corresponding subnetworks with the best trade-off (more discussionson 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 259 sub-spaces of SuperLM Base , SuperLM Small and SuperLM Tiny , respectively.
Notably all students areobtained 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 2 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 aggregatedover all the GLUE tasks.
We observe that the compressed model AutoDistil Agnostic generatedvia our task-agnostic SuperLM training leads to upto 3 x reduction in FLOPs over state-of-the-artdistilled models (e.g., MINILM [4], TinyBERT [6], DistilBERT [5]) that are hand-engineered whilematching 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 incurring5 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 languageinference, paraphrase detection and sentiment classification.
This depicts an interesting impact ofmassive 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 andparameters from corresponding papers in Table 2.
[Table] [Table]
[Table] [Table]
[Table]
AutoDistil outperforms all competing methods on aggregate for all sizes; except for small-sizedmodel; where it has marginally lower performance (0.1 points on avg) compared to AutoTinyBERT.
[Table]
comparison for generating students of similar FLOPs.
NR AutoTinyBERT 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
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 methodsthat use multi-stage training; where additional pre-training anddistillation is applied to NAS-generated candidates.
To better understand the impact of single-stage vs. multi-stagemethods on the training cost, we compare the overall cost of
NAS for AutoDistil and that reported in AutoTinyBERT 3for the small model segment in Table 3.
AutoDistil is muchcheaper due to its single-stage training protocol; where no additional pre-training or distillation is needed.
It is worth notingthat the overall SuperNet training cost of AutoDistil (themost expensive component of NAS) is less or comparable tothe additional training cost of re-training candidate modelsfor AutoTinyBERT.
Note that AutoTinyBERT does not reporttheir SuperNet training cost.
Additionally, AutoDistil hasa much faster search mechanism due to (1) inductive biasesbuilt into the search space definition to limit the number of student architectures and (2) task-agnosticsearch 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 NASmethods 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 fine-tuned on the MNLI task.
We observethe optimal AutoDistil models to outperform several competing methods.
[Table]
[Table]
We study different task-agnostic strategies for SuperLM training in AutoDistil .
Specifically, we compare three strate- [Table] gies in Table 4.
(i)
We replacing the KD loss in Eqn.
(4) withmasked language modeling (MLM) loss [1] to calculate gradients which is the most widely used task-agnostic pre-trainingand distillation strategy.
(ii) KD att
+Cont further continues training the searched compressed modelson the large language corpus. (iii) KD att is the strategy adopted in AutoDistil for self-attentiondistillation.
We evaluate subnetworks with the same architecture ( 6 layers, 768 hidden, 12 heads,
MLP ratio 4 ) from the trained SuperLM.
We fine-tune subnetworks on MRPC, RTE, and MNLI tasks,
[Figure]
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) the 3 sub-spacescorrespond to base, small and tiny model sizes; asused in prior work in CV; e.g. AutoFormer [17],(2) searching over different values of K is a veryresource-extensive process since it requires training
K SuperLMs for each choice of K , (3) As K increases, the search process becomes similar to theundesirable 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 withresults in Table 5.
For one-shot NAS, we consider a single search space containing varying number ofsubnetworks (e.g., 27 , 864 , 11232 ).
Few-shot NAS contains 256 subnetworks in each partition.
Wefine-tune subnetworks on RTE and MRPC tasks, and report accuracy and f1 respectively.
We observefewer subnetworks in a single search space for one-shot NAS result in better performance.
Thisresults from optimization interference and gradient conflicts as the number and size of subnetworksincrease in the space.
Finally, our design strategy performs the best while containing lesser numberof subnetworks demonstrating the benefit of few-shot NAS for language model distillation .
4.1.5 Comparing Search Strategies and Optimal Architectures
From Table 2, we observe that the student models AutoDistil Agnostic and AutoDistil Proxy Sobtained from SuperLM small by task-agnostic and task-proxy search strategies respectively obtain asimilar trade-off between performance and cost.
The task-proxy search results in a minor performancegain 0 .
3 over the fully task-agnostic search mechanism.
Table 6 shows the configuration of searchedoptimal architectures from AutoDistil with corresponding computational cost.
For reference, wealso 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 difficult to obtain by
[Table]
trial and error considering the large [Table] space Transformer architectures
[Table] We also observe that optimal compressed models have thin-and-deep [Table] structure consistent with findings thatthinner and deeper models performbetter [26] than wider and shallowerones.
While we use this as an inductive bias for sub-space partitioning, [Table] our search space (Table 1) also contains diverse subnetworks with different depth and width.
Non-maximal MLP ratio and attention heads for optimal compression indicatethat self-attention and feed-forward layers of Transformers are overparameterized [21, 22].
Related Work
Task-specific knowledge distillation.
Knowledge distillation (KD)
[42] is a widely used techniquefor model compression, which transfers knowledge from a large teacher to a smaller student.
Taskspecific KD aims to generate smaller students by using downstream task label information.
Typicaltask-specific KD works include BERT-PKD [43], BERT SMALL [40], TinyBERT [6], DynaBERT [12],and SparseBERT [44].
While task-specific KD often achieves good task performance, a typicaldrawback 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-specific KD, we explore task-agnostic
KDthat does not use any task label information.
The distilled task-agnostic models can be re-used bysimply fine-tuning on downstream tasks.
Task-agnostic KD leverages knowledge from soft targetprobabilities, hidden states, layer mappings and self-attention distributions of teacher to train studentmodels.
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-layerdistillation.
MiniLM transfers self-attention knowledge from the last layer of the teacher to that of thestudent.
These works rely on hand-designed architecture for the student models for KD that requiresseveral trials, and needs to be repeated for a new student with a different cost.
In contrast, we developtechniques 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 efficient sub-networks from the Transformer architecturefor machine translation tasks.
Some recent approaches closest to our method include, DynaBERT [12],
AutoTinyBERT [13] and NAS-BERT [14].
DynaBERT performs task-specific distillation.
NASBERT performs two-stage knowledge distillation with pre-training and fine-tuning of candidates.
Similar to above approaches, AutoTinyBERT also employs one-shot NAS with a single large searchspace containing millions of subnetworks that result in co-adaption and weight-sharing challengesfor SuperLM training.
Further it also uses a multi-stage training protocol for further pre-trainingand distillation of the NAS-generated candidates.
In contrast, AutoDistil employs few-shot NASwith a compact search space design with a single-stage task-agnostic training protocol.
This furtherallows us to do a lightweight search for the optimal student without re-training.
[16] studies few-shottask-specific NAS in the CV domain with CNN architectures, and different design and search spacesfrom 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 languagemodels into compressed students with variable computational cost.
To address the co-adaptionand weight-sharing challenges for SuperLM training, we partition the Transformer search spaceinto K compact sub-spaces covering important architectural components like its depth, width, andattention heads.
We leverage self-attention distillation for fully task-agnostic SuperLM training andlightweight optimal search without any re-training.
Obtained students can be re-used by simplyfine-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.