Robust Optimization for Multilingual Translation with Imbalanced Data

Abstract:
Multilingual models are parameter-efficient and especially effective in improving low-resource languages by leveraging crosslingual transfer. Despite recent advance in massive multilingual translation with ever-growing model and data, how to effectively train multilingual models has not been well understood. In this paper, we show that a common situation in multilingual training, data imbalance among languages, poses optimization tension between high resource and low resource languages where the found multilingual solution is often sub-optimal for low resources. We show that common training method which upsamples low resources can not robustly optimize population loss with risks of either underfitting high resource languages or overfitting low resource ones. Drawing on recent findings on the geometry of loss landscape and its effect on generalization, we propose a principled optimization algorithm, Curvature Aware Task Scaling (CATS), which adaptively rescales gradients from different tasks with a meta objective of guiding multilingual training to low-curvature neighborhoods with uniformly low loss for all languages. We ran experiments on common benchmarks (TED, WMT and OPUS-100) with varying degrees of data imbalance. CATS effectively improved multilingual optimization and as a result demonstrated consistent gains on low resources (+0.8 to +2.2 BLEU) without hurting high resources. In addition, CATS is robust to overparameterization and large batch size training, making it a promising training method for massive multilingual models that truly improve low resource languages.

heading: 1 Introduction
text: Multilingual models have received growing interest in natural language processing (NLP) [34, 41, 10, 26, 2, 19, 58]. The task of multilingual machine translation aims to have one model which can translate between multiple languages pairs, which reduces training and deployment cost by improving parameter efficiency. It presents several research questions around crosslingual transfer learning and multi-task learning (MTL) [23, 1, 2, 27, 45, 28].
Recent progress in multilingual sequence modeling, with multilingual translation as a representative application, been extending the scale of massive multilingual learning, with increasing number of languages [2, 10, 58] , the amount of data [2, 12], as well as model size [27, 13]. Despite the power-law scaling of (English-only) language modeling loss with model, data and compute [25], it has been found that multilingual models do not always benefit from scaling up model and data size, especially multilingual machine translation to multiple languages even after exploiting language proximity with external linguistic knowledge [12, 27, 51].
There has been limited understanding of the optimization aspects of multilingual models. Multilingual training is often implemented as monolithic with data from different languages simply combined.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Challenges were observed when training with imbalanced data [2, 26], which is common for multilingual NLP as only a few languages are rich in training data (high resource languages) while the rest of the languages in the world has zero or low training data (low resource languages) [24]. This has been mostly treated as a “data problem" with a widely used work-around of upsampling low resource languages’ data to make the data more balanced [2, 26, 10, 34, 58].
In this paper, we fill the gap by systematically study the optimization of multilingual models in the task of multilingual machine translation with Transformer architecture. Our contribution is twofold.
First, we reveal the optimization tension between high resource and low resource languages where low resources’ performance often suffer. This had been overlooked in multilingual training but have important implications for achieving the goal of leveraging crosslingual transfer to improve low resources. We analyze the training objectives of multilingual models and identify an important role played by local curvature of the loss landscape, where “sharpness" causes interference among languages during optimization. This hypothesis is verified empirically, where we found optimization tension between high and low resource languages. They compete to update the loss landscape during early stage of training, with high resource ones dominating the optimization trajectory during the rest of training. Existing approaches such as upsampling low resources implicitly reduce this tension by augmenting training distribution towards more uniform. We show that this approach is not robust to different data characteristics, where it suffers from either overfitting low resources or underfitting high resources.
Second, we propose a principled training algorithm for multilingual models to mitigate such tension and effectively improve all languages. Our algorithm explicitly learn the weighting of different languages’ gradients with a meta-objective of guiding the optimization to “flatter" neighborhoods with uniformly low loss (Curvature-Aware Task Scaling, CATS). Compared to static weighting implied by sampling probabilities of the data distribution, our method effectively reduced the optimization tension between high and low resource languages and improves the Pareto front of generalization. On common benchmarks of multilingual translation, CATS consistently improves low resources in various data conditions, +0.8 BLEU on TED (8 languages, 700K sentence pairs), +2.2 BLEU on WMT (10 languages, 30M sentence pairs), and +1.3 BLEU on OPUS-100 (100 languages, 50M sentence pairs) without sacrificing performance on high resources. Furthermore, CATS can effectively leverage model capacity, yielding better generalization in overparameterized models. The training algorithm is conceptually simple and efficient to apply to massive multilingual settings, making it a suitable approach for achieving equitable progress in NLP for every language.

