MixKD: Towards Efficient Distillation of Large-scale Language Models
Authors: Kevin J Liang
Weituo Hao
Dinghan Shen
Yufan Zhou
Weizhu Chen
Changyou Chen
Lawrence Carin
Venue: ICLR 2021 Poster (2021)
Keywords: Natural Language Processing
Representation Learning

Abstract: Large-scale language models have recently demonstrated impressive empirical performance. Nevertheless, the improved results are attained at the price of bigger models, more power consumption, and slower inference, which hinder their applicability to low-resource (both memory and computation) platforms. Knowledge distillation (KD) has been demonstrated as an effective framework for compressing such big models. However, large-scale neural network systems are prone to memorize training instances, and thus tend to make inconsistent predictions when the data distribution is altered slightly. Moreover, the student model has few opportunities to request useful information from the teacher model when there is limited task-specific data available. To address these issues, we propose MixKD, a data-agnostic distillation framework that leverages mixup, a simple yet efficient data augmentation approach, to endow the resulting model with stronger generalization ability. Concretely, in addition to the original training examples, the student model is encouraged to mimic the teacher's behavior on the linear interpolation of example pairs as well. We prove from a theoretical perspective that under reasonable conditions MixKD gives rise to a smaller gap between the generalization error and the empirical error. To verify its effectiveness, we conduct experiments on the GLUE benchmark, where MixKD consistently leads to significant gains over the standard KD training, and outperforms several competitive baselines. Experiments under a limited-data setting and ablation studies further demonstrate the advantages of the proposed approach.

Published as a conference paper at ICLR 2021
MIXKD: TOWARDS EFFICIENT DISTILLATION OF LARGE-SCALE LANGUAGE MODELS
Kevin J Liang': 2. Weituo Hao" DDnghan Shen', Yufan Zhou', Weizhu Chen
Changyou Chen', Lawrence Carin' DDuke University "Facebook AI Microsoff Dynamics 365 Al
'State University of New York at Buffalo
(kevin. liang, weituo. hww..aahhoaaae hho.eeuke..dde edu
ABSTRACT
Large-scale language models have recently demonstrated impressive empirical performance. Nevertheless, the improved results are attained at the price of bigger models, more power consumption, and slower inference, which hinder their appli- cability to low-resource (both memory and computation) platforms. Knowledge distillation (KD) has been demonstrated as effective framework for compress- ing such big models. However, large-scale neural network systems are prone to memorize training instances, and thus tend to make inconsistent predictions when the data distribution is altered slightly. Moreover, the student model has few op- portunities to request useful information from the teacher model when there limited task-specific data available. To address these issues, we propose MixKD, a data-agnostic distillation framework that leverages mixup, a simple yet efficient data augmentation approach, to endow the resulting model with stronger gener- alization ability. Concretely, in addition to the original training examples, the student model is encouraged mimic the teacher's behavior on the linear inter- polation of example pairs as well. We prove from a theoretical perspective that under reasonable conditions MixKI) gives rise to a smaller gap between the gen- eralization error and the empirical error. To verify its effectiveness, we conduct experiments on the GLUE benchmark, where MixKD consistently leads signif- icant gains over the standard KD training, and outperforms several competitive baselines. Experiments under a limited-dae setting and ablation studies further demonstrate the advantages of the proposed approach.
INTRODUCTION
Recent language models (LM) pre-trained on large-scale unlabeled text corpora in self-supervised manner have significantly advanced the state of the art across a wide variety of natural language
processing (NLP) tasks (Devlin al., 2018 Liu et 2019c Yang et al. 2019 Joshi et al. 2020 Sun et al. 2019b; Clark et al.) 2020 Lewis et al.) 2019 Bao et al. 2020). After the LM pre- training stage, the resulting parameters can fine-tuned to different downstream tasks. While these models have yielded impressive results, they typically have millions, if not billions, of parameters, and thus can be very expensive from storage and computational standpoints. Additionally, during deployment, such large models require a lot of time to process where computation may be limited (e.g. mobile, edge devices), such charactcristics may preclude such powerful models from deployment entirely.
single sample. In settings
One promising strategy to compress and accelerate large-scale language models is knowledge distil- lation (Zhao et al.) 2019; Tang et al. 2019: Sun al. 2020). The key idea is train smaller model (a "student") to mimic the behavior of the larger, stronger-performingg but perhaps less practical model (the "teacher"), thus achieving similar performance with faster, lighter-weight model. A simple but powerful method of achieving this is to use the output probability logits produced by the teacher model as soft labels for training the student (Hinton al. 2015). With higher entropy than one-hot labels, these soft labels contain more information for the student model to learn from.
"Equal contribution

---

Published as a conference paper at ICLR 2021
Previous efforts distilling large-scale LMs mainly focus designing better training objectives, such as matching intermediate representations (Sun et al. 2019a Mukherjee Awadallah 2019), learning multiple tasks together (Liu et al. 2019a), leveraging the distillation objective during the pre-training stage (Jiao et al. 2019 Sanh et al. 2019). However, much less effort has been made to enrich task-specific data, a potentially vital component of the knowledge distillation procedure. In particular, tasks with fewer data samples provide less opportunity for the student model to learn from the teacher. Even with a well-designed training objective, the student model is still prone to overfitting, despite effectively mimicking the teacher network on the available data.
In response to these limitations, we propose improving the value of knowledge distillation by us- ing data augmentation to generate additional samples from the available task-specific data. These augmented samples are further processed by the teacher network to produce additional soft labels, providing the student model more data to learn from large-scale L.M. Intuitively, this is akin to student learning more from a teacher by asking more questions to further probe the teacher's an- swers and thoughts. In particular, we demonstrate that mixup (Zhang et al., 2018 can significantly improve knowledge distillation's ' effectiveness, and we show with a theoretical framework why this is the case. We call our framework MixKD.
We conduct experiments on 6 GLUE datasets (Wang et al. 2019) across a variety of task types, demonstrating that MixKD significantly outperforms knowledge distillation (Hinton et al. 2015 and other previous methods that compress large-scale language models. In particular, we show that our method is especially effective when the number of available task data samples is small, substantially improving the potency of knowledge distillation. We also visualize representations learned with and without MixKD to show the value of interpolated distillation samples, perform series of ablation and hyperparameter sensitivity studies, and demonstrate the superiority of MixKD over other BERT data augmentation strategies.
RELATED WORK
2.1 MODEL COMPRESSION
Compressing large-scale language models, such BERT, has attracted significant attention recently. Knowledge distillation has been demonstrated as an effective approach, which can be leveraged during both the pre-training and task-specific fine-tuning stages. Prior research efforts mainly focus on improving the training objectives to benefit the distillation process. Specifically. Tuuc et al. (2019) advocate that task-specific knowledge distillation can be improved by first pre-traiming the student model. It is shown by Clark et al. (219)) that a multi-task BERT model can be learned by distilling from multiple single-task teachers. Liu et al. (2019b) propose learning a stronger student model by distilling knowledge from ensemble of BERT models. Patient knowledge distillation (PKD), introduced by |Snn et al. (2019a), encourages the student model to mimic the teacher's intermediate layers in addition to output logits. DiRT (anh et al. 2019) reduces the depth of BERT model by a factor of via knowledge distillation during the pre-training stage. In this work, we evaluate MixKD to the pre-training stage as well, which we leave for future work. Moreover, our method can be flexibly integrated with different KD training objectives (described above) to obtain even better results. However, we utilize the BERT-base model as the testbed in this paper without loss of generality.
the case of task-specific knowledge distillation. Notably, it can be extended
2.2 DATA AUGMENTATION IN NLP
Data augmentation (DA) has been studied extensively in computer vision as a powerful technique to incorporate prior knowledge of invariances and improve the robustness of learned models (Simard)
al. 1998 2003 Krizhevsky et al. 2012). Recently, it has also been applied and shown effective on natural language data. Many approaches can be categorized as label- aabel-preserving transformations, which essentially produce neighbors around a training example that maintain its original label. For example, EDA (Wei & Zou 2019) propose using various rule-based operations such as synonym replacement, word insertion, swap or deletion to obtain augmented samples. Back-translation (Yu) et al. 2018 Xie et al., 2019) is another popular approach belonging to this type, which relies on pre-trained translation models. Additionally, methods based on paraphrase generation have also been leveraged from the data augmentation perspective (Kumar et al. 2019). On the other hand, label-altering techniques like mixup (Zhang et al., (2018) have also been proposed for language (Guo)
et al. 2019 Chen et al. 2020), producing interpol lated inputs and labels for the models predict. The
2

