Self-supervised and Supervised Joint Training for
Resource-rich Machine Translation

Under the encoder-decoder paradigm  (Bahdanau et al., 2015; Gehring et al., 2017; Vaswani et al., 2017), the conditional probability $P(\mathbf{y}|\mathbf{x};\bm{\theta})$ of a target-language sentence $\mathbf{y}=y_{1},\cdots,y_{J}$ given a source-language sentence $\mathbf{x}=x_{1},\cdots,x_{I}$ is modeled as follows: The encoder maps the source sentence $\mathbf{x}$ onto a sequence of $I$ word embeddings $e(\mathbf{x})=e(x_{1}),...,e(x_{I})$. Then the word embeddings are encoded into their corresponding continuous hidden representations. The decoder acts as a conditional language model that reads embeddings $e(\mathbf{y})$ for a shifted copy of $\mathbf{y}$ along with the aggregated contextual representations $\mathbf{c}$. For clarity, we denote the input and output in the decoder as $\mathbf{z}$ and $\mathbf{y}$, i.e., $\mathbf{z}=\langle s\rangle,y_{1},\cdots,y_{J-1}$, where $\langle s\rangle$ is a start symbol. Conditioned on an aggregated contextual representation $\mathbf{c}_{j}$ and its partial target input $\mathbf{z}_{\leq j}$, the decoder generates $\mathbf{y}$ as:

The aggregated contextual representation $\mathbf{c}$ is often calculated by summarizing the sentence $\mathbf{x}$ with an attention mechanism (Bahdanau et al., 2015). A byproduct of the attention computation is a noisy alignment matrix $\mathbf{A}\in\mathbb{R}^{J\times I}$ which roughly captures the translation correspondence between target and source words (Garg et al., 2019).

Generally, NMT optimizes the model parameters $\bm{\theta}$ by minimizing the empirical risk over a parallel training set $(\mathbf{x},\mathbf{y})\in\mathcal{S}$:

where $\ell$ is the cross entropy loss between the model prediction $f(\mathbf{x},\mathbf{y};\bm{\theta})$ and $h(\mathbf{y})$, and $h(\mathbf{y})$ denotes the sequence of one-hot label vectors with label smoothing in the Transformer (Vaswani et al., 2017).

Pre-training sequence-to-sequence models for language generation has been shown to be effective for machine translation (Song et al., 2019; Lewis et al., 2019). These methods generally comprise two stages: pre-training and finetuning. The pre-training takes advantage of an abundant monolingual corpus $\mathcal{U}=\{\mathbf{y}\}$ to learn representations through a self-supervised objective called denoising autoencoder (Vincent et al., 2008) which aims at reconstructing the original sentence $\mathbf{y}$ from one of its corrupted counterparts.

Let $n(\mathbf{y})$ be a corrupted copy of $\mathbf{y}$ where the function $n(\cdot)$ adds noise and/or masks words. $(n(\mathbf{y}),\mathbf{y})$ constitutes the pseudo parallel data and is fed into the NMT model to compute the reconstruction loss. The self-supervised reconstruction loss over the corpus $\mathcal{U}$ is defined as:

The optimal model parameters $\bm{\theta}^{\star}$ are learned via the self-supervised loss $\mathcal{L}_{\mathcal{U}}(\bm{\theta})$ and used to initialize downstream models during the finetuning on the parallel training set $\mathcal{S}$.

This section introduces the crossover encoder-decoder (XEnDec). Different from a conventional encoder-decoder, XEnDec takes two training examples as inputs (called parents), shuffles the parents’ source sentences and produces a virtual example (called offspring) through a mixture decoder model. Fig. 1 illustrates this process.

Formally, let $(\mathbf{x},\mathbf{y})$ denote a training example where $\mathbf{x}=x_{1},\cdots,x_{I}$ represents a source sentence of $I$ words and $\mathbf{y}=y_{1},\cdots,y_{J}$ is the corresponding target sentence of $J$ words. In supervised training, $\mathbf{x}$ and $\mathbf{y}$ are parallel sentences. As we will see in Section 3.2, XEnDec can be carried out with and without supervision. We do not distinguish these cases for now and use generic notations to illustrate the idea.

Given a pair of examples $(\mathbf{x},\mathbf{y})$ and $(\mathbf{x}^{\prime},\mathbf{y}^{\prime})$ called parents, the crossover encoder shuffles the two source sequences into a new source sentence $\tilde{\mathbf{x}}$, calculated from:

where $\mathbf{m}=m_{1},\cdots,m_{I}\in\{0,1\}^{I}$ stands for a series of Bernoulli random variables with each taking the value 1 with probability $p$ called shuffling ratio. If $m_{i}=0$, then the $i$-th word in $\mathbf{x}$ will be substituted with the word in $\mathbf{x}^{\prime}$ at the same position. For convenience, the lengths of the two sequences are aligned by appending padding tokens to the end of the shorter sentence.

The crossover decoder employs a mixture model to generate the virtual target sentence. The embedding of the decoder’s input $\tilde{\mathbf{z}}$ is computed as:

where $e(\cdot)$ is the embedding function. $Z=\sum_{i=1}^{I}A_{(j-1)i}m_{i}+A^{\prime}_{(j-1)i}(1-m_{i})$ is the normalization term where $\mathbf{A}$ and $\mathbf{A}^{\prime}$ are the alignment matrices for the source sequences $\mathbf{x}$ and $\mathbf{x}^{\prime}$, respectively. Eq. (5) averages embeddings of $\mathbf{y}$ and $\mathbf{y}^{\prime}$ through the latent weights computed by $\mathbf{m}$, $\mathbf{A}$, and $\mathbf{A}^{\prime}$. The alignment matrix measures the contribution of the source words for generating a specific target word (Och & Ney, 2004; Bahdanau et al., 2015). For example, $A_{ji}$ represents the contribution score of the $i$-th word in the source sentence for the $j$-th word in the target sentence. For simplicity, this paper uses the attention matrix learned in the NMT model as a noisy alignment matrix (Garg et al., 2019).

Likewise, the label vector for the crossover decoder is calculated from:

The $h(\cdot)$ function projects a word onto its label vector, e.g., a one-hot vector. The loss of XEnDec is computed over its output ($\tilde{\mathbf{x}}$, $\tilde{\mathbf{y}}$) using the negative log-likelihood:

where $\tilde{\mathbf{z}}$ is a shifted copy of $\tilde{\mathbf{y}}$ as discussed in Section 2.1. Notice that even though we do not directly observe the “virtual sentences” $\tilde{\mathbf{z}}$ and $\tilde{\mathbf{y}}$, we are still able to compute the loss using their embeddings and labels. In practice, the length of $\tilde{\mathbf{x}}$ is set to $\mathop{\rm max}(|\mathbf{x}|,|\mathbf{x}^{\prime}|)$ whereas $\tilde{\mathbf{y}}$ and $\tilde{\mathbf{z}}$ share the same length of $\mathop{\rm max}(|\mathbf{y}|,|\mathbf{y}^{\prime}|)$.

The proposed method applies XEnDec to deeply fuse the parallel data $\mathcal{S}$ with nonparallel, monolingual data $\mathcal{U}$. As illustrated in Fig. 1, the first generation ($F_{1}$-XEnDec in the figure) uses XEnDec to combine monolingual sentences of different views, thereby incurring a self-supervised loss $\mathcal{L}_{F_{1}}$. We compute the loss $\mathcal{L}_{F_{1}}$ using Eq. (3). Afterward, the second generation ($F_{2}$-XEnDec in the figure) applies XEnDec to the offspring of the first generation ($n(\mathbf{y}^{u})$, $\mathbf{y}^{u}$) and a sampled parallel sentence ($\mathbf{x}^{p}$, $\mathbf{y}^{p}$), yielding a new loss term $\mathcal{L}_{F_{2}}$. The loss $\mathcal{L}_{F_{2}}$ is computed over the output of the $F_{2}$-XEnDec by:

where $(\tilde{\mathbf{x}},\tilde{\mathbf{y}})$ is the output of the $F_{2}$-XEnDec in Fig. 1.

The final NMT models are optimized jointly on the original translation loss and the above two auxiliary losses.

$\mathcal{L}_{F_{2}}$ in Eq. (9) is used to deeply fuse monolingual and parallel sentences at instance level rather than combine them mechanically. Section 4.4 empirically verifies the contributions of the $\mathcal{L}_{F_{1}}$ and $\mathcal{L}_{F_{2}}$ loss terms.

Algorithm 1 delineates the procedure to compute the final loss $\mathcal{L}(\bm{\theta})$. Specifically, each time, we sample a monolingual sentence for each parallel sentence to circumvent the expensive enumeration in Eq. (8). To speed up the training, we group sentences offline by length in Step 3 (cf. batching data in the supplementary document). For adding noise in Step 4, we can follow (Lample et al., 2017) to locally shuffle words while keeping the distance between the original and new position not larger than $3$ or set it as a null operation. There are two techniques to boost the final performance.

