Exploring All-In-One Knowledge Distillation Framework for Neural Machine Translation
Authors: Zhongjian Miao
Wen Zhang
Jinsong Su
Xiang Li
Jian Luan
Yidong Chen
Bin Wang
Min zhang
Venue: EMNLP 2023 Main (2023)
Keywords: neural machine translation
efficient knowledge distillation
multi-model scenario
TL;DR: In this work, we present AIO-KD, a novel knowledge distillation framework for neural machine translation, which generates multiple satisfactory student models efficiently.

Abstract: Conventional knowledge distillation(KD) approaches are commonly employed to compress neural machine translation(NMT) models. However, they only obtain one lightweight student each time. Consequently, we have to conduct KD multiple times when different students are required at the same time, which could be resource-intensive. Additionally, these students are individually optimized, and thus lack interactions with each other, leading to their potential not being fully exerted. In this work, we propose a novel All-In-One Knowledge Distillation(AIO-KD) framework for NMT, which generates multiple satisfactory students at once. Under AIO-KD, we first randomly extract fewer-layer subnetworks from the teacher as the sample students. Then, we jointly optimize the teacher and these students, where the students simultaneously learn the knowledge from the teacher and interact with other students via mutual learning. When utilized, we re-extract the candidate students, satisfying the specifications of various devices. Particularly, we adopt carefully-designed strategies for AIO-KD: 1) we dynamically detach gradients to prevent poorly-performed students from negatively affecting the teacher during the knowledge transfer, which could subsequently impact other students; 2) we design a two-stage mutual learning strategy, which alleviates the negative impacts of poorly-performed students on the early-stage student interactions. Extensive experiments and in-depth analyses on three benchmarks demonstrate the effectiveness and eco-friendliness of AIO-KD. Our source code is available at https://github.com/DeepLearnXMU/AIO-KD.

Exploring All-In-One Knowledge Distillation Framework for Neural Machine Translation
Zhongjian Miao Wen Zhang Jinsong Su Xiang Li²
Jian Luan² Yidong Chen Bin Wang? Min Zhang²
'School of Informatics, Xiamen University, China
Xiiooii AI Lab, China
Institute of Computer Science and Technology, Soochow University, China miaozhongjianest xmu. edu.cn, [jssu, yddhen edu.cn
Abstract
model performance, there have been many efforts in exploring wider or deeper Transformers (Wang et al., 2019; Xu et al., 2021; Wang et al., 2022). However, deploying such models with billions of parameters on edge devices (e.g., mobile phones) remains to be a challenge. To solve this prob- lem, various methods for model compression have been proposed (Choudhary al., 2020). Among them, knowledge distillation (KD) (Hinton et al., 2015) has been widely adopted due to its effec- tiveness and simplicity. In this regard, researchers explore many KD approaches for NMT, such as Word-KD (Kim and Rush, 2016) and Selective- KD (Wang et al., 2021).
Conventional knowledge distillation (KD) ap- proaches are commonly employed to com- press neural machine translation (NMT) mod- els. However, they only obtain one lightweight student each time. Consequently, we have to conduct KD multiple times when differ- ent students are required at the same time, which could be essurce-intensivve Addition- ally, these students are individually optimized, and thus lack interactions with each other, lead- ing to their potential not being fully exerted. In this work, we propose novel All-In-One Knowledge Distillation (ATO-KD) framework for NMT, which generates multiple satisfactory students at once. Under AIO-KD, we first ran- domly extract fewer-layer subnetworks from the teacher as the sample students. Then, we jointly optimize the teacher and these students, where the students simultaneously learn the knowledge from the teacher and interact with other students via mutual learning. When uti- lized, we re-extract the candidate students, sat- isfying the specifications of various devices. Particularly, we adopt carefully gies for AIO-KD: I) we dynamically detach gradients to prevent poorly perfo-pprfo students from negatively affecting the teacher during the knowledge transfer, which could subsequently impact other students; 2) we design a two- stage mutual learning strategy, which alleviates the negative impacts of poorly-performed stu- dents the early-stage student interactions. Extensive experiments and in-depth analyses three benchmarks demonstrate the effec- tiveness and fo-friendliness of AIO-KD. Our source code available https: //s::///itti thub. ccmm/DeepleeamMM///IO-KK0
Generally, in these conventional KD approaches, the knowledge of large teacher model is trans- ferred to only a compact student model, which usually adopts the same model architecture the teacher but with fewer parameters (Kim and Rush, 2016; Jafari al., 2021; Wang et al., 2021; Yang et al.. 2022). Nevertheless, due to hardware differ- ences, we are often required to deploy models of varying sizes on different devices (Sandler et al., 2018; Wang et al., 2020a; Tan et al., 2022). In this scenario, conventional KD approaches have two defects: 1) they have to be conducted mul- tiple times for different students, which leads substantial costs; 2) students are individually op- timized, and thus they are unable to interact with each other. Note that in the human learning process, communication between students will benefit their learning (Webb, 1989). In light of this, we believe that these students can be further improved through collaborative interactions.
strate-
In this paper, we propose novel All-Tn-One Knowledge Distillation (AIO-KD) framework for NMT, which constructs students from the teacher and jointly optimizes both the teacher and students from scratch. Employing AIO-KD, we regard the fewer-layer subnetworks extracted from the teacher as the candidate students. During training, we first
1 Introduction
In recent years, Transformer (Vaswani et al., 2017) has become the dominant architecture in neural machine translation (NMT). To further improve
Work was done when interning at Xiaomi Al Lab. Corrreppnnding Author.

---