---

Published as a conference paper at ICLR 2021
proposed MixKD framework leverages the ability of mixup facilitate the student learning more information from the teacher. It is worth noting that MixKD can be combined with arbitrary label- preserving DA modules. trckstranslttion is employed special case here, and we believe other advanced labee-pge transformations developed in the future can benefit the MixKD approach well.
2.3 MIXUP
Mixup Zhang et al. 2018 is popular data augmentation strategy increase model generalizability and robustness by training on convex combinations of pairs of inputs and labels (Ii ys) and (az. M5 )
x' Az; + (1 A)æ3
y = Ay: + (1 X)yı
(2)
with À [0, 1] and (x',8/) being the resulting virtual training example. This concept of interpo- lating samples was later generalized with Manifold mixup (Verma et al., (019a) and also found to be effective in semi-supervised learning settings (Verma et al.j 2019b/c Berthelot et al.] [209bbaa) Other strategies include mixing together samples resulting from chaining together other augmen- tation techniques (Hendrycks et al. 2020), or replacing linear interpolation with the cutting and pasting of patches (Yun et 2019).
METHODOLOGY
3.1 PRELIMINARIES
In NLP, input sample is often represented as vector of tokens w; {[ui, W1,2, ... Wi,)), with each token Wit E RV one-hot vector often representing words (but also possibly subwords, punetuation, or special tokens) and V being the vocabulary size. These discrete tokens are then mapped word embeddings x; {x,,1, Zi....... T, which serve input to the machine learning model f. For supervised classification problems, one-hot label yı E IR indicates the ground-truth class of x; out of possible classes. The parameters 0 f are optimized with some form of stochastic gradient descent so that the output of the model f fx,) IRC is as close to y/1 as possible, with cross cross-entropy as the most common loss function:
LMIE Σ yi (x;()(x:)
(3)
where n is the number of samples, and is the dot product.
3.2 KNOWLEDGE DISTILLATION FOR BERT
Consider models f and 9 parameterized by 0r and As, respectively, with |0r| >>> |0sl- Given enough training data and sufficient optimization, f is likely to yield better accuracy than g. due to higher modeling capacity, but may be too bulky slow for certain applications. Being smaller in size, g is more likely to satisfy operational constraints, but its weaker performance can be seen as disadvantage. To improve y. can use the output prediction (xx: input x; as extra supervision for to learn from, seeking to match (xx.) with fxx.). Given these roles, we refer to the student model and f as the teacher model.
While there are a number of recent large-scale language models driving the state of the art, we fo- cus here on BERT (Devlin et al., 2018) models. Following [Snn et al. (2019a), we use the notation BERT; & to indicate BERT model with & Transformer (Vaswani et al. 2017) layers. While pow- erful, BERT models also tend be quite large; for example, the default bert-base--ee-e.- -aase-uucasseeeeee (BERT 12) has ~ 100 parameters. Reducing the number of layers (e.g. using BERT;) makes such models significantly more portable and efficient, but at the expense of accuracy. With a knowledge distillation set-up, however, aim to reduce this loss in performance.
3.3 MIXUP DATA AUGMENTATION FOR KNOWLEDGE DISTILLATION
While knowledge distillation can be powerful technique, if the size of the available data is small, then the student has only limited opportunities to learn from the teacher. This may make it much harder for knowledge distillation to close the gap between student and teacher model performance. To correct this, we propose using data augmentation for knowledge distillation. While data aug- mentation Yuu et al. 2018 Xie et al. 2019 Yun et al. 2019 Kumar al. 2019 Hendryeks al.
3

---

Published as a conference paper at ICLR 2021
2020 Shen al. 2020 Qu al. 2020) is commonly used technique across machine learning for increasing training samples, robustness, and overall performance, a limited modeling capacity constrains the representations the student is capable of learning on its own. Instead, we propose using the augmented samples to further query the teacher model, whose large size often allows it to learn more powerful features.
While many different data augmentation strategies have been proposed for NLP, we focus on mixup (Zhagg et al. 2018) for generating additional samples learn from the teacher. Mixup's vicinal risk minimization tends to result in smoother decision boundaries and better generalization. while also being cheaper to compute than methods such as backtranslation (Yu et al., 2018 Xie et al.) 2019). Mixup was initially proposed for continuous data, where interpolations between data points remain in-domain; its efficacy was demonstrated primarily on image data, but examples in speech recognition and tabular data were also shown to demonstrate generality.
Directly applying mixup to NLP is not quite as straightforward as it is for images, as language commonly consists of sentences of variable length, each comprised of discrete word tokens. Since performing mixup directly on the word tokens doesn't 't result in valid language inputs, we instead perform mixup on the word embeddings at each time step Tg (Guo et al., (019). This can be
interpreted as a special case of Manifold mixup Verma et al. (2019a), where the mixing layer is set to the embedding layer. In other words, mixup samples are generated as:
,/,,- = AZiie + (1 A)a5, 11 3/ = Ay + (1-A)y
(4)
(5)
with À E [0, l]; random sampling of X from a Uniform or Beta distribution are common choices. Note that we index the augmented sample with regardless of the value of X. Sentence length variability can be mitigated by grouping mixup pairs by length. Alternatively, padding is common technique for setting a consistent input length across samples; thus, if x(() contains more word tokens than x6), then the extra word embeddings are mixed up with zero paddings. We find this approach be effective, while also being much simpler implement.
We query the teacher model with the generated mixup sample xt, producing output prediction f(x)). The student encouraged to imitate this prediction on the same input, by minimizing the objective:
LTMKD d(j d(f(((),g(x)))
where .(: is a distance metric for distillation, with temperature-adiuts adjusseed cross-entropy and mean square error (MSE) being common choices.
Since have the mixup samples already generated (with an easy-to- generate casy-to-generate interpolated pseu- dolabel yy), we can also train the student model on these augmented data samples in the usual way, with cross-entropy objective:
LsM Σ 3 log(g(x())
Our final objective for MixKD is sum of the original data cross-entropy loss, student cross-entropy loss on the mixup samples, and knowledge distillation from the teacher on the mixup samples:
L LMLE + CSMLSM @TMMKDTTоK
where aSM and @TMKD are hyperparameters weighting the loss terms.
3.4 THEORETICAL ANALYSIS
We develop a theoretical foundation for the proposed framework. We wish to prove that by adopting data augmentation for knowledge distillation, one can achieve i) a smaller gap between generaliza- tion error and empirical error, and ii) better generalization.
To this end, assume the original training data {x; ) are sampled i.i.d. from the true data distri- bution p(x), and the augmented data distribution by mixup is denoted as q(x) (apparently and are dependent). Let ff be the teacher function, and g E G be the learnable student function. Denote the loss function to learn g 1(. The population risk w.r.t. p(x) is defined as R(f,9. p)
'This is essentially the same as C. in equation 8 the data-wise arguments fxx and g(x).
We use different notation [[f(x). explicitly spell
4

---

Published as a conference paper at ICLR 2021
Exxp(xx ((xxx){(x))) and the empirical risk as Remp(f.8 (x.)/=) Σ.
(((()))
A classic statement for generalization is the following: with at least 1 - probability, we have
R 9p. p) - RemppSs90 App (x:)==)) e,
where e > O, and we have used 9p to indicate that the function learned based on p(x). Note different training data would correspond
different error in equation We cp denote the minimum value over all e's satisfying equation ( Similarly, we can replace p with q, and {x:]==1 with {x:}i=1 {x]]== in equation in the data-augmentattin case. In this case, the student function is learned based on both the training data and augmented data, which we denote as g. Similarly, we also have a corresponding minimum error. which we denote as e*. Consequently, our goal of better generalization corresponds to proving R(f.g". (' s R(f. 9p. p). and the goal of a smaller gap corresponds to proving S 6p- In our theoretical results, we will give conditions when these goals are achievable. First, we consider the following three cases about the joint data X A
{x; 1: U [x fi and the function class g:
Case 1:
There exists distribution P such that X are i..d. samples from i G is finite set.
Case 2: There exists ý such that X are i.i.d. samples from it; G is an infinite set.
Case 3: There does not exist a distribution such that X are i..d. samples from it.
Our theoretical results are summarized in Theorems which state that with enough augmented data, method can achieve smaller generalization
Proofs given in the Appendix.
Theorem Assume the loss function 1(:;)) constant > O such that if
) is upper bounded by M > Under Case I, there exists
M log(191/0))
b
then
S
where e" and Ep denote the minimal generalization gaps one can achieve with or without aug- mented data, with least 1 probability. If further assuming a better empirical risk with data augmentation (which is usually the case in practice), iee, Remp(Js g"" {x:] U (x())=)) Remplf ffap(f.pp: Ip: )a we have
R(S.9" S R(S.9p.P)
Theorem Assume the loss function 1(,; .) is upper bounded by M and Lipschitz continuous. Fix the probability parameter Under Case 2, there exists a constant e. >> 0 such that if
M log(1/6)
b 62
then
S
where e* and en denote the minimal generalization gaps one can achieve with or without augmented data, with at least å probability. If further assuming better empirical risk with data augmen- tation, i.e. Remp (f g. {x1)=1 U (x()))) I Remp [f Ips (xa 1 ) we have
RSS
R( 9p.2))
more interesting setting Case 3. Our result is based on Baxter 2000), which studies learning from different and possibly correlated distributions.
Theorem Assume the loss function 1(:,)) is upper bounded. Under Case 3, there exists constants C1, C2 C3 > 0 such that if
a log(4/8) and a 2
aja C2
We make such assumption because x; and x dependent, thus existence of ő unknown.

---

Published as a conference paper at ICLR 2021
Model
SST-2
MRPC
QQP
MNLI-m
QNLI
RTE
BERT 12
92.20
90.53/86.52
88.21/91.25
84.12
91.32
77.98
D
91.3
87.5/82.4
/88.5
82.2
89.2
59.9
BERT
90.94
88.54/83.82
87.16/90.43
81.28
88.25
66.43
B
91.63
88.93/83.82
86.60/90.27
81.49
88.71
65.34
B
91.17
89.30/84.31
87.19/90.56
82.02
88.63
65.34
B
91.74
89.60/84.80
87.06/90.39
82.10
87.68
67.51
BERTg-TMMKDBBY
91.86
89 89.52/84.56
87.15/90.59
82.17
88.38
69.98
BERT3-SM-TMKD+BB
92.09
89.22/84.07
87.57/90.78
82.53
88.82
67.87
3ERT
87.16
81.68/71.08
84.99/88.65
75.55
83.98
58.48
-TRKg
88.76
81.62/71.0
83.27/87.80
75.73
84.26
58.48
BERT3-MM+TMKD
88.99
81.73/71.08
84.47/88.37
75.52
84.24
59.57
3-RT3TFT+BT
88.88
83.36/74.26
85.31/88.81
76.88
83.67
59.21
89.79
84.46/75.74
85.17/89.00
77.19
84.68
62.82
BERT3-SM+TMKD-TT
90.37
84.14/75.25
85.56/89.09
77.52
84.83
60.65
Table I: GLUE dev set results. We report the results of our BERT teacher model, the 6-layer DistilBERT, and 3- and 6-layer MixKD student models with various ablations. DistilBERT results taken from Sanh et al. (2019). For MRPC and QQP, we report F1/Accuracy.
then
CS
where e" and denote the minimal generalization gaps one can achieve with or without augmented data, with least probability. If further assuming a better empirical risk with data augmen- tation, i.e, Remp (f.g", (x:); U ((x))_1) S Remp (f 9p: ) we have
R(S.9" S R(J P)
Remark 4 For Theorem hold, based Baxter (2000), it is enough ensure {xi, x²} and {x,:X,} to be independent for i + j. We achieve this by constructing x with x; and an extra random sample from the training data. Since all x)) and the extra random samples are inde- pendent, the resulting concatenation will also be independent.
EXPERIMENTS
We demonstrate the effectiveness of MixKD on a number of GLUE (Wagg et al. 2019) dataset tasks:
Stanford Sentiment Treebank (SST-2) (Socher et al.) (013). Microsoft Research Paraphrase Corpus
(MRPC) (Dolan & Brockett 2005), Quora Question Pairs (QQP Multi-Genre Natural Language
2018) Question Natural Language Inference (QNLI) (Rajpurkar et al.) 2016). and Recognizing Textual Entailment (RTE) (Dagan et al. [0055] Haim et al. 2006
Inference (MNLI) (Williams
Giampiccolo al. 2007: Bentivogl et al. (209). Note that MNLI contains both an in-domain (MNLI-m) and cross-domain (MNLI-mm) evaluation set. These datasets span sentiment analysis, paraphrase similarity matching, and natural language inference types of tasks. We use the Hugging Face Transfomeerfjimplemmntaa of BERT for our experiments.
4.1 GLUE DATASET EVALUATION
We first analyze the contributions of each component of our method, evaluating on the dev set of the GLUE datasets. For the teacher model, we fine-tune a separate 12 Transformer-layer pert-base-uncasee
(BERT12) for each task. We use the smaller BERT; and BERT, as the student model. We find that initializing the embeddings and Transformer layers of the student model from the first & layers the teacher model provides a significant boost to final performance. We use MSE as the knowledge distillation distance metric d(-,)). We generate mixup sample for each original sample in each minibatch (mixup ratio of 1), with X Beta(0.4, 0.4). We set hyperparameters weighting the components in the loss term in equation g as asM aTMKD 1.
http:/ggngg.gg......oooaaanrooa lease Serree-uuussioo----e0

---

Published as a conference paper at ICLR 2021
Model
SST-2
MRPC
QQP
MNLI-m
MNLI-mm
QNLI
RTE
BERT
93.5
88.9/84.8
71.2/89.2
84.6
83.4
90.5
66.4
BERT. 6
90.7
85.9/80.2
69.2/88.2
80.4
79.7
86.7
63.6
B KD
91.5
86.2/80.6
70.1/88.8
80.2
79.8
88.3
64.7
PKD
92.0
85.0/79.9
70.7/88.9
81.5
81.0
89.0
65.5
92.5
86.4/81.9
70.5/89.1
82.2
81.2
88.2
68.3
3
86.4
80.5/72.6
65.8/86.9
74.8
74.3
84.3
55.2
BERT 86.9
79.5/71.1
67.3/87.8
75.4
74.8
84.0
56.2
BERT,-PKD 87.5
80.7/72.5
68.1/87.8
76.7
76.3
84.7
58.2
BERTs-MixKD 89.5
83.3/75.2
67.2/87.4
77.2
76.8
84.4
62.0
Table GLUE test server results. We show results for the full variants the and 6-layer MixKD student models (SM+TMKD+BT) Knowledge distillation (KD) and Patient Knowledge Distillation (PKD) results are from Sun et al. (2019)).
As a baseline, we fine-tune the student model on the task dataset without any distillation or aug- mentation, which we denote as -FR We compare this against MixKD, with both knowledge distillation on the teacher's predictions (CTMKd) and mixup for the student (CsM), which we call BERTT-SMMTMKK We also evaluate an ablated version without the student mixup loss (BERT,, TMKD) to highlight the knowledge distillation component specifically. We note that our method can also easily be combined with other forms of data augmentation. For example, backtranslation (translating input sequence to the data space of another language and then translating back to the original language) tends to generate varied but semantically similar sequences; these sentences also droppdgo approaches. We show that method
tend to be of higher quality than masking has an additive effect with other techniques by also testing our method with the dataset augmented with German backtranslation, using the fa i rse eO al. 2019) neural machine translation code- base to generate these additional samples. We also compare all of the aforementioned variants with backtranslation samples augmenting the data; denote these variants with additional +BT.
We report the model accuracy (and F1 score, for MRPC and QQP) in Table [ We also show the performance of the full-scale teacher model (BERT12) and DistilBERT (Sanh al.) 2019), which performs basic knowledge distil- lation during BERT pre-training a 6-layer model. For our method, we observe that com- bination of data augmentation and knowledge distillation leads to significant gains in performance, with the best variant often being the combi- nation of teacher mixup knowledge distillation, student mixup, and backtranslation. In the case of SST-2, for example, BERT, S-MTTT-MMMMTMDDBB is able to capture 99.88% of the performance of the teacher model, closing 91.27% of the gap between the fine-tuned student model and the teacher, despite using far fewer parameters and having much faster inference speed (Table
Speed (samplex/sccod))
of
BERT Teacher BERT, Student BEKT', Student
115
283,77
252
66,956,546
15,692,930
Table 2: Computation cost comparison of teacher and student models on SST-2 with batch size of 16 on a Nvidia TITAN X GPU.
After analyzing the contributions of the components of our model on the dev set, we find the BM+TMKD++BT variant have the best performance overall and thus focus this variant. We submit this version of MixKD to the GLUE test server, reporting its results in comparison with fine-tuning (FT), vanilla knowledge distillation (KD) (Hinton al. 22055), and patient knowledge distillation (PKD) (Sun et al. 2019a) in Table Once again, we observe that our model outperforms the baseline methods on most tasks.
4.2 LIMITED-DAY DATA SETTINGS
One of the primary motivations for using data augmentation for knowledge distillation is to give the student more opportunities to query the teacher model. For datasets with large enough number of samples relative to the task's 's complexity, the original dataset may provide enough chances to learn from the teacher, reducing the relative value of data augmentation.
As such, we also evaluate MixKD with a BERT, student on downsampled versions of QQP, MNLI (matched and mismatched), and QNLI in Figure We randomly select 10% and 1 1% of the data from
7

