Order Matters in the Presence of Dataset Imbalance for Multilingual Learning
Authors: Dami Choi
Derrick Xin
Hamid Dadkhahi
Justin Gilmer
Ankush Garg
Orhan Firat
Chih-Kuan Yeh
Andrew M. Dai
Behrooz Ghorbani
Venue: NeurIPS 2023 poster (2023)
Keywords: Multitask Optimization
Multilingual
Pre-training
Language Models
Language Sampling
Low Resource Languages
Overfitting
TL;DR: We present a simple method for multitask optimization in the presence of data imbalance, and verify its efficacy via thorough empirical evaluations.

Abstract: In this paper, we empirically study the optimization dynamics of multi-task learning, particularly focusing on those that govern a collection of tasks with significant data imbalance. We present a simple yet effective method of pre-training on high-resource tasks, followed by fine-tuning on a mixture of high/low-resource tasks. We provide a thorough empirical study and analysis of this method's benefits showing that it achieves consistent improvements relative to the performance trade-off profile of standard static weighting. We analyze under what data regimes this method is applicable and show its improvements empirically in neural machine translation (NMT) and multi-lingual language modeling.

Order Matters in the Presence of Dataset Imbalance for Multilingual Learning
Dami Choi
Derrick Xin" Google Research dxin@google. com
U. Toronto & Vector Institute choi dami Acs. toronto..oeoooo edu
Hamid Dadkhahi
Justin Gilmer Google Deepmind gi 1mer@google. com
Ankush Garg Google Deepmind ankugargōgoogle.com com com
Google Research
hdadkhahi com
Orhan Firat Google Deepmind rchanf@google.cco com
Chih-Kuan Yeh Google Deepmind ch chibkuanyeh@google..c com
Andrew M. Dai Google Deepmind adai cdai@google.com
Behrooz Ghorbani OpenAl ghorbani@openai.comm com
Abstract
In this paper, we empirically study the optimization dynamics of multi-task learning, particularly focusing on those that govern collection of tasks with significant data imbalance. We present a simple yet effective method pre-training on high- resource tasks, followed by fine-tuning on a mixture of high/low We provide a thorough empirical study and analysis of this method's benefits showing that it achieves consistent improvements relative to the performance trade- off profile of standard static weighting. We analyze under what data regimes this method is applicable and show its improvements empirically in neural machine translation (NMT) and multi- Minlulllinull language modeling.
tasks.
Introduction
Over the past few years, large multi-task neural networks have emerged as a popular modeling paradigm in deep learning. The appeal behind these models is that they can leverage transfer learning among the tasks to outperform single-task models. Indeed, multi-task models have achieved state-of- language understanding [24 32
the-art performance in domains such machine translation 2 g and speech recognition 3.
Unfortunately, optimizing such multi-task models remains challenge. To effectively train these models, the different tasks need to be balanced during training. This often done by sampling each task with static probability.
Prior work [3] 20 shows evidence that when all tasks are in the data rich regime (high-resource), such static sampling approaches yield optimal results. However, when certain tasks are data sparse
'Equal contribution Work done as student researcher Google.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

flow-resource + which is quite common in real-world applications, the optimality of static sampling unclear.
The problem with static sampling in the presence of low-resource tasks is that it has difficulty dealing with overfitting on the low-resource tasks. This is because early stopping is not viable solution due to high-resouree tasks needing many more epochs to converge. The transfer learning scheme of pre-training on high-resource and fine-tuning on low-resource tasks (such as (33) provides solution to the overfitting problem, since the training of high and low-resource tasks are separated. Not only this, but the training of low-resource tasks can potentially benefit from positive transfer that comes from performing well on the high-resource tasks. The problem with this approach, however, is that during the fine-tuning phase, catastrophic forgetting of the pre-training tasks ensues.
In this paper, we introduce a simple training scheme that combines the best of static sampling and transfer learning: pre-train resource tasks. A pre-training and fine-tuning scheme effectively enables early stopping by allowing the training of low-resource tasks to happen for little as needed to prevent overfitting, while training the high-resource task for as long as needed. Furthermore, pre-training on a high-resource task will potentially enable positive transfer for low-resource tasks and result in faster convergence in the fine-tuning phase. Lastly, the fine-tuning phase on a mixture of high and low-resource tasks will not only remedy the catastrophic forgetting issue of fine-tuning only on low-resource tasks, but also enjoy further transfer learning among all the tasks.
high-resource task and fine-tune jointly on a mixture of high and low-
Through extensive empirical study, we find that the pre-training and joint fine-tuning scheme yields superior low-resource task performance compared to both static sampling and the transfer- learning scheme. We observed that the performance improvement on static sampling is driven by two mechanisms. The first is that pre-training initializes the fine-tuning phase at better starting point than random initialization due to positive transfer. The second is that higher sampling rates are more data-efficient than lower sampling rates. Because our method has two separate training phases, the low-resource-trainine training phase can be short. This in turn enables us to increase the low-resource sampling rate without risking overfitting. Indeed, our method is more data- efft-eefficent than static sampling in terms of the low-resource tasks throughout the entire fine-tuning phase, achieving better low-resource task performance while using only a fraction of the data seen by static sampling. We further observe that pre-training and joint fine-tuning seems to have a regularization effect. However, we find that regularization is not the main factor behind the performance improvement, since increased explicit regularization, such as dropout, does not improve the performance to the extent that our method does.
The contributions of this paper can be summarized as follows:
To the best of our knowledge, we are the first to show that it is possible to push the Pareto front of static sampling in the data-imbalanced regime.
We present a simple algorithm that can be readily used to boost low-resource tasks' perfor- mance in multilingual models.
We show on realistic workloads (up to 13B parameters) that scheme performs bet- ter than static sampling and transfer learning with respect to the low-resource language- pair/language.
2 Background
In our work, we focus on the supervised setting, where our model parameters 0 E R2 are trained on K different tasks, with the loss for task being (0).
We introduce the idea of Pareto optimality to better explain the trade-off effect that happens when training on many different tasks.
Definition (Pareto Optimality). E Rp Pareto dominates another 0' if VI S K, ()) < (0)))) and there exists task where (0) < (0'). 0 is Pareto optimal if it is not dominated by any other point. The collection of the Pareto optimal points is denoted as the Pareto front.
In this literature, data rich and data sparse tasks are often referred to as high-resource and low- respectively. Note that whether a task is high- righures or not depends on both the amount of training data and the model capacity.
2

---

A standard approach for optimizing multi-task models is scalarization S or static sampling:
(ww) - arg min Σ wi (0),
(I)
where is fixed vector of pre-determined task weights with 0 and Σ, w;= =
In our work, we follow convention and implement scalarization via proportional sampling, where data from task i is sampled with probability equal to wy In this case, the expected loss is equal to the loss from scalarization:
K
() = 1 [e(x;; 0)
Σ P(takk Exp ((xx0))
Σ wn (0).
Prior work [31] studied the performance trade-off behavior of scalarization and a variety of different multi-task optimization (MTO) methods in the two-task setting. They found that both in the high- resource case and in the data-imbalanced case, no MTO method improved upon the Pareto front of scalarization. In our work, we compare the performance trade-off behavior of scalarization and our proposed method, and find that the Pareto front of scalarization can be improved in the data-imbalanced regime.
Note that practically speaking, it is not feasible to determine whether truly Pareto optimal since we must check that it is not dominated by all 0' E IRP. Following [31], instead of considering all of IPP we consider only the parameters reachable by fixed set of hyperparameters.
3 Pre-training Joint Fine-tuning
Given K tasks, among which some low-resource, our goal is to optimize the performance of the low-resource tasks without sacrificing the performance of the remaining tasks. Static sampling is not ideal because all tasks seen constantly throughout the entirety of training, resulting in overfitting of low-resource tasks while high-resource tasks still need to be learned. Naively breaking up training into two phases and training low-resource tasks in the later phase results in catastrophic forgetting of earlier-trained tasks.
Assuming the existence of at least one high-resource task, we propose to first pre-train on high- resource task, and fine-tune the resulting model on the full mixture of K tasks. We call this method pre-training joint fine-tuning
In our preliminary experiments. we found that is important reset the learning rate schedule and optimizer state when switching over to the joint fine-tuning phase. This is because learning is extremely slow for tasks that are newly introduced when the learning rate has already decayed. In our evaluations, we additionally experiment with adding resetting to the scalarization baseline to ensure that improvements from our method are not purely from resetting. See Sections 4.1.2 and 4.2 for more detail.
Our two-stage training process introduces additional hyperparameters compared scalarization: the hyperparameters involved in the pre-training phase, and the length of the pre-training phase. However, we find that tuning is not much more difficult than scalarization, and in some cases easier to tune. The pre- pre-training phase only involves tuning for single task, which is much easier than tuning for multiple tasks. We also expect the joint fine-tuning phase to be shorter than the full training length of scalarization; therefore, tuning for the second phase should be around the same or easier than scalarization. Lastly, our results show that pre- treitrain does not hurt fine-tuning performance and longer pre- pre-training translates to better fine-tuning. From this, we recommend that if there is strict training budget, it better to be conservative and pre-trale for a shorter amount of time. However, if the goal is to obtain the best performance and there is no strict compute budget, we recommend pre-training for as long as possible before fine-tuning. See Section 4.3 for more details.
'e use the terms 'pre-trae pre-training' and fine-tuning" only to distinguish the two phases of training, and that the training objectives are the same for both phases. In other words, we do not suggest using any particular self supervised objective for the pre- training phase, or training on downstream tasks for the fine-tuning phase.
3

---

4 Experiments
In the following sections, we apply our proposed training scheme to NMT (where each task is a language-pair and multilingual training (where each task is language). In the NMT experiments, we show that pre-training joint fine-tuning pushes past the trade-off frontier of scalarization through significant improvements optimization methods were not able to achieve [31. In the language modeling experiments, we scale up the number of tasks, and show that our method retains the same benefits for the low-resource languages.
the low-resource task- feat that many popular gradient- multi-task
4.1 Neural Machine Translation
For our first experiment, we focus on a setting where we can trace out, and compare the trade-off frontiers obtained with and without pre- vre-training. As in prior work [I], we choose to work on the two-task setting due to the case of visualizing the performance trade-off curves.
We choose our high and low-resource language-pains from the WMT dataset, where frttish--Chhhneeee French) are the high-resource language pairs, and English (Rooaani,,, Hindi] are the low-resource language pairs. See Table [ for details on each language- section use pre-LayerNorm encoder-decoder transformer architecture [28). In the main paper, we present results on models with three encoder layers and three decoder layers. Results obtained with a larger model size are in Appendix A.2 Further details, including hyperparameters, are A.1
All models in this
In order trace out the trade-off frontiers for the pre-training joint fine-tuning method and the scalarization baseline, we adhere the follow- ing methodology. For scalarization, we iterate through grid of task weights (since there are only two tasks, grid is linear function of the granularity) and train on the two language pairs for N steps using proportional sampling accord- ing to the task weights. For the pre-training joint fine-tuning method, we first pre-train the high-resource language pair for N1 training steps. We then reset the optimizer state and the learning rate schedule and fine-tune mixture of high-resource and low-resource language pairs for N2 training steps such that N1 + N2 = N. For the fine-tuning phase, we iterate through grid of task weights as with scalarization. The grid of sampling rates will trace a performance trade-off front, which can be used to compare our method and scalarization.
Table 1: Overview of data sources used our NMT experiments. Our datasets are from WMT.
Language Pair
# Train Ex.
Eval Ex.
En-Fr '15
40, 853, 298
4,503
En-Zh '19
25, 986, 436
3, 3,981
En-Ro '16
610,320
1,999
En-Hi '14
313, 748
520
Lastly, we train a restart baseline in order to ablate the possibility that any improvements coming from pre-training joint fine-tuning are due to the resetting of optimizer state and learning rate schedules before fine-tuning. The restart baseline takes the model obtained via scalarization trained for Ni steps, resets optimizer states and the learning rate schedule, and continues to train it with the same sampling rate as in scalarization.
4.1.1 High-Resource and High- Righ-Resource:
We first start by highlighting that pre-training joint fine-tuning does not show benefits if all tasks are high-resource. Figure [ shows that in the English- (Chinese, French) translation tasks, the performance on each of the language-pairs are bounded by the amount of data seen from that pair. In other words, pre-training on En-FFr cannot act as proxy for En-tZh data, because if it could, the front would be improved. At the same time, pre-training does not negatively impact En-+Zh training. Figures 2T and 22 show that pre-training does not affect the learning efficiency for En-*Zh (slope of the curves are similar to one another), and also did not result in a worse initialization for En En->Zh.
4.1.2 High-Resource and Low-Resource
In the data-imbalance setting of English [Romanian, French), we pre-train for 400k steps and fine-tune for 50k steps to emphasize the computational benefits of pre-training fine-tuning. Although single full run of scalarization (N steps) and pre-training fine-tuning (N₁ + N2 N) take the same

---

Generalizatiinn Behavior Overview En (2h, Fr)
Optimization Behavior Overview En (2h, Fr)
En-(Zh.fr) (300k)
1:55
En-(22hFrr) (300K)
Valid Cross-Enrapy Loss (En FF
1.6
En-(Zh,fr) (600k)
1.50
En-(2h,Fr) (600k)
(300k) En En-(ZZ.Fr) (200k)
(300k) (x-(2h.F)) (300K)
En EniZh Rate
1.40
0.2
1.35
0.4
1.30
1.25
0.7
1.20
0.9
1.2
1.15
3.0
1.7
Velid Cross-Entrepy Loss (En 2h)
Train Cross- Cross-Entropyy Loss IEn 2h)
The trade-off front from pre-training does not improve upon the trade-off front from fully static sampling when all tasks are high-resource, The performance on each of the high-resource tasks are bounded by the amount of data seen for that task. We can also observe interference between the two tasks from how all 9 different sampling rates form the trade-off frontier. These observations hold for both testing (left) and training (right).
Figure 1:
Genaralization Bahavler Overview
Optimization Behavior
1.36
En- (Ro,Pr) 1400k)
17
En- (Ro.Fr) (400k)
Valid Cross-Entopp Loss (En -En
1.34
(Ro.Fr) (450k)
Trein CrossEEtropy Loss (En-Fr
(450K)
En-fr (400k) En- (r), (50k)
En-fr (400k) En- En-(Ro.Fr) (50k)
1.32
Best Scalartzation En Loss
0.1
0.2
1.30
0.4
1.28
0.6
1.26
2.3
1.24
0.9
1,75
1.00
1.05
3.90
1.95
2,00
2.05
Valid Cross-Enterr Enoos-Entropy..ooo Loss (En Ro)
Train Cross-Entropy Loss (En Ro)
Figure 2: (Left:) In the data-imbalanced case, the trade-off front from pre-training yields better low-resource task performance than the trade-off front of scalarization. The poor performance the restart baseline shows that the resetting of states is not why pre-training and fine-tuning performs well. Note that trade-off fronts consist of only a subset of the sampling ratios due to overfitting. which is different from the fully high-resource setting. Right: Pre-training results in noticeably worse performance on the training set, hinting that pre-training has a regularization effect on the low-resource task.
amount of compute, pre-training joint fine-tuning makes hyperparamter tuning much more efficient, since 1) tuning for pre-training is on a single task and therefore, easier to tune, and 2) tuning for fine-tuning is faster since N₂ < N.
In Figure we we can observe that pre-training joint fine-tuning able to achieve performance trade- off points that go beyond what is achievable via scalarization. Pre-training on high-resource language pair creates non-dominated points by yielding significantly better performance in the low-resource task (En~>Ro) without completely sacrificing performance in the high-resource task (En->Fr). Additionally, able to do this while seeing less overall Romanian tokens according to Figure
We see similar results for En (Hi, Fr), shown Figure 12 in the Appendix. This is surprising result since French and Hindi are less linguistically similar than French and Romanian. Finally, we can see from the sub-optimal performance the restart baseline in Figures and that the act of resetting is not the reason behind the success of the pre-training joint fine-tuning scheme. We provide BLEU score evaluations for En->(Ro, Fr) and En-(HHi Fr) Appendix A.5 validating that the improvements in loss translate to downstream metrics.
5

---

Valid Cross Entroyy Lass (En -Ba
En-(Nu,F) (409k)
Fn-(Ro..F1) (400K)
En-(Roffr) (450k)
En En-(Ro (1) (450K)
EniF (400K) En-(Ro.F)
(400K) En (Ro.F).))
Valid Lacs
10:
10"
102
10a
10
10a
102
10c
10'
10.
Num: Examples Seen (En Ro)
Num.
Examples Seen (En Fr)
Figure Pre-training joint fine-tuning has both better initialization and data-efficiency than scalarization. Each line corresponds to the datapoint that achieved the best En-*Ro validation loss in Figure among the different run groups
Volid Cross-Ente Ens--Entryyy Loss iEn flo)
Train Crozs Entropy Lass [En RD)
6
10
4
3
10a2
10
10
10
10
10
107
10
10
30"
10
10
10
10
107
10t7
Num. Examples (En- Ro)
Num. Examples Seen (En Rol
Figure 4: Each curve corresponds to a single scalarization trial with particular (static) sampling rate for En Rn. The rate which the training loss decreases is slower for lower En-Roo sampling rates than for higher sampling rates. At higher sampling rates, overfitting starts to happen.
4.1.3 Analysis
The performance improvement of pre-training joint fine-tuning stems from two main mechanisms.
Pre-training utilizes positive transfer between tasks, and initializes the fine-tuning phase at better starting point than random initialization. Figure 3 shows this effect for the En- Fr} translation tasks.
Higher sampling rates are more data-efficient than lower sampling rates. Figure 4 shows how optimization (training set performance) gets more and more data-efficient as the sampling rate increases. However, on the generalization side, increasing the sampling rate works only up until certain point, where overfitting kicks in.
By design, pre-training joint fine-tuning has two separate training phases which allows the low- shource-training phase to be short. This in turn enables us to increase the low-resource sampling rate, resulting in faster training. This effect can seen in Figure where the En- sampling rates that resulted in the best En->Ro performance was 0.4, while for pre-training joint fine-tuning, best rate is 0.5. Figure confirms that indeed after pre-training, fine-tuning on En-RRo is more data-efficient than not pre-training.
Joint fine-tuning is also an important piece in addition the two-stage setup. Only fine-tuning on the low-resource task, which is the classic transfer learning scheme, results in overfitting and catastrophic forgetting of the pre-training task shown in Figure 6
Lastly, Figure shows that pre-training joint fine-tuning yields worse training set performance, and therefore, could be seen having regularization effect. We show in Figure 5 that regularization by itself does not explain the superior performance of our scheme.
The results seen so far show that data order matters when training in the presence of a low-resource task, since seeing high-resource data first before seeing low-resource data later pushes the pareto front of seeing both types of data at the same time.

---

Generalization Behavior Overview En {Ro.
Optimization Behavior Overview En
1.55
En (Bo.FF)))) (450k, 0.1 de)
1.8
En-(Re.ft) (450k, 0.1 do)
Valid Cross-ntrppy Laxs (En-43
Er- (Ro.Fr) (450k, 0.2 de)
En-(Ro.Fr) (450k do)
1.50
(Ro,Fr) (450k, 0.3 do)
(Ro.) (450c,
En-Ro Rate
En-fr (400k) + (n-(Ro,Frr) (50k)
En-fr (400k) En- RR.Fr) (50k)
1.45
1.6
0.2
1.40
Croxs-Entropp
1.5
0.4
1.35
1.30
0.7
1.25
0.9
Valid Cross-Entropy Loss (En Ro]
Train Cross-Entropy Loss (En Ro!
Figure 5: pre-training joint fine-tuning has a regularization effect, but cannot be replaced by simply increasing regularization strength. The dropout rate used in pre-training joint fine-tuning is 0.1.
4.2 Multilingual Training
In this section, we expand from a two-task set- ting to many-task setting. We train five languages from the mC4 dataset Engllihh, Hindi, Gujarati, Swahili, and Gaelic- using the span corruption objective from T5 [24]. See Ta- ble for details on the dataset. Canonically the
Goneralization Behavior Overview
Envir (200K)
(200K) EniRa
valld Crosc Crose-Enrroy IEn
(200K) Rn,(Ro,R). (50K)
mC4 dataset is used the pre-training phase for
models (not to be confused by our pre-traing mraitranning joint fine-tuning method). These models are subsequently applied downstream tasks such as question answering. This multilingual pre- training phase also known the language balancing problem. Our goal is to show that our two stage method can effectively balance high- resource and low-resource languages, improving performance on low-resource languages beyond what is achievable by the conventional method of temperature sampling while not sacrificing performance on high-resource languages.
Valid Cross-Entrapp Loss (En Ro)
Figure 6: Fine-tuning solely the low-resource task (En-+Ro) leads to both catastrophic forgetting of the pre-trained task (En-+Fr) and worse low- resource task performance than fine-tuning on all tasks (En+ Fr)).
Note that in the mC4 corpus, English is 16745 times larger than the smallest language we use. This data imbalance un- derscores the necessity for effective language balancing, par-
Table Data used from mC4.
Language
# Chars (B)
En (English)
13, 396
ticularly in determining the proportion of each language to be used during training. This presents a highly challenging and computationally demanding problem, as it is not feasible to simply sweep the scalarization weights as one would two-task setting.
Hi (Hindi)
75
Gu (Gujarati)
3.6
Gd (Gaelic)
0.8
Sw (Swahili)
4.1
For our training setup we closely follow mT5 332 for the
model architecture and training procedure. Specifically, we use the mT5-XXL model (13B parame- ters), which an encoder-decoder transformer architecture. Additional training details are available in Appendix B
Temperature Sampling Because we increase the amount of tasks in this setting, detailing the full scalarization trade-off frontier would be computationally infeasible. Therefore, employ the widely used temperature sampling heuristic [1] 7 2]. Let D; data size of language or task i, we then
define the empirical distribution P for each task i as:
P(a: E task i) = D
(3)
Σ, Dj
7

---

Oo CrAS 37V
32t
Chars
Char Seen
Gevald Covos Enteey
r-3.2 (5:000)
<- 3.22 12.000 3-35 (300/)
En (20000 3.38 (3904)
30P
17
15"
Figure 8: Pre-training on English and joint fine-tuning on all 5 languages leads to better optima for Gujarati, Gaelic and Swahili, the low-resource languages. Pre-training also results better initialization and token-efficiency for all languages newly seen in the fine-tuning phase.
Temperature sampling then uses distribution + defined by temperature parameter T as follows:
Pie E task i)2/T
Q(x E task i)
(4)
E3 P(x task j)
The temperature parameter т controls the peakiness (or flatness) of the sampling distribution. Commonly used 7's in the literature are greater than 1, which essentially ua-samples low-resource tasks and down- samples high-resource tasks.
1.5
3.33 (300K)
1.4
En (200k)]
=m 3.33 1300k)
Bust Valia Erass-Emmey
1.3
1.1
Static Sampling Baseline Temperature sampling ubiquitous due to its simplicity and intuitiveness, but its performance varies greatly with т. For our static sampling baseline, we tuned т among com- monly used values in the literature (1.43, 2, 3.33. 5) at a smaller scale, and found that 7 = 3.33 performed the best in terms of oowrresuurce languages. We also tried a more intricate sampling strategy called Uni- Max (6), but found that the 5 languages chose, did not perform better than 3.33.
1.0
0.9
0.1-
0.0
En
Ci
6d
Figure 7: Pre-training joint fine-tuning yields the best performance in 4 out of languages, with significant improvements in the low- resource tasks.
Pre-training joint Fine-tuning For our pre-training joint fine-tuning setup, we first pre-train on English, reset the optimizer state and learning rate schedule, and then fine-tune on all languages using temperature sampling. We use the same sampling rates as the static sampling baseline (r 3.33) reduce the tuning overhead over static sampling.
As in the NMT experiments, we employ restart baseline to fully ablate the pre- pre-training fine-tuning scheme. The restart baseline resets the optimizer state and learning rate schedule in the middle of training for the static sampling baseline.
Results Figures and X show that while learning rate schedule restart helps performance, pre- training joint fine-tuning yields the best results on the low-resource tasks. Surprisingly, it not only improves the performance English. We note that due to the vast dataset imbalance, on the low-resource tasks before English has chance to converge. Consequently, pre-training joint fine-tuning can leverage the benefits mentioned the previous section-reeulaarizaion transfer, and reduced forgetting to achieve superior lower bound performance with higher token efficiency.
Gujarati, Gaelic, and Swahili, but also shows a slight enhancement on
temperature sampling baseline overfits

---

Generalization Behavior Overviow En (Ro, Fr)
1.6
1.36
(200k) Fn-(Ro,FF)
-- 3.33 (500k)
En-fr (400k) En- En-(Ho.F1) (50K)
(100k)
3:33 (400K)
Valid Entropy Loss itn
En-fr (600k) En- En-(Ro,Fr) (50K)
Cross-Ererpp
1.4
En [300k)
(200k) 3:33 (2001)
-.33,3 1200k)
1.32
Scalarization En-Ro
1.3
3.30
1.1
1.26
0.9
1.24
0.0
GU 5w
1.75
1,80
1.85
1.90
1.95
2.00
2.05
2.30
Valid Cruss-Entrapy Loss (En Rol
(a)
(b)
Figure 9: Left: For language modeling on mC4, longer pre-training leads to better eeee performance for the low-resource languages (Gu, Gd, Sw) despite the decreased length of fine- tuning. On the other hand, due to the decreased length of fine-tuning, high-resource languages do not enjoy the benefits of pre-training. Right: For NMT, when the training budget is not fixed, longer pre-training leads to better overall performance trade-off fronts.
4.3 Length of Pre-training
Our method is simple but comes with some choices to make, one of which is the number of steps to pre-train for. We investigate the effect of the number of pre-training steps in NMT and language modeling mC4 by pre-training with less, and more steps than in the previous sections. With the language modeling task, we fix the total training length to be 500k steps to emulate a compute- constrained scenario. We chose use smaller model LXL opposed to m used in Section 4.2 for faster training). With NMT, we fix the number of fine-tuning steps, but let the total training steps vary.
Figure [a) displays the effects of varying pre-training length in the mC4 experiments. We see that longer pre-training improves best achievable performance on the low-resource tasks of Gujarati, Gaelic, and Swahili. This is despite the fact that the number of fine-tuning steps decreased due to the fixed total step budget. In other words, for the 3 low-resource tasks, longer pre-training improves performance more than exposure to the tokens. On the other hand, performance on English and Hindi worsens with increased pre-training length. For English, this is due to the resetting of the learning rate schedule and the decreasing of fine- tuning steps. Resetting involves a learning rate warmup, which worsens English performance before improving again (see the panel corresponding to En for Figure B.. Decreasing fine-tuning steps gives English less time to recover its performance from pre-training. For Hindi, the worsened performance is simply because it is not low-resource task in this context, and therefore, less tokens seen translates to worse performance.
In Figure 9 we that in the NMT experiments, pre-training longer En translates to better overall trade-off fronts, not just for the low-resource task.
The implications of these results are that when there is a strict training budget, it is better to be conservative and pre-train for a shorter amount of time. However, if the goal to obtain the best performance with no strict compute budget, it is better to pre-train for as long as possible before fine-tuning. Note that longer overall training is an option for our method (by pre-training for longer) but not for static sampling because static sampling needs to constantly be training on the low-resource tasks, which will lead overfitting when training for too long.
5 Related Work
Multitask Learning Multitask learning has gained increased attention in being able to learn many tasks in an efficient way due to parameter sharing and transfer between tasks. the language domain, multilingual neural machine translation 12 [4] enables translation from multiple source languages to multiple target languages. Due to the transfer of information between language pairs, multilingual NMT has seen improvements in low www-resurre language-pal performance compared to training solely

---

on that language pair [12]. In addition to NMT, large multilingual pre-trained language models are used to fine-tune on variety of downstream tasks with different languages 32] Prior works on
intermediate training take advantage of cross-task [23] and cross-lingual [22 transfer to improve
downstream task performance. However, multilingual approaches there exists the problem of dataset imbalance, where low-resource languages tend to suffer in performance. Recently, [6] found that naive temperature sampling might lead to overfitting of low-count languages, and suggested epoch capping with a uniform distribution for high-count languages, showing improvements over temperature sampling. In multilingual NMT, to our knowledge, we are the first to show that a simple pre-training stage high-resource language pair improve the trade-off front of static sampling. Furthermore, our method is orthogonal to innovations in sampling strategies like [6], and can potentially show better results in conjunction with better sampling.
Transfer Learning in NMT The benefits of transfer learning to low-resource language- pange has been long known in the NMT literature 33
g 77 33 showed that pre-training on high-resource
language pair can improve performance compared to training from scratch. While most prior work on transfer learning in NMT focus on improving performance on low-resource bilingual data, recent work [21] used transfer learning to improve performance on multiple language pairs. Unlike the transfer learning literature in NMT [5), we show that pre-training can push the low-resource frontier in the multilingual setting, by testing a grid of sampling rates and hyperparameters to trace the trade-off front. Prior work in the literature study the relationship between the pre-training and fine-tuning language pairs [10) freezing different parts of the model during fine-tuning []. and experimenting with mtny-stage pre-training (9] this direction.
We expect to further benefit from research done in
Curriculum Learning Due to the imbalanced nature of multilingual datasets, a static sampling strategy is unsatisfactory. [30] used a hand-crafted temperature sampling schedule that samples more high-resource carlier in the training, and gradually samples more low-resoure languages. The performance boost from using such a schedule, compared to a static one, supports our observations from pre-training using a high-resource language pair. On the other hand, there many works that employ a more intricate strategy for an adaptive schedule 13
29 18. In comparison, our method is simple with little to no overhead. We include discussion on our experience, though preliminary, with trying an adaptive schedule in Appendix[[ Lastly, (26) showed that the ordering of data within a task affects catastrophic forgetting, which supports our observations.
6 Limitations and Future work
In our experiments, we focus on training on single high-resource task during the pre-training phase. It would be interesting future work to study pre-training with more than one language or language-pair. We also only experiment with fine-tuning all parameters of the pre-trained model. Studying the effect of freezing different parts of the model during fine-tuning, potentially as function of the relationship between pre-training and fine-tuning tasks, is left to future work.
7 Conclusion
In this work, we demonstrated the benefits of pre-train joint fine-tune setup for multi-objective optimization when there is a mixture of high and low-resource tasks. We show that in the presence of large data imbalance, the order at which tasks are introduced has significant impact on overall performance. We demonstrate through a variety of experimental settings that this methodology produces points that can go past the trade-off frontier achieved by scalarization. We show that a major weak point of scalarization in this regime is that it overfits on the low-resource task, being unable to early stop due to the high-resource task not converging. Our method both allows the high- resource task to converge during pre-training and prevents overfitting through joint fine-tuning. It also outperforms scalarization that muder-sample the low-resource task due to higher token efficiency. We also show that fine-tuning only on the low-resource task, popular scheme in the NMT literature, undesirable due to its inability to prevent forgetting. Our method simple natural strategy for avoiding the above failure modes. Given the significant performance boost we observe in our experiments, we believe that this training regime has the potential to become standard approach, particularly in the era of large language models.
10

---

Acknowledgments and Disclosure of Funding
We thank George E. Dahl, Wolfgang Macherey, and Macduff Hughes for their constructive comments on the initial version of this manuscript. Additionally, we thank Sourabh Medapati and Zachary Nado for their help in debugging our code base. Moreover, we are grateful to Soham Ghosh and Mojtaba Seyedhosseini for valuable discussions regarding the role of MTOs in large-scale models. Lastly, thank Chris J.H. Zhang for helpful discussions.
References
[1] Alham Fikri Aji, Nikolay Bogoychev, Kenneth Heafield, and Rico Sennrich. In neural machine translation, what does transfer learning transfer? Association for Computational Linguistics, 2020.
[2] Naveen Arivazhagan, Ankur Bapna, Orhan Firat, Dmitry Lepikhin, Melvin Johnson, Maxim Krikun, Mia Xu Chen, Yuan Cao, George Foster, Colin Cherry, et al. Massively multilingual neu- ral machine translation in the wild: Findings and challenges, arXiv preprint arXiv: 1907.5090, 2019.
[3] Ankur Bapna, Colin Cherry, Yu Zhang, Ye Jia, Melvin Johnson, Yong Cheng, Simran Khanuja, Jason Riesa, and Alexis Conneau. mslam: Massively multilingual joint pre-training for speech and text. arXiv preprint arXiv:2202.01374, 2022.
[4] Ankur Bapna, Yu-an Chung, Nan Wu, Anmol Gulati, Ye Jia, Jonathan H Clark, Melvin Johnson, Jason Riesa, Alexis Conneau, and Yu Zhang. Slam: A unified encoder for speech and language modeling via speech-text joint pre-training. arXiv preprint arXiv:2110.10329 2021.
[5] Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
[6] Hyung Won Chung, Xavier Garcia, Adam Roberts, Yi Tay, Orhan Firat, Sharan Narang, and Noah Constant. Unimax: Fairer and more effective language sampling for large-scale multilingual pretraining. In The Eleventh International Conference Learning Representations.
[7] Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wen- zek, Francisco Guzmán, Edouard Grave, Myle Ou, Luke Zettlemoyer, and Veselin Stoyanov. Unsupervised cross- coss-lingual representation learning at scale. arXiv preprint 2019.
[8] Marta R Costa-jussă, James Cross, Onur Çelebi, Maha Elbayad, Kenneth Heafield, Kevin Hefferman, Elahe Kalbassi, Janice Lam, Daniel Licht, Jean Maillard, et al. No language left behind: Scaling human-centered machine translation. arXiv preprint arXiv:2207.04677 2022.
[9] Raj Dabre, Atsushi Fujita, and Chenhui Chu. Exploiting multilingualism through multistage fine- tuning for low-resource neural machine translation. Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 1410-1416 1416, 2019.
[10] Raj Dabre, Tetsuji Nakagawa, and Hideto Kazawa. An empirical study of language relatedness for transfer learning in neural machine translation. In Proceedings of the 31 Pacific Asia conference on language, information and computation, pages 282-286, 2017.
[11] Jacob Devlin. Multilingual bert readme. httpss:/ www.giituu..ooooogleeeeeeee cow.r blob/master/mlltilnnee
md. 2018.
[12] Orhan Firat, Kyunghyun Cho, and Yoshua Bengio. Multi-way, multilingual neural machine translation with a shared attention mechanism. arXiv preprint 160iv: 2016.
[13] Sébastien Jean, Orhan Firat, and Melvin Johnson. Adaptive scheduling for multi-task learning. arXiv preprint 1rXin:1909.06434. 064344 2019.

---

[14] Melvin Johnson, Mike Schuster, Quoc V Le. Maxim Krikun, Yonghui Wu, Zhifeng Chen, Nikhil Thorat, Fernanda Viégas, Martin Wattenberg. Greg Corrado, et al. Google's multilingual neural machine translation system: Enabling zero-shot translation. Transactions of the Association for Computational Linguistics, 5:339-351, 2017.
[15] Yunsu Kim, Yingbo Gao, and Hermann Ney. Effective cross-lingual transfer of neural machine translation models without shared vocabularies. arXiv preprint arXiv:1905.05473 2019.
[16] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[17] Tom Koemi and Ondřej Bojar. Trivial transfer learning for low-resource neural machine translation. arXiv preprint 1809. 2018.
[18) Julia Kreutzer, David Vilar, and Artem Sokolov. Bandits don't 't follow rules: Balancing multi- facet machine translation with multi-armed bandits. arXiv preprint a1Xiv:2110.06997, 2021.
[19] Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv: 180.:v:888..0222:: 2018.
[20] Vitaly Kurin, Alessandro De Palma, Ilya Kostrikov, Shimon Whiteson, and M Pawan Ku- mar. In defense of the unitary scalarization for deep multi-tast learning. arXiv preprint arXiv:2201.04122 2022.
[21] Surafel M Lakew, Aliia Erofeeva, Matteo Negri, Marcello Federico, and Marco Turchi. Transfer learning in multilingual neural machine translation with dynamic vocabulary. arXiv preprint arXiv:1811.01137 2018
[22] Jason Phang, Iacer Calixto, Phu Mon Htut, Yada Pruksachatkun, Haokun Liu, Clara Vania, Katharina Kann, and Samuel Bowman, English intermediate-tass training improves zero-shot cross-lingual transfer too. In Proceedings of the Ist Conference of the Asia- Asia-Pacific Chapter of the Association for Computational Linguistics and the 10th International Joint Conference Natural Language Processing, pages 557-575, 2020.
[23] Jason Phang, Thibault Févry, and Samuel R Bowman. Sentence encoders on stilts: Supplemen- tary training on intermediate labeled-data tasks. arXiv preprini arXiv:1811.000 2018
[24] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with unified text-to-text transformer. J. Mach. Learn. Res., 21(140):1-61 2020.
[25] Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthomn,, Aitor Lewkowyez, Alex Salcianu, Marc van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bulian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee, Dan Garrette, James Lee-Thorp, Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Mlatinn-Shepall Noah Fiedel, Mark Omernick, Brennan Sacta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, and Andrea Gesmundo. Scaling up models and data with t5x and seqio. arXiv preprint 2rXiv:2203.17182 2022.
[26] Chenze Shao and Yang Feng. Overcoming catastrophic forgetting beyond continual learning: Balanced training for neural machine translation. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2023-2036. 2022.
[27] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[28] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V. Luxburg, Bengio, H. Wallach, Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
12

---

[29] Xinyi Wang, Yulia Tsvetkov, and Graham Neubig. Balancing training for multilingual neural machine translation. arXiv preprint 0667480 2020.
[30] Yiren Wang, ChangXiang Zhai, and Hany Hassan Awadalla. Multi-task learning for multilingual neural machine translation. arXiv preprint arXiv:2010.02523 2020.
[31] Derrick Xin, Behrooz Ghorbani, Justin Gilmer, Ankush Garg, and Orhan Firat. Do current multi-task optimization methods in deep learning even help? Advances in Neural Information Processing Systems, 35:13597-13609 2022.
[32] Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami AI-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mT5: A massively multilingual pre-trained text-to-text transformer. Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pages 483-498, Online, June 2021. Association for Computational Linguistics.
[33] Barret Zoph, Deniz Yuret, Jonathan May, and Kevin Knight. Transfer learning for low-resource neural machine translation, arXiv preprint 1iii:I6044.0222 02201, 2016.
13

---

A NMT Experiments: Additional Information
A.1 Detailed Training Setup
This section details the experimental setup used in Section 4.1 We use the pre-LN encoder-decoder transformer architecture. The experiments presented the main use three layers for both the encoder and decoder, but we also present results with 6 layers for the encoder and decoder. We follow the convention in NMT literature and train our models with 0. label smoothing and 0.1 dropout for feed-forwel and attention layers. See Table for complete architecture details.
Table 3: Transformer architecture details and common hyperparameters.
Hyperparameter
Feed-forward dim
2048
Model dim
512
Attention heads
Attention QKV dim
512
Label smoothing
0.1
Dropout
0.1
We use SentencePiece tokenization 19 to generate vocabulary of size 64,000 for each NMT problem (e.g. En-+ Fr)).
All models were trained using the Adam 16 optimizer with a batch size 1024. For all our NMT experiments, used linear warmup to the desired learning rate, followed by cosine decay schedule that decays to O. This is true for all legs of training for methods that use our scheme; during the pre-training phase, do linear warmup followed by cosine decay, and during the fine-tuning phase, after loading the pre-trained model, we do a linear warmup followed by cosine decay.
For the baseline experiments that do not do pre-training, and also for the pre-training portion, we warmup for 40k steps. For fine-tuning, we tune the warmup steps from within (10k, 20k, 30k, 40k) for all experiments other than for En En-, Fr), where warmup for 40k steps. The base number of training steps, and the number of fine-tuning steps are shown in Table 4 Note that for comparison's sake we also trained a baseline-without-oe wat--ttttt-------ranninnn pre-trainnng model for 'base + fine-tune' number of steps.
Table 4: Number of training steps for all NMT experiments.
3 3-layer
6-layer
base
fine-tune
base
fine-tune
En- (Zh, Fr)
300k
300k
300k
300k
En- ( Ro, Fr)
400k
50k
300k
50k
En (Hi, Fr)
300k
50k
275k
50k
For all experiments, sweep the base learning in the grid (2.5e-4, 5e-4, 2.5e-3, 5e-3,, 7.5e- 3). We also sweep the sampling rate for En-*Fr and En-tCs in the grid {i/10) which fully determines the sampling rate for the other language pair. All plotted points correspond the final measurement taken for each trial.
For all fine-tuning experiments, when loading the pre-trained model checkpoint, reset the optimizer state. We also trained all parameters of the model, and did not freeze anything.
14

---

A.2 Additional Performance Trade-Off Curves
In this section, we present the performance trade-off curves for (H-, Fr). as well as for 6-layer models on En->Zh, Fr). En Roo Fr]. and En- Fr). The black-bordered points in the generalization portion of Figures and 13 below correspond to the restart baseline.
Generalization Behavior Overview En - (2h, Fr)
Optimization Behavior Overview En (2h, Fr)
1.50
En (2h,Pr) (300k)
Train Entrogy (En=FF
En- [n- (300k)
Valid Cross Eles--Entooy Loss (En
1.45
En-[2h,F) (600K)
En-Fr (300k) + En-(2h-F)) (300k)
1.40
En-[2h,F/ (600K)
Er-Ft (300k) En-(Zh,R) (300K)
1.40
1.35
1.35
1.30
1.25
1.25
1.20
1.20
1.15
1.15
1.10-
2.0
Entropy Loss (En Zn)
Train Cross Entropy Loss [En Zh)
Figure 10: Performance trade-off behavior for En->(Zh, Fr] with 6-layer models. Each point corresponds the final performance of model. Similarly to the 3-layer-model case (Figure I
pre-training does not yield improvements.
Generalization Behavior Overview (Ro. Fi))
Optimization Behavior Overview (Ro, Fr)
Valid Cross Cross-Ettren Entsoyy Loss (En -50
1.28
En-(Ra.Fr Rr) (300k)
Train Cross Entropy Loss (En-FF)
1.6
En- (300K)
(o.Fr) (250k)
(Ro.Fr) (250k)
1.26
En-Fr (300k) En-(Ro.Fr (50K)
1.5
En-Ar (500K) En-(Ro,Fr) (50K)
Best Scalarization En-Ro Loss
1.24
Baseline
1.22
1.3
1.2
1.16
1.1
1.650
1.675
1.725
1.750
1.800
1.825
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
Valid Cross-Entrapy Loss (En Ro)
Train Cross Entropy Loss (En Ro)
Figure 11: Performance trade-off' behavior for En (Ro, Fr with 6-layer models. We see a similar behavior as with 3-layer models. In addition, we are able to further improve the performance on both En-RRo due to a larger model size.
Generalizetion Behavior Overview En- (Hi, Fr)
Optimization Behavior Overview En (Hi, Fr)
1.8
1.33
En-(Hi,F) (3004)
En-(H.Fr) (300k)
(H-(H)R) (350k)
En-(HLFr) (350K)
Valid Cross Entropy Loss [En
1.32
En-tr (300k) + En-(Hirr) 150R]
Train Cross Cross-Entrron Entrosy Loss [En
En-Fr (300k) + En-(HLFr) (50k)
1.31
Best Scalarization En-Hi
Restart Baseline
1.30
3.5
1.29
1.28
3.4
1.27
1.3
3.2
1.25
2.90
2.95
3.00
3.05
3.10
3.15
3.20
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Valid Cross-Entropy Loss Hi)
Train Cross-Fntropy Loss (Fn
Hi]
Figure 12: Performance trade-off behavior for En->(Hi, Fr) with 3-layer models. These results mirror those seen in Figure 2 than French and Romanian.
We note that here French and Hindi are more linguistically dissimilar
15

---

Generra Gentralizetion Behavior Overview En- (Mi,
Optimization Behavier Overview En (Hi,
Valid Crass-Entrop Loss (ois F0
1.26
En (n-(HiTr) (275g)
En (HLP) (325k)
Train Crass Entropy Lazs (En -FF
1.6
En (329K)
En-(HRrr) (275k)
(n-(HiRF)
(275k) En-(HLFF)
1.24
Bast Scalarization ErvHi Loss
Restart Baseline
1.22
1.3
1.20
1.1
2.75
2.80
2,85
2: 2.90
2.05
3.00
3.05
3.10
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Valid Entropy
Cross-EEnnnyg (E)+HI)
Figure 13: Performance trade-off behavior for En-+(Hi Fr) with 6-layer models. As with the 3-layer models, We observe similar improvement in both En- and En-+Fr performances, despite the dissimilarity of French and Hindi.
A.3 Performance Trade-Off Curves with Sampling Rate as Markers
In this section, we present the same performance trade-off curves as shown previously, but with the markers representing sampling rates for the lowerrresource language pair. We can see that in all but case (En-+ (Hi,Fr) 6-layer model; Figure T9). the model that performs the best in the low-resource language pair, samples the low-resource language pair at a higher rate than the baselines that do not use pre-training. The black-bordered points in the generalization portion of Figures T6 and
18
19 below correspond to the restart baseline.
Generalization
Vallo Cress-Enttoyy Loss (tn=ftt
1.6
En (2),F) 1200s)
Enx(Zh,h)) (600c)
1.55
En-(2h.f) (800k)
En-(2hFF) (300K)
Eisfi (300k) (=((2b.F) (300k)
--- Ensh (200k) + En-(2lur) (300k)
3.45
0.1
0.2
1.40
0.3
3.30
1.3
1.25
0.9
12
1.15
2.4
2.6
2.8
3.0
1.5
21 Train Cross Entropy Loss (En
17
18
2.0
25
Volid Cross Entropy Loss (En 2hl
Figure 14: Performance trade-off behavior for En (ZZ, Fr) with 3-layer models. We clearly
see that there is no optimal rate in this case, since we trace a Pareto front as we vary the En-+Zh sampling rates from 0.11 to 0.9.
Generalization Behavior Overview En (2h, Ir)
Optimization Behaaior Overview En (2h, Fr)
2 1.50
En-(2h.Fi Fr): (300)]
En- (2h.h) (300k)
1.45
(2t.22h.))
3.40
Fn-7h
Trai Crosse Entrrgy Losa
1.35
Crose Entroe
1.35
0.4
1.30
1.25
0.5
1.5
120
3.15
2.10
2.0
2.2
2.4
2.5
2.8
1.4
2.5
1.5
1.0
2.0
2.1
Valid Cross-Erempy Loss [Ein Zhi
Train Cross Entmayy Loss (En++ 25)
Figure 15: Performance trade-off' behavior for (Zh, Fr) with 6-layer models. We observe a similar behavior in the 3-layer case.
16

---

Generalizaton Behavior Overview En (Ro. Fr)
Optimization Behavior Overview En + (Ro, Fr)
1.36
En-(Ro.Fr) (600k)
Train Crosu Coos--trropp Loos IEn PP
Eny(Ro,Fr) (600k)
Vold Crosrtnttooy Loss (tm-
1.34
Ensft (1) (50k)
Fr-(Rtu,F))
Endr (403k) En-((u.5)) (56k)
En
1.32
Best Scalartzacion En-Ro Loos
Restwrt Bavelie
1.30
0.2
-4
0.4
0.5
126
1.75
1.80
1.55
1.90
1.05
3.00
05
0.2
0.3
0.4
0.5
0.6
0.7 0.8
0.9
Vali Cross Entropy Loss IEn Roi
Train Cross Entropy Loss IEn Rol
Figure 16: Performance trade-off behavior for En-, Fr) with 3-layer models. Unlike the En- Fr) case, we have a few sampling rates that are more optimal than the rest. Pre-training allows sampling Rn at higher rate without overfitting, than without pre-training.
Generalization Behavior Overview En (Ro. Fr)
Optimization BBehavior Overview En- (Ro, Pri
En- (Ru (200k)
1.6
En-(Ru.FF)) (300k)
En (350k)
En (Ro.F1 (350k)
1.76
7.5-
(n-(Bt.h) (50k)
PowRo Rate
Scale En-Ro Loss
Train Croes www.--ytooyyyooooo Look
Rostart Baseline
Veid Croe En0 We
1.22
0.4
3.3
0.5
1.20
3.2
1.6
1,600
1,675
1,700
1.725
1,750
2.775
1,000
1.825
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
Voild Cross-Petrogy Loss RRi--RR
Train Cross-Pntrooy
Figure 17: Performance trade-off behavior for (Ro, Fr) with 6-layer models. We see a similar behavior as in the 3-layer case.
Conorafization Behavior Overview En (Hi, Fr)
Optimization Behavier Overview En (Hi,
1.8
1.33
(H),8) (300k)
[HLFr) (300K)
En-[HiPP3 (350k]
En-[HLPrr (3506)
1:32
En Fr (300k) En (HIFF) (SOK)
Loss
En Fr (3001) En [HURQ (50R)
troyytttrgyy Loss
Restart Baseline
Best Scalorization En Losz
2.6
0.2
0.1
1.30
0.3
2.5
1.8
2.4
1.27
Train
0.8
0.9
1.25
2.50
3.05
3.10
3.35
3.20
0.00
0.25
0.50
1.50 1.75 Train Eroos-Enttrpy Loos (En-e Hil
0.75
1.00
2.00
Valid Cross-Entrooy Loss Hi))
Figure 18: Performance trade-off behavior for En- Fr) with 3-layer models. Like in the
En- Fr), pre-training allows sampling En-HHi at a higher rate without overfiting than without pre-training.
Coneralization Behavlor Overview En Fr)
Optimization Behavior Overview [HL
Valid treestrinnyy Loss [tre m
1.26
En-(HUPr) (325k)
(HLPr) (275k)
En-(Hirr) (325k)
(H.Fr) (275k)
En fr (275k) En (HUR) (50k)
Train Eroopyy Loss
En-Fr (275K) En (HiFr) 150K)
3.28
Best Scalarization En Hi Loss
Restart Baseline
1.5
0.2
0.1
2.4
0.3
1.22
2.3
1.20
2.2
0.81
0.9
1.18
2.85
2.90
2.55
3.00
3.05
3.10
0.00
0.25
0.50
0.75
1.00
3.25
1.50
1.75
2.00
Valid Cross-Entrop Looa H4)
Train Cross-Entrpy Loss (k))
Figure 19: Performance trade-off behavior for (n- Fr) with 6-layer models. In this case, pre-training still allows sampling En-HH at a higher rate, but the rate that yielded the best En-HH was surprisingly the same rate as the baseline without pre-training
17

---

A.4 Efficiency Plots
In this section, we plot the number of examples seen from one language pair against the validation at train step t
cross-entropy loss on that language pair. The number of XX-YY examples is computed by multiplying i, the batch size, and the sampling rate for XX- Each curve in given figure corresponds to the trial that achieved the best final validation performance low(er)-resource language pair within the method given by the legend (i.e. the blue curve in Figure 20 corresponds the trial that achieved the best final validation En-+Zh cross-entropy loss among all trials that did not use pre-training, and was trained for 300k steps.) For the curves corresponding to our proposed pre-training and fine-tuning scheme, we only show the fine-tuning portion of training.
the
Note that initial linear decay followed by a smooth decay is an artifact of evaluating on a linear-scale when the plots are in log-scale.
Valid Cross Entropy Lass (En- Zh)
10-
En-(Zh,Fr (300K)
Velid Cross: Entropy Loos (En-fr)
(n-(2),FF) (300k)
En-(22.rr) 1600k)
En-(2hiFr) (600k)
En (300k) En (2h.2) (300k)
(300k) En (2-(22h.F) (300K)
6-
4-
3-
10
10a
10a
10
10
10a
10r
30
10
10
Nam. Examples Seen (Fn 2h3
Num. Examples Seen (Fn Fr)
Figure 20: For the 3-layer model, pre-training does not provide any significant gains in training efficiency for En-*Zh when pre-training on En->Fr. Given that the blue and red curves coincide towards the end of training, we can anticipate that pre-training did not impair En Enn training (by providing a suboptimal initialization), and that if we were to train the red curve for 300k more steps, it would be able to catch up with the orange curve (best En 2n performance).
Valid Cross Entropy Loss (En 2h)
En-(Zh,F) (300K)
(Zh.Pr (600K)
Valid Cross- Entrapy Loss (En-11)
10-
En-(2h,Fr) (300k)
En-(2N (600K)
En-Ri (300k) En-(2h,Fr) (300)]
En-Fr (300k) En- En-[27,,r/ (300k)
6
4-
2
Num. Examples Soon (En 2h)
Num. Examples Seen (En
Figure 21: We observe a similar behavior with 6-layer models as with 3-layer models.
Vald Cross Entroyy Loss (En- Ro)
En.(Re,Fr) (400K)
10
En-(Ro.Fr) (400K)
En-(Ro,Fr; (450k)
En- F/) (450K)
En-(Ro.Fr)
En-(RouFr)
6
Entrapy
4
10
10
10 10
10
10
10
10
10
10
10
Num. Examples Seen (En Ro)
Num.
Examples Seen 1
Figure 22: On the 3-layer models, pre-training able to accelerate training on En-RRo when pre- trained on En- Even with less overall examples seen in En-RRo, we can perform better than the baselines that did not use pre-training
18

---

Valid Cross Entropy Lass (En -Ra)
En-(RR,FF) (300K)
1.0
Fn-(Ro.FF) (3006)
En (Ro.ft) (350k)
En (Ro (1) (350k)
En (300k) En-(Ro.Pr
(300k) En-(Ro,Fr (50K)
Valid Cross-Entroyy
100
10
30a
10c
10'
100
30
10a
10F
102
Num. Examples Seen [En Ro]
Num. Examples Seen (En-Fri
Figure 23: We observe similar efficiency boost with 6-layer models with 3-layer models.
Valid Cross Cross-Etrrpp Entropy Loss (En HHI
(HiF (300k)
Valid Loss P0
10
En-(HiFr)
10-
En-(HLFr) (350k)
(n-(HiFr) (3506)
En-Fr (300k) En-(HLFr) (50k)
Er-fr (300k) (H-(HLF) (506)
6-
102
Num. Examples Seen (En
Num. Examples Seen (En=FF)
Figure 24: On the 3-layer models, we observe a similar efficiency boost as with En-> (r)
Volld Cross Cross-Entrop Entropy Loss (En HH)
En-(HiFF) (275k)
Volid Entroyy (En Fr))
10
En-(HLFr) (275k)
10-
En-(HUFi)) (325k)
En-(HLFr)
En-Fr (275k) Fn-(HLFr) (50k)
Fr-fr (275k) + Fn-(HLFF) (50K)
Examples
Figure 25: On the 6-layer models, we observe similar efficiency boost as with 3-layer models.
A.5 BLEU Score Plots
Here, we present the performance trade-off curves for when the metric is BLEU score instead of cross-entropy loss, All translations are generated via Seam-Search with beam size of 4.
Ezationn Behavior Overview En - (2h, Fr)
Valid BLEU (En
35
34
En-(Zh,Fr) (300k)
33
En-(Zh.Fr) (600k)
En-Fr En-(Zn,Fr) (300k)
32
Valid BLEU Score (En Zh)
Figure 26: The BLEU score plot paints a better picture for pre-training than the cross-entropy plot (Figure []. since pre-training was able to improve the En-Zh BLEU score to be on par with the score of joint training for 600k steps. Results are with 3-layer models.
19

---

Genaralization Behavior Qverview En (Ro, Fr)
Generalization Behaviar Overview En (Ro, Fr)
36.2
Valld (En-FFi
35.0
said BLEU Score (En-FF)
36.0
34.5
34.0
35.6
(n-(Fo..r) (400k)
En-(Ro,F) (3004]
(n-(Ru./r) (450k)
En-(Ro,/t) (3504)
En-(Ro,Fr
Fr-Fr (300k) En-(Ro,F)
33.0
35.2
24.0
24.5
25.0
26.0
23.5
24.0
24.5
25.0
25.5
26.0
26.5
27.0
Valid BLEU Score (En Ro)
Valid ULEU
(En
Figure 27: Our proposed pre-training scheme improves upon the best BLEU score for En- without pre-training for both the 3-layer models (left) and 6-layer models (right).
Generalization Behaviar Overview En (Hi, Fr)
Generalization Behavior Overview En (Hi, Fr)
34.75
36.00
34.50
BLEU (En-aFr)
34.25
35.50
34.00
35.25
33.75
33:50
34.75
33:25
(H-(H)))) (300k)
En-(HI,R) (275k)
33.00
En-(HLFr) (250K)
34.50
En-(Hi,Fr (325k)
(300K) En-(MLR) (50K)
En.F (275k) (n-(H)
32.75
10.5
11.0
11.5
12.0
12.5
13.0
9.0
9.5
10.0
10.5 11.0 11.5
12.0
12.5
13.0
Valid BLEU Score Hi)
Vald BL Score (En-HH)
Figure 28: Our proposed pre-training scheme improves upon the best BLEU score for En->Hi without preiting for both the 3-layer models (left) and 6-layer models (right). The improvements are more substantial than for En Eno, Fr).
B Additional Training Details in Multilingual Training
We use an additionally processed version of the mC4 32 dataset as proposed in 6 (documents with language ID confidence below 0.95 were filtered).
The model architectures used are the same as mT5 models 322 except that relative position embed-
dings are not shared across layers. We also use the number of real target tokens as the effective loss normalization instead of using a loss normalization factor.
We use SentencePiece tokenization 19 to generate a vocabulary of size 64,000. The corpus used to generate the vocabulary sampled from the training data using temperature sampling with r 3.33.
We use the TSX library 255 to train the models. For all experiments, we use the Adafactor optimizer 27), where we momentum, and do not factorize the second moment of the Adafactor states. The baseline run without fine- mee-tunigg and the pre- ore-training phase of our proposed method, was run with constant learning rate of 0.01 in the first 10,000 steps and inverse square root decay afterwards. For the fine-tuning phase of our method, we reset the optimizer state, and do a 10,000-step linear warmup with inverse square root decay afterwards.
20

---

C Discussion on Sampling Rate Schedules
From our preliminary experiments on using schedules for the sampling rates in the NMT workloads, we find that the learning rate schedule must be tuned accordingly, which affects the overall perfor- mance of the run. For example, we find that cosine decay schedule performs better than inverse square root decay for scalarization. However, if we use cosine learning rate decay in conjunction with linear sampling rate decay (used by DDS, and defining sampling rate to for the high-resource language-pair), by the time the sampling rate for low-resource task is high enough, the learning rate has decayed rapidly (by nature of cosine decay), resulting little learning for the low-resource task. Using inverse square root learning rate decay solves this issue, but this results in overall worse performance due to the suboptimal learning rate schedule. In contrast, our method is free to use any scheduler that maximizes performance in each leg of training (pre-training and fine-tuning). Lastly, when tuning hyperparameters, using dynamic sampling rates requires executing the full training run many times. On the other hand, for our method, phase, (since the pre- pre-traine training phase has only one task, and is an easier optimization problem) which shorter than the total training time.
focus resources tuning the fine-tuning
21


Reviews:
[{"review_id": "micnjfhjqK", "reviewer": "Reviewer_yZ9n", "summary": "This paper presents an empirical finding: in the context of data imbalance in multi-task learning, pre-training on high-resource tasks then fine-tuning on a mixture of high/low-resource tasks can achieves superior results, compared to standard weighted sampling training. Authors applied the proposed training method *pre-training joint fine-tuning* to neural machine translation (NMT) and multi-lingual language modeling. Both experimental results shows that the proposed method allows the high-resource task to converge faster while preventing overfitting in low-resource tasks.", "questions": "Why temperature based sampling is only applied in the language modeling task not the NMT task? What is the advantage of temperature based sampling compared to scalarization in the large scale dataset? ", "limitations": "I am curious to see what if the authors applied this method in a more realistic setting: training LM on the whole mC4 dataset, and then evaluate the performance (low-resource language vs high-resource language) on different downstream tasks as in mT5.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Achieving a good balance between low-resource and high-resource languages during multitask training is a critical research area, particularly in the context of large-scale language modeling. The proposed approach of pre-training followed by joint fine-tuning is a timely solution that addresses this challenge.\n- The results obtained from the empirical study demonstrate the superiority of the pre-training joint fine-tuning approach over the conventional method of weighted sampling during training. The model trained using this approach consistently outperforms existing methods, showcasing its effectiveness in addressing the challenges posed by data imbalance in multi-task learning.\n- An additional ablation study showed that the pre-training joint fine-tuning has a regularization effect (to avoid overfiting to low resource language). This improvement cannot be achieved by simply increasing drop-out ratio.", "weaknesses": "- In the main paper, authors only reported perplexity metrics in the evaluation results. For multi-lingual language modeling task, I would like to see the evaluation on downstream tasks as in the mT5 paper. It would be helpful to demonstrate the impact of the proposed method if authors can show the improvement in downstream tasks on XQuAD or XTREME benchmark.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. Here is our response to the weaknesses and questions pointed out by the reviewer.\n\n**W1a:** *In the main paper, authors only reported perplexity metrics in the evaluation results.*\n\n**Response:** For the NMT experiments, we report BLEU score evaluations in Appendix A.5.\n\n\n**W1b:** *For multi-lingual language modeling task, I would like to see the evaluation on downstream tasks as in the mT5 paper.*\n\n**Response:** We agree with the reviewer on the importance of evaluating our proposed method on downstream tasks. We did not evaluate on downstream tasks for our paper because almost all languages included in the downstream tasks for evaluating cross-lingual generalization ability (e.g. XTREME) are high-resource in mC4 given our training budget (with the exception of Swahili and Yoruba). Furthermore, we believe that it will be more meaningful to evaluate on such downstream tasks after pre-training on many more languages for a longer time. We leave the scaling up of our method for future work to study, and we will expand on our limitations section to reflect on the reviewer’s comment.\n\n**Q1:** *Why temperature based sampling is only applied in the language modeling task not the NMT task?*\n\n**Response:** We do not need to do temperature sampling for the NMT tasks because we already search through the possible sampling rates in a grid. In other words, the sampling rates we have tried for NMT already include temperature sampling. \n\nIn the NMT experiments, we only have two tasks (language-pairs), which allows us to test sampling ratios in a grid (e.g. For En->{A, B}, we can try sampling rate x for En->A and 1-x for En->B, where x is in {0.1, 0.2, …, 0.9}). Doing so for more than 2 tasks will require testing exponentially many sampling ratios, which is why for the 5 task case in the language modeling experiment, we resort to temperature sampling.\n\n**Q2:** *What is the advantage of temperature based sampling compared to scalarization in the large scale dataset?*\n\n**Response:** Temperature sampling is a heuristic to obtain sampling rates to be used for scalarization (since for higher number of tasks, testing a grid of sampling rates is not feasible). Many prior work used temperature sampling with various temperatures, and its advantage is that it is simple, intuitive, and can be controlled with a single parameter. "}, {"review_id": "oJNHOAalw5", "reviewer": "Reviewer_wUUc", "summary": "This paper presents an analysis of multi-task training with the existence of low-resource tasks (in the context of the paper, they focused on multi languages). The paper argues that it is better to first \"pre-train\" on a high resource task and then \"fine-tune\" on a joint of high and low resource tasks. Experiments were conducted in neural machine translation and multilingual language modeling to support the hypothesis.", "questions": "The paper focuses on a specific multi-task setup of multilinguality, however, the finding/hypothesis doesn't seem to have strong ties with multilingualityThe paper presented some limitations on scaling the tasks and settings to be analyzed., and is therefore not necessarily confined to this setting. I am curious if the authors conducted any experiments on the the general multi-task setup.", "limitations": "The paper presented some limitations on scaling the tasks and settings to be analyzed.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The proposed hypothesis intuitively makes sense and the experiments have been conducted in a careful setting. The proposed solution is simple, but has its effectiveness verified in NMT and language modeling. The paper has a smooth flow and is easy to understand.", "weaknesses": "1. While the proposed solution is effective, the benefits of having a multilingual model (or in general a multi-task model) is that a single model can be used in various situations, without the need of (continue) training on different data splits and producing several models. I feel at the current stage the application of this finding might be limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. Our responses to weaknesses and questions are below:\n\n**W:** *[…]the benefits of having a multilingual model (or in general a multi-task model) is that a single model can be used in various situations, without the need of (continue) training on different data splits and producing several models.*\n\n**Response:** We completely agree with the reviewer about the benefits of multitask models, and how one model can be used in various situations. We believe that our proposed method does not deviate from this. Our method can be thought of as a specific scheduling strategy for sampling rates to balance the tasks. While the current convention is to train on all the tasks at once with a fixed ratio, we propose to use a 100% sampling rate for high-resource tasks in the beginning of training, before training on all the tasks with a fixed ratio. Our method yields a single model that can be used on all tasks trained, and does not need to be trained further on different data splits.\n\n**Q:** *I am curious if the authors conducted any experiments on the general multi-task setup.*\n\n**Response:** In our work, we frame multilingual learning as a multi-task optimization problem (as in many prior work [1, 2, 3, 4, 5]), and focus on utilizing cross-lingual transfer. We have yet to conduct experiments on the “general multi-task setup” (multiple NLP tasks that may or may not be multilingual), where the focus would be on utilizing cross-task transfer. We will update our future work section to reflect this. \n\n[1] Luong et al, 2015, Multi-task Sequence to Sequence Learning\n\n[2] Firat et al, 2016, Multi-Way, Multilingual Neural Machine Translation with a Shared Attention Mechanism\n\n[3] Arivazhagan et al, 2019, Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges\n\n[4] Jean et al, 2019, Adaptive Scheduling for Multi-Task Learning\n\n[5] Wang et al, 2020, Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models\n"}, {"review_id": "CrNXgb9W5h", "reviewer": "Reviewer_9ZBe", "summary": "The paper presents a multi-task training approach that involves pre-training on high resource tasks, followed by joint fine-tuning on the full set of tasks. The intuition is that high-resource tasks need a larger number of steps to reach convergence, whereas low-resource tasks will risk over-fitting if trained for that same number of steps.\n\nThe authors show that the proposed approach improves the performance for neural machine translation of low-resource languages (en-ro and en-hi, pretrained on en-fr) and similarly improves performance for multilingual language modelling for the case of low-resource languages.", "questions": "* Line 128 refers the reader to Figures 21 and 22 in the appendix to show that the proposed approach does not hurt performance in the high-high case. It's not clear to me how the said figures show this. The clearest evidence for me would have been a version of Figures 26 and 27 for the English+Chinese case. Would you be able to provide such a plot?", "limitations": "The authors do a good job of listing and addressing limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The authors show that _pretraining joint-finetuning_ is a simple approach to boost low-resource task performance in multi-task training. The experiments, which focus on language modelling and machine translation, are clearly designed and presented, with reasonable baselines.\n* In this proposed training approach, the pre-training phase ends up taking up most of the training time. Since pre-training is conducted on the subset of high-resource tasks, this makes it an easier optimisation problem compared to the standard approach where all tasks are trained. The potentially trickier joint finetuning step requires overall fewer steps. Hyperparameter sweeps could plausibly be limited just to this final step, making them cheaper and faster compared to the baseline joint training approach.\n* Between the paper and the appendices, enough details are given so that reproducing these run should present no issues.", "weaknesses": "* I would have loved to see some more realistic MT experimental settings. Pretraining on en-fr for either en-zh or en-hi (or viceversa) feels very unrealistic. There are much closer language pairs that could have been used, which would have made the results more convincing.\n* I am concerned about the results of the experiments of section 3.1.1. What's the effect on BLEU scores here? It would have been good to include these in Appendix A5, as was done for the other two language pairs.\n\nMinor:\n* Line 61: _if_ is attached to $\\forall$\n* Line 171: shows -> show\n* Nitpick: In the abstract, the wording \"experimental settings which range from neural machine translation (NMT) to multi-lingual language modeling\" seems a bit out of place. These are exactly the two settings on which the proposed approach is being evaluated, there aren't any other experimental settings that fall within these two.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and positive response. Our responses to the weaknesses and questions brought up are below:\n\n**W1:** *I would have loved to see some more realistic MT experimental settings. Pretraining on en-fr for either en-zh or en-hi (or vice versa) feels very unrealistic. There are much closer language pairs that could have been used, which would have made the results more convincing.*\n\n**Response:** In our MT experiments, we tried to choose task pairs with varying degrees of transfer-ability from En-Fr (Ro is more semantically closer to Fr than Hi and Zh are to Fr). If we only chose pairs that transferred well, it would be unclear whether our method showed improvement only because of the transferability, and whether it will work in settings where the tasks are not similar. We do agree that having more experiments in the realistic setting would make our arguments stronger.\n\n**W2:** *I am concerned about the results of the experiments of section 3.1.1. What's the effect on BLEU scores here?*\n\n**Response:** The key point we emphasize in section 3.1.1 is that for a high-resource task, the best performance achievable is bottlenecked by the amount of data seen for the task. The best En-Zh performance for pre-training (300k) joint fine-tuning (300k) is worse than the best En-Zh performance for joint training (600k) since the latter sees more En-Zh data than is possible for the former. At the same time, the best En-Zh performance for pre-training (300k) joint fine-tuning (300k) is equal to the best En-Zh performance seen by joint training (300k) which saw the same amount of En-Zh data. \n\nThe implications of these results is that when there is abundant data available for a task, and overfitting is not an issue, it is important to train on it as much as possible, and that pre-training on a different task cannot act as a proxy for more data on that task.\n\nLastly, we generated the BLEU score plot for En-{Zh,Fr} (see pdf attached in our global response). From the figure, we can see that the BLEU score version shows similar conclusions as the cross-entropy version, but with a more positive result– best performance for En-Zh is on-par with joint training for 600k steps.\n\n**Q:** *Line 128 refers the reader to Figures 21 and 22 in the appendix to show that the proposed approach does not hurt performance in the high-high case. It's not clear to me how the said figures show this.*\n\n**Response:** We apologize for the confusion. When we said the proposed approach does not hurt performance, we meant that our method does not worsen the best possible En-Zh performance achievable **given the same amount of En-Zh data** available for the model to see.\n\nElaborating on our response to W2, let’s imagine the Pareto front of pre-training on En–Fr for 300k steps and then joint training on En-{Zh, Fr} for 300k steps, and compare it to just joint training for 300k steps. The pareto front should be improved for the En-Fr task performance since it saw a lot more of the En-Fr task data. For the En-Zh task performance, we can expect one of 3 possibilities:\n1. The front is improved due to positive transfer despite seeing the same amount of En-Zh data\n2. The front is unchanged\n3. The front is worsened due to pre-training on En-Fr possibly resulting in worse initialization and/or slow training for En-Zh.\n\nFigure 1 shows that case 2 is indeed what happens– the best En-Zh performance for pre-training (300k) joint fine-tuning (300k) is equal to the best En-Zh performance seen by joint training (300k). Pre-training on En-Fr cannot act as a proxy for more En-Zh data, because if it could, then the front would be improved. At the same time, pre-training also does not negatively impact En-Zh training– Figures 21 and 22 show that pre-training does not affect the learning efficiency for En-Zh (slope of the curves are similar to one another), and also did not result in a worse initialization for En->Zh.\n\nWe hope that our response and our newly generated BLEU score plot (see pdf in global author response) addressed some of the reviewer’s concerns. We will update our writing for section 3.1.1 to make this point more clear.\n\n**Regarding the minor fixes**, thank you for the suggestion and for catching the typos– we will make the changes to the final version. \n"}, {"review_id": "3CKCTNRL2t", "reviewer": "Reviewer_V1qx", "summary": "The paper proposes a method called pretraining and joint-finetuning, which pretrains a model on a high-resource task than finetunes the model with joint high- and low-resource tasks, benefiting from both static sampling method and naïve transfer learning. The method is verified on multilingual translation and multilingual language modelling. Experiments show improvement on the two scenarios in term of validation loss.", "questions": "1. Definition of Pareto optimal points. According to the paper, “θ is Pareto optimal if it is not dominated by any other point”. However, for a parameter family for a certain architecture, it is impossible to test all parameters in the whole space. In this case, is it rigorous to say that a parameter θ is not dominated by any other parameters?\n2. Inadequate references. There are plenty of papers on the catastrophic forgetting phenomenon in multilingual/multitask learning. I think the claims in this paper might been seen from others. For example, the “order matters” claim has been found in “Overcoming Catastrophic Forgetting beyond Continual Learning: Balanced Training for Neural Machine Translation”. The authors should check previous literature more carefully.\n3. Scalarization implementation. In Line 72 of the paper, the static sampling is implementing by adding weights to the losses from different tasks (languages). Will the models trained with different proportions be fed different data? For example, suppose the En->Ro data is fixed, will a model with 0.1 En-Ro proportion be fed more En data than a model with 0.9 En-Ro proportion? If the answer is no (i.e., the data is the same, the only difference is the weights on losses), does the method have a same effect with sampling data in each batch?\n4. Why the Figures do not always show all results of the ten proportions (e.g., Fig. 2 Left)?\n\n\n\nMinor suggestions (mainly for presentation):\n- I think having a separate section about the proposed method will be better.\n- For figures of the loss that need a comparison (e.g., Fig. 3 and 8), the illustration will be clearer if the scale of the two subfigures is consistent.\n\n\n\nAfter Author Response: \nI strongly encourage the authors to address Questions 1 and 2, specifically the clarification of Pareto dominance and a more comprehensive discussion with previous literature. These enhancements would undoubtedly contribute to the rigor and solidity of the paper.", "limitations": "The authors consider an interesting possibility of multiple high-resource languages setting as a future work, which can be seen as a partition of tasks. However, they do not provide a set of hyperparameters effective across tasks, e.g., at least how many steps of joint-finetuning should be taken? Which sampling proportion should we use when training with new data/tasks?", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The proposed method is simple and easy to implement/reproduce.\n2. The method can achieve lower validation loss compared to static sampling and naïve transfer learning.\n3. The method verified on multilingual (> 2) setting and extremely resource-imbalanced (English is 16745 times larger than the smallest language) setting.", "weaknesses": "1. The algorithm is sensitive to the sampling proportion, resulting a need to search proportion for every model trained. In addition to searching the proportion of dataset size, there is a need for grid search of N1 and N2 (N1 + N2 = N). The hyperparameters (proportion = 0.4 for two tasks, tau = 3.33 for > 2 tasks, and joint-finetuning steps = 50k) are not verified among different tasks.\n2. Experiments are done only on multilingual setting, whose scope is a bit different from the title “multitask” learning. Is the proposed method still effective on multitask like NER + POS + sentiment analysis?\n3. The improvement on performance is not significant when the evaluation metric is BLUE (from Fig. 26).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. Here is our response:\n\n**W1a:** *The algorithm is sensitive to the sampling proportion, resulting a need to search proportion for every model trained.*\n\n**Response:**\nOur method is not any more sensitive to the sampling proportion than scalarization, which is our baseline. Furthermore, the best sampling proportions will always depend on the task at hand. Therefore, even for scalarization, ideally one would search for a good sampling proportion.\n\nIn our experiments, we follow two different schemes for testing sampling rates depending on what we aim to achieve. In our NMT experiments, the purpose of testing sampling rates in a grid was to investigate whether pre-training can improve the Pareto front of scalarization. We emphasize that we tested sampling rates in a grid for the scalarization baseline as well, so that the comparison is fair. In our language modeling experiments, our goal was to model reality as much as possible (by using temperature sampling), while still comparing against the best baseline. This is why we **tuned the temperature parameter for the scalarization baseline**, and used that same temperature for our method.\n\n**W1b:** *In addition to searching the proportion of dataset size, there is a need for grid search of N1 and N2 (N1 + N2 = N).* \n\n**Response:** Our method does indeed introduce an extra parameter to decide on compared to scalarization: scalarization only needs to decide on N, but we need to decide on N1 and N2. Regarding this, we would like to emphasize 3 things:\n- **We did not tune or perform a grid search in our experiments.** We chose values that were either the most beneficial to our baseline (scalarization), or seemed reasonable, and we continued using them without further tuning them.\n- **Parameters N1 and N2 are precisely how our method can do early stopping.** Let’s imagine using scalarization, and one of the tasks starts overfitting. It’s not ideal to do early stopping since the high-resource tasks can benefit from training longer. If we instead pre-train on the high-resource task, we can get a head-start on the high-resource task such that once the low-resource tasks start overfitting, we can stop training.\n- **Based on our experiments, we can recommend how one would decide on N1 and N2.** Our results show that pre-training doesn’t hurt performance, and the longer we pre-train for, the better performance we get. This means that if a practitioner is strictly bound with a compute budget, they should probably pre-train for only a fraction of their total budget to make sure that all tasks are able to converge. On the other hand, for a practitioner who mostly cares about best performance achievable, and does not have a strict computational budget, they should pre-train for as long as possible before joint fine-tuning.\n\n**W1c:** *The hyperparameters (proportion = 0.4 for two tasks, tau = 3.33 for > 2 tasks, and joint-finetuning steps = 50k) are not verified among different tasks.*\n\n**Response:** The performance of any given sampling rate depends on the task. Therefore, just like how it is unreasonable to believe that a single temperature parameter should be recommended for scalarization, we do not recommend a single proportion/temperature to be used for different tasks. We emphasize that we did not tune the sampling rate any more than the baseline scalarization.\n\nThe number of steps to fine-tune also depends on the task. Some tasks will take only 50k steps to converge (En-Ro, or En-Hi), while some tasks will take longer (as in the language modeling experiments).\n\n**W2:** *Experiments are done only on multilingual setting, whose scope is a bit different from the title “multitask” learning. Is the proposed method still effective on multitask like NER + POS + sentiment analysis?*\n\n**Response:** Please see to our global response.\n\n**W3:** *The improvement on performance is not significant when the evaluation metric is BLUE (from Fig. 26).*\n\n**Response:** Our improvement on En-{Ro, Fr} is indeed weak in terms of the BLEU score. The improvements are much better for En-{Hi, Fr}, so we expect the amount of improvement to vary depending on the task.\n\n**Q1:** *Is it rigorous to say that a parameter $\\theta$ is not dominated by any other parameters?*\n\n**Response:** You are right– we cannot rigorously guarantee that a given point is not dominated. We will update our writing to be more clear about this. However, we do believe that we did a reasonably good job of approximating the Pareto front by testing a grid of sampling rates and learning rates.\n\n**Q2:** *Inadequate references.*\n\n**Response:** Thank you for pointing this out, we will update our paper with more references on catastrophic forgetting.\n\n**Q3:** *Scalarization implementation. In Line 72 of the paper, the static sampling is implementing by adding weights to the losses from different tasks (languages). Will the models trained with different proportions be fed different data?*\n\n**Response:** We apologize for the confusion. We follow the convention in the NMT literature and implement scalarization via proportional sampling, where the average number of data corresponding to task i in a batch is proportional to w_i. So we do not change the loss function, only the way we sample the data. We will update our draft to include this point, thank you for pointing it out.\n\n**Q4:** *Why the Figures do not always show all results of the ten proportions (e.g., Fig. 2 Left)?*\n\n**Response:** This is because we are only plotting the points that were not dominated (considering only the space of parameters obtained through our grid searches). In the existence of low-resource tasks, the validation loss plot does not mirror the training loss plot like in the high-high case (Figure 1), and there are certain sampling rates that are more optimal than others.\n\nRegarding the minor suggestions, thank you for the feedback– we will update our draft to reflect them.\n"}, {"review_id": "MLTgtKZG6y", "reviewer": "Reviewer_evVX", "summary": "This paper presents a simple and effective multi-task learning strategy of a joint pretraining followed by fine-tuning, where pretraining is on the high-resource task and fine-tuning is on a mixture of high and low-resource tasks. This significantly improves performance on the low-resource tasks, while performing at par or even better sometimes on high-resource tasks. The authors show results on a machine translation task consisting of two high-resource and two low-resource language pairs and a language modeling task trained on five different languages.", "questions": "- In Figure 9a, pretraining longer on English seems to benefit Gujarati but not Hindi. Any thoughts on why this might be?\n\n- From line 260, \"Longer pre-training helps more in terms of performance than exposure to the tokens of the low-resource task\" -- This claim demands further probing. Is this dependent on the choice of high-resource language for pretraining? This claim holds for Gujarati but not Hindi. Are there features of the low-resource task (e.g., type-token ratio in the unlabeled text for the low-resource tasks) that makes transfer from the high-resource task more effective?\n\n- In the many-task setting, as in Section 3.2, would it be beneficial to introduce an intermediate pretraining stage between pretraining and finetuning? Such strategies have been shown to be effective in prior work (e.g., Phang et al., \" Sentence encoders on stilts: Supplementary training on intermediate labeled-data tasks\", 2018). Also, such works on intermediate labeled-data tasks should be cited in the related work section.\n\n- The authors should move up the BLEU results from the Appendix to the main draft. For NMT, it is important to show that validation loss reductions do translate to BLEU score improvements (especially for English->Hindi).\n\n- A suggestion: Since pretraining is usually associated with self-supervised objectives, it would be useful to explicitly clarify that the pretraining and finetuning objectives are both cross-entropy losses for the NMT task.", "limitations": "Limitations have been adequately addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed idea is attractive in its simplicity, and offers good validation loss reductions on translation and language modeling tasks.\n- The authors provide a fairly extensive empirical analysis of the proposed approach, and offer good baselines to compare against.\n- The paper is written clearly and the overall narrative flows well.", "weaknesses": "While the proposed approach has been empirically validated with multiple experiments, there remain unanswered questions about the initial pretraining on a high-resource language:\n- How important is the pretraining objective?\n- What is the influence of the chosen high-resource language on final performance on low-resource languages?\n- Why does pretraining longer improve the loss for high-resource NMT tasks but degrade the loss for high-resource language modeling tasks (c.f., Figure 9(a) and 9(b)).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback. We respond to the weaknesses and questions brought up below:\n\n**W:** *Why does pretraining longer improve the loss for high-resource NMT tasks but degrade the loss for high-resource language modeling tasks (c.f., Figure 9(a) and 9(b)).*\n\n**Response:** For all our experiments, when starting the fine-tuning phase of training we reset the cosine learning rate schedule with a warmup. Due to the warmup, the performance of the pre-trained high-resource task performance worsens a bit before improving again (see the panel corresponding to En for Figure 8). Because of this, the high-resource task needs enough training steps to recover its old performance, and to continue improving.\n\nIn the language modeling experiment (Figure 9(a)), we kept the total training budget the same, which meant increasing pre-training length resulted in decreased fine-tuning length. Less fine-tuning time means less time to recover for En, hence the worse performance.\nOn the other hand, for the NMT experiment (Figure 9(b)), we kept the fine-tuning length the same regardless of the pre-training length, which is why the En-Fr performance could keep improving.\n\nThe implications of these results is that when there is a strict training budget, it’s better to be conservative and pre-train for a shorter amount of time. However, if the goal is to obtain the best performance and there is no strict compute budget, it’s better to pre-train for as long as possible before fine-tuning. Note that longer overall training is an option for our method (by pre-training for longer) but not for scalarization because scalarization needs to always be training on the low-resource tasks, which will lead to overfitting when training for long.\n\nWe will update our final draft to be more clear about this.\n\n**Q1:** *In Figure 9a, pretraining longer on English seems to benefit Gujarati but not Hindi. Any thoughts on why this might be?*\n\n**Response:** In Figure 9a, as we pre-train for longer, we fine-tune for a shorter amount of time. In addition to this, Hindi is a high-resource task given our training budget. Therefore, overall lesser training for Hindi resulted in worse performance. Pre-training still benefits Hindi– Figure 8 shows that pre-training initialized Hi performance at a better position than random initialization, and resulted in more data-efficient training. \n\nAs an aside, we did not include Hindi in the pre-training phase because we wanted to keep our method simple. Including Hindi in the pre-training phase would require figuring out good sampling proportions not only for the fine-tuning phase, but for the pre-training phase. We can avoid under-training high-resource tasks by fine-tuning for longer.\n\n**Q2:** *\"Longer pre-training helps more in terms of performance than exposure to the tokens of the low-resource task\" -- This claim demands further probing. Is this dependent on the choice of high-resource language for pretraining? This claim holds for Gujarati but not Hindi. Are there features of the low-resource task (e.g., type-token ratio in the unlabeled text for the low-resource tasks) that makes transfer from the high-resource task more effective?*\n\n**Response:** We hope our response to Q1 helped answer a part of this question. In general, the relationship between the pre-training and fine-tuning tasks should have an impact on performance, although we are unsure what exactly about the relationship causes better transfer. We believe this is an important question that deserves separate investigation. In our work, we stuck to the simplest method and chose pre-training tasks (En-Fr for MT and En for language modeling) that are (and will always be) available at an abundance, therefore, are perfect to pre-train on.\n\n**Q3:** *In the many-task setting, as in Section 3.2, would it be beneficial to introduce an intermediate pretraining stage between pretraining and finetuning? Such strategies have been shown to be effective in prior work (e.g., Phang et al., \" Sentence encoders on stilts: Supplementary training on intermediate labeled-data tasks\", 2018). Also, such works on intermediate labeled-data tasks should be cited in the related work section.*\n\n**Response:** Thank you for the missed citation– we will include intermediate pre-training works in the related work section of our final draft.\n\nWe have not tested with intermediate training in our context. In Phang et al., Intermediate training on a different but related task seemed to help in the monolingual (English) setting, which means it could possibly improve our method further. In our paper, we focus on cross-lingual transfer, but we believe that pre-training to utilize cross-task transfer is an interesting question to be studied in the future.\n\n**Q4:** *The authors should move up the BLEU results from the Appendix to the main draft.*\n\n**Response:** Thank you for the suggestion, we will move the BLEU results to the main section in our final draft.\n\n**Q5:** *A suggestion: Since pretraining is usually associated with self-supervised objectives, it would be useful to explicitly clarify that the pretraining and finetuning objectives are both cross-entropy losses for the NMT task.*\n\n**Response:** Thank you for the suggestion, we will clarify this in the final draft.\n\nWe hope that we clarified some of the reviewer’s concerns– we will update our final draft to make our points more clear. Please let us know if there are any further questions.\n"}]

Author rebuttal:
**Q:** *Reviewers wUUC and V1qx both pointed out that our method focused on the multilingual setup, and were curious whether we ran experiments on the general multi-task setup.*

**Response**: In our work, we frame multilingual learning (NMT, language modeling) as a multi-task optimization problem as in many prior work [1, 2, 3, 4, 5]. Using translation as an example, learning to translate to/from multiple languages, and learning multiple “tasks” (e.g. question answering, and summarization) can both be seen as learning multiple functions at once.

Since our work focuses on utilizing cross-lingual transfer, we have yet to conduct experiments on the “general multi-task setup” (multiple NLP tasks that may or may not be multilingual), where the focus would be on utilizing cross-task transfer. We will update our future work section to reflect this.

[1] Luong et al, 2015, Multi-task Sequence to Sequence Learning

[2] Firat et al, 2016, Multi-Way, Multilingual Neural Machine Translation with a Shared Attention Mechanism

[3] Arivazhagan et al, 2019, Massively Multilingual Neural Machine Translation in the Wild: Findings and Challenges

[4] Jean et al, 2019, Adaptive Scheduling for Multi-Task Learning

[5] Wang et al, 2020, Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models

**Lastly**, we attach BLEU score plots for En-{Zh,Fr} as requested by Reviewer 9ZBe.

Meta-review: This paper investigates the strategies on how to best deal with data imbalance in multilingual learning. Their proposed method enables high-resource tasks to converge faster and prevents overfitting for low-resource tasks. The method is clearly described and evaluated on both NMT and language modeling. The paper would benefit from a more comprehensive discussion of prior work, further clarification of pareto dominance, additional evaluation on XQuAD or XTREME as pointed out by reviewers, and evaluation on a wider range of language for NMT.
Decision: Accept (poster)