SlowGuess/NAACL_2022_OCR huggingface dataset:
Vishakh Padmakumar  $^{1*}$  Leonard Lausen $^{2}$  Miguel Ballesteros $^{3}$  Sheng Zha $^{2}$  He He $^{12}$  George Karypis $^{2}$
$^{1}$ New York University,  $^{2}$ AWS AI,  $^{3}$ AWS AI Labs
vishakh@nyu.edu
{lausen, ballemig, zhasheng, hehea, gkarypis}@amazon.com
Recent work has found that multi-task training with a large number of diverse tasks can uniformly improve downstream performance on unseen target tasks. In contrast, literature on task transferability has established that the choice of intermediate tasks can heavily affect downstream task performance. In this work, we aim to disentangle the effect of scale and relatedness of tasks in multi-task representation learning. We find that, on average, increasing the scale of multi-task learning, in terms of the number of tasks, indeed results in better learned representations than smaller multi-task setups. However, if the target tasks are known ahead of time, then training on a smaller set of related tasks is competitive to the large-scale multi-task training at a reduced computational cost.
Following the wide success of unsupervised language model pre-training (Devlin et al., 2019; Liu et al., 2019b; Lewis et al., 2019), recent work on transfer learning has shown that additional supervised multi-task training further improves performance on various downstream NLP tasks (Raffel et al., 2019; Khashabi et al., 2020; Aghajanyan et al., 2021). There are two distinct ways in which the supervised data has been used: increasing the scale of the multi-task step to incorporate more tasks (Khashabi et al., 2020; Aghajanyan et al., 2021) and developing task similarity metrics to incorporate tasks related to the target task (Pruksachatkun et al., 2020; Vu et al., 2020).
Aghajanyan et al. (2021) show that a multi-task training step, or pre-finetuning step, with a sufficiently large, diverse set of tasks is an effective task-agnostic second stage of model pre-training before finetuning on target tasks. In particular, they find that using a large number of tasks (e.g., roughly 15
tasks) is crucial in achieving good downstream performance, while pre-finetuning with fewer tasks causes a small performance drop. Meanwhile, work on task transferability (Vu et al., 2020; Pruk-sachatkun et al., 2020) has shown that the choice of individual intermediate tasks significantly affects downstream fine-tuning performance—predicting the transfer is challenging and there is high variance depending on the choice of intermediate task.
This motivates us to ask the question if prefinetuning on a small group of tasks related to the target tasks can obtain comparable performance to large-scale multi-task training. In this work, we present an empirical study to answer this question by extending the task transferability experiment to groups of tasks.
We follow the two-step experimental pipeline from Aghajanyan et al. (2021), where a pre-trained model is first pre-finetuned on a set of tasks and then separately finetuned on various target tasks on which we report performance. In addition, we group our set of 29 pre-finetuning tasks based on task format into 3 groups—classification tasks, sequence labelling tasks and extractive question answering tasks (Figure 1). We perform model pre-finetuning on every combination of these task groups and report performance on target tasks that belong to each group. This allows us to systematically study how the size of the multi-task step and the choice of pre-finetuning tasks affects downstream task performance.
We observe that, on average, large-scale multi-task pre-finetuning results in improved performance on downstream target tasks. We also see that a model trained on related $^{1}$  pre-finetuning tasks obtains comparable downstream task performance to the large-scale model, at a reduced computational cost $^{2}$ , but pre-finetuning on an unrelated grouping
<sup>1</sup>In the rest of this work, when we say that two tasks are related it means that they belong to the same task group
2We say that a pre-finetuning run is of a cheaper than
* Work done during summer internship at AWS AI.
Human Language Technologies, pages 2542 - 2550
Figure 1: List of tasks grouped into sequence labelling (SL), question answering (QA) and classification (C).
Figure 2: High level workflow of the multi-task setup.
can result in a severe decline in performance.
Our findings show the interplay between multi-task scaling and task selection—when the target tasks are unknown then multi-task scaling is an effective, if expensive, intermediate step but if the goal is to improve performance on a specific target task set then multi-task training on a smaller set of related tasks is an effective cheaper alternative. These results also hint at the need to better study modeling techniques that mitigate negative transfer between tasks in large-scale multi-task training.
2 Multi-task Setup
Our multi-task experiments follow the two step approach from Aghajanyan et al. (2021), shown in Figure 2. A pre-trained model is first pre-finetuned on a set of tasks to obtain a single shared encoder. This model is then finetuned separately on various target tasks on which we report performance.
Grouping of Tasks Figure 1 shows our list of 29 pre-finetuning tasks and 12 unseen target tasks (listed in Appendix A). Informed by prior work (Vu et al., 2020; Ye et al., 2021; Sanh et al., 2021), we divide these datasets into three groups based on task format—sequence labelling (SL), extractive question answering (QA) and classification (C). As noted in Sanh et al. (2021), grouping of NLP tasks is an imperfect heuristic. Prior work (Achille et al., 2019; Vu et al., 2020) formalizes the notion of task similarity using learned task embeddings so an
another in terms of computational cost when it involves multitask training on fewer examples (Appendix B.3)
alternate formulation would be to divide the tasks into groups based on these learned embeddings. In this work we focus on a simple, intuitive grouping based on the task format, or output space, of the task.
Research Question We aim to study how the choice of pre-finetuning tasks and the size of the multi-task step, in terms of number of pre-finetuning tasks, affects target task performance. In order to do so, we compare pre-finetuning runs on all combinations of task groups, reporting performance on target tasks from each group.
3.1 Model Details
We use the pre-trained XLM-Roberta model (Connéau et al., 2020) for all of our experiments. During pre-finetuning, we learn a shared encoder for all tasks and a task-specific head for each pre-finetuning task. For downstream finetuning, we randomly initialize a new head for each target task. We use the Huggingface (Wolf et al., 2020) XLM-Roberta base pre-trained model. The various task specific heads are linear classifiers on the encoder output per the Huggingface implementation. More model details are provided in Appendix B.1.
During pre-finetuning, to ensure that the model does not overfit to any particular task, we follow the sampling approach from Aghajanyan et al. (2021) of ensuring that each batch consists of examples from multiple tasks while maintaining the empirical distribution of examples in the datasets. For our study, the loss function for all the different pre-finetuning tasks is cross entropy loss. Aghajanyan et al. (2021) recommends scaling the losses from each task-specific head based on the size of
the label space to ensure more stable training. Our preliminary experiments showed better results without loss scaling so we follow the same for all prefinetuning runs.
We run pre-finetuning on the powerset of the set of task groups from Figure 1 and for each pre-finetuning run we report performance on all the target tasks. Hyperparameters are kept uniform across pre-finetuning runs and we train with batch size 128 and an early stopping criteria based on validation loss. Once pre-finetuning is completed, we finetune the model on each target task and report the average validation set performance across 5 random seeds along with the associated standard deviation. More details on pre-finetuning and target task finetuning are included in Appendix B.2.
Notation We refer to the pre-finetuning runs with their initials—Only-QA means that only question answering tasks were used in pre-finetuning the model and QA+C means that all question-answering and classification tasks were used and so on. We have 3 task groups and run experiments on all 7 possible combinations—Only-SL, Only-QA, Only-C, QA+C, QA+SL, SL+C and SL+QA+C. The baseline is a pre-trained model directly finetuned on the target tasks.
Table 1 shows the results of each pre-finetuning run on all the target tasks. Table 2 contains the results aggregated by task group.
A large-scale target task-agnostic Multi-task step improves downstream performance From Table 1, we see that the multi-task setup containing all 29 pre-finetuning tasks  $(\mathrm{SL} + \mathrm{QA} + \mathrm{C})$  has the best average performance over all the target tasks as well as the best score in 5 individual tasks. This is consistent with observations reported in Aghajanyan et al. (2021) that increasing the scale of the multi-task step results in better downstream performance on average across all tasks. Our results show the same trend on a different set of tasks with a smaller batch size. We also show that this observation holds in a standard multi-task training regime, without the optimization tricks used
Aghajanyan et al. (2021), namely loss scaling and regularized finetuning (Aghajanyan et al., 2020).
Related tasks transfer better To identify the role of transferability, we aggregate the results on target tasks based on task groups in Table 2. Each row in Table 2 is the average of all the unseen tasks within that group from Table 1. From this, we see that the Only-SL and Only-QA setups are on-par with  $\mathrm{SL + QA + C}$  on unseen SL and QA tasks respectively indicating that pre-finetuning on a smaller set of related tasks obtains comparable performance to the large-scale multi-task model. We also see that selecting a mismatched set of pre-finetuning tasks significantly hurts downstream task performance. From Table 2, we see a drop of  $9.6\%$  compared to the baseline on SL tasks with Only-C pre-finetuning and a  $20.3\%$  drop on QA tasks with Only-SL. These results extend those observed in Pruksachatkun et al. (2020) and Vu et al. (2020) to transferability across task groups. With appropriate task selection, we can obtain comparable performance to the large-scale multi-task model at a reduced computational cost. The reduction in computational cost is mainly due to the change in number of pre-finetuning examples. We provide the comparison over various runs in Appendix B.3. Aghajanyan et al. (2021) reported that multi-task learning is detrimental to target task performance at a smaller scale ( $< 15$  tasks). In our study, we see that pre-finetuning on a single group of related tasks always outperforms the results from the mismatched pairwise setup—Only-QA outperforms  $\mathrm{SL + C}$  on unseen QA tasks, Only-SL outperforms  $\mathrm{QA + C}$  on SL tasks and Only-C outperforms  $\mathrm{SL + QA}$  on C tasks. Hence we conclude that at smaller scales, the particular pre-finetuning tasks selected significantly impacts downstream task performance, linking back to transferability literature.
Tasks interact differently, so selecting an optimal subset is hard When we look at prefinetuning runs on pairs of task groups taken together, we see that the  $\mathrm{SL + C}$  and  $\mathrm{QA + C}$  prefinetuning setups perform worse on QA and SL tasks than even the baseline model but the  $\mathrm{SL + QA}$  setup is competitive with the best  $\mathrm{SL + QA + C}$  setup across all unseen tasks. This shows that selecting an optimal combination of tasks can be challenging based on task group heuristics. Aribandi et al. (2021) also observed a similar result that, at larger scales, a random subset of tasks often outperforms
3For all sequence labelling target tasks, we report F1 score. Among classification tasks, we report F1 for CB and accuracy for the other tasks. For extractive question answering, we report F1. These choices was made based on the standard evaluation for the task from the Huggingface metrics module.
Dataset	Baseline	Only-SL	Only-C	Only-QA	SL+C	SL+QA	QA+C	SL+C+QA
WNUT17	63.239 0.33	63.124 0.65	43.461 1.85	60.623 1.85	59.729 1.25	64.261 0.79	57.560 3.62	64.359 0.57
Wikigold	80.409 1.43	83.449 0.16	72.174 8.05	78.173 1.04	81.403 0.51	83.114 0.57	76.023 1.21	81.635 0.99
NCBI Disease	86.849 0.46	87.069 0.55	85.412 0.59	86.456 0.86	87.037 0.43	87.458 0.32	86.642 0.70	87.378 0.73
MIT Movie	90.042 0.53	89.736 0.20	88.436 0.68	89.522 0.24	89.645 0.27	89.829 0.24	89.287 0.20	89.631 0.21
BoolQ	73.804 1.15	70.761 0.72	78.313 0.48	77.140 0.71	78.457 0.50	76.427 0.81	80.147 0.27	79.515 0.51
CB	86.333 4.20	82.121 5.17	87.161 1.07	68.930 10.65	84.336 4.20	85.003 3.92	80.638 10.67	89.565 2.16
Copa	52.166 4.09	53.500 2.56	55.333 3.39	51.833 3.76	53.666 3.20	53.666 2.49	54.333 1.59	57.333 1.24
WiC	60.136 7.45	63.244 1.01	64.812 0.51	54.937 6.56	65.229 1.45	65.177 0.61	65.151 1.35	65.674 1.05
Squad Adv	47.064 2.42	22.341 2.38	55.067 0.83	82.798 0.51	52.035 1.05	81.558 0.80	82.778 0.61	81.834 0.88
SubjQA	60.917 0.74	58.577 0.47	61.741 0.29	61.889 0.31	60.718 0.10	62.367 0.22	62.578 0.47	62.886 0.56
QED	37.779 5.49	36.544 1.16	51.604 1.34	74.643 0.44	46.533 1.21	77.196 0.51	73.910 1.42	76.312 0.82
XQuad-R	63.522 3.66	50.398 1.83	64.538 0.54	80.744 0.50	63.654 0.26	79.624 0.76	78.249 1.04	80.163 0.51
Average	68.312	64.475	68.298	72.385	69.147	75.564	74.160	76.483
Average Std. Dev.	2.662	1.406	1.637	2.285	1.202	1.003	1.929	0.852
Table 1: Results on all the target tasks (rows) for all the pre-finetuning schemes (columns). Each cell value is an average on 5 runs with different seeds and the corresponding subscript is the standard deviation over these values. We also report the average over all tasks and the average of the standard deviation values in separate rows for analysis. We observe the effect of scale by seeing that the  $\mathrm{SL + QA + C}$  setup has the best average performance across all tasks. We see that multi-task training results in reduced variability across multiple runs.
Baseline	Only-SL	Only-C	Only-QA	SL+C	SL+QA	QA+C	SL+C+QA
Unseen SL	80.134	80.844	72.370	78.693	79.453	80.165	77.378	80.750
Unseen C	68.109	67.406	71.404	63.21	70.422	70.068	70.067	73.021
Unseen QA	56.692	45.174	61.120	75.252	57.568	75.460	75.035	75.678
Average	68.312	64.475	68.298	72.385	69.147	75.564	74.160	76.483
Table 2: Results on all unseen tasks aggregated by task format (rows) for each pre-finetuning setup (columns). Each value in this table is an average of the 4 unseen tasks of that particular task format from Table 1. We see the effect of transferability where the Only-SL and Only-QA setups are competitive with  $\mathrm{SL + QA + C}$  on unseen SL and QA tasks but suffer significantly on mismatched task groups.
subsets selected using simple heuristics.
Multi-task training reduces the variability of downstream task performance In Table 1, we also report the standard deviation in target task performance across 5 random restarts. We see a trend that large-scale multi-task pre-finetuning reduces the variability across runs on all tasks—the  $\mathrm{SL} + \mathrm{QA} + \mathrm{C}$  setup has the lowest average and the pairwise setups average lower variation than the single task group setups. Phang et al. (2018) also reported similar findings that multi-task training reduces variability in performance across random restarts. Additionally we observe that the Only-SL, Only-C and Only-QA setups have lower variability on unseen tasks of the same group than other groups, indicating that the downstream performance is more reliable on tasks within the same group. We discuss some limitations of our setup in Appendix C
Large-Scale Multi-task Learning Post the wide success of unsupervised language model pretraining, Phang et al. (2018) showed that intermediate task training on large datasets results in performance improvements on the GLUE bench
mark. Liu et al. (2019a) showed an improvement over standard pre-training on multiple NLP benchmarks in the multi-task setting. T5 (Raffel et al., 2019) framed various NLP tasks in a text-to-text format and subsequent work (Khashabi et al., 2020; Paolini et al., 2021) and sequence labelling showed that adapting T5-style models to particular domains results in powerful multi-task models. Aghajanyan et al. (2020) found that increasing the scale of a multi-task pre-finetuning step results in uniform improvement across various unseen tasks. Recent work in prompting large LMs has also shown that multi-task training can improve zero-shot performance (Wei et al., 2021; Sanh et al., 2021). Ye et al. (2021) showed that the few-shot performance on unseen tasks can be improved via a supervised multi-task step and recommended further analysis on task similarity and transferability. Our work aims to address this gap and connect large-scale multi-task learning to work on transferability.
Exploring Relationships Between Tasks Wang et al. (2019a) and Pruksachatkun et al. (2020) performed extensive empirical studies to identify the most beneficial intermediate tasks that improve target task performance both yielding mixed results. Changpinyo et al. (2018) observed that jointly
learning 11 sequence tagging tasks, with task embeddings, improves performance in around half of them and that the learned task embeddings revealed interesting task relationships such as clusters of semantic and syntactic tasks. Kiperwasser and Ballesteros (2018) showed that learning to perform syntactic tasks such as dependency parsing and part-of-speech tagging along with translation in varying schedules improves translation performance—they used task embedding vectors to identify the tasks to the decoder model. Indeed the idea of using identifying tokens as language embeddings is known to improve translation (Johnson et al., 2017) and dependency parsing (Ammar et al., 2016) predates the widespread adoption of transformer models for these tasks.
More recently, Vu et al. (2020) proposed two methods to learn task embeddings capable of predicting transferability between source and target tasks—one by pooling the representations of the textual task data from BERT and the other by using the layer-wise gradients of a BERT model. Vu et al. (2021) learned task specific prompts that can benefit each other via prompt transfer. These works largely identify the single most suitable task for each target task, we extend the same to groups of tasks. Our work most closely relates to a concurrent study, Aribandi et al. (2021), that examined the transfer across various task families. Our results compliment theirs using a different base model—they use a T5 style formulation of tasks, we use a shared Roberta encoder approach, showing that the transferability phenomenon is independent to the model architecture. We differ from their work in that we compare transfer on individual task groups with pairs of task groups as well and present results on the variance of performance as a result of multi-task learning.
In this work, we bring together the lines of exploration on transferability of tasks and large-scale multi-task training. Our results show that when the target tasks are unknown then multi-task scaling offers an effective way to obtain good downstream performance but if the goal is to improve performance on a specific target task set then a smaller set of related tasks is an effective, cheaper alternative. We observe that task groups interact differently when combined and that selecting an optimum subset becomes harder as the size increases. We also
see that variability across multiple random restarts decreases on related target tasks and also reduces on increasing the size of the multi-task step.
We would like to thank Stefano Soatto for providing feedback on this draft as well as the anonymous reviewers for their helpful comments.
Alessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran, Subhransu Maji, Charless C Fowlkes, Stefano Soatto, and Pietro Perona. 2019. Task2vec: Task embedding for meta-learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6430-6439.
Armen Aghajanyan, Anchit Gupta, Akshit Shrivastava, Xilun Chen, Luke Zettlemoyer, and Sonal Gupta. 2021. Muppet: Massive multi-task representations with pre-finetuning. arXiv preprint arXiv:2101.11038.
Armen Aghajanyan, Akshit Shrivastava, Anchit Gupta, Naman Goyal, Luke Zettlemoyer, and Sonal Gupta. 2020. Better fine-tuning by reducing representational collapse. CoRR, abs/2008.03156.
Julio Cesar Salinas Alvarado, Karin Verspoor, and Timothy Baldwin. 2015. Domain adaption of named entity recognition to support credit risk assessment. In Proceedings of the Australasian Language Technology Association Workshop 2015, pages 84-90.
Waleed Ammar, George Mulcaire, Miguel Ballesteros, Chris Dyer, and Noah A Smith. 2016. Many languages, one parser. Transactions of the Association for Computational Linguistics, 4:431-444.
Vamsi Aribandi, Yi Tay, Tal Schuster, Jinfeng Rao, Huaixiu Steven Zheng, Sanket Vaibhav Mehta, Honglei Zhuang, Vinh Q. Tran, Dara Bahri, Jianmo Ni, Jai Gupta, Kai Hui, Sebastian Ruder, and Donald Metzler. 2021. Ext5: Towards extreme multi-task scaling for transfer learning.
Dominic Balasuriya, Nicky Ringland, Joel Nothman, Tara Murphy, and James R Curran. 2009. Named entity recognition in wikipedia. In Proceedings of the 2009 Workshop on The People's Web Meets NLP: Collaboratively Constructed Semantic Resources (People's Web), pages 10-18.
Johan Bos, Valerio Basile, Kilian Evang, Noortje J Venhuizen, and Johannes Bjerva. 2017. The grongingen meaning bank. In Handbook of linguistic annotation, pages 463-496. Springer.
Soravit Changpinyo, Hexiang Hu, and Fei Sha. 2018. Multi-task learning for sequence tagging: An empirical study. arXiv preprint arXiv:1808.04151.