randomly select from all candidate students ob- tain the sample ones at each training step. Then, we jointly optimize the teacher and the sample stu- dents. During this process, the students simulta- neously learn the knowledge from the teacher and interact with other students via mutual learning. When utilized, we re-extract the candidate students from the teacher, satisfying the specifications of various devices.
2021b). However, the majority of these works fo- cus on the Transformer encoder, while paying less attention to the overall Transformer architecture. Our work focuses on the latter, i.e., the overall Transformer with variable depths for NMT.
Knowledge Distillation in NMT. In recent years, model compression technologies have attracted much attention (Han et al., 2016; See et al., 2016; Choudhary et al., 2020). As a commonly-used technology for model compression, knowledge distillation (KD) (Hinton et al., 2015) has been widely used in many natural language processing tasks (Jiao et al., 2020; Wang et al., 2020b; Liu et al., 2020b; Zhou al., 2022c; Zhang al., 2023). In the community of NMT, Kim and Rush (2016) first apply KD to autoregressive NMT. Fur- ther, many studies explore more effective KD ap- proaches for NMT (Zeng et al., 2019; Wei et al., 2019; Zhang et al., 2019; Wang et al., 2021; Liang et al., 2022; Miao et al., 2022: Zhou et al., 2022b). Meanwhile, researchers focus on applying KD various aspects of NMT, including multilingual NMT (Tan et al., 2019; Do and Lee, 2022; Lu al., 2022; Huang et al., 2022b, 2023), unsupervised NMT (Sun et al., 2020; Nguyen et al., 2021), non- autoregressive NMT (Gu et al., 2018; Qian et al., 2021: Huang et al., 2022a; Wang et al., 2023), and ANN-NMT (Yang al., 2022).
To better accommodate AIO-KD, we carefully design the following strategies:
1) Dynamic Gradient Detaching, Under AIO- KD, the students are optimized jointly with the teacher, where the teacher and students mutually influence each other through the KD loss. When there exists significant performance gap between student and the teacher, the gradients of the KD loss specific to the student will harm the teacher, which further negatively affects other students. To address this issue, we measure the performance gap through the cross-entropy ratio of the student to the teacher. If this ratio exceeds a pre-defined thresh- old, we will not utilize these gradients to update the teacher's parameters.
2) Two-Stage Mutual Learning. As mentioned previously, introduce mutual learning fa- cilitate students. To avoid the negative impacts of poorly-performed students on student interac- tions at the early stage, we adopt a two-stage train- ing strategy. Concretely, we first only utilize the signals from the teacher to guide the training of students, and further introduce mutual learning to strengthen the interactions between students. Such multi-stage training strategy has been verified in previous study (Zhou al., 2022a).
Mutual Learning in NMT. Mutual learn- ing (Zhang et al., 2018) has been explored in NMT, with various techniques proposed to improve trans- lation quality. For example, Bi et al. (2019) propose multi-agent learning, where diverse students learn from each other, working together improve trans- lation quality. Liao et al. (2020) explore sentence- level and token-level mutual learning for NMT. Zhao et al. (2021) show the effectiveness of mutual learning in end-to-end speech translation.
Empirical experiments and in-depth analyses on three translation benchmarks demonstrate that AIO- KD is superior to conventional KD approaches in terms of translation quality and training costs. As a bonus, the teacher in AIOKKD is significantly enhanced through knowledge transfer with the stu- dents.
To the best of our knowledge, our work is the first attempt to incorporate both knowledge distil- lation and mutual learning into the variable-depth Transformer. Unlike conventional KD approaches, the candidate students in AIO-KD are the fewer- layer subnetworks derived from the teacher. Dur- ing training, we randomly select from all candidate students obtain the sample ones and jointly opti- mize them with the teacher from scratch, which in- volves knowledge transfer from teacher to students, as well as interactions between students via mu- tual learning. Additionally, we develop carefully-
2 Related Work
Our related works mainly include the following three lines of studies:
Transformer with Variable Depths. To reduce computation costs, plenty of researchers investigate the variable-depth Transformer architecture (Yu et al., 2019; Dehghani et al., 2019; Hou et al., 2020; Xin et al., 2020; Liu et al., 2020a; Fan et al., 2020; Elbayad et al., 2020; Cai et al., 2020; Liu et al.,

---

1-th Layer
Layer
Bl-4
DU-1
DL-5
EU-5
DL-5
EL
DL-4
The Stadest S(4.2)
EL-4
DL-4
DL-2
Cu
Lmi
BL-3
DL-3
EL-3
DL.1
FL-2
DL-2
2
D4-2
B1-1
D.
Fi-3
B.1
The Student S(4.2)
The Teacher
F1-2
The Teacher T
EL-1
The Student S(4,3)
Figure 1: Left: An illustration of AIO-KD with sampling 2 students, Right: An example of constructing the student S(4 2) from the teacher T. Employing AIO-KD, we optimize the sample students only at each training step, yet can still obtain plenty of satisfactory candidate students. S(la, la) refers to the student with L. encoder layers and la decoder layers. Note that the teacher is jointly trained with the sample students from scratch, with the students learning the knowledge from the teacher and interacting with other students via mutual learning. Lee, Lkd. and Lmm denote the cross-entropy loss, the knowledge distillation loss, and the mutual learning loss, respectively.
designed strategies for AIO-KD, which have been proven to be effective in subsequent experiments.
{S at each training step. Afterward, we jointly train the teacher and these students, where the students simultaneously learn the teacher's knowledge and interact with other students via mu- tual learning. Notice that during this process, we carefully develop strategies to accommodate AIO- KD, described in Section 3.3. When utilized, we re-extract |C| students from the teacher, satisfying the specifications of various devices.
3 All-In-One Knowledge Distillation
this section, we first give a brief overview of AIO-KD (Section 3.1), and then describe the train- ing objective (Section 3.2). Finally, we detail carefully-designed strategies (Section 3.3).
3.1 The Overview of AIO-KD
The left half of Figure provides overview of AIO-KD. In our work, we adopt the standard Transformer with A encoder/decoder layers as the teacher T. Inspired by recent studies (Wang et al., 2019; Kasai et al., 2021), we extract fewer- layer subnetworks with deep encoder and shal- low decoder from the teacher as the candidate students, which achieve satisfactory performance while maintaining efficient inference. Accord- ingly, all candidate students can be formalized as C=(S(1,, la) 1< la S < W where S(ee, la) refers to the student with L. encoder layers and la decoder layers. The right half of Figure gives an example of constructing the student S(4, 2), we extract the adjacent encoder and decoder layers of the teacher, starting from the first layer, to con- struct it, which shares not only architecture but also parameters with the teacher.
3.2 Training Objective
Overall, the training objective of AIO-KD consists of the following three parts:
L = Lev + aLkd + 3Lmi,
(1)
where Lees Lkd: and L denote the cross-entropy loss, the knowledge distillation loss, and the mutual learning loss, a and ,3 are two coefficients balanc- ing the effects of different losses, respectively.
Cross-Entropy Loss Loe As reported by previ- ous works (Zhang et al., 2018; Guo et al., 2020b), jointly training the teacher and student achieves better knowledge transfer. In this work, the stu- dents are optimized jointly with the teacher from scratch. Formally, decompose Zor into two parts as follows:
Lee = LT +
S
(2)
During training, we first randomly and uni- formly sample from C to obatin sample students
PIn early exploration, attempt explore intelligent sampling strategies for students. However, they do not perform well for AIO-KD.
'We follow previous studies (Sun et al., 2021; Ge et al., 2022) to exclude the models with I-layer deonder because they do not consistently perform well.
$1C = + + (N - -

---

where C and L represent the cross-entropy losses for the teacher T and the student Sk. re- spectively.
lows:
i OKL(PT (PPSS) L
S 1/
00T
Cle
(6)
0, otherwise.
Knowledge Distillation Loss Lkd Employing AIO-KD, we aim to transfer the teacher's knowl- edge to multiple students by aligning the students' output probability distributions with those of the teacher, and Ltd is formulated follows:
Two-Stage Mutual Learning. It is worth noting that the students at the early-stage training are often poorly-perlormes ones, which hinder student inter- actions. To better leverage the potential of mutual learning, we adopt two-stage training strategy. At the first training stage, we only utilize the sig- nals from the teacher to guide the training the students, formulating the training objective as
Ltd E (LPPTIIP))) 1=11
(3)
K.
where (-(-) the Kullback-eeibler distance func- tion, pd and pS, denote the output probability distributions of the teacher T and the student Sk. respectively.
L1 = Lee + aL
(7)
Thereafter, we further introduce Lmdl to optimize the students the second training stage, as shown below:
Mutual Learning Loss Lid To further promote the students, we incorporate mutual learning to fa- cilitate their interactions, with the loss Lidt defined
=oc + aLid + 3L
(8)
Experiments
mm K(K - 1) ISKN'K
Σ
ML PSr), (4)
4.1 Setup
Datasets. We conduct experiments on German- English (De-En), English-Romanian (En-Ro), and English-German (En-De) translation tasks. For the De-En task, use the IWSLT14 De-En corpus, where the training set comprises 160k sentence pairs extracted from TED talks. We use the com- bination of dev2010 and dev2012 as the valida- tion set, and the combination of kst2010, tst2011, and tst2012 as the test set, respectively. For the En-Ro task, we use the dataset of the WMT16 En-Ro as the training set, containing 610k sen- tence pairs. We separately choose newsdev2016 and newstest2016 the validation and test sets. For the En-De task, we use the WMT14 En-De dataset containing 4.5m sentence pairs for training, and we choose news test2013 3 and news test2211 as the validation and test sets, respectively. We employ Byte Pair Encoding (BPE) (Sennrich et al., 2016) to split words into subwords. Following common practices, we set the numbers of merging operations 10k, 32k, and 32k for the three tasks, respectively. Finally, we report case-sensitive kenized BLEU (Papineni et al., 2002) as well COMET (Rei et al., 2020).
Cee L
KL(PSy ||ps), otherwise.
(5)
Notice that for any two students, the one with the lower cross-entropy loss acts as the senior student, leading the process of mutual learning. As dis- cussed by Liao et al. (2020), such mutual learning is beneficial for NMT.
3.3 Carefully-Desinne Strategies
Dynamic Gradient Detaching, As mentioned in Section 3.2, the teacher and students mutually in- fluence each other through the KD loss. When the performance of the student S₁ much inferior to that of the teacher T, the gradients g= (KL(PT ((P5)) 8 By
of the Sy-related KD loss will harm the teacher and the latter may further negatively impact other students, where 0f represents the teacher's param- eters.
To deal with this issue, we calculate the cross- entropy ratio 2S//C' of the student S₁ to the teacher T, which measures their performance gap. If this ratio exceeds a pre-defined threshold 1), argue that there exists significant performance gap between S₁ and T, and thus we do not utilize the gradients g to update the teacher's parameters. To this end, we reformulate the gradients as fol-
Model Configuration. We develop AIO-KD and other baselines with fairseq (Ot et al., 2019). The standard Transformer with encoder/decoder layers is adopted as the teacher. We use the transformer jiwimeerrwwltttt www.l_lsdddee t__ee_n setting for the De-En

---

4.270
4,226
4,226
4,260
4.222
4.224
Avg.
4.250
4,222
4,218
4.240
-
4.220
-0-
4.214
4.0
4.5
5.0
5.5
6.0
0.1
0.3
0.5
0.7
0.9
1,005
1.01
1.05
1.1
w/o DGD
Figure Effects of a, [3, and y) on validation set of the En-Ro task. "Avg. Loss" refers to average cross-entropy losses of all candidate students the validation set, as defined in Equation 9. "w/o DGD" means that we remove dynamic gradient detaching from AIO-KD, with no students being detached.
task, and the transformer_wmt._e the En-Ro and En-De tasks, respectively.
setting for
the output probability distributions of the teacher the complicated words, which have higher cross-entropy losses estimated by the student.
To optimize models, use Adam (Kingma and Ba, 2015) optimizer with 3,=0.9, $2=0.98, and e=10-9 All experiments are conducted on train-
4.2 Effects of Hyper-Parameters
NVIDIA A100 GPUs with ing, where the batch sizes are individually set to 1k, 4k: and 32k tokens for the three tasks. We set the training steps to 300k, 300k, and 400k per stage for the De-En, En-Ro, and En-De tasks. For other KD baselines, set the training steps to 200k for each student, which are much longer than the aver- age training steps for each student in AIO-KD. We set the number IC of sample student as 2. The selec- tions syper-parameters o,, 3, and are discussed in Section 4.2.
We first investigate the effects of a. B, and 1), where and / are used to balance Led and Liss (See Equations 7 and 8), and controls which students'
gradients are not utilized to update the teacher's parameters during knowledge transfer (See Sec- tion 3.3).
Through our preliminary empirical studies, we find that changes in v have negligible impacts on the selection of a and B. Therefore, we tune a and 3 without dynamic gradient detaching, where is not involved. Concretely, first stage and then tune 3 at the second stage with fixed. Finally, we tune 1) after determining the optimal and
tune at the
Model Selection. In our work, we expect that all candidate students achieve satisfactory perfor- mance. However, it is impractical for each student to conduct beam search decoding on the validation set for model selection.
Figure shows the effects of a, B, and " on the En-Ro task, where we set (a,3, n) to (5.5, 0.5, 1.1). Similarly, we apply above pro- cedures to the De-En and En-De tasks, where (a, 3, n) are separately set to (5.5, 0.5, 1.1) and (4.5, 0.1, 1.01).
As efficient alternative, select the model according to the average cross-entropy losses of all candidate students on the validation set, which can be formulated as
C
4.3 Main Results
0" - arg min
To demonstrate the superiority of AIO-KD, we port translation quality all candidate students, well as training costs.
where 4* denotes the optimal parameters, which are essentially the teacher's parameters because all candidate students share them.
Translation Quality. Table presents BLEU scores of all candidate students on the three tasks.
Baselines. We compare our model with the fol- lowing baselines:
We can draw the following conclusions:
First of all, we observe that both Word-KD and Selective-Ko KD achieve remarkable improvements compared with Transformer, echoing the results reported in previous studies (Kim and Rush, 2016; Wang et al., 2021). Second, AI0-KD significantly outperforms other baselines across all tasks, indi- cating its effectiveness. Furthermore, the teachers in AIO-KD demonstrate impressive performance.
Transformer (Vaswani et al., 2017). It is the most dominant NMT model.
Word-KD (Kim and Rush. 2016). Under Word- KD, the student is optimized mimic the output probability distributions of the teacher.
Selective-KD (Wang et al., 2021). By using Selective-Kol KD, the student is optimized to mimic

---

Model
2-2
3-2
4-2
5-2
6-2
3-3
4-3
5-3
4-4
5-4
6-4
5-5
6-5
6-6
Avg.
IWSL114
De-En
Transformer
33.48
34.27
34.56
34.58
34.98
34.73
34.65
34.70
35.29
31.81
34.66
35.03
31.48
34.94
35.03
34.08
Word-KD
34.93
35.36
35.68
35.76
36.03
36.05
35.94
36.01
36.24
36,20
36.05
36.15
36.14
36.27
36.49
35.95
Selective-KD
34.69
34.98
35.32
35.30
35.82
35.33
35.46
35.77
35.72
35.56
35.76
36.18
18
35.61
36.21
35.76
35.56
AIO-KD (Ours)
34.64
35.99
36.62
36.73
36.85
36.43
37.15
37.22
37.25
37.25
37.36
37.45
37.53
37.45
37.69
36.915
WMT16
En-Ro
Transformer
30.30
30.97
31,20
30,99
31.06
31.31
30.86
31.77
31.71
31.49
31.92
31.31
31.41
31.72
32.01
31.34
Word-KD
33.48
34.23
33.67
33.42
34.07
33.96
34.10
33.94
34.22
34.07
34.3
34.0 69
34.30
34.25
34. 36
34.07
Selective-KD
33.00
32.51
32.51
32.48
32.73
33.18
33.18
33.00
32. 81
32.72
32.74
32.85
32.42
32.68
32. 50
32.76
AIO-KD (Ours)
33.97
34,50
34.84
34.80
34.77
34.87
35.17
35.15
35.11
35,29
35.40
35.32
35.39
35.47
35.44
35.038
WMT14 En-De
Transformer
26.08
26.25
27.05
27.56
27
26.14
27.12
27.41
27.66
27.39
27.
27.94
27.88
28.17
27:98
27.31
Word-KI)
26.01
26.63
63
27.
27.58
27.64
26
27
66
27.85
27 04
27.43
27 80
27.72
77.72
28
13
28 05.01
28
27.51
ScDective-KD
26.44
26.95
27.38
27.65
27 79
27
29
27.76
28.07
28 28.31
27.39
28.
28.45
28
51
28,
28.20
27.78
AIO-KD (Ours)
25.50
27.48
28.16
28.30
28.57
27.65
28.45
28.75
28.79
28.68
28.86
29.23
28.96
29.16
29.18
28.40:
Table 1: Comparisons of BLEU (%) scores of all candidate students. "Ig-la" refers to the candidate student with le encoder and la decoder layers. "Avg." refers to the average BLEU (%) score of all candidate students. The best results are highlighted in bold. We combine the translations from all candidate students for significance test (Koehn, 2004). where "4" means the improvements over Word- KD and Selective-KD KD are statistically significant with p<0.01.
Model
2-2
3-2
4-2
5-2
6-2
3.3
4.3
5.3
6-3
4-4
5-
6-4
5.5
6-5
6-6
Avg.
IWSLT14 De-En
Transformer
33.38
35.94
36.52
37.65
38.45
37.60
69
37.89
38.69
69 39.75
38.72
38.16
39.14
38.60
39.50
40.30
38.03
Word-kd
39.23
40.15
41.66
41.89
42.67
42.38
42.30
12
43.99
12.78
43.12
43.45
43.34
43.79
43.73
42.51
Selective-KD
38.90
40.15
41.25
41.35
42.75
11.11
41.70
42.69
43.02
41.52
42.31
43.24
42.98
43.80
42.54
41.96
AIO-KD
37.64
41.58
43.72
44.54
44.51
43.31
45,33
45.81
46.11
45.96
46.37
46.68
46.67
46.69
46.85
44.783
WMT16
En-Ro
Transformer
32.12
34.12
33.87
33.54
34.32
35.41
33.61
37.35
37.0
36.92
38.02
36.56
37.32
38.41
40.54
35.98
Word-kd
46.84
48.01
47.20
47.21
48.99
48. 42
50.90
50.11
50.40
50.14
51.67
32.83
50.17
50.92
50.95
49.65
Selective-KD
41.50
40.68
42.49
42.29
41.51
41.60
42.
43.46
45.54
45.43
44.18
43.82
43.21
43.49
44.78
43.09
AIO-KD
47.92
51.74
52.83
52.67
53.12
52.47
53,99
54.34
54.23
54.69
54.87
$4.94
55.08
55.17
55.31
53.562
WMT14
En-De
Transformer
36.30
39.46
41.10
42.75
43.47
41.10
42.24
44.27
44.71
43.82
44.87
45.79
45.36
46.97
46.45
43.24
Word-kd
38.68
42.13
43.80
45.02
43.21
43.91
45.02
45.75
45.82
46.42
46.72
46.72
45.65
47.45
46.15
44.96
Selective-KD
40.36
43.94
43.69
46.13
46.13
45.36
36
46.53
47.89
47.71
46.83
47.62
48.16
48.69
48.70
48.44
46.41
AIO-KD
37.20
43.37
45.58
46.61
46.83
45.32
47.86
48.88
47.93
47.93
48.75
48.95
48.97
49.08
49.46
46.85
Table 2: Comparisons of COMET (%) scores of all candidate students.
Specifically, the teachers achieve 37.69, 35.44, and 29.18 BLEU scores on the De-En, En-Ro, and En- De tasks, respectively, with improvements of +2.66, +3.43, and +1.20 BLEU scores over Transformer. We attribute the promising improvements of the teachers the interactions with the students, which will be explored in Section 5.3.
By contrast, AIO-KD is much more eco-friendly. Concretely, GPU hours of AIO-KD spent on train- ing are comparable to those of Transformer on the De-En (29.83 vs. 26.11) and En-Ro (34.83 vs. 33.06) tasks, and much less than those of Word- KD (218.67 vs. 456.67) and Selective-KD (218.67 vs. 406.67) on the En-De task. More encour- agingly, AIO-KD also demonstrates memory- friendliness compared with Transformer on the De- En (16.70 70 vs. 74.72), En-Ro (55.85 vs. 169.66), and En-De (123.67 vs. 221.22) tasks. Ultimately, AIO-KD saves only one model, i.e., the teacher, highlighting its storage-efficient.
Also, we report the results using COMET metric in Table 2, which support the above conclusions.
Training Costs. Apart from the satisfactory per- formance, the advantages of AIO-KD also owe to its training efficiency. To support our claim, we report the training time and memory usage of each approach, as displayed in Table 3.
4.4 Ablation Studies
First of all, we observe that although adopting both Word-KD and Selective-KI significantly im- proves model performance, it also brings enormous training costs.
To better investigate the effectiveness of the carefully-designel strategies, we compare AIO-KD with the following variants shown in Table 4: 1) wo DGD. In this variant, we discard the

---

Training time (GPU hours) / Memory usage (GB)
Model
IWSLT14 De-En
WMT16 En-Ro
WMT14 En-De
Transformer
26.11 / 74.72
33.06 / 169.66
114.44/ / 221.22
Word-KD
72.22 75.777
86.11/ 159.31
456.67 468.87
Selective-KD
67.22 / 80.65
81.67 250.86
106.67 / 193.33
AIO-KD (Ours)
29.83/ / 16.70
34.83 / 55.85
218.67 123.67
Table Comparisons of training time (GPU hours) and memory usage (GB). We sum up the training costs of all students for Transformer, Word-KD KD, and Selective-KI KD. The presented results are recorded on NVIDIA A100 GPUs.
Model
2-2
3-2
4-2
5-2
3-3
4-3
5-3
6-3
4-4
5-5
6-5
Avg.
AIO-KD
25,80
27.48
28.16
28.30
28.57
27.65
28.45
28.75
28.79
28.68
28.86
29.23
28.96
29.16
29.18
28.40
w/o DGD
24.96
26.70
27.56
27.81
27.70
27.11
27.82
27.91
28.04
27. 87.86
28.17
28.17
28.21
28.33
28.
27.64
w/o ML.
25.60
27.32
27.92
28.20
28.45
27.59
28.26
28.45
28.54
28.55
28.77
28.97
28.74
29.10
29.08
28.24
w/o TST
25.63
27.23
28.04
28.16
28.47
27.49
28.43
28.49
28.68
28.40
28.54
28.
80
28.59
28.70
28.87
28.17
Table 4:
Ablation studies on the En-De task. mutual learning. and "TST" denotes two-stage training.
"DGD" denotes dynamic gradient detaching, "ML" denotes
37.0
65.0
Avg. BLEF
GPU hours
R= 3
Avg. BLEU (%)
36.0
55.0
Avg. Lows
4.6
R=4 4
5
GPU
35.0
35.0
4,4
35,83
34,35
34,37
34,0
20.0
10
15
20
25
30
Number of Sample Student
Training Step (X10K)
Figure 3: Left: Effects of sample student number 1C on two metries: the average BLEU (%) scores of all candidate students on the test set, and GPU hours spent on training. Right: Loss curves of AIO-KD on the validation set during training. The results the En-Ro task reported.
Dynamic Gradient Detaching strategy. can be found that this removal leads to a significant de- generation of the teacher, with BLEU score drop- ping from 29.18 8 to 28.25. Morcover, other candi- date students suffer from varying degrees of perfor- mance degradation. These results support our claim that poorly-performe students harm the teacher, which further negatively affects other students.
Analysis
5.1 Effect of Sample Student Number
In previous experiments, set the number I. of sample students 2. question arises naturally: does increasing K. further improve the students? To answer this question, we experiment with K. ranging from 2 to 5, as illustrated in Figure 3.
2) w/o ML. To verify the benefits of interactions between students, we remove the Mutual Learning loss Lmt from Equation The results show that BLEU scores of all candidate students decrease, indicating that mutual learning indeed further pro- motes students.
In the left half of Figure 3, with an increase of K. the training time of AIO-KD grows from 33.33 to 59.17 GPU hours but does not lead to any performance improvements. Instead, the students degenerate. The right half of Figure 3 also displays loss curves on the validation set with different K,, showing that increasing I. leads to the over- wver-fitting problem of the students.
3) w/o TST. We employ a one-stage training strategy to train this variant, with the same total training steps as the original AIO-KD, The loss function of this variant is Le defined in Equation 8. As expected, AIO-KD benefits from the Two-Stage Training strategy across all candidate students, in- dicating that poorly-performed students at the early- stage training have negative impacts on mutual learning.
Regarding the above phenomenon, we attribute it to the gradient conftict problem (Yu et al., 2020; Liu et al., 2021a; Chai et al.. 2022; Yue al., 2023). Since different students share the parameters of the teacher, when increases, the conflict of their gradients becomes more severe during training, ul- timately leading to the decline in performance.

