Language-Aware Multilingual Machine Translation
with Self-Supervised Learning

Parameter sensitivity is a measure of the impact on the loss when a specific parameter of a model is zeroed-out. It is widely used in pruning as importance score (Ding et al., 2019; Molchanov et al., 2019; Lubana and Dick, 2021). A parameter can express different sensitivities depending on the language of the input data. Those parameters that have high sensitivity to a specific language but low sensitivity to others, are language-specific parameters. We define the $i^{\text{th}}$ parameter in a model parameterized by $\bm{\Theta}$ as $\theta_{i}\in\mathbbm{R}$. We further define $\bm{\Theta}_{i}=[0,\cdots,0,\theta_{i},0,\cdots,0]\in\mathbbm{R}^{|\bm{\Theta}|}$ and $\bm{\Theta}_{-i}=[\theta_{1},\cdots,\theta_{i-1},0,\theta_{i+1},\cdots,\theta_{|\bm{\Theta}|}]\in\mathbbm{R}^{|\bm{\Theta}|}$. The sensitivity of the $i^{\text{th}}$ parameter given input batch $b_{l}$ from language $l$ is formulated as

$$\mathcal{S}(\theta_{i},b_{l})=|\mathcal{L}(\bm{\Theta},b_{l})-\mathcal{L}(\bm{\Theta}_{-i},b_{l})|,$$

(1)

where $\mathcal{L}(\cdot)$ is the loss function given the input batch and parameters. Then, we use a first-order Taylor decomposition to approximate the sensitivity of any arbitrary parameters. Equation 1 then becomes

$$\mathcal{S}(\theta_{i},b_{l})\approx|\bm{\Theta}_{i}^{T}\nabla_{\bm{\Theta}}\mathcal{L}(\bm{\Theta},b_{l})|,$$

(2)

where $\nabla_{\bm{\Theta}}\mathcal{L}(\bm{\Theta},b_{l})$ is the gradient of the loss with respect to the model parameters. In our implementation, we randomly pick 500 batches and feed them to the model to retrieve the gradients and compute the average sensitivity. We then have

$$\mathcal{S}(\theta_{i},\mathcal{B}_{l})\approx\frac{1}{|\mathcal{B}_{l}|}\sum_{b_{l}\in\mathcal{B}_{l}}|\bm{\Theta}_{i}^{T}\nabla_{\bm{\Theta}}\mathcal{L}(\bm{\Theta},b_{l})|,$$

(3)

where $\mathcal{B}_{l}$ is a set containing 500 random $b_{l}$ batches.

Now, we propose to quantify the degree of language-specificity of $\theta_{i}$ with respect to language $l$ by measuring the relative sensitivity difference between language $l$ and the other languages as

$$D(\theta_{i},l)=\frac{\mathcal{S}(\theta_{i},\mathcal{B}_{l})-\mathcal{S}(\theta_{i},\mathcal{B}_{-l})}{\mathcal{S}(\theta_{i},\mathcal{B}_{-l})+\sigma},$$

(4)

where $\mathcal{B}_{-l}$ represents the set composed of mixed batches from all training languages except for the language $l$, and $\sigma$ is a very small positive constant³³3$\sigma$ is 1e-8 in our implementation.. The larger $D(\theta_{i},l)$ is, the more language-specific $\theta_{i}$ is to language $l$.

A model with more balanced parameter sensitivity distribution shows better generalization (Liang et al., 2022). Xu et al. (2022) propose intra-distillation (ID) as an effective task-agnostic training method, aiming to encourage all parameters to contribute equally, which improves performance when model size is fixed. However we argue that, in the multilingual setting, ID actually helps the model learn more language-specific parameters resulting in improved performance. Given an input batch, ID needs to forward pass the model $K$ times to obtain $K$ outputs and each time a random subset of parameters is zeroed out. The core idea of ID is to minimize the difference of these $K$ outputs to approximate minimizing the contribution gap of the parameters that are zeroed-out, because the $K$ outputs are forced to be the same with different zeroed parameters. Let $\{p_{1},\cdots,p_{i},\cdots,p_{K}\}$ denote the $K$ outputs. Note that the outputs are probability distributions in the translation and denoising task. The ID loss is then formulated by the X-divergence (Xu et al., 2022) to minimize the difference of $K$ outputs as

