Selective Knowledge Distillation for Neural Machine Translation

Through most of the rows, we observe noticeable gaps between the BLEU scores of the $\mathcal{S}_{High}$ and $\mathcal{S}_{Low}$, indicating there exists a clear difference of impact on medium of teacher knowledge. Specifically, for most of the criteria like cross-entropies or word frequency, the gap between two halves surpasses 0.35. In contrast, teacher $P_{golden}$ seems not useful for partitioning KD knowledge. We conjecture this is because no matter whether the teacher is convinced with the golden label or not, other soft labels could contain useful information Gou et al. (2020). Besides, we find teacher entropy is a good-enough criterion for partitioning KD data, which inlines with previous studies of dark knowledge Dong et al. (2019). Finally, we find that the KD is most sensitive (+0.64) with the Word CE criterion, which enjoys the adaptivity during the training phase and is a good representative for whether the student thinks the sample is difficult.

In conclusion, we regard the most suitable samples should have the following properties: higher Word CE, higher Sentence CE, higher Word Frequency, which probably benefits future studies of effective KD methods.

More interestingly, compared with ‘Distill-All’, which is the combination of the $\mathcal{S}_{High}$ and $\mathcal{S}_{Low}$, the $\mathcal{S}_{High}$ halves’ BLEU score even surpass the ‘Distill-All’, for Word CE, Sentence CE and Word Frequency criteria. This leads to two conclusions:

(1) Within some partitions, the $\mathcal{S}_{High}$ contributes most to the KD improvements.

(2) The amount of teacher knowledge is not the more, the better. The distillation knowledge of the $\mathcal{S}_{Low}$ does not directly combine with the $\mathcal{S}_{High}$, even hurts $\mathcal{S}_{High}$’s performance.

Given that distillation knowledge is most sensitive to Word CE, we conduct extra analysis on the Word CE. Figure 1 presents the results of varying the amount of knowledge for $\mathcal{S}_{High}$ and $\mathcal{S}_{Low}$. The consistent phenomenon is that the $\mathcal{S}_{High}$ perform significantly better than the $\mathcal{S}_{Low}$ when using the same amount of teacher’s knowledge. These results suggest that we should focus more on the $\mathcal{S}_{High}$ than on $\mathcal{S}_{Low}$. Besides, we notice that the model performance increases when we increase the knowledge in $\mathcal{S}_{High}$, but not the case for $\mathcal{S}_{Low}$. We conclude that the Word CE is distinguishable and a better indicator of teachers’ useful knowledge only for $\mathcal{S}_{High}$.

At the end of this section, we can summary the following points:

• •

  To find out the most suitable medium for transferring medium, we adopt a novel method of partition and comparison, which can easily be adopted to future studies.

• •

  The benefit of distillation knowledge drastically changes when applying to different mediums of knowledge.

• •

  Among all criteria, knowledge distillation is the most sensitive to Word CE. Distilling words with higher Word CE is more reliable than words with lower CE.

• •

  In some partitions, the distillation benefit of $\mathcal{S}_{Low}$ can not add to the $\mathcal{S}_{High}$, even hurts $\mathcal{S}_{High}$’s performance.

For WMT’14 En-De task, we use 4.5M preprocessed data, which is tokenized and split using byte pair encoded (BPE) (Sennrich et al., 2016) with 32K merge operations and a shared vocabulary for English and German. We use newstest2013 as the validation set and newstest2014 as the test set, which contain 3000 and 3003 sentences, respectively.

For the WMT’19 Zh-En task, we use 20.4M preprocessed data, which is tokenized and split using 47K/32K BPE merge operations for source and target languages. We use newstest2018 as our validation set and newstest2019 as our test set, which contain 3981 and 2000 sentences, respectively.

For evaluation, we train all the models with a maximum of 300K steps for WMT En-De’14 and WMT’19 Zh-En. We choose the model which performs the best on the validation set and report its performance on test set. We measure case sensitive BLEU calculated by multi-bleu.perl²²2https://github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/multi-bleu.perl and mteval-v13a.pl³³3https://github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/mteval-v13a.pl with significance test (Koehn, 2004) for WMT’14 En-De and WMT’19 Zh-En, respectively.

