Distiller: A Systematic Study of Model Distillation Methods in Natural Language Processing Haoyu He1∗, Xingjian Shi2 , Jonas Mueller2 , Sheng Zha2 , Mu Li2 , and George Karypis2 1 Northeastern University 2 Amazon Web Services he.haoy@northeastern.edu {xjshi,jonasmue,zhasheng,mli,gkarypis}@amazon.com arXiv:2109.11105v1 [cs.CL] 23 Sep 2021 Abstract Knowledge Distillation (KD) offers a natural way to reduce the latency and memory/energy usage of massive pretrained models that have come to dominate Natural Language Processing (NLP) in recent years. While numerous sophisticated variants of KD algorithms have been proposed for NLP applications, the key factors underpinning the optimal distillation performance are often confounded and remain unclear. We aim to identify how different components in the KD pipeline affect the resulting performance and how much the optimal KD pipeline varies across different datasets/tasks, such as the data augmentation policy, the loss function, and the intermediate representation for transferring the knowledge between teacher and student. To tease apart their effects, we propose Distiller, a meta KD framework that systematically combines a broad range of techniques across different stages of the KD pipeline, which enables us to quantify each component’s contribution. Within Distiller, we unify commonly used objectives for distillation of intermediate representations under a universal mutual information (MI) objective and propose a class of MIα objective functions with better bias/variance trade-off for estimating the MI between the teacher and the student. On a diverse set of NLP datasets, the best Distiller configurations are identified via large-scale hyper-parameter optimization. Our experiments reveal the following: 1) the approach used to distill the intermediate representations is the most important factor in KD performance, 2) among different objectives for intermediate distillation, MI-α performs the best, and 3) data augmentation provides a large boost for small training datasets or small student networks. Moreover, we find that different datasets/tasks prefer different KD algorithms, and thus propose a simple AutoDistiller algorithm that can recommend a good KD pipeline for a new dataset. ∗ vices. Work done while being an intern at Amazon Web Ser- 1 Introduction Recent advancements in Natural Language Processing (NLP) such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019), and ELECTRA (Clark et al., 2020) have demonstrated the effectiveness of Transformer models in generating transferable language representations. Pretraining over a massive unlabeled corpus and then fine-tuning over labeled data for the task of interest has become the state-ofthe-art paradigm for solving diverse NLP problems ranging from sentence classification to question answering (Raffel et al., 2019). Scaling up the size of these networks has led to rapid NLP improvements. However, these improvements have come at the expense of significant increase of memory for their many parameters and compute to produce predictions (Brown et al., 2020; Kaplan et al., 2020). This prevents these models from being deployed on resource-constrained devices such as smart phones and browsers, or for latency-constrained applications such as click-through-rate prediction. There is great demand for models that are smaller in size, yet still retain similar accuracy as those having a large number of parameters. To reduce the model size while preserving accuracy, various model compression techniques have been proposed such as: pruning, quantization, and Knowledge Distillation (KD) (Gupta and Agrawal, 2020). Among these methods, task-aware KD is a popular and particularly promising approach for compressing Transformer-based models (Gupta and Agrawal, 2020). The general idea is to first fine-tune a large model (namely the teacher model) based on the task’s labeled data, and then train a separate network that has significantly fewer parameters (namely the student model) than the original model to mimic the predictions of the original large model. A large number of task-aware KD algorithms have been proposed in the NLP regime, e.g., DistillBERT (Sanh et al., 2019), BERT-PKD (Sun et al., 2019), BERT-EMD (Li et al., 2020), TinyBERT (Jiao et al., 2020), DynaBERT (Hou et al., 2020), and AdaBERT (Chen et al., 2020), some of which can compress the teacher network by 10× without significant accuracy loss on certain datasets. Innovations in KD for NLP generally involve improvements in one of the following aspects: 1) the loss function for gauging the discrepancy between student and teacher predictions (Kim et al., 2021), 2) the method for transferring intermediate network representations between teacher and student (Sun et al., 2019; Li et al., 2020; Yang et al., 2021), 3) the use of data augmentation during student training (Jiao et al., 2020), and 4) multiple stages of distillation (Chen et al., 2020; Mirzadeh et al., 2020). Many research proposals have simultaneously introduced new variations of more than one of these components, which confounds the impact of each component on the final performance of the distillation algorithm. In addition, it is often unclear whether a proposed distillation pipeline will generalize to a new dataset or task, making automated KD challenging. For example, MixKD (Liang et al., 2021) has only been evaluated on classification problems and it is unclear if the method will be effective for question answering. To understand the importance of different components in KD, we undertake a systematic study of KD algorithms in NLP. Our study is conducted using a meta-distillation pipeline we call Distiller that contains multiple configurable components. All candidate algorithms in the search space of Distiller work for two types of NLP tasks: text classification and sentence tagging. Distiller unifies existing techniques for knowledge transfer from intermediate layers of the teacher network to the student network (i.e. intermediate distillation) as special cases of maximizing (bounds on) the Mutual Information (MI) between teacher and student representations. Based on this unification and recent progress in variational bounds of MI (Poole et al., 2019), we propose a new intermediate distillation objective called MI-α that uses a scalar α to control the bias-variance trade-off of MI estimation. Including MI-α in the search space of Distiller, we run extensive hyper-parameter tuning algorithms to search for the best Distiller-configuration choices over GLUE (Wang et al., 2019b) and SQuAD (Rajpurkar et al., 2016). This search helps us identify the best distillation pipelines and understand what impact different KD modules have on student performance in NLP. Using the observations of this large-scale study, we train a AutoDistiller model to predict the distillation ratio, which is defined as the fraction of the teacher’s performance achieved by the student, based on KD pipeline choices and characteristics of a dataset. Leave-one-out cross validation evaluation of AutoDistiller demonstrates that it is able to reliably prioritize high-performing KD configurations in most folds, and is able to suggest good distillation pipelines on two new datasets. The main contributions of this work include: • The meta KD pipeline Distiller used to systematically study the impact of different components in KD, including the: 1) data augmentation policy, 2) loss function for transferring intermediate representations, 3) layer mapping strategies for intermediate representations, 4) loss function for transferring outputs, as well as what role the task/dataset play. • Unification of existing objectives for distilling intermediate representations as instances of maximizing bounds of the mutual information between teacher and student representations. This leads us to propose the MI-α objective that outperforms the existing objectives. • Using the results collected from our systematic Distiller study, we fit a model that automatically predicts the best distillation strategy for a new dataset. On a never-seen dataset “BoolQ”(Wang et al., 2019a), predicted strategies achieve 1.002 distillation ratios (fraction of the student’s and the teacher’s performance) on average, outperforming random selected strategies with mean of 0.960. To the best of our knowledge, this is the first attempt towards automated KD in NLP. 2 Related Work Knowledge Distillation. The general KD framework was popularized by Buciluǎ et al. (2006); Hinton et al. (2014), aiming to transfer knowledge from an accurate but cumbersome teacher model to a compact student model by matching the class probabilities produced by the teacher and the student. Focusing on AutoML settings with tabular data, Fakoor et al. (2020) proposed a general KD algorithm for different classical ML models and ensembles thereof. Also hoping to identify good choices in the KD pipeline like our work, Kim et al. (2021) compared Kullback-Leibler divergence and mean squared error objectives in KD for image classification models, finding that mean squared error performs better. Recent KD research in the domain of NLP has investigated how to efficiently transfer knowledge from pretrained Transformer models. Sun et al. (2019) proposed BERT-PKD that transfers the knowledge from both the final layer and intermediate layers of the teacher network. Jiao et al. (2020) proposed the TinyBERT model that first distills the general knowledge of the teacher by minimizing the Masked Language Model (MLM) objective (Devlin et al., 2019), with subsequent task-specific distillation. Li et al. (2020) proposed a many-to-many layer mapping function leveraging the Earth Mover’s Distance to transfer intermediate knowledge. Our paper differs from the existing work in that we provide a systematic analysis of the different components in single-stage task-aware KD algorithms in NLP and propose the first automated KD algorithm in this area. Mutual Information Estimation. Mutual Information (MI) measures the degree of statistical dependence between random variables. Given random variables A and B, the MI between them, I(A, B), can be understood as how much knowing A will reduce the uncertainty of B or vice versa. For distributions that do not have analytical forms, maximizing MI directly is often intractable. To overcome this difficulty, recent work resorts to variational bounds (Donsker and Varadhan, 1975; Blei et al., 2017; Nguyen et al., 2010) and deep learning (Oord et al., 2018) to estimate MI. These works utilize flexible parametric distributions or critics that are parameterized neural networks (NNs) to approximate unknown densities that appear in MI calcuations. Poole et al. (2019) provides a review of existing MI estimators and proposes novel bounds that trade-off bias and variance. Kong et al. (2020) unified language representation learning objective functions from the MI maximization perspective. Data Augmentation in NLP. Data Augmentation (DA) is an effective technique for improving the accuracy of text classification models (Wei and Zou, 2019) and has also been shown to boost the performance of KD for NLP algorithms (Jiao et al., 2020) as well as KD in models for tabular data (Fakoor et al., 2020). Wei and Zou (2019) proposed the Easy Data Augmentation (EDA) technique that randomly replaces synonyms, inserts, swaps and deletes characters in the sentence. Jiao et al. (2020) proposed to utilize the pretrained BERT model and GloVe word embeddings (Pennington et al., 2014) to augment the input sentence via random wordlevel replacement. MixKD (Liang et al., 2021) adopts mixup (Zhang et al., 2018) and backtranslation (Edunov et al., 2018) in augmenting the text data to boost the performance of sentence classification models. Unlike these papers, we propose a novel search space for DA policies that supports stacking elementary augmentation operations such as EDA, mixup, and backtranslation. Thus, our considered DA module is similar to AutoAugment (Cubuk et al., 2019), except it is used for KD in NLP with different elementary operators. 3 Methodology Our study is structured around a configurable metadistillation pipeline called Distiller. Distiller contains four configurable components, namely: a data augmentation policy a(·, ·), a layer mapping configuration of intermediate distillation {mi,j }, an intermediate distillation objective linter (·, ·), and a prediction layer distillation objective lpred (·, ·). Assume the teacher network f T has M layers and the student network f S has N layers. Then for a given data/label pair (x, y) sampled from the dataset D, the student acquires knowledge from the teacher by minimizing the following objective: L = Ex̂,ŷ∼a(x,y);x,y∼D M X N X inter mi,j li,j (HiT , HjS ) i=1 j=1 + β1 l pred T S (f (x), f (x)) + γ1 lpred (y, f S (x)) + β2 lpred (f T (x̂), f S (x̂)) + γ2 lpred (ŷ, f S (x̂)), (1) Here mi,j ∈ [0, 1] represents the layer mapping weight between the i-th teacher layer and j-th student layer, HiT , HjS are the i-th and the j-th hidden states of the teacher and the student (i.e. their intermediate representations at layers i and j), β1 , β2 control the strength of distilling from class probabilities produced by the teacher, and γ1 and γ2 control the strength of learning from ground truth data (x, y) and synthesized (augmented) data (x̂, ŷ). In Appendix, we illustrated how previous model distillation algorithms (Jiao et al., 2020; Li et al., 2020; Liang et al., 2021) can be encompassed in the Distiller framework. Figure 1: Overview of the Distiller pipeline. All configurable components are colored. Algorithm 1: Data Augmentation Policy Params: A sequence of elementary data augmentation operations G, ∀Gj ∈ {CA, RA, BT, Mixup}. Input: Training Dataset Dtrain Output: Augmented dataset Dsynthesize 1 2 3 4 5 6 7 8 Initialize Dsynthesize ← {} foreach {xi , yi } ∈ Dtrain do for j ← 1 to len(G) do x̂i , ŷi = Gj (xi , yi ); xi , yi ← x̂i , ŷi ; end Dsynthesize ← Dsynthesize ∪ {xi , yi } end tions in Distiller is constructed by stacking these four elementary operations in an arbitrary order, as detailed in Algorithm 1. For contextual augmentation, we use the pretrained BERT model to do word level replacement by filling in randomly masked tokens. As in EDA (Wei and Zou, 2019), our random augmentation randomly swaps words in the sentence or replaces words with their synonyms. For backtranslation, we translate the sentence from one language (in this paper, English) to another language (in this paper, German) and then translate it back. Additionally, mixup can be used to synthesize augmented training samples. First proposed for image classification (Zhang et al., 2018), mixup constructs a synthetic training example via the weighted average of two samples (including the labels) drawn at random from the training data. To use it in NLP, Guo et al. (2019); Liang et al. (2021) applied mixup on the word embeddings at each sentence position xi,t with λ ∈ [0, 1] as the mixing-ratio for a particular pair of examples xi , xj : x̂i,t = λxi,t + (1 − λ)xj,t , ŷi = λyi + (1 − λ)yj , Here λ is typically randomly drawn from a Uniform or Beta distribution for each pair, yi , yj are labels in one-hot vector format, and (x̂i , ŷi ) denotes the new augmented sample. To apply mixup for sentence tagging tasks, in which each token has its own label, we propose calculating the weighted combination of the ground-truth target at each location t as the new target: x̂i,t = λxi,t +(1−λ)xj,t , ŷi,t = λyi,t +(1−λ)yj,t , 3.1 Data Augmentation Policy A key challenge is the limited data available to train students in KD. This can be mitigated via Data Augmentation (DA) to generate additional data samples. Unlike in supervised learning, where labels for synthetic augmented data may be unclear unless the augmentation is limited to truly benign perturbations, the labels for augmented data in KD are simply provided by the teacher which allows for more aggressive augmentation (Fakoor et al., 2020). Denoting the set of training samples of the down-stream task as Dtrain , the augmenter a(·, ·) will stretch the distribution from Ex,y∼Dtrain to Ex̂,ŷ∼a(x,y),x,y∼Dtrain . We consider various elementary DA operations including: 1) MLM-based contextual augmentation (CA) , 2) random augmentation (RA), 3) backtranslation (BT) and 4) mixup. The search space of possible augmenta- 3.2 Prediction Layer Distillation In traditional KD, the student network learns from the output logits of the teacher network, adopting these as soft labels for the student’s training data (Hinton et al., 2014). Here we penalize the discrepancy between the outputs of student vs. teacher via: Lpred = lpred (f T (x), f S (x)), (2) lpred (·, ·) where is the KD loss component whose search space in this work includes either: softmax Cross-Entropy (CE) or Mean Squared Error (MSE). 3.3 Intermediate Representation Distillation To ensure knowledge is sufficiently transferred, we can allow the student to learn from intermediate layers of the teacher rather than only the latter’s output predictions by minimizing discrepancies between selected layers from the teacher and the student. These high-dimensional intermediate layer representations constitute a much richer information-dense signal than is available in the low-dimensional predictions from the output layer. As teacher and student usually have different number of layers and hidden-state dimensionalities, it is not clear how to map teacher layers to student layers (mi,j ) and how to measure the discrepancy inter ). Previous works between their hidden states (li,j proposed various discrepancy measures (or loss functions) for intermediate distillation, including: Cross-Entropy (CE), Mean Squared Error (MSE), L2 distance, Cosine Similarity (Cos), and Patient Knowledge Distillation (PKD) (Sun et al., 2019). For these objectives, we establish the following result (the proof is relegated to the Appendix). Theorem 3.1 Minimizing MSE, L2, or PKD loss, and maximizing cosine similarity between two random variables X,Y are equivalent to maximizing lower bounds of the mutual information I(X; Y ). In our KD setting, X and Y correspond to the hidden state representations of our student and teacher model (for random training examples), respectively. Inspired by this result, we can use any lower bounds of MI as an intermediate objective function in KD. In particular, we consider the multisample MI lower bound of Poole et al. (2019), which estimates I(X; Y ) given the sample x, y from p(x, y) and another K additional IID samples z1:K that are drawn from a distribution independent from X and Y : # ef (x,y) αm(y; x, z1:K ) + (1 − α)q(y) " # ef (x,y) − Ep(x,z1:K )p(y) log +1 αm(y; x, z1:K ) + (1 − α)q(y) " I(X; Y ) ≥ Ep(x,z1:K )p(y|x) log , Iα . (3) In Iα , f (·, ·) and q(·) are critic functions for approximating unknown densities and m(·, ·) is a Monte-Carlo estimate of the partition function that appears in MI calculations. Typically, the space z and the sample x, y are from the same minibatch while training, that is K +1 equals to the minibatch size. Iα can flexibly trade off bias and variance, since increasing α ∈ [0, 1] will reduce the variance of the estimator while increasing its bias. We propose to use Iα as an objective for intermediate distillation and call it MI-α. Our implementation leverages a Transformer encoder (Vaswani et al., 2017) to learn f (·, ·) and q(·). To our knowledge, this is the first attempt to utilize complex NN architectures for critic functions in MI estimation; typically only shallow multilayer perceptrons (MLPs) are used (Tschannen et al., 2020). Our experiments (Table 4 in Appendix) reveal that Transformer produces a better critic function than MLP. Note that for intermediate distillation, objectives like MSE attempt to ensure the teacher and student representations take matching values, whereas objectives like MI (and tighter bounds thereof) merely attempt to ensure the information in the teacher representation is also captured in the student representation. The latter aim is conceptually better suited for KD, particularly in settings where the student’s architecture differs from the teacher (e.g. it is more compact), in which case forcing intermediate student representations to take the exact same values as teacher representations seems overly stringent and unnecessary for a good student (it may even be harmful for tiny student networks that lack the capacity to learn the same function composition used by the teacher). We emphasize that a high MI between student and teacher representations suffices for the teacher’s prediction to be approximately recovered from the student’s intermediate representation (assuming the teacher uses deterministic output layers as is standard in today’s NLP models). Given that high MI suffices for the student to match the teacher, we expect tighter MI bounds like MI-α can outperform looser bounds like MSE that impose additional requirements on the student’s intermediate representations beyond just their information content. 3.3.1 Layer Mapping Strategy We investigate three intermediate layer mapping strategies: 1) Skip: the student learns from every ⌊M/N ⌋ layer of the teacher, i.e., mi,j = 1 when j = i × ⌊M/N ⌋; 2) Last: the student learns from the last k layers of the teacher, i.e., mi,j = 1 when j = i + M − N ; and 3) EMD: a many-to-many learned layer mapping strategy (Li et al., 2020) based on Earth Mover’s Distance. In the Distiller pipeline, the intermediate loss with EMD mapping can be denoted as: PM PN H H i=1 j=1 wi,j di,j T S , LEMD (H1:N , H1:M ) = PM PN H i=1 j=1 wi,j (4) where DH = [dH ] is a distance matrix representi,j ing the cost of transferring the hidden states knowlH ] is the edge from H T to H S . And W H = [wi,j mapping flow matrix which is learned by minimiz- ing the cumulative cost required to transfer knowledge from H T to H S . In Distiller, the distance matrix is calculated via intermediate objective funcinter (H S , H T ). tion: dH i,j = l i j 3.4 AutoDistiller Our experiments indicate that the best distillation algorithm varies among datasets/tasks (see in particular the top-5 configurations listed in Table 9 in Appendix). This inspires us to train a prediction model that recommends a good KD pipeline given a dataset. To represent the distillation performance across datasets which are evaluated on different metrics, we define distillation ratio as the fraction of the teacher’s performance achieved by the student and use it as a general score for distillation performance. Then the prediction model can be trained to predict the distillation ratio based on features of the dataset/task as well as the features of each candidate distillation pipeline. Here we train our AutoDistiller performance prediction model via AutoGluon-Tabular, a simple AutoML tool for supervised learning (Erickson et al., 2020). To the best of our knowledge, our proposed method is the first attempt towards automated KD in NLP. 4 Experimental Setup To study the importance of each component described in the previous section, we randomly sample Distiller configurations in the designed search space while fixing the optimizer and other unrelated hyper-parameters. We apply each sampled distillation configuration on a diverse set of NLP tasks and different teacher/student architectures. All experiments are evaluated on GLUE (Wang et al., 2019b) and SQuAD v1.1 (Rajpurkar et al., 2016) that contain classification, regression, and sentence tagging tasks. Here, we view the question answering problem in SQuAD v1.1 as finding the correct answer span from the given context, which is essentially a sentence tagging task. We adopt the same metrics for these tasks as in the original papers (Wang et al., 2019b; Rajpurkar et al., 2016). Since Turc et al. (2019) finds initializing students with pretrained weights is better for distillation, we initialize student models with either weights obtained from taskagnostic distillation (Jiao et al., 2020) or pretrained from scratch (Turc et al., 2019). Three different pretrained models BERTBASE (Devlin et al., 2019), RoBERTaLARGE (Liu et al., 2019) and ELECTRALARGE (Clark et al., 2020) are considered as teacher models in our experiments after task-specific fine-tuning. As student models, we consider options like TinyBERT4 , ELECTRASMALL , as well as other models detailed in Table 6 in Appendix. Existing implementations of data augmentation in KD for NLP generally first generate an augmented dataset that is K times larger than the original one and then apply the distillation algorithm over the augmented dataset. Such implementation separates the process of DA from KD, leading to a time/storage-consuming and inflexible KD pipeline. In Distiller, we instead apply DA dynamically during training. In addition, we use the teacher network to compute the soft label ŷ assigned to any augmented sample x̂. To analyze the importance of different components in Distiller, we adopted fANOVA (Hutter et al., 2014), an algorithm for quantifying the importance of individual hyper-parameters as well as their interactions in determining down-stream performance. We use fANOVA to evaluate the importance of the four components in Distiller as well as their pairwise combinations: data augmentation, intermediate distillation objective, layer mapping strategy, and prediction layer distillation objective. 5 Experimental Results Under the previously described experimental setup, we conducted a random search over Distiller configurations on each dataset and collected more than 1300 data points in total. Each collected data point contains a particular Distiller configuration, the dataset/task, the teacher/student architectures, and the final performance of the student model. Analyzing the data reveals three major findings: 1. Design of the intermediate distillation module is the most important choice among all factors studied in Distiller. 2. Among different loss functions for intermediate distillation, MI-α performs the best. 3. DA provides a large boost when the dataset or the student model is small. Additionally, we observe that the best distillation policy varies among datasets/tasks as shown in Table 9 in Appendix. Thus we train a meta-learning model, AutoDistiller, that can recommend a good distillation policy on any new NLP dataset based on which configurations tended to work well for similar datasets in our study. Table 1: Comparison of evaluation results on GLUE test set. BERTBASE (G) and BERTBASE (T) indicate the finetuned BERTBASE from (Devlin et al., 2019) and the teacher model trained by ourselves, respectively. BERT-EMD4 and MI-α are both initialized from TinyBERT4 , the difference is that BERT-EMD4 is trained with “EMD” as intermediate layer mapping strategy and MSE as intermediate loss, our MI-α model is trained with “Skip” as intermediate layer mapping strategy and MI-α as intermediate loss. Model #Params BERTBASE (G) BERTBASE (T) BERT-PKD4 (Sun et al., 2019) BERT-EMD4 (Li et al., 2020) MI-α (α = 0.9, ours) 110M 110M 52M 14M 14M MNLI-m (393k) 84.6 84.5 79.9 82.1 81.9 MNLI-mm (393k) 83.4 83.6 79.3 80.6 80.6 Figure 2: As assessed via fANOVA, we report the individual importance of the four Distiller components in this figure and importance of interactions between any two of the four components in Figure 6 in Appendix. Four components are: linter for intermediate distillation objective, lpred for prediction layer objective, a for data augmentation and m for layer mapping strategy. Average importance for each component (across tasks) is listed in the legend. 5.1 Importance of Distiller Components Figure 2 illustrates that the objective function for intermediate distillation lpred has the highest individual importance out of all components in Distiller, and the combination of the intermediate distillation objective and layer mapping strategy has the highest joint importance. Thus these are the components one should most critically focus on when selecting a particular KD pipeline. One hypothetical explanation is that the teacher can provide token-level supervision to the student via intermediate distillation, which can better guide the learning process of the student. Our finding is also consistent with the previous observations (Sun et al., 2019; Li et al., 2020). 5.2 MI-α for Intermediate Distillation We submitted our MI-α model predictions to the official GLUE leaderboard to obtain test set results and also report the average scores over all tasks (the “AVG” column) as summarized in Table 1. The results show that the student model distilled QQP (364k) 71.2 71.7 70.2 69.3 69.8 QNLI (108k) 90.5 90.9 85.1 87.2 87.4 SST-2 (67k) 93.5 93.4 89.4 91.0 91.5 CoLA (8.5k) 52.1 49.3 24.8 25.6 25.9 MRPC (3.5k) 88.9 87.0 82.6 87.6 87.0 RTE (2.5k) 66.4 67.3 62.3 66.2 67.4 STS-B (5.7k) 85.8 84.7 82.3 82.3 84.0 AVG 79.6 79.2 72.9 74.7 75.1 Table 2: Ablation study of distillation performance on SQuAD v1.1 dev set. The first line shows the performance of the BERTBASE teacher. ELECTRASMALL , TinyBERT6 and TinyBERT4 are three student networks. ELECTRASMALL (FT) means to fine-tune without KD. TinyBERT6 † and TinyBERT4 † are results obtained from Jiao et al. (2020). Models that end with “(MSE)” are trained with the MSE loss. “+ MI-α” means to distill the student with MI-α (α=0.9) as the intermediate loss function. “+ mixup” means to further apply the mixup augmentation. Model BERTBASE (T) ELECTRASMALL (FT) ELECTRASMALL (MSE) + MI-α + mixup, MSE + mixup, MI-α TinyBERT6 † (Jiao et al., 2020) TinyBERT6 (MSE) + MI-α + mixup, MSE + mixup, MI-α TinyBERT4 † (Jiao et al., 2020) TinyBERT4 (MSE) + MI-α + mixup, MSE + mixup, MI-α SQuAD v1.1 EM F1 80.9 88.2 75.3 83.5 79.2 86.8 79.0 86.8 80.1 87.4 80.2 87.6 79.7 87.5 77.8 85.5 80.0 87.8 78.6 86.2 81.1 88.6 72.7 82.1 72.7 81.7 71.7 81.2 72.4 81.4 71.9 82.5 via the MI-α objective function outperforms previous student models distilled via MSE or PKD loss. To further verify the effectiveness of MI-α, we compare how different choices of linter affect the distillation performance. In detail, we first pick the top 5 best Distiller strategies according to the evaluation scores for every task and then count how many times each intermediate objective function appears in these strategies. Figure 3 shows that MI-α appears more frequently than all other objectives on both classification and regression tasks. And the results for SQuAD v1.1 in Table 2 indicate the MI-α also works well for sentence tagging. Figure 3: Intermediate objective functions used in the top-5 performing KD configurations on each dataset. The average count of each objective function is listed in the legend. Configurations are detailed in Appendix. Table 3: Student performance with(out) augmentation (augmenter initialized as CA+RA+mixup). We report the relative improvement for rows starting with “+ aug”. Model BERTBASE (T) TinyBERT6 + aug BERTMEDIUM + aug BERTSMALL + aug TinyBERT4 + aug ELECTRASMALL + aug BERTMINI + aug BERTTINY + aug #Params 110M 67M 41M 29M 14M 14M 11M 4M CoLA mcc 55.0 51.3 +0.1 44.1 +5.3 37.4 +5.0 23.6 +7.9 42.8 +16.2 11.2 +23.2 6.0 +6.6 MRPC f1/acc 89.6/85.0 92.5/89.7 -1.1/-1.8 89.3/84.8 -0.4/-0.7 86.8/80.6 +0.1/+0.8 88.9/83.8 +0.3/+0.0 88.3/83.8 +3.4/+3.7 86.1/80.1 +0.0/-0.1 83.2/73.3 +1.7/+3.7 RTE acc 65.0 75.5 -3.3 65.3 +4.4 64.6 +0.4 67.1 +2.2 66.4 +1.8 62.8 +3.3 60.0 +4.3 STS-B spearman/pearson 88.4/88.6 89.6/89.8 +0.2/+0.2 88.3/88.6 +0.6/+0.5 87.7/88.0 +0.3/+0.2 88.0/88.1 +0.7/+0.7 87.4/87.5 +1.0/+1.0 87.1/87.2 +0.2/+0.0 84.0/83.6 +0.1/+0.7 AVG 78.6 81.4 -1.0 76.7 +1.6 74.2 +1.1 73.3 +2.0 76.0 +4.5 69.1 +4.4 65.0 +2.9 5.3 Benefits of Data Augmentation Data augmentation in KD provides the student additional opportunities to learn from the teacher, especially for datasets of limited size. Thus, our experiments investigate the effect of DA on four data-limited tasks: CoLA, MRPC, RTE and STS-B. We also study whether students with different architectures/sizes benefit dissimilarly from DA. Table 3 demonstrates that DA generally provides a boost to student performance and is especially beneficial for small models (BERTMINI and BERTTINY ). 5.4 Performance of AutoDistiller Recall that we train our AutoDistiller performance prediction model on the previously collected experimental results via AutoGluon-Tabular. Once trained, AutoDistiller can recommend distillation pipelines for any down-stream dataset/task by fixing the dataset/task features and searching for configurations that maximize the predicted distillation ratio. AutoDistiller operates on features that represent the dataset domain, the task type, and the task complexity, which are detailed in Appendix. Figure 4: Distillation ratio of AutoDistiller (AD) top-5 KD strategies vs. 5 randomly selected strategies, for a fine-tuned BERTBASE teacher and TinyBERT4 student. Higher ratio indicates better distillation performance. Mean and standard deviation of the four groups of ratios are listed in the legend. We evaluate the performance of AutoDistiller on the 8 GLUE datasets via a leave-one-dataset-out cross-validation protocol. Figure 5 in Appendix shows that AutoDistiller achieves positive Spearman’s correlation coefficients for most datasets. Finally, we applied AutoDistiller on two unseen datasets “BoolQ” (Wang et al., 2019a) and “cloth” (Shi et al., 2021) not considered in our previous experiments. We compared the distillation ratio obtained by each of the top-N strategies suggested by AutoDistiller with the distillation ratio from each of N randomly selected strategies. The best strategy suggested by AutoDistiller achieves accuracy of 74.2 on “BoolQ” and 70.1 on “cloth”, close or superior to the teacher performance (73.4 on “BoolQ” and “71.2” on cloth). Figure 4 shows that AutoDistiller significantly outperforms random search, indicating its promise for automated KD. 6 Conclusion We provided a systematic study of KD algorithms in NLP to understand the importance of different components in the KD pipeline for various NLP tasks. Using data collected from our study, we fit a AutoDistiller to predict student performance under each KD pipeline based on dataset features, which is helpful for automatically selecting which KD pipeline to use for a new dataset. Here we unified the existing intermediate distillation objectives as maximizing the lower bounds of MI, leading to a new MI-α objective based on tighter bounds, which performs better on many datasets. References David M Blei, Alp Kucukelbir, and Jon D McAuliffe. 2017. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):859–877. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel HerbertVoss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language models are few-shot learners. In Advances in Neural Information Processing Systems. Cristian Buciluǎ, Rich Caruana, and Alexandru Niculescu-Mizil. 2006. Model compression. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 535–541. Rasool Fakoor, Jonas W Mueller, Nick Erickson, Pratik Chaudhari, and Alexander J Smola. 2020. Fast, accurate, and simple models for tabular data via augmented distillation. In Advances in Neural Information Processing Systems, volume 33. Hongyu Guo, Yongyi Mao, and Richong Zhang. 2019. Augmenting data with mixup for sentence classification: An empirical study. arXiv preprint arXiv:1905.08941. Manish Gupta and Puneet Agrawal. 2020. Compression of deep learning models for text: A survey. arXiv preprint arXiv:2008.05221. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2014. Distilling the knowledge in a neural network. In NIPS 2014 Deep Learning Workshop. Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. 2020. DynaBERT: Dynamic bert with adaptive width and depth. In Advances in Neural Information Processing Systems. Frank Hutter, Holger Hoos, and Kevin Leyton-Brown. 2014. An efficient approach for assessing hyperparameter importance. In ICML, pages 754–762. PMLR. Daoyuan Chen, Yaliang Li, Minghui Qiu, Zhen Wang, Bofang Li, Bolin Ding, Hongbo Deng, Jun Huang, Wei Lin, and Jingren Zhou. 2020. Adabert: Taskadaptive bert compression with differentiable neural architecture search. In IJCAI. Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020. TinyBERT: Distilling BERT for natural language understanding. In EMNLP. Kevin Clark, Minh-Thang Luong, Quoc V Le, and Christopher D Manning. 2020. ELECTRA: Pretraining text encoders as discriminators rather than generators. In ICLR. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361. Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. 2019. AutoAugment: Learning augmentation strategies from data. In CVPR, pages 113–123. Taehyeon Kim, Jaehoon Oh, NakYil Kim, Sangwook Cho, and Se-Young Yun. 2021. Comparing kullback-leibler divergence and mean squared error loss in knowledge distillation. In IJCAI. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL. Lingpeng Kong, Cyprien de Masson d’Autume, Wang Ling, Lei Yu, Zihang Dai, and Dani Yogatama. 2020. A mutual information maximization perspective of language representation learning. In ICLR. Monroe D Donsker and SR Srinivasa Varadhan. 1975. Asymptotic evaluation of certain markov process expectations for large time, i. Communications on Pure and Applied Mathematics, 28(1):1–47. Jianquan Li, Xiaokang Liu, Honghong Zhao, Ruifeng Xu, Min Yang, and Yaohong Jin. 2020. BERTEMD: Many-to-many layer mapping for bert compression with earth mover’s distance. In EMNLP. Sergey Edunov, Myle Ott, Michael Auli, and David Grangier. 2018. Understanding back-translation at scale. arXiv preprint arXiv:1808.09381. Kevin J Liang, Weituo Hao, Dinghan Shen, Yufan Zhou, Weizhu Chen, Changyou Chen, and Lawrence Carin. 2021. MixKD: Towards efficient distillation of large-scale language models. In ICLR. Nick Erickson, Jonas Mueller, Alexander Shirkov, Hang Zhang, Pedro Larroy, Mu Li, and Alexander Smola. 2020. Autogluon-tabular: Robust and accurate automl for structured data. arXiv preprint arXiv:2003.06505. Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. RoBERTa: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692. Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. 2020. Improved knowledge distillation via teacher assistant. In AAAI. Levy, and Samuel R Bowman. 2019a. Superglue: A stickier benchmark for general-purpose language understanding systems. arXiv preprint arXiv:1905.00537. XuanLong Nguyen, Martin J Wainwright, and Michael I Jordan. 2010. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11):5847–5861. Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. 2019b. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In ICLR. Aaron van den Oord, Yazhe Li, and Oriol Vinyals. 2018. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748. Jason Wei and Kai Zou. 2019. EDA: Easy data augmentation techniques for boosting performance on text classification tasks. In EMNLP. Jeffrey Pennington, Richard Socher, and Christopher D. Manning. 2014. Glove: Global vectors for word representation. In EMNLP, pages 1532–1543. Jing Yang, Brais Martinez, Adrian Bulat, and Georgios Tzimiropoulos. 2021. Knowledge distillation via softmax regression representation learning. In ICLR. Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. 2019. On variational bounds of mutual information. In ICML. Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. 2018. Mixup: Beyond empirical risk minimization. In ICLR. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2019. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv:1910.10683. Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. SQuAD: 100,000+ questions for machine comprehension of text. In EMNLP. Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108. Xingjian Shi, Jonas Mueller, Nick Erickson, Mu Li, and Alex Smola. 2021. Multimodal automl on structured tables with text fields. In 8th ICML Workshop on Automated Machine Learning (AutoML). Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019. Patient knowledge distillation for BERT model compression. In EMNLP. Michael Tschannen, Josip Djolonga, Paul K Rubenstein, Sylvain Gelly, and Mario Lucic. 2020. On mutual information maximization for representation learning. ICLR. Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv:1908.08962. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in neural information processing systems, pages 5998–6008. Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Appendix – Distiller: A Systematic Study of Model Distillation Methods in Natural Language Processing A Relationship with Other Distillation Methods Distiller is a generic meta-framework that encompasses various KD pipelines used in previous work. For example, Distiller with the following configurations corresponds to the KD pipeline used in each of the cited works: lpred = CE, linter = MSE, mi,j = Skip, a = CA (Jiao et al., 2020); lpred = CE, linter = MSE, mi,j = EMD (Li et al., 2020); lpred = CE, a = Mixup (Liang et al., 2021). B Proof of Theorem 3.1 We denote the Mutual Information (MI) between two random variables X and Y as I(X; Y ). Based on the results on variational bounds of MI (Poole et al., 2019), we derived that optimizing common knowledge distillation objectives, including Mean Squared Error (MSE), L2 distance, and cosine similarity between X and Y , can be viewed as maximizing certain lower bounds of I(X; Y ). Lemma B.1 (ITUBA ) . Assume that f (x, y) is an arbitrary neural network that takes x and y as inputs and outputs a scalar and a(y) > 0. The lower bound of x and y can be estimated by: I(X; Y ) ≥ Ep(x,y) [log ef (x,y) ef (x,y) ] − Ep(x)p(y) [ ] a(y) a(y) , ITUBA Proof. Based on the definition of MI, we have: p(x|y) ] p(x) p(y|x) = Ep(x,y) [log ] p(y) I(X; Y ) = Ep(x,y) [log (5) Replacing the intractable conditional distribution p(x|y) with a tractable variational distribution q(x|y) yields a lower bound on MI due to the nonnegativity of the KL divergence: where H(X) is the entropy of X. We choose an energy-based variational family that uses a critic f (x, y) and scaled by the data density p(x): p(x) f (x,y) e , Z(y) = Ep(x) [ef (x,y) ] Z(y) (7) Substituting this distribution into (6) gives a lower bound on MI: q(x|y) = I(X; Y ) ≥ Ep(x,y) [f (x, y)] − Ep(y) [log Z(y)] However, this objective is still intractable. To form a tractable bound, we can upper bound the log partition function by this inequality: log(x) ≤ xa + log(a) − 1 for all x, a > 0. Apply this inequality to get: I(X; Y ) ≥ Ep(x,y) [f (x, y)] − Ep(y) [log Z(y)] ≥ Ep(x,y) [f (x, y)] − Ep(y) [ Ep(x) [ef (x,y) ] = Ep(x,y) [log q(x|y) ] + Ep(y) [KL(p(x|y)||q(x|y))] p(x) (6) ≥ Ep(x,y) [log q(x|y)] + H(X) + log a(y) − 1] (8) = Ep(x,y) [f (x, y)] − Ep(x,y) [log a(y)] ef (x,y) ]+1 − Ep(x)p(y) [ a(y) ≥ Ep(x,y) [log ef (x,y) ef (x,y) ] − Ep(x)p(y) [ ] a(y) a(y) This bound holds for any a(y) > 0 Theorem 1. Minimizing the MSE, L2, PKD loss, and maximizing the cosine similarity between two random variables x,y can be viewed as maximizing lower bounds of I(X; Y ). In knowledge distillation, x and y are hidden states generated by student model and teacher model. Proof. We prove this theorem by constructing f (x, y) and a(y) in Lemma B.1 for each loss function. MSE LMSE (x, y) = ||x − y||22 , let f (x, y) = −||x − y||22 , a(y) = 1, we have: I(X; Y ) ≥ Ep(x,y) [log e ≥ Ep(x,y) −||x−y||2 2 ] − Ep(x)p(y) [e −||x−y||2 2] [log e = Ep(x,y) [log e p(x|y) q(x|y) ] + Ep(x,y) [log ] I(X; Y ) = Ep(x,y) [log p(x) q(x|y) a(y) −||x−y||2 2 −||x−y||2 2 ] 0 − Ep(x)p(y) [e ] ]−1 2 = Ep(x,y) [−||x − y||2 ] − 1. Thus, minimizing the MSE loss between x and y can be viewed as maximizing the lower bound of I(X; Y ). y x PKD Loss LPKD (x, y) = || ||x|| − ||y|| ||2 , let 2 2 2 y x − ||y|| ||2 , we have: f (x, y) = −|| ||x|| 2 2 2 I(X; Y ) ≥ Ep(x,y) [log e − Ep(x)p(y) [e −|| −|| ≥ Ep(x,y) [log e y x − ||2 ||x||2 ||y||2 2 ] y x − ||2 ||x||2 ||y||2 2 ] −|| y x − ||2 ||x||2 ||y||2 2 ] 0 − Ep(x)p(y) [e ] (9) y x − ||2 −|| ||x||2 ||y||2 2 ] Ep(x,y) [log e −1 x y 2 = Ep(x,y) [−|| − || ] − 1. ||x||2 ||y||2 2 = Thus, minimizing the PKD loss between x and y can be viewed as maximizing the lower bound of I(X; Y ). L2 Loss LL2 (x, y) = ||x − y||2 , let f (x, y) = −||x − y||2 , a(y)=1, we have: −||x−y||2 I(X; Y ) ≥ Ep(x,y) [log e −||x−y||2 ] − Ep(x)p(y) [e ] −||x−y||2 ]−1 ≥ Ep(x,y) [log e = Ep(x,y) [log e ] − Ep(x)p(y) [e −||x−y||2 ] 0 = Ep(x,y) [−||x − y||2 ] − 1 Thus, minimizing L2 loss between x and y is equivalent to maximizing the lower bound of I(X, Y ). Cosine Similarity The cosine similarity between two hidden states x and y is calculated as x·y x·y ||x||×||y|| , let f (x, y) = ||x||×||y|| − 1, a(y) = 1 x·y I(X, Y ) ≥ Ep(x,y) [log e ||x||×||y|| ≥ Ep(x,y) −1 x·y ] − Ep(x)p(y) [e ||x||×||y|| x·y −1 [log e ||x||×||y|| ] −1 D Benefits of Intermediate Distillation Table 5 shows the experimental results of distillation with(out) intermediate distillation across different tasks/datasets. All the experiments are trained 20 epochs and the intermediate distillation constitutes of “Skip” as intermediate mapping strategy and CE as intermediate distillation object. ] 0 − Ep(x)p(y) [e ] x·y [log e ||x||×||y|| 1 Ep(x,y) ]−1 e x·y 1 ]−1 = Ep(x,y) [ e ||x|| × ||y|| = Figure 5: Evaluating held-out AutoDistiller predictions on GLUE via leave-one-out estimates. We use datasetlevel cross-validation that holds out each GLUE dataset from AutoDistiller training. For each held-out dataset, the legend lists Spearman’s rank correlation between the predicted vs. actual distillation ratio across different KD pipelines. The average Spearman’s rank correlation value across the 8 datasets is 0.33. Thus, maximizing cosine similarity between x and y can be viewed as maximizing mutual information between x and y. E Dataset Embedding As features for the AutoDistiller model, we extract features from datasets and represent them as a fixed dimension embedding. Details are shown in Table 7. C Architecture of Teacher and Student Networks In Table 1, we use two baseline models BERT-PKD4 and BERT-EMD4 . As described in the original paper (Sun et al., 2019), we initialize BERT-PKD4 with the first 4 layers of parameters from pretrained BERTBASE . BERT-EMD4 is initialized from TinyBERT4 so they have the same architecture. We list the detailed configurations of the teacher and student architectures investigated in our study in Table 6. Figure 6: Importance of interactions between hyperparameters. Table 4: Comparison of evaluation results on GLUE test set. We compare the distillation performance when using MLP and Transformer as critic functions in MI-α respectively. BERTBASE (T) indicates the teacher model trained by ourselves. Both of the students are initialized with TinyBERT4 (Jiao et al., 2020) and distilled with “Skip” as intermediate layer mapping strategy and MI-α as intermediate objective functions. The difference is, TinyBERT4 (MLP) is trained with a 4-layer MLP with hidden state of 512 as critic function while TinyBERT4 (Transformer) uses a 2-layer Transformer with feed-forward hidden size 256. The result shows that a small Transformer architecture performs as a better critic function than an MLP in MI-α especially when the task is a token-level task (SQuAD v1.1). Model BERTBASE (T) TinyBERT4 (MLP) TinyBERT4 (Transformer) MNLI-m (393k) 84.5 81.7 81.9 MNLI-mm (393k) 83.6 80.6 80.6 QQP (364k) 71.7 69.7 69.8 QNLI (108k) 90.9 87.6 87.4 SST-2 (67k) 93.4 91.6 91.5 CoLA (8.5k) 49.3 24.9 25.9 MRPC (3.5k) 87.0 87.0 87.0 RTE (2.5k) 67.3 67.4 67.4 STS-B (5.7k) 84.7 81.9 84.0 SQuAD v1.1 (108k) 88.2 70.1 71.7 AVG 80.0 74.3 74.7 Table 5: Comparison of evaluation results on GLUE dev set. We compare the distillation performance with(out) intermediate distillation. A BERTBASE model is used as the teacher and TinyBERT4 is the student. TinyBERT4 (KD) represents using a vanilla knowledge distillation (student only learns from the outputs of teacher) and “+intermediate distillation” represents using vanilla KD and intermediate distillation. Model BERTBASE (T) TinyBERT4 (KD) +intermediate distillation MNLI-m (393k) acc 84.1 80.1 80.7 MNLI-mm (393k) acc 84.7 80.3 81.3 QQP (364k) f1/acc 88.0/91.1 86.4/89.7 87.0/90.2 QNLI (108k) acc 91.7 85.8 86.8 SST-2 (67k) acc 93.0 89.1 90.0 Table 6: Network architectures of the teacher and student models used in the paper. Model RoBERTaLARGE ELECTRALARGE BERTBASE TinyBERT6 BERT-PKD4 BERTMEDIUM BERTSMALL TinyBERT4 ELECTRASMALL BERTMINI BERTTINY #Params hunits nlayers Teacher architectures 335M 1024 24 335M 1024 24 110M 768 12 Student architectures 67M 768 6 52M 768 4 41M 512 8 29M 512 4 14M 312 4 14M 256 12 11M 256 4 4M 128 2 hmid nheads 4096 4096 3072 16 16 12 3072 3072 2048 2048 1200 1024 1024 512 12 12 8 8 12 4 4 2 F How much does larger/better teacher help? Table 8 shows the performance of different students distilled from teachers of different sizes and pretraining schemes. From the results, we found that although the teacher ELECTRALARGE has the best performance on average score, most of the students of ELECTRALARGE performs worse than students of RoBERTaLARGE . ELECTRASMALL is the only student that performs the best with ELECTRALARGE as teacher, that may be attributed to ELECTRASMALL and ELECTRALARGE are pretrained on the same pretraining task, so they have a similar knowledge representation scheme. And CoLA (8.5k) mcc 55.0 16.1 21.3 MRPC (3.5k) f1/acc 89.6/85.0 89.6/85.3 89.3/84.8 RTE (2.5k) acc 65.0 66.8 65.3 STS-B (5.7k) spearman/pearson 88.4/88.6 88.4/88.5 88.2/88.4 SQuAD v1.1 (108k) f1/em 88.2/80.9 77.3/66.8 79.4/69.4 AVG 83.8 77.9 78.7 also, for datasets (MNLI, QQP, QNLI and SST-2) which have abundant amount of data, students of BERTBASE perform better. G Computing Details All the experiments are performed on a single machine with 4 NVIDIA T4 GPUs. For hyperparameters in Equation 1, our experiments suggest that setting β1 and β2 to 1 produces the best overall performance so we fix their values to 1 in subsequent results. γ1 and γ2 are set to 0.5 when DA is used (otherwise β2 , γ1 and γ2 are set to 0). For controlled experiments, unless specified explicitly, we fix linter as MI-α (α = 0.9) and the layer mapping as “Skip”. Critic functions in MI-α are powered by a two-layer Transformer(hmid = 256, nheads = 8) and the comparison of using Transformer and MLP to estimate critic functions is described in Table 4. To reduce the hyper-parameter search space, we fix the batch size as 16 and the learning rate as 5e-5 for all experiments. We used automated mix precision for training. Maximum sequence length is set to 128 for sentence-pair tasks in GLUE, 64 for single sentence tasks and 320 for SQuAD. Most of the experiments are trained for 20 epochs except 50 epochs for the challenging task CoLA and 30 epochs for SQuAD. Table 7: We extract features from downstream datasets so that every task can be represented as a fixed-dimension embedding. The extracted embedding can be fed into AutoDistiller as dense features. In this table, we describe how the embedding is acquired. Feature Context Embedding Task Embedding Baseline Score Teacher Score Number of Examples Description Every document can be represented as a weighted average of the GloVe vectors, where the weights are defined by the TF-IDF scheme. Each down-stream dataset is viewed as a “document” in the TF-IDF scheme. 1 IDFw vw , Precisely, the embedding of a dataset s is vs = |s| where |s| denotes the number of words in the dataset, 1+N is the inverse document frequency of word w. and IDFw := log 1+N w N is the total number of datasets, and Nw denotes the number of datasets containing w. Intuitively, this feature represents the content of datasets. For the dataset, we collect their literal descriptions, usually one or two sentences. Then aggregate GloVe vectors of every words in these sentences and get a description embedding. This feature is used to represent what is the NLP task and how the data is formatted. We use a lite Bi-LSTM model as the baseline model and finetune it on the down-stream dataset. This feature aims to measure the difficulty of each task by measuring how well a simple architecture can perform on the specific dataset. The fine-tuned teacher score on the dataset. Comparing teacher score to aforementioned baseline score tells how much boost can a complex model has on this dataset. How many training samples in the dataset. Table 8: Performance comparison with different teacher and student models. We abbreviate three teacher models BERTBASE , RoBERTaLARGE and ELECTRALARGE as B, R, E. Results are evaluated on GLUE dev set and best results are in-bold. Model #Params BERTBASE (T) RoBERTaLARGE (T) ELECTRALARGE (T) BERTBASE 110M 335M 335M 110M TinyBERT6 67M BERTMEDIUM 41M BERTSMALL 29M TinyBERT4 14M ELECTRASMALL 14M BERTMINI 11M BERTTINY 4M Teacher R E B R E B R E B R E B R E B R E B R E B R E MNLI-m acc 84.1 90.2 90.5 84.5 84.4 83.9 83.5 83.0 82.6 80.9 81.0 81.0 78.7 78.6 81.1 80.0 80.0 82.7 82.3 82.0 78.5 76.6 76.3 72.8 71.5 71.2 MNLI-mm acc 84.7 90.1 90.4 84.6 84.6 84.0 83.5 83.0 83.0 81.4 81.3 81.0 78.6 78.8 81.6 80.7 80.2 83.8 83.2 82.7 79.7 77.3 77.1 73.4 72.0 71.8 QQP f1 acc 88.0 91.1 89.6 92.1 90.3 92.8 88.6 91.5 88.8 91.7 88.1 91.2 88.0 91.2 87.8 91.0 87.9 91.0 87.6 90.8 87.5 90.7 87.4 90.6 87.0 90.4 87.2 90.5 87.2 90.4 86.5 90.0 86.2 89.6 87.8 90.9 88.1 91.2 88.5 91.5 86.6 90.0 86.0 90.0 85.9 89.5 83.5 87.2 82.9 86.7 82.9 87.0 QNLI acc 91.7 94.7 95.1 91.7 91.6 91.3 90.8 90.6 90.0 89.0 89.0 87.3 87.0 87.0 87.4 86.0 85.9 89.7 89.5 89.3 84.9 84.6 84.2 81.3 80.2 80.0 H Distiller Search Space • lpred ∈ {MSE, CE} • {mi,j } ∈ {Skip, Last, EMD} CoLA mcc 55.0 64.6 67.4 59.3 59.5 50.5 48.0 48.6 48.3 50.5 51.0 43.1 41.2 43.0 12.3 24.6 21.8 60.6 58.6 60.6 20.1 32.2 33.8 0.0 6.2 6.7 MRPC f1 acc 89.6 85.0 91.3 88.0 91.7 88.5 91.6 88.0 91.9 88.7 90.3 86.5 91.9 88.7 91.6 88.5 88.9 84.1 88.9 84.6 89.5 85.3 87.8 82.4 89.1 84.1 88.7 84.1 89.4 85.0 90.4 86.5 90.9 86.8 91.3 87.7 91.3 87.5 92.3 89.0 87.0 81.6 86.6 81.1 87.0 81.6 84.5 75.7 84.9 76.2 85.2 77.0 RTE acc 65.0 78.7 84.5 66.4 69.3 75.5 72.6 76.2 65.0 64.3 64.3 63.5 64.3 63.9 66.4 67.9 68.6 60.6 67.5 69.7 61.0 65.3 65.0 58.5 60.3 62.8 STS-B spearman pearson 88.4 88.6 91.7 91.8 88.7 88.9 89.0 89.4 89.1 89.6 89.4 89.4 89.9 90.0 89.1 89.3 88.2 88.4 88.2 88.6 88.0 88.2 87.0 87.2 87.1 87.3 86.8 87.1 87.7 87.8 88.0 88.1 87.6 87.6 87.4 87.5 87.6 87.8 86.6 86.7 86.2 86.1 86.2 86.3 85.7 85.5 81.4 79.9 81.9 79.9 78.4 77.3 AVG 83.7 88.3 88.8 84.8 85.2 84.3 84.2 84.2 82.4 82.2 82.2 80.7 80.3 80.4 78.9 79.8 79.5 83.5 83.8 84.2 77.5 78.2 78.1 71.8 72.1 71.9 • Augmentation policy a is one or combinations of elementary augmentation operations in {CA, RA, BT, Mixup}. Here we recap the full search space considered for each stage of the KD pipeline in Distiller: • linter ∈ {MSE, L2, Cos, PKD, MI-α (α = 0.1, 0.5, or 0.9)} SST-2 acc 93.0 96.3 96.6 93.2 92.8 91.6 92.2 91.3 90.8 91.4 90.9 90.5 88.6 89.3 90.6 89.4 88.9 91.3 90.6 91.4 87.8 85.9 85.9 83.9 83.1 82.8 I Top Configurations in Distiller Here we list the top 5 configurations from the Distiller search space that performed best on each dataset in Table 9. Table 9: Top 5 configurations to distill a BERTBASE teacher to a TinyBERT4 student on every dataset. To reduce the search space, we only compare configurations that don’t use data augmentation. As Hyper-parameter α is only valid for MI-α, the value of α is set to N for other intermediate loss in the table. Task MNLI MNLI MNLI MNLI MNLI QQP QQP QQP QQP QQP QNLI QNLI QNLI QNLI QNLI SST-2 SST-2 SST-2 SST-2 SST-2 CoLA CoLA CoLA CoLA CoLA MRPC MRPC MRPC MRPC MRPC RTE RTE RTE RTE RTE STS-B STS-B STS-B STS-B STS-B SQuAD v1.1 SQuAD v1.1 SQuAD v1.1 SQuAD v1.1 SQuAD v1.1 Intermediate Loss MI-α MI-α MI-α MSE MSE MI-α MI-α CE MI-α MI-α CE MI-α MI-α MI-α MI-α Cos MSE MI-α CE MI-α MI-α MI-α MI-α MI-α MI-α MI-α MI-α CE MI-α CE MI-α MI-α MI-α MI-α MI-α MI-α MI-α MI-α MI-α PKD MSE CE MSE MI-α MI-α α 0.9 0.1 0.5 N N 0.9 0.1 N 0.1 0.1 N 0.5 0.5 0.9 0.1 N N 0.1 N 0.9 0.1 0.5 0.5 0.1 0.5 0.1 0.5 N 0.9 N 0.1 0.9 0.5 0.1 0.5 0.9 0.5 0.9 0.1 N N N N 0.9 0.9 Layer Mapping Strategy EMD EMD Skip EMD Skip EMD EMD Skip EMD Skip Last Skip Last Skip Skip Last Skip Last Skip Skip EMD EMD EMD Last Skip EMD EMD Skip EMD Last Skip Skip Skip Last Last Skip Skip EMD Last Skip Skip Skip EMD Skip Skip KD Loss MSE MSE MSE MSE MSE MSE MSE MSE MSE MSE MSE MSE CE CE CE MSE CE CE MSE CE MSE MSE CE MSE MSE CE CE CE CE MSE CE CE CE CE CE MSE MSE MSE MSE MSE MSE MSE MSE CE CE #Example 393000 393000 393000 393000 393000 364000 364000 364000 364000 364000 105000 105000 105000 105000 105000 67000 67000 67000 67000 67000 8500 8500 8500 8500 8500 3700 3700 3700 3700 3700 2500 2500 2500 2500 2500 7000 7000 7000 7000 7000 130000 130000 130000 130000 130000 Score 81.7 81.7 81.7 81.6 81.6 90.2 90.2 90.2 90.1 90.1 87.4 87.4 87.3 87.2 87.1 90.6 90.5 90.3 90.3 90.3 22.3 21.6 21.1 21.1 21.0 90.3 90.2 89.9 89.9 89.7 70.8 70.4 70.0 70.0 69.3 88.0 88.0 87.9 87.9 87.9 72.6 72.4 72.3 71.9 71.7