Distributionally Robust Multilingual Machine Translation (2021)

Abstract: Multilingual neural machine translation (MNMT) learns to translate multiple language pairs with a single model, potentially improving both the accuracy and the memory-efficiency of deployed models. However, the heavy data imbalance between languages hinders the model from performing uniformly across language pairs. In this paper, we propose a new learning objective for MNMT based on distributionally robust optimization, which minimizes the worst-case expected loss over the set of language pairs. We further show how to practically optimize this objective for large translation corpora using an iterated best response scheme, which is both effective and incurs negligible additional computational cost compared to standard empirical risk minimization. We perform extensive experiments on three sets of languages from two datasets and show that our method consistently outperforms strong baseline methods in terms of average and per-language performance under both many-to-one and one-to-many translation settings.

## introduction multilingual methods that process multiple languages with one single model have gained favor across a variety of nlp tasks @xcite @xcite @xcite because (1) training and deployment of one multilingual model is more computationally efficient compared to maintaining one model for each language @xcite , (2) training multilingually can improve accuracy, particularly for low-resource languages (lrls) @xcite @xcite . 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 @xcite . in this paper, we propose a new learning proce-dure for multilingual translation that automatically adjusts the training distribution of different languages using distributionally robust optimization (dro) @xcite . 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 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 @xcite 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 sagawa et al. with a more flexible uncertainty set, parameterized by the χ 2 -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 methodwhich we refer to as χ-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. ## preliminaries notation. throughout this paper, n denotes the training set size and d the number of parameters of the model. for @xmath0, ∆ m denotes the mdimensional simplex, i.e. ∆ m := {@xmath1, q @xmath2. the data lies in x × y where (x, y) ∈ x × y consists of a source and target sentence pair with @xmath3, . . . , x lx ) and @xmath4, . . . , y ly ). the function ℓ : (x × y) × r @xmath5. we consider maximumlikelihood estimation, i.e. for a target sentence y with l y tokens, we define ## 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 (2), 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. we have that l gdro (θ 1 ; d) = 1.1 but l gdro (θ 2 ; d) = 1.0. consequently, the group dro objective will prefer θ 2 to θ 1 while clearly one would pick θ 1 over θ 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 α-fraction. different choices of u will yield different objectives with different robustness properties. note that this is a general formulation as @xmath6, while u α = {q : q/p train ∞ ≤ 1/α} corresponds to the cvar-group dro of @xcite . we would like to choose 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 u: f -divergence balls @xcite of radius ρ > 0 around p train , namely with the χ 2 -uncertainty set, the objective rightly prefers θ 1 to θ 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. ## analysis 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. 3 and 4, we plot the best response of χ-ibr across epochs on the ted-diverse dataset for both translation directions. in addition, we also plot the historical losses in fig. 5 (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) χ-ibr demonstrates the desired behavior and, at the early stages of training, always down-weights hrls and up-weights lrls; (ii) somewhat counterintuitively, there is no direct correlation between |d i |, the amount of data in language i and the final value l(θ (t) ; d i ). the latter further evidences the limitations of sampling distributions only based on the amounts of training data |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. comparison to dro variants. we demonstrate the benefits of χ-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 @xcite , (2) cvar-group dro @xcite and ( 3 ) fastdro @xcite . 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. @xmath7, 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 @xcite , which is a hybrid of the two optimization algorithms with a primal update on θ 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 χ-ibr, none of the dro objectives are competitive with temperature-weighted erm. we also observe that for both uncertainty sets, iterated 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 ρ. we observe that in the ted dataset, baselineing with l( θ erm ; d i ) performs significantly better than baseline-ing with l( θ τ =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 τ = 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 ρ (e.g., comparing (c) and (d) versus (e) and ( f )). we consistently observe this for other translation directions and datasets. ## 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, χ-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.