Should We Be Pre-training? An Argument for End-task Aware Training as an Alternative

In most settings of practical concern , machine learning practitioners know in advance what end-task they wish to boost with auxiliary tasks . However , widely used methods for leveraging auxiliary data like pre-training and its continuedpretraining variant are end-task agnostic : they rarely , if ever , exploit knowledge of the target task . We study replacing end-task agnostic continued training of pretrained language models with end-task aware training of said models . We argue that for sufficiently important end-tasks , the benefits of leveraging auxiliary data in a task-aware fashion can justify forgoing the traditional approach of obtaining generic , end-task agnostic representations as with ( continued ) pre-training . On three different low-resource NLP tasks from two domains , we demonstrate that multi-tasking the end-task and auxiliary objectives results in significantly better downstream task performance than the widely-used task-agnostic continued pre-training paradigm of Gururangan et al . ( 2020 ) . We next introduce an online meta-learning algorithm that learns a set of multi-task weights to better balance among our multiple auxiliary objectives , achieving further improvements on endtask performance and data efficiency . 1 INTRODUCTION The increasingly popular pre-training paradigm ( Dai & Le , 2015 ; Devlin et al. , 2018 ; Gururangan et al. , 2020 ) involves first training a generalist model on copious amounts of easy-to-obtain data , e.g . raw text data in NLP , and then using this model to initialize training on a wide swath of downstream tasks . Generalist models like BERT ( Devlin et al. , 2018 ) , RoBERTa ( Liu et al. , 2019 ) , and GPT-3 ( Brown et al. , 2020 ) have a strong appeal ; a few institutions with significant resources incur the cost of training these large models whilst the rest of the research community enjoys a significant performance improvement at minimal computational overhead . However , the advantages of initializing a downstream task from a generalist model are not guaranteed . Previous work has shown that the benefits of pre-training depend heavily on the degree of domain overlap between the end-task data and the massive , heterogenous data on which the generalist model was trained ( Beltagy et al. , 2019 ; Gururangan et al. , 2020 ) . Notably , Gururangan et al . ( 2020 ) have demonstrated the benefits of continued pre-training of generalist models using data that is similar to that of the end-task . Their approach is formalized into two classes : Domain Adaptive Pre-training ( DAPT ) and Task Adaptive Pretraining ( TAPT ) where further stages of pre-training of generalist models are conducted on domain- and task-specific data , respectively . DAPT and TAPT exploit the fact that we often know the end-task beforehand , and so we can make specific choices about our pre-training regimen to improve end-task performance . However , in both pre-training for generalist models and continued pre-training , the training procedure itself does not explicitly incorporate the end-task objective function . Because of this , practitioners have to be careful with their choice of auxiliary tasks , the order in which they are trained on , and the early-stopping criteria for each pre-training stage so as to actually achieve good downstream end-task performance ( Gururangan et al. , 2020 ; Dery et al. , 2021 ) . In the absence of principled criteria to make these difficult design choices , it is common to instead resort to the computationally demanding heuristic of pre-training on as much data as possible for as long as possible . In this paper , we raise the following question : “ In settings where we have a particular end-task in mind , should we be pre-training at all ? ” . We define pre-training as any form of task-agnostic training that a model undergoes before it is finally fine-tuned on the end-task of interest . As a first milestone in addressing the larger question posed above , we explore the ubiquitous continued pretraining setting ( Gururangan et al. , 2020 ; Aghajanyan et al. , 2021 ) . Specifically , our paper questions the wisdom of having disjoint further pre-training then fine-tuning steps on a generalist model . In response , we advocate for an alternative approach in which we directly introduce the end-task objective of interest into the learning process . This results in a suite of end-task aware methods called TARTAN ( end-Task AwaRe TrAiniNg ) . Our formulations incorporate both unsupervised auxiliary objectives traditionally used in NLP pre-training ( such as masked language modeling as in Devlin et al . ( 2018 ) ) and the end-task objective , followed by an optional fine-tuning step on the end-task . We motivate TARTAN experimentally in the continued pre-training setting and based on this , we make the following contributions to the literature on leveraging auxiliary tasks and data : • In lieu of standard end-task agnostic continued pre-training , we suggest introducing the end-task objective into the training process via multi-task learning ( Caruana , 1997 ; Ruder , 2017 ) . We call this procedure Multi-Tasking end-Task AwaRe TrAiniNg ( MT-TARTAN ) ( Section 3.1 ) . MTTARTAN is a simple yet surprisingly effective alternative to task-agnostic pre-training . In Section 5 , we demonstrate that MT-TARTAN significantly improves performance and data efficiency over Gururangan et al . ( 2020 ) ’ s results . It also obviates the need for fickle hyper-parameter tuning through direct optimization of validation performance . • To allow more fine-grained control of the end-task over the auxiliary tasks , in Section 3.2 , we present an online meta-learning algorithm that learns adaptive multi-task weights with the aim of improving final end-task performance . Our META-learning end-Task AwaRe TrAiniNg ( METATARTAN ) allows us to robustly modulate between multiple objectives and further improves performance over MT-TARTAN . • A naive implementation of META-TARTAN based on first-order meta-learning analysis results in a sub-optimal algorithm that ignores all tasks except the end-task . We trace this problem to the use of a single model training head for computing both the end-task training loss and meta-objective ( end-task validation loss ) . To guard against this pathological solution , we introduce a separate model head for computing the meta-objective . In Section 3.3 , we justify this simple-to-implement fix and validate its practical efficacy in Section 5 . Our results suggest that TARTAN may be an attractive alternative to the continued pre-training paradigm , and further research into the place of pre-training in end-task aware settings is warranted . 2 FORMALIZING PRE-TRAINING AND CONTINUED PRE-TRAINING . Consider a dataset D = { ( xi , yi ) i∈ [ m ] } consisting of m labelled examples . We define a task as an objective function and dataset pair : T = { L ( · ) , D } . Mθ is a model parameterized by θ . The objective function L ( yi , Mθ ( xi ) ) evaluates how well a model prediction Mθ ( xi ) fits the true label yi , such as cross-entropy loss in the case of classification . Note that the task dataset , D , is typically decomposed into the sets ( Dtrain , Dval , Dtest ) . Dtrain is the set of examples used for model training whilst Dtest is used for final task evaluation . The validation set , Dval , is typically used for model selection but it is also frequently used in meta-learning to define the meta-objective – Lval . Given a specific end-task T ∗ , our aim is to improve performance on T ∗ ( as measured by the model loss on DtestT∗ ) by leveraging auxiliary tasks Taux = { T1 , . . . , Tn } . Note that we do not particularly care about the performance of any of the tasks in Taux . We are willing to sacrifice performance on Taux if it improves performance on T ∗ . From the perspective of model architecture , there are several ways to leverage Taux . We focus on the simple but widely-used parameter sharing setting . Here , all tasks share a model body θbody but each task Ti has its own head φi for prediction . We denote the head belonging to T ∗ as φ′ . Thus θ = [ θbody ; ( φ1 , . . . , φn , φ′ ) ] and θbody is reusable across new tasks . 2.1 PRE-TRAINING . Pre-training is when a model is first trained on Taux before performing a final fine-tuning phase on T ∗ . The motivation behind pre-training is that learning Taux first hopefully captures relevant information that can be utilized during training of T ∗ . This desire has led to the proliferation of generalist pre-trained models like BERT ( Devlin et al. , 2018 ) , RoBERTa ( Liu et al. , 2019 ) and GPT3 ( Brown et al. , 2020 ) that have been trained on copious amounts of data . Generalist models have been widely successful at improving downstream task performance when used as initilization . We can formalize the pre-training procedure as follows : θ0 = argminθ ( ∑ Ti∈Taux LTi ( θ ) ) ( 1 ) In Equation 1 , we seek a point θ0 that achieves minimal loss on the tasks in Taux . We hope that θ0 will be a good starting point for gradient descent on T ∗ . Let g ( θ0 ) represent the set of end-points of stochastic gradient descent on an initialization , θ0 . Stochastic gradient descent from the same initialization can produce different end-points due to differences in hyper-parameters like learning rate , batch size and order , as well as regularization strength . We can write the fine-tuning phase as : θ∗ = argmin { θ ∈ g ( θ0 ) } LT∗ ( θ ) ( 2 ) Note that pre-training is end-task agnostic : the pre-training Equation 1 occurs entirely before training on the end-task Equation 2 , and does not explicitly incorporate the end-task objective , T ∗ . Since there is no awareness of the end-task during pre-training it is important to carefully choose Taux so that pre-training actually results in improved performance on T ∗ ( Wang et al. , 2018a ) . For text data , past work has found left-to-right language modeling ( Peters et al. , 2017 ) and masked language modeling ( MLM ) ( Devlin et al. , 2018 ) to be good choices to include in Taux . 2.2 CONTINUED PRE-TRAINING . Recent work ( Beltagy et al. , 2019 ; Gururangan et al. , 2020 ; Lee et al. , 2020 ) showed that downstream performance on T ∗ can be improved by further adapting generalist models via continued pre-training on a more relevant set of auxiliary tasks . This is equivalent to sequentially performing multiple steps of Equation 1 , with different Taux , before finally performing Equation 2 on T ∗ . Domain and Task Adaptive Pre-training Gururangan et al . ( 2020 ) present Domain Adaptive PreTraining ( DAPT ) and Task Adaptive Pre-Training ( TAPT ) as methods for continued pre-training . During DAPT , a generalist model is further pre-trained on an unsupervised objective with large amounts of data from the same domain as the end-task . TAPT also pre-trains with the same unsupervised objective as DAPT , but on the actual dataset of the end-task . Gururangan et al . ( 2020 ) find that performance can be further improved by chaining objectives , DAPT first , followed by TAPT . Though TAPT and DAPT do not directly incorporate the end-task objective during training , it still indirectly informs both the choice of pre-training data and the order in which the pre-training tasks are trained on . Below , we explore stronger versions of this influence .

Summary: The paper makes the argument that generic pre-training (on auxiliary tasks) is inferior to task-specific pre-training. The authors argue that the final task should be learned together with the auxiliary tasks in a multi-task setup, which they call MT-TARTAN. They also propose a meta-learning algorithm META-TARTAN that uses meta-learning to mitigate the potential impact of updates from auxiliary tasks that detract from the main task.  The authors back this up through experiments. They consider three tasks on CS and Biomed papers, where they train 1) the main task on top of a RoBERTa checkpoint 2) the main task on top of a pre-trained model from Gururangan et al. (2020) and 3) the main task mixed-in with auxiliary tasks from (2) on top of a RoBERTa checkpoint. They observe that their approaches (3) outperform (1) and (2), both in terms of model accuracy and data efficiency. In a setting where the auxiliary tasks are potentially more noisy, they also observe an advantage of META-TARTAN over MT-TARTAN.  Overall their work heavily references Gururangan et al. (2020), who show that adapting a generic model to the task domain through more task-specific pre-training can improve model performance. They make it directly comparable by using their models and a subset of tasks considered in that paper.