Learn From One Specialized Sub-Teacher: One-to-One Mapping for Feature-Based Knowledge Distillation
Authors: Khouloud Saadi
Jelena Mitrović
Michael Granitzer
Venue: EMNLP 2023 Findings (2023)
Keywords: Knowledge Distillation
Compression
NLP
Large Language Models
Feature Distillation
TL;DR: We propose to break down the global feature distillation task into N local sub-tasks.

Abstract: Knowledge distillation is known as an effective technique for compressing over-parameterized language models. In this work, we propose to break down the global feature distillation task into N local sub-tasks. In this new framework, we consider each neuron in the last hidden layer of the teacher network as a specialized sub-teacher. We also consider each neuron in the last hidden layer of the student network as a focused sub-student. We make each focused sub-student learn from one corresponding specialized sub-teacher and ignore the others. This will facilitate the task for the sub-student and keep it focused. Our proposed method is novel and can be combined with other distillation techniques. Empirical results show that our proposed approach outperforms the state-of-the-art methods by maintaining higher performance on most benchmark datasets. Furthermore, we propose a randomized variant of our approach, called Masked One-to-One Mapping.
Rather than learning all the N sub-tasks simultaneously, we focus on learning a subset of these sub-tasks at each optimization step. This variant enables the student to digest the received flow of knowledge more effectively and yields superior results.

