T INY BERT: D ISTILLING BERT FOR N ATURAL L ANGUAGE U NDERSTANDING
Anonymous authors
Paper under double-blind review
A BSTRACT
Language model pre-training, such as BERT, has signiﬁcantly improved the performances of many natural language processing tasks.
However, the pre-trained language models are usually computationally expensive and memory intensive, so it is difﬁcult to effectively execute them on resource-restricted devices.
To accelerate inference and reduce model size while maintaining accuracy, we ﬁrstly propose a novel Transformer distillation method that is specially designed for knowledge distillation (KD) of the Transformer-based models.
By leveraging this new KD method, the plenty of knowledge encoded in a large “teacher” BERT can be well transferred to a small “student” TinyBERT.
Moreover, we introduce a new two-stage learning framework for TinyBERT, which performs Transformer distillation at both the pre-training and task-speciﬁc learning stages.
This framework ensures that TinyBERT can capture the general-domain as well as the task-speciﬁc knowledge in BERT.
TinyBERT 1 is empirically effective and achieves more than 96% the performance of teacher BERT BASE on GLUE benchmark, while being 7.5x smaller and 9.4x faster on inference.
TinyBERT is also signiﬁcantly better than state-of-the-art baselines on BERT distillation, with only ∼ 28% parameters and ∼ 31% inference time of them.
I NTRODUCTION
Pre-training language models then ﬁne-tuning on downstream tasks has become a new paradigm for natural language processing (NLP).
Pre-trained language models (PLMs), such as BERT (Devlin et al., 2018), XLNet (Yang et al., 2019), RoBERTa (Liu et al., 2019) and SpanBERT (Joshi et al., 2019), have achieved great success in many NLP tasks (e.g., the GLUE benchmark (Wang et al., 2018) and the challenging multi-hop reasoning task (Ding et al., 2019)).
However, PLMs usually have an extremely large number of parameters and need long inference time, which are difﬁcult to be deployed on edge devices such as mobile phones.
Moreover, recent studies (Kovaleva et al., 2019) also demonstrate that there is redundancy in PLMs.
Therefore, it is crucial and possible to reduce the computational overhead and model storage of PLMs while keeping their performances.
There has been many model compression techniques (Han et al., 2015a) proposed to accelerate deep model inference and reduce model size while maintaining accuracy.
The most commonly used techniques include quantization (Gong et al., 2014), weights pruning (Han et al., 2015b), and knowledge distillation (KD) (Romero et al., 2014).
In this paper we focus on knowledge distillation, an idea proposed by Hinton et al.
(2015) in a teacher-student framework.
KD aims to transfer the knowledge embedded in a large teacher network to a small student network.
The student network is trained to reproduce the behaviors of the teacher network.
Based on the framework, we propose a novel distillation method speciﬁcally for Transformer-based models (Vaswani et al., 2017), and use
BERT as an example to investigate the KD methods for large scale PLMs.
KD has been extensively studied in NLP (Kim & Rush, 2016;
Hu et al., 2018), while designing
KD methods for BERT has been less explored.
The pre-training-then-ﬁne-tuning paradigm ﬁrstly pre-trains BERT on a large scale unsupervised text corpus, then ﬁne-tunes it on task-speciﬁc dataset, which greatly increases the difﬁculty of BERT distillation.
Thus we are required to design an ef- fective KD strategy for both stages.
To build a competitive TinyBERT, we ﬁrstly propose a new
Transformer distillation method to distill the knowledge embedded in teacher BERT.
Speciﬁcally, we design several loss functions to ﬁt different representations from BERT layers: 1) the output of the embedding layer; 2) the hidden states and attention matrices derived from the Transformer layer; 3) the logits output by the prediction layer.
The attention based ﬁtting is inspired by the recent ﬁndings (Clark et al., 2019) that the attention weights learned by BERT can capture substantial linguistic knowledge, which encourages that the linguistic knowledge can be well transferred from teacher BERT to student TinyBERT.
However, it is ignored in existing KD methods of BERT, such as Distilled BiLSTM SOFT (Tang et al., 2019), BERT-PKD (Sun et al., 2019) and DistilBERT 2 .
Then, we propose a novel two-stage learning framework including the general distillation and the task-speciﬁc distillation .
At the general distillation stage, the original BERT without ﬁne-tuning acts as the teacher model.
The student TinyBERT learns to mimic the teacher’s behavior by executing the proposed Transformer distillation on the large scale corpus from general domain.
We obtain a general TinyBERT that can be ﬁne-tuned for various downstream tasks.
At the task-speciﬁc distillation stage, we perform the data augmentation to provide more task-speciﬁc data for teacherstudent learning, and then re-execute the Transformer distillation on the augmented data.
Both the two stages are essential to improve the performance and generalization capability of TinyBERT.
A detailed comparison between the proposed method and other existing methods is summarized in
Table 1.
The Transformer distillation and two-stage learning framework are two key ideas of the proposed method.
and Pred represent the knowledge from embedding layers, attention matrices, hidden states, and ﬁnal prediction layers, respectively.
KD Methods
KD at Pre-training Stage KD at Fine-tuning Stage
INIT Embd Attn
Hidn Pred
Embd Attn Hidn Pred DA
Distilled BiLSTM SOFT (cid:88) (cid:88)
BERT-PKD (cid:88) (cid:88)(cid:88)
DistilBERT (cid:88) (cid:88)(cid:88)
TinyBERT (our method) (cid:88) (cid:88) (cid:88) (cid:88) (cid:88) (cid:88) (cid:88) (cid:88)
The main contributions of this work are as follows: 1) We propose a new Transformer distillation method to encourage that the linguistic knowledge encoded in teacher BERT can be well transferred to TinyBERT.
We propose a novel two-stage learning framework with performing the proposed
Transformer distillation at both the pre-training and ﬁne-tuning stages, which ensures that TinyBERT can capture both the general-domain and task-speciﬁc knowledge of the teacher BERT. 3)
We show experimentally that our TinyBERT can achieve more than 96% the performance of teacher
BERT BASE on GLUE tasks, while having much fewer parameters ( ∼ 13.3%) and less inference time ( ∼ 10.6%), and signiﬁcantly outperforms other state-of-the-art baselines on BERT distillation.
P RELIMINARIES
We ﬁrstly describe the formulation of Transformer (Vaswani et al., 2017) and Knowledge Distillation (Hinton et al., 2015).
Our proposed Transformer distillation is a specially designed KD method for Transformer-based models.
2. T RANSFORMER L AYER
Most of the recent pre-trained language models (e.g., BERT, XLNet and RoBERTa) are built with Transformer layers, which can capture long-term dependencies between input tokens by selfattention mechanism.
Speciﬁcally, a standard Transformer layer includes two main sub-layers: multi-head attention (MHA) and fully connected feed-forward network (FFN).
Multi-Head Attention (MHA) .
The calculation of attention function depends on the three components of queries, keys and values, which are denoted as matrices Q , K and V respectively.
The attention function can be formulated as follows:
[Equation]
where d k is the dimension of keys and acts as a scaling factor, A is the attention matrix calculated from the compatibility of Q and K by dot-product operation.
The ﬁnal function output is calculated as a weighted sum of values V , and the weight is computed by applying softmax() operation on the each column of matrix A .
According to Clark et al.
(2019), the attention matrices in BERT can capture substantial linguistic knowledge, and thus play an essential role in our proposed distillation method.
Multi-head attention is deﬁned by concatenating the attention heads from different representation subspaces as follows:
[Equation]
where h is the number of attention heads, and head i denotes the i -th attention head, which is calculated by the Attention () function with inputs from different representation subspaces, the matrix
W acts as a linear transformation.
Position-wise Feed-Forward Network (FFN) .
Transformer layer also contains a fully connected feed-forward network, which is formulated as follows:
[Equation]
We can see that the FFN contains two linear transformations and one ReLU activation.
2. K NOWLEDGE D ISTILLATION
KD aims to transfer the knowledge of a large teacher network T to a small student network S .
The student network is trained to mimic the behaviors of teacher networks.
Let f T and f S represent the behavior functions of teacher and student networks, respectively.
The behavior function targets at transforming network inputs to some informative representations, and it can be deﬁned as the output of any layer in the network.
In the context of Transformer distillation, the output of MHA layer or FFN layer, or some intermediate representations (such as the attention matrix A ) can be used as behavior function.
Formally, KD can be modeled as minimizing the following objective function:
[Equation]
where L ( · ) is a loss function that evaluates the difference between teacher and student networks, x is the text input and X denotes the training dataset.
Thus the key research problem becomes how to deﬁne effective behavior functions and loss functions.
Different from previous KD methods, we also need to consider how to perform KD at the pre-training stage of BERT in addition to the task-speciﬁc training stage.
M ETHOD
In this section, we propose a novel distillation method for Transformer-based models, and present a two-stage learning framework for our model distilled from BERT, which is called TinyBERT.
3. T RANSFORMER D ISTILLATION
The proposed Transformer distillation is a specially designed KD method for Transformer networks.
Figure 1 displays an overview of the proposed KD method.
In this work, both the student and teacher networks are built with Transformer layers.
For a clear illustration, we ﬁrstly formulate the problem before introducing our method.
[Figure]
𝑑 ′
[Figure]
Input
[Figure]
Hidn 𝑙𝑜𝑠𝑠
Attn 𝑙𝑜𝑠𝑠 (ℝ ℎ𝑒𝑎𝑑∗𝑙∗𝑙 ) (ℝ ℎ𝑒𝑎𝑑∗𝑙∗𝑙 ) … …
(ℝ 𝑙∗𝒅 ) (ℝ 𝑙∗𝒅 ′ )
ADD & Norm
ADD
[Figure]
Norm
FFN
MHA
ADD
[Figure]
Norm
ADD [Figure] Norm
FFN
MHA
[Figure]
Problem Formulation . Assuming that the student model has M Transformer layers and teacher model has N Transformer layers, we choose M layers from the teacher model for the Transformerlayer distillation .
The function n = g ( m ) is used as a mapping function from student layers to teacher layers, which means that the m -th layer of student model learns the information from the nth layer of teacher model.
The embedding-layer distillation and the prediction-layer distillation are also considered.
We set 0 to be the index of embedding layer and M +1 to be the index of prediction layer, and the corresponding layer mappings are deﬁned as 0 = g (0) and N + 1 =
g ( M + 1) respectively.
The effect of the choice of different mapping functions on the performances will be studied in the experiment section.
Formally, the student can acquire knowledge from the teacher by minimizing the following objective:
[Equation]
(6) where L layer refers to the loss function of a given model layer (e.g., Transformer layer or embedding layer) and λ m is the hyper-parameter that represents the importance of the m -th layer’s distillation.
Transformer-layer Distillation .
The proposed Transformer-layer distillation includes the attention based distillation and hidden states based distillation , which is shown in Figure 1 (b).
The attention based distillation is motivated by the recent ﬁndings that attention weights learned by BERT can capture rich linguistic knowledge (Clark et al., 2019).
This kind of linguistic knowledge includes the syntax and coreference information, which is essential for natural language understanding.
Thus we propose the attention based distillation to encourage that the linguistic knowledge can be transferred from teacher BERT to student TinyBERT.
Speciﬁcally, the student learns to ﬁt the matrices of multihead attention in the teacher network, and the objective is deﬁned as:
[Equation]
where h is the number of attention heads, A i ∈ R l × l refers to the attention matrix corresponding to the i -th head of teacher or student, l is the input text length, and MSE() means the mean squared error loss function.
In this work, the (unnormalized) attention matrix A
i is used as the ﬁtting target instead of its softmax output softmax ( A i ) , since our experiments show that the former setting has a faster convergence rate and better performances.
In addition to the attention based distillation, we also distill the knowledge from the output of Transformer layer (as shown in Figure 1 (b)), and the objective is as follows:
[Equation]
[Equation]
[Figure]
where the matrices H S ∈ R l × d (cid:48) and H T ∈ R l × d refer to the hidden states of student and teacher networks respectively, which are calculated by Equation 4.
The scalar values d and d (cid:48) denote the hidden sizes of teacher and student models, and d (cid:48) is often smaller than d to obtain a smaller student network.
The matrix W h ∈ R d (cid:48)
× d is a learnable linear transformation, which transforms the hidden states of student network into the same space as the teacher network’s states.
Embedding-layer Distillation .
We also perform embedding-layer distillation, which is similar to the hidden states based distillation and formulated as:
[Equation]
where the matrices E S and H T refer to the embeddings of student and teacher networks, respectively.
In this paper, they have the same shape as the hidden state matrices.
The matrix W e is a linear transformation playing a similar role as W h .
Prediction-Layer Distillation .
In addition to imitating the behaviors of intermediate layers, we also use the knowledge distillation to ﬁt the predictions of teacher model (Hinton et al., 2015).
Speciﬁcally, we penalize the soft cross-entropy loss between the student network’s logits against the teacher’s logits:
[Equation]
where z S and z T are the logits vectors predicted by the student and teacher respectively, log softmax () means the log likelihood, t means the temperature value.
In our experiment, we ﬁnd that t = 1 performs well.
Using the above distillation objectives (i.e. Equations 7, 8, 9 and 10), we can unify the distillation loss of the corresponding layers between the teacher and the student network: [Equation] layer
[Equation]
In our experiments, we ﬁrstly perform intermediate layer distillation ( M ≥ m ≥ 0 ), then perform the prediction-layer distillation ( m = M + 1 ).
3. T INY BERT L EARNING
The application of BERT usually consists of two learning stages: the pre-training and ﬁne-tuning.
The plenty of knowledge learned by BERT in the pre-training stage is of great importance and should also be transferred.
Therefore, we propose a novel two-stage learning framework including the general distillation and the task-speciﬁc distillation , as illustrated in Figure 2.
General distillation helps student TinyBERT learn the rich knowledge embedded in teacher BERT, which plays an important role in improving the generalization capability of TinyBERT.
The task-speciﬁc distillation teaches the student the task-speciﬁc knowledge.
With the two-step distillation, we can further reduce the gap between teacher and student models.
General Distillation .
In general distillation, we use the original BERT without ﬁne-tuning as the teacher and a large-scale text corpus as the learning data.
By performing the Transformer distillation 5 on the text from general domain, we obtain a general TinyBERT that can be ﬁne-tuned for 5 In the general distillation, we do not perform prediction-layer distillation as Equation 10.
[Table] [Table]
downstream tasks.
However, due to the signiﬁcant reductions of the hidden/embedding size and the layer number, general TinyBERT performs relatively worse than BERT.
Task-speciﬁc Distillation .
Previous studies show that the complex models, ﬁne-tuned BERTs, suffer from over-parametrization for domain-speciﬁc tasks (Kovaleva et al., 2019).
Thus, it is possible for small models to achieve comparable performances to the BERTs.
To this end, we propose to derive competitive ﬁne-tuned TinyBERTs through the task-speciﬁc distillation.
In the task-speciﬁc distillation, we re-perform the proposed Transformer distillation on an augmented task-speciﬁc dataset (as shown in Figure 2).
Speciﬁcally, the ﬁne-tuned BERT is used as the teacher and a data augmentation method is proposed to expand the task-speciﬁc training set.
Learning more task-related examples, the generalization capabilities of student model can be further improved.
In this work, we combine a pre-trained language model BERT and GloVe (Pennington et al., 2014) word embeddings to do word-level replacement for data augmentation.
Speciﬁcally, we use the language model to predict word replacements for single-piece words (Wu et al., 2019), and use the word embeddings to retrieve the most similar words as word replacements for multiple-pieces words.
Some hyperparameters are deﬁned to control the replacement ratio of a sentence and the amount of augmented dataset.
More details of the data augmentation procedure are discussed in Appendix A.
The above two learning stages are complementary to each other: the general learning provides a good initialization for the task-speciﬁc distillation, while the task-speciﬁc learning further improves TinyBERT by focusing on the task-speciﬁc knowledge.
Although there is a big gap between BERT and TinyBERT in model size, by performing the proposed two-stage distillation, the TinyBERT can achieve competitive performances in various NLP tasks.
The proposed Transformer distillation and two-stage learning framework are the two most important components of the proposed distillation method.
E XPERIMENTS
In this section, we evaluate the effectiveness and efﬁciency of TinyBERT on a variety of tasks with different model settings.
4. M ODEL S ETUP
We instantiate a tiny student model (the number of layers M =4, the hidden size d (cid:48)
=312, the feedforward/ﬁlter size d (cid:48) i =1200 and the head number h =12) that has a total of 14.5M parameters.
If not speciﬁed, this student model is referred to as the TinyBERT.
The original BERT BASE (the number of layers N =12, the hidden size d =768, the feed-forward/ﬁlter size d i =3072 and the head number h =12) is used as the teacher model that contains 109M parameters.
We use g ( m ) =
× m as the layer mapping function, so TinyBERT learns from every 3 layers of BERT BASE .
The learning weight λ of each layer is set to 1, which performs well for the learning of our TinyBERT.
[Table]
4. E XPERIMENTAL R ESULTS ON GLUE
We evaluate TinyBERT on the General Language Understanding Evaluation (GLUE)
(Wang et al., 2018) benchmark, which is a collection of diverse natural language understanding tasks.
The details of experiment settings are described in Appendix B.
The evaluation results are presented in Tableand the efﬁciencies of model size and inference time are also evaluated in Table
[Table]
The experiment results demonstrate that: 1) There is a large performance gap between BERT SMALLand BERT BASE due to the big reduction in model size.
2) TinyBERT is consistently better than
BERT SMALL in all the GLUE tasks and achieves a large improvement of 6.3% on average.
This indicates that the proposed KD learning framework can effectively improve the performances of small models regardless of downstream tasks.
3) TinyBERT signiﬁcantly outperforms the state-ofthe-art KD baselines (i.e., BERT-PKD and DistillBERT) by a margin of at least 3.9%, even with only ∼ 28% parameters and ∼ 31% inference time of baselines (see Table 3).
Compared with the teacher BERT BASE , TinyBERT is 7.5x smaller and 9.4x faster in the model efﬁciency, while maintaining competitive performances.
5) TinyBERT has a comparable model efﬁciency (slightly larger in size but faster in inference) with Distilled BiLSTM SOFT and obtains substantially better performances in all tasks reported by the BiLSTM baseline.
For the challenging CoLA dataset (the task of predicting linguistic acceptability judgments), all the distilled small models have a relatively bigger performance gap with teacher model.
TinyBERT achieves a signiﬁcant improvement over the strong baselines, and its performance can be further improved by using a deeper and wider model to capture more complex linguistic knowledge as illustrated in the next subsection.
Moreover, BERT-PKD and DistillBERT initialize their student models with some layers of well pre-trained teacher BERT (see Table 1), which makes the student models have to keep the same size settings of Transformer layer (or embedding layer) as their teacher BERT.
In our two-stage distillation framework, TinyBERT is initialized by general distillation, so it has the advantage of being more ﬂexible in model size selection.
4. E FFECTS OF M ODEL S IZE
We evaluate how much improvement can be achieved when increasing the model size of TinyBERT on several typical GLUE tasks, where MNLI and MRPC are used in the ablation studies of Devlin et al.
(2018), and CoLA is the most difﬁcult task in GLUE.
Speciﬁcally, three wider and deeper variants are proposed and their evaluation results on development set are displayed in Table 4.
We can observe that: 1) All the three TinyBERT variants can consistently outperform the original smallest TinyBERT, which indicates that the proposed KD method works for the student models of various model sizes.
For the CoLA task, the improvement is slight when only increasing the number of layers (from 49.7 to 50.6) or hidden size (from 49.7 to 50.5).
To achieve more dramatic improvements, the student model should become deeper and wider (from 49.7 to 54.0).
Another interesting observation is that the smallest 4-layer TinyBERT can even outperform the 6-layers baselines, which further conﬁrms the effectiveness of the proposed KD method.
4. A BLATION S TUDIES
In this section, we conduct ablation studies to investigate the contributions of : 1) different procedures of the proposed two-stage TinyBERT learning framework (see Figure 2), and 2) different distillation objectives (see Equation 11).
[Table]
[Table] [Table]
Effects of different learning procedures.
The proposed two-stage TinyBERT learning framework (see Figure 2) consists of three key procedures: TD (Task-speciﬁc Distillation), GD (General
Distillation) and DA (Data Augmentation).
The effects of different learning procedures are analyzed and presented in Table 5.
The results indicate that all the three procedures are crucial for the proposed KD method.
The TD and DA has comparable effects in all the four tasks.
We can also ﬁnd the task-speciﬁc procedures (TD and DA) are more helpful than the pre-training procedure (GD) in all the four tasks.
Another interesting observation is that GD has more effect on CoLA than on
MNLI and MRPC.
We conjecture that the ability of linguistic generalization (Warstadt et al., 2018) learned by GD plays a more important role in the downstream CoLA task (linguistic acceptability judgments).
Effects of different distillation objectives.
We investigate the effects of distillation objectives on the TinyBERT learning.
Several baselines are proposed including the TinyBERT learning without the Transformer-layer distillation (No Trm), embedding-layer distillation (No Emb) and predictionlayer distillation (No Pred) 7 respectively.
The results are illustrated in Table 6 and show that all the proposed distillation objectives are useful for the TinyBERT learning.
The performance drops signiﬁcantly from 75.3 to 56.3 under the setting (No Trm), which indicates Transformer-layer distillation is the key for TinyBERT learning.
Furthermore, we study the contributions of attention (No Attn) and hidden states (No Hidn) in the Transformer-layer distillation.
We can ﬁnd the attention based distillation has a bigger effect than hidden states based distillation on TinyBERT learning.
Meanwhile, these two kinds of knowledge distillation are complementary to each other, which makes TinyBERT obtain the competitive results.
4. E FFECTS OF M APPING F UNCTION
We investigate the effects of different mapping functions n = g ( m ) on the TinyBERT learning.
Our original TinyBERT as described in section 4.1 uses the uniform-strategy, and we compare with two typical baselines including top-strategy ( g ( m ) = m
+ N − M ; 0 < m ≤ M ) and bottom-strategy ( g ( m ) = m ; 0 < m ≤ M ) .
The comparison results are presented in Table 7.
We ﬁnd that the top-strategy performs better than the bottom-strategy in MNLI, while being worse in MRPC and CoLA tasks, which conﬁrms the observations that different tasks depend on the different kinds of knowledge from BERT layers.
Since the uniform-strategy acquires the knowledge from bottom to top layers of BERT BASE , it achieves better performances than the other two baselines in all the four tasks.
Adaptively choosing layers for a speciﬁc task is a challenging problem and we leave it as the future work.
Other Experiments.
We also evaluate TinyBERT on the question answering tasks, and study whether we can use BERT SMALL as the initialization of the general TinyBERT.
The experiments are detailed in Appendix C and D.
C ONCLUSION AND F UTURE W ORK
In this paper we ﬁrstly introduce a new KD method for Transformer-based distillation, then we further propose a two-stage framework for TinyBERT learning.
Extensive experiments show that the TinyBERT achieves competitive performances meanwhile signiﬁcantly reducing the model size and shortening the inference time of original BERT BASE , which provides an effective way to deploy
BERT-based NLP applications on the edge devices.
In future work we would study how to effectively transfer the knowledge from wider and deeper teachers (e g., BERT LARGE and XLNet LARGE ) to student TinyBERT.
The joint learning of distillation and quantization/pruning would be another promising direction to further compress the pre-trained language models.
A PPENDIX A D ATA A UGMENTATION D ETAILS
In this section, we explain the proposed data augmentation method.
Speciﬁcally, we ﬁrstly mask a word in a sentence, then use BERT as a language model to predict M most-likely words in the corresponding position, while keeping other words unchanged.
By this way, we can easily get the candidates for each word under a speciﬁc context.
To induce new instances for a given sentence, we also use a threshold p t to determine whether we should replace the current word with a randomly selected candidate.
By repetitively performing this replacement operation for each word in a sentence, we can ﬁnally get a new augmented sentence.
In our preliminary experiments, we ﬁnd the quality of generated candidates for the words consisting of multiple sub-word pieces, is relatively low.
To alleviate this problem, we instead pick a similar word from GloVe (Pennington et al., 2014) word embeddings based on the cosine similarity.
We apply this data augmentation method N times to all the sentences of a downstream task.
In this work, we set p t = 0 .
4 , N = 20 , M = 15 for all our experiments.
The data augmentation procedure is illustrated as below:
Algorithm
The Proposed Data Augmentation
Input : x is a sequence of words p t , N, M are hyperparameters
Output : data aug , the augmented data 1: function DATA AUGMENTANTION ( x , p t , N ) 2: n ←3: data aug ← [ ] 4: while n < N do 5: x masked ← x 6: for i ← 1 to len ( x ) do 7: if x [ i ] is a single-piece word then 8:
Replace x masked [ i ] with [MASK] 9: candidates ← M most-likely words predicted by BertModel ( x masked )
[ i ] 10: else 11: candidates ← M similar words of x
[ i ] from GloVe 12: end if 13:
Sample p ∼ UNIFORM (0 , 1) 14: if p ≤ p t then 15:
Replace x masked [ i ] with a word from candidates randomly 16: else 17:
Keep x masked [ i ] as x [ i ] unchanged 18: end if 19: end for 20:
Append x masked to data aug 21: n = n +22: end while 23: return data aug 24: end function
B GLUE D ETAILS
TinyBERT setup.
TinyBERT learning includes the general distillation and the task-speciﬁc distillation.
For the general distillation, we use English Wikipedia (2,500 M words) as the text corpus and perform the intermediate layer distillation for 3 epochs with the supervision from a pre-trained
BERT BASE teacher and keep other hyper-parameters same as BERT pre-training (Devlin et al., 2018).
For the task-speciﬁc distillation, we ﬁrstly perform intermediate layer distillation on the augmented dataset for 10 epochs with batch size 32 and learning rate 5e-5 under the supervision of a ﬁne-tuned BERT teacher, and then perform prediction layer distillation for 3 epochs with batch size 32 and learning rate 3e-5.
For tasks like MNLI, QQP and QNLI which have ≥ 100 K training ex- amples, we distill intermediate layer knowledge for 5 epochs with batch size 256 on the augmented dataset.
Besides, for CoLA task, we perform 50 epochs of intermediate layer distillation .
Baselines setup.
We use BERT-PKD and DistilBERT as our baselines.
For a fair comparison, we ﬁrstly re-implemented the results of BERT-PKD and DistilBERT reported in their papers to ensure our implementation procedure is correct.
Then following the veriﬁed implementation procedure, we trained a 4-layer BERT-PKD and a 4-layer DistilBERT as the baselines.
The BERT SMALL learning strictly follows the same learning strategy as described in the original BERT work (Devlin et al., 2018).
Model efﬁciency evaluation.
To evaluate the inference speed, we ran inference procedure on the
QNLI training set with batch size of 128 and the maximum sequence length of 128.
The numbers reported in Table 3 are the average running time of 100 batches on a single NVIDIA K80 GPU.
The GLUE datasets are described as follows:
MNLI.
Multi-Genre Natural Language Inference is a large-scale, crowd-sourced entailment classiﬁcation task (Williams et al., 2018).
Given a pair of (cid:104) premise, hypothesis (cid:105) , the goal is to predict whether the hypothesis is an entailment, contradiction, or neutral with respect to the premise .
QQP.
Quora Question Pairs is a collection of question pairs from the website Quora.
The task is to determine whether two questions are semantically equivalent (Chen et al., 2018).
QNLI.
Question Natural Language Inference is a version of the Stanford Question Answering Dataset which has been converted to a binary sentence pair classiﬁcation task by Wang et al.
Given a pair (cid:104) question, context (cid:105) .
The task is to determine whether the context contains the answer to the question .
SST-2.
The Stanford Sentiment Treebank is a binary single-sentence classiﬁcation task, where the goal is to predict the sentiment of movie reviews (Socher et al., 2013).
CoLA.
The Corpus of Linguistic Acceptability is a
task to predict whether an English sentence is a grammatically correct one (Warstadt et al., 2018).
STS-B.
The Semantic Textual Similarity Benchmark is a collection of sentence pairs drawn from news headlines and many other domains (Cer et al., 2017).
The task aims to evaluate how similar two pieces of texts are by a score from 1 to 5.
MRPC.
Microsoft Research Paraphrase Corpus is a paraphrase identiﬁcation dataset where systems aim to identify if two sentences are paraphrases of each other (Dolan & Brockett, 2005).
RTE.
Recognizing Textual Entailment is a binary entailment task with a small training dataset (Bentivogli et al., 2009).
C SQ U AD 1.1 AND 2.
We also demonstrate the effectiveness of TinyBERT on the question answering (QA) tasks: SQuAD v1.1 (Rajpurkar et al., 2016) and v2.0 (Rajpurkar et al., 2018).
Following the learning procedure in the previous work (Devlin et al., 2018), we treat these two tasks as the problem of sequence labeling which predicts the possibility of each token as the start or end of answer span.
We follow the settings of task-speciﬁc distillation in GLUE tasks, except with 3 running epochs and a learning rate of 5efor the prediction-layer distillation on the original training dataset.
The results are shown in Table 8.
The results show that TinyBERT consistently outperforms the baselines in both the small and medium size, which indicates that the proposed framework also works for the tasks of token-level labeling.
Compared with sequence-level GLUE tasks, the question answering tasks depends on more subtle knowledge to infer the correct answer, which increases the difﬁculty of knowledge distillation.
We leave how to build a better QA-TinyBERT as the future work.
D BERT SMALL AS I NITIALIZATION OF G ENERAL T INY BERT
Initializing general TinyBERT with BERT SMALL is a straightforward idea.
However, BERT SMALL would derive mismatched distributions in intermediate representations (e.g., attention matrices
[Table] [Table]
and hidden states) with the teacher BERT BASE model, if without imitating the teacher’s behaviors at the pre-training stage.
Further task-speciﬁc distillation under the supervision of ﬁnetuned BERT BASE will disturb the learned distribution/knowledge of BERT SMALL , ﬁnally leading to poor performances in some less-data tasks.
The results in Table 9, show that the
BERT SMALL (MLM&NSP+TD) performs worse than the BERT SMALL in MRPC and CoLA tasks, which validates our hypothesis.
For the intensive-data task (e.g. MNLI), TD has enough training data to make BERT SMALL acquire the task-speciﬁc knowledge very well, although the pre-trained distributions have already been disturbed.
To make TinyBERT effectively work for all tasks, we propose General Distillation (GD) for initialization, where the TinyBERT learns the knowledge from intermediate layers of teacher BERT at the pre-training stage.
From the results of Table 9, we ﬁnd that GD can effectively transfer the knowledge from the teacher BERT to the student TinyBERT and achieve comparable results with
BERT SMALL (61.1 vs. 63.9), even without performing the MLM and NSP tasks.
Furthermore, the task-speciﬁc distillation boosts the performances of TinyBERT by continuing on learning the task-speciﬁc knowledge of ﬁne-tuned teacher BERT BASE .
E M ORE C OMPLETE C OMPARISONS WITH S AME S TUDENT A RCHITECTURE
For the easy and direct comparisons with prior works, we here also present the results of TinyBERT ( M =6; d (cid:48)
=
768; d (cid:48) i =3072) with the same architectures as used in the original BERT-PKD (Sun et al., 2019) and DistilBERT 2 .
Since in the original papers, the BERT-PKD is evaluated on the test set, and the DistilBERT is evaluated on the dev set.
Thus, for a clear illustration, we present the results in the following Tables 10 and 11, separately
[Table]
[Table]
[Table]
Thus, from the direct comparisons with the reported results in the original papers, we can see the
TinyBERT outperforms the baselines (DistilBERT and BERT-PKD) under the same settings of architecture and evaluation, the effectiveness of TinyBERT is conﬁrmed.
Moreover, since BERT-PKD and DistilBERT need to initialize their student models with some layers of pre-trained teacher BERT, they have the limitations that the student models have to keep the same size settings of hidden size and feedforward/ﬁlter size as their teacher BERT.
TinyBERT is initialized by general distillation, so it has the advantage of being more ﬂexible in model size selection.