Hint-Based Training for Non-Autoregressive Machine Translation

Abstract: Due to the unparallelizable nature of the autoregressive factorization, AutoRegressive Translation (ART) models have to generate tokens sequentially during decoding and thus suffer from high inference latency. Non-AutoRegressive Translation (NART) models were proposed to reduce the inference time, but could only achieve inferior translation accuracy. In this paper, we proposed a novel approach to leveraging the hints from hidden states and word alignments to help the training of NART models. The results achieve significant improvement over previous NART models for the WMT14 En-De and De-En datasets and are even comparable to a strong LSTM-based ART baseline but one order of magnitude faster in inference.

Neural machine translation has attracted much attention in recent years BIBREF0, BIBREF1, BIBREF2, BIBREF3. Given a sentence $x=(x_1, \dots ,x_{T_x})$ from the source language, the straight-forward way for translation is to generate the target words $y=(y_1, \dots , y_{T_y})$ one by one from left to right. This is also known as the AutoRegressive Translation (ART) models, in which the joint probability is decomposed into a chain of conditional probabilities:
While the ART models have achieved great success in terms of translation quality, the time consumption during inference is still far away from satisfactory. During training, the predictions at different positions can be estimated in parallel since the ground truth pair $(x,y)$ is exposed to the model. However, during inference, the model has to generate tokens sequentially as $y_{<t}$ must be inferred on the fly. Such autoregressive behavior becomes the bottleneck of the computational time BIBREF4.
In order to speed up the inference process, a line of works begin to develop non-autoregressive translation models. These models break the autoregressive dependency by decomposing the joint probability with
The lost of autoregressive dependency largely hurt the consistency of the output sentences, increase the difficulty in the learning process and thus lead to a low quality translation. Previous works mainly focus on adding different components into the NART model to improve the expressiveness of the network structure to overcome the loss of autoregressive dependency BIBREF5, BIBREF6, BIBREF7. However, the computational overhead of new components will hurt the inference speed, contradicting with the goal of the NART models: to parallelize and speed up neural machine translation models.
To tackle this, we proposed a novel hint-based method for NART model training. We first investigate the causes of the poor performance of the NART model. Comparing with the ART model, we find that: (1) the positions where the NART model outputs incoherent tokens will have very high hidden states similarity; (2) the attention distributions of the NART model are more ambiguous than those of ART model. Therefore, we design two kinds of hints from the hidden states and attention distributions of the ART model to help the training of the NART model. The experimental results show that our model achieves significant improvement over the NART baseline models and is even comparable to a strong ART baseline in BIBREF4.
In this section, we first describe the observations on the ART and NART models, and then discuss what kinds of information can be used as hints to help the training of the NART model. We follow the network structure in BIBREF8, use a copy of the source sentence as decoder input, remove the attention masks in decoder self-attention layers and add a positional attention layer as suggested in BIBREF5. We provide a visualization of ART and NART models we used in Figure FIGREF11 and a detailed description of the model structure in Appendix.
According to the case study in BIBREF5, the translations of the NART models contain incoherent phrases (e.g. repetitive words) and miss meaningful tokens on the source side, while these patterns do not commonly appear in ART models. After some empirical study, we find two non-obvious facts that lead to this phenomenon.
First, we visualize the cosine similarities between decoder hidden states of a certain layer in both ART and NART models for sampled cases. Mathematically, for a set of hidden states $r_1, \ldots , r_T$, the pairwise cosine similarity can be derived by $\cos _{ij} = {\left<r_i, r_j\right>}/{(\Vert r_i\Vert \cdot \Vert r_j\Vert )}.$ We then plot the heatmap of the resulting matrix $\cos $. A typical example is shown in Figure FIGREF4, where the cosine similarities in the NART model are larger than those of the ART model, indicating that the hidden states across positions in the NART model are “similar”. Positions with highly-correlated hidden states tend to generate the same word and make the NART model output repetitive tokens, e.g., the yellow area on the top-left of Figure FIGREF4(b), while this does not happen in the ART model (Figure FIGREF4(a)). According to our statistics, 70% of the cosine similarities between hidden states in the ART model are less than 0.25, and 95% are less than 0.5.
Second, we visualize the encoder-decoder attentions for sampled cases, shown in Figure FIGREF6. Good attentions between the source and target sentences are usually considered to lead to accurate translation while poor ones may cause wrong output tokens BIBREF0. In Figure FIGREF6(b), the attentions of the ART model almost covers all source tokens, while the attentions of the NART model do not cover “farm” but with two “morning”. This directly makes the translation result worse in the NART model. These phenomena inspire us to use the intermediate hidden information in the ART model to guide the learning process of the NART model.
Our study motivates us to leverage the intermediate hidden information from an ART model to improve the NART model. We focus on how to define hints from a well-trained ART teacher model and use it to guide the training process of a NART student model. We study layer-to-layer hints and assume both the teacher and student models have an $M$-layer encoder and an $N$-layer decoder, despite the difference in stacked components.
Without the loss of generality, we discuss our method on a given paired sentence $(x,y)$. In real experiments, losses are averaged over all training data. For the teacher model, we use $a_{t,l,h}^\mathit {tr}$ as the encoder-to-decoder attention distribution of $h$-th head in the $l$-th decoder layer at position $t$, and use $r_{t,l}^\mathit {tr}$ as the output of the $l$-th decoder layer after feed forward network at position $t$. Correspondingly, $a_{t,l,h}^\mathit {st}$ and $r_{t,l}^\mathit {st}$ are used for the student model. We propose a hint-based training framework that contains two kinds of hints:
The discrepancy of hidden states motivates us to use hidden states of the ART model as a hint for the learning process of the NART model. One straight-forward method is to regularize the $L_1$ or $L_2$ distance between each pair of hidden states in ART and NART models. However, since the network components are quite different in ART and NART models, applying the straight-forward regression on hidden states hurts the learning process and fails. Therefore, we design a more implicit loss to help the student refrain from the incoherent translation results by acting towards the teacher in the hidden-state level:
where $d_\mathit {st} = \cos (r_{s, l}^\mathit {st},r_{t, l}^\mathit {st})$, $d_\mathit {tr} = \cos (r_{s, l}^\mathit {tr},r_{t, l}^\mathit {tr})$, and $\phi $ is a penalty function. In particular, we let
where $-1\le \gamma _\mathit {st}, \gamma _\mathit {tr}\le 1$ are two thresholds controlling whether to penalize or not. We design this loss since we only want to penalize hidden states that are highly similar in the NART model, but not similar in the ART model. We have tested several choices of $-\log (1-d_\mathit {st})$, e.g., $\exp (d_\mathit {st})$, from which we find similar experimental results.
We observe that meaningful words in the source sentence are sometimes untranslated by the NART model, and the corresponding positions often suffer from ambiguous attention distributions. Therefore, we use the word alignment information from the ART model to help the training of the NART model.
In particular, we minimize KL-divergence between the per-head encoder-to-decoder attention distributions of the teacher and the student to encourage the student to have similar word alignments to the teacher model, i.e.
Our final training loss $\mathcal {L}$ is a weighted sum of two parts stated above and the negative log-likelihood loss $\mathcal {L}_\mathit {nll}$ defined on bilingual sentence pair $(x, y)$, i.e.
where $\lambda $ and $\mu $ are hyperparameters controlling the weight of different loss terms.
The evaluation is on two widely used public machine translation datasets: IWSLT14 German-to-English (De-En) BIBREF9, BIBREF1 and WMT14 English-to-German (En-De) dataset BIBREF4, BIBREF10. To compare with previous works, we also reverse WMT14 English-to-German dataset and obtain WMT14 German-to-English dataset.
We pretrain Transformer BIBREF8 as the teacher model on each dataset, which achieves 33.26/27.30/31.29 in terms of BLEU BIBREF11 in IWSLT14 De-En, WMT14 En-De and De-En test sets. The student model shares the same number of layers in encoder/decoder, size of hidden states/embeddings and number of heads as the teacher models (Figure FIGREF11). Following BIBREF5, BIBREF12, we replace the target sentences by the decoded output of the teacher models.
Hyperparameters ($\gamma _\mathit {st}, \gamma _\mathit {tr}, \lambda , \mu $) for hint-based learning are determined to make the scales of three loss components similar after initialization. We also employ label smoothing of value $\epsilon _\mathit {ls}=0.1$ BIBREF13 in all experiments. We use Adam optimizer and follow the setting in BIBREF8. Models for WMT14/IWSLT14 tasks are trained on 8/1 NVIDIA M40 GPUs respectively. The model is based on the open-sourced tensor2tensor BIBREF14. More settings can be found in Appendix.
During training, $T_y$ does not need to be predicted as the target sentence is given. During testing, we have to predict the length of the target sentence for each source sentence. In many languages, the length of the target sentence can be roughly estimated from the length of the source sentence. We choose a simple method to avoid the computational overhead, which uses input length to determine target sentence length: $T_y = T_x + C$, where $C$ is a constant bias determined by the average length differences between the source and target training sentences. We can also predict the target length ranging from $[(T_x+C)-B, (T_x+C)+B]$, where $B$ is the halfwidth. By doing this, we can obtain multiple translation results with different lengths. Note that we choose this method only to show the effectiveness of our proposed method and a more advanced length estimation method can be used to further improve the performance.
Once we have multiple translation results, we additionally use our ART teacher model to evaluate each result and select the one that achieves the highest probability. As the evaluation is fully parallelizable (since it is identical to the parallel training of the ART model), this rescoring operation will not hurt the non-autoregressive property of the NART model.
We compare our model with several baselines, including three ART models, the fertility based (FT) NART model BIBREF5, the deterministic iterative refinement based (IR) NART model BIBREF6, and the Latent Transformer BIBREF7 which is not fully non-autoregressive by incorporating an autoregressive sub-module in the NART model architecture.
The results are shown in the Table TABREF15. Across different datasets, our method achieves significant improvements over previous non-autoregressive models. Specifically, our method outperforms fertility based NART model with 6.54/7.11 BLEU score improvements on WMT En-De and De-En tasks in similar settings and achieves comparable results with state-of-the-art LSTM-based model on WMT En-De task. Furthermore, our model achieves a speedup of 30.2 (output a single sentence) or 17.8 (teacher rescoring) times over the ART counterparts. Note that our speedups significantly outperform all previous works, because of our lighter design of the NART model: without any computationally expensive module trying to improve the expressiveness.
We also visualize the hidden state cosine similarities and attention distributions for the NART model with hint-based training, as shown in Figure FIGREF4(c) and FIGREF6(c). With hints from hidden states, the hidden states similarities of the NART model decrease in general, and especially for the positions where the original NART model outputs incoherent phrases. The attention distribution of the NART model after hint-based training is more similar to the ART teacher model and less ambiguous comparing to the NART model without hints.
According to our empirical analysis, the percentage of repetitive words drops from 8.3% to 6.5% by our proposed methods on the IWSLT14 De-En test set, which is a 20%+ reduction. This shows that our proposed method effectively improve the quality of the translation outputs. We also provide several case studies in Appendix.
Finally, we conduct an ablation study on IWSLT14 De-En task. As shown in Table TABREF18, the hints from word alignments provide an improvement of about 1.6 BLEU points, and the hints from hidden states improve the results by about 0.8 BLEU points. We also test these models on a subsampled set whose source sentence lengths are at least 40. Our model outperforms the baseline model by more than 3 BLEU points (20.63 v.s. 17.48).
In this paper, we proposed to use hints from a well-trained ART model to enhance the training of NART models. Our results on WMT14 En-De and De-En significantly outperform previous NART baselines, with one order of magnitude faster in inference than ART models. In the future, we will focus on designing new architectures and training methods for NART models to achieve comparable accuracy as ART models.
This work is supported by National Key R&D Program of China (2018YFB1402600), NSFC (61573026) and BJNSF (L172037) and a grant from Microsoft Research Asia. We would like to thank the anonymous reviewers for their valuable comments on our paper.

Figure and table captions:
Figure 1: Case study: the above three figures visualize the hidden state cosine similarities of different models. The axes correspond to the generated target tokens. Each pixel shows the cosine similarities cosij between the last layer hidden states of the i-th and j-th generated tokens, where the diagonal pixel will always be 1.0.
Figure 2: Case study: the above three figures visualize the encoder-decoder attention weights of different models. The x-axis and y-axis correspond to the source and generated target tokens respectively. The attention distribution is from a single head of the third layer encoder-decoder attention, which is the most informative one according to our observation. Each pixel shows attention weights αij between the i-th source token and j-th target token.
Figure 3: Hint-based training from ART model to NART model.
Table 1: Performance on WMT14 En-De, De-En and IWSLT14 De-En tasks. “/” means non-reportable.
Table 2: Ablation studies on IWSLT14 De-En. Results are BLEU scores without teacher rescoring.
Table 3: Cases on IWSLT14 De-En.

Questions about this paper:
How slow is the unparallelizable ART model in the first place?  
What metric is used to measure translation accuracy?

Answers:
784 miliseconds
BLUE and the percentage of repetitive words