Source-linked AI summary

Imitation Learning for Non-Autoregressive Neural Machine Translation

Bingzhen Wei, Mingxuan Wang, Hao Zhou, Junyang Lin, Jun Xie, Xu Sun

arXiv:1906.02041v2cs.CL

TL;DR

Autoregressive translation is slowed by sequential decoding, while NAT’s parallel decoding poses training challenges from delayed supervision and a large search space. This paper trains NAT with an autoregressive demonstrator supervising its decoding states. The approach largely closes the performance gap across benchmarks while preserving NAT’s speed advantage.

  • Problem

    NAT training is challenging because parallel decoding creates a large search space and delays correct-word supervision until the final layer.

    Source: §1 · pp. 1–2
  • Method

    Imitate-NAT uses a knowledgeable autoregressive demonstrator to supervise every NAT decoding state across time steps and layers during training.

    Source: §1 · p. 2
  • Results

    Imitate-NAT largely closes the NAT–AT performance gap across benchmarks; on IWSLT16 En→De it scores 30.68 versus 30.85 for the autoregressive model.

    Source: §5 · p. 7, §7 · p. 9
  • Takeaways & Limitations

    Because the demonstrator is used only in training, imitate-NAT retains NAT’s high inference speed.

    Source: §1 · p. 2

    The paper leaves stronger or differently structured demonstrators and applications to other scenarios as future work.

    Source: §7 · p. 9

Abstract

from arXiv · show

Non-autoregressive translation models (NAT) have achieved impressive inference speedup. A potential issue of the existing NAT algorithms, however, is that the decoding is conducted in parallel, without directly considering previous context. In this paper, we propose an imitation learning framework for non-autoregressive machine translation, which still enjoys the fast translation speed but gives comparable translation performance compared to its auto-regressive counterpart. We conduct experiments on the IWSLT16, WMT14 and WMT16 datasets. Our proposed model achieves a significant speedup over the autoregressive models, while keeping the translation quality comparable to the autoregressive models. By sampling sentence length in parallel at inference time, we achieve the performance of 31.85 BLEU on WMT16 Ro$\rightarrow$En and 30.68 BLEU on IWSLT16 En$\rightarrow$De.

1 Introduction

Autoregressive decoding is slow because each word depends on previous words, while non-autoregressive translation is faster but faces training challenges from delayed supervision and a large search space. The paper introduces an autoregressive demonstrator to supervise NAT decoding states and reports strong benchmark results while retaining NAT’s speed advantage.

1 Introduction

  • Autoregressive decoders generate words sequentially conditioned on previous words, preventing parallel decoder computation and slowing translation.The paper notes that decoding with a large vocabulary is especially time-consuming.Source: §1 · p. 1
  • NAT generates all words simultaneously, enabling parallel decoding but typically yielding moderate accuracy, described as within 5 BLEU.NAT removes previous predicted words from the decoder’s next-step inputs.Source: §1 · p. 1
  • NAT training faces delayed supervision and a large search space because decoding states make interacting decisions across layers and time steps, with correct-word supervision arriving at the last layer.The paper attributes the performance gap between NAT and AT to these challenges.Source: §1 · pp. 1–2
  • A knowledgeable autoregressive demonstrator supervises every NAT decoding state across time steps and layers, while being used only during training.This design aims to improve translation performance without giving up NAT’s high inference speed.Source: §1 · p. 2
  • Imitate-NAT with LPD achieves the best translation performance on the benchmark datasets, approaching autoregressive results while retaining a 10-times speed advantage.The introduction reports that the model effectively closes the AT–NAT performance gap.Source: §1 · p. 2

2 Background

The background contrasts autoregressive models, which generate target words sequentially, with non-autoregressive models that generate them in parallel. It also describes how removing target-side dependencies can weaken prediction context and notes a fertility-based NAT approach.

2.1 Autoregressive Neural Machine Translation

  • Autoregressive NMT models target words sequentially from left to right and train by minimizing negative log-likelihood; the sequential inference dependency remains a bottleneck even in parallelizable architectures.T and T′ denote source and target sequence lengths, respectively.Source: §2.1 · p. 2

