SKDBERT: Compressing BERT via Stochastic Knowledge Distillation

In this paper, we propose Stochastic Knowledge Distillation (SKD) to obtain compact BERT-style language model dubbed SKDBERT. In each iteration, SKD samples a teacher model from a pre-defined teacher ensemble, which consists of multiple teacher models with multi-level capacities, to transfer knowledge into student model in an one-to-one manner. Sampling distribution plays an important role in SKD. We heuristically present three types of sampling distributions to assign appropriate probabilities for multi-level teacher models. SKD has two advantages: 1) it can preserve the diversities of multi-level teacher models via stochastically sampling single teacher model in each iteration, and 2) it can also improve the efficacy of knowledge distillation via multi-level teacher models when large capacity gap exists between the teacher model and the student model. Experimental results on GLUE benchmark show that SKDBERT reduces the size of a BERT$_{\rm BASE}$ model by 40% while retaining 99.5% performances of language understanding and being 100% faster.


A Multi-teacher Knowledge Distillation
For BERT-style language model compression, we verify the performance of MTBERT (Wu, Wu, and Huang 2021) whose objection function can be expressed as: where, N indicates the number of teacher models, CE(·, ·) is the cross-entropy loss, T denotes the temperature, y represents the ground-truth label, y i and y s refer to the outputs of i-th teacher model and the student model, respectively. We employ the student model used in Tiny-BERT (Jiao et al. 2020) which can be downloaded from https://huggingface.co/huawei-noah/TinyB ERT General 6L 768D. Furthermore, we choose five BERT-style models whose architecture information is shown in Table 1, as the teacher models from the BERT official website 1 . Furthermore, we fine tune the above teacher models on six downstream tasks of GLUE benchmark, i.e., MRPC, RTE, SST-2, QQP, QNLI and MNLI whose detailed introduction can be found in Section D. Moreover, the experimental settings of teacher models fine-tuning are shown as follows: • Learning Rate: For T 3 and T 4 , {6e-6, 7e-6, 8e-6, 9e-6} on MRPC and RTE tasks, and {2e-5, 3e-5, 4e-5, 5e-5} on other tasks. For other teacher models, {2e-5, 3e-5, 4e-5, 5e-5} on all tasks.
• Batch Size: {32} for teacher model fine-tuning on each task.
Other settings are following TinyBERT (Jiao et al. 2020). Subsequently, we employ the ensemble of the fine-tuned teacher models to compress the student model via (1) . Particularly, we only use the weighted multi-teacher distillation loss without the multi-teacher hidden loss and the taskspecific loss as in MT-BERT (Wu, Wu, and Huang 2021).

B SKD for Image Classification
To verify the effectiveness of general SKD on computer vision, we conduct a list of experiments on CIFAR-100 (Krizhevsky, Hinton et al. 2009) image classification dataset. Moreover, we employ ResNet (He et al. 2016) and Wide ResNet (WRN) (Zagoruyko and Komodakis 2016) with various architectures as the candidates of teacher team to distill a student model. Following the experimental settings in CRD (Tian, Krishnan, and Isola 2020), we train the student model and teacher model for 240 epochs, and employ SGD as the optimizer with a batch size of 64, a learning rate of 0.05 which is decayed by a factor of 0.1 when arriving at 150-th, 180-th, 210-th epoch and a weight decay of 5e-4, and show the results in Table 2. Similar to Tian, Krishnan, and Isola (2020), we show the test accuracy of the last epoch in Table 3 for a fair comparison. The proposed SKD also achieves novel performance on CIFAR-100. Compared to state-of-the-art CRD, our approach achieves the best performance in five out of six groups of distillation experiments.

C Why Does SKD Work?
In this section, taking uniform distribution based SKD as an example, we discuss why does it work.

C.1 Preliminaries: Knowledge Distillation
On classification task with C categories, vanilla KD employs the output of final layer, logits z z z = [z 1 , z 2 , . . . , z C ] T , to yield soft targets p p p = [p 1 , p 2 , . . . , p C ] T for knowledge transfer as where, T represents the temperature factor which is used to control the importance of p i . The total loss L with respect to input x x x and weights of student model W W W , consists of distilled loss L d and student loss L s as and L s (y y y, p(z z z s , T )) = C i=1 −y i ln(p(z si , T )) treat soft targets p p p(z z z t , T ) and ground truth y y y as their labels, respectively. Moreover, z z z t and z z z s represent the logits of teacher and student models, respectively. Particularly, an ensemble of multiple teacher models is used to provide

