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

1 Introduction

• Earlier non-autoregressive refinement methods iteratively improve an initial translation in discrete sentence space, while Shu et al. (2020) use a hybrid of continuous optimization and discrete search.Autoregressive decoding latency grows linearly with target length; non-autoregressive models seek faster generation with sufficient parallel computation.Source: §1 · p. 1
• The proposed method trains an inference network to estimate the marginal log-probability gradient from the latent variable, then follows that gradient to find a likely target sentence.Inference initializes the latent variable, for example at the prior mean.Source: §1 · pp. 1–2
• On WMT’14 En→De, the proposed approach achieved a 6.2× speedup over the autoregressive baseline during inference, at the expense of a 0.9 BLEU score.Source: §1 · p. 2
• Each refinement step is twice as fast as EM-like inference and yields higher marginal probabilities and BLEU scores with the same number of steps.Evaluations cover WMT’14 En→De, WMT’16 Ro→En, and IWSLT’16 De→En.Source: §1 · p. 2

2.1 Refinement in a Discrete Space

• Most existing non-autoregressive models generate an initial hypothesis and successively refine it, using objectives such as denoising, mask language modeling, or evidence lower-bound maximization.The lack of autoregression makes target-token dependencies difficult to model.Source: §2.1 · p. 2
• Discrete-space inference seeks a target maximizing log p(y|x), but intractable marginalization over refinement variables leads to approximate joint search over y and z1:L.Approximate search selects each refinement variable conditioned on earlier ones.Source: §2.1 · p. 2

2.2 Refinement in a Hybrid Space

• Continuous latent variables of dimensionality T × D capture target-token dependencies so the decoder can factorize over time, but exact posterior inference and learning are intractable for typical deep models.Shu et al. (2020) use a Dirac-delta posterior to form a proxy lower bound, then alternate posterior matching and target search.Source: §2.2 · pp. 2–3

3 Iterative Refinement in a Continuous Space

• The method replaces delta inference’s discrete target search with optimization only in the continuous latent space, avoiding T searches over a vocabulary of size V at each refinement step.The passage notes that this remains expensive even when the T searches are parallelized.Source: §3 · p. 3

3.1 Learning

• Because the most-likely target’s marginal log probability is nondifferentiable through an argmax, training uses proxy gradients from delta inference, with the change after k steps serving as an approximation to the gradient.The learned function is trained on differences between latent configurations, with configurations weighted according to the prior.Source: §3.1 · p. 3

3.2 Parameterization

• The gradient can be parameterized either through a scalar energy function or directly as a score function; direct score estimation avoids backpropagation at each refinement step.The paper compares both parameterizations experimentally.Source: §3.2 · p. 3

3.3 Inference

• Inference initializes the latent variable from the prior, such as at its mean or a sample, then repeatedly updates it using estimated gradients until a fixed-step or convergence condition is met.Each step optimizes only in continuous space, avoiding search over a large vocabulary.Source: §3.3 · p. 4

4.1 Datasets and Preprocessing

• The evaluation covers three widely used datasets: one with 197K training, 2K development, and 2K test pairs; WMT’16 Ro→En with 612K/2K/2K; and WMT’14 En→De with 4.5M/3K/3K sentence pairs.Source: §4.1 · p. 4

4.2 Models and Baselines

• The latent-variable model uses a time-factorized decoder and spherical Gaussian prior and approximate posterior; latent dimensionality is 8 across datasets.At inference, the target length is predicted once and fixed, so latent dimensionality does not change during refinement.Source: §4.2 · p. 4
• The autoregressive baselines use Transformer models, while the latent-variable procedures compare Delta, Energy, and Score under the same underlying model.The table caption defines speed relative to an autoregressive beam-width-4 model and reports average time per example on a Tesla V100 GPU.Source: §4.2 · p. 4, Table 1 · p. 5

4.3 Training and Optimization

• Inference networks are trained for 1M steps with the latent-variable model fixed, using prior samples refined by four Delta iterations to form training targets.A stochastic estimated-gradient update before target computation was also found to improve performance.Source: §4.3 · p. 5

