Few-shot Task-agnostic Neural Architecture Search for Distilling Large Language Models
Microsoft Research
Debadeepta Dey
Microsoft Research Microsoft Research
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 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 significant 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-specification of the 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 architectures as they are hand-engineered and to define 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 specification for the student architecture or computational cost for using it in a target environment.
author lead
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-specified 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 conflicts 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 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 pre-trained models (e.g., MobileNet [7], RegNet [19]) for task-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 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 fine-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 fine-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 fine-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 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, 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 ∈ [1 , L ] where L is the number of Transformer layers, X l ∈ 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]
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 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 ∈ 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 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 defined 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 significant 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 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 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 conflicts in weight-sharing.
[Table]
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 specific architecture α
= { l, d hid , r, h } , (i) we first 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) finally, 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 / √ d k ) , R Si = softmax( B i B ⊤ i / √ 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 first 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 sub-spaces A k ; initialized K
Su- perLMs S k on A k ; pre-trained teacher T ; unlabeled data D ; training epochs E ; sampling 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 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-specified 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 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) 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]) 3 .
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 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 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 ofand 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 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 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
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 classification.
This depicts an interesting impact of massive model compression on varying task types.
4.1.2 Comparison with Neural Architecture Search Baselines
[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] [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 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 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 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 .
Specifically, 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 fine-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 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
[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 fine-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 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 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 configuration 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 difficult to obtain by trial and error the large space
We also observe that optimal compressed models have thin-and-deep structure consistent with findings 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-specific 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.
Taskspecific KD aims to generate smaller students by using downstream task label information.
Typical task-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 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-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 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 efficient 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-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 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-specific 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 fine-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.
Appendix A Comparisons with Existing NAS and KD Methods A.1 Training Cost Comparisons with Traditional Task-agnostic KD Methods
Different hardwares (e.g., FPGA, CPU, GPU) have different resource constraints.
AutoDistil generates a gallery of fully trained compressed student models with variable resource constraints (e.g.,
FLOPs, parameters) using NAS.
One can simply choose a model from the trained pool given the resource constraint and only fine-tune on the downstream task.
In contrast, traditional task-agnostic knowledge distillation (KD) methods (e.g., MiniLM) target specific compression rate and needs to be trained repeatedly for different student configurations (corresponding to different resource constraints).
Therefore, AutoDistil has a much reduced amortized computation cost even considering traditional KD methods.
Further, traditional methods require several trial and errors to come up with a viable candidate architecture given a constraint before running the KD algorithm.
In practice, additional or continued training of the optimal student architecture has demonstrated increased task performance with increased computational cost as in AutoTinyBERT (cost comparison in Table 4).
The major advantage of AutoDistil is a single stage training scheme without additional training.
We perform an ablation in Table 5 where we continue training the searched model with self-attention distillation for additional steps referred as ‘KD att +Cont.’
(similar to MiniLM).
But we did not observe any significant gains on a subset of the tasks.
A.2 Comparisons between AutoDistil and Few-shot CV NAS [16]
Fully task-agnostic SuperNet training .
AutoDistil training is fully task-agnostic in contrast to [16] that uses task-specific NAS.
Task-agnostic NAS is challenging since we do not have access to task labels during training and we want to show generalization on evaluating diverse downstream NLU tasks in the GLUE benchmark.
AutoDistil leverages self-attention distillation that is an unsupervised training objective.
Incorporating self-attention loss for training and distillation in NAS is non-trivial as it requires aligning attention states of diverse student subnetworks and the large teacher model.
We develop an extraction and alignment strategy (Section 3.2) to address this challenge.
NLP vs. CV domain .
AutoDistil works on the NLP domain with the Transformer architecture (see Figure 2 under the pre-training and fine-tuning paradigm, while reference [16] works on the
CV domain with a CNN architecture with different design and search spaces.
Different from CV domain, NLP tasks have different objectives and evaluation metrics for classification (e.g., MNLI), regression (e.g., STS-B) and correlation (e.g., CoLA).
Overall, the search space design (Section 3.1), SuperNet training with distillation and sub-network extraction strategy (Section 3.2) and search strategy (Section 3.3) are all quite different.
A.3 Comparisons between AutoDistil and DynaBERT [12]
Compared with DynaBERT, (i) AutoDistil search space is more fine-grained.
For instance, we independently search for width, depth, heads, MLP ratio etc.
as opposed to searching for a constant depth ( m d ) or width multiplier ( m w ) in DynaBERT which only considers 12 possible combinations of m d and m w ; (ii) our training objective does not require labels and is fully task-agnostic with subnetwork attention state alignment for self-attention relation distillation; (iii) further, AutoDistil uses few-shot NAS (Table 1) to mitigate gradient conflicts in SuperNet training, while DynaBERT applies one-shot NAS; (iv) DynaBERT uses additional tricks like data augmentation and teacher assistant also specific to each task, whereas AutoDistil uses a single-stage task-agnostic training resulting in reduced computational cost.
A.4 Additional Quantitative Comparisons with Existing NAS and KD Methods
AutoDistil Agnostic is obtained by
[Table]
We compare different AutoDistil compressed models against state-of-the-art KD and NAS models distilled from the same teacher BERT BASE .
We present the relative performance improvement of
AutoDistil over several baselines in Table 8 with respect to the following measures: savings in computational cost in the form of (i) FLOPs and (ii) parameter reduction, along with (iii) improvement in the average task performance aggregated over all the GLUE tasks.
We also compare AutoDistil with state-of-the-art distilled models in terms of the trade-off between model size (#Para) and performance (accuracy).
The results are shown in Figure 4.
AutoDistil uses few-shot task-agnostic Neural Architecture Search to distill several compressed students with variable #Para (x-axis) from K = 3 SuperLMs (corresponding to each point cloud) trained on K sub-spaces of Transformer search space.
Each student extracted from the SuperLM is fine-tuned on MNLI with y-axis showing accuracy.
The best student from each SuperLM is marked in red.
Given any state-of-the-art distilled model, AutoDistil generates a better candidate with less #Para and improved task performance from corresponding search space.
B Layer Selection Strategies
[Table]
We study different strategies to construct subnetwork layers by selecting layers from the superLM model.
Alternate_Dropping is the strategy adopted in AutoDistil such that we drop alternating odd layers from the superLM model to construct subnetwork layers.
Top_Dropping means that we drop top layers of superLM to construct subnetwork layers.
Alternate_Top_Dropping means that we first perform Alternate_Dropping in superLM training stage and then perform Top_Dropping in fine-tuning stage (please refer to [47] for more details of different layer selection strategies).
For all strategies, we perform knowledge distillation between the last layer of the teacher model and the last layer of the subnetworks.
We evaluate the subnetworks with the same architecture (#layer=6, #hid=768, R=4, #heads=12) after superLM is trained.
We report accuracy and f1 for RTE and MRPC, respectively.
We report the results in Table 9.
We observe that the strategy of Alternate_Dropping achieves the best performance on both MRPC and RTE tasks, which demonstrates the effectiveness of the layer selection strategy used in AutoDistil .
Alternate_Top_Dropping performs the worst due to interference when different layer selection strategies are used in the superLM training stage and the fine-tuning stage of compressed models.
This indicates that the knowledge contained in the superLM model and the compressed model is structured and that it is non-trivial to select layers from superLM to extract subnetwork layers.
C Scaling of Training Data
[Table]
We investigate the effects of data sets of different sizes used for superLM training.
In particular, we compare MNLI [29], ParaNMT [48] (we sampled 5 million samples from the original 50 million data), Wiki, and Wiki+Book [49].
We report the size of each data set and the performance of
AutoDistil with each training data set in Table 10.
We observe that AutoDistil performs the best with Wiki+Book data set, and the larger the data set, the better the performance.
Moreover, we observe similar performance for MNLI and ParaNMT data sets, especially on MRPC task.
This is because MNLI is correlated to other GLUE tasks.
In addition, we observe that an increase in the amount of data does not guarantee to bring an equivalent increase in performance.
For example, Wiki data set is more than five times larger than ParaNMT data set, but our method performs only about 1% better With Wiki data set than with ParaNMT.
These observations illustrate that while using a larger data set does improve the performance of the method, the improvement could be quite limited.
D Search Space Design
In general, we partition the whole Transformer search space into K = 3 sub-spaces such that each sub-space covers different sizes of student models (by number of parameters) depicting Tiny, Small and Base model sizes.
Given a BERT-sized teacher model (109 M params), we roughly set the partition thresholds for Tiny, Small and Base sizes at 10 M , 40 M and 80 M params.
From Table 2 (#Params row), we observe that each partition contains compressed models from prior work – allowing us to fairly compare the models in each partition on accuracy vs. params/FLOPs.
For our search space, each partition still contains thousands of candidate subnetworks not all of which are useful.
Now, we leverage two primary heuristics: (i) we constrain all layers in a sampled student subnetwork to be homogeneous i.e., the same number of attention heads, hidden dimension, etc.
This not only reduces the search space, it is also more friendly to hardware and software frameworks [13].
(ii) Motivated by previous work [26, 27] showing that thinner and deeper neural networks have better representation capabilities and perform better than wider and shallower neural networks, we designed sub-spaces with deeper layers (e.g., 4 − 7 for Tiny, 9 − 12 for Small and Base) and compute the range of hidden dimensions to meet the overall model parameter budget in each partition.
Additional constraints arise from Transformer design principles, for instance, hidden size is always a multiple of the number of attention heads.
While the above steps require enumeration of different subnetwork architectures, this is typically fast given an algebraic expression to compute model parameters as a function of layers, heads, hidden size, etc.
(included in source code), does not require any training, and a one-time process depending only on the teacher model architecture.
The impact of network depth on model performance has been observed with both convolution architectures [26] and Transformers [27].
Figure 2(b) in [27] shows the impact of Transformer depth on MNLI accuracy given similar overall model parameters.
E Transferability of Optimal Architectures
The transferability of optimal architectures has been studied with regards to model pruning in the lottery ticket hypothesis work [41] for BERT.
They observe that transferability seems to correlate with the number of training examples available for the source task.
This is particularly beneficial with MNLI containing a large number of training examples as compared to other low-resource tasks in the GLUE benchmark.
Similar to [41], we also observe MNLI to transfer well to other tasks in the GLUE benchmark with AutoDistil-proxy even outperforming task-specific NAS method like
DynaBERT (Table 3) on both parameters (26.1 M vs. 37.7 M ) and average accuracy (79.9 vs. 77.4).
In general, a better teacher model leads to a better student model [50, 4] during distillation.
We adopted
BERT as teacher for a fair comparison with existing works with the same teacher.
We also compare with different training objectives like self-attention distillation and masked language modeling
(Table 5) and demonstrate the former to work better for our SuperNet training.
We demonstrate transferability by training the AutoDistil students in a task-agnostic manner and evaluating on different downstream tasks (Table 3).
Note that these tasks are quite diverse ranging from classification (e.g., MNLI), regression (e.g., STS-B) and correlation (e.g., CoLA).
We also demonstrate this to work better or comparable to task-specific NAS methods (e.g., DynaBERT, AutoTinyBERT) with further reduction of computational cost.
F Application of Few-shot Task-agnostic NAS Method to Other Domains
Most NAS works in computer vision (CV) (e.g., Once-for-all, One-Shot NAS) leverage hard class labels from a given task (e.g., image classification).
They often use similar training recipes for
SuperNets as in ImageNet-trained models (e.g., MobileNet, RegNet) for task-specific optimization with accuracy as an evaluation metric.
In contrast, the few-shot task-agnostic NAS strategy used in
AutoDistil training is fully task-agnostic and does not access task labels during SuperNet training.
A potential method to adopt this strategy for CV domain is to consider a self-supervised learning framework like SimCLR [51] that leverages data augmentation for consistency learning.
This requires both a self-supervised teacher like SimCLR and a self-supervised training objective (e.g., self-attention relation distillation for Transformers or architecture-agnostic consistency learning).
This forms an interesting direction for future work.
G Quantifying "Best Trade-off" between Task Performance (e.g., Accuracy) and Computational Cost (e.g., #FLOPs)
We describe how to search for the optimal sub-network in Section 3.3 and Section 4.1.
The "best trade-off" for optimal student selection is given by the sub-network with the least validation loss subject to the resource constraint as described in Eqn (6).
For instance, we set a constraint in Eqn.
(6) such that the #FLOPs of the optimal Base-sized task-agnostic compressed model is atleast 50% less than the teacher model.
Since the SuperNet training is task-agnostic, the obtained student models have to be fine-tuned on downstream tasks to report the final task performance (similar to pre-train and fine-tune paradigm of BERT-like language models).
H Why Are the Selected Models (Red in Figure 1) Not Always the Best Performing Models?
Note that our objective is to minimize the #FLOPs and maximize the accuracy (e.g., on MNLI) with the trade-off determined by the resource constraint for different partitions (Base, Small, Tiny).
Given a gallery of compressed models from AutoDistil with variable FLOPs and performance, we use strategy from Section G for optimal model selection.
Another potential reason why red is always not the best model is that we use the heldout validation set from the unlabeled training corpus (Wikipedia + BooksCorpus) for student selection and then evaluate them on MNLI (see Sections 3.3 and 4.1) which may not be optimal due to sample differences in the two datasets.
I Subnetwork Evaluation Strategy in Task-agnostic Search
The validation set contains 300 K instances.
We use 128 as the sequence length and batch size.
(ii)
Table 4 reports the search cost for Small-sized models from AutoDistil and AutoTinyBERT.
Note that this step does not require any training for AutoDistil.
We compute only the self-attention relation loss for all the 256 student subnetworks ( 5 x − 22 x speedups for Small-sized models) using Equation 4 with the teacher relations computed only once.
We use Equation 6 to select the subnetwork with desired trade-off with deterministic computation of the FLOPs.
The algebraic expression to compute
FLOPs as a function of layers, heads, hidden size etc. is included in the submitted source code.
In contrast, AutoTinyBERT performs task-specific search which requires fine-tuning the subnetworks on the task (e.g., MNLI) thereby increasing the search cost.
J Hyper-parameter Settings for Fine-Tuning
[Table]
We report the fine-tuning hyper-parameter settings of GLUE benchmark in Table 11.
AutoDistil and baselines follow the same settings.
K Limitations and Broader Impact
In this work, we introduce a framework for distilling large pre-trained neural language models with resource constraints.
This work is likely to increase the progress of NLP applications and drive the development of generalpurpose language systems for deployment environments with limited resources.
Our framework can be used for applications in finance, legal, healthcare, retail and other domains for edge scenarios where adoption of deep neural networks have been hindered due to resource efficiency concerns.
In our work, we propose a solution of a single-stage training combining NAS and distillation with no further pre-training or augmentation, and experiments on GLUE benchmark demonstrate our solution to outperform state-of-the-art KD and NAS methods with upto 3 x additional reduction in computational cost and negligible loss in task performance.
While our framework accelerates the progress of NLP, it may also suffer from similar concerns as with the use of large pre-trained models by malicious agents for propagating bias, misinformation and indulging in other nefarious activities.