G RADIENT V ACCINE : I NVESTIGATING AND I MPROVING M ULTI - TASK O PTIMIZATION IN M ASSIVELY M ULTILINGUAL M ODELS
A BSTRACT
Massively multilingual models subsuming tens or even hundreds of languages pose great challenges to multi -task optimization
While it is a common practice to apply a language-agnostic procedure optimizing a joint multilingual task objective, how to properly characterize and take advantage of its underlying problem structure for improving optimization efﬁciency remains under-explored
In this paper, we attempt to peek into the black-box of multilingual optimization through the lens of loss function geometry.
We ﬁnd that gradient similarity measured along the optimization trajectory is an important signal which correlates well with not only language proximity but also the overall model performance.
Such observation helps us to identify a critical limitation of existing gradient-based multi-task learning methods and thus we derive a simple and scalable optimization procedure, named Gradient Vaccine, which encourages more geometrically aligned parameter updates for close tasks.
Empirically, our method obtains signiﬁcant model performance gains on multilingual machine translation and XTREME benchmark tasks for multilingual language models.
Our work reveals the importance of properly measuring and utilizing language proximity in multilingual optimization and has broader implications for multi-task learning beyond multilingual modeling
I NTRODUCTION
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) Efﬁciency: 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 efﬁcacy, 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, mufﬂing 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.
Speciﬁcally, 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 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 ﬁrst 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 ﬁndings, 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 signiﬁcant 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 ﬁrst work to systematically study and improve loss geometries in multilingual optimization at scale.
I NVESTIGATING M ULTI - TASK O PTIMIZATION IN M ASSIVELY M ULTILINGUAL M ODELS
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; V ´ azquez 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 conﬂicts, deﬁned 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. E XPERIMENTAL S ETUP
For training multilingual machine translation models, we mainly follow the setup in Arivazhagan et al.
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 corpus 1 (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 languages 2 .
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 speciﬁc layers to the entire model.
O BSERVATIONS
We make the following three main observations.
Our ﬁndings are consistent across different model architectures and settings (see Appendix C and D for more results and additional discussions).
Gradient similarities reﬂect language proximities.
We ﬁrst examine if close tasks enjoy similar loss geometries and vice versa.
Here, we use language proximity (deﬁned 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.
Speciﬁcally, 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 ﬁne-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.
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 ﬁnd 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 pairs 5 .
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.
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
[Figure]
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 reﬂect language similarities.
In addition, we evaluate the difference between gradient similarities in the multilingual encoder and decoder in Figure 4(a).
We ﬁnd 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.
P ROPOSED M ETHOD
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 ﬁrst review a state-of-the-art multi-task learning method and show how the observation in Section 2 helps us to identify its limitation.
3. G RADIENT S URGERY
An existing line of work (Chen et al., 2018b; Sener & Koltun, 2018; Yu et al., 2020) has successfully utilized gradient-based techniques to improve multi-task models.
Notably, Yu et al.
[Figure] hypothesizes that negative cosine similarities between gradients are detrimental for multi-task optimization and proposes a method to directly project conﬂicting gradients (PCGrad) , also known as the Gradient Surgery.
As illustrated in the left side of Figure 5(a), the idea is to ﬁrst detect gradient conﬂicts and then perform a “surgery” to deconﬂict them if needed.
Speciﬁcally, for gradients g i and g j of the i -th and j -th task respectively at a speciﬁc training step, PCGrad (1) computes their cosine similarity to determine if they are conﬂicting, and (2) if the value is negative, projects g i onto the normal plane of g j as:
[Equation]
The altered gradient g (cid:48) i replaces the original g i and this whole process is repeated across all tasks in a random order.
For more details and theoretical analysis, we refer readers to the original work.
Now, we can also interpret PCGrad from a different perspective: notice that the gradient cosine similarity will always be zero after the projection, effectively setting a target lower bound.
In other words, PCGrad aims to align gradients to match a certain gradient similarity level, and implicitly makes the assumption that any two tasks must have the same gradient similarity objective of zero .
However, as we shown in Section 2, different language proximities would result in diverse gradient similarities.
In fact, many language pairs in our model share positive cosine similarities such that the pre-condition for PCGrad would never be satisﬁed.
This is shown in the left of Figure 5(b), where
PCGrad is not effective for positive gradient similarities and thus it is very sparse during training in the left of Figure 3.
Motivated by this limitation, we next present our proposed method.
3. G RADIENT V ACCINE
The limitation of PCGrad comes from the unnecessary assumption that all tasks must enjoy similar gradient interactions, ignoring complex inter-task relationships.
To relax this assumption, a natural idea is to set adaptive gradient similarity objectives in some proper manner.
An example is shown in the right of Figure 5(b), where two tasks have a positive gradient similarity of cos( θ ) =
φ ij .
While
PCGrad ignores such non-negative case, the current value of φ ij may still be detrimentally low for more similar tasks such as French versus Spanish.
Thus, suppose we have some similarity goal of cos( θ (cid:48) )
=
φ Tij > φ ij (e.g. the “normal” cosine similarity between these two tasks), we alter both the magnitude and direction of g i such that the resulting gradients match such gradient similarity objective.
In particular, we replace g i with a vector that satisﬁes such condition in the vector space spanned by g i and g j , i.e. a 1 · g i + a 2 · g j .
Since there are inﬁnite numbers of valid combinations of a 1 and a 2 , for simplicity, we ﬁx a 1 = 1 and by applying Law of Sines in the plane of g i and g j , we solve for the value of a 2 and derive the new gradient for the i -th task as 6 :
[Equation]
[Figure] This formulation allows us to use arbitrary gradient similarity objective φ Tij in [ − 1 , 1] .
The remaining question is how to set such objective properly.
In the above analysis, we have seen that gradient interactions change drastically across tasks, layers and training steps.
To incorporate these three factors, we exploit an exponential moving average (EMA) variable for tasks i, j and parameter group k
(e.g. the k -th layer) as:
[Equation]
where φ ( t ) ijk is the computed gradient similarity at training step t , β is a hyper-parameter, and ˆ φ (0) ijk = 0 .
The full method is outlined in Algorithm 1 (Appendix E).
Notice that gradient surgery is a special case of our proposed method
such that φ Tij = 0 .
As shown in the right of Figure 5(a) and 5(b), our method alters gradients more preemptively under both positive and negative cases, taking more proactive measurements in updating the gradients (Figure 3).
We therefore refer to it as Gradient
Vaccine (GradVac) .
Notice that the resulting models will have the same numbers of parameters for deploying as typical MNMT models and thus enjoy the same beneﬁts for memory efﬁciency, while the proposed method will have the same order of complexity with the original multi-task training paradigm as of computational efﬁciency.
E XPERIMENTS
We compare multi-task optimization methods with the monolithic approach in multilingual settings, and examine the effectiveness of our proposed method on multilingual NMT and multilingual language models.
4. G ENERAL S ETUP
We choose three popular scalable gradient-based multi-task optimization methods as our baselnes:
GradNorm (Chen et al., 2018b), MGDA (Sener & Koltun, 2018), and PCGrad (Yu et al., 2020).
For fair comparison, language-specifc gradients are computed for samples in each batch.
The sampling temperature is also ﬁxed at T=5 unless otherwise stated.
For the baselines, we mainly follow the default settings and training procedures for hype-parameter selection as explained in their respective papers.
For our method, to study how sensitive GradVac is to the distribution of tasks, we additionally examine a variant that allows us to control which languages are considered for GradVac.
Speciﬁcally, we search the following hyper-parameters on small-scale WMT dataset and transfer to our large-scale dataset: tasks considered for GradVac { HRL only, LRL only, all task } , parameter granularity { whole model, enc dec, all layer, all matrix } , EMA decay rate β { 1e-1, 1e-2, 1e-3 } .
We ﬁnd { LRL only, all layer, 1e-2 } to work generally well and use these in the following experiments (see Appendix F for more details and results).
4. R ESULTS AND A NALYSIS
WMT Machine Translation.
We ﬁrst conduct comprehensive analysis of our method and other baselines on a small-scale WMT task.
We consider two high-resource languages (WMT14 enfr, WMT19 en-cs) and two low-resource languages (WMT14 en-hi, WMT18 en-tr), and train two models for both to and from English.
Results are shown in Table 1.
[Table] First, we observe that while the naive multilingual baseline outperforms bilingual models on lowresource languages, it performs worse on high-resource languages due to negative interference (Wang et al., 2020b) and constrained capacity (Arivazhagan et al., 2019).
Existing baselines fail to address this problem properly, as they obtain marginal or even no improvement (row 3, 4 and 5).
In particular, we look closer at the optimization process for methods that utilize gradient signals to reweight tasks, i.e. GradNorm and MGDA, and ﬁnd that their computed weights are less meaningful and noisy.
For example, MGDA assigns larger weight for en-fr in the en-xx model, that results in worse performance on other languages.
This is mainly because these methods are designed under the assumption that all tasks have balanced data.
Our results show that simply reweighting task weights without considering the loss geometry has limited efﬁcacy.
By contrast, our method signiﬁcantly outperforms all baselines.
Compared to the naive joint training approach, the proposed method improves over not only the average BLEU score but also the individual performance on all tasks.
We notice that the performance gain on En → Any is larger compared to Any → En .
This is in line with our prior observation that gradients are less similar and more conﬂicting in En → Any directions.
We next conduct extensive ablation studies for deeper analysis: (1) GradVac applied to all layers vs. whole model (row 8 vs. 9): the all layer variant outperforms whole model, showing that setting ﬁne-grained parameter objectives is important.
(2) Constant objective vs. EMA (row 7 vs. 9): we also examine a variant of GradVac optimized using a constant gradient objective for all tasks (e.g. φ Tij = 0 .
5 , ∀ i, j ) and observe performance drop compared to using EMA variables.
This highlights the importance of setting task-aware objectives through task relatedness.
(3) GradVac vs. PCGrad (row 8-9 vs. 5-6): the two GradVac variants outperform their PCGrad counterparts, validating the effectiveness of setting preemptive gradient similarity objectives.
[Table]
Massively Multilingual Machine Translation.
We then scale up our experiments and transfer the best setting found on WMT to the same massive dataset used in Section 2.
We visualize model performance in Figure 6 and average BLEU scores are shown in Table 2.
We additionally compare with models trained with uniform language pairs sampling strategy (T=1) and ﬁnd that our method outperforms both multilingual models.
Most notably, while uniform sampling favor high-resource language pairs more than low-resource ones, GradVac is able to improve both consistently across all tasks.
We observe larger performance gain on high-resource languages, illustrating that addressing gradient conﬂicts can mitigate negative interference on these head language pairs.
On the other hand, our model still perform worse on resourceful languages compared to bilingual baselines, most likely limited by model capacity. [Figure] [Table]
XTREME Benchmark.
We additionally apply our method to multilingual language models and evaluate on the XTREME benchmark (Hu et al., 2020).
We choose tasks where training data are available for all languages, and ﬁnetune a pretrained multilingual BERT model (mBERT) (Devlin et al., 2018) on these languages jointly (see Appendix G for experiment details and additional results).
As shown in Table 3, our method consistently outperforms naive joint ﬁnetuning and other multi-task baselines.
This demonstrates the practicality of our approach for general multilingual tasks.
R ELATED W ORK
Multilingual models train multiple languages jointly (Firat et al., 2016;
Devlin et al., 2018; Lample & Conneau, 2019; Conneau et al., 2019; Johnson et al., 2017; Aharoni et al., 2019; Arivazhagan et al., 2019).
Follow-up work study the cross-lingual ability of these models and what contributes to it (Pires et al., 2019; Wu & Dredze, 2019; Wu et al., 2019; Artetxe et al., 2019; Kudugunta et al., 2019; Karthikeyan et al., 2020), the limitation of such training paradigm (Arivazhagan et al., 2019;
Wang et al., 2020b), and how to further improve it by utilizing post-hoc alignment (Wang et al., 2020c; Cao et al., 2020), data balancing (Jean et al., 2019; Wang et al., 2020a), or calibrated training signal (Mulcaire et al., 2019; Huang et al., 2019a).
In contrast to these studies, we directly investigate language interactions across training progress using loss geometry and propose a language-aware method to improve the optimization procedure.
On the other hand, multilingual models can be treated as multi-task learning methods (Ruder, 2017; Zamir et al., 2018).
Prior work have studied the optimization challenges of multi-task training (Hessel et al., 2019; Schaul et al., 2019), while others suggest to improve training quality through learning task relatedness (Zhang & Yeung, 2012), routing task-specifc paths (Rusu et al., 2016; Rosenbaum et al., 2019), altering gradients directly (Kendall et al., 2018; Chen et al., 2018b; Du et al., 2018;
Yu et al., 2020), or searching pareto solutions (Sener & Koltun, 2018; Lin et al., 2019).
However, while these methods are often evaluated on balanced task distributions, multilingual datasets are often unbalanced and noisy.
As prior work have shown training with unbalanced tasks can be prone to negative interference (Ge et al., 2014; Wang & Carbonell, 2018), we study how to mitigate it in large models trained with highly unbalanced and massive-scale dataset.
C ONCLUSION
In this paper, we systematically study loss geometry through the lens of gradient similarity for multilingual modeling, and propose a novel approach named GradVac for improvement based on our ﬁndings.
Leveraging the linguistic proximity structure of multilingual tasks, we validate the assumption that more similar loss geometries improve multi-task optimization while gradient conﬂicts can hurt model performance, and demonstrate the effectiveness of more geometrically consistent updates aligned with task closeness.
We analyze the behavior of the proposed approach on massive multilingual tasks with superior performance and we believe that our approach is generic and applicable beyond multilingual settings. A CKNOWLEDGMENTS
We want to thank Hieu Pham for tireless help to the authors on different stages of this project.
We also would like to thank Zihang Dai, Xinyi Wang, Zhiyu Wang, Jiateng Xie, Yiheng Zhou, Ruochen
Xu, Adams Wei Yu, Biao Zhang, Isaac Caswell, Sneha Kudugunta, Zhe Zhao, Christopher Fifty,
Xavier Garcia, Ye Zhang, Macduff Hughes, Yonghui Wu, Samy Bengio and the Google Brain team for insightful discussions and support to the work
B T RAINING D ETAILS
For both bilingual and multilingual NMT models, we utilize the encoder-decoder Transformer (Vaswani et al., 2017) architecture.
Following prior work, we share all parameters across all language pairs, including word embedding and output softmax layer.
To train each model, we use a single Adam optimizer (Kingma & Ba, 2014) with default decay hyper-parameters.
We warm up linearly for 30K steps to a learning rate of 1e-3, which is then decayed with the inverse square root of the number of training steps after warm-up.
At each training step, we sample from all language pairs according to a temperature based sampling strategy as in prior work (Lample & Conneau, 2019; Arivazhagan et al., 2019).
That is, at each training step, we sample each sentence from all language pairs to train proportionally to P i =
( L i (cid:80) j L j ) 1 T , where L i is the size of the training corpus for language pair i and T is the temperature.
We set T=5 for most of our experiments.
[Table]
[Table]
C A DDITIONAL R ESULTS ON WMT C. D ATA
We experiment with WMT datasets that are publicly available.
Compared to our dataset, they only contain a relatively small subsets of languages.
Therefore, we select 8 languages (16 language pairs) of 4 language families to conduct the same loss geometries analysis in Section 2.
These languages are detailed in Table x5: French (fr), Spanish (es), Russian (ru), Czech (cs), Latvian (lv), Lithuanian (lt), Estonian (et), Finnish (ﬁ).
We collect all available training data from WMT 13 to WMT 19, and then perform a deduplication process to remove duplicated sentence pairs.
We then use the validation sets to compute gradient similarities.
Notice that unlike our dataset, WMT validation sets are not multi-aligned.
Therefore, the semantic structures of these sentences may introduce an extra degree of noise.
C. V ISUALIZATION
As in Section 2, we compute gradients on the validation sets on all checkpoints and averaged across all checkpoints to visualize our results.
We use similar setups to our previous analysis, including model architectures, vocabulary sizes, and other training details.
The main results are shown in
Figure 8.
Similar to our ﬁndings in Section 2, gradient similarities cluster according to language proximities, with languages from the same language family sharing the most similar gradients on the diagonal.
Besides, gradients in the English to Any directions are less similar compared to the other direction, consistent with our above ﬁndings.
Overall, despite the scale being much smaller in terms of number of languages and sizes of training data, ﬁndings are mostly consistent.
[Figure]
C. V ISUALIZATION ON S MALLER M ODELS
Prior work has shown languages ﬁghting for capacity in multilingual models (Arivazhagan et al., 2019;
Wang et al., 2020b).
Therefore, we are also interested to study the effect of model sizes on gradient trajectory.
Since our larger dataset contains 25 language pairs in a Transformer-Large model, we additionally train a Transformer-Base model using the 8 language pairs of WMT.
We visualize it in Figure 9 and ﬁnd that our observed patterns are more evident in smaller models.
This ﬁnding is consistent across other experiments we ran and indicates that languages compete for capacity with small model sizes thereby causing more gradient interference.
It also shows that our analysis in this work is generic across different model settings.
D A DDITIONAL R ESULTS ON O UR D ATASET
In Figure 1 we show visualization on models trained using Any → En language pairs.
Here, we also examine models trained in the other direction,
En → Any .
As shown in Figure 10, we have similar observations made in Section 2 such that gradient similarities cluster strongly by language proximities.
However, the en-xx model has smaller scales in cosine similarities and more negative values.
For example, Nepali shares mostly conﬂicting gradients with other languages, except for those belonging to the same language family.
This is in line with our above discussion that gradient interference may be a source of optimization challenge, such that the en-xx model is harder to train than the xx-en model.
Moreover, while our previous models are trained using a large batch size for better performance (as observed in prior work (Arivazhagan et al., 2019)), we also evaluate gradients in a model trained with smaller batches (125k tokens) in Figure 11.
Compared to model trained with larger batch sizes, this model enjoy similar patterns but with smaller gradient cosine similarity values, indicating that gradients are less similar.
This presents an additional potential explanation of why larger batch sizes can be more effective for training large models: they may better reﬂect the correct loss geometries such that gradients are less conﬂicting in nature.
For our case, this means larger batches better reﬂect language proximities hence gradients of better quality.
Finally, these results also reveal that gradient similarities are mostly dependent on task relatedness, as even sentence pairs with identical semantic meanings can have negative cosine similarities due to language differences.
1: Require: EMA decay β , Model Components M = { θ k } , Tasks for GradVac G =
{T i } 2: Initialize model parameters 3: Initialize EMA variables ˆ φ (0) ijk = 0 , ∀ i, j, k 4:
Initialize time step t =5: while not converged do 6:
Sample minibatch of tasks B = {T i } 7: for θ k ∈ M do 8:
Compute gradients g
ik ← ∇ θ k L T i , ∀T i ∈ B 9:
Set g (cid:48)
ik ← g ik 10: for T i ∈ G ∩ B do 11: for T j ∈ B \ T i in random order do 12:
Compute φ ( t ) ijk ← g (cid:48)
ik · g jk (cid:107) g (cid:48)
ik (cid:107)(cid:107) g jk (cid:107) 13: if φ ( t ) ijk < ˆ φ ( t ) ijk then 14:
Set g (cid:48)
ik = g (cid:48)
ik + (cid:107) g (cid:48)
ik (cid:107) (ˆ φ ( t ) ijk (cid:113) 1 − ( φ ( t ) ijk ) 2 − φ ( t ) ijk (cid:113) 1 − (ˆ φ ( t ) ijk ) 2 ) (cid:107) g jk
(cid:107) (cid:113) 1 − (ˆ φ ( t ) ijk )·
g jk 15: end if 16:
Update ˆ φ ( t +1) ijk = (1 − β )ˆ φ ( t ) ijk + βφ ( t ) ijk 17: end for 18: end for 19:
Update θ k with gradient (cid:80) g (cid:48)
ik 20: end for 21:
Update t ← t +22: end while
E P ROPOSED M ETHOD D ETAILS
In this part, we provide details of our proposed method, Gradient Vaccine (GradVac).
We ﬁrst show how to derive our formulation in Eq. 2, followed by how we instantiate in practice.
And last, we also study its theoretical property.
E. M ETHOD D ERIVATION
As stated in Section 3, the goal of our proposed method is to align gradients between tasks to match a pre-set target gradient cosine similarity.
An example is shown in Figure 12, where we have two tasks i, j and their corresponding gradients g i and g j have a cosine similarity of φ ij , i.e. cos( θ ) =
[Equation]
=
[Equation]
Then, we want to alter their gradients, such that the resulting new gradients have gradient similarity of some pre-set value φ Tij .
To do so, we replace g i with a new vector in the vector space spanned by g i and g j , g (cid:48)
i = a 1 · g i + a 2 · g j .
Without loss of generality, we set a 1 = 1 and solve for a 2 , i.e. ﬁnd the a 2 such that cos( γ ) =
[Equation]
= φ Tij .
By using Laws of Sines, we must have that:
[Equation]
and thus we can further solve for a 2 as:
[Equation]
We therefore arrive at the update rule in Eq.
Our formulation allows us to set arbitrary target values for any two gradients, and thus we can better leverage task relatedness by setting individual gradient similarity objective for each task pair.
Notice that we can rescale the gradient such that the altered gradients will have the same norm as before.
But in our experiment we ﬁnd it is sufﬁcient to ignore this step.
On the other hand, we note that when φ Tij = 0 , we have that:
[Equation]
This is exactly the update rule of PCGrad in Eq. 1.
Thus, PCGrad is a special case of our proposed method.
E. A LGORITHM IN PRACTICE
Our proposed method is detailed in Algorithm 1.
In our experiments on multilingual NMT and multilingual BERT, we utilize a set of exponential moving average (EMA) variables to set proper pair-wise gradient similarity objectives, as shown in Eq. 3.
This is motivated by our observations in Section 2 such that gradients of different languages in a Transformer model evolve across layers and training steps.
Therefore, we conduct GradVac on different model components independently.
For example, we can do one GradVac on each layer in the model, or just perform a single GradVac on the entire model.
In addition, we also introduce an extra degree of freedom by controlling which tasks to perform GradVac.
This corresponds to selecting a of tasks G and only alter gradients for tasks within this set, as shown in line 10 in Algorithm 1.
Empirically, we ﬁnd performing GradVac by layers and on low-resource languages to work generally the best (See Appendix F for detailed discussion).
E. T HEORETICAL P ROPERTY
Finally, we analyze the theoretical property of our method.
Supper we only have two tasks, and their losses are L 1 and L 2 , and we denote their gradient cosine similarity at a given step as φ 12 .
When φ 12 is negative, our method is largely equivalent to PCGrad and enjoy PCGrad’s convergence analysis.
Thus, here we consider the other case when φ 12 is positive and show that:
Theorem 1.
Suppose L 1 and L 2 are convex and differentiable, and that the gradient of L is Lipschitz continuous with constant L > 0 .
Then, the GradVac update rule with step size 7 t < 2 L (1+ a 2 ) and t < 1 L , where a = sin( φ 12 − φ T 12 ) sin( φ T 12 ) ( φ 12 > 0 and φ T 12 ≥ φ 12 is some target cosine similarity), will converge to the optimal value L ( θ ∗ ) .
Proof.
Let g 1 = ∇L 1 and g 2 =
∇L 2 be gradients for task 1 and task 2 respectively.
Thus we have g = g 1 + g 2 as the original gradient and g (cid:48)
=
g + a
[Equation]
g [Equation] g 2 as the altered gradient by the GradVac update rule, such that:
[Equation]
where φ T 12 is some pre-set gradient similarity objective and φ T 12 ≥ φ 12 (thus a ≥ 0 since we only consider the angle between two gradients in the range of 0 to π ).
Then, we obtain the quadratic expansion of L as:
[Equation]
and utilize the assumption that ∇L is Lipschitz continuous with constant L, we have:
[Equation]
Thus, we plug in the update rule of GradVac to obtain:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] The last line implies that if we choose learning rate t to be small enough t < 2 L (1+ a 2 ) , we have that t − 1+ a 2 2 Lt 2 > 0 and thus L ( θ + ) <
L ( θ ) (unless the gradient has zero norm).
This tells us applying update rule of GradVac can reach the optimal value
L ( θ ∗ ) since the objective function strictly decreases.
[Table] [Table]
[Table]
F H YPER - PARAMETER S ETTINGS
Here, we show how we choose the best hyper-parameter setting for our method.
As discussed in
Appendix E, there are three hyper-parameter settings for our implementation: (1) which tasks to be considered for GradVac, (2) which layers to measure EMA and perform GradVac, (3) EMA decay rate.
Due to the scale of our model on the larger dataset, we use the smaller scale WMT dataset to ﬁnd the optimal setting and transfer to other experiments.
We do this by grid search using average perplexity on the validation set.
Below, we demonstrate part of our results for each hyper-parameter to choose from.
First, we examine the effect of what tasks to include for GradVac, i.e. G in Algorithm 1.
We consider three options: (1) HRL only: only perform GradVac on high-resource languages, (2) LRL only: only perform GradVac on low-resource languages, (3) all task: perform GradVac on all languages.
Results are shown in Table 6.
We ﬁnd that only conducting GradVac on a subset of languages obtain better performance while it is the best to conduct GradVac on low-resource language only.
This is probably because the effective batch sizes of low-resource languages are usually smaller due to the sampling strategy.
Next, we compare the effect of parameter granularity on model quality.
This corresponds to setting different model components for GradVac ( M in Algorithm 1).
We consider four possibilities, from coarse to ﬁne-grained: (1) whole model: only perform GradVac once on the entire model, (2) enc dec: perform separately for encoder and decoder, (3) all layer: perform individually for each layer in encoder and decoder, (4) all matrix: perform for each parameter matrix in the model.
As shown in Table 7, we ﬁnd that choosing proper parameter granularity is important, as neither too coarse nor too ﬁne-grained perform the best.
This is consistent with our observation made in Section 2.
However, we note that our settings are based on NLP tasks and Transformer networks, and therefore the best overall setting for problems of other domains may vary.
Finally, we study how sensitive our method is on the hyper-parameter β , i.e. the EMA decay rate.
Results in Table 8 illustrate that setting an effective “window” of 100 training steps work best for our problem setups.
This is expected, as setting a larger β value corresponds to conduct GradVac
[Table]
more aggressively, and vice versa.
In general, we ﬁnd our best settings to be consistent across tasks in this paper.
G XTREME E XPERIMENTS G. F INETUNING D ETAILS
We also conduct experiments on the XTREME benchmark (Hu et al., 2020) for cross-lingual transfer tasks.
While other work mostly focus on zero-shot cross-lingual transfer (ﬁnetune on English training data and then evaluate on the target language test data), we use a different setup of multitask learning such that we ﬁnetune multiple languages jointly and evaluate on all languages.
Notice that our goal is not to compare with state-of-the-art results on this benchmark but rather to examine the effectiveness of our proposed method on pre-trained multilingual language models.
We therefore only consider tasks that contain training data for all languages: named entity recognition (NER) and part-of-speech tagging (POS).
The NER task is from the WikiAnn (Pan et al., 2017) dataset, which is built automatically from
Wikipedia.
A linear layer with softmax classiﬁer is added on top of pretrained models to predict the label for each word based on its ﬁrst subword.
We report the F1 score.
Similar to NER, POS is also a sequence labelling task but with a focus on synthetic knowledge.
In particular, the dataset we used is from the Universal Dependencies treebanks (Nivre et al., 2018).
Task-speciﬁc layers are the same as in NER and we report F1.
We select 12 languages for each task randomly.
We use the multilingual BERT (Devlin et al., 2018) as our base model, which is a Transformer model pretrained on the Wikipedias of 104 languages using masked language modelling (MLM).
It contains 12 layers and 178M parameters.
Following Hu et al.
(2020), we ﬁnetune the model forepochs for NER and POS, and search the following hyperparameters: batch size { 16, 32 } ; learning rate { 2e-5, 3e-5, 5e-5 } .
G. POS R ESULT
We evaluate all multi-task baselines on the POS tasks in Table 9.
We ﬁnd that our proposed method outperforms other methods on average, consistent with results in other settings (Section 4).