Aligned Cross Entropy for Non-Autoregressive Machine Translation

Introduction

Non-autoregressive machine translation models can significantly improve decoding speed by predicting every word in parallel (Gu et al., 2018; Libovický & Helcl, 2018). This advantage comes at a cost to performance since modeling word order is trickier when the model cannot condition on its previous predictions. A range of semi-autoregressive models (Lee et al., 2018; Stern et al., 2019; Gu et al., 2019; Ghazvininejad et al., 2019) have shown there is a speed-accuracy tradeoff that can be optimized with limited forms of autoregression. However, increasing performance of the purely non-autoregressive models without sacrificing decoding speed remains an open challenge. In this paper, we present a new training loss for non-autoregressive machine translation that softens the penalty for word order errors, and significantly improves performance with no modification to the model or to the decoding algorithm.

Existing models (both autoregressive and non-autoregressive) are typically trained with cross entropy loss. Cross entropy is a strict loss function, where a penalty is incurred for every word that is predicted out of position, even for output sequences with small edit distances (see Figure 1). Autoregressive models learn to avoid such penalties, since words are generated conditioned on the sentence prefix. However, non-autoregressive models do not know the exact sentence prefix, and should (intuitively) focus more on root errors (e.g. a missing word) while allowing more partial credit for cascading errors (the right word in the wrong place).

To achieve this more relaxed loss, we introduce aligned cross entropy (AXE), a new objective function that computes the cross entropy loss based on an alignment between the sequence of token labels and the sequence of token distribution predictions. AXE uses dynamic programming to find the monotonic alignment that minimizes the cross entropy loss. It provides non-autoregressive models with a more accurate training signal by ignoring absolute positions and focusing on relative order and lexical matching. We efficiently implement AXE via matrix operations, and use it to train conditional masked language models (CMLM; Ghazvininejad et al., 2019) for machine translation. AXE only slightly increases training time compared to cross entropy, and requires no changes to parallel argmax decoding.

Extensive experiments on machine translation benchmarks demonstrate that AXE substantially boosts the performance of CMLMs, while having the same decoding speed. In WMT’14 EN-DE, training CMLMs with AXE (instead of the regular cross entropy loss) increases performance by 5 BLEU points; we observe similar trends in WMT’16 EN-RO and WMT’17 EN-ZH. Moreover, AXE CMLMs significantly outperform state-of-the-art non-autoregressive models, such as FlowSeq (Ma et al., 2019), as well as the recent CRF-based semi-autoregressive model with bigram LM decoding (Sun et al., 2019). Our detailed analysis suggests that training with AXE makes models more confident in their predictions, thus reducing multimodality, and alleviating a key problem in non-autoregressive machine translation.

Aligned Cross Entropy

Let $Y$ be a target sequence of $n$ tokens $Y_{1},\ldots,Y_{n}$, and $P$ be the model predictions, a sequence of $m$ token probability distributions $P_{1},\ldots,P_{m}$. Our goal is to find a monotonic alignment between $Y$ and $P$ that will minimize the cross entropy loss, and thus focus the penalty on lexical errors (predicting the wrong token) rather than positional errors (predicting the right token in the wrong place).

We define an alignment $\alpha$ to be a function that maps target positions to prediction positions, i.e. $\alpha:\{1,\ldots,n\}\rightarrow\{1,\ldots,m\}$. We further assume that this alignment is monotonic, i.e. $i\leq j$ iff $\alpha(i)\leq\alpha(j)$. Given a specific alignment $\alpha$, we define a conditional $AXE$ loss as:

The first term of this loss function is an aligned cross entropy between $Y$ and $P$, and the second term is a penalty for unaligned predictions. Epsilon ($\varepsilon$) is a special “blank” token in our vocabulary that appears in the probability distributions, but that does not appear in the final output string.

Now, the final $AXE$ loss is the minimum over all possible monotonic alignments of the conditional loss:

Finding the optimal monotonic alignment between two sequences is a well studied problem. For instance, dynamic time warping (DTW) (Sakoe & Chiba, 1978) is a well-known algorithm for finding the optimal alignment between two different time series. Here we have extended the idea to compute the optimal alignment between a sequence of target tokens and a sequence of prediction probability distributions. We use a simple dynamic program to find the optimal alignment while calculating the AXE loss.

Given a sequence of target tokens $Y=Y_{1}\ldots Y_{n}$ and a sequence of predictions $P=P_{1}\ldots P_{m}$ we propose a method to find the score of the optimal alignment between any prefix of these two sequences $Y_{1:i}=Y_{1}\ldots Y_{i}$ and $P_{1:j}=P_{1}\ldots P_{j}$, for any $i$ and $j$. The score of the optimal alignment for the full sequences is obtained at $i=n$ and $j=m$.

We start by defining a matrix $A$ of $n+1$ by $m+1$ dimensions, respectively corresponding to $Y$ and $P$, where $A_{i,j}$ represents the minimum loss value for aligning $Y_{1:i}$ to $P_{1:j}$ as defined in Equation 2. We initialize $A_{0,0}$ to be and then proceed to fill the matrix by taking the local minimum at each cell $A_{i,j}$ from three possible operators: Align, Skip Prediction , and Skip Target. Table 1 describes each operation and its update formula. Once the matrix is full, the cell $A_{n,m}$ will contain the cross entropy loss of the optimal alignment. Algorithm 1 lays out a straightforward implementation of AXE’s dynamic program.

According to Equation 2, the optimal alignment can be many-to-one, where multiple target positions can be mapped to a single prediction. This would be computed by aligning the first mapped token and skipping the rest of target tokens. To discourage skipping too many target tokens, we penalize skip target operators separately with a parameter $\delta$ as described in Table 1. Setting $\delta=1$ will result in the loss function defined in Equation 2, but as we show in our ablation study (Section 4.3), higher $\delta$ values yield better performance in practice.

Efficient Implementation

The implementation in Algorithm 1 has $O(n\cdot m)$ time complexity. However, multiple updates of the matrix $A$ can be parallelized on GPUs and other tensor-processing architectures. Rather than iterating over each cell, we iterate over each anti-diagonal, computing all the values along the anti-diagonal in parallel. In other words, we first compute the values of $[A_{0,1},A_{1,0}]$, followed by $[A_{0,2},A_{1,1},A_{2,0}]$, etc. Since the number of anti-diagonals is $n+m+1$, we arrive at a time complexity of $O(n+m)$. Since $m$ is typically on the same order of magnitude as $n$, the linear cost of computing AXE during training becomes negligible compared to forward and backward passes through the model.Batch implementation of this algorithm is straightforward. By doing so, we are able to achieve training times similar to (about 1.2 times slower than) training with cross entropy loss.

Example

Figure 2 depicts an example application of AXE. We see that the predictions are generally good, but start with a shift with respect to the target. This misalignment would cause the regular cross entropy loss to severely penalize the first three predictions, even though $P_{2}$ and $P_{3}$ are correct when aligned with $Y_{1}$ and $Y_{2}$. AXE, on the other hand, finds an alignment between the target and the predictions, which allows it to focus the penalty on the redundant prediction in $P_{1}$ and the missing token $Y_{3}$, i.e. the root errors.

Training Non-Autoregressive Models

We use AXE to train conditional masked language models (CMLMs) for non-autoregressive machine translation (Ghazvininejad et al., 2019).While in this work we apply AXE to CMLMs, the loss function can be used to train other models as well. We leave further investigation of this direction to future work.

A conditional masked language model takes a source sequence $X$ and a partially-observed target sequence $Y_{\text{obs}}$ as input, and predicts the probabilities of the masked (unobserved) target sequence tokens $Y_{\text{mask}}$. The underlying architecture is an encoder-decoder transformer (Vaswani et al., 2017).