---

Model
2-2
3-2
4-2
5-2
6-2
3-3
5-3
6-3
4-4
5-4
6-4
5-5
6-5
6-6
Avg.
AIO-KD
33.97
34.50
34.84
34.50
34.77
34.87
35.17
35.15
35
35
35.40
35.32
35.30
35.47
35.44
35.03
32.30
32.46
32.52
32.69
32.69
33.13
33.11
32.94
33.10
33.08
33.25
33.41
32.90
SDq-KD
It is evident that AIO-KD performs significantly better than Seq-KDD Moreover, when AIO-KD and Seq-KD are combined, i.e., AIO-KD+Seq-KD. achieves an average BLEU score of 35.27, surpass- ing both AIO-KD (35.27 vs. 35.03) and Seq-KD (35.27 vs. 32.90). These results confirm that AIO- KD and Seq-KD are compatible with each other. 11 (.188)
(-1.67)
(-2.04)
(-2.32)
(-2.11)
(-2.26)
(-2 48)
(-1.77)
98
(-2.88)
(-2.46)
(-2.22)
(-2,31)
(-2.22)
(-2.03)
(-2.13)
ATO-KD
34.34
31.99
31.95
35.12
35.10
35.14
35.38
35.42
35.33
35.40
35.60
35.56
35.55
35.55
35.56
35.27
+Seq-KD (+0.37) (40.49) (+0.11) (+0.32) (+0.33)
(+0.27)
(+0.21) (+0.27) (+0.22)
(+0.11)
(+0.20) (+0.23)
(+0.16) (+0.08) (+0.12) (+0.24)
Table 5: BLEU (%) scores of AIO-KD, Seq-KD, and AI0-KD4Seq-Kil on the En-Ro task. "AIO-KD+Seq-KDD
means that AIO-KD conducted the data provided by Seq-KD. The values in parentheses denote the performance gaps compared with AIO-KD.
IWSLT14
WMT16 En-Ro
WMT14
Model
De-En
En-De
Transformer
35.01
32.01
27.98
SegMixt (Guo et al., 2020a)
36.20
28.10
CutOfft (Shen et al., 2020)
37.60
29.10
PD-R= (Guo al., 2022)
34.93
AdMixt (Jin et al., 2022)
37.10
28.26
CipherDAug? (Kambhatla et al., 2022)
37.60
27.90
AIO-KD (Ours)
37.69
35.44
29.18
Table 6: BLEU (%) scores of the teachers in AIO-KD, "f" means that the results corresponding to the method are taken from the original papers.
36.0
6.0
we DGD
BLEU of Teacher
Detaching Step
Detaching Step (x 1000))
35.27
BLEU (%)
4,04
1.1
5-3
35.0
3.0
n
34.59
1.05
3-1
6-2
34.5
1.5
34.51
1,005
4-2
34.0
0.0
2-2
1.005
1.05
DGD
100
Proportion (%) Different Detaching Students
Figure 4: Left: Effects of metrics: the number of detaching steps, and BLEU (%) score of the teacher. Right: Proportion (%) of different detaching students at detaching steps. Dynamic gradient detaching varies 1) to protect teacher from some poorly-pes perrforme students at the training steps, where these training steps and students referred detaching steps and detaching students, respectively. "w/o DGD" means that remove dynamic gradient detaching from AIO-KD, with no students being detached.
5.2 Compatibility of AIO-KD with Seq-KD
the enhanced teacher with recently-proposed works NMT and observe that our model outperforms these strong baselines.
Seq-KD (Kim and Rush, 2016) also a widely- used KD approach, where student is trained on the teacher-gen genbrrree data. Hence, we explore whether AIO-KD and KD complement to each other. The results are shown in Table 5.
Furthermore, explore the enhanced teacher from the perspective of model interaction. Under AIO-KD, the teacher is optimized not only to align with labels but also to interact with the students via knowledge transfer. Therefore, we speculate that the teacher's improvements come from these interactions.
To gain deeper insights, we delve into the effects of in dynamic gradient detaching, as illustrated in Figure 4. By adjusting 7). the gradients of the KD loss specific some students at the training steps are not utilized to update the teacher's pa- rameters, where we refer to these training steps and students as detaching steps and detaching students, respectively.
5.3 Win-Win Knowledge Distillation
As discussed Section 4.3, AIO-KD enhances both the teacher and students, making a win-win KD technique. As shown in Table 6, we compare

