ExT5: Towards Extreme Multi-Task Scaling for Transfer Learning

1 INTRODUCTION . Transfer learning ( Schmidhuber , 1987 ; Pratt et al. , 1991 ; Caruana et al. , 1995 ) has been the cornerstone of recent progress in natural language processing ( Ruder et al. , 2019 ; Devlin et al. , 2019 ; Raffel et al. , 2020 ) . While self-supervised pre-training has been shown to be highly effective at exploiting large amounts of unlabeled data without relying on human annotation , there is still much to explore regarding transfer learning in a multi-task co-training setup . Prior seminal works like T5 ( Raffel et al. , 2020 ) and MT-DNN ( Liu et al. , 2019a ) have demonstrated a degree of promise in the paradigm of multi-task co-training ( Caruana , 1997 ) . However , the challenge of catastrophic forgetting remains . Tasks often have to be carefully selected in order to demonstrate positive affinity with regards to downstream transfer . In many cases , it is not unreasonable to expect negative transfer ( Rosenstein et al. , 2005 ; Vu et al. , 2020 ) . This makes the process of empirically curating a set of tasks to include in a transfer learning setup both computationally prohibitive and specific to downstream tasks . While standard pre-training typically employs a variant of the self-supervised language modeling objective ( Raffel et al. , 2020 ) , certain types of skills such as commonsense knowledge are only acquired at a slow rate even using massive amounts of unlabeled data ( Zhang et al. , 2021 ) . As ever larger models are trained , the development of much more sample-efficient pre-training settings becomes thus more important , and could be addressed via multi-task learning . For the first time , we explore and propose Extreme Multi-task Scaling — a new paradigm for multitask pre-training . Compared to the largest prior work ( Aghajanyan et al. , 2021 ) , our study doubles the number of tasks and focuses on multi-task pre-training rather than fine-tuning , which enables a direct comparison to standard pre-training . Our proposal is based on the insight that despite negative transfer being common during fine-tuning , a massive and diverse collection of pre-training tasks is generally preferable to an expensive search for the best combination of pre-training tasks . ∗Google AI Resident . †Equal contribution . Sebastian is now at Google Research . Sanket returned to CMU . To this end , we introduce EXMIX : a massive collection of 107 supervised NLP tasks to be included in a multi-task pre-training setup . We process all tasks in an encoder-decoder friendly format to readily support the sharing of all parameters across all tasks . We postulate that an ensembling effect across as many tasks , distributions and domains as possible results in a consistently net-positive outcome . This echoes early multi-task learning results ( Caruana , 1997 ; Baxter , 2000 ) indicating that a bias that is learned on sufficiently many tasks is likely to generalize to unseen tasks drawn from the same environment . Moreover , our experiments verify that our EXMIX mixture outperforms a best-effort mixture of manually curated tasks . Finally , we propose EXT5 : a T5 model ( Raffel et al. , 2020 ) pre-trained on a mixture of supervised EXMIX and self-supervised C4 span denoising . EXT5 outperforms state-of-the-art T5 models on well-established benchmarks such as SuperGLUE ( Wang et al. , 2019a ) , GEM ( Gehrmann et al. , 2021 ) , and Rainbow ( Lourie et al. , 2021 ) ; as well as Closed-Book QA ( Roberts et al. , 2020 ) tasks . Notably , our experimental findings also suggest that including EXMIX may reduce the number of pre-training steps required to achieve strong performance , bringing about substantial sample efficiency benefits . To summarize , this paper contributes the following : – We propose EXMIX ( §2 ) : a collection of 107 supervised NLP tasks for Extreme Multi-task Scaling , formatted for encoder-decoder training . EXMIX has approximately twice as many tasks as the largest prior study to date ( Aghajanyan et al. , 2021 ) , totaling 18M labeled examples across diverse task families . – Given this large collection of tasks , we conduct rigorous empirical studies evaluating transfer between common task families ( §2.1 ) . Our experiments show that curating a pre-training mixture based on fine-tuning transfer is not straightforward ( §2.2 ) . Hence , efficiently searching for the best subset of tasks to include in a multi-task pre-training setup is challenging and prohibitive . – Using EXMIX , we pre-train a model alongside the C4 span-denoising objective introduced by Raffel et al . ( 2020 ) , resulting in a new pre-trained model which we call EXT5 ( §3 ) . EXT5 outperforms state-of-the-art T5 on well-established benchmarks such as SuperGLUE , GEM , Rainbow , Closed Book Question Answering , and several other tasks that are outside of EXMIX ( §3.2 ) , while also being more sample-efficient ( §2.6 ) . 2 THE EXMIX TASK COLLECTION To explore the Extreme Task Scaling paradigm , we introduce EXMIX ( Extreme Mixture ) , a collection of 107 diverse English NLP tasks totaling 18M examples . Following Raffel et al . ( 2020 ) , we format all tasks as text-to-text examples to readily allow for multi-task training . This unified format also enables simple implementations without the need for taskspecific heads/losses , loss scaling , or explicit gradient accumulation for heterogenous batches as in prior works ( Liu et al. , 2019a ; Aghajanyan et al. , 2021 ) . When selecting examples from EXMIX , ex- amples from each dataset are sampled proportionate to the individual dataset ’ s size , with each dataset ’ s sampling rate capped at 3 × 105 maximum effective examples to ensure a balance between large and small datasets . We refer readers to Appendix A for a comprehensive breakdown of EXMIX . Additionally , we discuss future multilingual variants of EXMIX and EXT5 in §5 . 2.1 TRANSFER RELATIONS BETWEEN EXMIX TASKS . As discussed in §1 , our goal is to pre-train a model on EXMIX to improve downstream performance . One natural question to ask is which tasks have a negative impact on downstream performance ? Specifically , is there a subset of EXMIX that leads to better representations when used for multitask pre-training ? Obviously , testing all possible 2|EXMIX| combinations is impractical since the pre-training process is expensive to run . Instead , we experiment with the less expensive co-training procedure ( i.e. , multi-task fine-tuning ) , using representative subsets of similar tasks . Later , in §2.2 , we explore whether these results can be used to inform the task selection for multi-task pre-training . To study transfer amongst task-families in EXMIX , we construct subsets ( Table 1 ) of 3 tasks each that are partitioned along their task family . Using these subsets , we evaluate transfer among task families in a multi-task learning ( co-training ) setup . While other types of tasks are available in EXMIX , we did not include them because they were not diverse enough to be representative of a task family , and would scale the number of models needing to be trained at a quadratic rate . Experimental Setting We fine-tune a model on each pair of task families ( i.e. , 6 datasets at a time ) . To ensure a fair balance of tasks , we sample tasks proportional within their family , but uniformly between task families . For example , while evaluating how classification tasks and NLI tasks transfer amongst each other , the sampling ratio of MNLI : ANLI will be proportional ( approximately 2.4:1 ) , but the overall ratio of NLI examples to classification examples will be 1:1 . For reference , we also train a model on each individual task family using proportional example mixing ( Sanh et al. , 2019 ) . In total , this results in F+ ( F 2 ) models trained , where F is the number of task families . Our experiments use F = 8 as shown in Table 1 , resulting in 34 models trained in total . Each model is fine-tuned on top of the released T5.1.1BASE checkpoint for 200k steps using a batch size of 128 and a constant learning rate of 10−3 . Observations We summarize our results in Table 2 . We observe that although there exist particular task-family pairs that show positive transfer ( e.g. , co-training with NLI helps most other tasks ) , negative transfer is more common when training across task families compared to intra-family training . 21 out of the 56 inter-family relationships perform worse than intra-family models with the same data budget , which grows to 38 out of 56 for a fixed compute-budget . While the abundance of negative transfer among diverse task families is an expected result , interesting trends manifest in the individual relationships . For example , summarization tasks generally seem to hurt performance on most other task families ; and CBQA tasks are highly sensitive to multi-task fine-tuning . We also report correlations for intra-family datasets in Figure 2 using the same models as in Table 2 . In most cases , we see positive correlations between datasets in the same family . In a few cases , however , we observe an opposite trend . For example , fine-tuned models that performed better on the GEM schema-guided dialog dataset achieved lower scores on KILT Wizard-of-Wikipedia . This initial exploratory analysis highlights both the potential of EXMIX as a tool to systematically study task relationships , as well as the potential challenges in leveraging multi-task learning naively on top of pre-trained representations . 2.2 CAN FINE-TUNING TASK RELATIONSHIPS HELP CURATE A PRE-TRAINING MIXTURE ? . Our observations in §2.1 showed that multi-task co-training on top of existing pre-trained checkpoints is not straightforward , and often results in negative transfer . However , the uncovered task relationships might help efficiently search for an ideal subset of EXMIX for multi-task pre-training . To this end , we select a set of the most promising task families to be included in a multi-task pretraining setup , ranking task families by the average relative improvement they provide to other target families ( the last column in Table 2 ) . Specifically , we include NLI , commonsense , classification , and closed-book QA tasks from EXMIX to form a mixture of 48 tasks to include in a multi-task pre-training setup . We then fine-tune the resulting model on SuperGLUE , comparing it to T5.1.1 and EXT5 in Table 3 . While this model narrowly outperformed T5.1.1 , it did not yield better results than including all of EXMIX in the multi-task pre-training mixture , as we report in §3.2 . Moreover , it did not outperform a random selection of 55 tasks on average , as we report in our task scaling experiments ( §2.5 ) . We conclude that the negative transfer during multi-task fine-tuning does not necessarily inhibit pre-training . While we can not directly conclude that an ideal subset of EXMIX does not exist to be mixed with self-supervised pre-training for specific downstream tasks , our experiments show that randomly including more diverse pre-training tasks generally improves downstream performance . It must also be noted that the end-goal is to find a mixture that leads to a general pre-trained model that can be used for a large variety of downstream tasks , and that a setup to find a pre-training mixture tailored for SuperGLUE would be different . 2.3 MULTI-TASK PRE-TRAINING VS PRE-FINETUNING . Instead of pre-training on EXMIX , another way to leverage multi-task learning is as an intermediate step between pre-training and fine-tuning . This is referred to as pre-finetuning by Aghajanyan et al . ( 2021 ) . We conduct controlled experiments to compare pre-training with pre-finetuning . We begin with a standard T5 base checkpoint and pre-finetune it with EXMIX . After this phase , we fine-tune on SuperGLUE . Table 4 compares pre-finetuning and our proposed multi-task pre-training . We also report the total compute ( in total number of tokens processed ) by the model in both schemes . The results show that multitask pre-training is significantly superior to prefinetuning . A potential hypothesis is that multi-task pre-training narrows the gap between pre-training and finetuning data distributions , as the pre-training stage more closely resembles fine-tuning . Conversely , segregating pre-training and pre-finetuning into two different stages may induce catastrophic forgetting of the pre-training task . Hence , in EXT5 , we opt for multi-task pre-training over pre-finetuning .