In the original paper, CMLMs are used for machine translation where a random subset of $Y$ tokens are masked at training time. However, at inference all target tokens are masked ($Y=Y_{\text{mask}}$) and the length of $Y$ (the number of masked tokens) is unknown. To estimate the length of $Y$, an auxiliary task is introduced to predict the target length based on the source sequence $X$.See (Ghazvininejad et al., 2019) for further detail.

2 Adapting CMLMs to AXE

In our case, the model can also produce blank tokens ($\varepsilon$), which effectively shorten the predicted sequence’s length. To account for potentially skipped tokens during inference, we multiply the predicted length by a hyperparameter $\lambda$ (which is tuned on the validation set) before applying argmax decoding.

3 Adapting the Training Objectives to AXE

Since this work focuses on the purely non-autoregressive setting, the entire target sequence will be masked at inference time ($Y_{\text{mask}}=Y$). The same does not have to hold for training; we can utilize partially observed sequences in order to provide the learner with easier and more focused training examples. We experiment with three variations:

All the tokens in the target sequence are masked, and the model is expected to predict all of them. This is a direct replication of the task at inference time. While AXE allows for the number of masked tokens $m$ to be different from the length of the gold target sequence $n$, we found that setting $m=n$ produced better models in preliminary experiments.

Partially-Observed Input, Predict All

As in the original CMLM training process, a random subset of the target sequence is masked before being passed onto the model as input.The number of masked input tokens is distributed uniformly between $1$ and $n$. We then apply AXE on the entire sequence, regardless of which tokens were observed. When training on partially-observed inputs, we always set $m=n$ to avoid further alterations of the gold target sequence beyond masking.

Partially-Observed Input, Predict Masks

The straightforward application of AXE to CMLM training (which ignores whether each token was masked or observed) works well in practice. However, we can also allow AXE to skip the observed tokens when computing cross entropy, and focus the training signal on the actual task. We do so by setting $P_{i}(Y_{i})=1$ for every observed token $Y_{i}$; i.e. if the $i$-th token is observed and is aligned with the prediction corresponding to the same position ($P_{i}$), there is no penalty. Our ablation studies show that this modification provides a modest but consistent boost in performance (see Section 4.3). As a result, we use this setting for training our model.

Experiments

We evaluate CMLMs trained with AXE on 6 standard machine translation benchmarks, and demonstrate that AXE significantly improves performance over cross entropy trained CMLMs and over recently-proposed non-autoregressive models as well.

We evaluate our method on both directions of three standard machine translation datasets with various training data sizes: WMT’14 English-German (4.5M sentence pairs), WMT’16 English-Romanian (610k pairs), and WMT’17 English-Chinese (20M pairs). The datasets are tokenized into subword units using BPE (Sennrich et al., 2016).We run joint BPE for all language pairs except English-Chinese. We use the same data and preprocessing as Vaswani et al. (2017), Lee et al. (2018), and Wu et al. (2019) for WMT’14 EN-DE, WMT’16 EN-RO, and WMT’17 EN-ZH respectively. We evaluate performance with BLEU (Papineni et al., 2002) for all language pairs, except for translating English to Chinese where we use SacreBLEU (Post, 2018).SacreBLEU hash: BLEU+case.mixed+lang.en-zh +numrefs.1+smooth.exp+test.wmt17+tok.zh+version.1.3.7

Hyperparameters

Knowledge Distillation

Similar to previous work on non-autoregressive translation (Gu et al., 2018; Lee et al., 2018; Ghazvininejad et al., 2019; Stern et al., 2019), we use sequence-level knowledge distillation (Kim & Rush, 2016) by training CMLMs on translations generated by a standard left-to-right transformer model (transformer large for WMT’14 EN-DE and WMT’17 EN-ZH, transformer base for WMT’16 EN-RO). We report the performance of standard autoregressive base transformers trained on distilled data for WMT’14 EN-DE and WMT’17 EN-ZH.

2 Main Results

We first compare the performance of AXE-trained CMLMs to that of CMLMs trained with the original cross entropy loss. Table 2 shows that training with AXE substantially increases the performance CMLMs across all benchmarks. On average, we gain 5.2 BLEU by replacing cross entropy with AXE, with gains of up to 6.65 BLEU in WMT’17 EN-ZH.

