1 Introduction↩︎

Multilingual methods that process multiple languages with one single model have gained favor across a variety of NLP tasks [1]–[6] because (1) training and deployment of one multilingual model is more computationally efficient compared to maintaining one model for each language [7], (2) training multilingually can improve accuracy, particularly for low-resource languages (LRLs) [8]–[10].

However, in multilingual training, the amount and type of training data available varies greatly across languages. Because most models are trained using empirical risk minimization (ERM), which minimizes the average training loss on the training set, high-resource languages (HRLs) with large amounts of data contribute to the majority of the training objective. When model capacity is limited, this results in trade-offs or decreased performance on some languages, particularly LRLs [7], [11], [12]. To better control this trade-off, a common practice is to balance the training distribution by heuristic oversampling of LRLs [3], [7], [9]. Although simple data balancing can improve the performance on LRLs significantly, it is far from optimal. First, the sampling hyperparameters need to be adjusted for different datasets. Second, the use of simple heuristics does not consider the inherent level of difficulty in learning each language, the similarity between languages in the multilingual dataset, and other factors that affect cross-lingual transfer. Because of this, previous work has indicated the importance of learning strategies that are explicitly tailored for each multilingual learning scenario [13].

In this paper, we propose a new learning procedure for multilingual translation that automatically adjusts the training distribution of different languages using distributionally robust optimization (DRO) [14], [15]. In constrast to ERM, DRO casts learning as a game between the learner and an adversary, where the learner picks a model while the adversary picks the hardest data distribution for that model within an uncertainty set \(\mc{Q}\) of potential distributions we wish to perform well on (which typically contains the training distribution \(P_0\)).

We first demonstrate how to apply DRO to multilingual training by letting the adversary choose the relative weights of individual language pairs in the training objective. However, we empirically find that naively applying existing methods to multilingual learning yields inferior results to ERM, mostly because (1) standard DRO objectives tend to be overly conservative and only take into account language pairs with very large losses and (2) existing optimization algorithms for DRO essentially reweigh the gradients of examples in a mini-batch, which implicitly changes the scale of the learning rates. This hurts modern NLP models like Transformers [16] that are highly sensitive to learning rate schedules.

To remedy this, we propose both a novel training objective and a corresponding optimization algorithm amenable to the multilingual setting. Our objective is a variation on Group DRO of [17] with a more flexible uncertainty set, parameterized by the \(\cs\)-divergence. To efficiently solve the min-max game, we propose an iterated best response scheme that, at each epoch, re-samples the training data according to the worst weighting for the current model parameters, and then runs ERM training on the re-sampled dataset. Our method—which we refer to as \(\chi\)-IBR —incurs negligible additional computational cost compared to ERM.

While this method applies to essentially any multilingual task, we specifically demonstrate its benefit on three sets of language pairs from two multilingual machine translation datasets. We experimentally test these choices by comparing several objectives and optimization algorithms, and results show that our method consistently outperforms existing DRO procedures and various strong baselines.

2 Preliminaries↩︎

2.0.0.1 Notation.

Throughout this paper, \(n\) denotes the training set size and \(d\) the number of parameters of the model. For \(m \in \N\), \(\Delta^m\) denotes the \(m\)-dimensional simplex, i.e. \(\Delta^m \defeq \crl{q \in \R^m, q_i \ge 0~and~ \sum_i q_i = 1}\). The data lies in \(\mc{X}\times\mc{Y}\) where \((\vx, \vy) \in \mc{X}\times\mc{Y}\) consists of a source and target sentence pair with \(\vx = (\vx_1, \ldots, \vx_{L_{\vx}})\) and \(\vy = (\vy_1, \ldots, \vy_{L_{\vy}})\). The function \(\ell:(\mc{X}\times\mc{Y})\times\R^d \to \R\) refers to the loss. We consider maximum-likelihood estimation, i.e. for a target sentence \(\vy\) with \(L_\vy\) tokens, we define \[\label{eq:loss-function} \small \ell(\vx, \vy; \theta) = -\frac{1}{L_\vy}\sum_{i=1}^{L_\vy} \log p(\vy_i|\vx, \vy_{<i}; \theta)\nonumber\qquad{(1)}\]

2.1 Multilingual Machine Translation↩︎

In contrast to bilingual machine translation, which translates from a single source language \(S\) to a target language \(T\), multilingual neural MT (MNMT) learns a single model to translate between \(N\) language pairs \(\crl{\prn{S_1,T_1}, \ldots \prn{S_N, T_N}}\). The training data \(D_{\mathrm{train}}\) is the concatenation of the \(N\) parallel datasets, i.e. \(D_{\mathrm{train}} = [D_1; D_2, \cdots; D_N]\). We can then define the probability over each language pair \(p^\mathsf{train}\in \Delta^N\) as \(p^\mathsf{train}_i = \tfrac{\abs{D_i}}{\sum_j \abs{D_j}}\). We now describe two common training objectives for MNMT.

2.1.0.1 Empirical Risk Minimization (ERM).

