Diverse Machine Translation with a Single Multinomial Latent Variable

1 INTRODUCTION . Modeling uncertainty in the output distribution is a fundamental problem in machine learning . Some prediction tasks , in particular those involving text generation , may admit many valid outputs . However , in most cases the learner only observes one out of many correct outputs both at training and test time . This raises questions about how to properly fit the data distribution and how to properly measure such fitting . In this work , we investigate these questions for the task of machine translation ( MT ) , however the arguments and methods are general and can equally be applied to other models and tasks . Specifically , we aim at better modeling uncertainty for the purpose of effectively decoding a diverse set of hypotheses . The motivation is twofold . On one hand , we are interested in exploring the problem of modeling uncertainty , and we would like to better understand how different latent variable models capture uncertainty . MT serves as a well defined task for this investigation . On the other hand , we are driven by the potential practical applications of such models for MT , which could enable end-users to express their preference about how to translate a piece of text . Previous work has explored approaches for heuristically modifying the decoding algorithm to produce diverse hypotheses , e.g. , by enforcing diversity constraints during beam search ( Li & Jurafsky , 2016 ; Vijayakumar et al. , 2016 ) . Neverthless , it seems more natural to consider explicit ways to represent the multi-modal nature of this task by introducing latent variables in the model . Notably , Zhang et al . ( 2016 ) proposed a variational version of NMT with a latent Gaussian variable ( Kingma & Welling , 2013 ) . While they showed improvements due to the regularization effect of the Monte Carlo gradient estimate , it is unclear how well the corresponding latent space actually captures different output modes . We propose a much simpler approach that introduces a single multinomial latent variable to capture uncertainty in the target sentence . This model , which is a simplified instance of a mixture of experts ( MoE ) model ( Jacobs et al. , 1991 ) , is much easier to train and interpret than previous approaches . Our approach also enables fast parallel decoding from different latent values to generate a diverse set of hypotheses . We evaluate our method on three WMT benchmark datasets that uniquely provide test sets with multiple human references . These datasets enable us to quantitatively evaluate diversity , and we 1Code to reproduce this work is available at : anonymized URL . provide the first systematic comparison of the various methods that produce multiple hypotheses . Our results demonstrate that our proposed model is more robust and generates both high-quality and diverse translations , and compares favorably to variational NMT models , other variants of mixture of expert models , and diverse inference algorithms such as diverse beam search ( Vijayakumar et al. , 2016 ) and biased sampling ( Graves , 2013 ; Fan et al. , 2018 ) . Moreover , our study provides insights for best practices to avoid several failure modes of latent variable models , in particular when these are trained with modern architectures using dropout and shared parameters . 2 DIVERSE DECODING IN NMT . Let x be a source sentence consisting of words ( x1 , · · · , xL ) , and y = ( y1 , · · · , yT ) a given reference ( target ) translation . An NMT model has an encoder-decoder structure . The encoder maps ( x1 , · · · , xL ) to a sequence of representations h = ( h1 , · · · , hL ) that are then fed to the decoder , which generates an output sentence one word at a time conditioned on h. At each time step , the decoder additionally conditions its output on the previous outputs , resulting in an auto-regressive factorization p ( y|x ; θ ) = ∏T t=1 p ( yt|y1 : t−1 , x ; θ ) , where θ denotes the parameters of the network and is learned by maximum likelihood training . At test time , approximate inference algorithms such as beam search are commonly used to decode the most likely translations . Beam search decodes a sequence left-to-right while maintaining the K best hypotheses at each step . Although it is often effective in finding high quality hypotheses , these hypotheses typically have low diversity with only minor differences in the suffix . To remedy this , Vijayakumar et al . ( 2016 ) proposed diverse beam search , where they perform beam search sequentially for G iterations , penalizing the selection of words used in previous iterations . Another decoding strategy is to sample repeatedly from the model ’ s conditional distribution ŷt ∝ p ( ·|ŷ1 : t−1 , x ; θ ) . Unfortunately , this approach usually produces diverse but low quality outputs . Several heuristic methods have been proposed to obtain better quality , such as sampling with a temperature , biased sampling from the top k most likely words instead of all words at each step ( Graves , 2013 ; Fan et al. , 2018 ) , oversampling K ′ > K outputs and retaining the top K , adding noise to the hidden state ( Cho , 2016 ) , etc . 3 LATENT VARIABLE NMT . Rather than search for diverse outputs at test time , which may be slow and difficult to tune , we aim to explicitly model uncertainty during training . We introduce a multinomial latent variable z taking values in { 1 , · · · , K } to capture various ways to translate the source sentence , and train a latent variable model p ( y|x ; θ ) = ∑K z=1 p ( z|x ; θ ) p ( y|z , x ; θ ) . Intuitively , we want to push the uncertainty from p ( y|x ) into z , so that we can easily explore multiple modes during training and inference . For each value of z , let us consider p ( y|z , x ; θ ) to be an expert and p ( z|x ; θ ) to be a gating function over experts , i.e. , a source-dependent prior of using expert z . Training proceeds by minimizing the negative log-likelihood loss function : LSoft-MoE ( θ ) = E ( x , y ) ∼data [ − log K∑ z=1 p ( z|x ; θ ) p ( y|z , x ; θ ) ] ( 1 ) This is the Soft-Mixture of Experts ( Soft-MoE ) model . In our work , the encoder and experts are based on the Transformer architecture ( Vaswani et al. , 2017 ) , while the gating is implemented via mean-pooling over the encoder ’ s representations followed by a single layer feed-forward net to predict p ( z|x ; θ ) , similar to related work ( Zhang et al. , 2016 ) . 3.1 OUR APPROACH : HARD-MOE . A well-known challenge in training Soft-MoE models is that experts “ die ” during training ; in extreme cases the model collapses and only a single expert is trained ( Eigen et al. , 2014 ; Shazeer et al. , 2017b ) . We address this by introducing two simplifying assumptions . First , we assume that the prior p ( z|x ) = 1/K is uniform and source independent . This encourages the model to use multiple experts instead of collapsing into one with extreme p ( z|x ) values . Second , we assume that a particular Algorithm 1 Hard-MoE Training with K Experts Initialize model parameters θ repeat Sample a mini-batch of m examples { ( x ( i ) , y ( i ) ) } mi=1 Set the model in evaluation mode , turn off dropout and disable gradient computation for z in { 1 , · · · , K } do Compute − log p ( y ( i ) |z , x ( i ) ; θ ) end for Let z ( i ) = argminz − log p ( y ( i ) |z , x ( i ) ; θ ) Set the model in train mode , turn on dropout and enable gradient computation Do another forward-backward pass to compute loss L = 1m ∑m i=1− log p ( y ( i ) |z ( i ) , x ( i ) ; θ ) and update θ by gradient descent until Convergence translation y is well explained by exactly one expert , that is p ( y|z , x ) is large for only one value of z . This forces experts to specialize on different subsets of the data . Based on these assumptions , we can approximate the summation in Eq . 1 with the maximum term : p ( y|x ) = ∑ z p ( z|x ) p ( y|z , x ) ≥ max z p ( z|x ) p ( y|z , x ) = 1 K max z p ( y|z , x ) ( 2 ) Taking the negative logarithm , and omitting the constant logK term , we get our loss function : LHard-MoE ( θ ) = E ( x , y ) ∼data [ min z − log p ( y|z , x ; θ ) ] ( 3 ) To train with this objective , for each mini-batch we enumerate all K possible values of z ( typically K = 10 ) and compute their losses : − log p ( y|z , x ; θ ) . Next , for each ( x , y ) pair we select the value of z that yields the minimum loss , and use that value to back-propagate , i.e. , only one expert receives gradients per sentence pair . Since we make hard selections of the latent variable , we call this method Hard-Mixture of Experts ( Hard-MoE ) . This training algorithm can also be seen as a stochastic gradient descent version of the Hard-Expectation Maximization algorithm . 3.2 PRACTICAL CONSIDERATIONS . So far our discussion has been very general . Notice however that our motivation for training with a mixture of experts is different from existing MoE literature that typically aims to increase model capacity ( Collobert et al. , 2003 ; Shazeer et al. , 2017b ; Gross et al. , 2017 ) . Instead , we aim to better model multi-modal output distributions—a difference that has significant practical implications . First , even on large scale MT benchmarks , models are prone to overfit and require regularization ( e.g. , dropout , label smoothing , etc . ) ; the standard modeling choice of using independently parameterized ( decoder ) experts is likely to exacerbate this overfitting . Given our goal is not to increase model capacity , we choose instead to share almost all parameters among the experts , except for a single set of weights that we use to embed the different values of the latent variable . Accordingly , we require only a negligible number of additional parameters over a baseline model . More specifically , we found good results in pilot experiments by simply replacing the first dummy beginningof-sentence or 〈bos〉 token of the decoder with the embedded representation of our latent variable.2 This modeling choice further mitigates the latent variable collapse issue often reported in the MoE literature ( Eigen et al. , 2014 ; Shazeer et al. , 2017b ) , in which low quality experts are neglected and eventually “ die ” during training . Instead , by sharing parameters , even unpopular experts receive some gradients throughout training . Another potential problem is with dropout . Dropout noise can mask the effect of the change of the latent variable , making assignments inconsistent and eventually leading to a degenerate model 2We also tried other parameterizations , including adding or concatenating the latent variable embedding with the input word embeddings at each time step , and injecting it into each decoder layer . We found that they have similar results , therefore we adopt the simplest replacing 〈bos〉 strategy . that ignores the latent variable . Therefore , in practice we decompose the minimization over z in Eq . 3 into two separate steps . First we set the model in evaluation mode , turn off dropout , disable gradient computation and perform K forward passes to identify the optimal assignment of z.3 Once the optimal latent values are selected , we switch to train mode , turn on dropout , and do another forward-backward pass to update the parameters . This is crucial ; in Sec . 5.4 we show that if we instead leave dropout on during minimization , then the latent variable is ultimately ignored . The whole process is formally described in Algorithm 1 . All the decoding strategies for p ( y|x ; θ ) mentioned in Sec . 2 can be applied to the latent variable model p ( y|z , x ; θ ) as well . We adopt the most straightforward one : generating K hypotheses by first enumerating z and then greedily decoding ŷt = argmaxy p ( y|ŷ1 : t−1 , z , x ; θ ) . This tests whether our latent variable captures different modes in p ( y|x ) , and we can generate hypotheses efficiently in parallel .

Summary: The authors aim to increase diversity in machine translation using a multinomial latent variable that captures uncertainty in the target sentence. Modeling uncertainty with latent variables is of course relatively common in ML, and this work has similarities with latent variables models for MT [Zhang et al., 2016] and for other generation tasks such as dialogue [Serban et al., 2017; etc.]. The key difference is that the authors here use a Mixture of Expert (MoE) approach while most relevant prior works use variational approaches. Experiments show improvements in diversity over variational NMT [Zhang et al., 2016] and decoding-time approaches (e.g., diversity constraints [Vijayakumar et al., 2016]).