heading: 2 Related Work
text: Multilingual Learning and Multi-task Learning. Massive multilingual models have been gaining increasing research interest as a result of recent advancement in model and data scaling [17, 10, 23, 2, 19, 58]. However, multilingual models are often trained using the same monolithic optimization as is used in training single-language models. To deal with imbalanced training data, upsampling low resource languages (and downsampling high resource languages) was first proposed in massive multilingual machine translation in order to improve low resource performance [23, 2]. It was adopted in recent state-of-the-art multilingual pretraining [10, 34, 58].
Although recent work has looked into maximizing positive transfer across languages by learning parameter-sharing sub-networks conditional on languages [28, 31], few prior work had looked into the blackbox of multilingual optimization. Relevant efforts focus on dynamic data selection, such as adaptive scheduling [22] and MultiDDS (Differentiable Data Selection), which dynamically selects training examples based on gradient similarity with validation set [53]. Although they share the same motivation of treating multilingual training as a multi-objective optimization problem, data selection introduces additional computation and slows down training.
This work also adds to a growing interest in addressing interference in multilingual models, known as “the curse of multilinguality"[2, 10], which was initially hypothesized as “capacity bottleneck" in general multi-task learning [5]. Existing work have mainly focused on model architecture, e.g. via manually engineered or learnt parameter sharing across languages based on language proximity[61, 47, 46]. Gradient Vaccine is a recent work addressing interference from an optimization perspective by de-conflicting gradients via projection [54], a general approach MTL [60]. Although conflicting gradients are also examined in this work, they are used as evaluation metrics, where we
show that regularizing local curvature can prevent conflicting gradients from happening in the first place. Overall, we provide new understanding of interference in multilingual tasks by pointing out the optimization tension between high and low resources, which had not been studied before.
Multilingual training is an instance of multi-task learning (MTL) [9, 36, 11, 50, 44, 35]. Task balancing has been studied from both architecture and optimization perspectives [33, 6, 48, 60]. Among them, GradNorm [6] and MGDA [48] are closely related, where the GradNorm adjust gradient norms to be uniform across tasks, and MGDA adapts gradient-based multi-objective optimization to modern neural networks with efficient training.
Sharpness of the Loss Landscape and Generalization. Our analysis and proposed method are closely related to recent findings on the generalization of neural networks and optimization behaviors including geometric properties of loss landscape during training. It was observed that the “sharpness" of the loss landscape grows rapidly especially during the early phase of training by examining relevant metrics such as largest eigenvalues of Hessian, spectral norm of empirical Fisher Information matrix, etc. [8, 20]. Gradient alignment, measured as cosine similarity of gradients computed on different training examples, indicates the covariance of gradients, which was shown to be related to generalization performance [15, 21, 32]. Our work applies those insights from optimization literature to analyze the training dynamics of multilingual learning with the Transformer architecture. A closely-related concurrent work [14] has verified the effectiveness of regularizing loss sharpness in improving generalization, but in a single-task setting on computer vision benchmarks.

heading: 3 Demystifying Optimization Challenges in Multilingual Training
text: Notations. In multilingual translation task, we typically train one model with training data from a set of languagesN := {l1, ..., lN}, and measure its generalization performance on a set of evaluation languagesM. We use “language" and “task" interchangeably throughout the paper. We introduce the following notations:
• Dn := {(x(n)i , y (n) i )} refers to the set of labeled examples of language ln. A prominent property of
multilingual translation task is the highly imbalanced data, i.e. the distribution of |Dn|, denoted as p|Dn|, is usually heavy-tailed, with a few high resource (HiRes) languages, and a large number of low resource (LoRes) ones [2, 55].
• Let θ ∈ Rd be the parameters of the multilingual model. • Ln := E(x,y)∼Dn [L(ŷ, y)] as the expected loss as measured on language n, with token-level cross
entropy loss L(ŷ, y) = − ∑ i y |V| i log ŷ, where V is the vocabulary the target language(s) which the model translates to.
• Denote ∇ as gradient,∇2 as Hessian H, tr(.) as trace, and ‖.‖2 as Euclidean (l2) norm.