State of the Art

We compare the performance of CMLMs with AXE against nine strong baseline models: the fertility-based sequence-to-sequence model (Gu et al., 2018), transformers trained with CTC loss (Libovický & Helcl, 2018), the iterative refinement approach (Lee et al., 2018), transformers trained with auxiliary regularization (Wang et al., 2019), CMLMs trained with (regular) cross entropy loss (Ghazvininejad et al., 2019), Flowseq: a latent variable model based on generative flow (Ma et al., 2019), hint-based training (Li et al., 2019), bag-of-ngrams training (Shao et al., 2019), and the CRF-based semi-autoregressive model (Sun et al., 2019). All of these models except the last one are purely non-autoregressive, while the CRF-based model uses bigram statistics during decoding, which deviates from the purely non-autoregressive setting.CMLMs (Ghazvininejad et al., 2019) and the iterative refinement method (Lee et al., 2018) are presented as semi-autoregressive models that run in multiple decoding iterations. However, the first decoding iteration of these models is purely non-autoregressive, which is what we use as our baselines.

Table 3 shows that our system yields the highest BLEU scores of all non-autoregressive models. AXE-trained CMLMs outperform the best purely non-autoregressive model (FlowSeq) on both directions of WMT’14 EN-DE and WMT’16 EN-RO by 1.6 BLEU on average. Moreover, our approach achieves higher BLEU scores than the semi-autoregressive CRF decoder across all available benchmarks.

Raw Data

Finally, we compare the performance of AXE to other methods that train on raw data without knowledge distillation. Table 4 shows that AXE CMLMs still significantly outperform other non-autoregressive models in the raw data scenario. In addition, comparing raw data to knowledge distillation training follows previously-published results that demonstrate the importance of knowledge distillation for non-autoregressive approaches (Gu et al., 2018; Ghazvininejad et al., 2019; Zhou et al., 2019), although the gap is much smaller for WMT’16 EN-RO.

3 Ablation Study

Different Training Objectives Table 5 shows the effects of different training objectives (Section 3.3), in which all or part of the target tokens are masked and the loss function is calculated on all tokens or masked tokens only. We find that simulating the inference scenario, where all tokens are unobserved, is actually less effective than revealing a subset of the target tokens as input during training. We speculate that partially-observed inputs add easier examples to the training set, allowing for better optimization as in curriculum learning (Bengio et al., 2009). We also see that including only the masked tokens in the loss function gives us a modest but consistent boost in performance, possibly because the training signal is focused on the actual task.

Skip Target Penalty The hyperparameter $\delta$ acts as a coefficient for the penalty associated with skipping a target token (see Table 1 for a definition). We experiment with different values of $\delta$, and report our findings in Table 6. We observe that tuning $\delta$ can significantly improve performance with respect to the default of $\delta=1$. As intended, high values of $\delta$ discourage alignments that skip too many target tokens.

Length Multiplier The length multiplier $\lambda$ inflates the length predicted by a CMLM to account for extra blank tokens ($\varepsilon$) that the model could potentially generate (see Section 3.2 for more detail). Table 7 compares the effect of different length multiplier $\lambda$ values. Using the best length multiplier increases the performance by 0.53 BLEU on average for WMT’14 EN-DE and WMT’16 EN-RO.

Analysis

We provide a qualitative analysis to provide some insight where AXE improves over cross entropy, and potential directions for future research on non-autoregressive generation.

We first measure performance of cross entropy versus AXE-trained CMLMs for different sequence lengths. We use compare-mt (Neubig et al., 2019) to split the test sets of WMT’14 EN-DE and DE-EN into different buckets based on target sequence length and calculate BLEU for each bucket. Table 8 shows that the performance of models trained with cross entropy drops drastically as the sequence length increases, while the performance of AXE-trained models remains relatively stable. One explanation for this result is that the longer the sequence, the more likely we are to observe misalignments between the model’s predictions and the target; AXE realigns these cases, providing the model with a cleaner signal for modeling long sequences.