---

In the left half of Figure 4, we observe that when 7 decreases, the number of detaching steps gradu- ally increases. During this process, the teacher's performance experiences an initial improvement, however, it subsequently undergoes a decline. The above observations reveal that the significant im- pacts the teacher through the KD loss.
same model architecture the teacher. Besides, despite achieving impressive efficiency and per- formance, our work is only conducted based on Transformer. Thus, we plan to validate AIO-KD on more model architectures future work.
Ethics Statement
In the right half of Figure 4, we further present the proportion of different detaching students at detaching steps corresponding to different 1). We find that when r/11.1 and r=1.05, most of poorly- performed students are detached, thus positively impacting the teacher, which severally achieves 35.44 and 35.27 BLEU scores. Conversely, when we set to 1. 1.005, more well-performed students are detached, resulting in negative impact on the teacher, which achieves 31.51 BLEU score. The above results validate that the teacher benefits from interacting with well-performed students while suf- fering from the interactions with perffo-pe ones.
This work aims to explore an eco-friendly KD ap- proach for NMT, and we hope method can inspire future work. Our work does not involve any data collection. In the experiments, all the datasets are publicly available and commonly used in the NMT community. Besides, we develop AIO- KD and other baselines based on a widely- open-source tool fairseq (Ot et al., 2019). The comparisons in this work are conducted based on the same experimental settings and datasets.
A Acknowledgements
The project supported by National Natural Science Foundation of China (Nos. 62036004, 62276219, 6207621 1), Natural Science Foundation of Fujian Province China (No. 2020J06001), and University-Industry Cooperation Programs of Fu- jian Province of China (No. 2023H6001), We also thank the reviewers for their insightful comments.
Overall, our analyses suggest that the teacher can derive benefits from the weaker students, offering valuable insights for future research.
6 Conclusion
In this work, present AIO-KD, novel KD framework for NMT that constructs various can- didate students from the teacher itself. With AIO- KD, we jointly optimize the teacher and the sample students from scratch. During this process, the students learn from the teacher and interact with other students via mutual learning, resulting in mul- tiple satisfactory students with efficient training. strefully-deiinnee strategies are also introduced accommodate AIO-KD. Extensive experiments and in-depth analyses on three benchmarks demon- strate the superiority of our AIO-KD.
References
Tianchi Bi, Hao Xiong, Zhongjun He, Hua Wu, and Haifeng Wang. 2019. Multi-agent learning for neural machine translation. In Proc. of EMNLP.
Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han, 2020. Once-for-all: Train one net- work and specialize it for efficient deployment. In Proc. of ICLR.
Heyan Chai, Zhe Yin, Ye Ding, Li Liu, Binxing Fang. and Qing Liao. 2022. A model-agnostic approach to mitigate gradient interference for multi-task learning. IEEE Trans. on Cybern.
In the future, we plan to explore more compact subnetworks of teacher as students using parameter pruming methods. Additionally, we aim to extend ADO-KD to large language models (LLMs) to vali- date its generalizability.
Tejalal Choudhary, Vipul Kumar Mishra, Anurag Goswami, and Jagannathan Sarangapani. 2020. A comprehensive survey on model compression and acceleration. Artif. Intell. Rev.
Limitations
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Lukasz Kaiser. 2019. Universal transformers. In Proc. of ICLR.
As mentioned above, the students in AIO-KD derived from the teacher and they share parame- ters. Such design yields multiple high-quality students with significantly reduced training costs, compared with conventional KD approaches. How- ever, its limitation is that the students possess the
Heejin Do and Gary Geunbae Lee. 2022. Target-
oriented knowledge distillation with language-family- based grouping for multilingual nmt. ACM Journal of the ACM (JACM).