Summary: This paper introduces "ExMix," a collection of 107 pre-existing NLP datasets across 8 identified task-families. ExMix is proposed as a pre-training dataset, to be mixed with the typical self-supervised pre-training prevalent in NLP. The authors propose to pre-train a T5 model on a mixture of ExMix and C4 (the self-supervised pre-training dataset for T5), which they term ExT5. To do this, they convert every task to a text-to-text problem, so that all tasks, supervised and self-supervised, use standard sequence-to-sequence cross-entropy loss. They pre-train ExT5 for the same number of steps and on the same number of tokens as T5. They then evaluate how ExT5 and T5 perform when fine-tuned to individual tasks. They find that ExT5 largely outperforms T5 on the majority of downstream evaluated tasks, including tasks whose training data was included in ExMix as well as tasks which were not seen during pre-training.  The authors perform an experiment to identify whether or not a manually curated subset of ExMix can provide better pre-training generalization than the full set of tasks, due to negative transfer between tasks. To do this, they examine negative transfer between pairs of task families during fine-tuning. They select the task-families which exhibited the highest amount of fine-tuning transfer across all task-pairs, and use only 48 tasks from these task-families as a pre-training objective. They compare this model to a model trained on a slightly larger subset (55) of randomly selected tasks, as well as ExT5 (trained on all 107) tasks. The authors find that their method of task-selection is largely unhelpful for selecting pre-training tasks - the model trained on a random subset of 55 tasks outperforms the manually curated 48-task model, and ExT5 significantly outperforms both. This suggests that, for pre-training, selecting tasks that complement one-another with respect to downstream performance is not straight-forward, but increasing the number of tasks seems to alleviate negative transfer, rather than exacerbate it.  Finally, the authors perform some additional analysis of the training decisions made around ExT5. They find that: (a) pre-finetuning, which introduces multi-task learning as an intermediate step between self-supervised pre-training and fine-tuning, is not better than mixing all tasks for pre-training when using ExT5, (b) ExMix without self-supervised pre-training is significantly worse than self-supervised training without ExMix, suggesting that self-supervised pre-training is still crucial for pre-training, (c) downstream performance generally increases as the number of tasks scales up, and (d) pre-training with ExMix is significantly more sample-efficient than self-supervised pre-training along, i.e. it achieves strong downstream results very early, even 20k steps into pre-training.