Understanding and Improving Information Transfer in Multi-Task Learning

1 INTRODUCTION . Multi-task learning has recently emerged as a powerful paradigm in deep learning to obtain language ( Devlin et al . ( 2018 ) ; Liu et al . ( 2019a ; b ) ) and visual representations ( Kokkinos ( 2017 ) ) from large-scale data . By leveraging supervised data from related tasks , multi-task learning approaches reduce the expensive cost of curating the massive per-task training data sets needed by deep learning methods and provide a shared representation which is also more efficient for learning over multiple tasks . While in some cases , great improvements have been reported compared to single-task learning ( McCann et al . ( 2018 ) ) , practitioners have also observed problematic outcomes , where the performances of certain tasks have decreased due to task interference ( Alonso and Plank ( 2016 ) ; Bingel and Søgaard ( 2017 ) ) . Predicting when and for which tasks this occurs is a challenge exacerbated by the lack of analytic tools . In this work , we investigate key components to determine whether tasks interfere constructively or destructively from theoretical and empirical perspectives . Based on these insights , we develop methods to improve the effectiveness and robustness of multi-task training . There has been a large body of algorithmic and theoretical studies for kernel-based multi-task learning , but less is known for neural networks . The conceptual message from the earlier work ( Baxter ( 2000 ) ; Evgeniou and Pontil ( 2004 ) ; Micchelli and Pontil ( 2005 ) ; Xue et al . ( 2007 ) ) show that multi-task learning is effective over “ similar ” tasks , where the notion of similarity is based on the single-task models ( e.g . decision boundaries are close ) . The work on structural correspondence learning ( Ando and Zhang ( 2005 ) ; Blitzer et al . ( 2006 ) ) uses alternating minimization to learn a shared parameter and separate task parameters . Zhang and Yeung ( 2014 ) use a parameter vector for each task and learn task relationships via l2 regularization , which implicitly controls the capacity of the model . These results are difficult to apply to neural networks : it is unclear how to reason about neural networks whose feature space is given by layer-wise embeddings . To determine whether two tasks interfere constructively or destructively , we investigate an architecture with a shared module for all tasks and a separate output module for each task ( Ruder ( 2017 ) ) . See Figure 1 for an illustration . Our motivating observation is that in addition to model similarity which affects the type of interference , task data similarity plays a second-order effect after controlling model similarity . To illustrate the idea , we consider three tasks with the same number of data ∗Equal contribution . Correspondence to { senwu , hongyang , chrismre } @ cs.stanford.edu samples where task 2 and 3 have the same decision boundary but different data distributions ( see Figure 2 for an illustration ) . We observe that training task 1 with task 2 or task 3 can either improve or hurt task 1 ’ s performance , depending on the amount of contributing data along the decision boundary ! This observation shows that by measuring the similarities of the task data and the models separately , we can analyze the interference of tasks and attribute the cause more precisely . Motivated by the above observation , we study the theory of multi-task learning through the shared module in linear and ReLU-activated settings . Our theoretical contribution involves three components : the capacity of the shared module , task covariance , and the per-task weight of the training procedure . The capacity plays a fundamental role because , if the shared module ’ s capacity is too large , there is no interference between tasks ; if it is too small , there can be destructive interference . Then , we show how to determine interference by proposing a more fine-grained notion called task covariance which can be used to measure the alignment of task data . By varying task covariances , we observe both positive and negative transfers from one task to another ! We then provide sufficient conditions which guarantee that one task can transfer positively to another task , provided with sufficiently many data points from the contributor task . Finally , we study how to assign per-task weights for settings where different tasks share the same data but have different labels . Experimental results . Our theory leads to the design of two algorithms with practical interest . First , we propose to align the covariances of the task embedding layers and present empirical evaluations on well-known benchmarks and tasks . On 5 tasks from the General Language Understanding Evaluation ( GLUE ) benchmark ( Wang et al . ( 2018b ) ) trained with the BERTLARGE model by Devlin et al . ( 2018 ) , our method improves the result of BERTLARGE by a 2.35 % average GLUE score , which is the standard metric for the benchmark . Further , we show that our method is applicable to transfer learning settings ; we observe up to 2.5 % higher accuracy by transferring between six sentiment analysis tasks using the LSTM model of Lei et al . ( 2018 ) . Second , we propose an SVD-based task reweighting scheme to improve multi-task training for settings where different tasks have the same features but different labels . On the ChestX-ray14 dataset , we compare our method to the unweighted scheme and observe an improvement of 0.4 % AUC score on average for all tasks . In conclusion , these evaluations confirm that our theoretical insights are applicable to a broad range of settings and applications . 2 THREE COMPONENTS OF MULTI-TASK LEARNING . We study multi-task learning ( MTL ) models with a shared module for all tasks and a separate output module for each task . We ask : What are the key components to determine whether or not MTL is better than single-task learning ( STL ) ? In response , we identify three components : model capacity , task covariance , and optimization scheme . After setting up the model , we briefly describe the role of model capacity . We then introduce the notion of task covariance , which comprises the bulk of the section . We finish by showing the implications of our results for choosing optimization schemes . 2.1 MODELING SETUP . We are given k tasks . Letmi denote the number of data samples of task i . For task i , letXi ∈ Rmi×d denote its covariates and let yi ∈ Rmi denote its labels , where d is the dimension of the data . We have assumed that all the tasks have the same input dimension d. This is not a restrictive assumption and is typically satisfied , e.g . for word embeddings on BERT , or by padding zeros to the input otherwise . Our model assumes the output label is 1-dimensional . We can also model a multi-label problem with k types of labels by having k tasks with the same covariates but different labels . We consider an MTL model with a shared module B ∈ Rd×r and a separate output module Ai ∈ Rr for task i , where r denotes the output dimension of B . See Figure 1 for the illustration . We define the objective of finding an MTL model as minimizing the following equation over B and the Ai ’ s : f ( A1 , A2 , . . . , Ak ; B ) = k∑ i=1 L ( g ( XiB ) Ai , yi ) , ( 1 ) where L is a loss function such as the squared loss . The activation function g : R→ R is applied on every entry of XiB . In equation 1 , all data samples contribute equally . Because of the differences between tasks such as data size , it is natural to re-weight tasks during training : f ( A1 , A2 , . . . , Ak ; B ) = k∑ i=1 αi · L ( g ( XiB ) Ai , yi ) , ( 2 ) This setup is an abstraction of the hard parameter sharing architecture ( Ruder ( 2017 ) ) . The shared module B provides a universal representation ( e.g. , an LSTM for encoding sentences ) for all tasks . Each task-specific module Ai is optimized for its output . We focus on two models as follows . The single-task linear model . The labels y of each task follow a linear model with parameter θ ∈ Rd : y = Xθ + ε . Every entry of ε follows the normal distribution N ( 0 , σ2 ) with variance σ2 . The function g ( XB ) = XB . This is a well-studied setting for linear regression ( Hastie et al . ( 2005 ) ) . The single-task ReLU model . Denote by ReLU ( x ) = max ( x , 0 ) for any x ∈ R. We will also consider a non-linear model where Xθ goes through the ReLU activation function with a ∈ R and θ ∈ Rd : y = a ·ReLU ( Xθ ) +ε , which applies the ReLU activation on Xθ entrywise . The encoding function g ( XB ) then maps to ReLU ( XB ) . Positive vs. negative transfer . For a source task and a target task , we say the source task transfers positively to the target task , if training both through equation 1 improves over just training the target task ( measured on its validation set ) . Negative transfer is the converse of positive transfer . Problem statement . Our goal is to analyze the three components to determine positive vs. negative transfer between tasks : model capacity ( r ) , task covariances ( { X > i Xi } ki=1 ) and the per-task weights ( { αi } ki=1 ) . We focus on regression tasks under the squared loss but we also provide synthetic experiments on classification tasks to validate our theory . Notations . For a matrixX , its column span is the set of all linear combinations of the column vectors of X . Let X† denote its pseudoinverse . Given u , v ∈ Rd , cos ( u , v ) is equal to u > v/ ( ‖u‖ · ‖v‖ ) . 2.2 MODEL CAPACITY . We begin by revisiting the role of model capacity , i.e . the output dimension ofB ( denoted by r ) . We show that as a rule of thumb , r should be smaller than the sum of capacities of the STL modules . Example . Suppose we have k linear regression tasks using the squared loss , equation 1 becomes : f ( A1 , A2 , . . . , Ak ; B ) = k∑ i=1 ‖XiBAi − yi‖2F . ( 3 ) The optimal solution of equation 3 for task i is θi = ( X > i Xi ) †X > i yi ∈ Rd . Hence a capacity of 1 suffices for each task . We show that if r ≥ k , then there is no transfer between any two tasks . Proposition 1 . Let r ≥ k. There exists an optimum B ? and { A ? i } ki=1 of equation 3 where B ? A ? i = θi , for all i = 1 , 2 , . . . , k. To illustrate the idea , as long as B ? contains { θi } ki=1 in its column span , there exists A ? i such that B ? A ? i = θi , which is optimal for equation 3 with minimum error . But this means no transfer among any two tasks . This can hurt generalization if a task has limited data , in which case its STL solution overfits training data , whereas the MTL solution can leverage other tasks ’ data to improve generalization . The proof of Proposition 1 and its extension to ReLU settings are in Appendix A.1 . Algorithmic consequence . The implication is that limiting the shared module ’ s capacity is necessary to enforce information transfer . If the shared module is too small , then tasks may interfere negatively with each other . But if it is too large , then there may be no transfer between tasks . In Section 3.3 , we verify the need to carefully choose model capacity on a wide range of neural networks including CNN , LSTM and multi-layer perceptron .

Summary: This paper studies how to improve the multi-task learning from both theoretical and experimental viewpoints. More specifically, they study an architecture where there is a shared model for all of the tasks and a separate module specific to each task. They show that data similarity of the tasks, measured by task covariance is an important element for the tasks to be constructive or destructive. They theoretically find a sufficient condition that guarantee one task can transfer positively to the other; i.e. a lower bound of the number of data points that one task has to have. Consequently, they propose an algorithm which is basically applying a covariance alignment method to the input.