Learn From One Specialized Sub-Teacher: One-to-One Mapping for Feature-Based Knowledge Distillation
Khouloud Saadi, Jelena Mitrović, Michael Granitzer
University of Passau, Germany
[Khouloud. Velena.Mitrovis Michael.Gramiteee perrantzeereeui-paasaat
Abstract
al., 2019) and GPT-3 with billions of parameters (Brown et al., 2020), they are not highly adapted to real-world applications like mobile computing. Model compression an active field of research that focuses on effectively reducing the model size without significant performance degradation (Xu and McAuley, 2022; Frankle and Carbin, 2018; Jiao et al., 2020).
Knowledge distillation is known effective technique for compressing over wer-parameterized language models. In this work, we propose to break down the global feature distillation task into N local sub-tasks. In this new framework, we consider each neuron in the last hidden layer specialized sub-
of the teacher network teacher. We also consider each neuron in the last hidden layer of the student network focused sub-student. We make each focused sub-student learn from one corresponding spe- cialized sub-teacher and ignore the others. This will facilitate the task for the sub-student and keep focused. Our proposed method is novel and can be combined with other distillation techniques. Empirical results show that our pro- posed approach outperforms the state-of-the-n methods by maintaining higher performance on most benchmark datasets. Furthermore, we propose randomized variant of our approach, called Masked One-to-One Mapping. Rather than learning the sub-tasks simultane- ously, we focus on learning subset of these sub-tasks at each optimization step. This vari- ant enables the student to digest the received flow of knowledge more effectively and yields superior results.
Knowledge Distillation (KD) by (Hinton et al., 2015) is one of the effective compression tech- niques in NLP where the knowledge of highly capable large model, i.e., teacher, is transferred a smaller efficient model, i.c., student. KD essen- tially requires designing a loss function to mini- mize the distance of the output or the intermediate representations between the student and the teacher (Sanh et al., 2019; Jiao et al., 2020; Hinton et al., 2015). To distill the intermediate representations, previous research relied on the mean square error (MSE) as an objective function between the stu- dent and the teacher global representations (Sun et al., 2019; Jiao et al., 2020). However, this met- ric is sensitive to scale (Saadi and Taimoor Khan, 2022) and it is not accurate in high dimensional space (Aggarwal et al., 2001; Houle et al., 2010). Other works used the cosine distance tive, but, it also has several limitations (Zhou al., 2022: Schütze al., 2008) such as not performing well with sparse data.
alterna-
1 Introduction
Large language models, also known general purpose language models, have revolutionized the Brown et al., 2020; Radford et al., 2019; Clark et al., 2020). They are large architectures composed of several transformer blocks (Vaswani et al., 2017), typi- cally trained on large unlabeled corpora in self- supervised way (Devlin et al., 2019; Brown et al., 2020). They achieved state-of-the-art (SOTA) performance on downstream tasks through fine- tuning when the data scarce (Devlin al., 2019; Brown et al., 2020). However, since these models are typically large and computationally expensive, such as BERT with millions of parameters (Devlin
NLP domain (Devlin et al., 2019;
In this work, we focus on distilling the teacher's last hidden layer representation into the student's 's last hidden layer representation. To achieve that, we propose novel KD approach where we refor- mulate this feature distillation task from global problem to N local sub- nub-problems. Here, the term "local" refers to the consideration of each dimen- sion or neuron individually. In this work. our main assumption is that the student last hidden layer has the same dimension as the teacher last hidden layer.
We note that this is a common assumption in fea- ture distillation (Sanh et al., 2019).

---

In our proposed framework, we consider each neuron in the teacher's last hidden layer as spe- cialized sub-teacher and each neuron in the stu- dent's last hidden layer as a focused sub-student. Each specialized sub-teacher is charge of dis- tilling its knowledge to the corresponding focused sub-student. We call this one-to-one matching between the teacher and the student last hidden layers units. To accomplish this one-to-one distilla- tion, we propose a novel objective function where we maximize the per-batch correlation between the outputs of each specialized sub-teacher and its cor- responding focused sub-student. Empirical results show that studying the global feature distillation task from a local viewpoint helped the student to meet the global teacher's features representation. In addition, we propose an augmented variant of our approach where the model learns subset of sub-tasks at time. To sum up, our contributions are the following:
to perform computations (Lan et al., 2019; Reid et al., 2021). Quantization (Zhou al., 2017) is yet another strategy used in model compression, where the weights and activations in the network are represented using lower bit integers instead of higher precision floating- numbers (Kim et al., 2021; Prato et al., 2020).
Knowledge distillation (Hinton et al., 2015; Zhang et al., 2020; Sanh al., 2019), which is the focus of this paper, also a successful com- pression technique. It involves training a smaller, more efficient student model to mimic the behavior of a larger, knowledgeable teacher model. By lever- aging the knowledge of the teacher model to the student model, the student can achieve compara- ble performance to the teacher while maintaining lower size. Knowledge distillation has been proven to be an effective technique for large language model compression (Sun et al., 2019; Sanh et al., 2019). It can be applied during the pre-training stage to generate general-purposs distilled models (Jiao et al., 2020; Sanh et al., 2019) and during the fine-tuning stage to generate task-specific distilled models (Zhou et al., 2021; Liang et al., 2020; Sun et al.. 2019).
We reformulate the global feature distillation task into V local sub-tasks where we do a one- to-one mapping between the last hidden layers units of the teacher and the student models.
We propose a local correlation- objective function to achieve the One-to-One Mapping distillation.
In (Kovaleva et al., 2019), the authors show that large language models, e.g., BERT, suffer from over-parametrization in domain-specific tasks. Thus, previous work has been improving the task- specific distillation. Several methods focused on enhancing the objectives of the distillation process. These improvements mainly focused on which part of the teacher architecture can be distilled into the student architecture such as the attention matrices (Jiao et al., 2020), the different hidden states (Sun et al., 2019), and the prediction layer (Hinton et al., 2015).
We conduct extensive experiments: 1. Stand- alone experiments on the SQUAD V1 and the IMDB dataset. 2. Comparison with the com- peting methods on 8 GLUE datasets (Wang et al., 2018). Our approach performs the best in most cases.
We present an augmented variant of our pro- posed approach called: Masked One-to-On tone Mapping. Empirical results show the effec- tiveness of this new variant.
Coming up with effective objective functions distill the knowledge from any part of the teacher into the student is critical. For the logit-based KD, the KL divergence or the MSE are used as objective functions to minimize the distance between the logits of the student and the logits of the teacher (Hinton et al., 2015; Zhou et al., 2021). Following that, in (Zhao et al., 2022), the authors provided novel viewpoint to study the logit distillation by reformulating the classical KL divergence loss into two parts, which showed good improvement. In the feature-based KD, the MSE and the cosine distance are mainly used as objective functions to align the intermediate layers representations of the
2 Related Work
Model compression techniques are employed to effectively reduce the size of a neural network while maintaining good performance. Various approaches have been proposed such as pruning (Frankle and Carbin, 2018). It aims to remove unimportant structures, e.g., weights, neurons, and even entire layers from the model (Lagunas al., 2021; Prasanna et al., 2020). Another technique is weight sharing (Long et al., 2017), where differ- ent parts of the model use the same set of weights

---

teacher and the student (Sun et al., 2019; Jiao et al., 2020; Sanh et al., 2019).
to meet the global representation provided by the teacher. We note that the initial results of our work on the one-on-one mapping approach are presented in (Saadi et al., 2023).
In this work, we argue that relying on the MSE and the cosine distance as loss functions to align the intermediate representations between the student and the teacher is not an optimal choice. MSE is sensitive to scale and does not perform well in high- dimensional space. In a high dimensional space, which is the case in a neural network, the data tend be sparse and all the data points become uni- formly distant from each other (Aggarwal et al., 2001; Houle et al., 2010). In fact, the embeddings (features representation) provided by BERT-like models are high-dimensional and sparse tensors in nature (Li et al., 2022). Thus, using MSE or cosine distance to measure the distance between these sparse representations is not optimal as plained and advocated in (Aggarwal et al., 2001). So, splitting the task into sub-tasks and using one- to-one metric alleviates this problem and breaks the curse of dimensionality. In (Zhou et al., 2022), it is also shown that the cosine distance is not an accurate measurement of similarity between BERT embeddings. Moreover, the cosine distance mea- sure can be also affected by sparsity. In fact, in high-dimensional space, it can output large angles between two sparse vectors although they are sim- ilar in the non-zero components (Schütze et al., 2008). Another important limitation of the cosine distance is that it is a global metric and sensitive to unit permutation. For example, let W, be the output tensor of a given layer in the teacher and W2 be the output of the corresponding layer in the student. In the following example, we compute the cosine distance (CosD) from W1 and W2:
3 Our Approach: One-to-One Mapping
In this work, we provide a new viewpoint on how to study the feature distillation task. We break down the global KD task into multiple local KD sub-tasks. Specifically, in the last hidden layer of the teacher, we consider each unit specialized sub-teacher. In the last hidden layer of the student, we also treat each unit as focused sub-student, Each sub- student should concentrate and learn only from the corresponding specialized sub-teacher. We refer to this as One-to-One Mapping, which is reformu- lated as the cross tross-correlation between the outputs of each specialized sub-teacher and its correspond- ing focused sub-student. The mapping between each sub-student and each sub-teacher is fixed, i.e. based on an index, and thus remains consistent over batches and epochs. Although our approach can be applied to different cuudntt-tcacher hidden lay- ers, as in (Sun et al.. 2019), this work specifically focuses on the last hidden layer.
Typically, in KD framework, illustrated in Figure 1, we have the teacher network modeled by fe. which is an over-parameterized knowledgeable model. The student network is modeled by S which is an efficient model that has lower number of parameters compared to the teacher. The input batch X is fed to fo and f/ simultaneously to produce the batches of features representation Y, and Y,, respectively.
We assume that Y, and Y, are the features rep- resentation of the last hidden layer of the teacher hfast and the last hidden layer of the student haatt respectively. Y, and Y, are assumed to be mean- centered over the batch dimension. We assume that h/ and b.ast have N hidden units. The unit; in h.asst represents the specialized sub-teacher; and the unit; in h.ast represents the focused sub- student, The sub-task; is distilling the Knowledge from the specialized sub-teacher; to the focused sub-student; by reducing the distance between the features learned by each of them.
W1= [1,5]
W₂ = [5, 1]
W1 W2 0.6154
CosD(W1, = [Will ||22ll2
It is a high value but the layers actually learned the same representation.
In our work, instead of treating feature distilla- tion as global task (Sanh et ul., 2019; Sun et al., 2019) between teacher-student layers, we reformu- late it multiple local sub-tasks. Furthermore, for each per-dimension KD sub-task, instead of using MSE or cosine distance (Sanh et al., 2019; Jiao et al., 2020; Sun et al., 2019), we propose a noval loss function that utilizes a batch correla- tion function. This improves the student's ability
To simplify the task for the sub-student; and keep it focused, we force it to learn only from the sub-teacher, and ignore the other sub-teachers. We reformulate the objective function in this one-to- one mapping as maximizing the cross-corr

---

Y, 2
Student
4 (1 a(
Σ xf)
X
371
LKD
Σ 4
Lv C(vim
Σ Clusr 1/3)
YY Z
36N
Figure 1: One-to-One Mapping: X is the input batch. n is the batch size. T is given sample. C is the cross correlation function. The Student and the Teacher are modeled by for and fe. respectively. and Y, are the features representation of the last hidden layer of the student and the teacher, respectively.
between the two variables ył and yf- y and y are the output values the sub-teacher, and the sub- student, respectively. The variables y and y have n samples coming from the different examples in the input batch. Maximizing the correlation across batches between the two aforementioned variables, i.e., minimizing the following loss function:
sub-student; and the output of the sub-teacher, The second term is for minimizing the cross-correlation between the sub-student; and each sub-teacher, given j / i and j E {1, 2, NV.. Our end distilla- tion loss is the sum of the N local KD losses.
N
N
N ΣΣ
N
4 (1- 2
(1)
LKD Σ L,
Σ
C.)2
+
where Cf is the cross-correlation value between the variables yf and ył:
(5)
Additionally, we introduce X and X2 as the weights to control the contribution of the first term and the second term of the loss function, respec- tively:
Eb Vh. b=1 (wb )2 Eb (
Ca
N Σ
b is the index of given sample in the input batch X. n is the number of examples in the input batch X. i, index the output dimension of the last hidden layer in both the teacher and the student. In fact, i is the index of the den element in the output and it is also the index of the ith neuron in the last hidden layer, i.e., sub-teacher, sub-student,
LKD λ (1- A2
(6)
The final training loss of the original student is:
L = aLCE + 3LKD
(7)
As we want make the task easier for the sub- student, so it can effectively digest the received information, we force it to mimic only the teacher; and repel all knowledge coming from the other teachers. This results in minimizing the following term:
Where LeE is the classical cross entropy loss be- tween the student predictions and the ground truth labels:
1 n Σ y log (2x)
LEE
(8)
R₁
Σ jæ
(3) Where x; is an input sample and yi is its ground truth label. n is the number of samples per batch.
Thus, our final per-dimension KD loss function, as shown in Figure 1, is:
Empirical results will show that our One-to One-to-One Mapping approach effectively facilitates the align- ment of the student features representation with the global teacher feature representation. Our ap- proach can be implemented a low computational cost and can be combined with other knowledge distillation methods.
I 1, R C1)
Ca)2
Σ C
Li, the first term is for maximizing the cross- correlation over batches between the output of the

---

4 Experimental Results
tive function. Distilbert,MMSE stands for feature distillation with MSE loss.
In our experiments, the teacher is the BERT-base model, with 10 million parameters, after being fine-tuned on each of the datasets for 3 epochs. The student is the DistilBERT parameters. N, which is the number of neurons, in the last hidden layer of the teacher and the student, is equal to 768. All experiments are repeated for random seeds, the learning-rate rate set to 5e-5, the maximum sequence length is set to 128, and the batch size is set to 16.
As shown in Table I, our proposed method achieves the best results on the feature distillation task. It could outperform MSE and cosine distance on the squad and the imdb datasets. The results show that our approach performs the best when the distillation task is run for 3 epochs. It achieves an Exact Match (EM) score of 78.79% and an Fl score of 86.95% on the squad dataset. It also gives accuracy of 93.87%
with 66 million
the imdb dataset. In Table 2. For 10 epochs, while the performance of the other approaches decreased, ours effectively in- creased. It achieves 79.46% and 87.48% as EM and F1, respectively, on the squad dataset and 93.96% as accuracy on the imdb dataset. This proves that our approach yields better representation of the stu- dent features. Another noteworthy aspect is that the standard deviation of the results of our approach is low compared to the others. This indicates the stability and consistency of our proposed method.
4.1
Stand-Alone Experiments
4.1.1
Stand-Alone Results
In this stand-alone performance evaluation, we ex- periment with the Stanford Question Answering Dataset (SQuAD-VI) (Rajpurkar et al., 2016) and the Internet Movie Database dataset (IMDB) (Maas et al., 1). The SQUAD-VI V is a reading com- prehension dataset, consisting of questions posed by crowd workers set of Wikipedia articles, where the answer every question is segment of text, span, from the corresponding reading passage. The reported metrics for SQUAD are the Exact Match (EM) (Rajpurkar et al., 2016), which measures the percentage of predictions that match any one of the ground truth answers exactly, and the Fl-score (Rajpurkar et al., 2016), which mea- sures the average overlap between the prediction and ground truth answer. The IMDB dataset is sentiment analysis dataset consisting of 50,000 movie reviews labeled as positive or negative.
4.1.2 Sensitivity Analysis
In this subsection, we investigate the impact of the weight value assigned to each stand-alone loss the performance of the student model. As re- minder, in this stand-alone evaluation, the overall loss function for the student model computed as the weighted sum of the hard loss and each stand-alone KD loss. The KD stand-alone loss is computed between the last hidden layer of the student and the last hidden layer of the teacher. As KD stand-alone regularizers, we will compare the MSE, the cosine distance, and newly formu- lated loss function. We maintain fixed weight for the hard loss, i.e., 0.5, while varying the weight assigned to each stand-alone loss term. Here, the number of epochs is set to 3.
We distill the last hidden layer representation of the teacher into the last hidden layer of the student. We add our designed KD loss, the MSE as in (Sun et al., 2019), and the cosine distance as in (Sanh et al., 2019), stand-alone regularizers to the hard loss between the student predictions and the ground truth labels. This will show the effectiveness of our proposed objective function for the feature distil- lation task. We experiment with 3 and 10 epochs. The weight of each KD stand-alone loss and the weight of the hard loss are fixed to 0.5 (Sanh et al., 2019 Jiao et al., 2020).
As illustrated in Figure 2, on the IMDB dataset. the accuracy of the student model utilizing our stand-alone loss consistently surpasses those em- ploying MSE and cosine distance for all weight values. Additionally, as shown in Figure 3, on the squad dataset, utilizing our novel loss consistently yields the highest performance in terms of both the Exact Match (EM) and Fl metrics, across all hyper-parameter values, when compared to MSE and cosine distance. This highlights the robustness and effectiveness of our novel approach.
In Table 1, BERT and Distilbert, refer to the BERT-base pre-trained language model with 12 transformer blocks and to the DistilBERT-base model with 6 transformer blocks. respectively. Distilbertg-FT stands for fine-tuning the student without any distillation. Distilbertg-coss stands for feature distillation with cosine distance objec-

---

Approach
SQUAD-VI
IMDB ( (%)
BERT (teacher)
80.36/88.13
94.06
Distilbertg-IFT
77.43=0022/8567:10.111
93.1940.09
Distilbeerg-cosD
77.82=0029/85.2222222.
93.66++0.08
Distilbertg-MSE
77.72=
93.49:0.08
Distilbert,-ours
78.79-0.1276695-0000 93.87=0.01
Table 1: Stand-Alone regularizers. SQUAD-VI: The evaluation is reported Exact Match (EM) and Fl on the dev set. IMDB: The evaluation is reported as accuracy on the test set. Results are the average and the standard deviation 5 random seeds. Training for 3 epochs.
Approach
SQUAD-V1
IMDB (%)
BERT (teacher)
80.36/88.13
94.06
Distilberte-IFT
74.32:0.37/83.91100.11
92.71=0.08
Distilbert-cooDD
75.78=0.27784.5720.11
93.90=0.08
Distilbertg-MSE
75.17:10.18084:55H0.11 93.82:+0.07
Dist
79.46-00008877488-0006 93.96=0.01
Table 2: Stand-Alone regularizers. SQUAD-VL V1: The evaluation reported EM and F1 the dev set. IMDB: The evaluation reported as accuracy on the test set. Results are the average and the standard deviation of 5 random seeds. Training for 10 epochs.
Our
07.2
Oure
71.00
MSZ
C210
CoB
50.0
7830
00.4
75.00
20.5
25.6
Outs
37.50
37.25
00.0
082 0.2
0.5 010 0..0
0.7
0.e
Weipht stanch las
Woight of the standallcre les
Figure Variation of the model performance on the IMDB dataset in function of the weight of each stand- alone loss in the final loss. For each weight value, the experiments repeated for 5 seeds and the average is reported.
Figure 3: Variation of the model performance on the SQUAD-VI dataset in function of the weight of each stand-alone loss the final loss. For each weight value, the experiments repeated for 5 seeds and the average is reported.
4.2 Comparison With Competing Methods
is set to 5. For the rest of the baselines, similar to (Sanh et al., 2019; Zhou et al., 2021), the number of epochs is set to 3. For the vanilla-KI KD, the tempera- ture is set to 2 (Hinton et al., 2015). All weights in loss functions are fixed KD loss weight chosen from (0.4, 0.5, 0.6] (Sun et al., 2019) and 0.4 is the best match. In our ap- proach, A₁. X2, a, and / are set to 1, 5.1 0.5, and 5.10-3 respectively. After performing a man- ual tuning, we found out that the aforementioned hyper-parameter settings led to stable results. The number of epochs is set to 3.
this evaluation part, we compare the perfor- mance of our proposed approach with the compet- ing methods on the commonly used GLUE bench- mark dataset (Wang et al., 2018) for knowledge distillation in NLP. For comparison, we fine-tune the student on the tasks without distillation. We report the stand-alone results for MSE and cosine distance. We also generate the results of vanilla KD (Hinton al., 2015) and PKD (Sun et al., 2019). For the PKD approach, we use similar setting to (Sun al., 2019), where the number of epochs is set to 4, the distillation loss ratio is set to 100, the hard loss ratio is set to 0.5, and the temperature
except for the vanilla-
The GLUE benchmark dataset is composed of several sub-datasets for different tasks. QQP,

---

Approach
MRPC
RTE
CoLA
SST-2
STS-B
MNLI-m
QNLI
QQP
AVG
(
87.86
66.79
54.84
90.02
89.30
82.85
90.70
88.04
81.30
Distilbert,-FFT
84.48
56.17
44.28
89.54
85.40
80.41
86.71
87.79
76.85
Distilberte-Cool
85.38
63.75
46.82
89.79
85.59
80.51
88.62
87.78
78.53
Distilhert,,MMSR 86.04
62.38
47.35
89.79
85.42
80.12
88.13
87.90
78.39
Distilbertg-PKD
86.06
62.24
47.28
90.05
85.77
81.62
87.16
88.33
78.56
Distilbert,- 84.56
56.53
45.86
89.63
85.50
80.51
87.71
87.60
77.23
Distilbert,-ourrs 86.57
63.83
50.73
90.44
85.66
82.76
89.54
88.04
79.70
Table 3: Results on the GLUE dataset: Evaluation reported on the dev set as the average of random seeds. Training for 3 epochs. AVG is the average performance across the GLUE datasets. All values are in (%).
Approach
MRPC
RTE
CoLA
SST-2
STS-B
MNLI-m
QNLI QQP
AVG
|
87.86
66.79
54.84
90.02
89.30
82.85
90.70
88.04
81.30
Ours
86.03
64.19
53.21
89.91
85.88
82.29
89.09
88.28
79.84
Ours- Murs-Masked(80%)
85.79
64.55
53.16
90.18
85.94
82.41
89.44
88.32
79.97
Murs-Masked((70))
86.18
64.84
52.71
89.63
85.97
82.28
89.37
88.32
79.91
(Bes)))
86.45
65.05
53.16
90.18
85.97
82.42
89.44
88.37
80.13
Table Results on the GLUE dataset: Comparison between Ours and ours-Masked. Evaluation reported on the dev set the average of random seeds. AVG is the average performance across the 8 GLUE datasets, All values are in (%). Between parenthesises is the percentage (%) of the learned sub-tasks in each iteration. Training for epochs.
MRPC, and STS-B are for the paraphrase simi- larity matching task. SST-2 is for the sentiment classification task. MNLI-m (matched version of MNLI), QNLI, and RTE are for the natural lan- guage inference task. CoLA for the linguistics acceptability task. For MRPC and QQP we report the combined score from Fl and accuracy. For STS-B B we report the combined score from Pearson correlation and Spearman correlation. For CoLA we report the Matthew's 's correlation. For the rest of the tasks, accuracy is the metric.
gap between the teacher average performance, i.e., 81.30% and the student average performance, i.e., 79.70% on all the GLUE tasks.
5 Ablation Study
See results in Table 5 and as shown in the table both terms of the loss are important and their com- bination achieves the best results across most of the datasets. Thus, the two components of the loss have compound effect. LCE is the hard loss. is the first term in the LØD loss. R is the second term in the LKD loss (the repel term).
As shown Table 3, our approach outperforms KD, which stands for vanilla-KO KD (Hinton et al., 2015), and the PKD (Sun et al., 2019) baselines on most of the GLUE tasks. It is worth mentioning that for the SST-2 and QQP datasets, the resulting student model using our approach and the one using the PKD approach could exceed the performance of the teacher. This commonly observed in (Stanton et al., 2021; Furlanello et al., 2018). Furthermore, it consistently outperforms the MSE and the co- sine distance distillation objective functions, which are applied between the last hidden layer of the teacher and the last hidden layer of the student, on all the GLUE tasks. As shown in the last column of Table 3. all the GLUE tasks, our One-to-One Mapping approach achieves the highest average performance, i.e., 79.70%, with up 1.14% com- pared to the results the other approaches. Our proposed approach helped to effectively reduce the
Analayis
6.1 Layer Distillation: Our loss as an alternative for MSE and CosD
Rather than looking at our approach as a competing method to the KD state-of-the-an approaches, proposed to replace any MSE or cosine distance-based feature distillation methods between two hidden layers included in any knowledge distillation approach. To show that our approach can outperform MSE and cosine distance, applied globally, in the existing pipelines, experiment with SRRL (Yang al., 2020), one of the latest KD approaches. SRRL KD loss has two components: the feature loss, which is an MSE between the last layer's 's features of the teacher and the last layer's features of the student,

---

Table 5: Ablation Study (Ours).
All values in (%).
All the results are the average over 5 trials
approach
MRPC
RTE
CoLA
SST-2
STS-B
QNLI
MNLI-m
QQP
AVG
LOE +1
86.85
63.61
48.24
90.87
85.47
89.63
81.85
87.60
79.26
LOE + R
75.85
53.79
34.71
85.25
85.25
85.66
74.43
87.69
68.49
LCE + +R+1
86.57
63.83
50.73
90.44
85.66
89.54
82.76
88.04
79.70
and the logit loss. In this experiment, we replace the feature KD loss in SRRL with newly introduced correlation loss function. As shown in the Table 6, our feature distillation approach replacing MSE loss yields superior performance for SRRL on the two GLUE datasets.
enhance the student performance, we propose an augmented variant of our approach called Masked One-to-One Mapping. Instead of learning all N sub-tasks simultaneously, it is more effective for the student to focus on a subset of the total sub- tasks at a time. This approach enables the student to effectively absorb and manage the incoming flow of knowledge from the teacher.
Table 6: SRRL VS SRRL-OurS
approach
QNLI(%)
RTE(%)
As shown in Figure 4, during each iteration (op- timization step), we randomly mask subset of units in the last hidden layer of the teacher and in the last hidden layer of the student. Afterwards, we will have m (where m < N) neurons, i.e., m feature distillation sub-tasks left for the student to learn. We accomplish this by applying a same binary mask M to the last hidden layer features representation of both the student and the teacher simultancously.
SRRL
88. 88.04
63.68
SRRL-Ours
88.68
64. 10
6.2 Model Efficiency
Although our distillation objective is applied stand-alone between the last hidden layer of the student and the last hidden layer of the teacher, could outperform PKD. Note that PKD (Sun et al., 2019) distills different hidden-state features plus the logits. However, our approach distills only the last hidden feature of the teacher to the student. Our approach has lower computational, as shown in Table 7, our approach is trained in less time. In this experiment, we report the training time of PKD and Ours on 3 different datasets. Training was done on a single A 110 GPU with a fixed batch-size of 16.
Formally, given Y, E R Nx" and Y, E R RNxm the features representation of the student and the teacher, respectively. N the number of neurons in the last hidden layers of the teacher and the student. n is the batch size. To generate a ran- dom binary mask M, we create a vector of size N containing independent Binomial random vari- ables. Each variable in this vector takes the value with a probability p and 0 otherwise. We multiply element-wise this binary mask M by each column of Y, and Y, to generate the masked features rep- resentation Y; of the teacher and Y, of the student, respectively:
6.3 Sensitivity the Batch Size
Our proposed loss function depends on the batch- size. As shown in Table 8, in all the datasets, the performance of our approach is consistent and sta- ble across all the batch-size values. This illustrates the robustness of our proposed method and shows that it work well even with small batch-size values. However, it is clear when the batch-size increases, the results marginally increase on most of the datasets. This is logical because by increas- ing the number of samples, the correlation value become more accurate.
M ~ Bernoulli(p)
Y = (Y, C M. Y o M. C M...,, Y O M]
M] Where Y represents the first column (the fea- ture representation of the first sample in the input batch X) of the features representation Y, of the teacher. ② denotes the element wise multiplication. After removing the masked rows, i.c., rows with all elements zeros, from Y and Y'. obtain the new features representation of the teacher Y E mxn and the student Y E mm where m < N. In the next steps, we apply the same One-to-Onc Mapping approach described in section 3. How-
Y = IY M.Y. M,Y M,
7 Masked One-to-One Mapping
In our One-to-One Mapping feature distillation ap- proach, the goal for the student is to mimic the last hidden layer representation of the teacher. To achieve this, we break down the global feature dis- tillation task into local sub-tasks. To further

---

Table 7: Training time of PKD and Ours. The results are in hours (H)
approach MNLI-m (H) QNLI(H)
SST-2(H)
PKD
3.00
0.73
0.53
Ours
2.13
0.63
0.47
Table 8: The performance of our approach (Ours) across different batch-size size values. Average over 5 trials.
Batch-size
MNLI-m (%)
QNLI (%)
SST-2 (%)
RTE (%)
82.41
89.43
90.18
63.24
8
82.54
89.53
90.44
65,85
16
82.76
89.54
90.44
63.83
ever, instead of having N specialized sub-teachers and N focused sub-students, in our Masked One-to- One Mapping, we have m specialized sub-teachers and m focused sub-students. Furthermore, we have sub-tasks where each focused sub-student at- tempts mimic the feature representation of the corresponding specialized sub-teacher while disre- garding the others. To achieve this, we utilize the modified training objective function 6.
Diviitation Leariigg
Optimization Step;
Optmizetion Slep2
Optmization
LKD- - > A (a) + A2
Σ
(9)
Steps
j
Figure 4: Masked One-te to to-One Mapping: In each opti- mization step, the neurons are marked by the red color will be masked.
We conduct experiments on GLUE datasets us- ing the same setup as described in subsection 4.2. However, in these new experiments, we extend the training time to 5 epochs because as shown in Tables 3 and 4, the average performance the GLUE tasks of our student model is higher when trained for epochs compared when trained for 3 epochs. We introduce p as the percentage (%) of sub-tasks be learned at given optimization step. We conduct experiments for all values of E {10,20 20, 30, 40, 50, 70, 90). However, we report the results of the two best setups, in function of in terms of the average performance across all 8 GLUE datasets. Additionally, provide the best performance achieved for each individual dataset. For each dataset, the experiment is repeated for seeds and the average reported.
8 Conclusion and Future Work
In this paper, reformulated the global feature distillation problem into N local sub-problems. We proposed a one-to-one mapping between each neu- ron in the last hidden layer of the teacher, i.e., specialized sub-teacher, and each neuron in the last hidden layer of the student. i.e., focused sub- student. To achieve this goal, we proposed a local correlation-basee loss across batches between each specialized sub-teacher output and its correspond- ing focused sub-student output. Our approach only requires the teacher and the student to have the same last hidden layer size. Several experiments proved the effectiveness and consistency of our method. Our approach can be added to any KD method in NLP or vision. Moreover, we proposed an augmented variant of our One-to-One One Mapping approach called Masked One-to-One toee Mapping. In- stead of learning all the sub-tasks simultancously, make the student learn subset of the tasks time. Future work includes exploring the same distillation process with several intermediate layers and experimenting with our approach beyond the NLP tasks.
As illustrated in Table 4, the average perfor- mance on the GLUE datasets for the Masked One- to-One Mapping approach, referred to as Ours- Masked, is higher than the value obtained by the One-to-One Mapping, referred to as Ours, for the two p values 80% and 70%. When considering the best performance across all p values for each individual dataset, denoted as Ours-Masked(Bes)), the gap between the teacher performance (81.30%) and the student performance (80.13%) is effectively reduced.

---

Limitations
Jonathan Frankle and Michael Carbin. 2018. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint 3003::8883.0300
In this work, the proposed approach requires the student and the teacher to have the same last hid- den layer size, which might be a limitation. One possible solution could be to investigate the impact of adding projector after the last hidden layer of the teacher to make it match the last hidden layer size of the student.
Tommaso Furlanello, Zachary Lipton, Michael Tschan- nen, Laurent Itti, and Anima Anandkumar, 2018. Born again neural networks. In International Con- ference on Machine Learning, pages 1607-1616. PMLR.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015. Distilling the knowledge in neural network. arXiv preprint arXiv: 1503.02531.
Ethics Statement
This work is mainly about task-specific knowledge distillation in NLP. It proposes new viewpoint on how to study the feature distillation task. This new approach does not pose any ethical issues as the utilized datasets are publicly available and have already been used in previous research.
Michael E Houle, Hans-Peter Kriegel, Peer Kröger. Erich Schubert, and Arthur Zimek. 2010. Can shared- neighbor distances defeat the curse of dimensional- ity? In Scientific and Statistical Database Man- agement: 22nd International Conference, SSDBM 2010, Heidelberg, Germany, June 30-July 2, 2010. Proceedings 22, pages 482-500. Springer.
Acknowledgment
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang. and Qun Liu. 2020.
The project on which this report is based was funded by the German Federal Ministry of Education and Research (BMBF) under the funding code 011S20049. The author is responsible for the content of this publication.
TinyBERT: Distilling BERT for natural language un- derstanding. In Findings of the Association for Com- putational Linguistics: EMNLP 2020, pages 4163- 4174. Association for Computational Linguistics.
Fadera Ministty Educalan
Sehoon Kim, Amir Gholami, Zhewei Yao. Michael W Mahoney, and Kurt Keutzer. 2021. I-bert: Integer- only bert quantization. In International conference machine learning, pages 5506-5518. PMLR.
Olga Kovaleva. Alexey Romanov, Anna Rogers, and Anna Rumshisky. 2019. Revealing the dark secrets of BERT. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natu- ral Language Processing (EMNNP-IICNNLP), pages 1365-4374, Association for Computational Linguis- tics.
References
Charu C Aggarwal, Alexander Hinneburg, and Daniel A Keim. 2001. On the surprising behavior of dis- tance metrics in high dimensional space. In Interna- tional conference on database theory, pages 420-434. Springer.
François Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush. 2021. Block pruning for faster trans- formers. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10619-10629. Association for Computational Linguistics.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman. Kevin Gimpel, Piyush Sharma, and Radu Soricut. 2019. Albert: lite bert for self-supervised learn- ing of language representations. arXiv preprint 1XX9:1999.11944
Kevin Clark, Minh-Thang Luong, Quoe V Le, and Christopher Manning, 2020. Electra: Pre-training text encoders as discriminators rather than generators. arXiv preprint arXiv:2003. 10555.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language under- standing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, Volume I (Long and Short Papers), pages 4171 11766 Association for Computational Linguis- tics.
Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi. Ke Ye, Fe- lix Chern, Felix Yu, Ruiqi Guo, al. 2022. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In The Eleventh Interna- tional Conference on Learning Representations.
Kevin J Liang, Weituo Hao, Dinghan Shen, Yufan Zhou, Weizhu Chen, Changyou Chen, and Lawrence

---

Carin, 2020. Mixkd: Towards efficient distilla- tion of large-scale language models. arXiv preprint arXiv:201 1.00593.
Hinrich Schütze, Christopher D Manning, and Prab- hakar Raghavan. 2008. Introduction to information retrieval, volume 39. Cambridge University Press Cambridge.
Mingsheng Long, Zhangjic Cao, Jianmin Wang, and Philip S Yu. 2017. Learning multiple tasks with mul- tilinear relationship networks. Advances in neural information processing systems, 30.
Samuel Stanton, Pavel Izmailov, Polina Kirichenko, Alexander A Alemi, and Andrew G Wilson. 2021. Does knowledge distillation really work? Advances in Neural Information Processing Systems, 34:6906- 6919.
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. 2011. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the associ- ation for computational linguistics: Human language technologies, pages 142-150.
S. Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019. Patient knowledge distillation for bert model com- pression. In Conference on Empirical Methods in Natural Language Processing.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances neural information processing systems, 30.
Sai Prasanna, Anna Rogers, and Anna Rumshisky. 2020. When BERT Plays the Lottery, All Tickets Are Win- ning. In Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing (EMNLP), pages putational Linguisties.
Association for Com-
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2018, GLUE: A multi-task benchmark and analysis platform for nat- ural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLPP: Analyzing and Interpreting Neural Networks for NLP, pages 353 355. Association for Computational Linguistics.
Gabriele Prato, Ella Charlaix, and Mchdi Reza- gholizadeh. 2020. Fully quantized transformer for machine translation. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 1-14, Association for Computational Linguisties.
Canwen Xu and Julian McAuley. 2022. A survey on model compression for natural language processing. ArXiv, abs/2202.07105.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAl blog. 188):
Jing Yang, Brais Martinez, Adrian Bulat, and Georgios Tzimiropoulos. 2020. Knowledge distillation via softmax regression representation learning. In Inter- national Conference on Learning Representations.
Pranav Rajpurkar, Jian Zhang. Konstantin Lopyrev, and Percy Liang. 2016. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250.
Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. 2020. Ternarybert: Distillation-awar ultra-low bit bert. arXiv preprint arXiv:2009. 20002.0992..2....
Machel Reid, Edison Marrese- Maylese-Taylor, and Yutaka Mat- suo. 2021. Subformer: Exploring weight sharing for parameter efficiency in generative transformers. In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 4081-4090. Asso- ciation for Computational Linguistics.
Borui Zhao, Quan Cui, Renjie Song, Yiyu Qiu, and Jiajun Liang. 2022. Decoupled knowledge distilla- tion. In Proceedings of the IEEF/CVF Conference computer vision and pattern recognition, pages 11953- 11962.
Khouloud Saadi, Jelena Mitrović, and Michael Gran- itzer. 2023. Learn from one specialized sub-teacher: One-to-one mapping for feature- based knowledge distillation. In ICML 2023 Workshop Neural Com- pression:
Aojun Zhou, Anhang Yao, Yiwen Guo, Lin Xu, and Yurong Chen. 2017. Incremental network quanti- zation: Towards lossless cans with low- -ww-irccision weights. arXiv preprint arXiv: 7702.00000000 03044.
From Information Theory to Applications.
Khouloud Saadi and Muhammad Taimoor Khan. 2022. Effective prevention of semantic drift in continual deep learning. In Intelligent Data Engineering and Automated Learning- -DEAL 2022: 23rd Interna- tional Conference, IDEAL 2022, Manchester, UK November 24-26, 2022, Proceedings, pages 456-464. Springer.
Kaitlyn Zhou, Kawin Ethayarajh, Dallas Card, and Dan Jurafsky. 2022. Problems with cosine as a measure of embedding similarity for high frequency words. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume
2: Short Papers), pages 401 423. Association for Computational Linguistics.
Wangchunshu Zhou, Canwen Xu, and Julian McAuley. 2021. Bert learns to teach: Knowledge dis- tillation with meta learning. arXiv preprint arXiv:2106.0453 045700.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint


Reviews:
[{"review_id": "SUIGdA3toI", "reviewer": "Reviewer_pptK", "summary": "", "questions": "1. How do you ensure that the one-to-one mapping between the teacher and student neurons is consistent across different batches and epochs? Do you use any regularization or alignment techniques to prevent the mapping from changing over time?\n2. How do you handle the case where the student and teacher have different last hidden layer sizes? Do you use any projection or dimensionality reduction methods to make them compatible?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and your interesting feedback.  We include clarifications to individual points below. If there are any remaining questions, please let us know.\n\n>**Q1**:How do you ensure that the one-to-one mapping between the teacher and student neurons is consistent across different batches and epochs? Do you use any regularization or alignment techniques to prevent the mapping from changing over time?\n\nThe mapping between each sub-student and each sub-teacher is fixed, i.e. based on an index, and thus remains consistent over batches and epochs (see Equations and Figure 1). For example, the sub-teacher$_1$ (the first neuron in the teacher's last hidden layer) is mapped to the sub-student$_1$ with the same index (the first neuron in the student's last hidden layer), the sub-teacher$_2$ (the second neuron in the teacher's last hidden layer) is mapped to the sub-student$_2$ with the same index (the second neuron in the student's last hidden layer), etc. So, this mapping is fixed over epochs and batches.\n\n>**Q2:** How do you handle the case where the student and teacher have different last hidden layer sizes? Do you use any projection or dimensionality reduction methods to make them compatible?\n\nKindly refer to Line 079 of the paper and to the conclusion section. In our work, we assume that the teacher and student's last hidden states have the same dimension, as in [1]. We also mentioned that future work will include exploring the addition of a projector to address this limitation.\n\n>**Q3:** The paper does not provide a clear theoretical justification or analysis for why the proposed One-to-One Mapping approach works better than the global metrics such as MSE or cosine distance.\n\nWe have discussed the drawbacks using MSE (e.g. sensitive to scale Lines 065/184 ) and the cosine distance (e.g. sensitive to permutation Line 184) objective functions in the introduction (Line 060)  and the related work (Line 205), which we think is the relevant theoretical justification for our approach. We also linked different references and examples to support our claims.\n\n>**Q4:** The paper does not compare its approach with other recent methods for feature-based knowledge distillation, it needs to be supplemented.\n\nThe main contributions of our work, first, is identifying the drawbacks of the existing objective functions (MSE and cosine distance) used to perform the feature distillation task between two transformer layers. Second, breaking down the global layer distillation task into multiple local sub-tasks and introducing a correlation-based loss function to alleviate the limitation of previous loss functions (MSE and cosine distance-based). Thus, rather than looking at our approach as a competing method to the KD state-of-the-art approaches, it is proposed to replace any MSE or cosine distance-based feature distillation methods between two hidden layers included in any knowledge distillation approach. \\\nTo show that our approach can outperform MSE and cosine distance, applied globally, in the existing pipelines, we experiment now with SRRL[2], one of the latest KD approaches. \nSRRL KD loss has 2 components: the feature loss, which is an MSE between the last layer's features of the teacher and the last layer's features of the student, and the logit loss, which is an MSE between the teacher logits and the student feature after being passed through the teacher classifier.\\\nIn the experiment below, we replace the feature KD loss in SRRL with our newly introduced correlation loss function.\nAs shown in the Table below, our feature distillation approach replacing  MSE loss yields superior performance for SRRL on two GLUE datasets.\n\n| approach  | QNLI(\\%) | RTE(\\%) |\n|-----------|----------|---------|\n| SRRL      | 88.04    | 63.68   |   |\n| SRRL-Ours | 88.68    | 64.10   |   |\n\n**SRRL-Ours: SRRL with its feature distillation loss replaced by our proposed loss.**\\\nWe will include these results in the final version of the paper.\\\nTo sum this up, our approach can be incorporated in any KD approach to replace MSE/cosine distance, which are usually used in feature distillation [1-3] , with our correlation based loss. In this example, we applied it to SRRL. However, this can be incorporated in [1,3]. In this sense, our method is not competing with KD approaches. However, it is a new alternative for using MSE and cosine distance as global metrics in feature distillation.\n\n>**Q5:** The paper does not conduct ablation studies to investigate the impact of different components or hyperparameters of the proposed approach.\n\nWe thank the reviewer for the suggestion. Now, we add an ablation study to investigate the impact of different terms of our algorithm . We report the results in Table below. As can be seen, both terms of the loss are important and their combination achieves the best results across most of the datasets. Thus, the two components of the loss have a compound effect. We will include these results and discussion in the final version of the paper.\n\n| approach     | MRPC         | RTE            | CoLA           | SST-2          | STS-B          | QNLI           | MNLI-m         | QQP            | AvG            |\n|--------------|--------------|----------------|----------------|----------------|----------------|----------------|----------------|----------------|----------------|\n| $L_{CE}+l$   | **86.85** | 63.61          | 48.24          | **90.87** | 85.47          | **89.63** | 81.85          | 87.60          | 79.26          |\n| $L_{CE}+R$   | 75.85        | 53.79          | 34.71          | 85.25          | 85.25          | 85.66          | 74.43          | 87.69          | 68.49          |\n| $L_{CE}+R+l$ | 86.57        | **63.83**| **50.73** | 90.44          |**85.66**| 89.54          |**82.76** | **88.04** | **79.70**|\n\nAll the values in the table are in (\\%). Reported results are the average over 5 random seeds.\\\n$L_{CE}$: is the hard loss\\\n$l$: is the first term in the $L_{KD}$ loss.\\\n$R$: is the second term in the $L_{KD}$ loss (the repel term).\n\n\n>**Q6:** It also does not provide any error analysis or qualitative examples to illustrate the strengths and limitations of the proposed approach in different scenarios or tasks.\n\nIn terms of error analysis, we followed the standard approach found in related work, which usually does not contain qualitative examples. We evaluated our approach on 8 GLUE benchmark datasets. The GLUE dataset, which consists of different tasks, is the typical benchmark for Knowledge Distillation (KD) in NLP[4]. We also evaluated the proposed approach on the SQUAD.V1 and the IMDB datasets. We repeated the experiments for 5 random seeds and we reported the average (Table 3 and Table 4) or the average/standard deviation (Table 1 and Table 2).\n\n[1] Patient Knowledge Distillation for BERT Model Compression, S. Sun and Yu Cheng and Zhe Gan and Jingjing Liu, EMNLP2019.\\\n[2] Knowledge distillation via softmax regression representation learning, Jing Yang, Brais Martinez, Adrian Bulat, Georgios Tzimiropoulos. ICLR2021.\\\n[3] TinyBERT: Distilling BERT for Natural Language Understanding. Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, Qun Liu. EMNLP Findings 2020.\\\n[4] BERT Learns to Teach: Knowledge Distillation with Meta Learning. Wangchunshu Zhou, Canwen Xu, Julian McAuley. ACL2021.", "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": "1. The paper proposes a novel and effective approach for feature-based knowledge distillation, which is a technique for compressing large and complex neural networks into smaller and more efficient ones.\n2. The paper introduces a new perspective on feature distillation, where each neuron in the last hidden layer of the teacher network is considered as a specialized sub-teacher, and each neuron in the last hidden layer of the student network is considered as a focused sub-student. The paper argues that this perspective can better align the feature representations of the teacher and the student than using global metrics such as mean square error or cosine distance.\n3. The paper conducts extensive experiments on several natural language processing tasks and datasets, such as SQUAD-V1, IMDB, and GLUE, and shows that its proposed approach outperforms the state-of-the-art methods by maintaining higher performance on most benchmark datasets. The paper also presents a randomized variant of its approach, called Masked One-to-One Mapping, where the student learns a subset of sub-tasks at each optimization step. The paper reports that this variant yields superior results than learning all sub-tasks simultaneously.\n", "reasons_to_reject": "1. The paper does not provide a clear theoretical justification or analysis for why the proposed One-to-One Mapping approach works better than the global metrics such as MSE or cosine distance. It only relies on empirical results to support its claims, which might not be convincing enough for some reviewers or readers.\n2.The paper does not compare its approach with other recent methods for feature-based knowledge distillation, it needs to be supplemented.\n3. The paper does not conduct ablation studies to investigate the impact of different components or hyperparameters of the proposed approach. It also does not provide any error analysis or qualitative examples to illustrate the strengths and limitations of the proposed approach in different scenarios or tasks.\n\n", "questions_to_authors": "1. How do you ensure that the one-to-one mapping between the teacher and student neurons is consistent across different batches and epochs? Do you use any regularization or alignment techniques to prevent the mapping from changing over time?\n2. How do you handle the case where the student and teacher have different last hidden layer sizes? Do you use any projection or dimensionality reduction methods to make them compatible?", "paper_topic_and_main_contributions": "This paper is about knowledge distillation, which is a technique for compressing large and complex neural networks into smaller and more efficient ones. The paper addresses the problem of how to transfer the knowledge of the teacher network, which is the large model, to the student network, which is the small model, by using the intermediate feature representations of both networks. The main contributions of the paper are:\n\n1. It proposes a novel approach called One-to-One Mapping, where each neuron in the last hidden layer of the teacher network is considered as a specialized sub-teacher, and each neuron in the last hidden layer of the student network is considered as a focused sub-student. The paper argues that this approach can better align the feature representations of the teacher and the student than using global metrics such as mean square error or cosine distance.\n\n2. It introduces a local correlation-based objective function to achieve the One-to-One Mapping distillation, where each focused sub-student tries to maximize the correlation with its corresponding specialized sub-teacher and minimize the correlation with other sub-teachers. The paper claims that this objective function is more stable and accurate than using mean square error or cosine distance.\n\n3. It conducts extensive experiments on several natural language processing tasks and datasets, such as SQUAD-V1, IMDB, and GLUE, and shows that its proposed approach outperforms the state-of-the-art methods by maintaining higher performance on most benchmark datasets. The paper also presents a randomized variant of its approach, called Masked One-to-One Mapping, where the student learns a subset of sub-tasks at each optimization step. The paper reports that this variant yields superior results than learning all sub-tasks simultaneously.\n", "excitement": 2, "reproducibility": 2, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YeOwzCyEHL", "reviewer": "Reviewer_tgjk", "summary": "", "questions": "See above.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and your interesting feedback.  We include clarifications to individual points below. If there are any remaining questions, please let us know.\n\n>**Q1:** More advanced knowledge distillation methods should be compared. For example, Contrastive Representation Distillation, Decoupled KD, etc.\n\nThe main contributions of our work, first, is identifying the drawbacks of the existing objective functions (MSE and cosine distance) used to perform the feature distillation task between two transformer layers. Second, breaking down the global layer distillation task into multiple local sub-tasks and introducing a correlation-based loss function to alleviate the limitation of previous loss functions (MSE and cosine distance-based). Thus, rather than looking at our approach as a competing method to the KD state-of-the-art approaches, it is proposed to replace any MSE or cosine distance-based feature distillation methods between two hidden layers included in any knowledge distillation approach. \\\nTo show that our approach can outperform MSE and cosine distance, applied globally, in the existing pipelines, we experiment now the with SRRL[2], one of the latest KD approaches. \nSRRL KD loss has 2 components: the feature loss, which is an MSE between the last layer's feature of the teacher and the last layer's feature of the student, and the logit loss, which is an MSE between the teacher logits and the student feature after being passed through the teacher classifier.\\\nIn the experiment below, we replace the feature KD loss in SRRL with our newly introduced correlation loss function.\nAs shown in the Table below, our feature distillation approach replacing global MSE yields superior performance for SRRL.\n\n| approach  | QNLI(\\%) | RTE(\\%) |\n|-----------|----------|---------|\n| SRRL      | 88.04    | 63.68   |   |\n| SRRL-Ours | **88.68**    | **64.10**   |   |\n\n**SRRL-Ours: SRRL with its the feature distillation loss replaced by our proposed loss.**\\\nWe will include these results in the final version of the paper.\n\nTo sum this up, our approach can be incorporated in any KD approach to replace MSE/cosine distance, which are usually used [1-3] in feature distillation, with our correlation based loss. In this example, we applied it to SRRL. However, this can be incorporated in [1,3]. In this sense, our method is not competing with KD approaches. However, it is a new alternative for using MSE and cosine distance as global metrics in feature distillation.\n\n>(a) I personally find the novelty in this paper rely on the design of selecting the correlation as the metric, meanwhile reduce the correlation between different neurons. This has a similar insight with dropout by increasing the independence of every unit. I think the intuition behind these strategies can be more explored. I know in the relative work it has been discussed somehow, however, I think it can be more detailed, by showing why it is important to reduce the correlation between different neurons.\n\nThe reviewer makes an interesting point. The connection between our approach and dropout is interesting. However, from our perspective, our proposed approach is different compared to the dropout mechanism. It is true that dropout can reduce  the correlation between the units within the same layer. But, in our approach we reduce the correlation between student units and teacher units in order to facilitate the sub-task for the sub-student. In other words, the correlation is computed between the corresponding units from two different models.\n\n>(b) Why correlation as metric is superior. \n\nAs mentioned in the introduction and the related work, MSE and cosine distance suffer from various limitations such as sensitivity to scale and sensitivity to permutation. Kindly refer to Line 060 and Line 184  to see the drawbacks of using MSE and cosine distance.\n\n>(c) Plus, the paper writes \"In a high dimensional space, which is the case in a neural network, the data tend to be sparse and all the data points become uniformly distant from each other\", which is not the real case in deep neural network as far as I know, especially Transformer-based networks.\n\nThere seems to be a misunderstanding of our claim, what we wanted to say is that the embeddings (features representation) provided by BERT-like models are high-dimensional and sparse tensors in nature, see [4]. Thus, using MSE or cosine distance to measure the distance between these sparse representations is not optimal as explained and advocated in [5]. So, splitting the task into sub-tasks and using one-to-one metric alleviates this problem and breaks the curse of dimensionality. We will make this more clear in the final version of the paper. \n\n>Typos Grammar Style And Presentation Improvements:Please check if the equation at line 184 is mistaken.\n\nThank you for pointing out this. We will fix it in the final version of the paper.\n\n[1] Patient Knowledge Distillation for BERT Model Compression, S. Sun and Yu Cheng and Zhe Gan and Jingjing Liu, EMNLP2019.\\\n[2] Knowledge distillation via softmax regression representation learning, Jing Yang, Brais Martinez, Adrian Bulat, Georgios Tzimiropoulos. ICLR2021.\\\n[3] TinyBERT: Distilling BERT for Natural Language Understanding. Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, Qun Liu. EMNLP Findings 2020.\\\n[4] THE LAZY NEURON PHENOMENON: ON EMERGENCE OF ACTIVATION SPARSITY IN TRANSFORMERS ,Zonglin Li , Chong You , Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J. Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, and Sanjiv Kumar. ICLR2023.\\\n[5] On the Surprising Behavior of Distance Metrics in High Dimensional Space, Charu C. Aggarwal1, Alexander Hinneburg, and Daniel A. Keim2. ICDT 2001.", "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": "The proposed distillation method is novel. And the proposed idea is easy to follow. Besides, compared to the baseline, the experiments show that the proposed method have promising results.\n\nThe implementation is also straight, which may be practical in real scenarios. \n\nThe illustration about other metric's sensitive to unit permutation is intuitive.\n", "reasons_to_reject": "1. More advanced knowledge distillation methods should be compared. For example, Contrastive Representation Distillation, Decoupled KD, etc.\n2. I personally find the novelty in this paper rely on the design of selecting the correlation as the metric, meanwhile reduce the correlation between different neurons. This has a similar insight with dropout by increasing the independence of every unit. I think the intuition behind these strategies can be more explored. I know in the relative work it has been discussed somehow, however, I think it can be more detailed, by showing why it is important to reduce the correlation between different neurons, and why correlation as metric is superior. Plus, the paper writes \"In a high dimensional space, which is the case in a neural network, the data tend to be sparse and all the data points become uniformly distant from each other\", which is not the real case in deep neural network as far as I know, especially Transformer-based networks. ", "questions_to_authors": "See above.", "paper_topic_and_main_contributions": "The paper introduces a novel approach to knowledge distillation based on features by aligning every output of neurons from teacher and student. Instead of the usually used MSE, the paper proposes to maximize the correlation between mapped outputs and reduce that from unmapped ones. ", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "None", "Typos_Grammar_Style_and_Presentation_Improvements": "Please check if the equation at line 184 is mistaken.\n"}, "extra_text": {}}, {"review_id": "gn3o3QH5J5", "reviewer": "Reviewer_BBGK", "summary": "", "questions": "Please see my questions in the section \"Reasons to Reject.\"", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and your interesting feedback.  We include clarifications to individual points below. If there are any remaining questions, please let us know.\n\n>(a) Line 194 states the proposed loss function is more stable, but I do not find theoretical explanations or empirical evidence that supports this statement. In what aspects do we consider the proposed method more stable – less sensitive to unit permutation, robust to hyper-parameter choice, or something else?\n\nStability in this context is used in a sense that the results are stable which is reflected by a low value of standard deviation over the 5 random seeds. As you can see in Table 1 and Table 2, our results gave lower standard deviation with respect to different random seeds compared to the other approaches. We will make this clear in the final version of the paper. \n\n>(b) Line 295 claims the proposed method can be done with a low computation cost. However, there are no relevant statistics such as GPU hours to support the claim.\n\nBy low computation, we meant lower compared to previous approach, more specifically to PKD [1]. Note that, unlike PKD [1], which distills different hidden-state features plus the logits, our approach distills only the last hidden feature of the teacher to the student. Moreover, we can now confirm this with quantitative results by reporting the training time of each approach on 3 different datasets. We chose 3 big-size datasets from GLUE in order to have a fair comparison. Training was done on a single A100 GPU with a fixed batch-size of 16. As show in Table below, our approach is trained in less time. We will include these results in the final version of the paper. Thank you for your suggestion.\n\n\n| approach | MNLI-m (H)  | QNLI (H)       | SST-2 (H)      |   |\n|----------|-------------|---------------|---------------|---|\n| PKD      | 3.00        | 0.73          | 0.53          |\n| Ours     | **2.13** | **0.63** |**0.47** |\n\nThe results are in hours (H).\n\n>(c) Why does a higher EM, F1, or accuracy prove the proposed approach “facilitates the convergence of the student’s representation to the teacher’s representation” (line 348)? Why can’t it be the case that the student model represents the features vastly differently from the teacher model but has learned a better representation for the task? Although the claim is making sense, one would need to specifically look at the hidden representations in order to show this convergence.\n\nOur approach is formulated directly as maximizing the correlation between each unit in the student and its corresponding unit in the teacher. Thus, it is safe to assume that after training our model, the student's features representation is highly correlated to the teacher's features representation or in other words very similar to it. But we agree with the reviewer's point of view, this claim can be softened in the final version of the paper and we can replace it with \"yields better representation of the student features\"\n\n>(d) Line 462 states that masked one-to-one mapping “enables the student to effectively absorb and manage the incoming flow of knowledge from the teacher.” However, the empirical results in Table 3 (where no masking is applied) contain marginally higher accuracy for several tasks (e.g., MRPC, SST-2, etc.). This claim sounds untenable.\n\nThis claim is based on the fact that Ours-Masked outperforms ours in average across all the datasets. However, the reviewer has a point, on some datasets out of the 8, ours outperforms Ours-Masked. Interestingly this happens mostly when the dataset is of small size. We will include this analysis in the final version of the paper.\n\n>**Q2:** The paper does not provide sufficient details on the datasets used in the experiments. The authors should not assume that a reader is familiar with what task SQUAD-V1 or IMDB is addressing and what evaluation metrics are for the dataset.\n\nThank you for pointing this out.\n\n*Stanford Question Answering Dataset (SQuAD)* [2]\nis a reading comprehension dataset, consisting of questions posed by crowd workers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage. \\\n*Exact Match (EM)* metric [2]: measures the percentage of predictions that match any one of the ground truth answers exactly.\\\n*F1 score metric* [2]: measures the average overlap between the prediction and ground truth answer. \\\n*Internet Movie Database dataset (IMDB)* [3]  is a sentiment analysis dataset consisting of 50,000 movie reviews labeled as positive or negative. \\\nWe will include this description in the final version of the paper.\n\n>(a) Does each component of the loss function contribute to the effectiveness of the proposed method? ...\n\nOur approach is composed of two parts. The first term makes each sub-student to focus with the corresponding sub-teacher. The second term helps to make the task easy for each sub-student. We need to remember that the student has lower number of parameters compared to the teacher. To this end, we force each sub-student to ignore the redundant knowledge learnt by the other sub-teachers and to focus only on the distinctive information that it can get from its corresponding sub-teacher. As shown by the ablation study we add now, this second term helps to boost the results of most of the datasets.\n\n| approach     | MRPC         | RTE            | CoLA           | SST-2          | STS-B          | QNLI           | MNLI-m         | QQP            | AvG            |\n|--------------|--------------|----------------|----------------|----------------|----------------|----------------|----------------|----------------|----------------|\n| $L_{CE}+l$   | **86.85** | 63.61          | 48.24          | **90.87** | 85.47          | **89.63** | 81.85          | 87.60          | 79.26          |\n| $L_{CE}+R$   | 75.85        | 53.79          | 00.00          | 85.25          | 85.25          | 85.66          | 74.43          | 87.69          | 68.49          |\n| $L_{CE}+R+l$ | 86.57        |**63.83** | **50.73** | 90.44          | **85.66** | 89.54          |**82.76**|**88.04**| **79.70** |\n\nAll the values in the table are in (\\%). Reported results are the average over 5 random seeds.\\\n$L_{CE}$: is the hard loss.\\\n$l$: is the first term in the $L_{KD}$ loss.\\\n$R$: is the second term in the $L_{KD}$ loss (the repel term).\n\n>(b) The cross-correlation contains a variable, which is the batch size. Does the proposed method work equally well regardless of the batch size, or does it not work well if the n is small?\n\n(b) Indeed, our proposed loss function depends on the batch-size. In our experiments, we found that the results are not very sensitive to the batch size. To show that, now, we include the results in Table below with different batch-size values. The chosen datasets are with different sizes (big, medium, and small). As shown in Table below, in all the datasets, the performance of our approach is consistent and stable across all the batch-size values.  This illustrates the robustness of our proposed method and  shows that it can work well even with small batch-size values. However, we note when the batch-size increases, the results marginally increase on most of the datasets. This is reasonable because by increasing the number of samples, the correlation value become more accurate. We will include these results in the final version of the paper.\n\n| Batch-size | MNLI-m (\\%) | QNLI (\\%) | SST-2 (\\%) | RTE (\\%) |\n|------------|-------------|-----------|------------|----------|\n| 4          | 82.41       | 89.43     | 90.18      | 63.24    |\n| 8          | 82.54       | 89.53     | 90.44      | 65.85    |\n| 16         | 82.76       | 89.54     | 90.44      | 63.83    |\n\n\n\n>(a) The cosine distance formula contains typos in the subscripts.(b)The loss functions in Fig. 1 are missing the squaring operation around.\n\nThank you for mentioning this, we will fix them in the final version of the paper.\n\n[1] Patient Knowledge Distillation for BERT Model Compression, S. Sun and Yu Cheng and Zhe Gan and Jingjing Liu, EMNLP2019.\\\n[2] SQuAD: 100,000+ Questions for Machine Comprehension of Text. Rajpurkar, Pranav  and Zhang, Jian  and Lopyrev, Konstantin  and Liang, Percy. EMNLP2016.\\\n[3] Learning Word Vectors for Sentiment Analysis. Maas, Andrew L.  and Daly, Raymond E.  and Pham, Peter T.  and Huang, Dan  and Ng, Andrew Y.  and Potts, Christopher. ACL2011.", "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": "- The paper clearly illustrates the proposed idea with helpful visualizations and precise mathematical formulas.\n\n- The empirical results in the paper show that the proposed loss function for knowledge distillation leads to enhanced student model performance compared to existing methods. Results are repeated runs and thus are convincing in terms of statistical significance.", "reasons_to_reject": "1. Some claims are not well justified.\n\n    (a) Line 194 states the proposed loss function is more stable, but I do not find theoretical explanations or empirical evidence that supports this statement. In what aspects do we consider the proposed method more stable – less sensitive to unit permutation, robust to hyper-parameter choice, or something else?\n\n    (b) Line 295 claims the proposed method can be done with a low computation cost. However, there are no relevant statistics such as GPU hours to support the claim.\n\n    (c) Why does a higher EM, F1, or accuracy prove the proposed approach “facilitates the convergence of the student’s representation to the teacher’s representation” (line 348)? Why can’t it be the case that the student model represents the features vastly differently from the teacher model but has learned a better representation for the task? Although the claim is making sense, one would need to specifically look at the hidden representations in order to show this convergence. \n\n    (d) Line 462 states that masked one-to-one mapping “enables the student to effectively absorb and manage the incoming flow of knowledge from the teacher.” However, the empirical results in Table 3 (where no masking is applied) contain marginally higher accuracy for several tasks (e.g., MRPC, SST-2, etc.). This claim sounds untenable. \n\n2. The paper does not provide sufficient details on the datasets used in the experiments. The authors should not assume that a reader is familiar with what task SQUAD-V1 or IMDB is addressing and what evaluation metrics are for the dataset.\n\n3. The analysis of the proposed loss function is not thorough enough. \n\n    (a) Does each component of the loss function contribute to the effectiveness of the proposed method? Why is it necessary to have the second component with hyperparameter $\\lambda_2$ that aims to “repel all knowledge coming from the other teachers”? This could be potentially counter-productive if there exist features that are highly correlated with each other. In fact, I see that $\\lambda_2 = 5 \\times 10^{-3}$ is much smaller than $\\lambda_1=1$ in the experiment section (line 403).\n\n    (b) The cross-correlation contains a variable $n$, which is the batch size. Does the proposed method work equally well regardless of the batch size, or does it not work well if the $n$ is small?", "questions_to_authors": "Please see my questions in the section \"Reasons to Reject.\"", "paper_topic_and_main_contributions": "This paper proposes a feature distillation loss function that aligns the student model’s last hidden layer representation with the teacher’s neuron by neuron. This loss function is different from the existing MSE or cosine distance loss functions to align the hidden layer representations because MSE or cosine distance penalizes different representations per sample while the proposed one-to-one mapping in this paper encourages similar representations per feature (along each dimension across samples in the same batch). \n\nIn addition, the authors propose a variation of this one-to-one mapping, where only a subset of the neurons in the last hidden layers are optimized for the per-feature distillation loss function. The authors claim this variation introduces additional gains.\n\nOverall, the empirical results in the paper show that the idea is promising.", "excitement": 2, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "The paper is missing citations/links to the datasets used in the experiments (SQUAD, IMDB).", "Typos_Grammar_Style_and_Presentation_Improvements": "- The cosine distance formula contains typos in the subscripts.\n\n- The loss functions in Fig. 1 are missing the squaring operation around $(1-C)$."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work proposes a new distillation loss which maximizes the cross correlation between individual features at pre-logit activations. Authors apply their method on various BERT models fine-tuned on downstream tasks. Reviewers agree on the novelty of the method and find the results promising. Authors also provide further ablations and comparison to related work during the rebuttal.
Decision: Accept-Findings