Pea-KD: Parameter-efficient and accurate Knowledge Distillation
Authors: IKHYUN CHO
U Kang
Venue:  (2021)
Keywords: BERT
Deep Learning
Natural Language Processing
Transformer
Knowledge Distillation
Parameter Sharing

Abstract: How can we efficiently compress a model while maintaining its performance?  Knowledge Distillation (KD) is one of the widely known methods for model compression. In essence, KD trains a smaller student model based on a larger teacher model and tries to retain the teacher model's level of performance as much as possible. However, the existing KD methods suffer from the following limitations. First, since the student model is small in absolute size, it inherently lacks model complexity. Second, the absence of an initial guide for the student model makes it difficult for the student to imitate the teacher model to its fullest. Conventional KD methods yield low performance due to these limitations.

In this paper, we propose Pea-KD (Parameter-efficient and accurate Knowledge Distillation), a novel approach to KD. Pea-KD consists of two main parts: Shuffled Parameter Sharing (SPS) and Pretraining with Teacher's Predictions (PTP). Using this combination, we are capable of alleviating the KD's limitations. SPS is a new parameter sharing method that allows greater model complexity for the student model. PTP is a KD-specialized initialization method, which can act as a good initial guide for the student. When combined, this method yields a significant increase in student model's performance. Experiments conducted on different datasets and tasks show that the proposed approach improves the student model's performance by 4.4% on average in four GLUE tasks, outperforming existing KD baselines by significant margins.


Under review as a conference paper at ICLR 2021
PEA-KD: PARAMETER-ERFICIEN AND ACCURATE KNOWLEDGE DISTILLATION ON BERT
Anonymous authors
Paper under double-blind review
ABSTRACT
How can we efficiently compress a model while maintaining its performance? Knowledge Distillation (KD) is one of the widely known methods for model com- pression. In essence, KD trains a smaller student model based on a larger teacher model and tries to retain the teacher model's level of performance as much as possible. However, existing KD) methods suffer from the following limitations. First, since the student model smaller in absolute size, it inherently lacks model capacity. Second, the absence of initial guide for the student model makes it difficult for the student to imitate the teacher model to its fullest. Conventional KD methods yield low performance due to these limitations.
In this paper, we propose Pea-KD (Parameter-efficien and accurate Knowledge Distillation), a novel approach to KD. Pea-KD consists of two main parts: Shuf- fled Parameter Sharing (SPS) and Pretraining with Teacher's Predictions (PTP). Using this combination, we are capable of alleviating the KD's limitations. SPS is a new parameter sharing method that increases the student model capacity. PTP is a KD-specialized initialization method, which can act good initial guide for the student. When combined, this method yields a significant increase in student model's performance. Experiments conducted and tasks show that the proposed approach improves the student model's per- formance by 4.4% on average in four GLUE tasks, outperforming existing KD baselines by significant margins.
BERT with different datasets
INTRODUCTION
How can we improve the accuracy of knowledge distillation (KD) with smaller number of param- eters? KD uses a well-trained large teacher model to train smaller student model. Conventional KD method (Hinton et al.] (2006)) trains the student model using the teacher model's predictions as targets. That is, the student model uses not only the true labels (hard distribution) but also the teacher model's predictions (soft distribution) as targets. Since better KD accuracy is directly linked to better model compression, improving KD accuracy is valuable and crucial.
Naturally, there have been many studies and attempts to improve the accuracy of KD. Sun et al. (2019) introduced Patient KD, which utilizes not only the teacher model's final output but also the intermediate outputs generated from the teacher's layers. Jiao et al. (2019) applied additional KD in the pretraining step of the student model. However, existing KD methods share the limitation of students having lower model capacity compared to their teacher models, due to their smaller size. In addition, there is proper initialization guide established for the student models, which especially becomes important when the student model is small. These limitations lower than desired levels of student model accuracy.
In this paper, we propose Pea-KD (Parameter-cfficient and accurate Knowledge Distillation), a novel KD method designed especially for Transformer-based models (Vaswani al. (2017)) that signif- icantly improves the student model's accuracy. Pea-KD is composed of two modules, Shuffled Parameter Sharing (SPS) and Pretraining with Teacher's Predictions (PTP). Pea-KD is based on the following two main ideas.
We apply SPS in order to increase the effective model capacity of the student model with- out increasing the number of parameters. SPS has two steps: stacking layers that share

---