$$\begin{gathered}\mathcal{L}_{id}=\frac{1}{K}\sum_{i=1}^{K}\mathbbm{KL}(p_{i}\parallel\bar{p})+\mathbbm{KL}(\bar{p}\parallel p_{i})\\ \text{where }\bar{p}=\frac{1}{K}\sum_{i=1}^{K}p_{i}\end{gathered}$$

(5)

Let the original task loss be $\mathcal{L}_{i}$ for the $i^{\text{th}}$ pass. Then, the total loss is a combination of the original task losses and ID loss, given as

$$\min\frac{1}{K}\sum_{i=1}^{K}\mathcal{L}_{i}+\alpha\mathcal{L}_{id}$$

(6)

where $\alpha$ is a hyper-parameter to control the strength of ID. Similar to Xu et al. (2022), we use dropout to simulate zeroed-out parameters in all experiments.

Although the explanation for better performance after using ID is that the model parameters become more balanced, it is unclear how parameter contributions to different languages change after applying ID in a multilingual (multitask) setting. For instance, do parameters become more language-agnostic and shareable across all languages, or do they become more language-specific? We investigate this in more details in Section 3.2.

Next, we study whether parameter contributions are more language-specific or just shareable across all languages after ID. Given the $i^{\text{th}}$ language $l_{i}$, we compute the sensitivities (Equation 3) of all parameters and flatten them into a list. Then, we calculate the Pearson correlation coefficients (PCC) $p_{ij}$ between sensitivity lists of any arbitrary pair of languages $l_{i}$ and $l_{j}$. A lower $p_{ij}$ indicates that there are more contribution (sensitivity) disagreements between languages $l_{i}$ and $l_{j}$. We plot a heat map to visualize $p_{ij}$ for every language pair. Taking into account that the top 10% parameters usually dominate the contribution (Xiao et al., 2019; Sanh et al., 2020), we consider the performance of two groups of parameters, high-sensitive (top 10% most sensitive) and low-sensitive (the remaining 90%) parameters, respectively. Figure 2 shows that all $p_{ij}$ in both groups become lower, indicating there is lower sensitivity similarity between different languages for the same parameters, which means the model becomes more language-specific after ID. For instance, sensitivity similarity between zul and wol drops from 0.67 to 0.57 in the low-sensitive group. However, the $p_{ij}$ of low-sensitive parameters drops much more than high-sensitive ones, and high-sensitive parameters still hold high similarity (over 0.9). Thus, low-sensitive parameters mostly have language-specific properties while high-sensitive parameters tend to play ‘language-agnostic’ roles. Overall, parameters are more language-specific after ID⁵⁵5The overall parameter contribution is still more balanced as claimed in Xu et al. (2022). We leave further discussion on this to Appendix B.. In fact, learning more language-specific parameters through ID in MMT leads to better performance as seen in Section 3.1. These findings align with the results of recent studies which investigate language-specific parameters (Lin et al., 2021; Zhang et al., 2021; NLLB Team et al., 2022), indicating the importance of language-specific parameters.

Here, we study the reason why language-specific parameters are important and how much they contribute. To investigate this, we first measure the degree of language-specificity of all parameters based on Equation 4. We explore the contribution of language-specific parameters with respect to the BLEU scores. Then, we conduct one-shot unstructured pruning with respect to BLEU scores in order of the degree of language-specificity for both models with and without ID, starting with the least language-specific parameters⁶⁶6Note that, as shown in Figure 2(b), the 10% most sensitive parameters are highly language-agnostic. They are easy to classify as less language-specific and can be pruned, but pruning them would lead to near-random performance (BLEU$\approx$0), making it hard to evaluate the importance of more language-specific parameters. Thus, we keep the top 10% sensitive parameters and prune the rest of parameters that display a more language-specific behavior.. As more parameters are pruned, a slower performance drop means that a higher contribution comes from the remaining more language-specific parameters. Figure 3 shows the average BLEU drop across 8 languages versus the percentage of parameters pruned. After pruning the less language-specific parameters, the rest of the more language-specific parameters in the model with ID are able to preserve better performance, indicating the importance of more language-specific parameters.