---

Published as a conference paper at ICLR 2021
90
80
85
FT+ MixKD (wto
75
FT+ MixK (wto BT)
Owt BT)
Accuracy (%)
80
MiXKD
Accuracy (%)
70
MixM (w/
75
65
70
60
65
55
60
50
QQP
MNLI
MNLI-mm
QNLI
QQP
MNLI
MNLI-mm
QNLI
Dataset
Dataset
Figure 1: Results of limited data case, where both the teacher and student models are learned with only 10% (left) or 1 1 of the training data (right).
these datasets to train both the teacher and student models, using the same subset for experiments for fair comparison. In this data limited setting, observe substantial gains from MixKD over the fine-tuned model for QQP (+2.0% +3.0%), MNLI-m (+3.9%. +3.4%), MNLI-mm (+4.4% +3.3%), and QNLI (+2.4%, +4.1%) for 10% and 1% of the training data.
4.3 EMBEDDINGS VISUALIZATION
We perform qualitative examination of the ef- fect of the proposed MixKD by visualizing the latent space between positive and negative sam- ples as encoded by the student model with t- SNE plots (Maaten & Hinton 2008). In Fig- ure 2 we show the shift of the transformer features at the [CLS] token position, with and without mixup data augmentation from the teacher. We randomly select batch of 100 sentences from the SST-2 dataset, of which 50 are positive sentiment (blue square) and 50 are negative sentiment (red circle). The intermedi- mixup neighbours are indicated by triangles with color determined by the closeness to the positive group or negative group. From Fig- ure 2(a) to Figure 2(b) MixKD forces the lin- early interpolated samples to aligned with the manifold formed by the real training data and leads the student model to explore meaningful regions of the feature space effectively.
(a)
(b)
Latent space of randomly sampled training data and their mixup neighbours encoded by student model (a) learned by standard fine- tuning (b) learned with MixKD.
Figure 2:
4.4 HYPERPARAMETER SENSITIVITY & FURTHER ANALYSIS
Loss Hyperparameters Our final objective in equation g has hyperparameters asm and @TMKD- which control the weight of the student model's cross-entropy loss for the mixup samples and the knowledge distillation loss with the teacher's 's predictions on the mixup samples, respectively. We demonstrate that the model is fairly stable over a wide range by sweeping both aSM and OTMKDD over the range (0.1, 0.5. 1.0, 2.0, 10.03. We do this for a BERT, student and BERT teacher, with SST-2 as the task; we show the results this sensitivity study, both with and without German backtranslation, in Figure 3 Given the overall consistency, observe that method is stable over a wide range settings.
Mixup Ratio We also investigate the effect of the mixup ratio: the number of mixup samples gen- erated for each sample in a minibatch. We run a smaller sweep of aSM and @TMKD over the range (0.5, 1.0, 2.0} } for mixup ratios of 2 and for BERT, student SST-2, with and without German backtranslation, in Figure 3 We conclude that the mixup ratio does not have strong effect on over- all performance. Given that higher mixup ratio requires over which to compute the forward and backward pass), we find mixup ratio of 1 to be enough.
computation (due to
samples
8

---

Published as a conference paper at ICLR 2021
Accuracy (%)
89
Accuracy (%)
Accuracy (%)
89
88
877
66
Beswinn
Baseline w/ beck-transiation
Baseline
back-translation
W/ bacbstransiation
85
backetransiation
agu
Mixup Samples
Figure 3: Hyperparameter sensitivity analysis regarding the MixKD approach, with different choices of @TMKD, OSM and the ratio of mixup samples (w.r.. the original training data).
Comparing with TinyBERT's DA module TinyBERT (Jiao et al., 2019) also utilizes data augmentation for knowledge distillation. Specifically, they adopt conditional BERT contextual augmentation (Wu et al. 2019) strat- egy. To further verify the ellectiveness of our approach, we use TinyBERT's released code- base to generate augmented samples and make direct comparison with MixKD. As shown in Table our approach exhibits much stronger results for distilling a 6-layer BERT model (on both MILI and SST-2 datasets). Notably, TinyBERT's data augmentation module is much less efficient than mixup's 's simple operation, generating 20 times the original data as augmented samples, thus leading to massive computation overhead.
Methods
MNLI
SST-2
BERT
8T.3
90.9
BERT, + TEyBERT DA module
81.5
91.3
BERT, MixKD
82.5
92.1
Table 4: We compare approach with the data augmentation module proposed by TinyBert (Jiao et al. 2019).
5 CONCLUSIONS
We introduce MixKD, method that uses data augmentation to significantly increase the value of knowledge distillation for compressing large-scale language models. Intuitively, MixKD allows the student model additional queries to the teacher model, granting more opportunities absorb the latter's richer representations. We analyze MixKD from theoretical standpoint, proving that our approach results smaller gap between generalization error and empirical error, as well as better generalization, under appropriate conditions. Our approach's 's success on a variety of GLUE tasks demonstrates its broad applicability, with a thorough set of experiments for validation. We also believe that the MixKI) framework can further reduce the gap between student and teacher models with the incorporation of more recent mixup and knowledge distillation techniques (Lee al. 2020 Wang et al. 2020, Mirzadeh et al., 2019), and we leave this to future work.
ACKNOWLEDGMENTS
CC is partly supported by the Verizon Media FREP program.
REFERENCES
Hangbo Bao, Dong, Furu Wei, Wenhui Wang, Nan Yang, Xiaodong Liu, Yu Wang, Songhao Piao, Jianfeng Gao, Ming Zhou, al. UniLMv2: Pscudo-masked masked Language Models for Unified Language Model Pre-training. arXiv preprint arXiv:2002. 128044 2020.
J. Baxter. A Model of Inductive Bias Learning. Journal of Artificial Intelligence Research, 2000.
Luisa Bentivogli, Peter Clark, Ido Dagan, and Danilo Giampiccolo. The Fifth PASCAL. Recognizing Textual Entailment Challenge. TAC, 2009.
David Berthelot, Nicholas Carlini, Ekin Cubuk, Alex Kurakin, Kihyuk Sohn, Han Zhang, and Colin Raffel. ReMixMatch: Semi-Supervised Learning with Distribution Alignment and Aug- mentation Anchoring. International Conference on Learning Representations, 2019a.
FinyBERT

---

Published as a conference paper at ICLR 2021
David Berthelot, Nicholas Carlini, Ian Goodfellow, Nicolas Papernot, Avital Oliver, and Colin A Raffel. MixMatch: A Holistic Approach to Sami-Supervised Learning. Neural Information Processing Systems, 2019b.
Jiaao Chen, Zichao Yang, and Diyi Yang. MixText: Linguistically-Informed Interpolation of Hidden Space for Semi-Supervised Text Classification. Association for Computational Linguistics, July 2020.
Kevin Clark, Minh-Thang Luong, Urvashi Khandelwal, Christopher 1) Manning, and Quoe V Le. BAM! Born-again Multi-task Networks for Natural Language Understanding. arXiv preprint MrXiv:1907.04829 2019.
Kevin Clark, Minh-Thar Thang Luong, Quoe V Le, and Christopher D Manning. ELECTRA: Pre-training Text Encoders as Discriminators Rather than Generators. arXiv preprint arXiv:2003. 10555, 2020.
Ido Dagan, Oren Glickman, and Bemarro Magnini. The PASCAL Recognising Textual Entailment Challenge. Machine Learning Challenges Workshop, 2005.
Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv preprint arXiv:1810.044: 04805, 2018.
William B Dolan and Chris Brockett. Automatically Constructing Corpus of Sentential Para-
phrases. International Workshop on Paraphrasing, 2005.
Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The Third PASCAL Recog- nizing Textual Entailment Challenge. ACL-PASCAL Workshop Textual Entailment and Para- phrasing, 2007.
Hongyu Guo, Yongyi Mao, and Richong Zhang. Augmenting Data with Mixup for Sentence Clas- sification: An Empirical Study. arXiv preprint arXiv: /0X.v91905.08941 2019.
R Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Szpektor. The Second PASCAL Recognising Textual Entailment Challenge. PASCAL Challenges Workshop Recognising Textual Entailment, 2006.
Dan Hendrycks, Norman Mu, Ekin Dogus Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshmi- narayanan. AugMix: Simple Method to Improve Robustness and Uncertainty under Data Shift. International Conference Learning Representations, 2020.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in Neural Network. arXiv preprint XXI::15033025311 2015.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang. Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. TinyBERT: Distilling BERT for Natural Language Understanding. arXiv preprint arXiv:/909 10351, 2019.
Mandar Joshi, Danqi Chen, Yinhan Liu, Daniel S Weld, Luke Zettlemoyer, and Omer Levy. Span- BERT: Improving Pre-training by Representing and Predicting Spans. Association for Computa- tional Linguistics, 2020.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet Classification with Deep Con- volutional Neural Networks. Neural Information Processing Systems, 2012.
Ashutosh Kumar, Satwik Bhattamishra, Manik Bhandari, and Partha Talukdar. Submodular Optimization-basel Diverse Paraphrasing and Its Effectiveness in Data Augmentation. North American Chapter of the Association for Computational Linguistics; Human Language Tech- nologies, 2019.
Saehyung Lee, Hyungyu Lee, and Sungroh Yoon. Adversarial Vertex Mixup: Toward Better Adver- sarially Robust Generalization. Computer Vision and Pattern Recognition, 2020.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. BART: Denoising Sequence-to-sequence Pre- training for Natural Language Generation, Translation, and Comprehension. arXiv preprint arXiv: 1910. 13461, 2019.
10

