Robust Optimization for Multilingual Translation with Imbalanced Data

1 Introduction . 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 . 2 Related Work . 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 . 3 Demystifying Optimization Challenges in Multilingual Training . 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 . 3.1 Optimization Objective . 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 .

Summary: 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.