Under review as a conference paper at ICLR 2021
parameters and 2) shuffling the parameters between shared pairs of layers. Doing in- creases the model's effective capacity which enables the student to better replicate the teacher model (details in Section 3.2).
2. We apply a pretraining task called PTP for the student. Through PTP, the student model learns general knowledge about the teacher and the task. With this additional pretraining, the student more efficiently acquires and utilizes the teacher's knowledge during the actual KD process (details Section 3.3
Throughout the paper we use Pea-KD applied on BERT model (PeaBERT) as an example to inves- tigate our proposed approach. We summarize our main contributions as follows:
Novel framework for KD. We propose SPS and PTP, parameter sharing method and KD-specialized initialization method. These methods serve as new framework for KD to significantly improve accuracy.
Performance. When tested on four widely used GLUE tasks, PeaBERT improves student's accuracy by 4.4% on average and up 14.8% maximum when compared to the original BERT model. PeaBERT also outperforms the existing state-of-the-art KD baselines by 3.5% on average.
Generality. Our proposed method Pea-KD can be applied to any transformer-based models and classification tasks with small modifications.
Then conclude in Section 5,
2 RELATED WORK
Pretrained Language Models. The framework of first pretraining language models and then fine- tuning for downstream tasks has now become the industry standard for Natural Language Processing (NLP) models. Pretrained language models, such as BERT (Devlin et al. (018)), XLNet (Yang et al. (201 RoBERTa (Liu al. (2019) and ELMo (Peters et (2018) prove how powerful pretrained language models can be. Specifically, BERT is a language model consisting of multiple transformer encoder layers. Transformers (Vaswani al. (2017) can capture long-term dependencies between input tokens by using a self-attention mechanism. Self-attention calculates attention function using three components, query, key, and value, each denoted matrices Q, K, and V. The attention function is defined follows:
A QKT
VTT
Attention(Q2 (Q,K, K, V) = (ttx(())))))
(2)
is known that through pretraining using Masked Language Modeling (MLM) and Next Sentence Prediction (NSP), the attention matrices in BERT capture substantial linguistic knowledge. BERT has achieved the state-of-the-art performance on a wide range of NLP tasks, such as the GLUE benchmark (Wang al. (2018)) and SQuAD (Rajpurkar al. (2016)).
However, these modern pretrained models are very large in size and contain millions of parameters, making them nearly impossible to apply on edge devices with limited amount of resources.
Model Compression. As deep learning algorithms started getting adopted, implemented, and re- searched in diverse fields. high computation costs and memory shortage have started to become challenging factors. Especially in NLP, pretrained language models typically require a large set of parameters. This results in extensive cost of computation and memory. As such, Model Compres- sion has now become important task for deep learning. There have already been many attempts to tackle this problem, including quantization (Gong et al. (2014)) and weight pruning (Han et al. (2015)). Two promising approaches are KD (Hinton et al. (2015)) and Parameter Sharing, which tocus on this paper.
Knowledge Distillation (KD). As briefly covered in Section I KD transfers knowledge from well-trained and large teacher model to a smaller student model. KD uses the teacher models pre- dictions on top of the true labels train the student model. It is proven through many experiments
2

---

Under review as a conference paper at ICLR 2021
that the student model learns to imitate the soft distribution of the teacher model's predictions and ultimately performs better than learning solely from the original data. There have already been many attempts compress BERT using KD. Patient Knowledge Distillation (Sun et al. (2019)) ex- tracts knowledge not only from the final prediction of the teacher, but also from the intermediate layers. TinyBERT (Jiao al. (20T9)) uses a two-stage learning framework and applies KD in both pretraining and task-specific finetuning, DistilBERT (Sanh et al. (019)) uses half of the layers of BERT-base model and applies KD during pretraining and finetuning of BERT. Insufficient capac- ity and the absence of a clear initialization guide are some of the existing KD method's area of improvement.
Parameter Sharing. Sharing parameters across different layers is widely used idea for model compression. There have been several attempts to apply parameter sharing in transformer architec- ture and BERT model. However, the existing parameter sharing methods exhibit a large tradeoff between model performance and model size. They reduce the model's size significantly but also suffers from great loss in performance as a result.
PROPOSED METHODS
In the following, we provide an overview of the main challenges faced in KD and our methods to address them in Section We then discuss the precise procedures of SPS and PTP in Sections 3.2 and 3.3 Lastly, we explain our final method, PeaBERT and the training details in Section 3.4
3.1 OVERVIEW
BERT-base model contains over 109 million parameters. Its extensive size makes model deployment often infeasible and computationally expensive in many cases, such mobile devices. As result, industry practitioners commonly use a smaller version of BERT and apply KD. However, the existing KD methods face the following challenges:
1. Insufficient model capacity of the student model. Since the student model contains fewer number of parameters than the teacher model, its model capacity also lower. The smaller and simpler the student model gets, the gap between the student and the teacher grows, making it increasingly difficult for the student to replicate the teacher model's accuracy. The limited capacity hinders the student model's performance. How can we enlarge the student model's ' capacity while maintaining the same number of parameters?
2. Absence of proper initial guide for the student model. There no widely accepted and vetted guide to selecting the student's initial state of the KD process. In most cases, truncated version of pretrained BERT-base model is used. reality, this hinders the student from reproducing the teacher's results. Can we find better method for the student's 's KD initialization?
We propose the following main ideas address these challenges:
Shuffled Parameter Sharing (SPS): increasing "effective' model capacity the stu- dent. SPS consists of two steps. Step 1 increases the student's effective model capacity by stacking parameter-shared layers. Step 2 further improves the model capacity by shuffling the shared parameters. Shuffling enables more efficient use of the parameters by leamn ing diverse set of information. As a result, the SPS-applied student model achieves much higher accuracy while still using the same number of parameters (details in Section 2. Pretraining with Teacher's Predictions (PTP): a novel pretraining task utilizing teacher's predictions for student initialization. To address the limitation of the initial guide, we propose PTP, novel pretraining method for the student by utilizing teacher model's predictions. Through PTP, the student model pre-learns knowledge latent in the teacher's softmax output which is hard to learn during the conventional KD process. This helps the student better acquire and utilize the teacher's knowledge during the KD process (details in Section 3.3).
The following subsections describe the procedures of SPS, PTP, and PeaBERT in detail.
3

---

Under review as a conference paper at ICLR 2021
3.2
SHUFFLED PARAMETER SHARING (SPS)
Layer
Layer
Q 6+
Q6-K
Q
Shared
Layer
Layer
Q+ K
Q6+X
& Shuffled
Step
Layer
Shared Layers
Layer
Step
Q K
Shared
SPS
Layers
Layer
Layer
Layer
Shuffled
Layers
Layer
Layer
Layer
Layer
Layer
Layer
(a) SPS step
(b) SPS step 1 + step 2
(c) SPS for 6-layer stud.
Figure I: Graphical representation of SPS: (a) first step of SPS (b) second step of SPS, and (c)
modified SPS for a 6-layer student.
SPS improves student's model capacity while using the same number of parameters, addressing the capacity limitations of a typical KD. SPS is composed the following two steps.
Step1. Paired Parameter Sharing. We start with doubling the number of layers in the student model. We then share the parameters between the bottom half and the upper half of the model, as graphically represented in Figure La
Step2. Shuffling. We shuffle the Query and Key parameters between the shared pairs. That is, for the shared upper half of layers use the original Query parameters Key parameters, and the original Key parameters as Query parameters.
We call this architecture SPS, which is depicted in Figure 1b modify the architecture as shown in Figure Tc We apply SPS on the top 3 layers only.
For the 6-layer student case slightly
Motivation behind step 1: Increased model capacity due to applying step 1, the student model now has more layers while maintaining the same number of parameters used. This leads to a higher model complexity, which we expect will increase model capacity and performance.
complex structure. By
Motivation behind step 2: Increased parameter efficiency and additional regularization effect. The intuition behind shuffling becomes clear when we compare the model with and without step 2. Let us consider the case depicted in Figurq[b) In step the first layer's Query parameter is used as a Query parameter in both the first layer and in the shared fourth layer. Since this parameter is used only as a Query parameter throughout the model, will only learn the information relevant to Query. In step 2, on the other hand, this parameter's ' function changes due to shuffling. The first layer's Query parameter is used as Key parameter the shared fourth layer. Thus it has the opportunity to learn the important features of both the Query and the Key, gaining more diverse and wider breadth of knowledge. The shuffling mechanism allows the parameters to learn diverse information, thereby maximizing learning efficiency. We can also view this type of regularization learning. Since these parameters need to learn the features of both the Query and the Key, this prevents overfitting to either one. Therefore, expect the shuffling mechanism improve the generalization ability of the model.
multi-task
3.3 PRETRAINING WITH TEACHER'S PREDICTIONS (PTP)
There can be several candidates for KD-specialized initialization. We propose pretraining approach called PTP, and experimentally show that it improves KD accuracy.
Most of the previous studies on KD do not elaborate on the initialization the student model. There some studies that use a pretrained student model
initial state, but those pretraining tasks are irrelevant to either the teacher model or the downstream task. To the best of our knowledge, our study is the first case that pretrains the student model with task relevant to the teacher model and its downstream task. PTP consists of the following two steps.
Step 1. Creating artificial data based the teacher's predictions (PTP labels).
4

---

Under review as a conference paper at ICLR 2021
We first input the training data in the teacher model and collect the teacher model's predictions. We then define "confidence" as the following. We apply softmax function to the teacher model's predictions, and the maximum value of the predictions is defined the confidence. Next, with specific threshold "t" (a hyperparameter between 0.5 and 1.0), we assign new label to the training data according the rules listed in Table We call these new artificial labels PTP labels.
Table 1: Assigning new PTP labels to the training data.
Teacher's prediction correct
confidence >
PTP label
True
True
confidently correct
True
False
unconfidently correct
False
True
confidently wrong
False
False
unconfidently wrong
Step 2. Pretrain the student model to predict the PTP labels. Using the artificial PTP labels (data n, PTP label) we created, we now pretrain the student model to predict the PTP label when æ is provided as input. In other words, the student model is trained to predict the PTP labels given the downstream training dataset. We train the student model until convergence.
Once these two steps complete, use this PTP- TTP-pretrained student model the initial state for the KD process.
Motivation behind PTP: A relatively easy way of learning the teacher's generalized knowl- edge. The core idea is to make PTP labels by explicitly expressing important information from the teachers' softmax outputs, such whether the teacher model predicted correctly how confident the teacher model is. Pretraining using these labels would help the student acquire the teacher's 8 gen- eralized knowledge latent in the teacher's softmax output. This pretraining makes the student model better prepared for the actual KD process. For example, if a teacher makes an incorrect prediction for a data instance x, then we know that data instance x is generally a difficult one to predict. Since this knowledge is only obtainable by directly comparing the true label with the teachers output, it would be difficult for the student to acquire this in the conventional KD process. Representing this type of information through PTP labels and training the student to predict them could help the stu- dent acquire deeply latent knowledge included the teacher's output much easily, Intuitively, we expect that student model that has undergone such pretraining session is better prepared for the actual KD process and will likely achieve better results.
3.4
PEABERT: SPS AND PTP COMBINED
3.4.1 OVERALL FRAMEWORK OF PEABERT
PeaBERT applies SPS and PTP together on BERT for maximum impact on performance. Given student model, PeaBERT first transforms it into an SPS model and applies PTP. Once PTP completed, we use this model as the initial state of the student model for the KD process. The overall framework of PeaBERT is depicted in Figure 2
3.4.2 LEARNING DETAILS OF PEABERT
For the starting point of the KD process, well-finetuned teacher model should be used. We the 12 layer BERT-base model as the teacher. The learned parameters are denoted as:
0' arg min
Σ C iEN
0t
The 0 denotes parameters of the teacher, a denotes the softmax function, x, denotes the training data, 2 denotes the teacher model's output predictions, y: denotes the true labels, and Leg denotes cross-entropy loss.
We then pretrain the student model with PTP labels using the following loss:
arg min Σ LOE PTP (2, #*)))
(4)
PTP
5

---

Under review as a conference paper at ICLR 2021
True Label
Teacher
Predictions
Teacher
Predictions
Labels
Training
Training Data
Transformer Distillation
Data
PTP-trained
CE
True Label
SPS Student
Predictions
SPS
Predictions
Loss
Student
Figure 2: Overall framework of PeaBERT. The left half represents applying SPS and PTP to student model. The right half represents the learning baseline.
where u!TP denotes the PTP labels and the subscript s denotes the student model. When PTP is completed, we use the 0r the initial state of the KD process. Our KD process uses three loss terms total: pross-entropy loss between student model's final predictions and ground-truth labels, a Kullback- Lallba divergence loss between the final predictions of student and teacher models, and a mean-squared error loss between the intermediate layers student and teacher models.
The loss function is as follows:
= (1- (1-a) (y),
+ (oaa (r)(a) 0)))++
Σ (21
1 1|2-1-2 -2 2x
(5)
where LX. denotes the Kullback-Leibler divergence loss, K. denotes the index pairs of layers we use for intermediate layer-wise distillation, and ² denotes the output logits of the k-th layer. a and 3 are hyperparameters.
Note that during the KD process, we use a softmax- tempertter which controls the softness of teacher model's output predictions introduced in (Hinton et al. (2015)).
ea/T
0T(2))
(6)
Σ E
4 EXPERIMENTS
We discuss experimental results to assess the effectiveness of our proposed method. Our goal is to answer following questions.
Q1. Overall performance. How does PeaBERT perform compared to the currently exist- ing KD methods? (Section[4. 4.2)
Q2. Effectiveness of SPS. To what extent does SPS improve the effective capacity of the student model without increasing the number of parameters? (Section 4.3)
Q3. Effectiveness of PTP. Is PTP good initialization method? Compared the conventionally-usee truncated version of the BERT-base model, what is the impact of PTP on the model's performance? (Section
4.1 EXPERIMENTAL SETTINGS
Datasets. We use four of the most widely used datasets in the General Language Understanding Evaluation (GLUE) benchmark (Wang al, (2018)): SST-2] QNL RTH3] and MRPC For sen- timent classification, we use the Stanford Sentiment Treebank (SST-2) (Socher et al. (201)). For

---

Under review as a conference paper at ICLR 2021
Table 2: Overall results of PeaBERT compared to the state-of-the-art KD baseline, PatientKD. The results are evaluated on the test of GLUE official benchmark. The subscript numbers denote the number of independent layers of the student.
Method
RTE (Acc)
MRPC (F1)
SST-2 (Acc)
QNLI (Acc)
Avg
BERT-base
66.4
88.9
93.5
90.5
84.8
BERT,-Patientie
52.8
80.6
83.6
64.0
70.3
PeaBERT,
53.0
81.0
86.9
78.8
75.0
SERT.-PatienttD
53.5
80.4
87.0
80.1
75.2
PeaBERT
64.1
82.7
88.2
86.0
80.3
BERTg-PatienKKD PeaBERT3
58,4
81.9
88.4
85.0
78.4
64.5
85.0
90.4
87.0
81.7
Table 3: PeaBERT in comparison to other state-of-the-ard competitors dev set. The cited results of the competitors dimensions are fixed to six layers across all models compared.
from the official papers of each method. For accurate comparison, model
Method
# of parameters
RTE (Acc)
MRPC (FI)
SST-2 (Acc)
QNLI (Acc)
Avg
DistiBBERT
42.6M
59.9
87.5
91.3
89.2
82.0
TinyBERT
42.6M
70.4
90.6
93.0
91.1
86.3
PeaBERT 42.6M
73.6
92.9
93.5
90.3
87.6
natural language inference, we QNLI(Rajpurkar al. (2016)) and RTE. For paraphrase similar- ity matching, we use Microsoft Research Paraphrase Corpus (MRPC) (Dolan & Brockett (2005), ) Specifically, SST-2 is movie review dataset with binary annotations where the binary label indi- cates positive and negative reviews. QNLI is task for predicting whether pair of question and an answer is an entailment or not. RTE is based on a series of textual entailment challenges and MRPC contains pairs of sentences and corresponding labels, where the labels indicate the semantic equivalence relationship between the sentences in each pair.
Competitors. We use Patient Knowledge Distillation (PatientKD, Sun et al. (2099) our baseline learning method to compare and quantify the effectiveness PeaBERT. PatientKD is one of most widely used baselines and is variant of the original Knowledge Distillation method (Hin-) ton et al.] (2006)). We conduct experiments on BERT model (Devlin et al. (2018) and compare the results of PeaBERT to the original BERT. In addition, compare the results with other state- of-the-art BERT-distillation models, including (Ssh et al. (2019)) and TinyBERT(jia al. (2019)).
Training Details. We use the 12-layer original BERT model (Devlin et al. (2018) as a teacher model and further finetune the teacher for each task independently. The student models created using the same architecture as the original BERT, but the number of layers is reduced to either 1, 2, depending on experiments. That is, initialize the student model using the first n-layers of parameters from the pretrained original BERT obtained from Google's 's official BERT repe F We use the standard baseline PatientKD and the following hyperparameter settings: training batch size from {32, 64). learning rate from {1, 2, 3, 53-10-5. number of epochs from {4, 6, 10), a between {0.1 and 0.7}. and 3 between {0 and 500).
4.2 OVERALL PERFORMANCE
We summarize the performance of PeaBERT against the standard baseline PatientKD in Table We also compare the results of PeaBERT the competitors DistilBERT and TinyBERT in Table We observe the following from the results.
First, we see from Table 2 that PeaBERT consistently yields higher performance in downstream tasks across all three model sizes. Notably, PeaBERT shows improved accuracy by an average of 4.7% for the 1-layer student, 5.1% for the 2-layer student, and 3.3% for the 3-layer student (Note that we use F1 score for MRPC when calculating the average accuracy throughout the experiments).

---

Under review as a conference paper at ICLR 2021
A maximum improvement of 14.8% is achieved in MAABERTi-NNLY These results validate the effectiveness of PeaBERT across varying downstream tasks and student model sizes.
Second, using the same number of parameters, PeaBERT outperforms the state-of-the KD base- lines DistilBERT and TinyBERT, by 5.6% and 1.3% on average. We use a 6-layer student model for this comparison. An advantage of PeaBERT is that it achieves remarkable accuracy improve- ment just by using the downstream dataset without touching the original pretraining tasks. Unlike its competitors DistiBBERT and TinyBERT, PeaBERT does not touch the original pretraining tasks, Masked Language Modeling (MLM) and Next Sentence Prediction (NSP). This reduces training time significantly. For example, DistilBERT took upproximately 90 hours with eight 16GB V100 GPUs, while PeaBERT took a minimum of one minute (PeaBERT, with RTE) to maximum of one hour (PeaBERT, with QNLI) using just two NVIDIA T4 GPUs.
Finally, another advantage of PeaBERT is that it can be applied to other transformer- models with minimal modifications. The SPS method can be directly applied to any transformer-base models, and the PTP method can be applied to any classification task.
4.3 EFFECTIVENESS OF SPS
Table 4: An ablation study to validate each step of SPS. The results are derived using GLUE dev set.
Method
# of parameters
RTE (Acc)
MRPC (F1)
SST-2 (Acc)
QNLI (Acc)
Avg
BERT,
21.3M
61.4
84.3
89.4
84.8
80.0
SPS-1
21.3M
63.5
85.8
89.6
85.5
81.1
SPS-2
21.3M
68.6
86.8
90.2
86.5
83.0
We perform ablation studies verify the effectiveness of SPS. We compare three models BERT,, SPS-1, and SPS-2. BERT; is the original BERT model with 3 layers, which applies none of the SPS steps. SPS-1 applies only the first SPS step (paired parameter sharing) to BERT,, SPS-2 applies both the first step and the second step (shuffling) to BERT3-
The results are summarized in Table 4 Compared the original BERT SPS-1 shows improved accuracy in all the downstream datasets with an average of 1.1%, verifying our first motivation. Comparing SPS-1 with SPS-2, we note that SPS-2 consistently shows even better performance with an average of 1.9%, which validates our second motivation. Based on these results, we conclude that both steps of the SPS process work as intended to increase the student model's capacity without increasing the number of parameters used.
4.4 EFFECTIVENESS OF PTP
Table 5: An ablation study to verify PTP. The results are derived using GLUE dev set.
Model
RTE (Ace)
MRPC (F1)
SST-2 (Acc)
QNLI (Acc)
Avg
SSPS
68.6
86.8
90.2
86.5
83.0
BERT,+SPS+PTP-1
69.0
88.0
90.4
86.6
83.5
BERT,+SPS+PTP-2
69.3
88.3
90.9
86.8
83.8
BERT,+SPS+PTP
70.8
88.7
91.2
87.1
84.5
We perform an ablation study to validate the effectiveness of using PTP the student model. We use B which is SPS applied on BERT,, as our base model. We compare the results of PTP to its variants PTP-1 and PTP-2. Note that PTP uses four labels constructed from two types of information latent in the teacher model's softmax prediction: (1) whether the teacher predicted correctly and (2) how confident the teacher is. PTP-1 is a variant of PTP that uses only two labels that state whether the teacher predicts correctly or not. PTP-2 is another PTP variant that uses two labels to state whether the teacher's 's predictions confident or not. PTP-1 and PTP-2 contain different type of teacher's information in their labels, respectively. From the results summarized in Table 5) we see that PTP-1 and PTP-2 increase the student's 's accuracy by 0.5% and 0.8% on average, respectively. Combining PTP-1 1 and PTP-2, which is essentially our proposed PTP method, further improves model accuracy by an average of 1.5%. This study supports the efficacy of our PTP method.
initial guide for
8

---

Under review as a conference paper at ICLR 2021
These results prove two things: (1) The student that goes through only the conventional KD process does not fully utilize the knowledge included in the teacher models softmax outputs, and (2) PTP does help the student better utilize the knowledge included in teachers softmax outputs. This firmly supports the efficacy of our PTP method and also validates our main claim that initializing student model with ID-specialized method prior to applying KID can improve accuracy. As existing KI) methods do not place much emphasis on the initialization process, this finding highlights poten- tially major, undiscovered path improving model accuracy. Further and deeper researches related to KD KD-specialized initialization could be promising.
CONCLUSION
In this paper, we propose Pea-KD, new KD method for transformer- distillation, and show its efficacy. Our goal is to address and reduce the limitations of the currently available KD methods: insufficient model capacity and absence of proper initial guide for the student. We first introduce SPS, new parameter sharing approach that uses a shuffling mechanism, which enhances ca- pacity of the student model while using the same number of parameters. We then introduce PTP, KD-specific initialization method for the student model. Our proposed PeaBERT comes from ap- plying these methods SPS and PTP on BERT. Through extensive experiments conducted using multiple datasets and varying model sizes, we show that our method improves KD accuracy by an average of 4.4% the GLUE test set. We also show that PeaBERT works well across different datasets, and outperforms the original BERT as well as other state-of-the-art baselines on BERT distillation by average of 3.5%.
As a future work, we plan to delve deeper into the concept of KD-specialized initialization of the student model. Also, since PTP and SPS are independent processes on their own, we plan combine PTP and SPS with other model compression techniques, such as weight pruning and quantization.
REFERENCES
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800 04800, 2018.
William B Dolan and Chris Brockett Automatically constructing corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005), 2005.
Yunchao Gong, Liu Liu, Ming Yang, and Lubomir Bourdev. Compressing deep convolutional net- works using vector quantization. arXiv preprint 4rXiv:1412.6115. 2014.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. In Advances in neural information processing systems, pp. 1 1135-1143, 2015.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge neural network. arXiv 2015.
preprint
Geoffrey E. Hinton, Simon Osindero, and Yee Whye Teh. A fast learning algorithm for deep belief nets. Neural Computation, 18:552711 2006.
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 1909. 2019.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1692, 2019.
Matthew E Peters, Mark Neumann, Mohit lyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations, arXiv preprint arXiv: /80..003550. 2018.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint 1606vv1606.0222 2016.

---

Under review as a conference paper at ICLR 2021
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.
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. In Proceedings of the 2013 conference on empirical methods in natural language pro- cessing, pp. 1631-1642, 2013.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprint 1oXiv:1908.0990 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information 2017.
processing systems, pp.
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. arXiv preprint arXiv: 1804: 07460, 2018.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in neural information processing systems, pp. 5753-5763, 2019.
10


Reviews:
[{"review_id": "RbZxI_6CYC", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This paper proposes a parameter-efficient KD which consists of two main parts: Shuffled Parameter Sharing (SPS) and Pretraining with Teacher’s Predictions (PTP). This work explores some new framework like parameters-sharing and shuffling in KD and obtain promising results. \n\n\nStrengths:\n1. The paper is well-written and easy to follow.\n2. The experiment part explores the overall performance, the effects of SPS, and PTP separately, which is clear and makes readers easy to understand the effects of each part (SPS: step1+step2, PTP) in Pea-KD.  \n\nWeaknesses:\n1. When applying SPS, the number of layers in the student model is double than the normal case, even the parameters are the same as the counterpart, I guess the FLOPs will increase or be doubled than the original one. I think for a small student model, the number of parameters should not be the only metric but also consider the FLOPs. Why this paper doesn’t discuss this? This is my main concern.\n2. From Table4, we can find that the shuffling in SPS has great effects. step1+step2 would have much greater improvement than only applying step1, how about only conduct shuffling without parameters sharing (only step2 by shuffling the original parameters without step1)?  On the other hand, what is the potential reason why shuffling can enrich the model capacity?\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "JhWFWPoSzy", "reviewer": "AnonReviewer2", "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 proposes a distillation method for BERT. The work is based on two-fold main ideas. First, as the student model is usually smaller in the number of parameters, the model capacity is limited. The authors propose to stack the layers that share parameters to counter this limitation. Second, the authors argue that the initialization of the student model is crucial, so they propose an pre-training strategy for boosting the student's performance.\n\nPros:\nThe idea of learning good initializations for the student model is interesting.\n\n\nCons:\nThis paper has several writing problems which need be carefully addressed before its publication.\n\n- As all the experiments are conducted on GLUE and only BERT models are considered in the paper, the proposed method seems to be tailored for BERT. The authors should stress this point clearly and early in the paper (in the title, abstract or introduction). Otherwise, the authors should provide some experimental results on other models or tasks, e.g., the typical image recognition task in computer vision.\n\n- In the section of PROPOSED METHODs (section 3.1), the motivation and the main idea of the paper are introduced again. As no any new information is provided here compared to the abstraction and the introduction sections, it seems very redundant.\n\n- The introduction of SPS in Section 3.2 is quite confusing. What do Key  and Query parameters stand for?  I have no idea what the author is talking about here. Maybe it is because I have little background in NLP and BERT. However, even so the authors are still responsible for making the paper easy to follow for readers like me.\n\n\nAs for the motivation and the proposed method, my comments are as follows.\n\n- The motivation and the proposed method are somewhat problematic. Firstly, the authors argue that small student with few parameters are limited in capacity, so they propose to stack repeating layers to enlarge the model capacity. However, stacking repeating layer will introduce much more computation cost, which violates the goal of distillation.\n\n-  Secondly, the authors propose to pre-train the student model with the teacher predictions to initialize the student. However, it is odd to view this step as pre-training as it actually adopts the teacher predictions to train the student. It is actually doing distillation! The improvement  in performance may simply come from the more training epochs.\n\n======================\n\npost-rebuttal:\n\nI have read all the comments from other reviewers and replies from the authors. The revised version partially addresses my concerns so I raise the score from 4 to 5. However, my concerns about the motivation of the work still exist, so I am still slightly leaning to reject this paper.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "CLijjy3CItp", "reviewer": "AnonReviewer1", "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": "This paper proposed a framework for knowledge distillation with smaller number of parameters.The authors proposed a new parameter sharing method that allows a greater model complexity for the student model. Another contribution is that a KD-specialized initialization method named Pretraining with Teacher’s Predictions can improve the student's performance. The author combined these two methods to improve the performance of the student model on existing tasks, which has surpassed the existing knowledge distillation baseline.\n\nThe SPS method is a very natural idea of extending the student model. It expands the model complexity of the student model and improves the representation ability of the model without increasing the number of parameters. The idea of PTP is an excellent initializing method. At the same time, the teacher model's generalization of knowledge is given to students through initialization, and then fine-tuned.\n\nThe experimental part of this article is also quite sufficient. The ablation experiment illustrates the effectiveness of the two steps for the optimization of results. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "3TKjVowcbP8", "reviewer": "AnonReviewer4", "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 claims to present an improved form of knowledge distillation which tackles the following perceived weaknesses of existing kd systems:\n- student's expressive power is less than original teacher model, and\n- unclear how to initialize the student model weights in a principled way\n\nThe paper proposes two methods to improve kd:\n- using stacked layers, with weight sharing between the weights, and keys and queries swapped in the upper layers\n   - this is proposed because it means the effective student geometry matches that of the teacher, or at least is significantly larger than without the shared layers, whilst the number of unique parameters remains small\n- getting the student to predict the confidence of the teacher on each example, rather than just the softmax output of the teacher\n   - the confidence is classified as either 'strong' or 'weak'\n   - the student must also predict whether the teacher gets the example correct or not (which we know, since this is for training set examples, for which we have ground truth)\n\nStrong points of the paper:\n- SPS sounds novel, but I felt it was only very weakly motivated\n    - in addition it sounds to me like the inference time will be similar to the original model? I feel that reducing the number of parameters in the student is not the only goal: the student model should be fast to execute.\n    - the paper does admittedly not claim to provide fast inference times as a goal; nevertheless I have a hard time imagining a BERT model running on an edge device with a limited amount of resources in a reasonable time, so I think that inference time is a critical metric which I feel this paper does not consider, or measure\n- PTP sounds to me only weakly novel, since typically the student will be trained to predict the full softmax output of the teacher, whereas in PTP, the student must predict whether the highest value of the softmax is high or low. I feel that little insight or motivation is given as to why this auxiliary task was used, or is better than predicting the full softmax distribution. Nor was a rigorous comparison carried out to compare using just this pre-training approach with just the softmax-matching approach, where both using appropriately tuned hyper-parameters, I feel. There was an experiment in 4.4 where the PTP task is dropped, but I got the impression that this was after tuning the hyper-parameters to rely on PTP task? The hyper-parameters would I feel need to be retuned in the absence of the PTP task I think?\n\nI mostly like the paper. I found it interesting. However, I have two main concerns:\n- the approaches used in the paper feel to me only weakly motivated. Little insight is given into why the approaches were chosen, and why they should work. For example, flipping the keys and queries was not well justified I felt. Nor was the PTP task relative to traditional softmax matching.\n- the paper avoids discussing inference times, and it's unclear to me whether a distillation approach that results in few parameters but long inference times would be useful in practice?\n    - of course, the underlying techniques can still be interesting from a theoretical point of view\n\nWhilst I was reading through the paper, I wrote down some more detailed reactions:\n\nIntroduction:\n\n\"stacking the layers that share parameters\" => \"stacking layers that share parameters\" (the former form of language is using 'that share parameters' as a filter to choose from existing layers; the latter form implies that we create new layers, which we then stack, and which we configure to have shared parameters)\n\n3.1 Overview\n\nSo, it looks like SPS means the number of unique parameters is kept low whilst the effective model complexity is unchanged compared to original BERT. How does this affect inference speed? It sounds to me like the inference speed under SPS will be unchanged from a BERT model of equivalent model complexity? Often, a use-case of distillation is to reduce inference time, and requirements in terms of power, eg for mobile devices, or production inference. Assuming that the inference time is unchanged from the original model of equal complexity, what is the target use-case of SPS/Pea-KD?\n\n3.2 SPS\n\nNo insight or intuition is given as to why swapping query and key is likely to be a useful thing to do. I think it would be good to provide such insight and intuition. Like, you say that it increases the expressive power, but you don't justify why you feel this is true.\n\nFor the six-layer student, I couldnt quite undersatnd why it looks like the 6-layer student has 9 layers? Please consider clarifying this point:\n- why do you call it '6 layer' if it has 9 effective layers? Perhaps consider renaming it to 9-layer student?\n- why do you make the 6 layer model become 9 layers, and not 12 layers? ie, following the initial recipe detailed in section 3.2, of doubling all the layers.\n\n3.3 PTP\n\nI feel that the column for 'PTP label' in table 1 should be on the right hand side:\n- the inputs should be I feel on the left, ie 'teachers prediction correct', and 'confidence > t',\n- and then we can read off the 'output' of the table as the right-hand column\n- (currently I find the table hard to read)\n\nThe pretraining task itself seems interesting to me.\n\nHow do you get the student to predict this result? Do you put a linear layer on the output from the student? Why do you use a classification task (4 classes) instead of using a regression task to predict the confidence? did you try a regression task to predict the confidence? Why use a classification into 4 classes, rather than two classifiers into two classes each?\n\n4. Experiments:\n\nFor table 2, I wasn't really sure how these numbers compare to two key baselines:\n- a simple bilstm with attention, without any pretraining, and\n- a BERT-base model\n\nSo, I dug these out, and here are the peabert numbers from table 2, put into the context of these two baselines (which I feel are lower and upperbound really for what we'd expect to see from PeaBERT):\n\n     BiLSTM+Attn\n     (single-task training)   BERT-base   PeaBERT1  PeaBERT2  PeaBERT3\n    RTE: 51.9                     66.4        53.0     64.1       64.5\n    MRPC 68.5                     88.9        81.0     82.7        85.0\n    SST2 85.9                     93.5        86.9     88.2       90.4\n    QNLI 77.2                     90.5        78.8      86.0       87.0\n\nI got the BiLSTM+Attn numbers from GLUE paper, and\nBERT numbers from https://arxiv.org/pdf/1910.03176.pdf\n\nI kind of feel that I might not be the only person who might want to see the PeaBERT numbers in the context of such lower and upperbound baselines? Maybe consider adding these baselinse into table 2?\n\nBasically, my take away from this is that PeaBERT1 is barely better than the simple BiLSTM+Attn baseline, but PeaBERT3 approaches the score of a full BERT-base model?\n\nI kind of think the sentence \"For example, DistilBERT took approximately 90 hours with eight 16GB V100 GPUs while PeaBERT took a minimum of one minute (PeaBERT1 with RTE) to a maximum of one hour (PeaBERT3 with QNLI) using just two NVIDIA T4 GPUs\" should be highlighted in a table somewhere somehow perhaps, rather than buried in text? Not sure if that's a good idea, just occurred to me though.\n\nQuestion: why only show results on a subset of the GLUE tasks? eg patient-kd paper shows results on additionally:\n- QQP\n- MNLI-m\n- MNLI-mm\n(whilst also showing results for: SST-2, MRPC, QNLI and RTE, as here)\n\nIt is unclear from this whether you ran against all, and only showed the four tasks that show a benefit, or whether you simply didn't have time to run on all 7 tasks. Preference to show results for all 7 tasks that Patient-KD paper reports results for.\n\nFor ablation studies, I feel it would be interesting to see an ablation study that removes each of various losses in equation 3 in turn.\n\nImportantly, none of the experiments mention inference time, which I feel is a key metric to report for distillation?\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper got mixed reviews. One for acceptance, one for reject and two borderline. After the rebuttal, AR2 raises the review to borderline.  AR1 gives the highest recommendation but did not provide detailed supporting evidence. Other reviews provide comment on the strength and also share the concerns. Most of the concerns concentrate on the motivation (whether the proposed method is violating the objective of knowledge distillation) and the brought additional computation overhead. Also the scope of this paper was defined wider than the actual one. The authors only did experiments for BERT but did not consider and compare with existing KD method. Overall, AC read the paper and also has the similar concerns, the novelty is limited. the brought increase in inference time is violating the KD objective and the scope of this paper was not defined clearly. The authors should improve the submission in these aspects. At its current status, AC does not recommend acceptance. 
Decision: Reject