---

Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. 2020. Depth-adaptive transformer. In Proc. of ICI.R
Chang Jin, Shigui Qiu, Nini Xiao, and Hao Jia. 2022. Admix: A mixed sample data augmentation method for neural machine translation. Proc. of IJCAI.
Angela Fan, Edouard Grave, and Armand Joulin. 2020. Reducing transformer depth on demand with struc- tured dropout. In Proc. of ICLR.
Nishant Kambhatla, Logan Born, and Anoop Sarkar. 2022. Cipherdaug: Ciphertext based data augmen- tation for neural machine translation. In Proc. of ACL.
Tao Ge, Si-Qing Chen, and Furu Wei. 2022. Edge- former: A parameter-efficient transformer for on- device seq2seq generation. In Proc. of EMNLP.
Jungo Kasai, Nikolaos Pappas, Hao Peng, James Cross, and Noah A. Smith. 2021. Deep encoder, shallow decoder: Reevaluating aon-autoregressive machine translation. In Proc. of ICLR.
Jiatao Gu, James Bradbury, Caiming Xiong, Victor O. Li, and Richard Socher. 2018. Non-autoreerresive neural machine translation. In Proc. of ICLR.
Yoon Kim and Alexander M. Rush. 2016. Sequence- level knowledge distillation. In Proc. of EMNLP.
Demi Guo, Yoon Kim, and Alexander M Rush. 2020a. Sequence-level mixed sample data augmentation. In Proc. of EMNLP.
Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In Proc. of ICLR.
Philipp Koehn. 2004. Statistical significance tests for machine translation evaluation. In Proc. of EMNLP.
Dengji Guo, Zhengrui Ma, Min Zhang, and Yang Feng. 2022. Prediction difference regularization against perturbation for neural machine translation. In Proc. of ACL.
Xiaobo Liang, Lijun Wu, Juntao Li, Tao Qin, Min Zhang, and Tie-Yan Liu. 2022. Multi-teacher distilla- tion with single model for neural machine translation. IEEF/ACM TASLP.
Qiushan Guo, Xinjiang Wang, Yichao Wu, Zhipeng Yu, Ding Liang, Xiaolin Hu, and Ping Luo. 2020b. On- line knowledge distillation via collaborative learning. In Proc. of CVPR.
Baohao Liao, Yingbo Gao, and Hermann Ney. 2020. Multi-agent mutual learning at sentence-level and token-level for neural machine translation. In Proc. of EMNLP Findings.
Song Han, Huizi Mao, and William Dally. 2016. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman cod- ing. In Proc. of ICLR.
Bo Liu, Xingehao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. 2021a. Conflict-avesse gradient descent for multi-task learning. In Proc. of NeurIPS.
Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. 2015. Distilling the knowledge in neural network. Arxiv.
Weijie Liu, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. 2020a. Fasthert: self- distilling BERT with adaptive inference time. Proc. of ACL.
Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. 2020. Dynabert: Dynamic BERT with adaptive width and depth. In Proc. of NeurIPS.
Xin Liu, Kai Liu, Xiang Li, Jinsong Su, Yubin Ge, Bin Wang, and Jiebo Luo. 2020b. An iterative multi- source mutual knowledge transfer framework for ma- chine reading comprehension. In Proc. of IJCAI.
Chenyang Huang, Hao Zhou, Osmar Zaïane, Lili Mou, and Lei Li. 2022a. Non-autoreeressive translation with layer-wise prediction and deep supervision. In Porc. of AAAI.
Yijin Liu, Fandong Meng. Jie Zhou, Yufeng Chen, and Jinan Xu. 2021b. Faster depth-adaptive transformers. In Proc. of AAAI.
Yichong Huang. Xiaocheng Feng, Xinwei Geng, Bao- hang Li, and Bing Qin. 2023. Towards higher pareto frontier in multilingual machine translation. Proc. of ACL.
Ziyao Lu, Xiang Li, Yang Liu, Chulun Zhou, Jianwei
Cui, Bin Wang. Min Zhang, and Jinsong Su. 2022. Exploring multi-stage information interactions for multi-source neural machine translation. TEEE/ACM TASLP.
Yichong Huang, Xiaocheng Feng, Xinwei Geng, and Bing Qin. 2022b. Unifying the convergences in mul- tilingual neural machine translation. In Proc. of EMNLP.
Zhongjian Miao, Xiang Li, Liyan Kang, Wen Zhang, Chulun Zhou, Yidong Chen, Bin Wang, Min Zhang, and Jinsong Su. 2022. Towards robust neural ma- chine translation with iterative scheduled data-switch training. In Proc. of COLING.
Aref Jafari, Mehdi Rezagholizadeh, Pranav Sharma, and Ali Ghodsi. 2021. Annealing knowledge distillation. In Proc. of EACL.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. 2020. Tinybert: Distilling BERT for natural language un- derstanding. In Proc. of EMNLP Findings.
Xuan-Phi Nguyen, Shafiq Joty, Thanh-Tung Nguyen, Kui Wu, and Ai Ti Aw. 2021. Cross-model back- translated distillation for unsupervised machine trans- lation, In Proc. of ICML.

