Hint-based Training for Non-Autoregressive Translation

1 INTRODUCTION . Neural machine translation has attracted much attention from the research community ( Bahdanau et al. , 2014 ; 2016 ; Kalchbrenner et al. , 2016 ; Gehring et al. , 2016 ) and has been gradually adopted by industry in the past several years ( Wu et al. , 2016 ) . Despite the huge variety of model architectures ( Bahdanau et al. , 2014 ; Gehring et al. , 2017 ; Vaswani et al. , 2017 ) , given a source sentence x = ( x1 , ... , xTx ) and a target sentence y = ( y1 , ... , yTy ) , most neural machine translation models decompose and estimate the conditional probability P ( y|x ) in an universal autoregressive manner : P ( y|x ) = ΠTyt=1P ( yt|y < t , x ) , ( 1 ) where y < t represents the first t − 1 words of y . During inference , given an input sentence , those models generate the translation results sequentially , token by token from left to right . We call all such models AutoRegressive neural machine Translation ( ART ) models . A state-of-the-art ART model , Transformer ( Vaswani et al. , 2017 ) , is shown in the left part of Figure 1 . A well-known limitation of the ART models is that the inference process can hardly be parallelized , and the inference time is linear with respect to the length of the target sequence . As a result , the ART models suffer from long inference time ( Wu et al. , 2016 ) , which is sometimes unaffordable for industrial applications . Consequently , people start to develop Non-AutoRegressive neural machine Translation ( NART ) models to speed up the inference process ( Gu et al. , 2017 ; Kaiser et al. , 2018 ; Lee et al. , 2018 ) . These models use the general encoder-decoder framework : the encoder takes a source sentence x as input and generates a set of contextual embeddings and predicted length Ty ; conditioned on the contextual embeddings , the decoder takes a transformed copy of x as input and predicts the target tokens at all the positions independently in parallel according to the following decomposition : P ( y|x , Ty ) = Π Ty t=1P ( yt|Ty , x ) . ( 2 ) While the NART models achieve significant speedup during inference ( Gu et al. , 2017 ) , their accuracy is considerably lower than their ART counterpart . Most of the previous works attribute the poor performance to this unavoidable conditional independence assumption of the NART model . To tackle this issue , they try to improve the expressiveness and accuracy of the decoder input in different ways : Gu et al . ( 2017 ) introduce fertilities from statistical machine translation models into the NART models , Lee et al . ( 2018 ) base the decoding process of their proposed model on an iterative refinement process , and Kaiser et al . ( 2018 ) take a step further to embed an autoregressive submodule that consists of discrete latent variables into their model . Although such methods provide better expressiveness of decoder inputs and improve the final translation accuracy , the inference speed of these models will be hurt due to the overhead of the introduced modules , which contradicts with the original purpose of introducing the NART models , i.e. , to parallelize and speed up neural machine translation models . Different from previous works that develop new submodules for decoder input , we improve the translation model from another perspective . We aim to provide more guided signals during optimization . That is , we do not introduce any new prediction submodule but introduce better regularization . The reason we tackle the problem from this perspective lies in two points : First , the encoder input ( source words ) contains all semantic information for translation , and the decoder input in the NART model can be considered as a middle layer between input and output . It is not clear how much gain can be achieved by developing a sophisticated submodule for a middle layer in a deep neural network . Second , the encoder-decoder-based NART model is already over-parameterized . We believe that such neural network still has great ability and space to be better optimized if we can provide it with stronger and richer signals , for example , from a much better ART model : Once we have a well-trained ART model , we actually know rich information about the contexts to make the prediction at each time step and the natural word alignments between bilingual sentences . All the information could be invaluable towards the improved training of a NART model . To well leverage an ART model , we use the hint-based training framework ( Romero et al. , 2014 ; Chen et al. , 2017 ) , in which the information from hidden layers of teacher model ( referred as hints ) are used to guide the training process of a student model . However , hint-based training was developed for image classification models and it is challenging to define and use hints for translation . First , the translation model is composed of stacked encoder layers , attention layers , and stacked decoder layers . It is not clear how to define hints in such an encoder-decoder framework . Second , the NART and ART models are of different architectures on the decoding stage . It is not obvious how to leverage hints from the teacher to the training of student with a different architecture . We find that directly applying hints used in the classification tasks fails . In this paper , we first investigate the causes of the bad performance of the NART model , and then define hints targeting to solve the problems . According to our empirical study , we find that the hidden states of the NART model differ from the ART model : the positions where the NART model outputs incoherent tokens will have very high hidden states similarity . Also , the attention distributions of the NART model are more ambiguous than those of ART model . Based on these observations , 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 . We have conducted experiments on the widely used WMT14 English-to-German/German-toEnglish ( En-De/De-En ) task and IWSLT14 German-to-English task . For WMT14 En-De task , our proposed method achieves a BLEU score of 25.20 which significantly outperforms the nonautoregressive baseline models and is even comparable to a strong ART baseline , Google ’ s LSTMbased translation model ( 24.60 Wu et al . ( 2016 ) ) . For WMT14 De-En task , we also achieve significant performance gains , reaching 29.52 in terms of BLEU . 2 RELATED WORKS . 2.1 AUTOREGRESSIVE TRANSLATION . Given a sentence x = ( x1 , . . . , xTx ) from the source language , the straight-forward way for translation is to generate the words in the target language y = ( y1 , . . . , yTy ) one by one from left to right . This is also known as the autoregressive factorization in which the joint probability is decomposed into a chain of conditional probabilities , as in the Eqn . ( 1 ) . Deep neural networks are widely used to model such conditional probabilities based on the encoder-decoder framework . The encoder takes the source tokens ( x1 , . . . , xTx ) as input and encodes x into a set of context states c = ( c1 , . . . , cTx ) . The decoder takes c and subsequence y < t as input and estimates P ( yt|y < t , c ) according to some parametric function . There are many design choices in the encoder-decoder framework based on different types of layers , e.g. , recurrent neural network ( RNN ) -based ( Bahdanau et al. , 2014 ) , convolution neural network ( CNN ) -based ( Gehring et al. , 2017 ) and recent self-attention based ( Vaswani et al. , 2017 ) approaches . We show a self-attention based network ( Transformer ) in the left part of Figure 1 . 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 ground truth pair ( x , y ) is exposed to the model , and thus the prediction at different positions can be estimated in parallel based on CNN or self-attention networks . However , during inference , given a source sentence x , the decoder has to generate tokens sequentially , as the decoder inputs y < t must be inferred on the fly . Such autoregressive behavior becomes the bottleneck of the computational time ( Wu et al. , 2016 ) . 2.2 NON-AUTOREGRESSIVE TRANSLATION . In order to speed up the inference process , a line of works begin to develop non-autoregressive translation models . These models follow the encoder-decoder framework and inherit the encoder structure from the autoregressive models . After generating the context states c by the encoder , a separate module will be used to predict the target sentence length Ty and decoder inputs z = ( z1 , . . . , zTy ) by a parametric function : ( Ty , z ) ∼ fz ( x , c ; θ ) , which is either deterministic or stochastic . The decoder will then predict y based on following probabilistic decomposition P ( y|x , Ty , z ) = Π Ty t=1P ( yt|z , c ) . ( 3 ) Different configurations of Ty and z enable the decoder to produce different target sentence y given the same input sentence x , which increases the output diversity of the translation models . Previous works mainly pay attention to different design choices of fz . Gu et al . ( 2017 ) introduce fertilities , corresponding to the number of target tokens occupied by each of the source tokens , and use a non-uniform copy of encoder inputs as z according to the fertility of each input token . The prediction of fertilities is done by a separated neural network-based module . Lee et al . ( 2018 ) define z by a sequence of generated target sentences y ( 0 ) , . . . , y ( L ) , where each y ( i ) is a refinement of y ( i−1 ) . Kaiser et al . ( 2018 ) use a sequence of autoregressively generated discrete latent variables as inputs of the decoder . While the expressiveness of z improved by different kinds of design choices , the computational overhead of z will hurt the inference speed of the NART models . Comparing to the more than 15× speed up in Gu et al . ( 2017 ) , which uses a relatively simpler design choice of z , the speedup of Kaiser et al . ( 2018 ) is reduced to about 5× , and the speedup of Lee et al . ( 2018 ) is reduced to about 2× . This contradicts with the design goal of the NART models : to parallelize and speed up neural machine translation models . 3 HINT-BASED TRAINING FOR NON-AUTOREGRESSIVE TRANSLATION . In this section , we introduce the proposed hint-based training algorithm that leverages a well-trained ART model to train the NART model . Our model mostly follows Transformer ( Vaswani et al. , 2017 ) , with an additional positional attention layer proposed by Gu et al . ( 2017 ) , as shown in the right part of Figure 1 . To avoid overhead , we use simple linear combinations of source token embeddings as z , which has no learnable parameters . Details about the model can be found in the appendix . We first describe the observations we find about the ART and NART models , and then discuss what kinds of information can be used as hints and how to use them to help the training of the NART model .