---

Published as a conference paper at ICLR 2021
Linqing Liu, Huan Wang. Jimmy Lin, Richard Socher, and Caiming Xiong, Attentive Student Meets Multi-task Teacher: Improved Knowledge Distillation for Pretrained Models. arXiv preprint arXiv:1911.03588, 2019a.
Xiaodong Liu, Pengcheng He, Weizhu Chen, and Jianfeng Gao. Multi-task Deep Neural Networks for Natural Language Understanding. arXiv preprint
2019b.
Yinhan Liu, Myle Out Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A Robustly Optimized BERT Pre- training Approach. arXiv preprint 190. 2019c.
Laurens van der Maaten and Geoffrey Hinton. Visualizing Data using t-SNE. Journal of Machine Learning Research, 2008.
Seyed-Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Has- san Ghasemzadeh. Improved Knowledge Distillation via Teacher Assistant. arXiv preprint arXiv:1902.03339, 2019.
Mehryar Mohri, Afshin Rostamizadch, and Ameet Talwalkar. Foundations of Machine Learning. 2018.
Subhabrata Mukherjee and Ahmed Hassan Awadallah. Distilling Transformers into Simple Neural Networks with Unlabeled Transfer Data. arXiv preprint 1xXx.0991000776 2019.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. fairseq: A Fast, Extensible Toolkit for Sequence Modeling. North Ameri- can Chapter of the Association for Computational Linguistics: Human Language Technologies: Demonstrations, 2019.
Yanru Qu, Dinghan Shen, Yelong Shen, Sandra Sajeev, Jiawei Han, and Weizhu Chen. CoDA: eonareet-enhanced and Diversity-promotie Data Augmentation for Natural Language Under- standing. arXiv preprint arXiv:2010.08867 0807000000000,,. 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ Questions for Machine Comprehension of Text. Empirical Methods in Natural Language Processing, 2016.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistiBBERT, A Distilled Version of BERT: Smaller, Faster, Cheaper and Lighter. arXiv preprint arXiv: 19/ 0910.01 108, 2019.
Dinghan Shen, Mingzhi Zheng, Yelong Shen, Yanru Qu, and Weizhu Chen. A Simple but Tough-to- Beat Data Augmentation Approach for Natural Language Understanding and Generation. arXiv preprint arXiv:2009.1388 13818, 2020.
Patrice Y Simard, Yann A LeCun, John S Denker, and Bernard Victorri. Transformation Invariance in Pattern Iettogniiion-Tangeen Distance and Tangent Propagation. Neural Networks: Tricks of the Trade, 1998.
PY Simard, D Steinkraus, and JC Platt. Best Practices for Convolutional Neural Networks Applied to Visual Document Analysis. International Conference on Document Analysis and Recognition, 2003.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive Deep Models for Semantic Compositionality over a Sentiment Treebank. Empirical Methods in Natural Language Processing, 2013.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu, Patient Knowledge Distillation for BERT Model Compression. arXiv preprint arXiv: 093550, 2019a.
Yu Sun, Shuohuan Wang, Yukun Li, Shikun Feng. Xuyi Chen, Han Zhang, Xin Tian, Danxiang Zhu, Hao Tian, and Hua Wu. ERNIE: Enhanced Representation through Knowledge Integration. arXiv preprint arXiv:9904 09223, 2019b.
Zhiqing Sun, Hongkun Yu, Xiaodan Song. Renjie Liu, Yiming Yang, and Denny Zhou. Mo- bileBERT: A Compact Task-agnostic BERT for Resource-limite Devices. arXiv preprint arXiv:2004.02984 2020.
11