---

Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. 2019. fairseq: A fast, extensible toolkit for sequence modeling. In Proc. of NAACL.
Hongyu Wang, Shuming Ma, Dong, Shaohan Huang. Dongdong Zhang, and Furu Wei. 2022. Deepnet: Scaling transformers to I, 000 layers. Arxiv.
Qiang Wang, Bei Li. Tong Xiao, Jingbo Zhu, Changliang Li, Derek F. Wong, and Lidia S. Chao. 2019. Learning deep transformer models for machine translation. In Proc. of ACL.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei- Jing Zhu. 2002. Bleu: method for automatic evalu- ation of machine translation. In Proc. of ACL.
Lihua Qian, Hao Zhou, Yu Bao, Mingxuan Wang. Lin Qiu, Weinan Zhang, Yong Yu, and Lei Li. 2021. Glancing transformer for non-autoregressive neural machine translation. In Proc. of ACL.
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. 2020b. Minilm: Deep self- attention distillation task-agnostic compression of pre-trained transformers. In Proc. of NewrIPS.
Ricardo Rei, Craig Stewart, Ana C. Farinha, and Alon Lavie. 2020. COMET: A neural framework for MT evaluation. In Proc. of EMNLP.
Zhihao Wang, Longyue Wang. Jinsong Su. Junfeng Yao, and Zhaopeng Tu. 2023. Revisiting non- autoregressive translation scale. In Proc. of ACL Findings.
Mark Sandler, Andrew Howard, Menglong Zhu, An- drey Zhmoginov, and Liang-Chine Chen. 2018. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proc. of CVPR.
Noreen M Webb. 1989, Peer interaction and learning in small groups International journal of Educational research (IJER).
Abigail See, Minh-Thang Luong, and Christopher D Manning. 2016. Compression of neural machine translation models via pruning. In Proc. of CoNLL.
Hao-Ran Wei, Shujian Huang, Ran Wang, Xin- Yu Dai, and Jiajun Chen. 2019. Online distilling from check- points for neural machine translation. In Proc. of ACL.
Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016. Neural machine translation of rare words with subword units. Proc. ACL.
Ji Xin, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. 2020. Deebert: Dynamic early exiting for accelerating BERT inference. In Proc. ACL.
Dinghan Shen, Mingzii Zheng, Yelong Shen, Yanru Qu, and Weizhu Chen. 2020. A simple but tough-to- beat data augmentation approach for natural language understanding and generation. Arxiv.
Peng Xu, Dhruv Kumar, Wei Yang, Wenjie Zi, Keyi Tang, Chenyang Huang, Jackie Chi Kit Cheung, Si- mon J. D. Prince, and Yanshuai Cao. 2021. Optimiz- ing deeper transformers on small datasets. Proc. of ACL.
Haipeng Sun, Rui Wang, Kehai Chen, Masao Utiyama, Eicchiro Sumita, and Tiejun Zhao. 2020. Knowledge distillation for multilingual unsupervised neural ma- chine translation. In Proc. of ACL.
Zhixian Yang, Renliang Sun, and Xiaojun Wan. 2022. Nearest neighbor knowledge distillation for neural machine translation. In Proc. of NAACL.
Xin Sun, Tao Ge, Furu Wei, and Houfeng Wang. 2021. Instantaneous grammatical enror correction with shal- low aggressive decoding. In Proc. of ACL.
Xu Tan, Yi Ren, Di He, Tao Qin, Zhou Zhao, and Tie- Yan Liu. 2019. Multilingual neural machine transla- tion with knowledge distillation. In Proc. of ICLR.
Jiahui Yu, Linjie Yang. Ning Xu, Jianchao Yang, and Thomas S. Huang. 2019. Slimmable neural networks. In Proc. of ICLR.
Zhixing Tan, Zeyuan Yang, Meng Zhang, Qun Liu, Maosong Sun, and Yang Liu. 2022. Dynamic multi- branch layers for on-device neural machine transla- tion. IEEE/ACM TASLP.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. 2020. Gradient surgery for multi-task learning. In Proc. of NeurIPS.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proc. of NeurIPS.
Zhixiong Yue, Yu Zhang, and Jie Liang. 2023. Learning conflict-noticed architecture for multi-task learning. In Proc. of AAAI.
Fusheng Wang, Jianhao Yan, Fandong Meng, and Jie Zhou. 2021. Selective knowledge distillation for neural machine translation. In Proc. of ACL.
Jiali Zeng, Yang Liu, Jinsong Su, Yubing Ge, Yaojie Lu, Yongjing Yin, and Jiebo Luo. 2019. Iterative dual domain adaptation for neural machine translation. In Proc. of EMNLP.
Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han. 2020a. Hat: Hardware-awan transformers for efficient natural language processing. In Proc, of ACL.
Biao Zhang, Deyi Xiong, Jinsong Su, and Jiebo Luo. 2019. Future-aware knowledge distillation for neural machine translation. JEEF/ACM TASLP.

