[Paper Review] Aligned Cross Entropy for Non-Autoregressive Machine Translation

Marjan Ghazvininejad, Vladimir Karpukhin+2 more|arXiv (Cornell University)|Apr 3, 2020

Natural Language Processing Techniques30 references68 citations

TL;DR

AXE is a new training loss for non-autoregressive MT that aligns target tokens with model predictions via a differentiable dynamic program, improving BLEU significantly over cross-entropy trained CMLMs while maintaining fast decoding.

ABSTRACT

Non-autoregressive machine translation models significantly speed up decoding by allowing for parallel prediction of the entire target sequence. However, modeling word order is more challenging due to the lack of autoregressive factors in the model. This difficultly is compounded during training with cross entropy loss, which can highly penalize small shifts in word order. In this paper, we propose aligned cross entropy (AXE) as an alternative loss function for training of non-autoregressive models. AXE uses a differentiable dynamic program to assign loss based on the best possible monotonic alignment between target tokens and model predictions. AXE-based training of conditional masked language models (CMLMs) substantially improves performance on major WMT benchmarks, while setting a new state of the art for non-autoregressive models.

Motivation & Objective

• Motivate and address the difficulty of modeling word order in non-autoregressive MT under standard cross-entropy loss.
• Introduce Aligned Cross Entropy (AXE) as a differentiable, alignment-based loss.
• Demonstrate AXE's effectiveness by training conditional masked language models (CMLMs) and achieving state-of-the-art non-autoregressive MT results on major benchmarks.
• Analyze how AXE affects model confidence, multimodality, and handling of long sequences.

Proposed method

• Define AXE as a loss computed over the best monotonic alignment between target Y and predictions P.
• Use a three-operator dynamic programming update (Align, Skip Prediction, Skip Target) to compute AXE for any alignment.
• Minimize AXE over all monotonic alignments; implement efficiently with anti-diagonal parallelization for O(n+m) time per sequence.
• Adapt AXE to CMLMs by allowing blank tokens (epsilon) and predicting target length, with length inflation via a hyperparameter lambda at decoding.
• Experiment with training variations: Unobserved Input vs Partially-Observed Input, and predicting all vs predicting masks, to identify effective training setups.

Experimental results

Research questions

• RQ1 Can AXE reduce the harsh penalties on word order misalignment inherent in standard cross-entropy training for non-autoregressive MT?
• RQ2 Does AXE lead to improved translation quality (BLEU) on standard WMT benchmarks using CMLMs?
• RQ3 How do AXE-trained models compare to other non-autoregressive approaches and autoregressive baselines in both raw and distilled data settings?
• RQ4 What training choices (e.g., delta skip penalty, length multiplier lambda, partially-observed inputs) most influence AXE performance?

Key findings

• AXE-trained CMLMs achieve substantial BLEU gains over cross-entropy CMLMs (average +5.2 BLEU across benchmarks).
• AXE CMLMs outperform state-of-the-art non-autoregressive models such as FlowSeq and competitive semi-autoregressive systems on WMT benchmarks.
• AXE maintains decoding speed identical to non-autoregressive decoding with a modest training-time overhead (~1.2x).
• AXE reduces multimodality and repetitions, and yields more confident position-wise predictions, especially on longer sequences.
• Experiments show AXE improves longer sequence handling and reduces reliance on neighboring token probabilities compared to cross-entropy training.
• Ablations indicate that partially-observed inputs with loss computed on masked tokens, and tuned delta and lambda values, enhance performance.