heading: 3.1 Optimization Objective
text: Our analysis starts by taking a closer look at the training and generalization objective of multilingual models. Generalization performance is measured on a set of evaluation languages M weighted by their importance wn, LMulti := ∑|M| n=1 wnLn(x, y|θ). For multilingual translation, wn = 1|M| , n = 1, ...|M| indicating that all translation directions have equal importance. This population loss is usually optimized by minimizing empirical risk, i.e. a training loss L̂Multi := minθ ∑|N | n=1 αnL̂n(x, y;θ), where αn corresponds to the weight of language n’s loss during training.
We point out that as an important optimization choice, α := {αn} is either left implicit or chosen heuristically in existing work of multilingual training. A common choice for α resembles the sampling probability of language ln in a mini-batch according to training data distribution p 1 T {|Dn|} with a temperature hyperparameter T . For example, T = 1 corresponds to sampling proportional to training data sizes [|D1|, ...|DN |]; T =∞ corresponds to uniformly sampling across languages. For most multilingual tasks, training data is highly imbalanced with a few high resource languages accounting for majority of the training data and a long tail of low resource languages. Practitioners manually tune T to make minimizing training loss L̂Multi a proxy of optimizing the population loss
LMulti. Massive multilingual translation and pretraining with real-world datasets of 100 languages usually upsample low resource languages (e.g. T = 5) to improve the averaged generalization performance [2, 34, 10, 58].
3.2 Local Curvature and Robust Multi-task Optimization
We analyze how local curvature affects optimization efficiency of multilingual training. Task interference has been observed to pose optimization challenges in general multi-task learning and continual learning, where improvement in one task hurts performance on another task [45, 48, 30, 60, 38].
In the context of training multilingual models, we define interference from the optimization perspective, as the change of empirical risk L1 as a result of an update to θ dominated by empirical risk minimization for another language’s loss L2. In Figure 1, we illustrate how curvature of the loss landscape affects interference similar to the analysis which has been done in continual learning [38]. After a gradient step driven by ∇L2, the shared model parameters moved from θ∗1 to θ2. The corresponding changes in L1, i.e. L1(θ2)− L1(θ∗1) is affected by the curvature.
I(L1,L2) ∆ = L1(θ2)− L1(θ∗1) ≈ (θ2 − θ∗1)>∇L1(θ∗1) +
1 2 ∇2L1(θ∗1)‖θ2 − θ∗1‖22 (1)
This relationship is also summarized in Eq. (1). We make an assumption that in a small neighborhood of θ∗1, the loss surface is almost convex[7, 16]. Then we can apply a second-order Taylor expansion of L1(θ2) in Eq. (2), and derive the connection between I(L1,L2) and the Hessian H(θ) = ∇2L(θ) which indicates the local curvature of loss surface.
L1(θ2) ≈ L1(θ∗1) + (θ∗2 − θ∗1)>∇L1(θ∗1) + 1
2 (θ2 − θ∗1)>∇2L1(θ∗1)(θ2 − θ∗1) (2)
Algorithm 1 Curvature Aware Task Scaling (CATS). Input: a mini-batch B with K languages; number of inner loop updates m. Output: converged multilingual model θ.
1: Initialize θ, α = 1 , λ = 0. 2: while not converged do 3: for i from 1 to m do 4: for n = 1, ...,K do 5: Compute gn = ∇θL̂n(Bn) 6: g̃n = clone(gn) 7: α̃n = clone(αn) 8: end for 9: Update θ with ∑K n=1 α̃ngn
10: end for 11: Compute ∇αL̂metaα ,∇λL̂metaα ac-
cording to Eq. (4) using g̃n for ∇θL̂n 12: Update α with gradient descent 13: Update λ gradient ascent 14: end while
Since∇L1(θ∗1) ≈ 0 at critical point for L1, the major contributing factor to interference is local curvature, measured by the spectral norm of H(θ∗1), as well as the magnitude of parameter update ‖θ2 − θ∗1‖22. We hypothesize that the optimization tension defined above affects low resource languages more than high resource ones, assuming that low-resource tasks have lower sample complexity, and thus are more likely to reach local minima during early stage of training.