The simplest and most common approach for MNMT is to minimize the empirical loss over data points, which we will refer to as ERM. More precisely, we define the average loss on a parallel dataset \(D_i\) as \[\nonumber \small\mathcal{L}(\theta;D_i) \defeq \frac{1}{\abs{D_i}}\sum_{(\vx, \vy) \in D_i} \ell(\vx, \vy;\theta). \] ERM for multilingual models then corresponds to simply minimizing the loss over \(D\), i.e. over all the aggregated parallel sentence pairs. This yields \[ \what{\theta}_n^\mathsf{ERM} \in \argmin_{\theta} \mathcal{L}(\theta;D) = \sum_{i\le N} p^\mathsf{train}_i \mathcal{L}(\theta;D_i). \] Classical results in learning theory guarantee that, under mild assumptions, as \(n\) goes to infinity, \(\what{\theta}_n^\mathsf{ERM}\) will show good performance on test sets with the same distribution as \(D\). However, this does not guarantee that our model will perform adequately on individual parallel datasets. To remedy this issue, several works propose varying the sampling distribution—or equivalently the weighting of the parallel datasets in ERM—in order to encourage more uniform performance across language pairs.

2.1.0.2 Weighted Risk Minimization and Sampling Strategies.

The amount of training data can vary significantly across language pairs. As a result, in ERM training—i.e. optimizing for the average loss across sentence pairs—HRLs contribute most of the objective, resulting in poor performance on LRLs. Balancing the objective—or equivalently, the usage of training data—between HRLs and LRLs is important to maintain good performance across all languages [4], [7]. A commonly adopted approach in multilingual training is temperature-based sampling [6], [7] where the probability of sampling data from \(D_i\) is proportional to its data size exponentiated by a temperature term \(\tau\), i.e. \(p_{\tau,i} = \tfrac{\abs{D_i}^{1/\tau}}{\sum_j \abs{D_j}^{1/\tau}}\) (referred to as ERM with \(\tau\) in §4). This is equivalent to optimizing the re-weighted objective \[\nonumber \small \mathcal{L}_\tau(\theta;D) = \sum_{i\le N}p_{\tau,i}\mathcal{L}(\theta;D_i). \] As a result, \(\tau=1\) corresponds to ERM where most of the contribution comes from the HRLs and \(\tau=\infty\) corresponds to sampling language pairs uniformly at random, i.e. with data from LRLs being over-sampled. This approach comes with three major drawbacks (1) \(\tau\) is an extra hyper-parameter that requires tuning for each MNMT instance to balance the performance across both HRLs and LRLs, (2) this heuristic sampling method does not consider the training dynamics of each language and how the optimal sampling distribution might evolve during the training process and (3) the parameterization of \(p_\tau\) is not only very constrained (essentially one degree of freedom), it is also only a function of the quantity of training data, which is too rigid to achieve the desired performance.

To resolve some of the above issues, the recently proposed MultiDDS [13] uses a gradient-based meta-learning approach to learn the sampling distribution over language pairs to maximize gradient similarity with a multilingual development set. However, due to the necessity to calculate and store extra gradients, their approach comes at an increased computational and memory cost. In contrast, \(\chi\)-IBR enjoys the same computational complexity as ERM, and as we show in experiments it also largely outperforms MultiDDS.

2.2 Distributionally Robust Optimization↩︎

In contrast to ERM and related sampling strategies which optimize for a fixed training distribution, DRO aims to find a model \(\theta\) that performs well on an entire collection of potential test distributions \(\mc{Q}\) (the uncertainty set). Formally, we wish to \[\label{eq:dro} \minimize_{\theta}\sup_{Q \in \mc{Q}} \E_{(\vx, \vy)\sim Q}\brk{\ell(\vx, \vy; \theta)}.\tag{1}\]

Originating from operations research [14], [18]–[20], DRO is a promising way to tackle robustness in a variety of machine learning and NLP problems [21]–[23].

We present here a recent variant, Group DRO, developed by [17] which incorporates additional information about the data distribution to define more meaningful uncertainty sets. Abstractly, this method assumes a collection of distributions over subpopulations \(\crl{P_g}_{g\in\mc{G}}\) such that the training distribution is a mixture of these subpopulations. Importantly, it assumes that this group structure is observed. The Group DRO objective then minimizes the worst-case loss over these groups, which is equivalent to 1 with \(\mc{Q} = \crl{\sum_{g\in\mc{G}} q_gP_g: q\in\Delta^{\abs{\mc{G}}}}\), or equivalently, all possible mixtures of the distributions over subpopulations. In MNMT, the \(N\) language pairs at our disposal naturally correspond to groups; thus the Group DRO objective can be defined as \[\label{eq:gdro}\mathcal{L}^{\mathsf{GDRO}}(\theta;D) = \max_{i\in \brk{N}} \mathcal{L}(\theta;D_i).\qquad{(2)}\] In other words, Group DRO wishes to find a model \(\theta\) that performs well for the worst language pair. [22] propose a related but less conservative objective, CVaR-Group DRO at level \(\alpha\in\brk{0, 1}\) which, considers instead the average of the \(\ceil{\alpha N}\) largest group losses.

3 Methods for Distributionally Robust Multilingual Learning↩︎

As we previewed in §2, Group DRO is a natural objective for the multilingual setting. However, in experiments we found that naively applying existing DRO objectives fails to achieve performance on par with strong baselines, often improving results on language pairs with high losses but sacrificing too much performance overall. Our main contribution is showing how to successfully apply DRO to the MNMT setting, and to the best of our knowledge, our work is the first to do so. To that end, our methodological contributions are two-fold: (i) we first describe the shortcomings of the Group DRO objective ?? , then propose a related training criterion that addresses these issues, (ii) we describe an optimization algorithm to solve the min-max optimization problem that is amenable to the MNMT setting.

3.1 \(\cs\)-Group DRO↩︎

3.1.0.1 Shortcomings of Group DRO.

