[Paper Review] Non-Autoregressive Machine Translation with Disentangled Context Transformer

Jungo Kasai, James Cross+2 more|arXiv (Cornell University)|Jan 15, 2020

Natural Language Processing Techniques43 references51 citations

TL;DR

This paper introduces the DisEntangled Context (DisCo) transformer for non-autoregressive translation, using an attention-masking objective and a parallel easy-first inference to generate all tokens simultaneously with reduced decoding steps while maintaining competitive BLEU.

ABSTRACT

State-of-the-art neural machine translation models generate a translation from left to right and every step is conditioned on the previously generated tokens. The sequential nature of this generation process causes fundamental latency in inference since we cannot generate multiple tokens in each sentence in parallel. We propose an attention-masking based model, called Disentangled Context (DisCo) transformer, that simultaneously generates all tokens given different contexts. The DisCo transformer is trained to predict every output token given an arbitrary subset of the other reference tokens. We also develop the parallel easy-first inference algorithm, which iteratively refines every token in parallel and reduces the number of required iterations. Our extensive experiments on 7 translation directions with varying data sizes demonstrate that our model achieves competitive, if not better, performance compared to the state of the art in non-autoregressive machine translation while significantly reducing decoding time on average. Our code is available at https://github.com/facebookresearch/DisCo.

Motivation & Objective

• Motivate reducing decoding latency in neural machine translation by moving away from left-to-right autoregressive decoding.
• Propose a DisEntangled Context (DisCo) transformer that predicts every target token conditioned on an arbitrary subset of other tokens.
• Develop a parallel easy-first inference algorithm to refine all tokens in parallel with iteration-based convergence.
• Show that DisCo achieves competitive BLEU scores with substantial decoding-time reductions across multiple language directions and data sizes.

Proposed method

• Introduce the DisCo transformer with attention masking to attend only to observed tokens when predicting each target position.
• Define the DisCo objective: predict Y_n given X and an arbitrary subset Y_obs^n of other target tokens, enabling single-pass computation of conditional probabilities for all positions.
• Describe how to stack DisCo layers while avoiding leakage by decontextualizing keys/values from previous layers.
• Train with random subsets of observed tokens and include a length prediction loss to enable parallel decoding.
• Propose parallel easy-first inference: in each iteration, predict all positions and update tokens in order of increasing uncertainty, allowing variable numbers of iterations.
• Utilize distillation from a strong autoregressive teacher and standard transformer hyperparameters; assess with BLEU on multiple WMT directions.

Experimental results

Research questions

• RQ1 Can a non-autoregressive Transformer using disentangled context achieve competitive BLEU compared to state-of-the-art NAT and autoregressive models?
• RQ2 Does the DisCo objective enable efficient single-pass conditioning and effective parallel decoding?
• RQ3 How does the parallel easy-first inference compare to mask-predict in terms of BLEU and number of iterations across data sizes?
• RQ4 What is the impact of data size and distillation on DisCo’s performance relative to baselines?
• RQ5 How do different decoding strategies affect speed and quality for NAT on WMT tasks?

Key findings

• DisCo with parallel easy-first achieves competitive to better BLEU than CMLM-based Mask-Predict while using significantly fewer iterations (e.g., en→de 4.82 steps; ro→en 3.10 steps).
• On EN-DE / EN-RO, DisCo+Easy-First reaches BLEU scores comparable to or better than strong NAT baselines, with large gains when data are plentiful (EN-ZH, EN-FR).
• Distillation consistently benefits non-autoregressive models, with DisCo gaining more from distillation than CMLM under the same inference settings.
• Decoding speed shows substantial wall-clock gains; average iterations correlate with speedup, with DisCo achieving about a 4–5x reduction in iterations versus autoregressive baselines depending on direction and setup.
• DisCo with contextless keys/values can preserve performance even in autoregressive settings, suggesting broader applicability of the approach.
• Training variants that more closely align training and inference (easy-first training) did not outperform random-sampling training, indicating random masking provides useful exploration.