Understanding Knowledge Distillation in Non-autoregressive Machine Translation

1 INTRODUCTION . Traditional neural machine translation ( NMT ) systems ( Bahdanau et al. , 2015 ; Gehring et al. , 2017 ; Vaswani et al. , 2017 ) generate sequences in an autoregressive fashion ; each target token is predicted step-by-step by conditioning on the previous generated tokens in a monotonic ( e.g . left-to-right ) order . While such autoregressive translation ( AT ) models have proven successful , the sequential dependence of decisions precludes taking full advantage of parallelism afforded by modern hardware ( e.g . GPUs ) at inference time . In contrast , non-autoregressive translation ( NAT ) models ( Gu et al. , 2018 ; Lee et al. , 2018 ) predict the whole sequence or multi-token chunks of the sequence simultaneously , alleviating this problem by trading the model ’ s capacity for decoding efficiency . Such a non-autoregressive factorization assumes that the output tokens are independent from each other . However , this assumption obviously does not hold in reality and as a result NAT models generally perform worse than standard AT models . One key ingredient in the training recipe for NAT models that is used in almost all existing works ( Gu et al . ( 2018 ) ; Lee et al . ( 2018 ) ; Stern et al . ( 2019 ) , inter alia ) is creation of training data through knowledge distillation ( Hinton et al. , 2015 ) . More precisely , sequence-level knowledge distillation ( Kim & Rush , 2016 ) – a special variant of the original approach – is applied during NAT model training by replacing the target side of training samples with the outputs from a pre-trained AT model trained on the same corpus with a roughly equal number of parameters . It is usually assumed ( Gu et al. , 2018 ) that knowledge distillation ’ s reduction of the “ modes ” ( alternative translations for an input ) in the training data is the key reason why distillation benefits NAT training . However , this intuition has not been rigorously tested , leading to three important open questions : ∗Equal Contribution . Most work was done during Chunting ’ s internship at FAIR . 1Code is released at https : //github.com/pytorch/fairseq/tree/master/examples/ nonautoregressive_translation . • Exactly how does distillation reduce the “ modes ” , and how we could we measure this reduction quantitatively ? Why does this reduction consistently improve NAT models ? • What is the relationship between the NAT model ( student ) and the AT model ( teacher ) ? Are different varieties of distilled data better for different NAT models ? • Due to distillation , the performance of NAT models is largely bounded by the choice of AT teacher . Is there a way to further close the performance gap with standard AT models ? In this paper , we aim to answer the three questions above , improving understanding of knowledge distillation through empirical analysis over a variety of AT and NAT models . Specifically , our contributions are as follows : • We first visualize explicitly on a synthetic dataset how modes are reduced by distillation ( §3.1 ) . Inspired by the synthetic experiments , we further propose metrics for measuring complexity and faithfulness for a given training set . Specifically , our metrics are the conditional entropy and KL-divergence of word translation based on an external alignment tool , and we show that these metrics are correlated with NAT model performance ( §3.2 ) . • We conduct a systematic analysis ( §4 ) over four AT teacher models and six NAT student models with various architectures on the standard WMT14 English-German translation benchmark . These experiments find a strong correlation between the capacity of an NAT model and the optimal dataset complexity that results in the best translation quality . • Inspired by these observations , we propose approaches to further adjust the complexity of the distilled data in order to match the model ’ s capacity ( §5 ) . We also show that we can achieve the state-of-the-art performance for NAT models and largely match the performance of the AT model . 2 BACKGROUND . 2.1 NON-AUTOREGRESSIVE NEURAL MACHINE TRANSLATION . In order to model the joint probability of the output sequence y , NMT models usually generate each output token conditioned on the previously generated ones p ( y|x ) = ∏T t=1 p ( yt|y < t , x ) . This is known as the autoregressive factorization . To generate a translation from this model , one could predict one token at a time from left to right and greedily take arg max over each output probability distribution , or use beam search to consider a fixed number of hypotheses . In this work , we study non-autoregressive translation ( NAT ) , a special subset of NMT models with an additional restriction ( the zeroth-order Markov assumption ) upon the output predictions or a subset thereof . The simplest formulation of an NAT model independently factors the conditional distribution : p ( y|x ) = ∏T t=1 p ( yt|x ) . Standard NAT models ( Gu et al. , 2018 ) adopt an architecture similar to the Transformer ( Vaswani et al. , 2017 ) and make non-autoregressive predictions for the entire sequence with one forward pass of the decoder . However , because multiple translations are possible for a single input sentence ( the so-called multi-modality problem ; Gu et al . ( 2018 ) ) , vanilla NAT models can fail to capture the dependencies between output tokens . As a result , they tend to make egregious mistakes such as outputting tokens repeatedly . To improve the model ’ s ability to handle multi-modality , recent works have incorporated approaches including ( 1 ) relaxing the fully non-autoregressive restriction and adopting K decoding passes ( instead of just one ) to iteratively refine the generated outputs ( Lee et al. , 2018 ; Ghazvininejad et al. , 2019 ; Wang et al. , 2018 ; Stern et al. , 2018 ; 2019 ; Gu et al. , 2019 ) ; ( 2 ) using latent variables ( Kaiser et al. , 2018 ; Ma et al. , 2019 ; Shu et al. , 2019 ) or structured information such as syntax trees ( Akoury et al. , 2019 ) to capture translation variation ; ( 3 ) training NAT models with objectives other than maximum likelihood ( Wang et al. , 2019 ; Wei et al. , 2019 ; Shao et al. , 2019 ) which ameliorates the effects of multi-modality . However , to achieve competitive performance with the autoregressive model , almost all existing NAT models rely on training using data distilled from a pre-trained AT model instead of the real parallel training set , as described below . 2.2 SEQUENCE-LEVEL KNOWLEDGE DISTILLATION . Knowledge distillation ( Liang et al. , 2008 ; Hinton et al. , 2015 ) was originally proposed for training a weaker student classifier on the targets predicted from a stronger teacher model . A typ- ical approach is using the label probabilities produced by the teacher as “ soft targets ” qi = exp ( zi/τ ) /∑ j exp ( zj/τ ) for training the student model , where qi and zi are the probability and the logit of class i respectively and τ is the temperature . Prior work has shown the effectiveness of adopting knowledge distillation in adversarial defense ( Papernot et al. , 2016 ) , neural network compression ( Howard et al. , 2017 ) , and fast inference for speech synthesis ( Oord et al. , 2018 ) . In the context of sequence generation , Kim & Rush ( 2016 ) extend knowledge distillation to the sentence level using “ hard targets ” from a pretrained large teacher model to train a small sequence generation model . More precisely , the teacher distribution q ( t|x ) is approximated by its mode : q ( t|x ) ≈ 1 { t = arg maxt∈T q ( t|x ) } with the following objectives : Lseq-KD = −Ex∼data ∑ t∈T q ( t|x ) log p ( t|x ) ≈ −Ex∼data , ŷ=argmax t∈T q ( t|x ) [ log p ( t = ŷ|x ) ] , ( 1 ) where t ∈ T is the space of possible target sequences . This can also be seen as a special case of standard distillation over the sentence space when the temperature τ approaches 0 , which is equivalent to taking the arg max over all feasible translations . While the “ hard target ” ŷ is the most likely translation predicted by the teacher , in practice we use beam search as an approximation . As mentioned earlier , almost all the existing literature trains NAT models using sequence-level knowledge distillation from a pre-trained AT model to achieve competitive performance . Particularly , it is common to train the teacher model as a standard autoregressive Transformer ( Vaswani et al. , 2017 ) with a roughly equal number of trainable parameters as the desired NAT model on the real data . Next , we will first study how this knowledge distillation process affects the behavior of NAT models . 3 HOW DOES DISTILLATION IMPROVE NAT ? . In this section , we start from an introductory example to illustrate how NAT models fail to capture the multi-modality of data . Then we propose a metric to assess the multi-modality of a data set and use it to test our hypothesis about how knowledge distillation affects NAT models . 3.1 SYNTHETIC EXPERIMENT FOR MULTI-MODALITY . Dataset . We start by investigating NAT ’ s difficulties in modeling multi-modality in output data using a synthetic setup where we explicitly include multiple modes in the training data . More specifically , we utilize three language pairs – English-German ( En-De ) , English-French ( En-Fr ) , and English-Spanish ( En-Es ) – from the Europarl parallel corpus.2 We extract sentences that have aligned sentences for all languages , and create a multi-target En-De/Es/Fr corpus . In this case every English input sentence always corresponds to target sentences in three different languages , which forms three explicit output modes . Notably , this is similar to the one-to-many translation setting in Johnson et al . ( 2017 ) but in our case we do not have an explicit signal ( e.g . target language tag ) to tell the NMT model which target language to translate to . Models . We train both the AT and NAT models on this concatenated data set , then compare the distributions of translations with each other . We use the standard Transformer ( base ) model ( Vaswani et al. , 2017 ) as the AT model , and a simplified version of Gu et al . ( 2018 ) as the NAT model where the decoder ’ s inputs are monotonically copied from the encoder embeddings and a length predictor is learned to predict the target sentence length . Both models are trained for 300 , 000 steps using maximum likelihood . After training , we use both models to translate the English sentences in the validation and test sets . Visualization of AT Outputs . The synthetic setup enables us to better understand and visualize the modes in the outputs more easily . First , we visualize the outputs from the AT model . For every translated sentence , we visualize the estimated probability distribution of language classes as a point in Fig . 1 ( a ) . This probability is calculated as the average of the posterior probability of each token , and it is estimated based on the Bayes ’ law : p ( li|y ) ≈ 1 T T∑ t=1 p ( li|yt ) = 1 T T∑ t=1 p ( yt|li ) p ( li ) ∑ k p ( yt|lk ) p ( lk ) ( 2 ) 2https : //www.statmt.org/europarl/ where li denotes the language class i , and p ( yt|li ) is the token frequency of yt in language li . We assume p ( li ) follows a uniform distribution . As shown in Fig . 1 ( a ) , points of the AT outputs are clustered closely to each vertex of the simplex , indicating that the AT model prefers to generate the whole sequence in one language . This phenomenon verifies our assumption that decoding with the AT model ( distillation ) is essentially selecting “ modes ” over the real data . Visualization of NAT Outputs . We visualize outputs for the NAT model trained on the same data in Fig . 1 ( b ) . In contrast to the AT results , the NAT points are scattered broadly inside the simplex , indicating that the NAT model fails to capture the mode of language types . Instead , it predicts tokens mixed with multiple languages , which corroborates our hypothesis that the NAT model has trouble consistently selecting a single mode when multiple modes exist . Next , we create two datasets that have fewer modes than the original dataset . First , we randomly select a single target sentence from one of the three languages for each source sentence . Second , we perform distillation , decoding from the AT model trained on the combined training set . As noted in the AT results , distillation will also roughly be selecting a language mode , but we conjecture that this selection may be more systematic , selecting a particular language for a particular type of training sentence . As shown in Fig . 1 ( c ) ( d ) , NAT models trained on both of these datasets are more likely to choose one mode ( language ) when generating translations , showing that training with reduced modes is essential for NAT model . Furthermore , points in Fig . 1 ( d ) are clearly clustered better than ( c ) indicating that modes selected by AT models are indeed likely more systematic and easy to capture than those generated by randomly assigning a language for each sentence .

Summary: The paper analyses recent distillation techniques for non-autoregressive machine translation models (NAT). These models use a autoregressive teacher (AT), which typically perform better. However, AT models can not be parallelized that easily such as the NAT models. The distillation has the effect of removing modes from the dataset which helps the NAT models as they suffer from the averaging effect of maximum likelihood solutions. The authors analyze why such distillation is needed and what the effect of the complexity of the training set is and further propose 3 methods to adjust the complexity of the teacher to the complexity of the NAT model.