Source-linked AI summary
Muppet: Massive Multi-task Representations with Pre-Finetuning
Armen Aghajanyan, Anchit Gupta, Akshat Shrivastava, Xilun Chen, Luke Zettlemoyer, Sonal Gupta
TL;DR
The paper addresses whether supervised multi-task learning can reliably improve representations after language-model pre-training without selecting specific transfer tasks. It introduces large-scale pre-finetuning with balanced heterogeneous training and finds consistent downstream gains, especially in low-resource fine-tuning, while task count determines whether representations degrade or improve.
Problem
The effect of multi-task learning after language-model pre-training is not fully understood, and it can be difficult to know which intermediate tasks transfer best.
Method
Pre-finetuning inserts a massive multi-task learning stage between pre-training and fine-tuning, using loss scaling and task-heterogeneous batches across around 50 tasks.
Results
Pre-finetuned RoBERTa and BART models consistently outperform vanilla pretrained counterparts across several tasks, with particularly strong gains on smaller datasets.
Takeaways & Limitations
Large-scale multi-tasking can refine more robust representations, but fewer tasks may degrade them until a critical task-count point is reached.
Abstract
from arXiv · showhide
We propose pre-finetuning, an additional large-scale learning stage between language model pre-training and fine-tuning. Pre-finetuning is massively multi-task learning (around 50 datasets, over 4.8 million total labeled examples), and is designed to encourage learning of representations that generalize better to many different tasks. We show that pre-finetuning consistently improves performance for pretrained discriminators (e.g.~RoBERTa) and generation models (e.g.~BART) on a wide range of tasks (sentence prediction, commonsense reasoning, MRC, etc.), while also significantly improving sample efficiency during fine-tuning. We also show that large-scale multi-tasking is crucial; pre-finetuning can hurt performance when few tasks are used up until a critical point (usually above 15) after which performance improves linearly in the number of tasks.
1 Introduction
The paper introduces pre-finetuning, a large-scale supervised multi-task stage between language-model pre-training and task-specific fine-tuning. It shows that scale, balanced training, and heterogeneous tasks improve representations, downstream performance, and fine-tuning efficiency.
- 1 Introduction: Pre-finetuning adds a massive multi-task stage between standard pre-training and fine-tuning, using around 50 tasks and 4.8 million training examples.The tasks span classification, summarization, question answering, and commonsense reasoning.
- 1 Introduction: Loss scaling and task-heterogeneous batches balance competing tasks, improving multi-task training stability and overall performance.The scheme addresses instability in standard multi-tasking and is designed for learning at scale.
- 1 Introduction: Pre-finetuning yields consistent improvements for RoBERTa and BART, including state-of-the-art results on RTE and HellaSWAG.The gains are particularly strong when downstream fine-tuning has relatively little labeled data.
- 1 Introduction: Pre-finetuning consistently improves pretrained representations across downstream tasks and requires less data during fine-tuning.The reported gains apply to both discriminative and generative pretrained models across multiple task types.
- 1 Introduction: Beyond a critical point, usually above 15 tasks, adding more tasks improves generalizable representations, whereas fewer tasks can degrade them.The observed improvement continues linearly with task count as far as the experiments scale.
2 Related Work
Prior work established mixed evidence about multi-task learning after language-model pre-training and often used smaller or narrower setups. This paper argues that much larger, more diverse multi-tasking can consistently improve representations across language benchmarks.
- 2 Related Work: MT-DNN reported gains from multi-task learning, whereas T5 found that adding multi-task learning to larger models did not improve standardized pre-training and fine-tuning.These differing results leave the effect of multi-task learning across pre-training methods unresolved.
- 2 Related Work: Khashabi et al. improved T5 through cross-dataset transfer across QA tasks, but used a smaller scale, a different batching strategy, and a seq2seq formulation.Their work focused solely on improving QA tasks, unlike this paper’s broader task mix.
- 2 Related Work: This paper studies multi-task learning at much larger scale and reports consistent improvements across several language benchmarks and domains.It contrasts with T5 by reporting benefits from a secondary multi-task stage.
3 Pre-Finetuning Through Massive Multitask Learning
Pre-finetuning adds a large-scale multi-task stage between pre-training and fine-tuning, using diverse tasks and balancing methods to stabilize representation learning. The approach combines task-specific losses, heterogeneous batches, loss scaling, and robustness techniques while preserving natural dataset distributions.
- Training Challenges: Standard multi-task learning can struggle to balance competing task losses, especially when tasks are numerous and heterogeneous.Upsampling may overfit low-resource tasks, while downsampling may impair learning on specific tasks.
- Diverse Tasks: Pre-finetuning uses over 4.8 million supervised samples across classification, commonsense reasoning, machine reading comprehension, and summarization tasks.The setup is designed to learn general language representations across diverse domains.
- Multi-task Optimization: Task-specific heads optimize separate losses, whose gradients are averaged after loss scaling to produce each model update.This aggregates information from multiple task objectives into a unified representation.
- Multi-task Optimization: The training scheme combines heterogeneous batches with loss scaling so gradient contributions from competing tasks are more evenly balanced.The authors introduce this scheme to improve training stability and overall performance.
- Better Finetuning: Pre-finetuning initializes from self-supervised pretrained representations and uses R3F/R4F to encourage robustness to small input perturbations.R3F was pivotal for making MUPPET work with BART, while other training used standard SGD.
- Loss Scaling: Static loss scaling normalizes data-point losses so tasks would have equivalent values under uniform class distributions and model predictions.The method accounts for differing numbers of predictions, such as binary classes versus generation vocabularies.
- Dataset Sampling: Dataset sampling was consistently detrimental because it caused unmanageable overfitting and stability issues, so experiments retained natural dataset distributions.This choice avoided upsampling smaller datasets and downsampling larger ones.
4 Empirical Results
Pre-finetuning improves pretrained RoBERTa and BART representations across diverse downstream tasks, with especially strong gains on smaller datasets. The improvements include state-of-the-art results and extend beyond tasks included during pre-finetuning.
- Results Across Tasks: Pre-finetuned BART outperforms all other summarization baselines, while results also improve on sentence prediction and commonsense tasks.The reported tables include GLUE, machine reading comprehension, other sentence prediction tasks, commonsense tasks, and summarization.
- Finetuning Choices: Reusing pretrained task heads is generally beneficial for machine reading comprehension, commonsense, and sentence prediction tasks with small datasets.For SQuAD, the task head from pre-finetuning was reused.
- Overall Results: Across the board, pre-finetuned representations outperform standard pretrained representations, with larger gains on smaller fine-tuning datasets.The pre-finetuned RoBERTa-BASE model improves RTE by close to 9 points, while RoBERTa-LARGE reaches new state-of-the-art performance on RTE.
- Overall Results: Pre-finetuned models consistently outperform their base counterparts across task types and datasets, including the challenging ANLI dataset.The comparison covers multiple model variants and task families.
5 Understanding Multi-Task at Scale
The study examines how multi-task learning scale and batching strategy affect downstream representations. Performance can initially decline with few tasks, but improves after a task-count threshold, while heterogeneous batches support more effective training.
- Importance of Scale: Multi-task scale is evaluated by pre-finetuning RoBERTa models with progressively larger, nested sets of datasets.Seven models use increasing dataset counts, with each larger set containing the datasets from earlier points.
- Importance of Scale: Performance initially degrades in the low-task regime, then improves after a dataset-dependent critical point.The critical point is roughly between 10 and 25 tasks, while RTE improves monotonically with scale.
- Importance of Scale: The observed threshold suggests that earlier multi-task learning limitations may reflect insufficient scale rather than fundamental constraints.The paper reports that representations become more generalizable beyond the critical point.
- Importance of Heterogenous Batches: The batching study compares dataset-homogeneous, batch-homogeneous, and batch-heterogeneous sampling schemes.Batch-heterogeneous updates combine sub-batches from multiple datasets and tasks before aggregation.
- Importance of Heterogenous Batches: Batching strategy materially affects multi-task effectiveness, with heterogeneous batches producing the strongest reported performance.The comparison trains RoBERTa on 35 tasks and evaluates the resulting models on the same five downstream datasets.
- Low-Resource Fine-Tuning: Larger-scale pre-finetuning provides more data-efficient downstream fine-tuning than standard pre-training.At the 34/40 scale, models reach higher evaluation accuracies sooner than the base RoBERTa model.
6 Conclusion
The paper concludes that pre-finetuning refines pretrained representations through large-scale multi-task learning before end-task fine-tuning. Across several tasks and models, this approach improves performance, especially when fine-tuning data are limited, and depends on appropriate scaling and batching.
- Conclusion: Pre-finetuning adds a multi-task learning stage after pre-training and before end-task fine-tuning.The stage is intended to further refine pretrained representations.
- Conclusion: Loss scaling and heterogeneous batches are reported as critical for stable, effective multi-task learning at scale.The conclusion links these design choices to leveraging better representations.
- Conclusion: The batching comparison evaluates Dataset Homogeneous, Batch Homogeneous, and Batch Heterogeneous strategies across five datasets.The figure caption reports that heterogeneous batches outperform the other strategies by a significant margin.
- Conclusion: MUPPET improves over vanilla pretrained representations across MNLI, SQuAD, BoolQ, and Common Sense QA.The paper also reports improvements over prior state-of-the-art methods on RTE and HellaSWAG.
- Conclusion: Held-out low-resource experiments show that 34–40-task pre-finetuned models reach higher evaluation accuracies with much less data than RoBERTa.These experiments use datasets not included during pre-finetuning.
3. MRPC (Dolan and Brockett, 2005)
The listed multi-task datasets include mathematics, open-book question answering, commonsense reasoning, reading comprehension, and related tasks. The passages enumerate datasets used in the paper’s broader task collection.
- Task Collection: The collection includes Math QA and Openbook QA.These are listed as datasets 32 and 33, respectively.
- Task Collection: The collection includes SWAG, HellaSWAG, and RACE.These are listed as datasets 34 through 36.
- Task Collection: The collection includes CommonSense QA, Cosmos QA, and AI2 ARC - Easy.These are listed as datasets 37 through 39.
- Task Collection: The collection includes AI2 ARC - Challenge, SCIQ, and SQUAD.These are listed as datasets 40 through 42.
- Task Collection: The collection includes NQ, DROP, RECORD, Hotpot, and TriviaQA.These are listed as datasets 43 through 47.