2.2 Non-Autoregressive Neural Machine Translation

  • NAT removes autoregressive connections and factorizes the target distribution into conditionally independent per-step distributions, allowing target distributions to be computed in parallel.Removing target-side dependencies can leave the decoder with weaker target context and a larger modeling gap.Source: §2.2 · pp. 2–3
  • A described NAT approach uses a latent fertility model to copy source embeddings as decoder inputs, then samples fertilities and rescoring with an autoregressive model to obtain its best results.Its reported speed is 2–15 times that of a comparable autoregressive model, depending on the number of fertility samples.Source: §2.2 · p. 3

3 Proposed Method: imitate-NAT

imitate-NAT uses an autoregressive demonstrator to supervise a non-autoregressive learner across decoding states, aiming to close the performance gap between NAT and AT. Its method combines action-distribution regularization, SoftCopy initialization, and learning from AT experts.

3 Proposed Method: imitate-NAT

  • Figure 3 depicts the imitation module shared between the AT demonstrator and NAT learner.The framework describes the learner as imitating the demonstrator across decoding states.Source: Figure 3 · p. 4, §3.1 · p. 4
  • The framework trains a NAT learner to imitate an autoregressive demonstrator at each decoding state, while using the demonstrator only during training.This is intended to retain NAT's inference speed without its relatively lower performance.Source: §3 · pp. 3–4, §3.1 · p. 4

3.2 AT Demonstrator

  • DAT emits actions at every decoding step to provide sequential guidance, and its next-layer states use both decoder output and emitted action state.The action distribution can serve as the NAT model's training signal.Source: §3.2 · p. 4

3.2.1 Action Distribution Regularization

  • Action-distribution regularization treats NAT action prediction as unsupervised clustering and adjusts probabilities toward more balanced category assignments.The method uses cumulative activation levels to redistribute action probabilities.Source: §3.2.1 · pp. 4–5

3.3.1 Soft Copy

  • SoftCopy replaces unstable, non-differentiable Round-based input initialization with differentiable, distance-based weights between source and target positions.A trainable τ adjusts the focus of copying; top-most hidden states are used as source representations to cache global context.Source: §3.3.1 · p. 5

3.3.2 Learning from AT Experts

  • Because conditional independence limits NAT's ability to capture multimodal target distributions, the learner imitates AT action distributions step by step.The AT distribution serves as the posterior supervisory signal, and the NAT minimizes cross-entropy between the two policies.Source: §3.3.2 · p. 5

3.4 Training

  • Training combines action-distribution regularization for the AT demonstrator and imitation learning for the NAT learner with their commonly used cross-entropy losses.The two regularization weights are set to 0.001 in the experiments.Source: §3.4 · p. 5

4 Experiments

Experiments evaluate the method on three public translation corpora, with model settings varying by dataset. Inference uses length parallel decoding to generate candidates at different target lengths and rescore them with a pretrained autoregressive model.

Datasets

  • The evaluation covers IWSLT16 En-De (196K pairs), WMT14 En-De (4.5M pairs), and WMT16 En-Ro (610K pairs).The datasets use Moses tokenization and 32k BPE subword symbols.Source: §4 · p. 6

Model Settings

  • The autoregressive demonstrator is trained first and frozen while imitate-NAT is trained; the NAT model is initialized from corresponding expert parameters.The WMT tasks use the base Transformer settings, while IWSLT16 uses a smaller model; sequence-level distillation uses beam size 4.Source: §4 · p. 6

Length Parallel Decoding

  • Length parallel decoding predicts a target length, generates candidates across a range around it, then uses a pretrained autoregressive model to select the best translation.The range is [T̂ − ΔT, T̂ + ΔT].Source: §4 · pp. 6–7

5 Results and Analysis

Imitate-NAT with LPD achieves the best translation performance on the benchmark datasets, approaching autoregressive quality while retaining a speed advantage. The analyses report gains from imitation, soft copy, length sampling, knowledge distillation, and category redistribution.

