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: In the paper the authors propose TARTAN, methods to enable end-task aware pre-training. MT-TARTAN simply combines the pre-training objectives and the end-task objective as multi-task learning. META-TARTAN learns a set of weights for the pre-training objective and the end-task objective, using meta-learning. MT-TARTAN and META-TARTAN shows improved performance and data efficiency in a set of three low-resource text classification tasks.