A weakness of the objective ?? is that apart from the language pair with largest loss, the objective does not take into account the value of the loss on the other language pairs. To illustrate this, consider this example with \(N=3\) language pairs and suppose that there exists two parameters \(\theta_1\) and \(\theta_2\) with the following loss:

\[\begin{align} \mathcal{L}(\theta_1;D_1) = 0.1, \, \mathcal{L}(\theta_1;D_2) = 0.1, \, \mathcal{L}(\theta_1; D_3) = 1.1 \\ \mathcal{L}(\theta_2;D_1) = 1.0, \, \mathcal{L}(\theta_2; D_2) = 1.0, \, \mathcal{L}(\theta_2; D_3) = 1.0 \end{align}\] We have that \(\mathcal{L}^{\mathsf{GDRO}}(\theta_1;D) = 1.1\) but \(\mathcal{L}^{\mathsf{GDRO}}(\theta_2;D) = 1.0\). Consequently, the Group DRO objective will prefer \(\theta_2\) to \(\theta_1\) while clearly one would pick \(\theta_1\) over \(\theta_2\) in most practical cases. The aforementioned CVaR-Group DRO also exhibits this behavior and ignores the values of the language pairs not in the largest \(\alpha\)-fraction.

To address this issue, let us rewrite the objective \[\mathcal{L}^{\mathsf{GDRO}}(\theta;D) = \max_{q\in\Delta^N}\sum_{i\le N}q_i\mathcal{L}(\theta;D_i),\] where the equality holds because the optimal weighting just puts all the mass on the language with largest loss. A natural way to make the objective less conservative is to instead take the maximum over a subset of the simplex \(\mc{U} \subset \Delta^N\). This leads to the following objective \[\label{eq:gen-gdro} \mathcal{L}^{\mc{U}}(\theta;D) = \sup_{q\in\mc{U}}\sum_{i\le N}q_i\mathcal{L}(\theta;D_i).\tag{2}\]

Different choices of \(\mc{U}\) will yield different objectives with different robustness properties. Note that this is a general formulation as \(\mc{U} = \crl{p^\mathsf{train}}\) reduces to the ERM objective, while \(\mc{U}_\alpha = \crl{q: \norms{q / p^\mathsf{train}}_\infty \le 1/\alpha}\) corresponds to the CVaR-Group DRO of [22]. We would like to choose \(\mc{U}\) such that optimizing this objective results in models with better performance on language pairs with large losses (typically LRLs) without significant degradation of average performance.

To this end, we turn to a common and flexible choice for \(\mc{U}\): \(f\)-divergence balls [24] of radius \(\rho > 0\) around \(p^\mathsf{train}\), namely \[\mc{U}^f_{\rho} \defeq \crl{q: D_f(q, p^\mathsf{train}) \le \rho},\] where \(D_f(q, p) \defeq \sum_{i\le N}p_i f(q_i/p_i)\). In particular, we propose using the \(\cs\)-divergence which corresponds to \(f(t) = \tfrac{1}{2}(t-1)^2\) and define \(\cs(q, p) = \tfrac{1}{2}\sum_i p_i(q_i/p_i - 1)^2\) with its corresponding uncertainty set \(\mc{U}^{\cs}_{\rho}\). The \(\cs\)-divergence has a long history [19] and previous work shows that minimizing the robust loss with the \(\cs\)-uncertainty set enjoys favorable statistical properties such as optimally trading-off bias and variance [25] or guaranteeing robustness and fairness [21], [26]. We refer to the objective \(\mathcal{L}^{\mc{U}^{\cs}_{\rho}}\) as the \(\cs\)-Group DRO. Going back to the toy example, setting \(\rho = 0.1\), yields that \(\mathcal{L}^{\mc{U}^{\cs}_{\rho}}(\theta_1;D) = 0.64\) while \(\mathcal{L}^{\mc{U}^{\cs}_{\rho}}(\theta_2;D) = 1.0\). With the \(\cs\)-uncertainty set, the objective rightly prefers \(\theta_1\) to \(\theta_2\) and takes into account all the losses and not only the largest. We further confirm these intuitions and show in §4 and §5 that this is a suitable choice of uncertainty set for MNMT.

3.2 Optimization algorithm↩︎

3.2.0.1 Desiderata of the optimization algorithm.

Minimizing the objective 2 effectively corresponds to a min-max optimization problem. Even in the relatively simple convex-concave setting, these are generally harder to solve than convex minimization problems. Recall that we want to \[\minimize_\theta \sup_{q:\cs(q, p^\mathsf{train}) \le \rho} \sum_{i\le N}q_i \mathcal{L}(\theta;D_i).\]

Due to the architectures we consider in MNMT, we wish for an algorithm that effectively changes the sampling distribution over mini-batches of data instead of importance-weighting the gradients. Indeed, standard MT architectures such as the Transformer [16] are extremely sensitive to learning rate schedules and we empirically observe that importance-weighting the gradients result in poor performance.

The canonical way to solve min-max problem is via primal-dual methods (PD) [27] (see background in Appendix 8), where at each step \(t\), one keeps two vectors \((\theta_t, q_t)\) and alternates between a gradient descent step on \(\theta_t\) and a gradient ascent step on \(q_t\). One can perform these updates efficiently as they only require unbiased stochastic gradient estimate of the loss w.r.t. \(\theta_t\) and \(q_t\). To obtain unbiased gradient estimate of the loss w.r.t. \(\theta_t\), one either has to, at each step, sample a mini-batch of examples from \(\mathsf{Multinomial}(q_t)\) and return the gradient of the loss or sample a mini-batch from \(\mathsf{Multinomial}(p^\mathsf{train})\) and importance-weight the gradient. As previously mentioned, the latter is not suitable for Transformer-type architectures. The former option is not ideal as it is more convenient for an algorithm to decide the sequence of mini-batches every epoch rather than every optimizer step as this integrates much more smoothly with data loaders in deep learning frameworks, especially when doing distributed training. As a result, we posit that primal-dual algorithms are not an adequate choice for optimizing DRO-type objectives in our setting. We further discuss this point in §5.