Summary: This work proposes a non-autoregressive Neural Machine Translation model which the authors call NART, as opposed to an autoregressive model which is referred to as an ART model. The main idea behind this work is to leverage a well trained ART model to inform the hidden states and the word alignment of NART models. The joint distribution of the targets y given the inputs x, is factorized into two components as in previous works on non-autoregressive MT: an intermediate z which is first predicted from x, which captures the autoregressive part, while the prediction of y given z is non-autoregressive. This is the approach taken e.g., in Gu et al, Kaiser et al, Roy et al., and this also seems to be the approach of this work.  The authors argue that improving the expressiveness of z (as was done in Kaiser et al, Roy et al), is expensive and so the authors propose a simple formulation for z. In particular, z is a sequence of the same length as the targets, where the j^{th} entry z_j is a weighted sum of the embedding of the inputs x (the weights depend in a deterministic fashion on j) . Given this z, the model predicts the targets completely non-autoregressively. However, this by itself is not entirely sufficient, and so the authors also utilize "hints": 1) If the pairwise cosine similarity between two successive hidden states in the student  NART model is above a certain threshold, while the similarity is lower than another threshold in the ART model, then the NART model incurs a cost proportional to this similarity 2) A KL term is used to encourage the distribution of attention weights of the student ART model to match that of the teacher NART model. These two loss terms are used in different proportions (using additional hyperparameters) together with maximizing the likelihood term.