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: The paper discusses an approach that learns to weight data from different tasks in pretraining or mutli-task learning. It also gives a VC/empirical-processes style analysis, which gives guarantees for the algorithm and insights about sample complexity. Finally it describes experiments on a number of NLP problems.