Self-supervised learning objectives usually involve sentence denoising either on the encoder side, such as MLM (Devlin et al., 2019), or on the decoder side, such as DAE (Liu et al., 2020). Jointly denoising sentences on both the encoder and the decoder sometimes is better than a single denoising objective (Wang et al., 2020; Kim et al., 2021) for MMT, but the training cost is doubled as we need to calculate the loss for the same monolingual sentence twice (masked in two different ways). We propose concurrent denosing, a self-supervised task that denoises a single masked sentence both on the encoder and decoder sides, which not only reduces the training time but also improves the language understanding of the model to result in better MMT performance.

We add noise to the monolingual data by whole-word masking (Devlin et al., 2019), where we randomly replace $r_{m}\%$ words with the special token <mask>. During the replacement process, each word has a 10% chance not to be masked, and another 10% chance to be replaced with other random tokens. The encoder and decoder use a shared output layer to reconstruct the original sentence. The loss for the encoder and decoder side are denoted as $\mathcal{L}_{e}$ and $\mathcal{L}_{d}$ respectively⁷⁷7Unlike DAE training on the decoder side, we zero out the losses which predict non-masked tokens.. The total training loss combining translation loss $\mathcal{L}_{MMT}$ and two self-supervised losses is

$$\mathcal{L}=\mathcal{L}_{MMT}+\mathcal{L}_{e}+\mathcal{L}_{d}.$$

(7)

Concurrent denoising is illustrated in Figure 4. Two key differences between our concurrent denoising method and regular MLM or DAE methods are worth highlighting.

We investigate whether ID helps concurrent denoising to improve overall performance. We apply ID to the co-training of CD and MMT tasks. Following Equation 6 and 7, our final loss is

$$\mathcal{L}=\frac{1}{K}(\sum_{i=1}^{K}\mathcal{L}_{MMT_{i}}+\sum_{i=1}^{K}\mathcal{L}_{e_{i}}+\sum_{i=1}^{K}\mathcal{L}_{d_{i}})+\\ \alpha(\mathcal{L}_{id\_MMT}+\mathcal{L}_{id\_e}+\mathcal{L}_{id\_d}),$$

(8)

where $\mathcal{L}_{id\_MMT}$, $\mathcal{L}_{id\_e}$ and $\mathcal{L}_{id\_d}$ respectively represent the ID loss for translation, encoder denoising and decoder denoising (i.e., $\mathcal{L}_{id\_e}$ minimizes the difference of the $K$ encoder outputs based on Equation 5, etc.). The $i$ index in $\mathcal{L}_{e_{i}}$, $\mathcal{L}_{MMT_{i}}$ and $\mathcal{L}_{d_{i}}$ indicates that these losses are for the $i^{\text{th}}$ forward pass.

We consider three strong baselines. All baselines are our own implementation following the settings from the original papers.

In addition to the M8 dataset described in Section 3.1, we also build a larger dataset (M15), covering 15 languages. In composing this dataset, we take into account linguistic diversity and data size. The resulting dataset has languages from 6 linguistic families and a balanced number of high-resource, low-resource and very low-resource languages. Detailed information on this dataset is in Appendix D. We randomly sample at most 3M monolingual samples per language for M8, and 15M for M15. The distribution of monolingual data and parallel data for M8 and M15 is shown in Figure 5. Note that we also use parallel data for self-supervised learning, so the true monolingual data size includes bitext data. We use the Flores-200 dataset for evaluation. All datasets come from the primary bitext and monolingual data used for the NLLB-200 model (NLLB Team et al., 2022).

We use a data sampling temperature of $T=1$ suggested by NLLB Team et al. (2022) to train on the MMT objective. For monolingual data, we use a temperature of $\frac{10}{7}$ to balance the SSL training, as suggested by Liu et al. (2020). During co-training, we mix the two sources in an equal ratio (50% monolingual data (including bitext used for SSL training) with self-supervision and 50% parallel data).