5 Results and Analysis

  • Imitate-NAT with LPD achieves the best translation performance on the benchmarks, scoring 30.68 vs. 30.85 on IWSLT16 En→De and 31.81 vs. 32.59 on WMT16 Ro→En against autoregressive models.The authors describe these results as close to the autoregressive model.Source: §5 · p. 7
  • Imitation improves translation quality by over almost 3 BLEU on average versus a competitor trained with the same steps, initialization, distillation, and Soft-Copy.The imitation module is the only stated difference between the systems.Source: §5 · p. 7
  • 9.7× speedup over the autoregressive model accompanies imitate-NAT’s stronger performance and faster inference than NAT-FT and NAT-IR.The passage reports refinement-related speedup drops from 15.6× to 2.36× for NAT-FT and from 8.9× to 1.5× for NAT-IR.Source: §5 · p. 7

5.1 Ablation Study

  • 25.34 vs. 20.71: differentiable Soft-Copy outperforms Uniform-Copy, which hard-copies source embeddings without global information.The authors say hard copying increases the decoder’s learning burden.Source: §5.1 · p. 7, Table 2 · p. 8
  • On IWSLT16 dev set, imitation learning improves BLEU by around 3 points (28.41 vs. 25.34) over non-imitation learning; without IL, NAT degenerates into a normal NAT model.The authors connect NAT’s training difficulty to delayed supervision or rewards and a large search space.Source: §5.1 · p. 7, Table 2 · p. 8
  • 30.68 vs. 28.41: Length Parallel Decoding improves performance over greedy beam search by around 2 BLEU points.The authors attribute the observed gain to sampling from the length space.Source: §5.1 · p. 7, Table 2 · p. 8

Complementary with Knowledge Distillation

  • +3.3 BLEU: imitation substantially improves non-imitation NAT with knowledge distillation; sequence-level distillation alone yields +4.2 BLEU for NAT models.Imitate-NAT without distillation scores 23.56 BLEU, comparable to non-imitation NAT with distillation.Source: Complementary with Knowledge… · p. 7, Table 2 · p. 8
  • Category redistribution makes action assignments more balanced; without it, usage concentrates in a few spatial locations.The paper evaluates action space usage on the IWSLT16 development set.Source: Figure 4 · p. 8, Complementary with Knowledge… · p. 8

6 Related Work

Prior NAT work sought parallel decoding speed while addressing translation-quality losses through changes to architectures, inputs, or latent representations. This paper instead introduces imitation learning to close the NAT–AT performance gap.

6 Related Work

  • Gu et al. (2017) introduced parallel-output NAT for faster inference, but removing intrinsic target dependencies sacrifices translation quality.Later approaches explored iterative refinement, connectionist temporal classification, discrete latent variables, phrase-table inputs, and dual-task reconstruction.Source: §6 · p. 8
  • Earlier NAT methods addressed quality or latency through iterative refinement, discrete latent variables, phrase-table inputs, or dual-task reconstruction.The passage also describes a connectionist temporal classification approach as improving latency.Source: §6 · p. 8
  • Unlike work modifying NAT architecture or decoder inputs, this paper introduces imitation learning to close the performance gap with autoregressive translation.The authors state that, to their knowledge, this is the first application of imitation learning to these problems.Source: §6 · p. 8

7 Conclusion

Imitate-NAT uses an autoregressive demonstrator to supervise NAT decoding states, and largely closes the performance gap between NAT and autoregressive models on benchmark datasets. The authors identify stronger demonstrators and other applications as future directions.

7 Conclusion

  • A knowledgeable autoregressive demonstrator supervises NAT decoding states across different time steps and layers.Source: §7 · p. 8
  • Imitate-NAT largely closes the performance gap between NAT and autoregressive models on several benchmark datasets.The conclusion describes this as a result of the proposed imitation-learning framework.Source: §7 · pp. 8–9
  • Future work could use more powerful demonstrators, such as a right-to-left model, or apply the framework to simultaneous interpretation.Source: §7 · p. 9
Loading 1906.02041v2…