heading: 3.3 Meta-learn α with curvature regularization
text: Motivated by previous section’s analysis, we propose a principled optimization procedure for multilingual translation with a meta objective of regularizing local curvature of the loss landscape of shared parameters θ. We explicitly learn the task weighting parameters α := [αn]n=1,...,N so as to minimize the trace norm of empirical Fisher information tr(F ) = E(x,ŷ)[‖ ∑N n=1 αn∇θL̂n‖2], which is an approximation of tr(H) as was proposed in optimization literature [20, 52]. To leverage distributed training in modern deep learning, we estimateα from a mini-batch B with K languages as is shown in Eq. (3):
min α1,...,αK ‖ K∑ n=1 αn∇θL̂n(θ)‖2 s.t. K∑ n=1 αn = 1, αn ≥ 0, ∀n (3)
We treat solving Eq. (3) along with minimizing empirical risk L̂Multi as a multi-objective optimization problem, and optimizing the corresponding Lagrangian [4]. The overall training objective is as follows:
L̂CATS(θ,α,λ) = K∑
n=1
αnL̂n − λc( − ‖ K∑
n=1
αn∇θL̂n‖2) + λs( K∑
n=1
αn − 1)2 − K∑
n=1
λn(αn − ) (4)
We learn both model parameters θ task weighting parameters α simultaneously with bi-level optimization, where update θ in the inner loop and update α and Lagrangian multipliers λ := [λc, λs, λ1, ..., λN ] in the outer loop. We refer to proposed algorithm as CATS (Curvature-Aware Task Scaling) with the detailed training procedure described in Algorithm 1.
4 Experiments
4.1 Experimental setup
Datasets. We experiment on three public benchmarks of multilingual machine translation with varying characteristics of imbalanced data as is shown in Table 1. They are representative in terms of the number of languages |N |, total volume of training data |Dn| measured as the number of sentence pairs in millions (M), and the entropy of data distribution H|Dn| indicating the degree of data imbalance. For example, distributions with multiple high resource languages are covered by experiments on the TED and
OPUS100 datasets, while experiments on WMT dataset cover a unique scenario of extremely skewed distribution with one high resource language and a long tail of low resource ones, which is not uncommon in real world applications. Additional details and statistics of the datasets are provided in Appendix A.
Models. We use the Transformer architecture, and the same model size and configuration as were used in corresponding baselines. For experiments on TED, the model is a 6-layer encoder-decoder, 512 hidden dimension, 1024 FFN dimension, and 4 attention heads [53]. For experiments on OPUS100, the model is the Transformer-base configuration as is used in [62]. We use Transformer-base for WMT experiments. We use the same preprocessed data by the MultiDDS baseline authors [53], and followed the same procedure to preprocess OPUS-100 data released by the baseline [62]. All models are trained with the same compute budget for comparison. We provide detailed training hyperparameters in Appendix B.
Baselines. We compare to strong baselines used in state-of-the-art multilingual translation and relevant approaches in generic multi-task learning:
• Proportional sampling. This is a straightforward yet common approach used in practice by training on combined training data from all languages [62], which corresponds to αn = p 1 T , T = 1.
• Upsampling low resources T = 5. This has been adopted in state-of-the-art multilingual transformers as is discussed in Section 3.1.
• MultiDDS[53]: A recently proposed method to balance training losses in multilingual training. It learns to select training examples among different languages based on gradient similarity with validation set. Although it does not address the optimization challenges studied in this work, it shares the same interest of improving generalization performance and utilizing training examples in a dynamic and adaptive fashion.
• GradNorm[6]: We also compare to a closely related approach proposed for general multi-task learning. The key difference is that the objective in GradNorm is to rescale each task’s gradients to be closer to the average gradients norm (i.e. only based on ∇Li), while the objective in our approach is to minimize the second-order derivative∇2Li which corresponds to the curvature of the loss landscape.
Evaluation of Optimization. For the analysis of loss landscape, we look into the top eigenvalue (λmax) of H computed from examples in a mini-batch with the power iteration method described in [59] for each language. To evaluate optimization efficiency, we analyze gradient similarity using the same metrics introduced in [60]. At each training step, with g1 = ∇L1, g2 = ∇L2 computed from training examples of two languages, we measure the following metrics as indicators of gradient similarity:
• Gradient direction similarity (alignment): cosφ(g1, g2) where φ is the angle between g1 and g2.
• Gradient magnitude (L2 norm) similarity: γ(g1, g2) = 2‖g1‖2‖g2‖2‖g1‖22+‖g2‖22 .
Evaluation of Generalization. We verify whether improved optimization leads to better generalization. We report both token-level loss (negative log-likelihood, NLL↓) and BLEU scores (↑) on hold-out datasets. we choose the best checkpoint by validation perplexity and only use the single best model without ensembling.
We conduct most experiments in multilingual one-to-many task as it is a more challenging task than many-to-one, and represents the core optimization challenges underlying many-to-many task. Additional experiments with many-to-one are provided in Appendix C.
4.2 Results
4.2.1 Robust Optimization
E.SA E.FC D.SA D.EA D.FC
M2O,HiLo,R O2M,HiLo,R O2M,HiLo,D
O2M,HiHi O2M,LoLo
E.SA E.FC D.SA D.EA D.FC
0.00
0.01
0.02
0.03
0.7
0.8
Figure 3: Gradients similarity of different Transformer parameters (x-axis) in common multilingual translation tasks (y-axis), measured as gradients direction similarity (Left), and gradients norm similarity (Right).
Abrasive Gradients between HiRes and LoRes. First, we illustrate the optimization tension between high resource languages and low resource languages (HiLo). We examine abrasive gradients in Figure 3, which shows a fine-grained view of gradients similarity in terms of both direction (left) and magnitude (right) for different parameters in the Transformer architecture: encoder self attention (E.SA), encoder feed-
forward (E.FC), decoder self-attention (D.SA), decoder encoder attention (D.EA), and decoder feedforward (D.FC). We can see that gradients are more similar when training with balanced data, i.e. HiHi and LoLo, compared to HiLo for two fundamental multilingual translation tasks one-to-many (O2M) and many-to-one (M2O), indicating that the amount of training data is less a cause of the problem compared to the distribution of training data across languages. Furthermore, contrary to
common heuristics which use language family information to determine which languages share parameters [12], language proximity, i.e. related (R) vs. diverse (D), has a smaller effect on gradient similarity compared to data imbalance.
Next, we look into optimization tension through the lens of curvature. Figure 2 (top) plots the top eigenvalues λmax when training a multilingual model in the HiLo setting. We can see that with uniform sampling (equal weighting, T = ∞), LoRes updates the shared parameter space in the direction towards higher curvature during the early stage of training, while HiRes dominates the optimization trajectory (with the LoRes having negative top eigenvalue) for the rest of the training. We found that common heuristics of T = 5 reduces the early-stage sharpness although the loss landscape is still dominated by HiRes. In contrast, the proposed optimization approach (CATS αn) mitigated the abrasive optimization and increased gradients alignment as is illustrated Figure 2 (bottom).
Training without layer normalization. We take a closer look at gradients alignment for different types of parameters in the Transformer architecture. We found that gradient alignments for layer normalization parameters (LayerNorm, LN) are much noisier than other parameters as is shown in Figure 4a. Extensive work has provided empirical and theoretical results of LayerNorm be being critical for training stability of Transformers[3, 57, 56]. Our work is the first to show that it is a doubleedged sword in multilingual training and exacerbates the tension between HiRes and LoRes, which competes to set the gain and bias parameters in LayerNorm, where the final weights highly depend on training data distributions. However, simply removing the gain and bias parameters causes training instability[32]. In Figure 4b, we show that with CATS optimization in place, LayerNorm can be removed yet stable training can still be achieved without further adjusting other hyperparameters such as decreasing learning rate, etc. As a result, the generalization gap (difference between training and validation loss) for low resource is greatly reduced possibly due to that the learnt LayerNorm parameters would otherwise be biased towards HiRes.

