SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dami Choi*†
U. Toronto & Vector Institute
choidami@cs.toronto.edu
Derrick Xin*
dxin@google.com
Hamid Dadkhahi
hdadkhahi@google.com
Justin Gilmer
gilmer@google.com
Ankush Garg
ankugarg@google.com
orhanf@google.com
Chih-Kuan Yeh
chihkuanyeh@google.com
Andrew M. Dai
adai@google.com
Behrooz Ghorbani
ghorbani@openai.com
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.
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-the-art performance in domains such as machine translation [2, 8], language understanding [24, 32], and speech recognition [4, 3].
Unfortunately, optimizing such multi-task models remains a challenge. To effectively train these models, the different tasks need to be balanced during training. This is often done by sampling each task with a static probability.
Prior work [31, 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 a student researcher at Google.
(low-resource) $^2$ , which is quite common in real-world applications, the optimality of static sampling is 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 a viable solution due to high-resource 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 in [33]) provides a 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 on a high-resource task and fine-tune jointly on a mixture of high and low-resource tasks. A pre-training and fine-tuning scheme effectively enables early stopping by allowing the training of low-resource tasks to happen for as 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.
Through an 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 a 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-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-efficient 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.
- 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' performance in multilingual models.
- We show on realistic workloads (up to 13B parameters) that our scheme performs better than static sampling and transfer learning with respect to the low-resource language-pair/language.
In our work, we focus on the supervised setting, where our model parameters  $\pmb{\theta} \in \mathbb{R}^p$  are trained on  $K$  different tasks, with the loss for task  $i$  being  $\mathcal{L}_i(\pmb{\theta})$ .
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).  $\pmb{\theta} \in \mathbb{R}^p$  Pareto dominates another  $\pmb{\theta}'$  if  $\forall 1 \leq i \leq K$ ,  $\mathcal{L}_i(\pmb{\theta}) \leq \mathcal{L}_i(\pmb{\theta}')$  and there exists a task  $j$  where  $\mathcal{L}_j(\pmb{\theta}) < \mathcal{L}_j(\pmb{\theta}')$ .  $\pmb{\theta}$  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-resource respectively. Note that whether a task is high-resource or not depends on both the amount of training data and the model capacity.
A standard approach for optimizing multi-task models is scalarization [5] or static sampling:
$$
\hat {\boldsymbol {\theta}} (\boldsymbol {w}) = \arg \min  _ {\boldsymbol {\theta}} \sum_ {i = 1} ^ {K} \boldsymbol {w} _ {i} \mathcal {L} _ {i} (\boldsymbol {\theta}), \tag {1}
$$
where  $\pmb{w}$  is a fixed vector of pre-determined task weights with  $\pmb{w} > 0$  and  $\sum_{i}\pmb{w}_{i} = 1$ .
In our work, we follow convention and implement scalarization via proportional sampling, where data from task  $i$  is sampled with probability equal to  $\pmb{w}_i$ . In this case, the expected loss is equal to the loss from scalarization:
$$
\mathcal {L} (\boldsymbol {\theta}) = \mathbb {E} _ {\boldsymbol {x}} [ \ell (\boldsymbol {x}; \boldsymbol {\theta}) ] = \sum_ {i = 1} ^ {K} \mathbb {P} (\text {t a s k} i) \mathbb {E} _ {\boldsymbol {x} \sim \text {t a s k} i} [ \ell (\boldsymbol {x}; \boldsymbol {\theta}) ] = \sum_ {i = 1} ^ {K} \boldsymbol {w} _ {i} \mathcal {L} _ {i} (\boldsymbol {\theta}). \tag {2}
$$
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  $\pmb{\theta}$  is truly Pareto optimal since we must check that it is not dominated by all  $\pmb{\theta}' \in \mathbb{R}^p$ . Following [31], instead of considering all of  $\mathbb{R}^p$  we consider only the parameters reachable by a fixed set of hyperparameters.
3 Pre-training Joint Fine-tuning
Given  $K$  tasks, among which some are 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 are 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 on 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 a 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<sup>3</sup>.
In our preliminary experiments, we found that it is important to 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 to 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 it is easier to tune. The pre-training phase only involves tuning for a 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-training does not hurt fine-tuning performance and longer pre-training translates to better fine-tuning. From this, we recommend that if 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 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.
3We use the terms '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.
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 a language). In the NMT experiments, we show that pre-training joint fine-tuning pushes past the trade-off frontier of scalarization through significant improvements on the low-resource task-- a feat that many popular gradient-based multi-task 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.
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-training. As in prior work [31], we choose to work on the two-task setting due to the ease of visualizing the performance trade-off curves.
We choose our high and low-resource language-pairs from the WMT dataset, where English  $\rightarrow$  {Chinese, French} are the high-resource language pairs, and English  $\rightarrow$  {Romanian, Hindi} are the low-resource language pairs. See Table 1 for details on each language-pair. All models in this section use a 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 in A.1.
In order to trace out the trade-off frontiers for the pre-training joint fine-tuning method and the scalarization baseline, we adhere to the following methodology. For scalarization, we iterate through a grid of task weights (since there are only two tasks, a grid is a linear function of the granularity) and train on the two language pairs for  $N$  steps using proportional sampling according to the task weights. For the pre-training joint fine-tuning method, we first pre-train on the high-resource language pair for  $N_{1}$  training
Table 1: Overview of data sources used in 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,981
En-Ro '16	610,320	1,999
En-Hi '14	313,748	520
steps. We then reset the optimizer state and the learning rate schedule and fine-tune on a mixture of high-resource and low-resource language pairs for  $N_{2}$  training steps such that  $N_{1} + N_{2} = N$ . For the fine-tuning phase, we iterate through a 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.
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  $N_{1}$  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-Resource:
We first start by highlighting that pre-training joint fine-tuning does not show benefits if all tasks are high-resource. Figure 1 shows that in the English  $\rightarrow$  {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  $\rightarrow$  Fr cannot act as a proxy for En  $\rightarrow$  Zh data, because if it could, the front would be improved. At the same time, pre-training does not negatively impact En  $\rightarrow$  Zh training. Figures 21 and 22 show that pre-training does not affect the learning efficiency for En  $\rightarrow$  Zh (slope of the curves are similar to one another), and also did not result in a worse initialization for En  $\rightarrow$  Zh.
4.1.2 High-Resource and Low-Resource
In the data-imbalanced setting of English  $\rightarrow$  {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 a single full run of scalarization ( $N$  steps) and pre-training fine-tuning ( $N_{1} + N_{2} = N$ ) take the same
Figure 1: 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 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 of the restart baseline shows that the resetting of states is not why pre-training and fine-tuning performs well. Note that the 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 a 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 hyperparameter 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_{2} \ll N$ .
In Figure 2 we can observe that pre-training joint fine-tuning is able to achieve performance tradeoff points that go beyond what is achievable via scalarization. Pre-training on a 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, it is able to do this while seeing less overall Romanian tokens according to Figure 3.
We see similar results for  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$ , shown in Figure 12 in the Appendix. This is a surprising result since French and Hindi are less linguistically similar than French and Romanian. Finally, we can see from the sub-optimal performance of the restart baseline in Figures 2 and 12 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  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{Fr}\}$  and  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$  in Appendix A.5, validating that the improvements in loss translate to downstream metrics.
Figure 3: Pre-training joint fine-tuning has both better initialization and data-efficiency than scalarization. Each line corresponds to the datapoint that achieved the best  $\mathrm{En} \rightarrow \mathrm{Ro}$  validation loss in Figure 2 among the different run groups.
Figure 4: Each curve corresponds to a single scalarization trial with a particular (static) sampling rate for  $\mathrm{En}\rightarrow \mathrm{Ro}$ . The rate at which the training loss decreases is slower for lower  $\mathrm{En}\rightarrow \mathrm{Ro}$  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 a better starting point than random initialization. Figure 3 shows this effect for the  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{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 a certain point, where overfitting kicks in.
By design, pre-training joint fine-tuning has two separate training phases which allows the low-resource-training phase to be short. This in turn enables us to increase the low-resource sampling rate, resulting in faster training. This effect can be seen in Figure 2, where the  $\mathrm{En}\rightarrow \mathrm{Ro}$  sampling rates that resulted in the best  $\mathrm{En}\rightarrow \mathrm{Ro}$  performance was 0.4, while for pre-training joint fine-tuning, the best rate is 0.5. Figure 3 confirms that indeed after pre-training, fine-tuning on  $\mathrm{En}\rightarrow \mathrm{Ro}$  is more data-efficient than not pre-training.
Joint fine-tuning is also an important piece in addition to 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 as shown in Figure 6.
Lastly, Figure 2 shows that pre-training joint fine-tuning yields worse training set performance, and therefore, could be seen as having a 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.
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 setting to a many-task setting. We train on five languages from the mC4 dataset [32]–English, Hindi, Gujarati, Swahili, and Gaelic– using the span corruption objective from T5 [24]. See Table 2 for details on the dataset. Canonically the mC4 dataset is used in the pre-training phase for models (not to be confused by our pre-training joint fine-tuning method). These models are subsequently applied to downstream tasks such as question answering. This multilingual pretraining phase is also known as 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.
Figure 6: Fine-tuning solely on the low-resource task  $(\mathrm{En} \rightarrow \mathrm{Ro})$  leads to both catastrophic forgetting of the pre-trained task  $(\mathrm{En} \rightarrow \mathrm{Fr})$  and worse low-resource task performance than fine-tuning on all tasks  $(\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{Fr}\})$ .
Note that in the mC4 corpus, English is 16745 times larger than the smallest language we use. This data imbalance underscores the necessity for effective language balancing, particularly 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 in a two-task setting.
Table 2: Data used from mC4.
Language	#Chars (B)
En (English)	13,396
Hi (Hindi)	75
Gu (Gujarati)	3.6
Gd (Gaelic)	0.8
Sw (Swahili)	4.1
For our training setup we closely follow mT5 [32] for the
model architecture and training procedure. Specifically, we use the mT5-XXL model (13B parameters), which is 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, we employ the widely used temperature sampling heuristic [11, 7, 2]. Let  $D_{i}$  be data size of language or task  $i$ , we then define the empirical distribution  $\mathbb{P}$  for each task  $i$  as:
$$
\mathbb {P} (\boldsymbol {x} \in \text {t a s k} i) = \frac {D _ {i}}{\sum_ {j} D _ {j}}. \tag {3}
$$
Figure 8: Pre-training on English and joint fine-tuning on all 5 languages leads to better optima for Gujarati, Gaelic and Swahili, the 3 low-resource languages. Pre-training also results in better initialization and token-efficiency for all languages newly seen in the fine-tuning phase.
Temperature sampling then uses a distribution  $\mathbb{Q}$  defined by a temperature parameter  $\tau$  as follows:
$$
\mathbb {Q} (\boldsymbol {x} \in \operatorname {t a s k} i) = \frac {\mathbb {P} (\boldsymbol {x} \in \operatorname {t a s k} i) ^ {1 / \tau}}{\sum_ {j} \mathbb {P} (\boldsymbol {x} \in \operatorname {t a s k} j) ^ {1 / \tau}} \tag {4}
$$
The temperature parameter  $\tau$  controls the peakiness (or flatness) of the sampling distribution. Commonly used  $\tau$ 's in the literature are greater than 1, which essentially up-samples low-resource tasks and down-samples high-resource tasks.