AXE Increases Position Confidence

We also study how confident each model is about the position of each generated token. Ideally, we would like each predicted token to have a high probability at the position in which it was predicted and a very low probability in the neighboring positions. After applying argmax decoding, we compute the probability assigned to each generated token in all positions of the sequence and average these probabilities based on the relative distance (positive or negative) to the generated position. Figure 3 plots these averaged probabilities for both short ($<10$ tokens) and long ($>30$ tokens) target sequences.

Both models are rather confident in their predictions for short sequences (Figure 3(a)): the probability has a high peak at the generated position and drops rapidly as we move further away. However, for longer sentences (Figure 3(b)), we observe that the plot for cross entropy has lost its sharpness. Specifically, the immediate neighbors of the prediction position ($\pm 1$) receive about $0.14$ probability on average, almost a third of the peak probability. Meanwhile, the probabilities predicted by the AXE-trained model are significantly sharper, assigning negligible probabilities to the generated token in neighboring positions when compared to the center.

On way to explain this result is that cross entropy training encourages predictions to have some probability mass of their neighbors, in order to “hedge their bets” in case the predictions are misaligned with the target. Since AXE finds the best alignment before computing the actual loss, spreading the probability mass of a token among its neighbors is no longer necessary.

AXE Reduces Multimodality

We further argue that AXE reduces the multimodality problem in non-autoregressive machine translation (Gu et al., 2018). Due to minimal coordination between predictions in many non-autoregressive models, a model might consider many possible translations at the same time. In this situation, the model might merge two or more different translations and generate an inconsistent output that is typically characterized by token repetitions. We therefore use the frequency of repeated tokens as a proxy for measuring multimodality in a model.

Table 9 shows the repetition rate for cross entropy and AXE-trained CMLMs. Replacing cross entropy with AXE drastically reduces multimodality, decreasing the number of repetitions by a multiplicative factor of 12.

Related Work

Advances in neural machine translation techniques in recent years has brought an increasing interest in breaking the autoregressive generation bottleneck in translation models.

Semi-autoregressive models introduce partial parallelism into the decoding process. Some of these techniques include iterative refinement of translations based on previous predictions (Lee et al., 2018; Ghazvininejad et al., 2019, 2020; Gu et al., 2019; Kasai et al., 2020) and combining a lighter autoregressive decoder with a non-autoregressive one (Sun et al., 2019).

Building a fully non-autoregrssive machine translation model is a much more challenging task. One branch of prior work approaches this problem by modeling with latent variables. Gu et al. (2018) introduces word fertility as a latent variable to model the number of generated tokens per each source word. Ma et al. (2019) uses generative flow to model complex distribution of latent variables for parallel decoding of target. Shu et al. (2019) proposes a latent-variable non-autoregressive model with continuous latent variables and a deterministic inference procedure.

There is also work that develops other alternative loss functions for non-autoregressive machine translation. Libovický & Helcl (2018) use the Connectionist Temporal Classification training objective, a loss function from the speech recognition literature that is designed to eliminating repetitions. Li et al. (2019) uses the learning signal provided by hidden states and attention distributions of an autoregressive teacher. Yang et al. (2019) improves the decoder hidden representations by adding the reconstruction error of source sentence from these representations as an auxiliary regularization term to the loss function. Finally, Shao et al. (2019) introduce the bag-of-ngrams training objective to encourage the model to capture target-side sequential dependencies.

Conclusion

We introduced Aligned Cross Entropy (AXE) as an alternative loss function for training non-autoregressive models. AXE focuses on relative order and lexical matching instead of relying on absolute positions. We showed that, in the context of machine translation, a conditional masked language model (CMLM) trained with AXE significantly outperforms cross entropy trained models, setting a new state-of-the-art for non-autoregressive models.

Acknowledgements

We thank Abdelrahman Mohamed for sharing his expertise on non-autoregressive models, and our colleagues at FAIR for valuable feedback.

References