Towards Understanding and Improving Knowledge Distillation
for Neural Machine Translation

Given a source sentence with $M$ tokens $\mathbf{x}=\{x_{1},x_{2},\ldots,x_{M}\}$ and the corresponding target sentence with $N$ tokens $\mathbf{y}=\{y_{1},y_{2},\ldots,y_{N}\}$, NMT models are trained to maximize the probability of each target token conditioning on the source sentence by the cross-entropy (CE) loss:

where $y_{j}^{\ast}$ and $\mathbf{y}_{<j}$ denote the ground-truth target and the target-side previous context at time step $j$, respectively. And $\theta$ is the model parameter.

Word-level KD Kim and Rush (2016) aims to minimize the KL divergence between the output distributions of the teacher model and the student model on each target token. Formally, given the probability distribution $q(\cdot)$ from the teacher model, the KL divergence-based loss is formulated as follows:

where $\theta_{t}$ and $\theta_{s}$ denote the model parameters of the teacher and the student, respectively.

Then, the overall loss function of word-level KD is the linear interpolation between the CE loss and the KL divergence loss:

Sequence-level KD Kim and Rush (2016) encourages the student model to imitate the sequence probabilities of the translations from the teacher model. To this end, it optimizes the student model through the following approximation:

where $\mathcal{Y}$ denotes the hypothesis space of the teacher and $\widehat{\mathbf{y}}$ is the approximate result through the teacher’s beam search.

In word-level KD, the relative probabilities between negative candidates in the soft targets from the teacher contain rich correlation information, which is previously deemed to carry knowledge from the teacher Hinton et al. (2015); Tang et al. (2020); Jafari et al. (2021). However, in practice, strong teachers usually have high confidence in their top-1 predictions while retaining little probability mass for other candidates. Hence, to study the mystery of KD, it is necessary to first investigate the real effects of the correlation information and the top-1 prediction information during KD and then figure out which one actually determines the performance of KD.

To this end, during word-level KD, we separately remove the top-1 information and the correlation information from the original soft targets of the teacher (as depicted in Fig.1) and then observe the corresponding performance. Besides the BLEU score, we also introduce a new metric, namely the Top-1 Agreement (TA) rate, which calculates the overlap rate of the top-1 predictions between the student and the teacher on each position under the teacher-forcing mode. As shown in Tab.2, the performance slightly increases when we remove the probabilities of all other candidates except for the top-1 ones in soft targets (see Fig.1(b))⁴⁴4Considering the regularization effect, we do not add a uniform distribution to complement the removed probability. Please refer to Appendix B for more detailed explanations.. However, when we only remove the top-1 information and keep the remaining correlation information (see Fig.1(c))⁵⁵5Note that we do not simply remove the probability of the top-1 prediction, but add this probability to the ground-truth token to maintain the correctness of the distribution, i.e., the soft target is unchanged if its top-1 prediction is correct., the performance of KD drops close to the baseline without any KD. Moreover, we observe that the TA rates are well correlated with the final BLEU scores among these students. Therefore, we conjecture that the top-1 information is the one that actually determines the performance of word-level KD (answer to Q1).

To further confirm the above conjecture, we examine whether the student models have successfully learned the correlation information of the teacher during KD. To achieve this, we design two metrics to measure the ranking similarities between token rankings from the student and the teacher, named top-$k$ edit distance and top-$k$ ranking distance.

As the importance of the top-1 information for transferring knowledge in word-level KD has been validated, we further investigate whether more knowledge can be exploited by extending top-1 information to top-$k$ information⁸⁸8Equivalent to vanilla word-level KD when $k=|V|$.. Similar to Fig.1(b), we keep the top-$k$ probabilities in the original soft target and remove others to extract its top-$k$ information. However, the results in Tab.4 give a negative answer that more information does not bring significantly more knowledge. Thus, we can believe that almost all the knowledge of the teacher in word-level KD comes from the teacher’s top-1 information, even though the whole distribution is distilled to the student.

Although the previous results have coarsely located the knowledge in word-level KD on the top-1 information of the teacher, it is still not clear whether this holds for all types of soft targets, especially when the teacher is under-confident in its top-1 predictions. Towards this end, we divide the soft targets of the teacher into three intervals Wang et al. (2021) based on their top-1 probabilities: $(0.0,0.4]$, $(0.4,0.7]$, and $(0.7,1.0)$. Then we separately conduct the same KD processes as described in Fig.1, only using the soft targets in one of these intervals. Surprisingly, the results in Fig.2 show that even when the teacher is not so confident (i.e., $q_{\rm max}\leq 0.7$) in its top-1 predictions, using only the top-1 information (i.e., the blue bars) still achieves better performance than using the full information in corresponding soft targets. However, in these cases, removing the top-1 information in soft targets largely degrades the performance of the students. We conjecture that these under-confident top-1 predictions of the teacher can serve as hints for students to learn the difficult ground-truth labels, while the correlation information in these cases carries more noise than real knowledge for students.