heading: T =∞ 5.4 9.6 24.6 22.5 39.3 19.9 22.2 15.9 19.9 15.5 12.7 10.8 14.8 22.7 30.3 31.3 37.2 16.9 22.1 15.3

heading: T = 5 5.5 10.4 24.2 22.3 39.0 19.6 22.5 16.0 19.9 15.6 13.4 10.4 14.7 23.8 30.4 32.3 37.8 17.6 22.5 15.6
text: annotated with underscore. * indicates the improvements are statistically significant with p < 0.05.
Robust to different training data distributions p|Dn|. We show that the improved optimization is reflected in generalization performance. In Figure 5, we report the training and validation loss of HiRes and LoRes languages throughout training on two representative benchmarks: TED (top) and WMT (bottom) with different degrees of data imbalance. We can see that the optimal temperature hyperparameters T vary given the distinct distributions of training data p|Dn| in these two benchmarks. For example, on WMT dataset where the training data distribution across languages is more skewed, upsampling LoRes (T = 5,∞) is beneficial, as was observed in other large-scale multilingual
datasets [2, 10, 34]. However, T = 5 easily leads to overfitting for LoRes on a more balanced dataset such as the TED dataset. Compared to existing approaches of using a hyperparameter T , CATS is more robust as a principled way to address varying distributions of training data. It does not need manual tuning of a hyperparameter T nor suffer from overfitting or underfitting. On both benchmarks, it consistently improves LoRes without sacrificing HiRes.
4.2.2 Generalization LoRes MidRes HiRes Avg.
kk vi tr ja nl it ar ro hi de All Lo |Dn| (M) 0.1 0.1 0.2 0.2 0.2 0.3 0.3 0.6 0.8 27.8 T = 1 1.1 27.3 13.5 13.3 28.6 25.2 14.0 29.8 14.1 26.8 19.4 17.6 T = 5 1.2 31.3 15.5 14.9 30.9 30 16.9 31.1 14.7 22.3 20.9 20.1 CATS α 2.0 33.1 19.2 16.7 33.3 32.1 19.6 35.4 18.0 25.9 23.5* 22.3*
Table 3: Detailed performance on the WMT benchmark. Compared to the strongest baselines (T = 1, 5), CATS drastically improves low and mid resource languages. Multilingual training which achieves the best BLEU is in bold and the best baseline is annotated with underscore.
We report translation quality on three representative benchmarks whose detailed characteristics are described in Table 1.
TED. Table 2 summarizes translation quality improvement on the TED corpus. To control for language similarity, we evaluate on the 8-language
benchmark with both Related (4 LoRes and 4 HiRes from the same language family) and Diverse (4 LoRes and 4 HiRes without shared linguistic properties) settings as is used in [53]. We found simple mixing (T = 1) is a very strong baseline and complex training methods such as GradNorm
and MultiDDS do not generalize well (TED is a smaller dataset and more prone to overfitting). CATS achieves +0.4 ∼ 0.7 BLEU on average for LoRes and +0.7 ∼ 0.8 BLEU overall.
WMT. On a larger and more imbalanced 10-language WMT datset with 30.5M sentence pairs, we found regularizing local curvature with meta learnt α (CATS) is more beneficial. CATS improved low and mid resource languages across the board as is shown in Table 3. Compared to the strongest baseline (T = 5), CATS achieved +2.2 average BLEU for low resource languages and +2.6 BLEU across all languages. Furthermore, CATS is more sample efficient as is shown in Figure 6.
OPUS-100. Finally, we test the scalability of CATS in massive multilingual setting using the OPUS-100 benchmark[62]. Results are summarized in Table 4. Consistent with previous results on OPUS-100 [62, 61] as well as other massive multilingual model in [2], upsampling low resource with T = 5 improves BLEU scores for low resource but at the cost of accuracy for high resource ones. The trade-off is the opposite for proportional sampling (T = 1). CATS achieves the the best of both worlds, especially +1.3 BLEU on low resource, +1.4 BLEU on mid resource, and +1.0 BLEU on high resource compared to the strong baseline of T = 5. We also compare to a recent state-of-the-art model, conditional language-specific routing (CLSR) [61], which adds additional language specific parameters (157M parameters in total compared to 110M
parameters in our experiments). By improving optimization without increasing model capacity, CATS still outperforms CLSR on low resource by +0.9 BLEU.
5 Analysis
Robust to overparameterization. Scaling up model size has been of central interest in recent development of massive multilingual models such as GShard and Switch Transformer with trillions of parameters [27, 13]. Training overparameterized models a shorter amount of time have shown to be more efficient than training a smaller model for longer time [29, 39, 25]. Therefore, we are interested in understanding CATS’ performance in training overparameterized models. Figure 7 plots change in generalization (measured by BLEU score difference) as we increase model capacity. CATS can effectively benefit from larger model capacity, especially in the overaparameterized regime (300M parameters for TED dataset) where performance begins to degradate with standard training.
Performance with large batch size. Large batch size is an effective way to speed up training without sacrificing performance[18, 49]. However, heterogeneous training data do not necessarily benefit from large batch size training[37]. This is a challenge for multilingual training as is shown in Figure 8, where increasing batch size hurts generalization for LoRes with a common practice of upsampling (T = 5). This is likely due to LoRes is prone to overfitting (illustrated in Figure 5) and larger batch size exacerbates it. The same batch size without upsampling leads to improved generalization (T = 1). In comparison, CATS can effectively benefit from larger batch size due to its adaptive rescaling of gradients.
Ablation layer normalization. To further understand the effect from removing layer normalization which mitigates LoRes overfitting as was shown in Section 4.2.1, we provide an ablation study on its role when combined with CATS. We ran experiments on TED Diverse dataset with the strongest baseline (T = 1). In Table 5, we can see that CATS brings additional +0.5 BLEU to low resources on top of the +0.3 BLEU improvement from removing layer normalization parameters (-LN).

