Relational Multi-Task Learning: Modeling Relations between Data and Tasks

1 INTRODUCTION . The general idea of learning from multiple tasks has been explored under different settings , including multi-task learning ( Caruana , 1997 ) , meta learning ( Finn et al. , 2017 ) , and few-shot learning ( Vinyals et al. , 2016 ) . While these learning settings have inspired models that can utilize relationships among tasks ( Chen et al. , 2019 ; Zamir et al. , 2018 ; Sener & Koltun , 2018 ; Lin et al. , 2019 ; Ma et al. , 2020 ) , they are not able to capture the full complexity of real-world machine learning applications . Concretely , when learning from multiple tasks , current approaches assume that the test data points have no access to the labels from other tasks when making predictions on a new task . However , this assumption oversimplifies potential useful knowledge in many applications . In reality , when making predictions for a given task , one may have access and want to utilize any known auxiliary task labels to further improve the prediction performance . Intuitively , imagine we want to predict whether it is safe to use a chemical compound x in the production of a recently discovered drug . In the meantime if ( 1 ) we know the compound x has positive results on two toxicology tests , and ( 2 ) the two toxicology tests have a high correlation with the task/test we are studying , then we could make a more accurate prediction with higher confidence by leveraging the additional knowledge . The example above shows how modeling and leveraging auxiliary tasks may benefit model performance . However , current deep learning architectures can not model such knowledge transfer from auxiliary task labels to the task of interest . Naively concatenating the known labels to the input features has its limitations especially since auxiliary task labels are sparsely available and it is also unclear how to use the approach for new and unseen tasks . Another potential solution to model such flexible and conditional inference is through generative models ( Dempster et al. , 1977 ; Koller & Friedman , 2009 ) . Although generative models are powerful , they are notoriously data-hungry , thus it is very difficult to construct and train a generative model for high dimensional data ( Turhan & Bilge , 2018 ) . State-of-the-art deep generative models have diverged quite heavily from state-of-the-art discriminative architectures in terms of accuracy on most of the downstream tasks . Here we propose a new multi-task learning setting called relational multi-task learning , where the goal is to achieve data efficiency through leveraging labels of a given data point on other already known auxiliary tasks . This setting is especially important for biochemical domains where knowledge is scattered in multiple datasets and tasks , e.g. , solving the molecule toxicity prediction problem we mentioned earlier . To tackle the relational multi-task learning , we propose MetaLink , a general discriminative model that can explicitly incorporate the knowledge from auxiliary tasks . Our key innovation is to build a knowledge graph that connects different tasks tj and data points x ( i ) ( Figure 1 ) . The first step of our approach is to take input data points x ( i ) and the feature extractor ( i.e. , neural network ) fθ to get its embedding z ( i ) . Then we build the knowledge graph which consists of two types of nodes : data nodes x ( i ) and task nodes tj . A data node x ( i ) connects to a task node tj if data point x ( i ) participates in task tj and the edge is annotated with the label y ( i ) j of x ( i ) on task tj . We initialize data node features to be the last layer embedding z ( i ) in the feature extractor fθ , and task node tj features are instantiated as the last layer ’ s weights wj . Given our knowledge graph , we reformulate the multi-task learning problem as a link-label prediction problem between data nodes and task nodes . This means that at the inference time MetaLink is able to use all the information about a given data point x ( i ) ( including its labels on auxiliary tasks { tj } ) to predict its label on a new task tn . We solve this link label prediction learning task via a Graph Neural Network ( GNN ) ( Hamilton et al. , 2017 ; He et al. , 2019 ; Xue et al. , 2021 ) . Unlike previous works , e.g. , ML-GCN ( Chen et al. , 2019 ) , that solely model relationships among tasks , MetaLink allows flexible and automatic modeling for data-task , data-data and task-task relationships . We evaluate MetaLink on six benchmark datasets in both biochemical and vision domains under various settings . Empirical demonstrate that MetaLink can successfully utilize the relations among different tasks , outperforming the state-of-the-art methods under the proposed relational multi-task learning setting , with up to 27 % improvement in ROC AUC . 2 RELATIONAL MULTI-TASK LEARNING SETTINGS . Here we first formally introduce the settings for relational multi-task learning . Suppose we have m machine learning tasks { tj } j∈T , where T = { 1 , 2 , ... , m } is integers between 1 and m. We propose to categorize different settings in two dimensions : ( 1 ) whether the task is a relational task , i.e. , if auxiliary task labels can be used at inference time ; and , ( 2 ) whether the task is a meta task , i.e. , if the task at test time has been seen at the training time . Altogether , there are four possible task settings , which are illustrated in Figure 2 and below . Standard supervised setting . Let x ( i ) denote the input and y ( i ) j denote the corresponding label associated with task tj , i.e. , y ( i ) j ∼ tj . Standard supervised multi-task learning can be represented as Train : { x ( i ) → { y ( i ) j ∼ tj } j∈T } Test : { x ( i ) → { y ( i ) j ∼ tj } j∈T } where→ connects the input and the output . Training and test sets cover non-overlapping data points . We use y ( i ) j to concisely represent y ( i ) j ∼ tj later on . Relational setting . In relational setting , in addition to input x ( i ) , we assume we also have access to auxiliary task labels when making predictions . Taux and Ttest are partitions of integers ( T ) that relates to subsets of tasks . Specifically , Taux refers to the indices of tasks that input x has access to , and Ttest are the indices of tasks that we wish to predict ; these two sets are non-overlapping , i.e. , Taux ∩ Ttest = ∅ , and input-dependent , i.e. , the partition T ( 1 ) test and T ( 2 ) test can be different . The input-output pairs are now in the form of Train : { ( x ( i ) , { y ( i ) j } j∈T ( i ) aux ) → { y ( i ) j } j∈T ( i ) test } Test : { ( x ( i ) , { y ( i ) j } j∈T ( i ) aux ) → { y ( i ) j } j∈T ( i ) test } Meta setting . In the meta setting , we want to learn to predict unseen tasks at the test time . Formally , let Ts , Tu denote the set of partitions for seen tasks ( used at training time ) and unseen tasks ( used at test time ) , where Ts ∩ Tu = ∅ . Here , we have access to a batch of samples with labels as the support set S and the objective is to correctly predict samples in the query set Q . Train : given S = { ( x ( i ) , { y ( i ) j } j∈Ts ) } , predict Q = { x ( i ) → { y ( i ) j } j∈Ts } Test : given S = { ( x ( i ) , { y ( i ) j } j∈Tu ) } , predict Q = { x ( i ) → { y ( i ) j } j∈Tu } Relational meta setting . Relational meta setting combines the features of relational setting and meta setting . Similar to the meta setting , we aim to predict unseen tasks Tu at test time ; meanwhile , similar to the relational setting , we also assume having labels on a limited number of auxiliary tasks Taux ⊆ Ts to harness . Formally , we have a support set and a query set in the form of Train : given S = { ( x ( i ) , { y ( i ) j } j∈T ( i ) s ) } , predict Q = { ( x ( i ) , { y ( i ) j } j∈T ( i ) aux ) → { y ( i ) j } j∈T ( i ) s \T ( i ) aux } Test : given S = { ( x ( i ) , { y ( i ) j } j∈T ( i ) u ) } , predict Q = { ( x ( i ) , { y ( i ) j } j∈T ( i ) aux ) → { y ( i ) j } j∈T ( i ) u } 3 METALINK FRAMEWORK . Next , we describe our MetaLink framework , which allows us to formulate the above multi-task learning settings in a single framework . In particular , MetaLink formulates them as a link label prediction task on a heterogeneous knowledge graph , this way , MetaLink can harness the relational information about data and tasks . 3.1 BUILD A KNOWLEDGE GRAPH ON TASK HEADS . We first recap the general formulation of a neural network . Given data points { x ( i ) } ni=1 and labels { { y ( i ) j } j∈T } ni=1 , a deep learning model can be formulated as a parameterized embedding function fθ ( which is deep ) and a task head fw , consisting of only a single weight matrix w. fθ maps a data point x ( i ) to a vector embedding space , fθ ( x ( i ) ) = z ( i ) ∈ RD . fw then maps an embedding z ( i ) to prediction ŷ ( i ) ∈ R1 , fw ( z ( i ) ) = wT z ( i ) = ŷ ( i ) . When a task head involves multi-layer transformation , we have fw ( z ( i ) ) = wT g ( z ( i ) ) = ŷ ( i ) , where g ( · ) can be an arbitrary function . In multi-task learning settings , people usually assign multiple task heads to a neural network . Suppose we have m tasks { tj } mj=1 , then there are m task heads such that ŷ ( i ) j = fwj ( z ( i ) ) . Here , our observation is that the weights in task heads w1 , . . . , wj , and the data embeddings z ( i ) play symmetric roles in a multi-task prediction task ( due to the dot product ) . Therefore , instead of viewing weights w1 , . . . , wj as parameters in a neural network , we propose to represent w1 , . . . , wj as another type of input that supports the prediction . Essentially , we reformulate a task head from ŷ ( i ) j = fwj ( z ( i ) ) to ŷ ( i ) j = fφ ( wj , z ( i ) ) . This new perspective , where both task weights wj and embedding z ( i ) are viewed as input , enables us to build a more sophisticated predictive model fφ . fφ in the sense that it contains two main steps , i.e. , GraphConv ( · ) and EdgePred ( · ) . In general , EdgePred ( · ) has a similar model complexity as fwj , whereas GraphConv ( · ) provides additional expressiveness . In MetaLink , we propose to build a knowledge graph over task weights wj and embedding z ( i ) . By building this knowledge graph , we can succinctly represent relationships between data points and tasks , as well as different multi-task learning settings . Concretely , a knowledge graph helps us easily express any data-task relationship ( e.g. , a data point has a label on a given task ) , data-data relationship ( e.g. , two data points are similar or not ) , or task-task relationship ( e.g. , hierarchy of different tasks ) . Moreover , a knowledge graph greatly simplifies the representation of all the multi-task learning settings that we outlined in Section 2 ; in fact , all the settings can be viewed as link label prediction tasks where different portions of the knowledge graph can be constructed , as illustrated in Figure 2 . We define the knowledge graph as G = { V , E } , where V is the node set and E ⊆ V × V is the edge set . We define two types of nodes , data nodes Vd = { x ( 1 ) , . . . , x ( n ) } and task nodes Vt = { t1 , . . . , tm } . We can then define edges between data and task nodes Edt ⊆ Vd × Vt , within data nodes Edd ⊆ Vd×Vd , and within task nodes Ett ⊆ Vt×Vt . MetaLink framework can work with all three types of edges ; however , since most benchmark datasets do not have information on datadata or task-task relationship , we focus on data-task relationship Edt in the remaining discussions . Specifically , we define Edt based on task labels , Edt = { ( x ( i ) , tj ) ∼ y ( i ) j∈Taux } , i.e. , we connect a data node x ( i ) with a task node tj if label y ( i ) j exists .

Summary: The paper exploits the multi-task modeling using a heterogenous Graph Neural Network. The key contribution of the work is having both data (sample) and task nodes in the same graph, focusing on data-task relation, accommodating for sparse task labels by design.  Experimental results show relevant improvements with the proposed relational model on biomedical datasets and Imagenet splits for few-shot.