Inspired by the analyses on word-level KD, we move on to sequence-level KD and decompose its loss function in Eq.(2.3) into a word-level form:

where $\widehat{y}_{j}$ is the teacher-decoded target for students at time step $j$. Considering the similar word-level form, it is intuitive to speculate that the top-1 information may also matter in sequence-level KD. To verify this, we divide the targets $\widehat{y}_{j}$ into the top-1 and the non-top-1 predictions of the teacher⁹⁹9There are about 70% top-1 predictions and 30% non-top-1 predictions selected by the teacher’s beam search during decoding. and investigate the respective effects of these targets by separately using them during sequence-level KD. As shown in Tab.5, there is only a negligible performance change when we only use the top-1 targets for KD (row 1 vs. row 2). However, if we only use the non-top-1 targets, the BLEU score drastically drops (row 1 vs. row 3). Moreover, considering the different proportions of the two kinds of targets in the teacher’s translations (i.e.,70% vs. 30%), we also use a fixed part (the same amount as the non-top-1 targets) of the top-1 targets for a fair comparison, and the performance is still steady (row 2 vs. row 4) and much better than using only the non-top-1 targets (row 3 vs. row 4). Interestingly, by adding additional word-level top-1 information to the non-top-1 part, the performance of sequence-level KD further improves (row 1 vs. row 5). Therefore, we can also confirm the importance of the top-1 information in sequence-level KD.

Through the above analyses, we verify the importance of the teacher’s top-1 information on both KD techniques, which actually reflects a potential connection between them. A brief theoretical analysis on this connection is provided in Appendix A. In short, the two kinds of techniques share a unified objective that imparts the teachers’ top-1 predictions to student models at each time step. Thus, we believe that they are well connected on their similar working mechanisms (answer to Q2).

Further, we revisit word-level KD from this perspective and find two inherent issues. Firstly, the KL divergence-based objective in vanilla word-level KD directly optimizes whole distributions of students, while lacking specialized learning of the most important top-1 information. Secondly, since the top-1 predictions of the teacher mostly overlap with the ground-truth targets, the knowledge from the teacher is largely covered by the ground-truth information, which largely limits the potential of word-level KD. Therefore, we claim that the performance of the current word-level KD approach is far from perfect and the solutions to these problems are urgently needed.

To help student models better grasp the top-1 information during distillation, we design a new loss named hierarchical ranking loss. To gently achieve this goal, we first encourage the student to rank the teacher’s top-$k$ predictions as its own top-$k$ predictions and then rank the teacher’s top-1 prediction over these top-$k$ predictions. Formally, given the student’s top-$k$ predictions as $[y_{j}^{s_{1}},...,y_{j}^{s_{k}}]$ and the teacher’s top-$k$ predictions as $[y_{j}^{t_{1}},...,y_{j}^{t_{k}}]$, the hierarchical ranking loss $\mathcal{L}_{hr}$ can be expressed as:

where $p(\cdot)$ and $q(\cdot)$ are the probabilities from the student model and the teacher model, respectively. And $\mathbbm{1}\{\cdot\}$ is an indicator function.

In this way, the student model can be enforced to rank the top-1 predictions of the teacher to its own top-1 places, and thus it can explicitly enhance the learning of the knowledge from the teacher. Then, we add this loss to the original KL divergence loss, i.e., Eq.(2), forming a new loss for KD:

Given that the large overlap between the top-1 predictions and ground-truth targets limits the amount of additional knowledge from the teacher during word-level KD, introducing data without ground-truth targets for KD could be helpful to mitigate this issue. Inspired by previous studies on decoder-side data manipulation Zhang et al. (2019); Goodman et al. (2020); Liu et al. (2021a, b); Xie et al. (2021), we design an iterative knowledge distillation procedure to expose more target-side data for KD.

Specifically, as shown in Algorithm 1, at each training step, we conduct KD for $N$ iterations (line 3), by using the predictions of the student in the current iteration as the decoder-side inputs for KD in the next iteration (line 8). Generally, these predictions can be regarded as similar but new inputs compared to the original target inputs. Meanwhile, there is no ideal ground-truth target for these inputs since they are usually not well-formed sentences. Then during each iteration, we collect the loss of KD according to Eq.(7) (lines 4$\sim$7) and average it across all the iterations (line 10). Since all the supervision signals are from the teacher after the first iteration, the knowledge of the teacher model will be more exploited during the following iterations and thus the potential of word-level KD can be more released.