To circumvent this issue, we consider a different optimization algorithm which we refer to as iterated best response (IBR), where, instead of doing a single gradient descent and ascent step, we iterate between (approximately) solving the maximization (resp. minimization) on \(q\) (resp. \(\theta\)), while keeping \(\theta_t\) (resp. \(q_t\)) fixed. This is similar in spirit to algorithms in the game theory literature where individuals play the optimal strategy (best response) assuming everyone else’s strategies remain constant. Under some mild assumptions, this procedure converges to the equilibrium of the game [28]. Formally, we alternate between \[\begin{align} \small\theta^{t+1} & \leftarrow \argmin_{\theta} \sum_i q^{t}_i \mathcal{L}(\theta;D_i)\tag{3}\\ q^{t+1} & \leftarrow \argmax_{q:\cs(q, p^\msf{train})\le \rho} \sum_i q_i \mathcal{L}(\theta^{t+1};D_i)\tag{4}.\end{align}\]

3.2.0.2 Practical implementation.

As we show in Appendix 7, given the values of the loss, the \(q\)-update 4 is computed to accuracy \(\epsilon\) in \(O(N\log(1/\epsilon))\) steps. Indeed, by taking the dual [29] of 4 , we transform the \(N\)-dimensional problem into a one-dimensional root-finding procedure over the dual variable which we efficiently solve with a bisection. We provide the details in Appendix 7. Note that computation cost is negligible compared to computing the gradient of the loss. We implement the \(\theta\)-update of 3 by running a training epoch on a re-sampling of the training set \(D\) according to \(q^{t+1}\).

To compute the loss values \(\crl{\mathcal{L}(\theta_t;D_i)}_{i\le N}\), necessary to perform the \(q\)-update, one needs to compute the loss \(\ell(\vx, \vy;\theta_t)\) for every single example \((\vx, \vy)\in D\). This is prohibitively expensive to compute at every epoch. To avoid this, we keep track of the (approximate) historical values of the token-level loss \(\hat{\mathcal{L}}_k\) on each language pair \(k\) with an exponential moving average (EMA) (see line 14 in Algorithm 2). We precisely describe our implementation in Algorithm 2. We see that it respects our desiderata and comes at no computational cost. In §5, we compare primal-dual and iterated best response for various uncertainty sets.

3.2.0.3 Subtracting the baseline.

[22] propose subtracting a per-group scalar—which we refer to as a baseline—to each group loss before taking the maximum over \(q\). They learn this baseline using a generative bi-gram model on each group. Recall that \(\what{\theta}^{\mathsf{ERM}}\) is the parameter we obtain when we minimize the average loss and define \(\what{\theta}^\tau\) when optimizing \(\mathcal{L}_\tau\). In this work, we propose using \(b_i = \mathcal{L}(\what{\theta}^{\mathsf{ERM}};D_i)\) or \(b_i = \mathcal{L}(\what{\theta}^\tau;D_i)\). Intuitively, the baseline corresponds to the minimum performance we wish for our model on the given group and as such the loss obtained with ERM and its temperature variants are natural candidates. We show in §5 that these yield significant improvement and conveniently make our method more robust to the choice of \(\rho\). We leave different (potentially learned) choices of baseline to future work.

4 Experiments↩︎

4.1 Datasets↩︎

We evaluate the proposed method on two datasets: the 58-languages TED talk corpus [30] and WMT datasets. For the TED corpus, we evaluate on two sets of languages with varying levels of language diversity following [13]: (1) related includes 4 LRLs (aze, bel, glg, slk) and their corresponding related HRL (tur, rus, pos, ces). (2) diverse includes 8 languages with varying amount of data without considering linguistic similarities (bos, mar, hin, mkd, ell, bul, fra, kor)³. Both of the related and diverse sets have around 760K sentences of training data.

For WMT, we consider 2 HRLs (German:deu and French:fra) and 2 LRLs (Tamil:tam and Turkish:tur). We subsample around 5M training sentences from the parallel corpus provided by the WMT shared task. Specifically, the training data of deu-eng, fra-eng is from WMT14, tam-eng is from WMT20 and tur-eng is from WMT18. We use the corresponding test and dev sets from each shared task for evaluation and validation.

We evaluate both en-to-any (translate English to a target language) and any-to-en (translate a source language into English) directions for all language sets. We provide dataset statistics in Appendix 9.

4.2 Experimental setup↩︎

For the translation models, we adopt the encoder-decoder Transformer [16] architecture with the implementation provided in fairseq [31]. For both datasets, we use a Transformer-base architectures that also has 6 encoder and decoder layers with hidden dimension size being 512 and 8 attention heads.⁴ The model is trained for 200K and 300K steps for TED and WMT respectively with the batch size of 65,536 tokens. For both datasets, we learn the sentencepiece [32] vocabulary for the English and the combined corpus of other languages respectively. We use beam search with beam size 5 for decoding and report the SacreBLEU score [33], [34] on test sets for evaluation. For the TED and WMT datasets respectively, the constraint size \(\rho\) for the chi-square ball is set to be 0.05 and 0.3, and for the baseline losses we use the average token-level loss on each \(D_i\) computed from the ERM model with \(\tau=1\) and \(\tau=100\)—see §5 for more analyses of these choices. We provide additional pre-processing and training details in Appendix 10.

