Sequential Reptile: Inter-Task Gradient Alignment for Multilingual Learning

1 INTRODUCTION . Multilingual language models ( Devlin et al. , 2019 ; Conneau & Lample , 2019 ; Conneau et al. , 2020 ; Liu et al. , 2020 ; Lewis et al. , 2020a ; Xue et al. , 2021 ) have achieved impressive performance on a variety of multilingual natural language processing ( NLP ) tasks . Training a model with multiple languages jointly can be understood as a multi-task learning ( MTL ) problem where each language serves as a distinct task to be learned ( Wang et al. , 2021 ) . The goal of MTL is to make use of relatedness between tasks to improve generalization without negative transfer ( Kang et al. , 2011 ; Kumar & Daumé III , 2012 ; Lee et al. , 2016 ; Wang et al. , 2019b ; 2020b ) . Likewise , when we train with a downstream multilingual MTL objective , we need to maximize knowledge transfer between the languages while minimizing negative transfer between them . This is achieved by developing an effective MTL strategy that can prevent the model from memorizing task-specific knowledge not easily transferable across the languages . Such MTL problem is highly related to the gradient alignment between the tasks , especially when we finetune a well-pretrained model like multilingual BERT ( Devlin et al. , 2019 ) . We see from the bottom path of Fig . 1 that the cosine similarity between the task gradients ( gradients of MTL losses individually computed for each task ) tend to gradually decrease as we finetune the model with the MTL objective . It means that the model gradually starts memorizing task-specific ( or languagespecific ) knowledge not compatible across the languages , which can cause a negative transfer from one language to another . In case of finetuning the well-pretrained model , we find that it causes catastrophic forgetting of the pretrained knowledge . Since the pretrained model is the fundamental knowledge shared across all NLP tasks , such catastrophic forgetting can severely degrade the performance of all tasks . Therefore , we want our model to maximally retain the knowledge of the pretrained model by finding a good trade-off between minimizing the downstream MTL loss ∗Equal contribution and maximizing the cosine similarity between the task gradients , as illustrated in the upper path of Fig . 1 . In this paper , we aim to solve this problem by developing an MTL method that can efficiently align gradients across the tasks while finetuning the pretrained model . There has been a seemingly related observation by Yu et al . ( 2020 ) that the conflict ( negative cosine similarity ) between task gradients makes it hard to optimize the MTL objective . They propose to manually alter the direction of task gradients whenever the task gradients conflict to each other . However , their intuition is completely different from ours . They manually modify the task gradients whenever the gradient conflicts happen , which leads to more aggressive optimization of MTL objective . In case of finetuning a well-pretrained model , we find that it simply leads to catastrophic forgetting . Instead , we aim to make the model converge to a point where task gradients are naturally aligned , leading to less aggressive optimization of the MTL objective ( See the upper path of Fig . 1 ) . Then a natural question is if we can alleviate the catastrophic forgetting with early stopping . Our observation is that whereas early stopping can slightly increase cosine similarity to some extent , it is not sufficient to find a good trade-off between minimizing MTL objective and maximizing cosine similarity to improve generalization ( See Fig . 1 ) . It means that we may need either an implicit or explicit objective for gradient alignment between tasks . Also , Chen et al . ( 2020 ) recently argue that we can mitigate catastrophic forgetting by adding ` 2 regularization to AdamW optimizer ( Loshchilov & Hutter , 2019 ) . They argue that the resultant optimizer penalizes ` 2 distance from the pretrained model during the finetuning stage . However , unfortunately we find that their method is not much effective in preventing catastrophic forgetting in the experimental setups we consider . On the other hand , Reptile ( Nichol et al. , 2018 ) implicitly promotes gradient alignment between mini-batches within a task . Reptile updates a shared initial parameter individually for each task , such that the task gradients are not necessarily aligned across the tasks . In continual learning area , MER ( Riemer et al. , 2019 ) and La-MAML ( Gupta et al. , 2020 ) propose to align the gradients between sequentially incoming tasks in order to maximally share the progress on their objectives . However , as they focus on continual learning problems , they require explicit memory buffers to store previous task examples and align gradients with them , which is complicated and costly . Further , their methods are rather inefficient in that the inner-optimization is done with batch size set to 1 , which takes significantly more time than usual batch-wise training . Therefore , their methods are not straightforwardly applicable to multilingual MTL problems we aim to solve in this paper . In this paper , we show that when we finetune a well-pretrained model , it is sufficient to align gradients between the currently given downstream tasks in order to retain the pretrained knowledge , without accessing the data used for pretraining or memory buffers . Specifically , during the finetuning stage , we sequentially sample mini-batches from all the downstream tasks at hand to perform a single inner-optimization , followed by a Reptile outer update . Then , we can efficiently align the gradients between tasks based on the implicit dependencies between the inner-update steps . This procedure , which we call Sequential Reptile , is a simple yet effective method that can largely improve the performance of various downstream multilingual tasks by preventing negative transfer and catastrophic forgetting in an efficient manner . We summarize our contributions as follows . • We show that when finetuning a well-pretrained model , gradients not aligned between tasks can cause negative transfer and catastrophic forgetting of the knowledge acquired from the pretraining . • To solve the problem , we propose Sequential Reptile , a simple yet effective MTL method that efficiently aligns gradients between tasks , thus prevents negative transfer and catastrophic forgetting . • We extensively validate our method on various MTL and zero-shot cross-lingual transfer tasks , including question answering , named entity recognition and natural language inference tasks , in which our method largely outperforms all the baselines . 2 RELATED WORKS . Multi-task Learning The goal of MTL is to leverage relatedness between tasks for effective knowledge transfer while preventing negative interference between them ( Zhang & Yeung , 2010 ; Kang et al. , 2011 ; Lee et al. , 2016 ) . GradNorm ( Chen et al. , 2018 ) tackles task imbalance problem by adaptively weighting each task loss . Another line of literature propose to search Pareto optimal solutions which represent trade-offs between the tasks ( Sener & Koltun , 2018 ; Lin et al. , 2019 ) . Recently , Yu et al . ( 2020 ) and Wang et al . ( 2021 ) propose to manually resolve the conflict between task gradients to more aggressively optimize the given MTL objective . However , their goal is completely different from ours because here we focus on preventing negative transfer by finding a model that can naturally align the task gradients without such manual modifications . Multilingual Language Model Training a multilingual language model is a typical example of multi-task learning . Most of the previous works focus on jointly pretraining a model with hundreds of languages to transfer common knowledge between the languages ( Devlin et al. , 2019 ; Conneau & Lample , 2019 ; Conneau et al. , 2020 ; Liu et al. , 2020 ; Lewis et al. , 2020a ; Xue et al. , 2021 ) . Some literature show the limitation of jointly training the model with multilingual corpora ( Arivazhagan et al. , 2019 ; Wang et al. , 2020b ) . Several follow-up works propose to tackle the various accompanying problems such as post-hoc alignment ( Wang et al. , 2019c ; Cao et al. , 2019 ) , data balancing ( Wang et al. , 2020a ) and loss curvature-aware optimization to improve the performance of low resource languages ( Li & Gong , 2021 ) . In this paper , we focus on how to finetune a well pretrained multilingual language model by preventing catastrophic forgetting of the pretrained knowledge . Zero-shot Cross Lingual Transfer Zero-shot cross-lingual transfer is to train a model with monolingual labeled data and evaluate it on some unseen target languages without further finetuning the model on the target languages . Nooralahzadeh et al . ( 2020 ) utilize meta-learning to learn how to transfer knowledge from high resource languages to low resource ones . Hu et al . ( 2021 ) and Pan et al . ( 2021 ) leverage a set of paired sentences from different languages to train the model and minimize the distance between the representation of the paired sentences . Instead , we partition the monolingual labeled data into groups and consider them as a set of tasks for multi-task learning . 3 APPROACH . The goal of multi-task learning ( MTL ) is to estimate a model parameter φ that can achieve good performance across all the given T tasks , where each task t = 1 , . . . , T has task-specific data Dt . We learn φ by minimizing the sum of task losses . min φ T∑ t=1 L ( φ ; Dt ) + λΩ ( φ ) ( 1 ) where Ω ( φ ) is a regularization term and λ ≥ 0 is an associated coefficient . Reptile We briefly review Reptile ( Nichol et al. , 2018 ) , an efficient first-order meta-learning method suitable for large-scale learning scenario . We show that Reptile has an approximate learning objective of the form in Eq . 1 . Although Reptile is originally designed for learning a shared initialization , we can use the initialization φ for actual predictions without any adaptation ( Riemer et al. , 2019 ) . Firstly , given a set of T tasks , we individually perform the task-specific optimization from φ . Specifically , for each task t = 1 , . . . , T , we perform the optimization by sampling mini-batches B ( 1 ) t , . . . , B ( K ) t from task data Dt and taking gradient steps with them . θ ( 0 ) t = φ , θ ( k ) t = θ ( k−1 ) t − α ∂L ( θ ( k−1 ) t ; B ( k ) t ) ∂θ ( k−1 ) t ( 2 ) for k = 1 , . . . , K , where α is an inner-learning rate and θ ( k ) t denotes the task-specific parameter of task t evolved from φ by taking k gradient steps . After performing K gradient steps for all T tasks , we meta-update φ as follows : φ← φ− η · 1 T T∑ t=1 MGt ( φ ) , where MGt ( φ ) = φ− θ ( K ) t ( 3 ) where η denotes an outer-learning rate . Nichol et al . ( 2018 ) show that expectation of MGt ( φ ) over the random sampling of batches , which is the meta-gradient of task t evaluated at φ , can be approximated as follows based on Taylor expansion : E [ MGt ( φ ) ] ≈ ∂ ∂φ E K∑ k=1 L ( φ ; B ( k ) t ) − α 2 K∑ k=1 k−1∑ j=1 〈 ∂L ( φ ; B ( k ) t ) ∂φ , ∂L ( φ ; B ( j ) t ) ∂φ 〉 ( 4 ) where 〈· , ·〉 denotes a dot product . We can see that E [ MGt ( φ ) ] approximately minimizes the taskspecific loss ( first term ) and maximizes the inner-products between gradients computed with different batches ( second term ) . The inner-learning rate α controls the trade-off between them . However , the critical limitation is that the dot product does not consider aligning gradients computed from different tasks . This is because each inner-learning trajectory consists of the batches B ( 1 ) t , . . . , B ( K ) t sampled from the same task data Dt .

Summary: The paper proposed an optimization method for training multi-task learning models by increasing the task alignments. Their key assumption is that standard MTL can result in catastrophic forgetting of pretrained knowledge and lead to local task optimum in the finetuning stages. To avoid this, the paper proposed a method to jointly optimize the task losses and gradient alignments between tasks. Their key contributions are: - The paper empirically examined the question of whether worse alignments of task gradients can lead to negative transfer, though extensive experiments on multilingual tasks. - Based on such observation, the paper proposes a simple and effective optimization method for MTL that jointly optimizes the task losses and gradient alignments between tasks.