Computing $\mathbf{A}$: The alignment matrix $\mathbf{A}$ is obtained by averaging the cross-attention weights across all decoder layers and heads. We also add a temperature to control the sharpness of the attention distribution, the reciprocal of which was linearly increased from $0$ to $2$ during the first $20K$ steps. To avoid overfitting when computing $e(\tilde{\mathbf{z}})$ and $h(\tilde{\mathbf{y}})$, we apply dropout to $\mathbf{A}$ and stop back-propagating gradients through $\mathbf{A}$ when calculating the loss $\mathcal{L}_{F_{2}}(\bm{\theta})$.

Computing $h(\tilde{\mathbf{y}})$: Instead of interpolating one-hot labels in Eq. (6), we use the prediction vector $f(\mathbf{x},\mathbf{y};\hat{\bm{\theta}})$ on the sentence pair $(\mathbf{x},\mathbf{y})$ estimated by the model where $\hat{\bm{\theta}}$ indicates no gradients are back-propagated through it. However, the predictions made at early stages are usually unreliable. We propose to linearly combine the ground-truth one-hot label with the model prediction using a parameter $v$, which is computed as $vf_{j}(\mathbf{x},\mathbf{y};\hat{\bm{\theta}})+(1-v)h(y_{j})$ where $v$ is gradually annealed from $0$ to $1$ during the first $20K$ steps ¹¹1These two annealing hyperparameters in computing both $\mathbf{A}$ and $h(\tilde{\mathbf{y}})$ are the same for all the models and not elaborately tuned.. Notice that the prediction vectors are not used in computing the decoder input $e(\tilde{\mathbf{z}})$ which can be clearly distinguished from schedule sampling (Bengio et al., 2015).

This subsection shows that XEnDec, when fed with appropriate inputs, yields learning objectives identical to two recently proposed self-supervised learning approaches: MASS (Song et al., 2019) and BART (Lewis et al., 2019), as well as a supervised learning approach called Doubly Adversarial (Cheng et al., 2019). Table 1 summarizes the inputs of XEnDec to recover these approaches.

XEnDec can be used for self-supervised learning. As shown in Table 1, the inputs to XEnDec are two pairs of sentences $(\mathbf{x},\mathbf{y})$ and $(\mathbf{x}^{\prime},\mathbf{y}^{\prime})$. Given arbitrary alignment matrices, if we set $\mathbf{x}^{\prime}=\mathbf{y}^{u}$, $\mathbf{y}^{\prime}=\mathbf{y}^{u}$, and $\mathbf{x}$ to be a corrupted copy of $\mathbf{y}^{u}$, then XEnDec is equivalent to the denoising autoencoder which is commonly used to pre-train sequence-to-sequence models such as in MASS (Song et al., 2019) and BART (Lewis et al., 2019). In particular, if we allow $\mathbf{x}^{\prime}$ to be a dummy sentence of length $|\mathbf{y}^{u}|$ containing only “$\langle mask\rangle$” tokens ($\mathbf{y}^{u}_{mask}$ in the table), Eq. (7) yields the learning objective defined in the MASS model (Song et al., 2019) except that losses over unmasked words are not counted in the training loss. Likewise, as shown in Table 1, we can recover BART’s objective by setting $\mathbf{x}=\mathbf{y}^{u}_{noise}$ where $\mathbf{y}^{u}_{noise}$ is obtained by shuffling tokens or dropping them in $\mathbf{y}^{u}$. In both cases, XEnDec is trained with a self-supervised objective to reconstruct the original sentence from one of its corrupted sentences. Conceptually, denoising autoencoder can be regarded as a degenerated XEnDec in which the inputs are two views of its source correspondence for a monolingual sentence, e.g., $n(\mathbf{y})$ and $\mathbf{y}_{mask}$ for $\mathbf{y}$.

XEnDec can also be used in supervised learning. The translation loss proposed in (Cheng et al., 2019) is achieved by letting $\mathbf{x}^{\prime}$ and $\mathbf{y}^{\prime}$ be two “adversarial inputs”, $\mathbf{x}^{p}_{adv}$ and $\mathbf{y}^{p}_{adv}$, both of which consist of adversarial words at each position. For the construction of $\mathbf{x}^{p}_{adv}$, we refer to Algorithm 1 in (Cheng et al., 2019). In this case, the crossover encoder-decoder is trained with a supervised objective over parallel sentences.