4.2.0.1 Baselines.

We compare with (1) temperature-based sampling method described in §2.1 in three standard settings (\(\tau=1/5/100\)), where \(\tau=100\) approximates uniform sampling over language pairs, and (2) MultiDDS described in §2.1. In addition, we also perform extensive empirical studies over different DRO uncertainty sets and optimization procedures in §5.

4.3 Main Results↩︎

We present the BLEU scores of en\(\rightarrow\)any and any\(\rightarrow\)en translation directions on TED and WMT data in Tab. 1 and 2 respectively. First, for both TED and WMT datasets, \(\chi\)-IBR outperforms all the other baseline methods in terms of average BLEU score over all language pairs. By taking a closer look at the BLEU scores for each individual language pairs, \(\chi\)-IBR improves over almost all the language pairs for both translation directions compared to ERM. Secondly, as expected from temperature-based sampling methods, different values of \(\tau\) achieve different trade-offs between the performances on HRLs and LRLs. As we explained in §2.1, large values of \(\tau\) favor LRLs as this results in data being sampled with equal probability from each language pair while small values of \(\tau\) approach ERM and will benefit HRLs. As a result, \(\tau\) needs to be carefully tuned to achieve adequate performance on both HRLs and LRLs.

Importantly, \(\chi\)-IBR achieves a significantly better trade-off than the sampling method for any value of \(\tau\). We show in Fig. 3 the quantitative improvements of \(\chi\)-IBR over the best \(\tau\) for various datasets and in both translation directions. Surprisingly, while the improvements are larger on LRLs in most cases, we consistently observe improvements on HRLs. This indicates that finding the right sampling distribution over language pairs facilitates cross-lingual transfer. We further observe that \(\chi\)-IBR achieves more significant improvements in the en\(\rightarrow\)any direction than in the any\(\rightarrow\)en direction. This further supports our hypothesis. Indeed, it is attested in previous work on MNMT [7] that it is harder to decode to multiple languages than encode from multiple languages and as such, the en\(\rightarrow\)any direction is a significantly harder multi-task learning problem. As such, this is where optimal cross-lingual transfer would yield the larger gains, which is what we observe in practice. We also note that our method incurs negligible computational overhead compared to ERM.

5 Analysis↩︎

5.0.0.1 The importance of the sampling distribution.

An advantage of our method over temperature-based sampling methods is that it dynamically adjusts the training distribution as the model evolves and does not compute it solely as a function of amount of training data. Our hypothesis is that this is important to achieve good performance across language pairs and that different sampling distributions will be adequate at different stages of training. We empirically check our hypothesis by studying how the training distribution \(q\) (the so-called best response) changes across training epochs. In Fig. 4 and 5, we plot the best response of \(\chi\)-IBR across epochs on the TED-diverse dataset for both translation directions. In addition, we also plot the historical losses in Fig. 6 (Appendix). Our first observation is that the optimal \(q\) noticeably evolves across epochs which further showcases the need for dynamically adjusting the sampling distribution. We make the following observations (i) \(\chi\)-IBR demonstrates the desired behavior and, at the early stages of training, always down-weights HRLs and up-weights LRLs; (ii) somewhat counter-intuitively, there is no direct correlation between \(\abs{D_i}\), the amount of data in language \(i\) and the final value \(\mathcal{L}(\theta^{(t)};D_i)\). The latter further evidences the limitations of sampling distributions only based on the amounts of training data \(\abs{D_i}\). Indeed, while kor is a HRL, it is typologically much farther from English so there is more inherent uncertainty in the task. Consequently, it has larger losses and is consistently up-weighted throughout training. On the other hand, while hin is a LRL, it achieves low loss after being up-weighted during the early stages of training and is consequently down-weighted after that.

5.0.0.2 Comparison to DRO variants.

We demonstrate the benefits of \(\chi\)-IBR over other DRO objectives by extensively evaluating a range of robust objectives and associated optimization algorithms. In terms of objective, we compare against (1) Group DRO [17], (2) CVaR-Group DRO [22] and (3) FastDRO [23]. In terms of algorithms, we experiment with primal-dual methods and our proposed iterated best response procedure which we both described in §3.2. Note that in the case of Group DRO (i.e. \(\mc{U} = \Delta^N\)), iterated best response is not a sensible choice as it would result in each training epoch being spent on a single language pair. In the case of CVaR Group DRO, we follow the implementation of [22], which is a hybrid of the two optimization algorithms with a primal update on \(\theta\) and a best response update on \(q\). We compare the performance of these methods on the WMT dataset. For fairness, we baseline losses in the same way for all the DRO objectives. We first observe that, outside of \(\chi\)-IBR , none of the DRO objectives are competitive with temperature-weighted ERM. We also observe that for both uncertainty sets, iterated best response convincingly outperforms the same objective trained with primal-dual. We finally note that, for a fixed optimization algorithm, \(\cs\)-Group DRO outperforms the CVaR objective on all but one language pairs. This validates both our choice of uncertainty set and of optimization procedure.

5.0.0.3 The effects of baselined losses.

