S EQUENTIAL R EPTILE : I NTER -T ASK G RADIENT A LIGNMENT FOR M ULTILINGUAL L EARNING
A BSTRACT
Multilingual models jointly pretrained on multiple languages have achieved remarkable performance on various multilingual downstream tasks.
Moreover, models ﬁnetuned on a single monolingual downstream task have shown to generalize to unseen languages.
In this paper, we ﬁrst show that it is crucial for those tasks to align gradients between them in order to maximize knowledge transfer while minimizing negative transfer.
Despite its importance, the existing methods for gradient alignment either have a completely different purpose, ignore inter-task alignment, or aim to solve continual learning problems in rather inefﬁcient ways.
As a result of the misaligned gradients between tasks, the model suffers from severe negative transfer in the form of catastrophic forgetting of the knowledge acquired from the pretraining.
To overcome the limitations, we propose a simple yet effective method that can efﬁciently align gradients between tasks.
Speciﬁcally, we perform each inner-optimization by sequentially sampling batches from all the tasks, followed by a Reptile outer update.
Thanks to the gradients aligned between tasks by our method, the model becomes less vulnerable to negative transfer and catastrophic forgetting.
We extensively validate our method on various multi-task learning and zero-shot cross-lingual transfer tasks, where our method largely outperforms all the relevant baselines we consider.
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 & Daum ´ e 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-speciﬁc knowledge not easily transferable across the languages.
Such MTL problem is highly related to the gradient alignment between the tasks, especially when we ﬁnetune 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 ﬁnetune the model with the
MTL objective.
It means that the model gradually starts memorizing task-speciﬁc (or languagespeciﬁc) knowledge not compatible across the languages, which can cause a negative transfer from one language to another.
In case of ﬁnetuning the well-pretrained model, we ﬁnd 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 ﬁnding a good trade-off between minimizing the downstream MTL loss [Figure] and maximizing the cosine similarity between the task gradients, as illustrated in the upper path of
Fig.
In this paper, we aim to solve this problem by developing an MTL method that can efﬁciently align gradients across the tasks while ﬁnetuning the pretrained model.
There has been a seemingly related observation by Yu et al.
that the conﬂict (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 conﬂict to each other.
However, their intuition is completely different from ours.
They manually modify the task gradients whenever the gradient conﬂicts happen, which leads to more aggressive optimization of MTL objective.
In case of ﬁnetuning a well-pretrained model, we ﬁnd 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 sufﬁcient to ﬁnd 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 (cid:96) 2 regularization to AdamW optimizer (Loshchilov & Hutter, 2019).
They argue that the resultant optimizer penalizes (cid:96) 2 distance from the pretrained model during the ﬁnetuning stage.
However, unfortunately we ﬁnd 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 inefﬁcient in that the inner-optimization is done with batch size set to 1, which takes signiﬁcantly 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 ﬁnetune a well-pretrained model, it is sufﬁcient 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.
Speciﬁcally, during the ﬁnetuning 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 efﬁciently 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 efﬁcient manner.
We summarize our contributions as follows.
• We show that when ﬁnetuning 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 efﬁciently 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
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 conﬂict 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 ﬁnding a model that can naturally align the task gradients without such manual modiﬁcations.
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 ﬁnetune 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 ﬁnetuning 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.
A PPROACH
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-speciﬁc data D t .
We learn φ by minimizing the sum of task losses.
[Equation]
where Ω( φ ) is a regularization term and λ ≥ 0 is an associated coefﬁcient.
Reptile We brieﬂy review Reptile (Nichol et al., 2018), an efﬁcient ﬁrst-order meta-learning method suitable for large-scale learning scenario.
We show that Reptile has an approximate learning objective of the form in Eq.
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-speciﬁc optimization from φ .
Speciﬁcally, for each task t = 1 , . .
, T , we perform the optimization by sampling mini-batches
B (1) t , . . .
, B ( K ) t from task data D t and taking gradient steps with them.
[Equation]
for k = 1 , . . .
, K , where α is an inner-learning rate and θ ( k ) t denotes the task-speciﬁc parameter of task t evolved from φ by taking k gradient steps.
After performing K gradient steps for all T tasks, [Figure] we meta-update φ as follows:
[Equation]
where MG t
[Equation]
φ
[Equation]
where η denotes an outer-learning rate.
Nichol et al.
(2018) show that expectation of MG t ( φ ) 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:
[Equation]
where (cid:104)· , ·(cid:105) denotes a dot product.
We can see that E [ MG t ( φ )] approximately minimizes the taskspeciﬁc loss (ﬁrst 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 D t .
3. S EQUENTIAL R EPTILE
In order to consider gradient alignment across tasks as well, we propose to let the inner-learning trajectory consist of mini-batches randomly sampled from all tasks , which we call Sequential Reptile .
Unlike Reptile where we run T task-speciﬁc inner-learning trajectory (in parallel), now we have a single learning trajectory responsible for all T tasks (See Figure 2).
Speciﬁcally, for each inner-step k , we randomly sample a task index t k ∈ { 1 , . . .
, T } and a corresponding mini-batch B ( k ) t
k , and then sequentially update θ ( k ) as follows.
[Equation]
Cat ( p 1 , . . .
, p T ) is a categorical distribution parameterized by p 1 , . .
. , p T , the probability of each task to be selected.
For example, we can let p t ∝ ( N t ) q where N t is the number of training instances for task t and q is some constant.
After K gradient steps with Eq.
5, we update φ as follows
[Equation]
Again, based on Taylor expansion, we have the following approximate form for expectation of the meta-gradient MG ( φ ) over the random sampling of tasks (See derivation in Appendix C).
[Equation]
Note that the critical difference of Eq.
7 from Eq.
4 is that the dot product between the two gradients is computed from the different tasks, t k and t j .
Such inter-task dependency appears as we randomly sample batches from all the tasks sequentially and compose a single learning trajectory with them.
As a result, Eq. 7, or Sequential Reptile promotes the gradient alignments between the tasks, preventing the model from memorizing language speciﬁc knowledge.
It also means that the model can ﬁnd a good trade-off between minimizing the MTL objective and inter-task gradient alignment, thereby effectively preventing catastrophic forgetting of the knowledge acquired from pretraining. [Figure]
E XPERIMENTS
We ﬁrst verify our hypothesis with synthetic experiments.
We then validate our method by solving multi-task learning (MTL) and zero-shot cross-lingual transfer tasks with large-scale real datasets.
Baselines We compare our method against the following relevant baselines.
STL: Single Task Learning (STL) model trained on each single language.
MTL: Base MTL model of which objective is Eq.
1 without the regularizer Ω( φ ) .
GradNorm (Chen et al., 2018):
This model tackles task imbalance problem in MTL.
It prevents the training from being dominated by a single task by adaptively weighting each task gradient.
PCGrad (Yu et al., 2020):
This model aims to optimize MTL objective more aggressively by resolving gradient conﬂicts.
Speciﬁcally, it projects a task gradient onto the other task gradients if the inner product between them is negative.
GradVac (Wang et al., 2021): Similarly to PCGrad, this model alters the task gradients to match the empirical moving average of cosine similarity between the task gradients.
RecAdam (Chen et al., 2020): A model trained with RecAdam optimizer to prevent catastrophic forgetting by penalizing (cid:96) 2 distance from the the pretrained model.
Reptile (Nichol et al., 2018):
A ﬁrst-order meta-learning method suitable for large-scale learning scenario.
Unlike our method, Reptile performs inner-optimization individually for each task.
Sequential Reptile: Our method that can align gradients across the tasks by composing the inner-learning trajectory with all the tasks.
4. S YNTHETIC E XPERIMENTS
We ﬁrst verify our hypothesis with the following synthetic experiments.
We deﬁne three local optima x 1 = (0 , 10) , x 2 = (0 , 0) and x 3 = (10 , 0) in a 2-dimensional space.
Then, we deﬁne three tasks as minimizing each of the following loss functions w.r.t φ ∈ R 2 .
[Equation]
for i = 1 , 2 , 3 .
MTL objective is deﬁned as (cid:80) 3 i =
L i ( φ ) , which we optimize from the initialization (20 , 5) .
Results and analysis Fig.
3 shows the MTL loss surface and the learning trajectory of each method.
We observe that except for Reptile and Sequential Reptile, all the other baselines converge to one of the MTL local minima, failing to ﬁnd a reasonable solution that may generalize better across the tasks.
While Reptile can avoid such a minimum, the resultant solution has very low cosine similarity (See Fig. 3h) because it does not enforce gradient alignments between tasks.
On the other hand, Figure 3h shows that our Sequential Reptile tends to ﬁnd a reasonable trade-off between minimizing MTL loss and maximizing cosine similarity, thanks to the implicit enforcement of inter-task gradient alignment in Eq.
[Table]
[Table]
[Table]
M ULTI -T ASK L EARNING
We next verify our method with large-scale real datasets.
We consider multi-task learning tasks such as multilingual Question Answering (QA) and Named Entity Recognition (NER).
For QA, we use “Gold passage” of TYDI-QA (Clark et al., 2020) dataset where a QA model predicts a start and end position of answer from a paragraph for a given question.
For NER, we use WikiAnn dataset (Pan et al., 2017) where a model classiﬁes each token of a sentence into three classes.
We consider each language as a distinct task and train MTL models.
Implementation Details For all the experiments, we use multilingual BERT (Devlin et al., 2019) base model as a backbone network.
We ﬁntune it with AdamW (Loshchilov & Hutter, 2019) optimizer, setting the inner-learning rate α to 3 · 10 − 5 .
We use batch size 12 for QA and 16 for NER, respectively.
For our method, we set the outer learning rate η to 0 .
1 and the number inner-steps K to 1000.
Following Wang et al.
(2021), for all the baselines, we sample eight tasks proportional to p t ∝ ( N t ) 1 / 5 where N t is the number of training instances for task t .
For our Sequential Reptile, we set the parameter of the categorical distribution in Eq.
5 to the same p t .
Results We compare our method, Sequential Reptile against the baselines on QA and NER tasks.
Table 1 shows the results of QA task.
We see that our method outperforms all the baselines including
STL on most of the languages.
Interestingly, all the other baselines but ours underperform STL on
Arabic language which contains the largest number of training instances ( 2 ∼ 3 times larger than the other languages.
See Appendix B for more information about data statistics).
It implies that the baselines suffer from negative transfer while ours is relatively less vulnerable.
We can observe essentially the same tendency for NER task, which is a highly imbalanced such that the low-resource languages can have around 100 training instance, while the high-resource languages can have around 5,000 to 20,000 examples.
Table 2 shows the results of NER task.
We see that all the baselines but ours highly degrade the performance on high resource languages — de, en, es, hi, mr and tl, which means that they fail to address the imbalance problem properly and suffer from severe negative transfer.
Even GradNorm is not effective in our experiments, which is developed to tackle the task imbalance problem by adaptively scaling task losses. [Figure] [Figure]
Analysis We next analyze the source of performance improvements.
Our hypothesis is that our method can effectively ﬁlter out language speciﬁc knowledge when solving the downstream tasks, which prevents negative transfer and helps retain linguistic knowledge acquired from the pretraining.
[Table]
Firstly
[Table]
we quantitatively measure the cosine similarity between the gradients computed from different tasks in Fig.
4a (QA) and Fig. 4b (NER).
We see from the ﬁgure that our Sequential Reptile shows the highest cosine similarity as expected, due to the approximate learning objective in Eq.
Such high cosine similarity between the task gradients implies that the model has captured the common knowledge well transferable across the languages.
This is in contrast to the other baselines whose task gradients can have even negative cosine similarity with high probability.
Such different gradient directions mean that the current model has memorized the knowledge not quite transferable across the languages, which can cause negative transfer from one language to others.
As a result of such negative transfer, we see from Fig.
5a and Fig.
5b that the baselines suffer from catastrophic forgetting.
In those ﬁgures, high MLM losses on seen (Fig. 5a) and unseen languages (Fig.
5b) mean that the models have forgotten how to reconstruct the masked sentences, which is the original training objective of BERT.
On the other hand, our method shows relatively lower MLM loss, demonstrating its effectiveness in preventing negative transfer and thereby alleviating the catastrophic forgetting.
We further conﬁrm this tendency in Fig.
5c by measuring the (cid:96) 2 distance from the initially pretrained BERT model.
We see that the distance is much shorter for our method than the baselines, which implies that ours can better retain the linguistic knowledge than the baselines.
We can actually test if the ﬁnetuned models have acquired knowledge transferable across the languages by evaluating on some unseen languages without further ﬁnetuning.
In Table 3, we test the ﬁnetuned QA models on ﬁve unseen languages from MLQA dataset (Lewis et al., 2020b).
Again,
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
Sequential Reptile outperforms all the baselines.
The results conﬁrm that the gradient alignment between tasks is key for obtaining common knowledge transferable across the languages.
4. Z ERO - SHOT C ROSS L INGUAL T RANSFER
We next validate our method on zero-shot cross lingual transfer task, motivated from the previous results that our method can learn common knowledge well transferable across the languages.
We train a model only with English annotated data and evaluate the model on target languages without further ﬁnetuning on the target datasets.
To utilize the methods of MTL, we cluster the data into four groups with Gaussian mixture model.
We focus on QA and natural language inference (NLI) task.
For QA, we train the model on SQuAD (Rajpurkar et al., 2016) dataset and evaluate the model on six languages from MLQA dataset (Lewis et al., 2020b).
For NLI, we use MNLI (Williams et al., 2018) dataset as a source training dataset and test the model on fourteen languages from XNLI (Conneau et al., 2018) as a target languages.
Baselines As well as the baselines from the previous experiments, we additionally include
SWEP (Lee et al., 2021).
It learns to perturb word embeddings and uses the perturbed input as extra training data, which is empirically shown to be robust against out-of distribution data.
Implementation Detail
We ﬁnetune multilingual BERT-base model with AdamW. For QA, we use the same hyperparameter in multi-task learning experiment.
For NLI, we use batch size 32 and choose the learning rate 3 · 10 − 5 or 5 · 10 − 5 with AdamW optimizer based on the performance on the validation set.
For our model, we set the outer learning rate η to 0 .
1 and the number inner steps K to 1000.
In order to construct multiple tasks from a single dataset, we cluster concatenation of questions and paragraphs from SQuAD or sentences from MNLI into four groups.
Following
Aharoni & Goldberg (2020), we encode the paragraphs or sentences into hidden representations with pretrained multilingual BERT.
Then, we reduce the dimension of hidden representations with
PCA and run Gaussian mixture model to partition them into disjoint four clusters.
Since the number of training instances are almost evenly distributed for each task, we sample mini-batches from all the four tasks for all the baselines and set p t = 1 / 4 for our model.
Results and
Analysis The results of zero-shot cross lingual transfer show essentially the same tendency as those of multi-task learning in the previous subsection.
Firstly, our model largely outperforms all the baselines for QA and NLI tasks as shown in Table 4 and 5.
To see where the performance improvements come from, we compute the average pairwise cosine similarity between the gradients computed from different tasks in Fig.
4c and 4d.
Again, Sequential Reptile shows much higher similarity than the baselines, which implies that our method can effectively ﬁlter out
[Figure] [Figure]
task-speciﬁc knowledge incompatible across the languages and thereby prevent negative transfer.
As a result, Fig.
6a and 6b show that our method can better retain the linguistic knowledge obtained from the pretraining in terms of relatively lower MLM loss.
Further, Fig.
6c and 6d conﬁrm this tendency with shorter (cid:96) 2 distance from the initially pretrained BERT model.
4. F URTHER A NALYSIS
(a) Trade-off shown in Fig.
In Fig.
7a, MTL loss and cosine similarity decrease as we ﬁnetune the model from top-right to the bottom-left corner.
Meanwhile, Sequential Reptile shows much higher cosine similarities between task gradients at the points of similar MTL losses, achieving better trade-off than Reptile.
It explains why simple early stopping cannot outperform Sequential
Reptile (see Fig.
7c) that directly enforces gradient alignment across tasks.
(b) Effect of the strength of gradient alignment: Then the next question is, how can we further control the trade-off to maximize performance?
According to Eq.
7, we can strengthen or weaken the gradient alignment by increasing or decreasing the inner-learning rate α , respectively.
Fig.
7b shows that while we can control the cosine similarity by varying α as expected, the best-performing α is around 3 · 10 − 5 , which is indeed the most commonly used value for ﬁnetuning the BERT model.
(c) Computational efﬁciency: Lastly, one may suspect training Sequential Reptile takes signiﬁcantly longer wall clock time because inner steps are not parallelizable as in Reptile (See Fig. 2).
This is not true.
Fig.
7c shows that whereas base MTL requires around 40 K gradient computations to achieve 64 EM score, Sequential Reptile requires only around 15 K. As a result, although we run
Sequential Reptile with a single GPU at a time, the wall-clock time becomes even comparable to the base MTL that we run in parallel with 8 GPUs.
Please see wall clock comparison on Appendix D.
CONCLUSION
We showed that when ﬁnetuning a well-pretrained language model, it is important to align gradients between the given set of downstream tasks to prevent negative transfer and retain linguistic knowledge acquired from the pretraining.
We proposed a simple yet effective method aligning gradients between tasks with efﬁciency.
Speciﬁcally, instead of performing multiple inner-optimizations separately for each task, we performed a single inner-optimization by sequentially sampling batches from all the tasks, followed by a Reptile outer update.
We extensively validated the efﬁcacy of our method on various MTL and zero-shot cross-lingual transfer tasks, where ours largely outperformed all the baselines we considered.
Acknowledgements This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.2019-000075, Artiﬁcial Intelligence Graduate School Program(KAIST)), the Engineering Research Center
Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT (NRF-2018R1A5A1059921), Samsung Research Funding Center of Samsung Electronics (No.
SRFC-IT1502-51), Samsung Electronics (IO201214-08145-01), Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT)
(No. 2021-0-02068, Artiﬁcial Intelligence Innovation Hub), the National Research Foundation of Korea (NRF) funded by the Ministry of Education (NRF-2021R1F1A1061655) and KAIST-NAVER Hypercreative AI Center A ALGORITHM
We provide the pseudocode for Sequential Reptile described in section 3.1:
Algorithm
1: Input:
2: while not converged do 3: θ (0) ← φ 4: for k = 1 to k
= K do 5:
Sample a task t
[Table]
∼ Cat ( p 1 , . .
. , p T ) 6:
Sample a mini-batch B [Table] k ) t k from the dataset D t k 7: θ ( k ) ← θ ( k − 1) − α∂ L
( θ ( k − 1) ; B ( k ) t k )
∂θ ( k − 1) 8: end for 9:
MG ( φ ) =
φ − θ ( K ) 10: φ ← φ − η · MG ( φ )
[Table]
: end while
[Table] [Table]
B D ATASET
TYDI-QA (Clark et al., 2020)
It is multilingual question answering (QA) dataset coveringlanguages, where a model retrieves a passage that contains answer to the given question and ﬁnd start and end position of the answer in the passage.
Since we focus on extractive QA tasks, we use “Gold Passage” 1 in which ground truth paragraph containing answer to the question is provided.
Since some of existing tools break due to the lack of white spaces for Thai and Japanese, the creators of the dataset does not provide the Gold Passage of those two languages.
Following the conventional preprocessing for QA, we concatenate a question and paragraph and tokenize it with
BertTokenizer (Devlin et al., 2019) which is implemented in transformers library (Wolf et al., 2020).
We split the tokenized sentences into chunks with overlapping words.
In Table 6, we provide the number of preprocessed training and validation instances for each language.
WikiAhn (Pan et al., 2017)
It a multilingual NER dataset which is automatically constructed from
Wikipedia.
We provide the number of instances for training and validation in Table 7.
XNLI (Conneau et al., 2018)
It is the dataset for multilingual NLI task, which consists oflanguages other than English.
Since it targets for zero-shot cross-lingual transfer, there is no training instances.
It provides 7,500 human annotated instances for each validation and test set.
[Table]
MLQA (Lewis et al., 2020b)
It is the dataset for zero-shot multilingual question answering task.
As XNLI dataset, it only provides only validation and test set for 6 languages other than English.
In
Table 8, we provide data statistics borrowed from the original paper (Lewis et al., 2020b)
Common Crawl 100 (Conneau et al., 2020;
Wenzek et al., 2020)
It is multilingual corpus consisting of more than 100 language that is used for pretraining XLM-R (Conneau et al., 2020) model.
We download the preprocessed corpus 2 provided by Wenzek et al.
We sample 5,000 instances for each language and evaluate masked language modeling loss.
C I MPLICIT G RADIENT ALIGNMENT OF SEQUENTIAL R EPTILE
In this section, we provide derivation of implicit gradient alignment of Sequential Reptile in the equation 7.
Firstly, we deﬁne the following terms from Nichol et al.
[Equation] [Equation]
( gradient obtained during SGD with mini-batch
[Equation]
[Equation]
( sequence of parameter vectors ) (10)
[Equation]
gradient at [Equation] point
[Equation]
mini-batch
[Equation] [Equation]
where α denotes a learning rate for inner optimization.
t k ∈ { 1 , . . .
, T } is a task index and B ( k ) t k is a mini-batch sampled from the task t k .
Corollary.
After K steps of inner-updates, expectation of (cid:80) Kk =1 g t k over the random sampling of tasks approximates the following:
[Equation]
where (cid:104)· , ·(cid:105) denotes a dot-product in R d .
Proof.
We borrow the key idea from the theorem of Reptile (Nichol et al., 2018) that task speciﬁc inner-optimization of Reptile implicitly maximizes inner products between mini-batch gradients within a task.
First, we approximate the gradient g t k as follows.
[Equation]
[Figure]
After K gradient steps, we approximate the expectation of the meta-gradient MG ( φ ) over the task t 1 , . .
. , t K , where t k ∼ Cat ( p 1 , . . .
, p T ) for k = 1 , . . .
, K .
E [ MG ( φ )] =
E (cid:34) K (cid:88) k
= g t k (cid:35) ≈
E  
K (cid:88) k
= g t k − α
K (cid:88) k
= k − 1 (cid:88)
j =
H t k g t j   =
E (cid:34) K (cid:88) k
= g t k (cid:35) − E  
α 
K (cid:88) k
= k − 1 (cid:88)
j =
H t k g t j
+
K (cid:88) k
= k − 1 (cid:88)
j =
H t j g t k   
 = E (cid:34) K (cid:88) k
= g t k (cid:35) − E  
α
K (cid:88) k
= k − 1 (cid:88)
j = ∂ (cid:10) g t k , g t j (cid:11)
∂φ 
 =
E 
[Equation]
Similarly, Smith et al.
(2021) show implicit regularization of SGD.
After one epoch of SGD, it implicitly biases a model to minimize difference between gradient of each example and full batch gradient.
As a result, it approximately aligns gradient of each instance with the full-batch gradient
[Equation]
D F URTHER A NALYSIS
Computational efﬁciency In Figure 8a, we plot test Exact Match score as a function of wall clock time.
Although training Sequential Reptile is not parallelizable, it shows tolerable computational efﬁciency compared to MTL model.
For MTL, it takes about 1 hour and 15 minutes while Sequential
Reptile takes 1 hour and 45 minutes to reach 64 EM score.
Inner steps As shown in Figure 8b, Sequential Reptile is robust to the number of steps for the inner optimization.
It shows consistent performance with little variance.
[Table] [Table]
E A DDITIONAL E XPERIMENTS
In order to show that our model Sequential Reptile is generally applicable to various multi-task learning problems, we additionally perform experiments on monolingual text classiﬁcation and image classiﬁcation.
Text classiﬁcation Following (Pilault et al., 2021), we train BERT base model on 7 text classiﬁcation tasks — CoLA, MNLI, MRPC, QNLI,QQP, RTE, SST2 and one text similarity score regression — STSB from GLUE dataset (Wang et al., 2019a).
We share the BERT encoder across all the task and add linear layer on top of the encoder for each task.
For Reptile, we use learned initialization with each task speciﬁc head for prediction at test time.
For STSB task, we use Pearson correlation coefﬁcient to measure the performance of the baselines and ours.
For the other 7 task, we evaluate all the models with accuracy.
As shown in Table 9, Sequential Reptile outperforms the Reptile and
MTL on 4 tasks – CoLA, MRPC, RTE, SST2 and shows comparable performance on the other tasks.
Image classiﬁcation Following the experimental setup from Shin et al.
(2021), we use 7 datasets — Tiny-ImageNet (Le & Yang, 2015), CIFAR100 (Krizhevsky et al., 2009), Stanford Dogs (Khosla et al., 2011), Aircraft (Maji et al., 2013), CUB (Wah et al., 2011), Fashion-MNIST (Xiao et al., 2017), and SVHN (Netzer et al., 2011).
We class-wisely divide Tiny-ImageNet into two splits which are denoted as TIN-1 and TIN-2, respectively, and consider each split as a distinct task, which results in total 8 tasks for multi-task learning image classiﬁcation.
We ﬁnetune ResNet
(He et al., 2016) which is pretrained on ImageNet (Deng et al., 2009) with a randomly initialized linear classiﬁer for each task.
For all the models, the pretrained ResNet is shared across tasks.
As the previous experiments, we use shared initialization of Reptile with task speciﬁc linear classiﬁers.
We measure accuracy of each image classiﬁcation task and report the average score.
As shown in Table 10, Sequential Reptile outperforms MTL and Reptile with large margin other than Fashion MNIST (F-MNIST).
We observe that Reptile fails on image classiﬁcation, which is contrast to text classiﬁcation tasks where it shows reasonable accuracy compared to MTL model.
We conjecture that Reptile needs adaptation to the target task since the set of image datasets are more heterogeneous than the set of text datasets.
F V ISUALIZATION OF LEARNING TRAJECTORY
In Figure 9, 10, and 11, we visualize learning trajectory of the MTL models and cosine similarity between task gradients with three different initialization as described in section 4.1.
The similar pattern holds for all different initialization.
All the MTL baselines except Reptile fall into one of local optima.
On the other hand, Sequential Reptile avoid such local minima while maximizing cosine similarities of task gradients.
[Figure]
[Figure] [Figure]