---

Published as a conference paper at ICLR 2021
Raphael Tang. Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin. Distilling Task- specific Knowledge from BERT into Simple Neural Networks. arXiv preprint trXiv:1903.121333 2019.
Tulia Ture, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read Students Learn Better: On the Importance of Pre-training Compact Models. arXiv preprint
2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All You Need. Neural Information Process- ing Systems, 2017.
Vikas Verma, Alex Lamb, Christopher Beckham, Amir Najafi, Ioannis Mitliagkas, David Lopez- Paz, and Yoshua Bengio. Manifold Mixup: Better Representations by Interpolating Hidden States. International Conference Machine Learning, 2019a.
Vikas Verma, Alex Lamb, Juho Kannala, Yoshua Bengio, and David Lopez-Paz. Interpolation Con- sistency Training for Semi-supervised Learning. International Joint Conference on Artificial In- telligence, 2019b.
Vikas Verma, Meng Qu, Alex Lamb, Yoshua Bengio, Juho Kannala, and Jian Tang. GraphMix: Improved Training of GNNs for Semi-Supervised Learning. arXiv preprint arXiv: 11715, 2019c.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE: A Multi-task Benchmark and Analysis Platform for Natural Language Understanding. International Conference on Learning Representations, 2019.
Dongdong Wang, Yandong Li, Liqiang Wang, and Boqing Gong. Neural Networks Are More Pro- ductive Teachers Than Human Raters: Active Mixup for Data-Effficiet Knowledge Distillation from Blackbox Model. Computer Vision and Pattern Recognition, 2020.
Jason Wei and Kai Zou. EDA: Easy Data Augmentation Techniques for Boosting Performance on Text Classification Tasks. arXiv preprint arXiv:1901.11 2019.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A Broad-coverage Challenge Corpus for Sentence Understanding through Inference. North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2018.
Xing Wu, Shangwen Lv, Liangjun Zang, Jizhong Han, and Songlin Hu. Conditional BERT Contex- tual Augmentation. International Conference on Computational Science, 2019.
Qizhe Xic, Zihang Dai, Eduard Hovy, Minh-Thang Luong, and Quoc V Le. Unsupervised Data Augmentation for Consistency Training. arXiv preprint arXiv:/904 /2848, 2019.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. XLNet: Generalized Autoregressive Pretraining for Language Understanding, Neural Informa- tion Processing Systems, 2019.
Adams Wei Yu, David Dohan, Minh- Than-TThng Luong, Rui Zhao, Kai Chen, Mohammad Norouzi, and Quoe V Lee. QANet: Combining Local Convolution with Global Self-attention for Reading Comprchension. arXiv preprint 1rXiv:1804.00 005409941 2018.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. CutMix: Regularization Strategy Train Strong Classifiers with Localizable Features. Interna- tional Conference on Computer Vision, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond Empiri- cal Risk Minimization. International Conference on Learning Representations, 2018.
Sanqiang Zhao, Raghav Gupta, Yang Song, and Denny Zhou. Extreme Language Model Compres- sion with Optimal Subwords and Shared Projections. arXiv preprint 1909 2019.
12