The above connections to existing works illustrate the power of XEnDec when it is fed with different kinds of inputs. The results in Section 4.4 show that XEnDec is still able to improve the baseline with alternative inputs. However, our experiments show the best configuration found so far is to use the $F_{2}$-XEnDec in Algorithm 1 to deeply fuse the monolingual and parallel sentences.

Datasets. We evaluate our approach on two representative, resource-rich translation datasets, WMT’14 English-German and WMT’14 English-French across four translation directions, English$\rightarrow$German (En$\rightarrow$De), German$\rightarrow$English (De$\rightarrow$En), English$\rightarrow$French (En$\rightarrow$Fr), and French$\rightarrow$English (Fr$\rightarrow$En). To fairly compare with previous state-of-the-art results on these two tasks, we report case-sensitive tokenized BLEU scores calculated by the multi-bleu.perl script. The English-German and English-French datasets consist of 4.5M and 36M sentence pairs, respectively. The English, German and French monolingual corpora in our experiments come from the WMT’14 translation tasks. We concatenate all the newscrawl07-13 data for English and German, and newscrawl07-14 for French which results in 90M English sentences, 89M German sentences, and 42M French sentences. We use a word piece model (Schuster & Nakajima, 2012) to split tokenized words into sub-word units. For English-German, we build a shared vocabulary of 32K sub-words units. The validation set is newstest2013 and the test set is newstest2014. The vocabulary for the English-French dataset is also jointly split into 44K sub-word units. The concatenation of newstest2012 and newstest2013 is used as the validation set while newstest2014 is the test set. Refer to the supplementary document for more detailed data pre-processing.

Model and Hyperparameters. We implement our approach on top of the Transformer model (Vaswani et al., 2017) using the Lingvo toolkit (Shen et al., 2019). The Transformer models follow the original network settings (Vaswani et al., 2017). In particular, the layer normalization is applied after each residual connection rather than before each sub-layer. The dropout ratios are set to $0.1$ for all Transformer models except for the Transformer-big model on English-German where $0.3$ is used. We search the hyperparameters using the Transformer-base model on English-German. In our method, the shuffling ratio $p_{1}$ is set to $0.50$ while $0.25$ is used for English-French in Table 6. $p_{2}$ is sampled from a Beta distribution $Beta(2,6)$. The dropout ratio of $\mathbf{A}$ is $0.2$ for all the models. For decoding, we use a beam size of $4$ and a length penalty of $0.6$ for English-German, and a beam size of $5$ and a length penalty of $1.0$ for English-French. We carry out our experiments on a cluster of $128$ P100 GPUs and update gradients synchronously. The model is optimized with Adam (Kingma & Ba, 2014) following the same learning rate schedule used in (Vaswani et al., 2017) except for warmup_steps which is set to $4000$ for both Transform-base and Transformer-big models.

Training Efficiency. When training the vanilla Transformer model, each batch contains $4096\times 128$ tokens of parallel sentences on a $128$ P100 GPUs cluster. As there are three losses included in our training objective (Eq. (9)) and the inputs for each of them are different, we evenly spread the GPU memory budget into these three types of data by letting each batch include $2048\times 128$ tokens. Thus the total batch size is $2048\times 128\times 3$. The training speed is on average about $60\%$ of the standard training speed. The additional computation cost is partially due to the implementation of the noise function to corrupt the monolingual sentence $\mathbf{y}^{u}$ and can be reduced by caching noisy data in the data input pipeline. Then the training speed can accelerate to about $80\%$ of the standard training speed.

Table 2 shows the main results on the English-German and English-French datasets. Our method is compared with the following strong baseline methods. (Ott et al., 2018) is the scalable Transformer model. Our reproduced Transformer model performs comparably with their reported results. (Cheng et al., 2019) is a NMT model with adversarial augmentation mechanisms in supervised learning.  (Nguyen et al., 2019) boosts NMT performance by adopting multiple rounds of back-translated sentences. Both (Zhu et al., 2020) and (Yang et al., 2019) incorporate the knowledge of pre-trained models into NMT models by treating them as frozen input representations for NMT. We also compare our approach with MASS (Song et al., 2019) and BART (Lewis et al., 2019). As their goals are to learn generic pre-trained representations from massive monolingual corpora, for fair comparisons, we re-implement their methods using the same backbone model as ours, and jointly optimize their self-supervised objectives together with the supervised objective on the same corpora.