heading: 6 Conclusion
text: In this work, we look under the hood of monolithic optimization of multilingual models. We unveil optimization challenges arising from imbalanced training data, where low resource languages have been sub-optimally optimized. We proposed a principled optimization algorithm for multilingual training with adaptive gradient rescaling of different languages where the scaling is learnt with a meta-objective of guiding the optimization to solutions with low local curvature.
We evaluated the proposed method on three representative benchmarks (TED, WMT and OPUS100) which cover a wide range of imbalanced data distributions commonly seen in real word multilingual datasets. Compared to existing methods of simply mixing the data or manually augmenting the data distribution to be more “balanced” with a temperature hyperparameter, the proposed training method demonstrates robust optimization and consistently improves generalization for low resource languages. Further analysis shows that the proposed approach is suitable for the realistic training settings of large-scale multilingual learning such as overparameterized models, large batch size, highly imbalanced data, as well as large number of languages etc., paving the way for advancing massive multilingual models which truly benefit low resource languages.
Broader Impact. Recent progress in NLP enabled by scaling up model size and data is widening the gap of technology equity between high resource languages (such as English) and low resource languages. Multilingual model is a promising approach to close this gap. However, current language agnostic multilingual models does not effectively improve low resources for various reasons to be understood. Our investigation in optimization is one step towards building truly inclusive multilingual models.
This work has several limitations which we hope to extend in future work. We did not run experiments on translation between non-English directions including zero-shot. The absolute size of the overparameterized model in our analysis is relatively small compared to the state-of-the-art of model scaling of billions of parameters. As an language generation application, machine translation model could produce unsafe output or hallucination.