We study the effect of the choice of baseline on the performance across languages. In Tab. 4, we empirically evaluate different baseline choices and uncertainty sizes \(\rho\). We observe that in the TED dataset, baseline-ing with \(\mathcal{L}(\what{\theta}^{\mathsf{ERM}};D_i)\) performs significantly better than baseline-ing with \(\mathcal{L}(\what{\theta}^{\tau=100};D_i)\) ((e) versus (g) while it is reversed for WMT. We explain this by observing that the LRLs in TED consist of very small amounts of data (on the order of a few thousands) and using \(\tau=100\) results in a severe oversampling of LRLs, which the model then fits perfectly. As a result, recall the intuition that the baseline sets a lower bound on the performance we wish to achieve but because of the small training data and overfitting, the model disproportionately up-weighs the LRLs, which harms overall performance. This does not occur in WMT and uniform sampling across language pairs sets a good target performance for DRO methods. Finally, we see that with the right baseline loss, our method is more robust to different choices of \(\rho\) (e.g., comparing (c) and (d) versus (e) and (f)). We consistently observe this for other translation directions and datasets.

6 Conclusion↩︎

We showed how to successfully apply DRO to the MNMT setting and automatically adjust the sampling distribution over language pairs resulting in sizeable improvements in performance. We posit that this approach would also be successful in other multilingual scenarios. Our work raises a few questions: (i) what are the right baseline losses? (ii) surprisingly, \(\chi\)-IBR also improves performance on HRLs; under what circumstances does cross-lingual transfer happen and which languages does it benefit most? We hope our work could inspire better distributionally robust learning methods for multilingual training in the future.

Acknowledgements↩︎

This work was supported in part by a Facebook SRA Award and the NSF/Amazon Fairness in AI program under grant number 2040926.

7 Best response↩︎

In this section, we describe the bisection procedure we use to solve the best response and update \(q\) shown in 4 . This derivation generically exists in the literature (e.g. Appendix A.1.2 in [23]) but we specialize it to the \(\cs\)-ball centered at \(p^\mathsf{train}\) and include it here for completeness.

For \(v\in\R^m\), the optimization problem we wish to solve is \[\label{eq:best-response-cs} \begin{align} \maximize_{q\in \Delta^m} & \qquad q^\top v \\ \subjectto & \qquad \cs(q, p^\mathsf{train}) \le \rho, \end{align}\tag{5}\]

Let us consider the Lagrangian of this problem \[\Lambda(q;\nu, \eta) \defeq q^\top v - \eta(\ones^\top q - 1) - \lambda(\cs(q, p^\mathsf{train}) - \rho)\] Maximizing over \(q\) yields that the solution as a function of \(\eta\) is \(q^\star(\eta)_i \propto p^\mathsf{train}_i (v_i - \eta)_+\), where \((u)_+ \defeq \max\crl{0, u}\). Since the objective is linear, it holds that the maximum is attained on the extreme points of the constraint sets. Consequently, we need to find \(\eta^\star\) such that \(\cs(q^\star(\eta^\star), p^\mathsf{train}) = \rho\). This is a simple root finding procedure that we solve to accuracy \(\epsilon\) in \(\log (1/\epsilon)\) steps with a bisection. Given that each evaluation of \(q^\star(\eta)\) requires \(O(m)\) operations, the runtime of the algorithm is \(O(m\log(1/\epsilon))\).

8 Primal-dual methods↩︎

Primal-dual algorithms [27], [35] are the methods of choice to efficiently solve min-max problems.

8.1 The primal-dual algorithm↩︎

Let us assume that \(\mc{X} \subset \R^d\) and \(\mc{Y} \subset \R^p\) are closed, bounded convex sets and let \(F:\mc{X}\times\mc{Y}\to \R\) be a function such that \(F(x, \cdot)\) is concave for all \(x\in\mc{X}\) and \(F(\cdot, y)\) is convex for all \(y\in\mc{Y}\). We wish to solve the following min-max problem \[\label{eq:minmax} \min_{x\in\mc{X}} \max_{y\in\mc{Y}} F(x, y).\tag{6}\]

This is a well-studied problem of optimization and there the literature provides optimal algorithms for many choices of \(F\), \(\mc{X}\) and \(\mc{Y}\). A standard approach are the so-called primal-dual methods, where one keeps a pair of iterates \((x_t, y_t)\) performs a gradient descent (resp. ascent) step on \(x_t\) (resp. \(y_t\)). In all generality, these updates are often mirror descent updates to properly exploit the geometric structures of \(\mc{X}, \mc{Y}\) and \(F\). Let \(h^\msf{x}:\R^d\to\R\) (resp. \(h^\msf{y}:\R^p\to\R\)) be a \(1\)-strongly-convex function w.r.t. a given norm \(\norms{\cdot}_\msf{x}\) (resp. \(\norms{\cdot}_\msf{y}\)). We denote \(D_{h^\msf{x}}\) and \(D_{h^\msf{y}}\) their associated Bregman divergences. The primal-dual algorithms initialize \(x_0 \in \mc{X}, y_0 \in\mc{Y}\) and for a stepsize \(\eta > 0\), iterates \[\begin{align} x_{t+1} & = \argmin_{x\in\mc{X}} \crl*{ g_{\msf{x}, t}^\top x + \frac{1}{\eta}D_{h^\msf{x}}(x, x_t) } \\ y_{t+1} & = \argmax_{y\in\mc{Y}} \crl*{ g_{\msf{y}, t}^\top y - \frac{1}{\eta}D_{h^\msf{y}}(y, y_t), }, \end{align}\] where \(g_{\msf{x}, t} \in \partial_x F(x_t, y_t)\) and \(g_{\msf{y}, t} \in \partial_y F(x_t, y_t)\). After \(T\) steps, return \((\bar{x}_T, \bar{y}_T)\) with \(\bar{z}_T \defeq 1/T\sum_{t\le T}z_t\). Assuming \(F\) is appropriately Lipschitz and that \(\mc{X}\) and \(\mc{Y}\) are bounded, one finds an \(\epsilon\)-approximate saddle point in \(O(\epsilon^{-2})\) steps. Importantly, these guarantee still holds even when only having stochastic unbiased estimates of \(g_\msf{x}\) and \(g_\msf{y}\) [27] which is essential in large-scale settings.

8.2 Primal-dual algorithms for Group DRO↩︎

The problem 2 can be cast as an instance of 6 by setting \(\mc{Y} = \mc{U}\) and \(F(\theta, q) \defeq \sum_{i\le N}q_i\mathcal{L}(\theta;D_i)\). Note that the problem is unconstrained in \(\theta\) but it is generally not an issue in practice. The gradient in \(\theta\) and \(q\) are \[\begin{align} \nabla_\theta F(\theta, q) = \sum_{i\le N}q_i \nabla_\theta \mathcal{L}(\theta;D_i) \\ \brk*{\nabla_q F(\theta, q)}_i = \mathcal{L}(\theta;D_i). \end{align}\]

8.2.1 Obtaining stochastic gradients↩︎

To run primal-dual, we require stochastic gradient estimates \(\tilde{g}_{\theta}\) and \(\tilde{g}_{q}\). First of all, note that an unbiased estimate of \(\nabla_\theta \mathcal{L}(\theta;D_i)\) is just \(\nabla_\theta \ell(\vx, \vy;\theta)\) where \((\vx, \vy)\) is sampled at random from \(D_i\). To obtain stochastic gradient estimates, we have two choices: sampling from \(q\) or sampling from an arbitrary \(p_0 \in \Delta^N\) and importance-weighting.

8.2.1.1 Sampling from \(q\)

Let \(B\) be the mini-batch size and \(\prn{I_1, \ldots, I_B}\) be \(B\) indices sampled from \(q\) and \((\vx_{I_j}, \vy_{I_j})\) be randomly sampled examples from \(D_{I_j}\). In this case, \[\begin{align} \tilde{g}_\theta & = \frac{1}{B}\sum_{j\le B} \nabla_\theta \ell(\vx_{I_j}, \vy_{I_j};\theta) \\ \brk{\tilde{g}_q}_i & = \frac{1}{B}\sum_{j\le B} \frac{1}{q_{I_j}}\ell(\vx_{I_j}, \vy_{I_j};\theta), \end{align}\] are clearly unbiased stochastic gradient estimates for \((\nabla_\theta F(\theta, q), \nabla_q F(\theta, q))\).

8.2.1.2 Sampling from an arbitrary \(p_0 \in \Delta^N\)

Assume \(p_0\) is in \(\Delta^n\) and that \(p_{0, i} > 0\) for \(i\in\brk{N}\). As previously, assume we sample a mini-batch of indices \(\prn{J_1, \ldots, J_B}\) from \(p_0\) and that \((\vx_{J_j}, \vx_{J_j})\) is a randomly sampled example from \(D_{J_j}\). We can obtain stochastic gradient estimates by importance-weighting (IW) \[\begin{align} \tilde{g}^{\msf{IW}}_\theta & = \frac{1}{B}\sum_{j\le B} \frac{q_{J_j}}{p_{0, J_j}}\nabla_\theta \ell(\vx_{J_j}, \vy_{J_j};\theta) \\ \brk{\tilde{g}^{\msf{IW}}_q}_i & = \frac{1}{B}\sum_{j\le B} \frac{1}{p_{0, I_j}}\ell(\vx_{I_j}, \vy_{I_j};\theta), \end{align}\]

In terms of implementation, it is often impractical to change the distribution of batches at each iteration of the optimizer.

8.2.2 Choice of Bregman divergence↩︎

8.2.2.1 \(\theta\)-update.

The update in \(\theta\) is unconstrained and it is standard to optimize the objective with Stochastic Gradient Descent so as a result, we pick \(h^\theta(\theta) = \frac{1}{2}\norm{\theta}_2^2\). This results in the standard SGD update \[\theta_{t+1} = \theta_t - \eta \tilde{g}_\theta,\] where \(\tilde{g}_\theta\) is an unbiased stochastic gradient estimate that we compute following the previous section.

8.2.2.2 \(q\)-update for \(\mc{U} = \Delta^N\).

In the case where \(\mc{U}\) is the full simplex, it is standard to choose \(h^\msf{q}(q) = \sum_{i\le N}q_i \log q_i\) (the negative Shannon entropy). This choice yields the familiar Exponentiated Gradient update \[q_{t+1, i} \propto q_{t, i}\exp(\eta \brk{\tilde{g}_q}_i)\]

8.2.2.3 \(q\)-update for \(\mc{U} = \crl{q\in\Delta^N: \cs(q, p^\mathsf{train}) \le \rho}\).

As the \(\cs\)-uncertainty set is essentially the intersection between a (weighted)-norm-2 ball and the simplex, it is a natural choice to pick \(h^\msf{q}(q) = \frac{1}{2}\norm{q}_2^2\). This leads to the following update \[q_{t+1} = \argmin_{q\in\Delta^N:\cs{q, p^\mathsf{train}}\le \rho}\norm*{(q_t + \eta \tilde{g}_q) - q}_2^2,\] or in other words, the projection of the gradient ascent step \((q_t + \eta \tilde{g}_q)\) onto the \(\cs\)-ball. We explain in Appendix 8.3 how to efficiently compute this projection for arbitrary \(p^\mathsf{train}\).

8.2.2.4 \(q\)-update for \(\mc{U} = \mc{U}_\alpha\).

We follow the implementation of [22] which runs a hybrid between primal-dual methods and best response and thus we do not need to explicit include the projection onto the CVaR uncertainty set. We discuss the option here for completeness. For the CVaR uncertainty set, it is standard to also use the negative Shannon entropy \(h^\msf{q} = \sum_i q_i \log q_i\). We refer to Appendix F.6.2 of [23] and their provided code for more details on this projection.

8.3 Projecting on the \(\cs\)-ball↩︎

To run the primal-dual algorithm, the gradient update on \(q\) requires projecting on the \(\cs\)-ball centered at \(p^\mathsf{train}\). For \(v\in\R^m\), we wish to solve the following \[\label{eq:proj-cs} \begin{align} \minimize_{q\in \Delta^m} & \qquad \norm{q - v}_2^2 \\ \subjectto & \qquad \cs(q, p^\mathsf{train}) \le \rho, \end{align}\tag{7}\] where \(\Delta^m\) is the \(m\)-dimensional simplex and \(\cs(q, p^\mathsf{train}) \defeq \tfrac{1}{2}\sum_{i\le m} p_i (q_i / p_i - 1)^2\), the \(f\)-divergence [24] corresponding to \(t\mapsto \tfrac{1}{2}(t-1)^2\).

While this projection is standard in the literature, it is often derived in the case of \(p^\mathsf{train}= \ones / m\) (see e.g. [36]). We show here how to efficiently do the projection for arbitrary \(p^\mathsf{train}\in \Delta^m\).

First, for \(\lambda \ge 0, \eta\in\R\), the Lagrangian of 7 is \[\begin{align} \Lambda(q, \lambda, \eta) \defeq & \frac{1}{2}\norm{q - v}_2^2 \\ & + \lambda (\cs(q, p^\mathsf{train}) - \rho) + \eta(q^\top \ones - 1). \end{align}\] Taking the partial dual \(g(\lambda, \eta) \defeq \inf_{q \succeq 0} \Lambda(q, \lambda, \eta)\) yields \[\begin{align} g(\lambda, \eta) = -\inf_{\lambda \ge 0, \eta \in \R} & \frac{1}{2} \sum_{i\le m} \frac{p_i}{p_i + \lambda}(v_i - \eta)_+ \\ & + \frac{\lambda}{2}(1+2\rho) + \eta \end{align}\] In contrast to the uniform case (i.e.\(p^\mathsf{train}= \ones/m\)), one cannot derive the optimal dual variable \(\lambda^*\) in closed-form and we have to solve for both dual variables. Finding an \(\epsilon\)-accurate solution takes order \(m\log(1/\epsilon)\) time using cutting plane-type methods when the dimension is \(O(1)\) [37]. In the large-scale applications we consider, this is negligible in comparison to computing the gradient of the loss with respect to the network parameters and thus the primal-dual algorithm incurs (almost) no additional computational overhead.

In practice, we implement this with two nested bisections; while this adds an additional \(\log(1/\epsilon)\) factor, this is significantly more convenient to implement.

9 Data Statistics↩︎

We present the number of training sentences in the TED datasets in Tab. 5. The number of training sentences in the WMT dataset is 2.5M (deu), 1.8M (fra), 512,608 (tam) and 195,762 (tur).

10 Preprocessing and Training Details↩︎

We describe the preprocessing and training details in this section.

10.0.0.1 Preprocessing

We download WMT datasets from the WMT official websites⁵. The TED data is downloaded from the link⁶ provided by [13]. We learn sentencepiece [32] for English and other languages respectively. Specifically, we combine all the English sentences from individual parallel datasets. Following [7], we also resample the sentences of other languages based on the temperature-based distribution (\(\tau=5\)) to learn the vocabulary. For TED, the vocabulary sizes for English and other languages are set to be 10K and 30K respectively. For WMT, the vocabulary sizes for English and other languages are set to be 24K and 34K respectively. We also remove sentences in the training data that are longer than 250 tokens after bpe. For validation set, to ensure the validation loss are fairly comparable across different languages, we cap the number of the sentences in the validation set to be the same for each language (800 for TED and 1,500 for WMT).

10.0.0.2 Training details

The hyperparameters of Transformer models are described in Tab. 6. For all experiments, we adopt the Adam optimizer [38] using \(\beta_1=0.9, \beta_2=0.98, \epsilon=1e-8\). We set the label smoothing as \(0.1\)⁷. For ERM models, we use the default learning rate scheduling. The learning rate is scheduled using inverse_sqrt with a maximum learning rate of \(5e-4\) and \(2e-4\) (following [13]) for WMT and TED respectively. The warmup steps is set to be 4000 for both datasets. For \(\chi\)-IBR , because the training distribution is dynamically changing across epochs, we use step learning rate decay to retain a higher learning rate during training. Specifically, we use the same maximum learning rate and warm up steps as ERM, but decay the learning rate by half every 100K training steps for both datasets. Recall that we use exponential moving average to compute the historical loss values, and we set the hyperparameter \(\lambda\) to be 0.1 throughout. We train on 4 and 8 V100 GPUs for TED and WMT experiments respectively. The training time for one experiment takes around 1 - 2 days.

References↩︎