---

Liang Zhang, Jinsong Su, Zijun Min, Zhongjian Miao, Qingguo Hu, Biao Fu, X. Shi, and Yidong Chen. 2023. Exploring self-distillation based relational rea- soning training for document-level relation extraction. In Proc. of AAAI
Ying Zhang, Tao Xiang, Timothy M Hospedales, and Huchuan Lu. 2018. Deep mutual learning. In Proc. of CVPR.
Jiawei Zhao, Wei Luo, Boxing Chen, and Andrew Gilman. 2021. Mutual-learning improves end-to-end speech translation. In Proc. of EMNLP.
Chulun Zhou, Yunlong Liang, Fandong Meng. Jie Zhou, Jinan Xu, Hongji Wang, Min Zhang, and Jinsong Su. 2022a. multi-task multi-stage transitional training framework for neural chat translation. IEEE Transac- tions on Pattern Analysis and Machine Intelligence.
Chulun Zhou, Fandong Meng. Jie Zhou, Min Zhang, Hongji Wang, and Jinsong Su. 2022b. Confidence based bidirectional global context aware training framework neural machine translation. In Proc. ACL.
Wangchunshu Zhou, Canwen Xu, and Julian J. McAuley. 2022c. BERT learns teach: Knowledge distillation with meta learning. In Proc. of ACL.