Review summaries:
This paper studies the problem of robustness optimization in multilingual machine translation models. In response to this problem, the author empirically found that in the loss optimization of multilingual machine translation models under unbalanced data scenarios, "sharpness" of local curvature in the loss landscape causes interference among languages, which caused instability in training. Inspired by the robust optimization work of Multi-task Learning, the paper proposes the Curvature Aware Task Scaling (CATS) optimization algorithm, which alleviates the phenomenon that different translation directions competitive update the loss landscape during the early stage of training in multi-language translation training, while high resource ones dominating the optimization trajectory during the rest of training.

Paper presents a curvature aware bi-level optimization algorithm applied to multilingual/multi-task machine translation. The proposed algorithm is designed to encourage the loss is reduced along the directions of low-curvature (with the implicit inductive bias of flat minimas being better at generalization). Akin to "Sharpness-Aware Minimization (Foret et al. 2020)" but in a multi-task environment. Experiments on three different datasets tests the effectiveness of the proposed method (called Curvature Aware Task Scaling - $CATS$). Results show clear trend that the proposed method is highly effective at improving the quality of low-resource languages. Further analysis provided to understand the tradeoffs, sensitivity and dynamics of the algorithm.   - Making use of curvature information during the optimization of multi-task models for translation is novel and expected to be impactful in the long run (while the paper is not demonstrating convincing empirical results, it is showing the potential of using higher order signals during optimization, which I found to be an important contribution).

This work looks at the optimization of multilingual NMT models. Typically, the data across different language pairs is highly skewed and prior work either upsamples low-resource languages such that it's uniform or uses a temperature (T=5) to upsample lowRes and downsample HiRes. This work proposes a new optimization algorithm called CATS which encourages solutions with low local curvature. The proposed technique has been shown to be effective when compared to 3 temperature-based baselines (T=1,5,infinity) and two dynamic methods multiDDS and GradNorm. Experiments are performed on 3 public datasets: TED, WMT and OPUS-100. Robustness of the technique to overparameterization and to large batches has been studied.

This paper theoretically motivates and then implements a curvature-based, online, per-language learning rate adjustment to improve the performance of low-resource languages in multilingual machine translation (MT) models. Experiments in three scenarios varying in multilinguality and data size show that the technique works consistently, outperforming a widely-adopted temperature-based sampling baseline. Through a number of analysis and ablation experiments, they are able to provide insight into how and why the optimization technique works.