Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models

1 INTRODUCTION . Modern multilingual methods , such as multilingual language models ( Devlin et al. , 2018 ; Lample & Conneau , 2019 ; Conneau et al. , 2019 ) and multilingual neural machine translation ( NMT ) ( Firat et al. , 2016 ; Johnson et al. , 2017 ; Aharoni et al. , 2019 ; Arivazhagan et al. , 2019 ) , have been showing success in processing tens or hundreds of languages simultaneously in a single large model . These models are appealing for two reasons : ( 1 ) Efficiency : training and deploying a single multilingual model requires much less resources than maintaining one model for each language considered , ( 2 ) Positive cross-lingual transfer : by transferring knowledge from high-resource languages ( HRL ) , multilingual models are able to improve performance on low-resource languages ( LRL ) on a wide variety of tasks ( Pires et al. , 2019 ; Wu & Dredze , 2019 ; Siddhant et al. , 2020 ; Hu et al. , 2020 ) . Despite their efficacy , how to properly analyze or improve the optimization procedure of multilingual models remains under-explored . In particular , multilingual models are multi-task learning ( MTL ) ( Ruder , 2017 ) in nature but existing literature often train them in a monolithic manner , naively using a single language-agnostic objective on the concatenated corpus of many languages . While this approach ignores task relatedness and might induce negative interference ( Wang et al. , 2020b ) , its optimization process also remains a black-box , muffling the interaction among different languages during training and the cross-lingual transferring mechanism . In this work , we attempt to open the multilingual optimization black-box via the analysis of loss geometry . Specifically , we aim to answer the following questions : ( 1 ) Do typologically similar languages enjoy more similar loss geometries in the optimization process of multilingual models ? ( 2 ) If so , in the joint training procedure , do more similar gradient trajectories imply less interference between tasks , hence leading to better model quality ? ( 3 ) Lastly , can we deliberately encourage ∗Work done during an internship at Google . more geometrically aligned parameter updates to improve multi-task optimization , especially in real-world massively multilingual models that contain heavily noisy and unbalanced training data ? Towards this end , we perform a comprehensive study on massively multilingual neural machine translation tasks , where each language pair is considered as a separate task . We first study the correlation between language and loss geometry similarities , characterized by gradient similarity along the optimization trajectory . We investigate how they evolve throughout the whole training process , and glean insights on how they correlate with cross-lingual transfer and joint performance . In particular , our experiments reveal that gradient similarities across tasks correlate strongly with both language proximities and model performance , and thus we observe that typologically close languages share similar gradients that would further lead to well-aligned multilingual structure ( Wu et al. , 2019 ) and successful cross-lingual transfer . Based on these findings , we identify a major limitation of a popular multi-task learning method ( Yu et al. , 2020 ) applied in multilingual models and propose a preemptive method , Gradient Vaccine , that leverages task relatedness to set gradient similarity objectives and adaptively align task gradients to achieve such objectives . Empirically , our approach obtains significant performance gain over the standard monolithic optimization strategy and popular multi-task baselines on large-scale multilingual NMT models and multilingual language models . To the best of our knowledge , this is the first work to systematically study and improve loss geometries in multilingual optimization at scale . 2 INVESTIGATING MULTI-TASK OPTIMIZATION IN MASSIVELY MULTILINGUAL MODELS . While prior work have studied the effect of data ( Arivazhagan et al. , 2019 ; Wang et al. , 2020a ) , architecture ( Blackwood et al. , 2018 ; Sachan & Neubig , 2018 ; Vázquez et al. , 2019 ; Escolano et al. , 2020 ) and scale ( Huang et al. , 2019b ; Lepikhin et al. , 2020 ) on multilingual models , their optimization dynamics are not well understood . We hereby perform a series of control experiments on massively multilingual NMT models to investigate how gradients interact in multilingual settings and what are their impacts on model performance , as existing work hypothesizes that gradient conflicts , defined as negative cosine similarity between gradients , can be detrimental for multi-task learning ( Yu et al. , 2020 ) and cause negative transfer ( Wang et al. , 2019 ) . 2.1 EXPERIMENTAL SETUP . For training multilingual machine translation models , we mainly follow the setup in Arivazhagan et al . ( 2019 ) . In particular , we jointly train multiple translation language pairs in a single sequenceto-sequence ( seq2seq ) model ( Sutskever et al. , 2014 ) . We use the Transformer-Big ( Vaswani et al. , 2017 ) architecture containing 375M parameters described in ( Chen et al. , 2018a ) , where all parameters are shared across language pairs . We use an effective batch sizes of 500k tokens , and utilize data parallelism to train all models over 64 TPUv3 chips . Sentences are encoded using a shared source-target Sentence Piece Model ( Kudo & Richardson , 2018 ) with 64k tokens , and a < 2xx > token is prepended to the source sentence to indicate the target language ( Johnson et al. , 2017 ) . The full training details can be found in Appendix B . To study real-world multi-task optimization on a massive scale , we use an in-house training corpus1 ( Arivazhagan et al. , 2019 ) generated by crawling and extracting parallel sentences from the web ( Uszkoreit et al. , 2010 ) , which contains more than 25 billion sentence pairs for 102 languages to and from English . We select 25 languages ( 50 language pairs pivoted on English ) , containing over 8 billion sentence pairs , from 10 diverse language families and 4 different levels of data sizes ( detailed in Appendix A ) . We then train two models on two directions separately , namely Any→En and En→Any . Furthermore , to minimize the confounding factors of inconsistent sentence semantics across language pairs , we create a multi-way aligned evaluation set of 3k sentences for all languages2 . Then , for each checkpoint at an interval of 1000 training steps , we measure pair-wise cosine similarities of the model ’ s gradients on this dataset between all language pairs . We examine gradient similarities at various granularities , from specific layers to the entire model . 1We also experiment on publicly available dataset of WMT and obtain similar observations in Appendix C. 2In other words , 3k semantically identical sentences are given in 25 languages . 2.2 OBSERVATIONS . We make the following three main observations . Our findings are consistent across different model architectures and settings ( see Appendix C and D for more results and additional discussions ) . 1 . Gradient similarities reflect language proximities . We first examine if close tasks enjoy similar loss geometries and vice versa . Here , we use language proximity ( defined according to their memberships in a linguistic language family ) to control task similarity , and utilize gradient similarity to measure loss geometry . We choose typological similarity because it is informative and popular , and we leave the exploration of other language similarity measurements for future work . In Figure 1 , we use a symmetric heatmap to visualize pair-wise gradient similarities , averaged across all checkpoints at different training steps . Specifically , we observe strong clustering by membership closeness in the linguistic family , along the diagonal of the gradient similarity matrix . In addition , all European languages form a large cluster in the upper-left corner , with an even smaller fine-grained cluster of Slavic languages inside . Furthermore , we also observe similarities for Western European languages gradually decrease in West Slavic→South Slavic→East Slavic , illustrating the gradual continuum of language proximity . 2 . Gradient similarities correlate positively with model quality . As gradient similarities correlate well with task proximities , it is natural to ask whether higher gradient similarities lead to better multi-task performance . In Figure 2 ( a ) , we train a joint model of all language pairs in both En→Any and Any→En directions , and compare gradient similarities between these two . While prior work has shown that En→Any is harder and less amenable for positive transfer ( Arivazhagan et al. , 2019 ) , we find that gradients of tasks in En→Any are indeed less similar than those in Any→En . On the other hand , while larger batch sizes often improve model quality , we observe that models trained with smaller batches have less similar loss geometries ( Appendix D ) . These all indicate that gradient interference poses great challenge to the learning procedure . To further verify this , we pair En→Fr with different language pairs ( e.g . En→Es or En→Hi ) , and train a set of models with exactly two language pairs5 . We then evaluate their performance on the En→Fr test set , and compare their BLEU scores versus gradient similarities between paired two tasks . As shown in Figure 2 ( b ) , gradient similarities correlate positively with model performance , again demonstrating that dissimilar gradients introduce interference and undermine model quality . 3 . Gradient similarities evolve across layers and training steps . While the previous discussion focuses on the gradient similarity of the whole model averaged over all checkpoints , we now 4Western European includes Romance and Germanic . 5To remove confounding factors , we fix the same sampling strategy for all these models . study it across different layers and training steps . Figure 4 ( c ) shows the evolution of the gradient similarities throughout the training . Interestingly , we observe diverse patterns for different gradient subsets . For instance , gradients between En→Fr and En→Hi gradually become less similar ( from positive to negative ) in layer 1 of the decoder but more similar ( from negative to positive ) in the encoder of the same layer . On the other hand , gradient similarities between En→Fr and En→Es are always higher than those between En→Fr and En→Hi in the same layer , consistent with prior observation that gradients reflect language similarities . In addition , we evaluate the difference between gradient similarities in the multilingual encoder and decoder in Figure 4 ( a ) . We find that the gradients are more similar in the decoder ( positive values ) for the Any→En direction but less similar ( negative values ) for the En→Any direction . This is in line with our intuition that gradients should be more consistent when the decoder only needs to handle one single language . Moreover , we visualize how gradient similarities evolve across layers in Figure 4 ( b ) . We notice that similarity between gradients increase/decrease as we move up from bottom to top layers for the Any→En/En→Any direction , and hypothesize that this is due to the difference in label space ( English-only tokens versus tokens from many languages ) . These results demonstrate that the dynamics of gradients evolve over model layers and training time . Our analysis highlights the important role of loss geometries in multilingual models . With these points in mind , we next turn to the problem of how to improve multi-task optimization in multilingual models in a systematic way . 3 PROPOSED METHOD Following our observations that inter-task loss geometries correlate well with language similarities and model quality , a natural question to ask next is how we can take advantage of such gradient dynamics and design optimization procedures superior to the standard monolithic practice . Since we train large-scale models on real-world dataset consisting of billions of words , of which tasks are highly unbalanced and exhibit complex interactions , we propose an effective approach that not only exploits inter-task structures but also is applicable to unbalanced tasks and noisy data . To motivate our method , we first review a state-of-the-art multi-task learning method and show how the observation in Section 2 helps us to identify its limitation .

Summary: Taking multilingual NMT (MNMT) into account, this work, investigates better model optimization alternative, that is in part can be attributed as a multi-task optimization problem. MNMT's are quite beneficial from different perspectives (improving low-resource languages, efficiency, etc). However, their inherently multi-task nature requires more focus on how to gist out the best possible learning for each of the languages pairs. With a potential impact on the optimization of other multi-task models, this work asks how model the similarity between model gradients is crucial in multi-task settings, and how to best optimize MNMT models focusing on the typologically similarity of languages. By analyzing the geometry of the NMT model objective function, authors indicate that computing similarity along gradient provides information on the relationship between languages and the overall model performance. Authors argue the analysis of the gradient helps to identify the point of limitation in multi-task learning, which the work aims to address, by focusing the parameter updates for tasks that are similar or close in terms of geometrical alignment (also known as Gradient Vaccine /GradVac/).