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: The paper studies the behaviour of gradient similarities across languages in multilingual NMT models. They find gradient similarities mirror language similarity. Hence, they look at method to gradient-based methods for multilingual NMT. They apply PCGrad to the multilingual NMT task and also extend this method to address the cases when the task gradients have only weak similarity.