Reviews:
[{"review_id": "DwtjPBzprR", "reviewer": "Reviewer_gH6v", "summary": "", "questions": "[Table 2] Is the training time of Transformer, Word-KD and Selective-KD the sum of the training time of each student (2-2 to 6-6)?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{1. [Reply to Q1 in Reasons To Reject]}$ We compare our AIO-KD with Word-KD, Seq-KD, and Selective-KD because they are widely used in neural machine translation (NMT). As discussed in our paper, our main contribution is proposing an efficient multi-student KD framework that is compatible with various KD techniques, including recent ones. In future, we will attempt to incorporate more KD techniques into our framework and validate it on various NLP tasks beyond NMT.\n\n$\\textbf{2. [Reply to Q1 in Question For The Authors]}$ Thank you for the question. Yes, we sum up the training time of all candidate students as the training time of Transformer, Word-KD, and Selective-KD. Since each baseline obtains one student at a time, we determine its training costs by adding up the training costs of all candidate students. Overall, our AIO-KD is much more efficient, as it only needs to be executed once to obtain multiple high-quality students.  In the revision, we will clarify these points in the caption of Table 2.", "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 proposed framework seems sound.\n2. The experimental analysis given is sufficient and can support their claims.\n3. The paper is well-organized and easy to understand.", "reasons_to_reject": "The compared baselines, Word-KD, Seq-KD and Selective-KD are not the latest. More recent works should be compared. ", "questions_to_authors": "[Table 2] Is the training time of Transformer, Word-KD and Selective-KD the sum of the training time of each student (2-2 to 6-6)?", "paper_topic_and_main_contributions": "Conventional knowledge distillation (KD) approaches have to be conducted multiple times to train student models of varying sizes that fit different devices. To address this issue, this paper proposes a framework based on the Transformer to train various student models at the same time and consider the correlation among them as well. To further improve the KD performance, two strategies, dynamic gradient detaching and two-stage mutual learning, are also proposed to mitigate the negative impact caused by poorly-performed students.", "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": {}, "extra_text": {}}, {"review_id": "jOsvXtn7ui", "reviewer": "Reviewer_TCJa", "summary": "", "questions": "A. In line 229, why 200k training steps for student models obtained using other distillation algorithms are considered “much longer than the average training steps for each student in AIO-KD”? The student models in AIO-KD are obtained with the same number of training steps as the teacher model. It does not seem to make sense to take the average. Does this different number of training steps make the comparison of GPU hours to evaluate computational costs unfair?\n\nB. What is the student model architecture (for Word-KD and Selective-KD) used to obtain training time and memory usage statistics in Table 2?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{1. [Reply to Q1 in Reasons to Reject]}$ Thanks for your kind suggestions.  Due to page limitations, the introduction to neural machine translation is omitted. We will supplement related preliminaries in the later revision.\n\n$\\textbf{2. [Reply to Q2 in Reasons to Reject]}$ We will check the paper again in light of your feedback and supplement more conclusions in the analysis. Regarding hyper-parameter selection, it is time-consuming to determine 3 hyper-parameters in AIO-KD using grid search. Instead, we turn to adopting a more efficient heuristic method based on our preliminary empirical results (See Section 4.2).  In fact, our AIO-KD is relatively insensitive to these hyper-parameters. We will supplement the impacts of hyper-parameters of all translation tasks in the later revision to support our claim. \n\n$\\textbf{3. [Reply to Q1 in Questions For The Authors]}$ There may be some misunderstandings about our work. At each training step, the teacher is always involved, while only $\\mathcal{K}$ students are sampled from |$\\mathcal{C\n}$| candidate students for training. This means that each candidate student does not obtain the same number of training steps as the teacher. Since the sampling procedure is random and independent, the expected number of training steps for each candidate student in our AIO-KD can be represented as $\\mathcal{T}$ $\\times$  $\\mathcal{K}$ $/$ |$\\mathcal{C\n}$|, where $\\mathcal{T}$ refers to the total number of training steps including the first and second stages.  For the De-En and En-Ro tasks, the expected number of training steps per student is (300K+300K) $\\times$ 2/15 = 80K steps, which is less than 200K steps, and for the En-De task, it is  (400K+400K) $\\times$ 2/15 $\\approx$ 107K steps, which is also less than 200K steps. These results support our claim in Section 4.1 that the average number of training steps for each student in our AIO-KD is much shorter than that of the student in other KD baselines. \n\n$\\textbf{4. [Reply to Q2 in Questions For The Authors]}$ Following previous work [1,2], both the students and teacher in our work adopt the dominant Transformer architecture, with the only difference being that the students have fewer encoder and decoder layers compared with the teacher.\n\n$\\textbf{References}$\n\n[1] Wang et al. Selective Knowledge Distillation for Neural Machine Translation. 2021. In Proc. of ACL.\n\n[2] Liang et al. Multi-Teacher Distillation with Single Model for Neural Machine Translation. 2022. IEEE/ACM Trans. Audio, Speech, Language Process.", "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 engineering part of the paper seems to be well thought out and neatly executed. The empirical results demonstrate the effectiveness of AIO-KD.\n\n2. Most analyses of the proposed algorithm are interesting and insightful. The authors have conducted extensive experiments to understand the influence of a wide range of hyper-parameters or variables in AIO-KD (e.g., the number of student models, and ablation studies for each component in the overall loss function). These analyses are helpful for us to understand the advantages and limitations of the proposed method.", "reasons_to_reject": "1. The paper has a very strong assumption that a reader is familiar with the task of neural machine translation, and hence, some preliminaries like input representations, model outputs, precise loss functions, etc. are omitted. This raises a barrier for a wider range of audience to understand the paper well.\n\n2. While the paper presents a lot of interesting analyses, some are left without a thorough explanation or a clear conclusion. For example, the trends of some hyper-parameters seem rather random in Sec 4.2. Is this randomness a result of statistical noises alone? How does Fig. 2 teach us a strategy for hyper-parameter selection?\n\n3. A few technical details could be further clarified. Otherwise, some claims may not be well supported. Kindly refer to the \"Questions for the Authors\" section.", "questions_to_authors": "A. In line 229, why 200k training steps for student models obtained using other distillation algorithms are considered “much longer than the average training steps for each student in AIO-KD”? The student models in AIO-KD are obtained with the same number of training steps as the teacher model. It does not seem to make sense to take the average. Does this different number of training steps make the comparison of GPU hours to evaluate computational costs unfair?\n\nB. What is the student model architecture (for Word-KD and Selective-KD) used to obtain training time and memory usage statistics in Table 2?", "paper_topic_and_main_contributions": "The paper proposes a distillation framework where a few student models are trained simultaneously with the teacher model in order to obtain a series of smaller models at the end of knowledge distillation. Each of the student models is a sub-part of the teacher model (i.e. they share the exactly same weight parameters for the sub-part and are stored in the same place in memory). In the proposed method AIO-KD, knowledge distillation does not only take place between the teacher model and each student model, but also between the student models via mutual learning. \n\nThe authors present a few challenges when the algorithm attempts to train the teacher model and student models simultaneously, such as the potential bad influence of immature student models in the early training stage. The paper then proposes corresponding solutions that dynamically detach gradients and employ mutual learning only during the second stage in a two-stage training scheme.", "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": {}, "extra_text": {}}, {"review_id": "D9wpY4AoeL", "reviewer": "Reviewer_s6c5", "summary": "", "questions": "- Am I reading this correctly that increasing the number of students *degrades* performance? ", "limitations": "", "rating": 3, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{1. [Reply to Q1 and Q2 in Reasons to Reject]}$ It is worth noting that our contribution lies in proposing an efficient knowledge distillation framework for multi-student scenarios. The innovation of our framework far exceeds the simple combination of knowledge distillation (KD) and mutual learning (ML), as our framework is not limited to any specific KD or ML techniques.\n\n$\\textbf{2. [Reply to Q3 in Reasons to Reject]}$ We appreciate your valuable suggestions. In our early exploration, we also attempt to explore intelligent sampling strategies for students. For instance, we sort all the candidate students according to the ratio of  the student's parameter count to that of the teacher, and then gradually sample the candidate students in descending/increasing order of this ratio. After exploring some sampling strategies, we find that they do not consistently perform well for our AIO-KD. Hence, we default to sampling students uniformly at random during training. We will clarify the student sampling strategy in the revision and explore more effective student sampling strategies in future work.\n\n$\\textbf{3. [Reply to Q1 in Questions For The Authors]}$ Thank you for the question. Yes, increasing the number of sampled students degrades performance. In Section 5.2, we investigate the impact of the sample student number $\\mathcal{K}$, i.e., the number of the students sampled at each training step. In the left half of Figure 3, we vary $\\mathcal{K}$ from 2 to 5 and observe that increasing $\\mathcal{K}$ does not lead to further performance improvements but instead results in a significant increase in training costs.", "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": "- This paper combines online KD and Mutual Learning and adds some nice tricks that seem to improve empirical results compared to Seq-KD and incorporating both simultaneously brings additional minor gains.\n- Ablation studies are sound and appreciated the analysis of varying the number of students.\n", "reasons_to_reject": "- Online distillation (where the teacher and student are jointly optimized on some loss Lce) has been extensively explored in the past (see section 5.3 in https://arxiv.org/pdf/2006.05525.pdf)\n- Mutual learning (wherein students learn from each-other) has also been explored.\n- The authors should have also considered more intelligent ways to sample from the subnetworks of a larger teacher model rather than just randomly sampling encoder and decoder layers.", "questions_to_authors": "- Am I reading this correctly that increasing the number of students *degrades* performance? ", "paper_topic_and_main_contributions": "The authors propose All-in-One Knowledge Distillation (AIO-KD) which (1) randomly extract subnetworks from the teacher as students (2) jointly optimize the teacher and students. They leverage some tricks like stopping training loss from being affected by poorly performing students. ", "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": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper provides a framework for jointly distilling a teacher to multiple student networks. It presents a few challenges eg the potential bad influence of immature student models in the early training stage. The paper then proposes to dynamically detach gradients and employ mutual learning only during the second stage in a two-stage training scheme. The empirical results demonstrate the effectiveness of AIO-KD. The analyses of the proposed algorithm are interesting and insightful.
Decision: Accept-Main