C.2 The Convergence of SKD
In order to optimize W W W , we should evaluate the crossentropy gradient of total loss with regard to z si , Here, the second term of the right side of (2) of SKD is identical with vanilla KD, so that we only discuss the first term. The cross-entropy gradient of distilled loss with regard to z si can be computed as For ∂ln(p(zsj ,T )) ∂zsi , there are two situations (i.e., i = j and Subsequently, we replace ∂ln(p(zsj ,T )) ∂zsi in (3) with (4), and it can be rewritten as On the one hand, in vanilla multi-teacher KD framework with an ensemble of N teacher models, z ti is the average of ensemble as so that its soft target can be computed as .
On the other hand, SKD randomly employs a single teacher from N teacher models to obtain the soft target at n-th iteration as where, the sampled probability of z z z (n) t subjects to uniform distribution U (1, N ).
We assume that the number of iterations is L. In the total training process of multi-teacher KD, the amount of gradient update on z i dubbed G (i) Avg can be represented as Beyond that, the amount of gradient update of SKD on z i dubbed G (i) SKD can be obtained by Obviously, the main difference between G  (5) as According to (7) and the term b i in (6), we can find that their relationship is similar to the relationship between Batch Gradient Descent (BGD) and Stochastic Gradient Descent (SGD). Moreover, BGD employs total data to obtain the gradient in each iteration, but SGD uses single one. Based on stochastic approximation (e.g. ROBBINS-MONRO's theory (Robbins and Monro 1951)) and convex optimization, the convergence of SGD has been proven (Turinici 2021). MRPC belongs to a sentence similarity task where system aims to identify the paraphrase/semantic equivalence relationship between two sentences.

D Details of GLUE Benchmark
RTE belongs to a natural language inference task where system aims to recognize the entailment relationship of given two text fragments.
CoLA belongs to a single-sentence task where system aims to predict the grammatical correctness of an English sentence.
STS-B belongs to a sentence similarity task where system aims to evaluate the similarity of two pieces of texts by a score from 1 to 5.
SST-2 belongs to a single-sentence task where system aims to predict the sentiment of movie reviews.
QQP belongs to a sentence similarity task where system aims to identify the semantical equivalence of two questions from the website Quora.
QNLI belongs to a natural language inference task where system aims to recognize that for a given pair <question, context>, the answer to the question whether contains in the context.
MNLI belongs to a natural language inference task where system aims to predict the possible relationships (i.e., entailment, contradiction and neutral) of hypothesis with regard to premise for a given pair <premise, hypothesis>,.
WNLI belongs to a natural language inference task where system aims to determine the referent of a sentence's pronoun from a list of choices.

E Evaluation Metrics and Hyper-parameters E.1 Evaluation Metrics
In order to choose the best model, we use appropriate metrics on GLUE-dev as shown in Table 4.

E.2 Fine-tuning Hyper-parameters
In this paper, we need to perform teacher model fine-tuning, student model fine-tuning and student model distillation. For the above three cases, the used hyper-parameters are shown in Table 5.

F Student Model and Teacher Team Details
The architecture information of each student model and teacher model is shown in Table 6. On the one hand, we direct treating the pre-trained model of TinyBERT 4 2 and 2 https://huggingface.co/huawei-noah/TinyBE RT General 4L 312D  TinyBERT 6 3 as the student models of SKDBERT. On the other hand, we choose 14 BERT models with various capabilities as the candidates for teacher team. Moreover, each pre-trained teacher model can be downloaded from official implementation of BERT 4 . Furthermore, the results of the student model and the teacher model on GLUE-dev are shown in Table 7 which can be used to obtain the teacherrank sampling distribution.

G Distillation Performance of SKDBERT with Various Teacher Models
In order to determine the student-rank sampling distribution for SKD, we employ each teacher model to distill  SKDBERT 4 and SKDBERT 6 on GLUE-dev, and show the results in Table 8.

H The Used Sampling Distributions for SKDBERT with Various Teacher Teams on The GLUE Benchmark
We implement extensive ablation experiments for SKD-BERT using various teacher teams: {T 04 -T 06 , T 07 -T 09 , T 10 -T 12 , T 12 -T 14 } for SKDBERT 4 and {T 04 -T 06 , T 07 -T 09 , T 10 -T 12 , T 04 -T 09 , T 04 -T 12 , T 01 -T 14 , T 04 -T 14 , T 07 -T 14 , T 09 -T 14 , T 12 -T 14 , T 13 -T 14 } for SKDBERT 6 . We show the used dis-tributions for the above teacher teams on various tasks in Figure 1 to 15. Particular, 0 indicates that the teacher models are not employed for stochastic knowledge distillation.