Following the setting in Vaswani et al. (2017), we carry out our experiments on standard Transformer (Vaswani et al., 2017) with the fairseq toolkit Ott et al. (2019). By default, we use Transformer (Base), which contains six stacked encoder layers and six stacked decoder layers as both teacher model and student model. To verify our approaches can be applied to a stronger teacher and student models, we further use deep Transformers with twelve encoder layers and six decoder layers. In training processing, we use Adam optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.98$, learning rate is 7e-4 and dropout is 0.1. All experiments are conducted using 4 NVIDIA P40 GPUs, where the batch size of each GPUs is set to 4096 tokens. And we accumulate the gradient of parameters and update every two steps. The average runtimes are 3 GPU days for all experiments.

There are two hyper-parameters in our experiment, i.e., distil rate $r\%$ and global queue size ${Q}_{size}$. For distil rate $r\%$, the search space is [10%, 30%, 50%, 70%, 90%]. The search result of $r\%$ is shown in Figure 2, we can find that the performance is sensitive to the value of r%. When the ratio is smaller than 50%, the increase of ratio is consistent with the BLEU score increases, and the best performance peaks at 50%. We directly apply the distil rate $r\%$ to the WMT’19 Zh-En task without extra searching. Besides, We set the $\mathcal{Q}_{size}=$ 30K for WMT’14 En-De. For larger dataset WMT’19 Zh-En, we enlarge the $Q_{size}$ to from 30K to 50K and keep word rate unchanged. The hyper-parameter search of $\mathcal{Q}_{size}$ can be found in Section 6.4.

We compare our method with several existing NMT systems (KD and others):

• •

  Word-KD (Kim and Rush, 2016). Word-KD is a standard method that distills knowledge equally for each word. The detailed description is in Section 3.2.

• •

  Seq-KD (Kim and Rush, 2016). Sequence-KD uses teacher generated outputs on training corpus as an extra source. The training loss can be formulated as:

  	$\displaystyle\mathcal{L}_{seq\_kd}=-\sum_{j=1}^{J}\sum_{k=1}^{|V|}\mathbbm{1}\{\hat{y}_{j}=k\}$
  	$\displaystyle\times\log p(y_{j}=k|\bm{\hat{y}}_{<j},\bm{x};\theta),$		(6)

  where $\bm{\hat{y}}$ denotes the sequence predicted by teacher model from running beam search, $J$ is the length of target sentence.

• •

  Bert-KD (Chen et al., 2020b). This method leverages the pre-trained Bert as teacher model to help NMT model improve machine translation quality.

• •

  Other Systems. We also include some existing methods based on Transformer(Base) for comparison, i.e., Zheng et al. (2019); So et al. (2019); Tay et al. (2020).

The results on WMT’14 En-De are shown in Table 2. In this experiment, both the teacher model and student model are Transformer (Base). We also list our implementation of word-level distillation and sequence level distillation (Kim and Rush, 2016) method.

Firstly, compared with the Transformer (Base), our re-implemented word-level and the sequence-level distillation show similar improvements with the BLEU scores up from 27.29 to 28.14 and 28.15, respectively. Secondly, compared with these already strong baseline methods, our batch-level selective approach further extends the improvement to 28.42, proving the selective strategy’s effectiveness. Thirdly, our global-level distillation achieves a 28.57 BLEU score and outperforms all previous methods, showing that the better evaluation of words’ CE distribution with FIFO global queue helps selection. It is worth noting that our strategy also significantly improves translation quality over all others methods including Word-KD. Finally, our methods show comparable/better performance than other existing NMT systems and even surpass the Transformer (Big), with much fewer parameters.

The first question is that why our methods surpass the Word-KD with more knowledge. To answer this question, we collect the statistics on the gradient difference between knowledge distillation loss and cross-entropy loss on the ground-truth label for $\mathcal{S}_{Hard}$ and $\mathcal{S}_{Easy}$.

