Balancing Training for Multilingual Neural Machine Translation (2020)

Abstract: When training multilingual machine translation (MT) models that can translate to/from multiple languages, we are faced with imbalanced training sets: some languages have much more training data than others. Standard practice is to up-sample less resourced languages to increase representation, and the degree of up-sampling has a large effect on the overall performance. In this paper, we propose a method that instead automatically learns how to weight training data through a data scorer that is optimized to maximize performance on all test languages. Experiments on two sets of languages under both one-to-many and many-to-one MT settings show our method not only consistently outperforms heuristic baselines in terms of average performance, but also offers flexible control over the performance of which languages are optimized.

## introduction multilingual models are trained to process different languages in a single model, and have been applied to a wide variety of nlp tasks such as text classification @xcite , syntactic analysis @xcite , named-entity recognition @xcite , and machine translation (mt) @xcite . these models have two particularly concrete advantages over their monolingual counterparts. first, deploying a single multilingual model is much more resource efficient than deploying one model for each language under consideration @xcite . second, multilingual training makes it possible to transfer knowledge from high-resource languages (hrls) to improve performance on lowresource languages (lrls) @xcite @xcite @xcite . a common problem with multilingual training is that the data from different languages are both heterogeneous (different languages may exhibit very different properties) and imbalanced (there may be wildly varying amounts of training data for each language). thus, while lrls will often benefit from transfer from other languages, for languages where sufficient monolingual data exists, performance will often decrease due to interference from the heterogeneous nature of the data. this is especially the case for modestly-sized models that are conducive to efficient deployment @xcite . to balance the performance on different languages, the standard practice is to heuristically adjust the distribution of data used in training, specifically by over-sampling the training data from lrls @xcite @xcite . for example, @xcite sample training data from different languages based on the dataset size scaled by a heuristically tuned temperature term. however, such heuristics are far from perfect. first, @xcite find that the exact value of this temperature term significantly affects results, and we further show in experiments that the ideal temperature varies significantly from one experimental setting to another. second, this heuristic ignores factors other than data size that affect the interaction between different languages, despite the fact that language similarity has been empirically proven important in examinations of cross-lingual transfer learning @xcite . in this paper, we ask the question: "is it possible to learn an optimal strategy to automatically balance the usage of data in multilingual model training?" to this effect, we propose a method that learns a language scorer that can be used throughout training to improve the model performance on all languages. our method is based on the recently proposed approach of differentiable data selection @xcite , a general machine learning method for optimizing the weighting of different training examples to improve a pre-determined objective. in this work, we take this objective to be the average loss from different languages, and directly optimize the weights of training data from each language to maximize this objective on a multilingual development set. this formulation has no heuristic temperatures, and enables the language scorer to consider the interaction between languages. based on this formulation, we propose an algorithm that improves the ability of dds to optimize multiple model objectives, which we name multi-dds. this is particularly useful in the case where we want to optimize performance on multiple languages simultaneously. specifically, multidds (1) has a more flexible scorer parameterization, (2) is memory efficient when training on multiple languages, and (3) stabilizes the reward signal so that it improves all objectives simultaneously instead of being overwhelmed by a single objective. while the proposed methods are model-agnostic and thus potentially applicable to a wide variety of tasks, we specifically test them on the problem of training multilingual nmt systems that can translate many languages in a single model. we perform experiments on two sets of languages (one with more similarity between the languages, one with less) and two translation directions (one-to-many and many-to-one where the "one" is english). results show that multidds consistently outperforms various baselines in all settings. moreover, we demonstrate multidds provides a flexible framework that allows the user to define a variety of optimization objectives for multilingual models. ## multilingual training preliminaries monolingual training objective a standard nmt model is trained to translate from a single source language s to a target language t . the parameters of the model are generally trained by preparing a training dataset d train , and defining the empirical distribution of sentence pairs x, y sampled from d train as p . we then minimize the empirical risk j(θ, p ), which is the expected value of the loss function ℓ(x, y; θ) over this distribution: multilingual training formulation a multilingual nmt model can translate n pairs of languages {s 1 -t 1 , s foot_1 -t 2 , ..., s n -t n }, from any source language s i . to its corresponding target t i . to train such a multilingual model, we have access to n sets of training data d @xmath0, d 2 train , . . . , d n train , where d i train is training data for language pair s i -t i . from these datasets, we can define p i , the distribution of sentences from s i -t i , and consequently also define a risk j(θ, p i ) for each language following the monolingual objective in eq. 1. however, the question now becomes: "how do we define an overall training objective given these multiple separate datasets?" several different methods to do so have been proposed in the past. to discuss all of these different methods in a unified framework, we further define a distribution p d over the n sets of training data, and define our overall multilingual training objective as in practice, this overall objective can be approximated by selecting a language according to ĩ ∼ p d (i), then calculating gradients with respect to θ on a batch of data from d ĩ train . evaluation methods another important question is how to evaluate the performance of such multilingual models. during training, it is common to use a separate development set for each language d @xmath1, d 2 dev , ..., d n dev to select the best model. given that the objective of multilingual training is generally to optimize the performance on all languages simultaneously @xcite , we can formalize this objective as minimizing the average of dev risks 2 : relation to heuristic strategies this formulation generalizes a variety of existing techniques that define p d (i) using a heuristic strategy, and keep it fixed throughout training. uniform: the simplest strategy sets p d (i) to a uniform distribution, sampling minibatches from each language with equal frequency @xcite . proportional: it is also common to sample data in portions equivalent to the size of the corresponding corpora in each language @xcite . temperature-based: finally, because both of the strategies above are extreme (proportional underweighting lrls, and uniform causing overfitting by re-sampling sentences from limited-size lrl datasets), it is common to sample according to data size exponentiated by a temperature term τ @xcite : when τ = 1 or τ = ∞ this is equivalent to proportional or uniform sampling respectively, and when a number in the middle is chosen it becomes possible to balance between the two strategies. as noted in the introduction, these heuristic strategies have several drawbacks regarding sensitivity to the τ hyperparameter, and lack of consideration of similarity between the languages. in the following sections we will propose methods to resolve these issues. ## differentiable data selection now we turn to the question: is there a better way to optimize p d (i) so that we can achieve our final objective of performing well on a representative development set over all languages, i.e. minimizing j dev (θ, d dev ). in order to do so, we turn to a recently proposed method of differentiable data selection @xcite , a general purpose machine learning method that allows for weighting of training data to improve performance on a separate set of held-out data. specifically, dds uses a technique called bilevel optimization @xcite , that learns a second set of parameters ψ that modify the training objective that we use to learn θ, so as to maximize the final objective j dev (θ, d dev ). specifically, it proposes to learn a data scorer p (x, y; ψ), parameterized by ψ, such that training using data sampled from the scorer optimizes the model performance on the dev set. to take the example of learning an nmt system to translate a single language pair i using dds, the general objective in eq. 1 could be rewritten as dds optimizes θ and ψ iteratively throughout the training process. given a fixed ψ, the update rule for θ is simply to update the data scorer, dds uses reinforcement learning with a reward function that approximates the effect of the training data on the model's dev performance ## dds for multilingual training in this section, we use the previously described dds method to derive a new framework that, instead of relying on fixed heuristics, adaptively optimizes usage of multilingual data for the best model performance on multiple languages. we illustrate the overall workflow in fig. 1 . first, we note two desiderata for our multilingual training method: 1) generality: the method should be flexible enough so that it can be utilized universally for different multilingual tasks and settings (such as different translation directions for nmt). 2) scalablity: the method should be stable and efficient if one wishes to scale up the number of languages that a multilingual model supports. based on these two properties, we introduce multi-dds, an extension of the dds method tailored for multilingual training. method multidds directly parameterizes the standard dataset sampling distribution for multilingual training with ψ: and optimizes ψ to minimize the dev loss. notably, unlike standard dds we make the design decision to weight training datasets rather than score each training example x, y directly, as it is more efficient and also likely easier to learn. we can thus rewrite the objective in eq. 2 to incorporate both ψ and θ as: in other words, while the general dds framework evaluates the model performance on a single dev set and optimizes the weighting of each training example, our multilingual training objective evaluates the performance over an aggregation of n dev sets and optimizes the weighting of n training sets. the reward signal for updating ψ t is intuitively, eq. 10 implies that we should favor the training language i if its gradient aligns with the gradient of the aggregated dev risk of all languages. implementing the scorer update the pseudocode for the training algorithm using multidds can be found in line 25. notably, we do not update the data scorer ψ on every training step, because it is too computationally expensive for nmt training @xcite . instead, after training the multilingual model θ for a certain number of steps, we update the scorer for all languages. this implementation is not only efficient, but also allows us to re-estimate more frequently the effect of languages that have low probability of being sampled. in order to do so, it is necessary to calculate the effect of each training language on the current model, namely r(i; θ t ). we estimate this value by sampling a batch of data from each d i train to get the training gradient for θ t , and use this to calculate the reward for this language. this process is detailed in line 11 of the line 25. unlike the algorithm in dds which requires storing n model gradients, foot_2 this approximation does not require extra memory even if n is large, which is important given recent efforts to scale multilingual training to 100+ @xcite or even 1000+ languages @xcite . ## stabilized multi-objective training in our initial attempts to scale dds to highly multilingual training, we found that one challenge was that the reward for updating the scorer became unstable. this is because the gradient of a multilingual dev set is less consistent and of higher variance than that of a monolingual dev set, which influences the fidelity of the data scorer reward. 4 algorithm 1: training with multidds input :d train ; m: amount of data to train the multilingual model before updating ψ; output :the converged multilingual model θ * initialize p d (i, ψ) to be proportional to dataset size thus, instead of using the gradient alignment between the training data and the aggregated loss of n dev sets as the reward, we propose a second approach to first calculate the gradient alignment reward between the data and each of the n dev sets, then take the average of these as the final reward. this can be expressed mathematically as follows: to implement this, we can simply replace the standard reward calculation at line 11 of line 25 to use the stable reward. we name this setting multidds-s. in § 6.6 we show that this method has less variance than the reward in eq. 10. ## related work our work is related to the multilingual training methods in general. multilingual training has a rich history @xcite @xcite , but has become particularly prominent in recent years due the ability of neural networks to easily perform multi-task learning @xcite @xcite . as stated previously, recent results have demonstrated the importance of balancing hrls and lrls during multilingual training @xcite 2019), which is largely done with heuristic sampling using a temperature term; multidds provides a more effective and less heuristic method. @xcite ; @xcite choose languages from multilingual data to improve the performance on a particular language, while our work instead aims to train a single model that handles translation between many languages. @xcite @xcite propose improvements to the model architecture to improve multilingual performance, while multidds is a model-agnostic and optimizes multilingual data usage. our work is also related to machine learning methods that balance multitask learning @xcite . for example, @xcite proposes to weigh the training loss from a multitask model based on the uncertainty of each task. our method focuses on optimizing the multilingual data usage, and is both somewhat orthogonal to and less heuristic than such loss weighting methods. finally, our work is related to meta-learning, which is used in hyperpa-rameter optimization @xcite , model initialization for fast adaptation @xcite , and data weighting @xcite . notably, @xcite apply meta-learning to learn an nmt model initialization for a set of languages, so that it can be quickly fine-tuned for any language. this is different in motivation from our method because it requires an adapted model for each of the language, while our method aims to optimize a single model to support all languages. to our knowledge, our work is the first to apply meta-learning to optimize data usage for multilingual objectives. ## conclusion in this paper, we propose multidds, an algorithm that learns a language scorer to optimize multilingual data usage to achieve good performance on many different languages. we extend and improve over previous work on dds @xcite , with a more efficient algorithmic instantiation tailored for the multilingual training problem and a stable reward to optimize multiple objectives. mul-tidds not only outperforms prior methods in terms of overall performance on all languages, but also provides a flexible framework to prioritize different multilingual objectives. notably, multidds is not limited to nmt, and future work may consider applications to other multilingual tasks. in addition, there are other conceivable multilingual optimization objectives than those we explored in § 6.4.