Weighted Training for Cross-Task Learning

In this paper , we introduce Target-Aware Weighted Training ( TAWT ) , a weighted training algorithm for cross-task learning based on minimizing a representationbased task distance between the source and target tasks . We show that TAWT is easy to implement , is computationally efficient , requires little hyperparameter tuning , and enjoys non-asymptotic learning-theoretic guarantees . The effectiveness of TAWT is corroborated through extensive experiments with BERT on four sequence tagging tasks in natural language processing ( NLP ) , including part-of-speech ( PoS ) tagging , chunking , predicate detection , and named entity recognition ( NER ) . As a byproduct , the proposed representation-based task distance allows one to reason in a theoretically principled way about several critical aspects of cross-task learning , such as the choice of the source data and the impact of fine-tuning.1 1 INTRODUCTION . The state-of-the-art ( SOTA ) models in real-world applications rely increasingly on the usage of weak supervision signals ( Pennington et al. , 2014 ; Devlin et al. , 2019 ; Liu et al. , 2019 ) . Among these , cross-task signals are one of the most widely-used weak signals ( Zamir et al. , 2018 ; McCann et al. , 2018 ) . Despite their popularity , the benefits of cross-task signals are not well understood from a theoretical point of view , especially in the context of deep learning ( He et al. , 2021 ; Neyshabur et al. , 2020 ) , hence impeding the efficient usage of those signals . Previous work has adopted representation learning as a framework to understand the benefits of cross-task signals , where knowledge transfer is achieved by learning a representation shared across different tasks ( Baxter , 2000 ; Maurer et al. , 2016 ; Tripuraneni et al. , 2020 ; 2021 ; Du et al. , 2021 ) . However , the existence of a shared representation is often too strong an assumption in practice . Such an assumption also makes it difficult to reason about several critical aspects of cross-task learning , such as the quantification of the value of the source data and the impact of fine-tuning ( Kalan & Fabian , 2020 ; Chua et al. , 2021 ) . In this paper , we propose Target-Aware Weighted Training ( TAWT ) , a weighted training algorithm for efficient cross-task learning . The algorithm can be easily applied to existing cross-task learning paradigms , such as pre-training and joint training , to boost their sample efficiency by assigning adaptive ( i.e. , trainable ) weights on the source tasks or source samples . The weights are determined in a theoretically principled way by minimizing a representation-based task distance between the source and target tasks . Such a strategy is in sharp contrast to other weighting schemes common in machine learning , such as importance sampling in domain adaptation ( Shimodaira , 2000 ; Cortes et al. , 2010 ; Jiang & Zhai , 2007 ) . 1Our code is publicly available at http : //cogcomp.org/page/publication_view/963 . The effectiveness of TAWT is verified via both theoretical analyses and empirical experiments . Using empirical process theory , we prove a non-asymptotic generalization bound for TAWT . The bound is a superposition of two vanishing terms and a term depending on the task distance , the latter of which is potentially negligible due to the re-weighting operation . We then conduct comprehensive experiments on four sequence tagging tasks in NLP : part-of-speech ( PoS ) tagging , chunking , predicate detection , and named entity recognition ( NER ) . We demonstrate that TAWT further improves the performance of BERT ( Devlin et al. , 2019 ) in both pre-training and joint training for cross-task learning with limited target data , achieving an average absolute improvement of 3.1 % on the performance . As a byproduct , we propose a representation-based task distance that depends on the quality of representations for each task , respectively , instead of assuming the existence of a single shared representation among all tasks . This finer-grained notion of task distance enables a better understanding of cross-task signals . For example , the representation-based task distance gives an interpretable measure of the value of the source data on the target task based on the discrepancy between their optimal representations . Such a measure is more informative than measuring the difference between tasks via the discrepancy of their task-specific functions ( e.g . linear functions ) as done in previous theoretical frameworks ( Tripuraneni et al. , 2020 ) . Furthermore , the representation-based task distance clearly conveys the necessity of fine-tuning : if this distance is non-zero , then fine-tuning the representation becomes necessary as the representation learned from the source data does not converge to the optimal target representation . Finally , we compare our work with some recent attempts in similar directions . Liu et al . ( 2020 ) analyze the benefits of transfer learning by distinguishing source-specific features and transferable features in the source data . Based on the two types of features , they further propose a meta representation learning algorithm to encourage learning transferable and generalizable features . Instead of focusing on the distinction between two types of features , our algorithm and analyses are based on the representation-based task distance and are thus different . Chua et al . ( 2021 ) present a theoretical framework for analyzing representations derived from model agnostic meta-learning ( Finn et al. , 2017 ) , assuming all the tasks use approximately the same underlying representation . In contrast , we do not impose any a priori assumption on the proximity among source and target representations , and our algorithm seeks for a weighting scheme to maximize the proximity . Our work is also different from task weighting in curriculum learning . This line of work tends to learn suitable weights in the stochastic policy to decide which task to study next in curriculum learning ( Graves et al. , 2017 ) , while TAWT aims to learn better representations by assigning more suitable weights on source tasks . Compared to heuristic weighting strategies in multi-task learning ( Gong et al. , 2019 ; Zhang & Yang , 2021 ) , we aim to design a practical algorithm with theoretical guarantees for cross-task learning . 2 TAWT : TARGET-AWARE WEIGHTED TRAINING 2.1 PRELIMINARIES . Suppose we have T source tasks , represented by a collection of probability distributions { Dt } Tt=1 on the sample space X × Y , where X ⊆ Rd is the feature space and Y ⊆ R is the label space . For classification problems , we take Y to be a finite subset of R. We have a single target task , whose probability distribution is denoted as D0 . For the t-th task , where t = 0 , 1 , . . . , T , we observe nt i.i.d . samples St = { ( xti , yti ) } nti=1 from Dt . Typically , the number of samples from the target task , n0 , is much smaller than the samples from the source tasks , and the goal is to use samples from source tasks to aid the learning of the target task . Let Φ be a collection of representations from the feature space X to some latent space Z ⊆ Rr . We refer to Φ as the representation class . Let F be a collection of task-specific functions from the latent space Z to the label space Y . The complexity of the representation class Φ is usually much larger ( i.e. , more expressive ) than that of the task-specific function class F . Given a bounded loss function ` : Y ×Y → [ 0 , 1 ] , the optimal pair of representation and task-specific function of the t-th task is given by ( φ ? t , f ? t ) ∈ argmin φt∈Φ , ft∈F Lt ( φt , ft ) , Lt ( φt , ft ) : = E ( X , Y ) ∼Dt [ ` ( ft ◦ φt ( X ) , Y ) ] . ( 2.1 ) Note that in general , the optimal representations of different tasks are different . For brevity , all proofs for the theory part are deferred to the Appx . A . 2.2 DERIVATION OF TAWT Under the assumption that the optimal representations { φ ? t } Tt=0 are similar , a representation learned using samples only from the source tasks would perform reasonably well on the target task . Consequently , we can devote n0 samples from the target task to learn only the task specific function . This is a much easier task , since the complexity of F is typically much smaller than that of Φ . This discussion leads to a simple yet immensely popular two-step procedure as follows ( Tripuraneni et al. , 2020 ; Du et al. , 2021 ) . First , we solve a weighted empirical risk minimization problem with respect to the source tasks : ( φ̂ , { f̂t } Tt=1 ) ∈ argmin φ∈Φ , { ft } ⊂F T∑ t=1 ωtL̂t ( φ , ft ) , L̂t ( φ , ft ) : = 1 nt nt∑ i=1 ` ( ft ◦ φ ( xti ) , yti ) , ( 2.2 ) where ω ∈ ∆T−1 is a user-specified vector lying in the T -dimensional probability simplex ( i.e. , ∑T t=1 ωt = 1 and ωt ≥ 0 , ∀1 ≤ t ≤ T ) . In the second stage , we freeze the representation φ̂ , and seek the task-specific function that minimizes the empirical risk with respect to the target task : f̂0 ∈ argmin f0∈F L̂0 ( φ̂ , f0 ) . ( 2.3 ) In practice , we can allow φ̂ to slightly vary ( e.g. , via fine-tuning ) to get a performance boost . In the two-step procedure ( 2.2 ) – ( 2.3 ) , the weight vector ω is usually taken to be a hyperparameter and is fixed during training . Popular choices include the uniform weights ( i.e. , ωt = 1/T ) or weights proportional to the sample sizes ( i.e. , ωt = nt/ ∑T t′=1 nt′ ) ( Liu et al. , 2019 ; Johnson & Khoshgoftaar , 2019 ) . This reveals the target-agnostic nature of the two-step procedure ( 2.2 ) – ( 2.3 ) : the weights stay the same regardless the level of proximity between the source tasks and the target task . Consider the following thought experiment : if we know a priori that the first source task D1 is closer ( compared to other source tasks ) to the target task D0 , then we would expect a better performance by raising the importance of D1 , i.e. , make ω1 larger . This thought experiment motivates a target-aware procedure that adaptively adjusts the weights based on the proximity of source tasks to the target . A natural attempt for developing such a task-aware procedure is a follows : min φ∈Φ , f0∈F , ω∈∆T−1 L̂0 ( φ , f0 ) subject to φ ∈ argmin ψ∈Φ min { ft } ⊂F T∑ t=1 ωtL̂t ( ψ , ft ) . ( OPT1 ) That is , we seek for the best weights ω such that solving ( 2.2 ) with this choice of ω would lead to the lowest training error when we subsequently solve ( 2.3 ) . Despite its conceptual simplicity , the formulation ( OPT1 ) is a complicated constrained optimization problem . Nevertheless , we demonstrate that it is possible to transform it into an unconstrained form for which a customized gradient-based optimizer could be applied . To do so , we let ( φω , { fωt } ) be any representation and task-specific functions that minimizes ∑T t=1 ωtL̂t ( φ , ft ) over φ ∈ Φ and { ft } ⊂ F . Equivalently , φω minimizes ∑T t=1 ωt minft∈F L̂t ( φ , ft ) over φ ∈ Φ . With such notations , we can re-write ( OPT1 ) as min f0∈F , ω∈∆T−1 L̂0 ( φω , f0 ) . ( OPT2 ) The gradient of the above objective with respect to the task-specific function , ∇f L̂0 ( φω , f0 ) , is easy to calculate via back-propagation . The calculation of the gradient with respect to the weights requires more work , as φω is an implicit function of ω . By the chain rule , we have ∂∂ωt L̂0 ( φ ω , f0 ) = [ ∇φL̂0 ( φω , f0 ) ] > ∂∂ωtφ ω . Since φω is a minimizer of φ 7→ ∑T t=1 ωt minft∈F L̂t ( φ , ft ) , we have F ( φω , ω ) = 0 , ∀ω ∈ ∆T−1 , F ( φ , ω ) : = ∇φ T∑ t=1 ωt min ft∈F L̂t ( φ , ft ) . ( 2.4 ) By implicit function theorem , if F ( · , · ) is everywhere differentiable and the matrix ∂F ( φ , ω ) /∂φ is invertible for any ( φ , ω ) near some ( φ̃ , ω̃ ) satisfying F ( φ̃ , ω̃ ) = 0 , then we can conclude that the Algorithm 1 : Target-Aware Weighted Training ( TAWT ) Input : Datasets { St } Tt=0 . Output : Final pair of representation and task-specific function ( φ̂ , f̂0 ) for the target task . Initialize parameters ω0 ∈ ∆T−1 , φ0 ∈ Φ , { f0t } Tt=0 ⊂ F ; for k = 0 , . . . , K − 1 do Starting from ( φk , { fkt } Tt=1 ) , run a few steps of SGD to get ( φk+1 , { fk+1t } Tt=1 ) ; Use the approximate gradient∇f L̂0 ( φk+1 , f0 ) to run a few steps of SGD from fk0 to get fk+10 ; Run one step of approximate mirror descent ( 2.7 ) – ( 2.8 ) from ωk to get ωk+1 ; end return φ̂ = φK , f̂0 = fK0 map ω 7→ φω is a locally well-defined function near ω̃ , and the derivative of this map is given by ∂ ∂ωt φω = − ( ∂F ( φ , ω ) ∂φ ∣∣∣∣ φ=φω ) −1 ( ∂F ( φ , ω ) ∂ωt ∣∣∣∣ φ=φω ) . ( 2.5 ) To simplify the above expression , note that under regularity conditions , we can regard ∇φL̂t ( φ , fωt ) as a sub-gradient of the map φ 7→ minft∈F L̂t ( φ , ft ) . This means that we can write F ( φω , ω ) =∑T t=1 ωt∇φL̂t ( φω , fωt ) . Plugging this expression back to ( 2.5 ) and recalling the expression for ∂L̂0 ( φω , f0 ) /∂ωt derived via the chain rule , we get ∂ ∂ωt L̂0 ( φω , f0 ) = − [ ∇φL̂0 ( φω , f0 ) ] > [ T∑ t=1 ωt∇2φL̂t ( φω , fωt ) ] −1 [ ∇φL̂t ( φω , fωt ) ] . ( 2.6 ) Now that we have the expressions for the gradients of L̂0 ( φω , f0 ) with respect to f0 and ω , we can solve ( OPT2 ) via a combination of alternating minimization and mirror descent . To be more specific , suppose that at iteration k , the current weights , representation , and task-specific functions are ωk , φk , and { fkt } Tt=0 , respectively . At this iteration , we conduct the following three steps : 1 . Freeze ωk . Starting from ( φk , { fkt } Tt=1 ) , run a few steps of SGD on the objective function ( φ , { ft } Tt=1 ) 7→ ∑T t=1 ω k t L̂t ( φ , ft ) to get ( φk+1 , { fk+1t } Tt=1 ) , which is regarded as an approxi- mation of ( φω k , { fωkt } Tt=1 ) ; 2 . Freeze ( φk+1 , { fk+1t } Tt=1 ) . Approximate the gradient∇f L̂0 ( φω k , f0 ) by∇f L̂0 ( φk+1 , f0 ) . Using this approximate gradient , run a few steps of SGD from fk0 to get f k+1 0 ; 3 . Freeze ( φk+1 , { fk+1t } Tt=0 ) . Approximate the partial derivative ∂L̂0 ( φω k , fk+10 ) /∂ωt by gkt : = − [ ∇φL̂0 ( φk+1 , fk+10 ) ] > [ T∑ t=1 ωt∇2φL̂t ( φk+1 , fk+1t ) ] −1 [ ∇φL̂t ( φk+1 , fk+1t ) ] . ( 2.7 ) Then run one step of mirror descent ( with step size ηk ) from ωk to get ωk+1 : ωk+1t = ωkt exp { −ηkgkt } ∑T t′=1 ω k t′ exp { −ηkgkt′ } . ( 2.8 ) We use mirror descent in ( 2.8 ) , as it is a canonical generalization of Euclidean gradient descent to gradient descent on the probability simplex ( Beck & Teboulle , 2003 ) . Note that other optimization methods , such as projected gradient descent , can also be used here . The update rule ( 2.8 ) has a rather intuitive explanation . Note that gkt is a weighted dissimilarity measure between the gradients ∇φL̂0 and∇φL̂t . This can further be regarded as a crude dissimilarity measure between the optimal representations of the target task and the t-th source task . The mirror descent updates ωt along the direction where the target task and the t-th source task are more similar . The overall procedure is summarized in Algorithm 1 . A faithful implementation of the above steps would require a costly evaluation of the inverse of the Hessian matrix ∑T t=1 ωt∇2φL̂t ( φk+1 , f k+1 t ) ∈ Rr×r . In practice , we can bypass this step by replacing2 the Hessian-inverse-weighted dissimilarity measure ( 2.7 ) with a consine-similarity-based dissimilarity measure ( see Section 4 for details ) . The previous derivation has focused on weighted pre-training , i.e. , the target data is not used when defining the constrained set in ( OPT1 ) . It can be modified , mutatis mutandis , to handle weighted joint-training , where we change ( OPT1 ) to min φ∈Φ , f0∈F , ω∈∆T L̂0 ( φ , f0 ) subject to φ ∈ argmin ψ∈Φ min { ft } ⊂F T∑ t=0 ωtL̂t ( ψ , ft ) . ( 2.9 ) Compared to ( OPT1 ) , we now also use the data from the target task when learning the representation φ , and thus there is an extra weight ω0 on the target task . The algorithm can also be easily extended to handle multiple target tasks or to put weights on samples ( as opposed to putting weights on tasks ) . The algorithm could also be applied to improve the efficiency of learning from cross-domain and cross-lingual signals , and we postpone such explorations for future work .

