SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Static Sampling Baseline Temperature sampling is ubiquitous due to its simplicity and intuitiveness, but its performance varies greatly with  $\tau$ . For our static sampling baseline, we tuned  $\tau$  among commonly used values in the literature (1.43, 2, 3.33, 5) at a smaller scale, and found that  $\tau = 3.33$  performed the best in terms of low-resource languages. We also tried a more intricate sampling strategy called UniMax [6], but found that on the 5 languages we chose, it did not perform better than  $\tau = 3.33$ .
Figure 7: Pre-training joint fine-tuning yields the best performance in 4 out of 5 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 5 languages using temperature sampling. We use the same sampling rates as the static sampling baseline ( $\tau = 3.33$ ) to reduce the tuning overhead over static sampling.
As in the NMT experiments, we employ a restart baseline to fully ablate the 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 7 and 8 show that while a learning rate schedule restart helps performance, pretraining joint fine-tuning yields the best results on the low-resource tasks. Surprisingly, it not only improves the performance on Gujarati, Gaelic, and Swahili, but also shows a slight enhancement on English. We note that due to the vast dataset imbalance, the temperature sampling baseline overfits on the low-resource tasks before English has a chance to converge. Consequently, pre-training joint fine-tuning can leverage the benefits mentioned in the previous section-regularization, transfer, and reduced forgetting-to achieve a superior lower bound performance with higher token efficiency.
(b)
Figure 9: Left: For language modeling on mC4, longer pre-training leads to better best-achievable performance for the 3 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 on 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 to use a smaller model (mT5-XL as opposed to mT5-XXL 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 9(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 8). 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 a low-resource task in this context, and therefore, less tokens seen translates to worse performance.
In Figure 9(b) we see that in the NMT experiments, pre-training longer on  $\mathrm{En} \rightarrow \mathrm{Fr}$  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 is 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 to overfitting when training for too long.
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. In the language domain, multilingual neural machine translation [12, 14] 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-resource language-pair 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 a 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, in 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 on a high-resource language pair can 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-pairs has been long known in the NMT literature [33, 9, 17]. [33] showed that pre-training on a 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 [21, 15], 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 [1], and experimenting with many-stage pre-training [9]. We expect to further benefit from research done in this direction.
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 earlier in the training, and gradually samples more low-resource 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 are 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 C. 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 a 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 a function of the relationship between pre-training and fine-tuning tasks, is left to future work.
In this work, we demonstrated the benefits of a 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 under-samples the low-resource task due to higher token efficiency. We also show that fine-tuning only on the low-resource task, a popular scheme in the NMT literature, is undesirable due to its inability to prevent forgetting. Our method is a 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 a standard approach, particularly in the era of large language models.
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, we thank Chris J.H. Zhang for helpful discussions.
[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 neural machine translation in the wild: Findings and challenges. arXiv preprint arXiv:1907.05019, 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 First, Sharan Narang, and Noah Constant. Unimax: Fairer and more effective language sampling for large-scale multilingual pretraining. In The Eleventh International Conference on Learning Representations.
[7] Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Unsupervised cross-lingual representation learning at scale. arXiv preprint arXiv:1911.02116, 2019.
[8] Marta R Costa-jussà, James Cross, Onur Çelebi, Maha Elbayad, Kenneth Heafield, Kevin Heffernan, Elahe Kalbassi, Janice Lam, Daniel Licht, Jean Maillard, et al. No language left behind: Scaling human-centered machine translation. arXiv preprint arXiv:2207.04672, 2022.
[9] Raj Dabre, Atsushi Fujita, and Chenhui Chu. Exploiting multilingualism through multistage fine-tuning for low-resource neural machine translation. In 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, 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 31st Pacific Asia conference on language, information and computation, pages 282-286, 2017.
[11] Jacob Devlin. Multilingual bert readme. https://github.com/google-research/bert/blob/master/multilingual.md, 2018.
[12] Orhan First, Kyunghyun Cho, and Yoshua Bengio. Multi-way, multilingual neural machine translation with a shared attention mechanism. arXiv preprint arXiv:1601.01073, 2016.
[13] Sébastien Jean, Orhan First, and Melvin Johnson. Adaptive scheduling for multi-task learning. arXiv preprint arXiv:1909.06434, 2019.
[29] Xinyi Wang, Yulia Tsvetkov, and Graham Neubig. Balancing training for multilingual neural machine translation. arXiv preprint arXiv:2004.06748, 2020.
[30] Yiren Wang, ChengXiang 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 First. 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 Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mT5: A massively multilingual pre-trained text-to-text transformer. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: 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 arXiv:1604.02201, 2016.
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 in the main text 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.1 label smoothing and 0.1 dropout for feed-forward and attention layers. See Table 3 for complete architecture details.
Table 3: Transformer architecture details and common hyperparameters.
Hyperparameter
Feed-forward dim	2048
Model dim	512
Attention heads	8
Attention QKV dim	512
Label smoothing	0.1
Dropout	0.1
We use SentencePiece tokenization [19] to generate a vocabulary of size 64,000 for each NMT problem (e.g.  $\mathrm{En} \rightarrow \{\mathrm{Zh}, \mathrm{Fr}\}$ ).
All models were trained using the Adam [16] optimizer with a batch size of 1024. For all our NMT experiments, we used a linear warmup to the desired learning rate, followed by a cosine decay schedule that decays to 0. This is true for all legs of training for methods that use our scheme; during the pre-training phase, we do a linear warmup followed by a 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  $40\mathrm{k}$  steps. For fine-tuning, we tune the warmup steps from within  $\{10\mathrm{k}, 20\mathrm{k}, 30\mathrm{k}, 40\mathrm{k}\}$  for all experiments other than for  $\mathrm{En} \rightarrow \{\mathrm{Zh}, \mathrm{Fr}\}$ , where we warmup for  $40\mathrm{k}$  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-pre-training model for 'base + fine-tune' number of steps.
Table 4: Number of training steps for all NMT experiments.
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, we sweep the base learning rate in the grid  $\{2.5\mathrm{e - 4},5\mathrm{e - 4},2.5\mathrm{e - 3},5\mathrm{e - 3},7.5\mathrm{e - }$  3}. We also sweep the sampling rate for  $\mathrm{En}\rightarrow \mathrm{Fr}$  and  $\mathrm{En}\rightarrow \mathrm{Cs}$  in the grid  $\{i / 10\}_{i = 1}^{9}$ , which fully determines the sampling rate for the other language pair. All plotted points correspond to the final measurement taken for each trial.
For all fine-tuning experiments, when loading the pre-trained model checkpoint, we reset the optimizer state. We also trained all parameters of the model, and did not freeze anything.
A.2 Additional Performance Trade-Off Curves
In this section, we present the performance trade-off curves for  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$ , as well as for 6-layer models on  $\mathrm{En} \rightarrow \{\mathrm{Zh}, \mathrm{Fr}\}$ ,  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{Fr}\}$ , and  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$ . The black-bordered points in the generalization portion of Figures 11 and 13 below correspond to the restart baseline.
Figure 10: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Zh}, \mathrm{Fr}\}$  with 6-layer models. Each point corresponds to the final performance of a model. Similarly to the 3-layer-model case (Figure 1), pre-training does not yield improvements.
Figure 11: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{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  $\mathrm{En} \rightarrow \mathrm{Ro}$  due to a larger model size.
Figure 12: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$  with 3-layer models. These results mirror those seen in Figure 2. We note that here French and Hindi are more linguistically dissimilar than French and Romanian.
Figure 13: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$  with 6-layer models. As with the 3-layer models, We observe a similar improvement in both  $\mathrm{En} \rightarrow \mathrm{Hi}$  and  $\mathrm{En} \rightarrow \mathrm{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 lower-resource language pair. We can see that in all but one case  $(\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\})$  6-layer model; Figure 19), 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 16, 17, 18, and 19 below correspond to the restart baseline.
Figure 14: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Zh}, \mathrm{Fr}\}$  with 3-layer models. We can clearly see that there is no optimal rate in this case, since we trace a Pareto front as we vary the  $\mathrm{En} \rightarrow \mathrm{Zh}$  sampling rates from 0.1 to 0.9.
Figure 15: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Zh}, \mathrm{Fr}\}$  with 6-layer models. We observe a similar behavior as in the 3-layer case.
Figure 16: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{Fr}\}$  with 3-layer models. Unlike the  $\mathrm{En} \rightarrow \{\mathrm{Zh}, \mathrm{Fr}\}$  case, we have a few sampling rates that are more optimal than the rest. Pre-training allows sampling  $\mathrm{En} \rightarrow \mathrm{Ro}$  at a higher rate without overfitting, than without pre-training.
Figure 17: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{Fr}\}$  with 6-layer models. We see a similar behavior as in the 3-layer case.
Figure 18: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$  with 3-layer models. Like in the  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{Fr}\}$ , pre-training allows sampling  $\mathrm{En} \rightarrow \mathrm{Hi}$  at a higher rate without overfitting than without pre-training.
Figure 19: Performance trade-off behavior for  $\mathrm{En} \rightarrow \{\mathrm{Hi}, \mathrm{Fr}\}$  with 6-layer models. In this case, pre-training still allows sampling  $\mathrm{En} \rightarrow \mathrm{Hi}$  at a higher rate, but the rate that yielded the best  $\mathrm{En} \rightarrow \mathrm{Hi}$  was surprisingly the same rate as the baseline without pre-training.
A.4 Efficiency Plots
In this section, we plot the number of examples seen from one language pair against the validation cross-entropy loss on that language pair. The number of XX  $\rightarrow$  YY examples seen at train step  $t$  is computed by multiplying  $t$ , the batch size, and the sampling rate for XX  $\rightarrow$  YY. Each curve in a given figure corresponds to the trial that achieved the best final validation performance on the low(er)-resource language pair within the method given by the legend (i.e. the blue curve in Figure 20 corresponds to the trial that achieved the best final validation  $\mathrm{En} \rightarrow \mathrm{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.
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.
Figure 20: For the 3-layer model, pre-training does not provide any significant gains in training efficiency for  $\mathrm{En} \rightarrow \mathrm{Zh}$  when pre-training on  $\mathrm{En} \rightarrow \mathrm{Fr}$ . Given that the blue and red curves coincide towards the end of training, we can anticipate that pre-training did not impair  $\mathrm{En} \rightarrow \mathrm{Zh}$  training (by providing a suboptimal initialization), and that if we were to train the red curve for  $300\mathrm{k}$  more steps, it would be able to catch up with the orange curve (best  $\mathrm{En} \rightarrow \mathrm{Zh}$  performance).
Figure 21: We observe a similar behavior with 6-layer models as with 3-layer models.
Figure 22: On the 3-layer models, pre-training is able to accelerate training on  $\mathrm{En} \rightarrow \mathrm{Ro}$  when pretrained on  $\mathrm{En} \rightarrow \mathrm{Fr}$ . Even with less overall examples seen in  $\mathrm{En} \rightarrow \mathrm{Ro}$ , we can perform better than the baselines that did not use pre-training.
Figure 23: We observe a similar efficiency boost with 6-layer models as with 3-layer models.
Figure 24: On the 3-layer models, we observe a similar efficiency boost as with  $\mathrm{En} \rightarrow \{\mathrm{Ro}, \mathrm{Fr}\}$
Figure 25: On the 6-layer models, we observe a 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 Beam-Search with beam size of 4.