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