For English-German, our approach achieves significant improvements in both translation directions over the standard Transformer model. Even compared with the strongest baseline on English$\rightarrow$German, our approach obtains a $+0.72$ BLEU gain. More importantly, when we apply our approach to a significantly larger dataset, English-French with $36$M sentence pairs (vs. English-German with $4.5$M sentence pairs), it still yields consistent and notable improvements over the standard Transformer model.

The single-stage approaches (Joint Training with MASS & BART) perform slightly better than the two-stage approaches (Zhu et al., 2020; Yang et al., 2019), which substantiates the benefit of jointly training supervised and self-supervised objectives for resource-rich translation tasks. Among them, BART performs better with stable improvements on English-German and English-French and faster convergence. However, they still lag behind our approach. This is mainly because the $\mathcal{L}_{F_{2}}$ term in our approach can deeply fuse the supervised and self-supervised objectives instead of simply summing up their training losses. See Section 4.4 for more details.

Furthermore, we evaluate our approach and the best baseline method (Joint Training with BART) in Table 3 in terms of two additional evaluation metric, BLEURT (Sellam et al., 2020) and YiSi (Lo, 2019), which claim better correlation with human judgement. Results in Table 3 corroborate the superior performance of our approach compared to the best baseline method on both English$\rightarrow$German and English$\rightarrow$French.

Effect of Monolingual Corpora Sizes. Table 4 shows the impact of monolingual corpora sizes on the performance for our approach. We find that our approach already yields improvements over the baselines when using no monolingual corpora (x0) as well as when using a monolingual corpus with size comparable to the bilingual corpus (1x). As we increase the size of the monolingual corpus to 5x, we obtain the best performance with $30.46$ BLEU. However, continuing to increase the data size fails to improve the performance any further. A recent study (Liu et al., 2020a) shows that increasing the model capacity has great potential to exploit extremely large training sets for the Transformer model. We leave this line of exploration as future work.

Finetuning vs. Joint Training. To further study the effect of pre-trained models on the Transformer model and our approach, we use Eq. (3) to pre-train an NMT model on the entire English and German monolingual corpora. Then we finetune the pre-trained model on the parallel English-German corpus. Models finetuned on pre-trained models usually perform better than models trained from scratch at the early stage of training. However, this advantage gradually vanishes as training progresses (cf. Figure 1 in the supplementary document). As shown in Table 5, Transformer with finetuning achieves virtually identical results as a Transformer trained from scratch. Using the pre-trained model over our approach impairs performance. We believe this may be caused by a discrepancy between the pre-trained loss and our joint training loss.

Back Translation as Noise. One widely applicable method to leverage monolingual data in NMT is back translation (Sennrich et al., 2016b). A straightforward way to incorporate back translation into our approach is to treat back-translated corpora as parallel corpora. However, back translation can also be regarded as a type of noise used for constructing $\mathbf{y}^{u}_{noise}$ in $F_{2}$-XEnDec (shown in Fig. 1 and Step 4 in Algorithm 1), which can increase the noise diversity. As shown in Table 6, for English$\rightarrow$German trained on the Transformer-big model, our approach yields an additional $+1.9$ BLEU gain when using back translation to noise $\mathbf{y}^{u}$ and also outperforms the back-translation baseline. When applied to the English-French dataset, we achieve a new state-of-the-art result over the best baseline (Edunov et al., 2018). In contrast, the standard back translation for English-French hurts the performance of Transformer, which is consistent with what was found in previous works, e.g. (Caswell et al., 2019). These results show that our approach is complementary to the back-translation method and performs more robustly when back-translated corpora are less informative although our approach is conceptually different from works related to back translation (Sennrich et al., 2016b; Cheng et al., 2016; Edunov et al., 2018).

Robustness to Noisy Inputs. Contemporary NMT systems often suffer from dramatic performance drops when they are exposed to input perturbations (Belinkov & Bisk, 2018; Cheng et al., 2019), even though these perturbations may not be strong enough to alter the meaning of the input sentence. In this experiment, we verify the robustness of the NMT models learned by our approach. Following (Cheng et al., 2019), we evaluate the model performance against word perturbations which specifically includes two types of noise to perturb the dataset. The first type of noise is code-switching noise (CS) which randomly replaces words in the source sentences with their corresponding target-language words. Alignment matrices are employed to find the target-language words in the target sentences. The other one is drop-words noise (DW) which randomly discards some words in the source sentences. Figure 2 shows that our approach exhibits higher robustness than the standard Transformer model across all noise types and noise fractions. In particular, our approach performs much more stable for the code-switching noise.

Table 7 studies the contributions of the key components and verifies the design choices in our approach.