Towards Understanding and Improving Knowledge Distillation for Neural Machine Translation (2023)

Abstract: Knowledge distillation (KD) is a promising technique for model compression in neural machine translation. However, where the knowledge hides in KD is still not clear, which may hinder the development of KD. In this work, we first unravel this mystery from an empirical perspective and show that the knowledge comes from the top-1 predictions of teachers, which also helps us build a potential connection between word- and sequence-level KD. Further, we point out two inherent issues in vanilla word-level KD based on this finding. Firstly, the current objective of KD spreads its focus to whole distributions to learn the knowledge, yet lacks special treatment on the most crucial top-1 information. Secondly, the knowledge is largely covered by the golden information due to the fact that most top-1 predictions of teachers overlap with ground-truth tokens, which further restricts the potential of KD. To address these issues, we propose a new method named Top-1 Information Enhanced Knowledge Distillation (TIE-KD). Specifically, we design a hierarchical ranking loss to enforce the learning of the top-1 information from the teacher. Additionally, we develop an iterative KD procedure to infuse more additional knowledge by distilling on the data without ground-truth targets. Experiments on WMT’14 English-German, WMT’14 English-French and WMT’16 English-Romanian demonstrate that our method can respectively boost Transformer _{base} students by +1.04, +0.60 and +1.11 BLEU scores and significantly outperforms the vanilla word-level KD baseline. Besides, our method shows higher generalizability on different teacher-student capacity gaps than existing KD techniques.

## introduction in recent years, neural machine translation (nmt) has made marvelous progress in generating highquality translations (bahdanau et al., 2014; gehring * yufeng chen is the corresponding author. @xcite @xcite , especially with some exquisite and deep model architectures @xcite @xcite . despite their amazing performance on translation tasks, high computational and deployment costs still prevent these models from being applied in real life. on this problem, knowledge distillation (kd) @xcite @xcite @xcite @xcite is regarded as a promising solution for model compression, which aims to transfer the knowledge from these strong teacher models into compact student models. generally, there are two categories of kd techniques, i.e., word-level kd @xcite @xcite and sequence-level kd @xcite . (1) word-level kd is conducted on each target token, where it shrinks the kullback-leibler (kl) divergence @xcite between the predicted distributions from the student and the soft targets from the teacher. in these soft targets, the knowledge was previously deemed to come from the probability relationship between negative candidates (i.e., the correlation information) @xcite @xcite . to answer these two questions, we conduct an empirical study that starts from word-level kd to find out where the knowledge hides in the teacher's soft targets and then explore whether the result can be expanded to sequence-level kd. as a result, we summarize several intriguing findings: these findings sufficiently prove that 1) the knowledge actually comes from the top-1 information of the teacher during kd in nmt, and 2) the two kinds of kd techniques can be connected from the perspective of the top-1 information. on these grounds, we further point out that there are two inherent issues in vanilla word-level kd. firstly, as the source of teachers' knowledge, the top-1 information receives no special treatment in the training objective of vanilla word-level kd since the kl divergence directly optimizes the entire distribution. secondly, since most top-1 predictions of strong teachers overlap with ground-truth tokens (see the first row of tab.1), the additional knowledge from teachers beyond the golden information is poor and the potential of word-level kd is largely limited (see the second row of tab.1). to address these issues, we propose a new kd method named top-1 information enhanced knowledge distillation (tie-kd) for nmt. specifically, we first design a hierarchical ranking loss that can enforce the student model to learn the top-1 information through ranking the top-1 predictions of the teacher as its own top-1 predictions. moreover, we develop an iterative kd procedure to expose more input data without ground-truth targets for kd to exploit more knowledge from the teacher. ## probing the knowledge of kd in nmt in this section, we start from word-level kd and offer exhaustive empirical analyses on 1) the determining information in word-level kd ( §3.1); 2) whether the correlation information has been learned ( §3.2); 3) whether there are more benefits when extending the top-1 to top-k information ( §3.3) and 4) the importance of the top-1 information on soft targets with different confidence ( §3.4). then we expand the conclusion to sequence-level kd ( §3.5) and lastly revisit kd for nmt from a novel view ( §3.6). ## top-1 information enhanced knowledge distillation for nmt to address the aforementioned issues in wordlevel kd, in this section, we introduce our method named top-1 information enhanced knowledge distillation (tie-kd), which includes a hierarchical ranking loss to boost the learning of the top-1 information from the teacher ( §4.1) and an iterative knowledge distillation procedure to exploit more knowledge from the teacher ( §4.2). ## analysis 6.1 ablation study to separately verify the effectiveness of our solutions for the two issues in vanilla word-level kd, we conduct an ablation study on wmt'14 en-de task and record the results in tab.7. when only adding hierarchical ranking loss to vanilla wordlevel kd, the bleu scores and the ta rates gain by +0.3/+0.22 and +0.32/+0.47 on the validation/test set, respectively. it reflects that kl divergence only provides a loose constraint on the learning of the top-1 information from the teacher, while our hierarchical ranking loss helps to explicitly grasp this core information. when only using iterative kd, the student also improves by +0.36/+0.25 bleu scores and +0.18/+0.28 ta rates. it indicates that our iterative kd can effectively release the potential of word-level kd by introducing data without ground-truth targets. when combined together, the two solutions finally compose our tie-kd and can yield further improvement on both metrics. therefore, the two issues in word-level kd are orthogonal and our proposed solutions are complementary to each other. ## conclusion in this paper, we explore where the knowledge hides in kd for nmt and unveil that it comes from the top-1 information of the teacher. this finding reflects the connection between word-and sequencelevel kd and reveals the common essence of both kd techniques in nmt. from this perspective, we further propose a top-1 information enhanced knowledge distillation (tie-kd) to address the two issues in vanilla word-level kd. experiments on three wmt tasks prove the effectiveness of our method. besides, we investigate the performance of the existing kd techniques in nmt and our method under different teacher-student capacity gaps and show the stronger generalizability of our method on various gaps.