Here, we study gradients over the output distributions, which are directly related to the model’s performance. Particularly, decoder maps target sentences $\bm{y}=(y_{1}^{*},...,y_{m}^{*})$ to their corresponding hidden representation $\bm{h}=(h_{1},...,h_{m})$. For words in target sequence, the prediction logits $l\in R^{d_{model}\times|V|}$ is given by:

where $h\in R^{d_{model}}$ is the layer output of transformer decoder, $\mathcal{W}\in R^{d_{model}\times|V|}$ is projection matrix. Then, the gradient respect to $l$ from golden cross-entropy loss can be denotes as $\nabla_{l}{\mathcal{L}_{ce}}$. The gradient from distillation loss can be denotes as $\nabla_{l}{\mathcal{L}_{kd}}$. Next, we calculate the probability that $\nabla_{l}{\mathcal{L}_{ce}}$ and $\nabla_{l}{\mathcal{L}_{kd}}$ share the same direction.

Figure 3 presents the results with the probability that gradients agree with each other during training. We observe that $\mathcal{S}_{Easy}$ (green line) is consistently lower than distillation with all words (blue line) and $\mathcal{S}_{Hard}$ (red line), which means $\mathcal{S}_{Easy}$ has more inconsistency with ground-truth. Combining with the BLEU performances, we argue this consistency leads to the risk of introducing noise and disturbs the direction of parameter updating.

Besides, the agreement of Distill-All (blue line in Fig) lies in the middle of two halves. It proves that $\mathcal{S}_{Easy}$ and $\mathcal{S}_{Hard}$ compromise with each other on some conflicts. It also proves that there exist some conflicts between the knowledge in $\mathcal{S}_{Easy}$ and $\mathcal{S}_{Hard}$.

In our approaches, we select the transferring samples from the student model’s point of view. However, in previous literature, they commonly consider knowledge from the teacher’s perspective. Hence, in this section, we study the correlation between these two perspectives.

Because previous studies commonly regard teacher’s soft-labels contain dark knowledge (Dong et al., 2019), we take the entropy of teacher’s prediction as a proxy. Concretely, we randomly select 100K tokens in the training set and calculate the entropy of distribution predicted by the teacher model for both $\mathcal{S}_{Hard}$ and $\mathcal{S}_{Easy}$. As shown in Figure 4, we notice that the $\mathcal{S}_{Easy}$’s entropy distribution is more concentrated in range (0, 4) and peaks around 1.2. In contrast, the $\mathcal{S}_{Hard}$’s entropy distribution is more spread out. The overall distribution shifts to higher entropy, which indicates $\mathcal{S}_{Hard}$ tends to provide a smoother supervision signal. Consequently, we conclude that even though our selective strategy comes from the student’s perspective, it also favors samples with abundant dark knowledge in teacher’s perspective. To some extent, this explains why the $\mathcal{S}_{Hard}$’ knowledge benefits distillation performance more.

We also conduct experiments on the larger WMT’19 Zh-en dataset (20.4M sentence pairs) to ensure our methods can provide consistent improvements across different language pairs.

As shown in Table 3, our method still significantly outperforms the Transformer (Base) with +0.89. Compared with the Word-KD, our approach consistently improves with +0.41 BLEU points. Besides, we also find that Seq-KD with our methods extends the improvement of BLEU score from 27.27 to 27.61. This indicates that our selective strategy is partially orthogonal to the improvement of Seq-KD and maintains generalizability. In summary, these results suggest that our methods can achieve consistent improvement on different sized datasets across different language pairs.

Here, we investigate how our method is well-generalized to larger models. We use a deep transformer model with twelve encoder layers and six decoder layers for our larger model experiments. As shown in Table 4, Deep Transformer (12 + 6) and Word-KD have already achieved strong performance with up to 28.90 BLEU points, and our method still outperforms these baselines (29.12 BLEU). It proves our methods’ generalizability to larger models.