Iterative Refinement in the Continuous Space for Non-Autoregressive Neural Machine Translation (2020)

Abstract: We propose an efficient inference procedure for non-autoregressive machine translation that iteratively refines translation purely in the continuous space. Given a continuous latent variable model for machine translation (Shu et al., 2020), we train an inference network to approximate the gradient of the marginal log probability of the target sentence, using the latent variable instead. This allows us to use gradient-based optimization to find the target sentence at inference time that approximately maximizes its marginal probability. As each refinement step only involves computation in the latent space of low dimensionality (we use 8 in our experiments), we avoid computational overhead incurred by existing non-autoregressive inference procedures that often refine in token space. We compare our approach to a recently proposed EM-like inference procedure (Shu et al., 2020) that optimizes in a hybrid space, consisting of both discrete and continuous variables. We evaluate our approach on WMT’14 En→De, WMT’16 Ro→En and IWSLT’16 De→En, and observe two advantages over the EM-like inference: (1) it is computationally efficient, i.e. each refinement step is twice as fast, and (2) it is more effective, resulting in higher marginal probabilities and BLEU scores with the same number of refinement steps. On WMT’14 En→De, for instance, our approach is able to decode 6.2 times faster than the autoregressive model with minimal degradation to translation quality (0.9 BLEU).

## introduction most neural machine translation systems are autoregressive, hence decoding latency grows linearly with respect to the length of the target sentence. for faster generation, several work proposed nonautoregressive models with sub-linear decoding latency given sufficient parallel computation @xcite @xcite . as it is challenging to precisely model the dependencies among the tokens without autoregression, many existing non-autoregressive models first generate an initial translation which is then iteratively refined to yield better output @xcite @xcite . while various training objectives are used to admit refinement (e.g. denoising, evidence lowerbound maximization and mask language modeling), the generation process of these models is similar in that the refinement process happens in the discrete space of sentences. meanwhile, another line of work proposed to use continuous latent variables for non-autoregressive translation, such that the distribution of the target sentences can be factorized over time given the latent variables @xcite . unlike the models discussed above, finding the most likely target sentence under these models requires searching over continuous latent variables. to this end, @xcite proposed an em-like inference procedure that optimizes over a hybrid space consisting of both continuous and discrete variables. by introducing a deterministic delta posterior, it maximizes a proxy lowerbound by alternating between matching the delta posterior to the original approximate posterior (continuous optimization), and finding a target sentence that maximizes the proxy lowerbound (discrete search). in this work, we propose an iterative inference procedure for latent variable non-autoregressive models that purely operates in the continuous space. 1 given a latent variable model, we train an inference network to estimate the gradient of the marginal log probability of the target sentence, using only the latent variable as input. at inference time, we find the target sentence that approximately maximizes the log probability by (1) initializing the latent variable e.g. as the mean of the prior, and (2) following the gradients estimated by the inference network. we compare the proposed approach with the emlike inference @xcite on three machine translation datasets: wmt'14 @xmath0, wmt'16 @xmath1. the advantages of our approach are twofold: (1) each refinement step is twice as fast, as it avoids discrete search over a large vocabulary, and (2) it is more effective, giving higher marginal probabilities and bleu scores with the same number of refinement steps. our procedure results in significantly faster inference, for instance giving 6.2× speedup over the autoregressive baseline on wmt'14 @xmath2.9 bleu score. ## background: iterative refinement for non-autoregressive translation we motivate our approach by reviewing existing refinement-based non-autoregressive models for machine translation in terms of their inference procedure. let us use v, d, t and l to denote vocabulary size, latent dimensionality, target sentence length and the number of refinement steps, respectively. most machine translation models are trained to maximize the conditional log probability log p(y|x) of the target sentence y given the source sentence x, averaged over the training data consisting of sentence pairs {(x n , y n )} n @xmath3. to find the most likely target sentence at test time, one performs maximum-a-posteriori inference by solving a search problem ŷ = argmax y log p(y|x). ## iterative refinement in a continuous space while the delta inference procedure is an effective inference algorithm for machine translation models with continuous latent variables, it is unsatisfactory as the m-step requires searching over v tokens t times for each refinement step. as v is large for most machine translation models, this is an expensive operation, even when the t searches can be parallelized. we thus propose to replace the delta inference with continuous optimization in the latent space only, given the underlying latent variable model. ## conclusion we propose an efficient inference procedure for non-autoregressive machine translation that refines translations purely in the continuous space. given a latent variable model for machine translation, we train an inference network to approximate the gradient of the marginal log probability with respect to the target sentence, using only the latent variable. this allows us to use gradient based optimization to find a target sentence at inference time that approximately maximizes the marginal log probability. as we avoid discrete search over a large vocabulary, our inference procedure is more efficient than previous inference procedures that refine in the token space. we compare our approach with a recently proposed delta inference procedure that optimizes jointly in discrete and continuous space on three machine translation datasets: wmt'14 @xmath4, wmt'16 @xmath5. with the same underlying latent variable model, the proposed inference procedure using a learned score function has following advantages: (1) it is twice as fast as delta inference, and (2) it is able to find target sentences resulting in higher marginal probabilities and bleu scores. while we showed that iterative inference with a learned score function is effective for spherical gaussian priors, more work is required to investigate if such an approach will also be successful for more sophisticated priors, such as gaussian mixtures or normalizing flows. this will be particularly interesting, as recent study showed latent variable models with a flexible prior give high test loglikelihoods, but suffer from poor generation quality as inference is challenging @xcite .