4.4 Inference

• Inference can use either the target-length argmax or parallel decoding of top length candidates, and can refine multiple prior samples in parallel before selecting by log probability.For latent search, initialization may use a prior sample or its mean; the experiments fix the random seed during sampling.Source: §4.4 · p. 5

5.1 Translation Quality and Speed

• The learned score is comparable to or better than Delta in BLEU on all datasets, and four refinement steps improve IWSLT’16 De→En BLEU by 2.1.Often, one refinement step gives comparable translation quality to four; parallel length and latent-sample search adds at least 1 BLEU on all datasets.Source: §5.1 · pp. 5–6
• The learned-score procedure is twice as fast as Delta on IWSLT’16 De→En and WMT’16 Ro→En, and almost four times as fast on WMT’14 En→De.Four learned-score steps on WMT’14 En→De take a time close to one Delta step; learned Energy is slower, with wall-clock time similar to Delta.Source: §5.1 · p. 6
• 6.2× speedup over the autoregressive baseline on WMT’14 En→De came with 0.9 BLEU score degradation using one learned-score refinement step and parallel search.The approach is described as highly effective and efficient overall.Source: §5.1 · p. 6

5.2 Log Probability Comparison

• Marginal log probability improves with additional refinement steps for Delta and both proposed procedures.The comparison estimates marginal log probability by importance sampling with 500 approximate-posterior samples.Source: §5.2 · p. 6, Figure 1 · p. 6

5.3 Token Statistics

• Learned-score refinement produces greater edit distance from the initial output and fewer token repetitions than Delta on IWSLT’16 De→En.Figure 2 compares these statistics across 1, 2, 4, and 8 refinement steps.Source: §5.3 · p. 6, Figure 2 · p. 6

6.1 Visualization of learned gradients

• For Token 4, the score-function local optimum is far from the posterior mean, yet both map to reference translation (“my”), showing multiple latent variables yield the reference output.Source: §6.1 · p. 7

6.2 Sample translations

• One refinement reorders the German main clause and prepositional phrase, and changes “es gibt” to the grammatical “gibt es.”The clause order places the prepositional phrase before the main clause.Source: §6.2 · p. 7
• In another example, refinement exchanges the two monetary amounts and inserts “aus den” between them.The example concerns 1.2 billion and 6.9 billion dollars.Source: §6.2 · p. 7
• A further example combines “Weisheit in Bedouin” into “Bedouin-Weisheit” and swaps two phrases to better resemble the reference.The phrases are “Der erste ...” and “mit dieser ...”.Source: §6.2 · p. 7

Learning

• The training objective is related to score matching but uses a proxy gradient for the data log density and considers samples only from the prior.Prior work cited in this discussion includes denoising interpretations and image generative modeling applications.Source: §7 · p. 7

Inference

• Other inference work relaxes discrete target-token search with Gumbel-softmax and trains networks to maximize pretrained-model log probability.The cited relaxation methods include Hoang et al. (2017), Gu et al. (2018b), and Tu et al. (2020).Source: §7 · p. 7
• The cited continuous-relaxation line of work has also been applied to automatic post-editing.The passage names Novak et al. (2016) and Grangier and Auli (2017).Source: §7 · p. 8

• The proposed inference procedure refines non-autoregressive machine translation translations purely in continuous space.An inference network estimates the gradient of the marginal log probability.Source: §8 · p. 8
• Effectiveness is demonstrated for spherical Gaussian priors, while success with Gaussian mixtures or normalizing flows remains uninvestigated.Source: §8 · p. 8
• Flexible-prior latent-variable models have shown high test loglikelihoods but poor generation quality when inference is challenging (Lee et al., 2020).Source: §8 · p. 8

• Eq. 5 follows from a linear approximation.Source: A Training objective · p. 11
• The ∥z−z∥2 term in Eq. 6 can be eliminated because dividing the objective by a positive constant does not change its solution.Source: A Training objective · p. 11
• The second term in Eq. 7 is ignored because it is constant with respect to ψ.Source: A Training objective · p. 11