Aligned Cross Entropy for Non-Autoregressive Machine Translation

1. Introduction

• Cross entropy penalizes every out-of-position word, even when the output has a small edit distance; non-autoregressive models lack the sentence prefix that helps autoregressive models avoid these penalties.The introduction argues they should focus more on root errors while allowing partial credit for cascading errors.Source: §1 · p. 1
• AXE assigns cross-entropy loss using the best monotonic alignment between target tokens and prediction distributions, emphasizing relative order and lexical matching over absolute positions.A dynamic program finds the alignment that minimizes loss.Source: §1 · p. 1
• Parallel prediction significantly speeds decoding, and AXE only slightly increases training time without changing parallel argmax decoding.The paper presents the loss as improving performance without sacrificing decoding speed.Source: §1 · p. 1
• AXE-trained CMLMs gain 5 BLEU points on WMT’14 EN-DE and significantly outperform state-of-the-art non-autoregressive models, with similar trends on WMT’16 EN-RO and WMT’17 EN-ZH.AXE CMLMs also outperform FlowSeq and a CRF-based semi-autoregressive model with bigram LM decoding.Source: §1 · pp. 1–2

• AXE uses dynamic programming to find the minimum-loss monotonic alignment between target tokens and prediction distributions, focusing penalties on lexical rather than positional errors.The score for the full sequences is the final cell of a matrix of prefix-alignment losses.Source: §2 · p. 2
• Each dynamic-programming cell chooses among aligning tokens, skipping a prediction, or skipping a target; δ weights target skips and permits many-to-one alignments.A target skip is implemented by aligning the first mapped token and skipping the remaining targets.Source: §2 · pp. 2–3
• Higher δ values improve performance over the default δ = 1 in the reported ablation, while discouraging alignments that skip many target tokens.The parameter controls the penalty associated with skipping a target token.Source: §2 · p. 3
• Anti-diagonal parallelization reduces computation from O(n·m) sequential updates to O(n+m); batch training is about 1.2 times slower than cross-entropy training.The paper notes that m is typically on the same order as n, making AXE's training cost negligible relative to model passes.Source: §2 · p. 3
• In the illustrated shifted sequence, AXE skips the redundant first prediction and missing target, while aligning the correctly ordered words around them.This example contrasts AXE's focus on root errors with cross-entropy penalties on the first three shifted predictions.Source: §2 · p. 3, Figure 2 · p. 3

3. Training Non-Autoregressive Models

• AXE is used to train conditional masked language models (CMLMs) for non-autoregressive machine translation.Source: §3 · p. 3

3.1. Conditional Masked Language Models

• A CMLM takes a source sequence and a partially observed target as input and predicts masked target tokens with an encoder-decoder Transformer.At inference, all target tokens are masked, so an auxiliary task predicts target length from the source.Source: §3.1 · pp. 3–4
• Although AXE can train other model types, this work applies it only to CMLMs and leaves broader investigation for future work.Source: §3.1 · p. 3

3.2. Adapting CMLMs to AXE

• At inference, blank-token predictions can shorten the output, so the predicted length is multiplied by a validation-tuned λ before argmax decoding.Source: §3.2 · p. 4

3.3. Adapting the Training Objectives to AXE

• Training compares fully masked inputs with partial masking; for partial inputs, AXE either predicts all tokens or focuses its loss on masked tokens.The masked-token-focused variant sets the probability of an observed token at its matching position to 1, and the authors use it after finding a modest, consistent performance boost.Source: §3.3 · p. 4

4.1. Setup

• Evaluation covers both directions of WMT’14 English-German, WMT’16 English-Romanian, and WMT’17 English-Chinese, using BLEU except SacreBLEU for English-to-Chinese.The experiments also use sequence-level knowledge distillation from standard left-to-right transformers to train CMLMs.Source: §4.1 · pp. 4–5

4.2. Main Results

• Across benchmarks, AXE-trained CMLMs improve over cross-entropy CMLMs, with gains up to 6.65 BLEU on WMT’17 EN-ZH.Table 2 reports test-set BLEU comparisons; both models use a single forward pass during argmax decoding.Source: §4.2 · p. 5, Table 2 · p. 5
• AXE CMLMs achieve the highest BLEU scores among non-autoregressive models, beat FlowSeq by 1.6 BLEU on average on WMT’14 and WMT’16 directions, and exceed the semi-autoregressive CRF decoder across available benchmarks.Table 3 compares test-set BLEU with non-autoregressive methods and reports standard autoregressive transformer results for reference.Source: §4.2 · pp. 5–6, Table 3 · p. 5
• On raw data without knowledge distillation, AXE CMLMs still significantly outperform other non-autoregressive models; the raw-versus-distilled gap is smaller for WMT’16 EN-RO.Table 4 compares test-set BLEU for raw-data models and includes distilled AXE CMLM as a reference.Source: §4.2 · p. 6, Table 4 · p. 6

4.3. Ablation Study

• Partially observed target tokens during training outperform fully masked inputs, and computing loss only on masked tokens gives a modest, consistent boost.The authors suggest partially observed inputs add easier examples, while masked-token loss focuses the training signal on the task.Source: §4.3 · p. 6, Table 5 · p. 6
• Tuning skip penalty δ improves performance over δ = 1, while the best length multiplier adds 0.53 BLEU on average across WMT’14 EN-DE and WMT’16 EN-RO.Higher δ discourages alignments that skip too many target tokens; the ablations use validation sets.Source: §4.3 · p. 6, Table 6 · p. 6

Length multiplier analysis

• The analysis also evaluates the effect of tuning the length multiplier λ on validation-set BLEU.Table 7 presents the effect of λ tuning on performance.Source: Table 7 · p. 7

AXE Handles Long Sequences Better

• As target sequences grow longer, cross-entropy CMLM performance drops drastically while AXE performance remains relatively stable.The authors suggest AXE realigns prediction-target misalignments, providing a cleaner signal for modeling long sequences.Source: §5 · p. 7

AXE Increases Position Confidence

• For long sequences, cross-entropy assigns about 0.14 probability to immediate neighboring positions, while AXE assigns negligible probabilities there.The analysis averages each generated token’s probability by its relative distance from the generated position, comparing short (< 10 tokens) and long (> 30 tokens) sequences; Figure 3 plots cross-entropy as a dashed blue line and AXE as a solid red line.Source: §5 · p. 7, Figure 3 · p. 7

AXE Reduces Multimodality

• AXE reduces the number of repeated tokens by a multiplicative factor of 12, using repetition frequency as a proxy for multimodality.Table 9 reports repetition percentages on WMT’14 EN-DE and DE-EN test sets.Source: §5 · p. 8, Table 9 · p. 8

Semi-autoregressive models

• Semi-autoregressive approaches use iterative refinement based on previous predictions or combine a lighter autoregressive decoder with a non-autoregressive one.These approaches introduce partial parallelism into decoding.Source: §6 · p. 8

Latent-variable models

• Latent-variable approaches include word fertility, generative flows, and continuous latent variables with deterministic inference.Word fertility models the number of generated tokens per source word.Source: §6 · p. 8

Alternative loss functions

• Other proposed objectives include Connectionist Temporal Classification, source-reconstruction regularization, and bag-of-ngrams training.The source-reconstruction term regularizes decoder hidden representations, while bag-of-ngrams encourages target-side sequential dependencies.Source: §6 · p. 8

Conclusion

• AXE-trained CMLMs significantly outperform cross-entropy-trained models and set a new state of the art for non-autoregressive machine translation.AXE focuses on relative order and lexical matching instead of relying on absolute positions.Source: §7 · p. 8