---

Published as a conference paper at ICLR 2021
A PROOFS
Proof [Proof of Theorem First of all, {x.]==1 (x]
can be regarded as drawn from distribu-
ap(x) (x(x)
tion r(x)
a a+6
Given 9 is finite, we have the following theorem
Theorem (Mohri et al. 2018) Let be a bounded loss function, hypothesis set G is finite. Then for any ó > with probability at least 1 8, the following inequality holds for all 9:
log (gg(()1/))
R(f.g.) (f (x:)==1) S MM
2n
Thus have in our case:
log (191/8)
R(f gp: p) Remp (f (xili=1) 5 c 1 My
2n
and
R(f,g" .p) - Remp ( (x
=R(fg Remp(f: (x: U (x+)==)) + / (1F(()) - r(x))dx
-R(fgg"" Remp (xili= U (x)) + a ((r(())) (x))(p(x) (x))(()) =xx))dxx
SR(S, rr) RemplS. (x: (x)++ + ((U(()), x(x))dxx
(9/8)
SM
(10)
2(a +b) b)
where f(((((((()))))))))))))(()) (x))(((( - g(xx))xxx
M (og(((1)))))
b
2(Ep A)
then
M Molog((G1/8)
2(a + b)
(Cp A)2
M log((11))))
(ep 2) 2
2(a+1)
log(191/5)
cp 2 M
+ A
2(a + b)
Substitute into equation 10 we have
R(J,g" "f"g"""" ,p) - Remp {x: U (x S cp
Recall the definition of 6", which is the minimum value of all possible satisfying
(f.g" g", {x1) U (x) < R(f.9" .P) Remp(f
know that ex Ep- Let c 2(E, 2)2, we can conclude the theorem.
Proof [Proof of Theorem
First of all, {x,] U {x{ can be regarded as drawn from distribu-
tion r(x) ap(x) bg(x)
a a+6
13

---

Published as a conference paper at ICLR 2021
Theorem 6 (Mohri al. 2018 Let be non-negative loss function upper bounded by M 0, and for any fixed y. ly.y
L-Lipschitz, for some L > then with probability at least 6.
(og(1/8)
R(f.g.) Remp (x)]==))= S 2L M
2n
Thus we have
(1/6)
R(Sgg.P) - Remp(s (x: S cp S (9) + M
2n
where Rg)
Rademacher complexity over all samples of size samples from p(x).
We also have
R(f.g" p) (x. (x](i))
=Rffgg" Remplf.g" (xi]i- U (x()==)+ + / ((f(xx). - r(x)))xx
=R(f. "0f.g"",) ,r) Remp(f. (x:]; u (x))))) + b
a+b. ( o q(x)))xx
<R(f.g" Remplf.g (xiii U (x))+)+++++ / ((f(xx). g(x)))xx
(2) M1 (1/()/)) +A 2(a + b)
(11)
where == - 9" (x)) dx (9) are Rademacher complexity over all samples ap(x) + bx(x)
of size (a+ b) samples from r(x)
a
If
M (1/)
b2 2
2(ep A 2L9,, (G))z
then:
M log(1/8)
2(a + b) 2
(Ep A - 2LD,, (g))²
log(118)
Eo A ( 2 M
2(a + b)
c 2 M log(1/5)
A + (9)
2(a + b)
Substitute into equation T1 we have:
R(f.g",P) p)) - Remp {x: U (x)+ S cp
Recall the definition of e*, which is the minimum value of all possible satisfying
R(f.g" p) Rfmp(f g", {x;] U (x(} <
we know that e" S Ep. Let c 2(E, A (G)) we can conclude the theorem.
Proof [Proof of Theorem 3 Similar to previous theorems, we write
R(f.g" p) Ramp (f.g", fxi U fx())_)
=R(f.g" ap + bq, + b
+ /
ap(x) + (b((x)
R Remp(f.g j U (x()+))++++
)d
a+
=R(f.g', ap + bq
Remplf.g g", {xie U (x]
a - (((()) (x))))))))) xxx))xx
b
a + b
<R(f.g' ap + bq
(f.g (x U [x] +4
(12)
b
14

---

Published as a conference paper at ICLR 2021
where A = (x))(x) - qxx))xxx
For notation consistency, write x. However, (x U {x))== are not
R ap + bq = ['((()) s(x)) ap(x) + bx(x),
a
a b
drawn from the same distribution (which is r(x) ap(x) + by(xx) in previous cases). a +
a [-- we split (x,)==1 U (x)i= into y parts that don't overlap with each other. The
Let y =
first part is |x all the other parts has at least a elements from (x]
Let
64 log (/6) + log C(g)
64
b
where C(g) is space capacity defined in Definition 4 in Baxter (2000), which depends on E" and G.
By Theorem in Baxter (2000).
R(Sg"" ap + bg a a Remy (Fig" {x:1i=1 (x))=n)
64
4C
16
maxt
log(
7a
8
By Theorem 5 in Baxter (2000).
64 aa
log( 4C(g))
64
< (log( 4
64
(log( 5) ) +og((((()))
y log(C(()))) s X2
ya
ya
a
The last inequality comes from b I ya, which is because of y =
Then we have
a
R(S.9" ap + bq Remp (S,g" {x:]i_ U (x))=n))
64
4C
16
S max max 16
I max
log
+
7a
64 1og(4/6)
2 (Ep
4
64 log C(9)/a
Then
64
1 (cp A)2 64 log C(G)/a
xa s
log ((9) 64 6olog(
a
8
641og(4/8)
x² $ (Ep - 22
(13)
If
16
(e, 4) S
then
S (ep - 2)
(14)
Combine equation 13 and equation 14 we have
ap + by Remp(f. fxil U (x S - -A + b
R(f.g",
Substitute into equation 12 we have:
R(f, g", p) Remp {x:}a_ U {x s Ep
Recall the definition of E", which is the minimum value of all possible satisfying
R(f.g" . p) &fmp(f.g) g", fx {x)+ Sc
we know that c* S Ep-
15

---

Published as a conference paper at ICLR 2021
B VARIANCE ANALYSIS
For the purpose of getting a sense of variance, we run experiments with additional random seeds on MRPC and RTE, which relatively smaller datasets, and MNLI and QNLI, which relatively larger datasets. Mean and standard deviation on the dev set of these GLUE datasets are reported in Table 5 relatively larger datasets.
We observe the variance of the same model's performance be small, especially on the
Model
MRPC
MNLI-m
QNLI
RTE
82.05+0.111
88.42-40.06
69.37+0..50
BERT,- SM-TKKDBB
$9,64110.888444430...6
82.41
88.76:10.15
68.02:10.11
BEIT;:TTMKD+BT
84.79:0033375.82210444
77.16::0.03
84.60+0.007
62.47-10.36
BERTT--M 8455:10.3997588-16.060
77.42:10111
84.88:10.06
60.83:10.18
Table 5: Mean and variance reported for
TMKD+BT and
16


Reviews:
[{"review_id": "x9WQNdDhu2e", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper applies mixup (Zhang et al., 2018) to augment training data to improve knowledge distillation in NLP tasks . Mixup was originally proposed to augment data for continuous data. To apply mixup to textual data, this paper applies mixup to the word/token embeddings instead of the tokens themselves. Some theoretical analysis has been done, and the experimental results show improved metrics over baseline methods such as DistillBERT.\n\n\n*Strong points*\n- This paper is well-written and very easy to follow.\n- Theoretical results have been provided and the experiment has been carefully done.\n\n*Weak points*\n- The proposed method is a direct application of mixup (Zhang et al., 2018). The overall idea is quite straightforward.\n- The empirical results of the proposed method are not impressive (worse than the metrics reported in TinyBERT (Jiao et al., 2019)?).\n- It is true that for benchmarks like GLUE, the datasets are quite small so that data augmentation is important for knowledge distillation. This is however not true for most real-world applications where we often have almost unlimited unlabelled data (e.g., from logs, extract from webs etc). In other words, text data augmentation for knowledge distillation is not necessarily an important problem for real-world applications.\n\n*Questions & Other comments*\n- In Table 3, the metrics for TinyBERT is different from what has been reported in TinyBERT (Jiao et al., 2019). Can you add more explanation here?\n- “Notably, TinyBERT’s data augmentation module is much less efficient than mixup’s simple operation” It is not an issue to have some computation overhead in data augmentation, as it is a one-off operation.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "_DX0z-z_dsw", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes combining the MixUp data augmentation method with teacher-student distillation to improve the fine-tuned performance of BERT on benchmark NLP tasks (GLUE). The problem is important, well-motivated and of interest to a broad base of NLP researchers and practitioners. The paper is clear, and generally well-written, although the idea itself is not surprisingly novel (somewhat of a low-hanging fruit), from the experimental results, the method improves upon baselines, and so the real-world impact could be high, especially given its simple implementation.\n\nPros:\n* Well-written, solid experiments / ablations on a canonical benchmark (bonus points for reporting GLUE test set performance), some theoretical contributions.\n* Reports performance in labeled-data-limited setting (by subsampling GLUE dataset)\n\nCons:\n* In the theory, a bounded loss function is assumed, but in practice the unbounded cross-entropy is used (right?)\n* Optimal hyper-parameters $\\alpha_\\text{SM}$ and $\\alpha_\\text{TMKD}$ for \"MixKD\" on the GLUE dev test are not reported in the main text\n* \"MixKD\" includes a back-translation (BT) loss term as well as a student loss on mixup samples (SM). While the impact of each term is studied on the GLUE dev set, I'm interested in seeing its performance on GLUE test.\n* In the theorems, $\\delta$ should be re-introduced.\n* The embedding visualization section needs improvement. What is being visualized here? 2-dimensional logits / probabilities from the model or input embeddings projected down to 2D? Furthermore, it is hard to distinguish circles from triangles without zooming in 3-5x.", "extra_scores": {}, "extra_text": {}}, {"review_id": "YzU82qRBZcP", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper describes a distillation framework that leans heavily on Mixup, an effective data augmentation technique that is typically applied to images, to achieve impressive results on a range of GLUE benchmarks (). By using SM+TMKD+BT, the student model (BERT_6) is able to capture 99.88% of the performance of the teacher model (BERT_12).\n\nThe authors provide good motivation for the problem of distilling large-scale language models and conduct extensive experiments across multiple GLUE benchmarks, including ablation and hyperparameter sensitivity studies, and the results are quite convincing. Figure 3 nicely demonstrates that the framework is not overly sensitive to any of its hyperparameters.\n\nQuestions:\n- Given that the paper motivates distillation as a means of reducing the computational cost of state-of-the-art language models (less power, less memory, lower latency), could you also describe the number of parameters that the student models use and their latency at inference time (vs. that of the teacher models)?\n\nWhile the paper is generally well-written, there are a number of small issues / inconsistencies:\n- Nit: Mixup is inconsistently capitalized / hyphenated.\n- Since Mixup is central to this distillation framework, it would be valuable to briefly describe the method in the body of the paper, before describing how it was adapted to the domain of language.\n- “that a a multi-task BERT model…”\n- “around an training example”\n- Nit: “vicinal” is an extremely rare word, so it might be worth using a more accessible word such as “neighboring”\n- Nit: “tabular data were also shown, to demonstrate generality” (remove comma)\n- Nit: “the extra word embeddings are mixuped with zero paddings” (maybe “mixed up” is better? :))\n- In Table 1, for ease of reading, it would be good to describe in the caption the meaning of the first and second ‘/’ delimited number in each cell.\n- For Figure 2, please describe the dimensionality reduction method used to visualize the latent space of training data and augmented examples. The pattern looks compelling, but I don’t see you describe anywhere how the data was projected. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "LFc_qGo-k4Y", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Nice paper. The main idea in this paper is to use a specific kind of data augmentation, Mixup (Manifold Mixup), in order to improve the effectiveness of the KD process and obtain better performing student models, especially in cases where not enough data is available on the target dataset and task.\n\nWhile the idea is interesting in itself, I think what is proposed in this paper, is very relevant to this paper (noisy student): https://arxiv.org/abs/1911.04252\n\nSo, methodologically it's not super novel, but it should also be taken into account that many of the developments in this area have happened very recently. Moreover, the paper still has an added value since it is applying these techniques for a different input modality and in a slightly different setup. \n\nSummary of the experiments:\n-12 layer BERT is fine tuned as the teacher (different fine-tuned teacher for each task).\n- Two different architectures are used for student models (6 layer BERT and 3 layer BERT).\n- Student models are initialised by copying the lower layers of the teachers.\n- The teacher is used to provide pseudo labels for the student model on the target dataset, while the target dataset is augmented with the manifold mixup approach on the embedding layer (Proposed Approach).\n- The proposed approach is compared with plain KD and no KD (just fine-tuning on target).\n \nMy Questions:\n- During KD, do you feed original examples from target to the teacher then apply mixup, or do you first apply mixup then feed the generated examples to the teacher and get the pseudo labels from the teacher?\n-Most importantly, It is not clear to me from the experiments whether there is still an advantage of doing this, if the teacher is trained well enough e.g., with the same type of augmentation (intuitively, if the teacher is trained well enough and KD process is well tuned, all the information that the model can gain by doing additional data augmentations should already be transferrable through the soft targets from the teacher.)\n\nSome positive point about the paper:\n- Ablation experiments are conducted to separate the gain from simple KD with KD+Mixup.\n\nSome of my concerns:\n- The improvements in the accuracies reported seem marginal and I think there are indication of the significancy of the results in the paper (e.g., mean and variance over several trials?).\n\nSome minor points:\n- Figure 2: Mixup and Non Mixup examples (circles and triangles) are not clearly differentiable (just hard to see).\n-I think the type of data augmentation applied is a special case of Manifold Mixup, so it would be nice to cite this paper as well:  https://arxiv.org/abs/1806.05236\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work explores the distillation of language models using MixUp for data augmentation. Distillation with MixUp seems to be novel in the narrow context of distilling language models, although it has been used before in different contexts as the reviewers point out. The results of the experimental validation are encouraging, and the application is valuable and of wide interest to the ICLR audience. I therefore recommend accepting this paper for a poster presentation.
Decision: Accept (Poster)