Summary: This paper introduces Target-Aware Weighted Training (TAWT), a cross-task learning algorithm. In the two-step procedure popular for cross-task training, the weight vector $\omega$ defining the relative weight of each source task is usually exogenous (i.e: a hyperparameter). In contrast, in TAWT it is part of the optimization procedure, with the weights depending on the proximity (throughout training) of each source task to the target task (hence the “target aware” name).  The authors first derive an algorithm that enables learning when $ \omega $ is made endogenous in the case of (weighted) pre-training. They then show how it can naturally be extended for weighted joint-training by considering the source tasks as one of the source tasks we learn representations from. They then provide theoretical performance guarantees for TAWT.  The authors apply TAWT on 4 NLP tasks, using BERT as their base model. They evaluate their method on both full-data and limited-data settings on each target task, using the 3 other tasks as source tasks. They show that TAWT-pretraining and TAWT-joint training significantly outperform their unweighted counterparts across tasks. These improvements are more marked when target task data is scarce.  They also show that weight initialization can be particularly important when the different source datasets have different # examples (e.g: when training jointly on data abundant source tasks and the data-scarce target task). This specific finding is not very novel, but it is great to see TAWT-training also helping when initialization weights are chosen correctly.   Critically, the authors also show the importance of varying task weights throughout training. Indeed, fixing the task weights to be the final weights from TAWT leads to worse results than training with TAWT. This suggests that it is important that task-weights vary throughout the training process.