All experiments consider both the eng$\rightarrow$xxx and xxx$\rightarrow$eng directions and use the Transformer architecture (Vaswani et al., 2017). We use Transformer${}_{\text{big}}$ (242M parameters, 6 layers, 16 heads, 1,024 hidden dimension, 4,096 FFN dimension) for M8 experiments. For M15 experiments, we double the layers of Transformer${}_{\text{big}}$ (418M parameters). We use a vocabulary of size 32k for both M8 and M15 with SentencePiece (Kudo and Richardson, 2018). The batch size is 30K tokens. We warm-up for the first 8K steps. We set the total training steps to 100K and 300k for M8 and M15 respectively, with patience set to 10 for early stopping. We forward pass the model twice ($K$=2) to conduct ID. We set the ID weight $\alpha=5$. During concurrent denoising, the masking ratio is set to $r_{m}=30\%$. We also show the effect of masking ratio in Appendix E. During generation, we use beam search with a beam size of 5 and a length penalty of 1.0. All models are evaluated with sacreBLEU (spm tokenizer).

The overall results for the xxx$\rightarrow$eng and eng$\rightarrow$xxx directions are shown in Tables 2 and 3. For both M8 and M15, and both translation directions, concurrent denosing is better than all aforementioned baselines, and combining it with ID further improves upon the baselines by an even larger margin. For instance, our method outperforms MASS by 11.3% and 3.7% on M8 and M15 respectively, averaged across all languages and directions. We also show the effectiveness of ID on other objectives like DAE in Section 6.1, but the results are subpar compared to CD+ID.

Aligned with the findings of Wang et al. (2020); Kim et al. (2021), we observe that DAE+MLM is better than DAE alone in M8 xxx$\rightarrow$eng, but the improvements become very minor when it comes to M8 eng$\rightarrow$xxx or when scaling to 15 languages. MASS performs similarly or better than DAE in the eng$\rightarrow$xxx but worse in the xxx$\rightarrow$eng.

In M15, high-resource languages perform slightly worse with SSL methods compared to the MMT only baseline, but improves other categories, similar to the observations of NLLB Team et al. (2022). It does not occur on M8, possibly due to the smaller dataset size allowing for sufficient model capacity to learn from additional monolingual data.

Note that the effectiveness of SSL such as DAE and MASS is not as pronounced as reported by Wang et al. (2020) and Siddhant et al. (2022). However, it is necessary to consider the for domain mismatch between the training and evaluation data. As demonstrated by Siddhant et al. (2022), a significant decline in performance can occur when either monolingual or bitexts diverge from the evaluation domain. In our study, the training data is sourced from NLLB-200 and FLORES-200, which encompasses a wide range of domains. we hypothesize that this contributes to the observed lessened effectiveness of SSL techniques in our experiments.

The final loss, described in Equation 8, has 6 loss terms. Except for the translation loss, we ablate the relative contribution of all the other 5 loss terms to the translation task performance. In Table 4, we show the results of this ablation study on M8 xxx$\rightarrow$eng directions. Method is the regular MMT model and method is ID training only for MMT (the same result as in Section 3.1). Method is the same as the MMT+DAE method. With the help of ID for the decoder denoising (method ) and an additional ID for translation (method ), translation performance can respectively obtain +0.41 and +0.98 BLEU on average compared to . Note that method is the MMT+DAE+ID method. Compared to our MMT+CD+ID method, it substantially underperforms our method (17.73 vs. 18.69), which shows that our method could better stimulate the potential of ID. The results for methods , and indicate the effectiveness of encoder denoising with CD and applying ID. Overall, the translation performance improves by including all the loss terms.

In Section 3, we observed that ID helps MMT learn more language-specific parameters and improve model generalization. We are also interested in understanding 1) whether the model also learns more language-specific parameters for the SSL task (here we investigate CD), and 2) what is the relationship of parameter contribution between MMT and SSL tasks for the same language. We use the xxx$\rightarrow$eng direction of the M8 dataset as an example to study these questions.

In Figure 6, we plot a heat map to illustrate the PCC of all parameter sensitivities between every language pair. As expected, parameter sensitivity similarity becomes lower for all languages, which means there are more language-specific parameters when we train SSL methods with ID. For the second question, in Figure 7 we show the parameter sensitivity similarity between the MMT and CD tasks for each language. The contribution similarity becomes higher between the two tasks for every language with ID. This is expected, since the losses of MMT and CD have the same objective on the decoder side, i.e., text generation conditioned on another text. This is also another reason why SSL tasks can help multilingual translation.