We conduct experiments on three commonly-used WMT tasks, i.e., the WMT’14 English to German (En-De), WMT’14 English to French (En-Fr) and WMT’16 English to Romanian (En-Ro). For all these tasks, we share the source and the target vocabulary and segment words into subwords using byte pair encoding (BPE) Sennrich et al. (2016) with 32k merge operations. More statistics of the datasets can be found in Appendix C.1.

All our experiments are conducted based on the open-source toolkit fairseq Ott et al. (2019) with FP16 training Ott et al. (2018). By default, we follow the big/base setting Vaswani et al. (2017) to implement the teacher/student models in our experiments. More training and evaluation details can be referred to Appendix C.2. For word-level KD-based methods, we set the $\alpha$ in Eq.(3) to 0.5 following Kim and Rush (2016). For our method, we set top-$k$ in Sec.4.1 to 5 and iteration time $N$ in Sec.4.2 to 3 on all three tasks. The selection of top-$k$ and $N$ are shown in Appendix D.

We compare our proposed method with existing KD techniques in NMT (the detailed description of these compared techniques can be referred to Appendix C.3) on three WMT tasks. To make the results more convincing, we report both BLEU and COMET Rei et al. (2020) scores in Tab.6. Using Transformer_big as the teacher, our method can boost the Transformer_base students by +1.04/+0.60/+1.11 BLEU scores and +4.52/+2.57/+4.80 COMET scores on three tasks, respectively. Compared to the vanilla Word-KD baseline, our method can outperform it significantly on all translation tasks, which verifies the effectiveness of our proposed solutions. Additionally, as a word-level KD method, our TIE-KD can outperform Seq-KD on all three tasks and even achieves fully competitive results with the teacher on En-Ro, which demonstrates that the potential of Word-KD can be largely released by our method.

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 word-level 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.

According to Kim and Rush (2016), word-level KD can be well combined with sequence-level KD and yields better performance. As a word-level KD approach, our TIE-KD can also theoretically be combined with sequence-level KD. We verify this on the WMT’14 En-De task and list the results in Tab.8. Like Word-KD, our TIE-KD can also achieve better performance when combined with Seq-KD and is also better than “Word-KD + Seq-KD”, indicating the superiority of our method and its high compatibility with sequence-level KD.

Among the prior literature on KD Cho and Hariharan (2019); Jin et al. (2019); Mirzadeh et al. (2020); Guo et al. (2020); Jafari et al. (2021); Qiu et al. (2022), a general consensus is that a large teacher-student capacity gap may harm the quality of KD. We also check this problem in NMT by using teachers of three model sizes. Besides the default configuration (i.e., Transformer_big) in our experiments above, we also add Transformer_base setting as the weaker teacher and Transformer${}_{deep\raisebox{0.0pt}{-}big}$ setting with 18 encoder layers and 6 decoder layers as the stronger teacher¹⁰¹⁰10To stably train a deeper Transformer, we use Admin Liu et al. (2020) in layer normalization.. We compare our method with word- and sequence-level KD under these teachers in Fig.3 and draw several conclusions:

1. (1)

   The stronger teacher can bring improvement to sequence-level KD but fails to word-level KD, where the reason may be the less additional knowledge from the stronger teacher due to its higher top-1 accuracy (68%$\rightarrow$70%).

2. (2)

   As a word-level KD method, our TIE-KD instead brings conspicuous improvement with the stronger teacher, indicating that our method can exploit more knowledge from the teacher.

3. (3)

   Under the weaker teacher, the student from our method even significantly surpasses the teacher, while other methods are largely limited by the performance of the teacher, demonstrating the high generalizability of our TIE-KD to different teacher-student capacity gaps.

The decoding process of language generation models can be regarded as a sequential decision-making process Yu et al. (2017); Arora et al. (2022). As mentioned in Sec.3.5, during decoding, beam search tends to pick the top-1 predictions of the NMT model on each beam and finally selects the most probable beam. Thus, the top-1 information (including both the top-1 word index and its corresponding probability) of the teacher model largely represents its decision on each decoding step, which is exactly what we expect the student model to learn from the teacher through KD in NMT. Therefore, the top-1 information can be seen as the embodiment of the knowledge of the teacher model in NMT tasks and should be emphatically learned by the student models.