R ELATIONAL M ULTI -T ASK L EARNING : M ODELING R ELATIONS BETWEEN D ATA AND T ASKS
A BSTRACT
A key assumption in multi-task learning is that at the inference time the multi-task model only has access to a given data point but not to the data point’s labels from other tasks.
This presents an opportunity to extend multi-task learning to utilize data point’s labels from other auxiliary tasks, and this way improves performance on the new task.
Here we introduce a novel relational multi-task learning setting where we leverage data point labels from auxiliary tasks to make more accurate predictions on the new task.
We develop MetaLink, where our key innovation is to build a knowledge graph that connects data points and tasks and thus allows us to leverage labels from auxiliary tasks.
The knowledge graph consists of two types of nodes: (1) data nodes, where node features are data embeddings computed by the neural network, and (2) task nodes, with the last layer’s weights for each task as node features.
The edges in this knowledge graph capture data-task relationships, and the edge label captures the label of a data point on a particular task.
Under MetaLink, we reformulate the new task as a link label prediction problem between a data node and a task node.
The MetaLink framework provides ﬂexibility to model knowledge transfer from auxiliary task labels to the task of interest.
We evaluate MetaLink on 6 benchmark datasets in both biochemical and vision domains.
Experiments 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.
I NTRODUCTION
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; Cao et al., 2020a;b).
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 oversimpliﬁes potential useful knowledge in many real-world applications.
For example, multi-task learning studies simultaneously learning multiple predictive tasks to exploit commonalities between the tasks.
At the test time the multi-task model predicts labels of a given data point for the tasks of interests, e.g., predicting whether the chemical compound x is non-toxic. At the same time, one may also have access to the data point’s labels for some other auxiliary tasks, e.g., the compound x has a positive result on two toxicology tests.
Such auxiliary task labels could greatly help with improved predictions.
However, current deep learning architectures cannot model such knowledge transfer between auxiliary tasks/labels and the target tasks.
Naively concatenating the known labels to the input features has its limitations, especially since such 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 ﬂexible 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 difﬁcult to construct and train a generative model for high dimensional data (Turhan & Bilge, 2018).
Featureextractor 𝑓 !
Input Data
MetaLink 𝑥 (#) 𝑥 (%) 𝑥 (&) 𝑡 # 𝑡 % 𝑡 & 𝑤 !
𝑡 # 𝑡 % 𝑡 & 𝑡 # 𝑥 (!) 𝑧 (!)
𝑥 ($) 𝑧 ($) 𝑥
(%) 𝑧 (%) 𝑡 % 𝑡 & 𝑤 !
𝑤 $ 𝑤 %
Default Approach 𝑧 ($) 𝑧 (!)
𝑧 (%) 𝑤 $ 𝑤 %
Here we propose a new multi-task learning setting called relational multi-task learning . In our setting we distinguish between the target tasks (e.g., predicting molecule toxicity), which are tasks we aim to predict, and auxiliary tasks, which are tasks for which data point’s labels are available at the inference time.
Note that under our setting, each data point might have labels available for a different subset of auxiliary tasks. The goal then is to achieve strong predictive performance through leveraging labels of a given data point on some subset of auxiliary tasks.
To tackle the relational multi-task learning, we propose MetaLink 1 , 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 t j and data points x ( i ) (Figure 1).
The ﬁrst 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 t j .
A data node x ( i ) connects to a task node
t
j if data point x ( i ) participates in task t
j and the edge is annotated with the label y
( i ) j of x ( i ) on task t j .
We initialize data node features to be the last layer embedding z ( i ) in the feature extractor f θ , and task node t j features are instantiated as the last layer’s weights w j .
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 { t j } ) to predict its label on a new task t n .
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 ﬂexible 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.
R ELATIONAL M ULTI - TASK L EARNING S ETTINGS
Here we ﬁrst formally introduce the settings for relational multi-task learning.
Suppose we have m machine learning tasks { t j } 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 [Figure] 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 t j , i.e. , y ( i ) j ∼ t j .
Standard supervised multi-task learning can be represented as
[Equation]
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 ∼ t j 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.
T aux and T test are partitions of integers ( T ) that relates to subsets of tasks.
Speciﬁcally, T aux refers to the indices of tasks that input x has access to, and T test are the indices of tasks that we wish to predict; these two sets are non-overlapping, i.e. , T aux ∩ T test = ∅ , 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
[Equation]
Meta setting.
In the meta setting, we want to learn to predict unseen tasks at the test time.
Formally, let T s , T u denote the set of partitions for seen tasks (used at training time) and unseen tasks (used at test time), where T s ∩ T u = ∅ .
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 =
(cid:110)
( x ( i ) , { y ( i ) j } j ∈ T s ) (cid:111) , predict Q
=
(cid:110) x
( i ) →
{ y ( i ) j } j ∈ T s (cid:111) Test: given S
=
(cid:110)
( x ( i ) , { y ( i ) j } j ∈ T u ) (cid:111) , predict Q
=
(cid:110) x
( i ) → { y ( i ) j } j ∈ T u (cid:111)
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 T u at test time; meanwhile, similar to the relational setting, we also assume having labels on a limited number of auxiliary tasks
T aux ⊆
T s to harness.
Formally, we have a support set and a query set in the form of
Train: given S =
(cid:110) ( x ( i ) , { y ( i ) j } j ∈ T ( i ) s ) (cid:111)
predict Q
=
(cid:110)
( x ( i ) , { y ( i ) j } j ∈ T ( i ) aux )
→ { y ( i ) j } j ∈ T
[Equation]
i )
[Equation]
[Table]
M ETA L INK F RAMEWORK
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. B UILD A K NOWLEDGE G RAPH ON T ASK H EADS
We ﬁrst 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 f w , 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 ) ∈ R D .
f w then maps an embedding z ( i ) to prediction ˆ y ( i ) ∈ R 2 , f w ( z ( i ) )
=
w T z ( i ) =
ˆ y ( i ) .
When a task head involves multi-layer transformation, we have f w ( z ( i ) ) =
w T 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 { t j } mj =1 , then there are m task heads such that ˆ y ( i )
j
= f w j
( z ( i ) ) .
Here, our observation is that the weights in task heads w 1 , . . . , w j , 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 w 1 , . . .
, w j as parameters in a neural network, we propose to represent w 1 , . . . , w j as another type of input that supports the prediction.
Essentially, we reformulate a task head from ˆ y ( i )
j = f w j
( z ( i ) ) to ˆ y
( i )
j = f φ ( w j , z
( i ) ) .
This new perspective, where both task weights w j 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 f w j , whereas GraphConv ( · ) provides additional expressiveness.
In MetaLink, we propose to build a knowledge graph over task weights w j 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 simpliﬁes 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 deﬁne the knowledge graph as G =
{ V, E } , where V is the node set and E ⊆
V × V is the edge set.
We deﬁne two types of nodes, data nodes V d = { x (1) , . . .
, x ( n ) } and task nodes
V t = { t 1 , . . . , t m } .
We can then deﬁne edges between data and task nodes
E dt ⊆
V d × V t , within data nodes E dd ⊆
V d × V d , and within task nodes E tt ⊆
V t × V t .
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 E dt in the remaining discussions.
Speciﬁcally, we deﬁne E dt based on task labels, E dt
= { ( x ( i ) , t j ) ∼ y
( i ) j ∈ T aux } , i.e. , we connect a data node x ( i ) with
a task node t
j if label y
( i ) j exists.
3. L EARN FROM THE K NOWLEDGE G RAPH VIA A H ETEROGENEOUS GNN
Given the constructed graph, we then discuss how MetaLink learns from the built knowledge graph.
Initialize node/edge features.
First, we initialize features for the knowledge graph.
Concretely, we initialize data node features to be the data embeddings z ( i ) computed from the feature extractor f θ ( · ) , h (0) i =
z ( i ) .
We initialize the seen task node embeddings using the weights in task heads w j , h (0)
j = w j .
In meta settings where unseen task nodes appear during test time, we will initialize those nodes by a constant vector of 1 ; such initialization ensures that MetaLink can generalize to unseen tasks, since the node feature construction process is inductive.
Predict via a heterogeneous GNN.
We implement the predictive model f φ over data nodes and task nodes as a GNN.
The goal of GNN is to learn expressive node embeddings h v based on an iterative aggregation of local network neighborhoods.
The l -th iteration of GraphConv ( · ) , or the l -th layer, can be written as:
[Equation]
where h ( l ) v is the node embedding after l iterations, h (0) v have been initialized as explained above, and N G ( v ) denotes the direct neighbors of v .
A GG is the acronym for aggregation function and M SG is the acronym for message function.
We perform L GNN layers on top of the knowledge graph G that we have built.
After updating data and task node embeddings, we can make predictions on the given task through EdgePred ( · ) in the form of ˆ y ( i ) j = MLP (cid:16) C ONCAT ( h ( L ) i , h ( L ) j ) (cid:17) .
In general, MetaLink should work with any GNN architecture that follows the deﬁnition in Equation 1.
We use the GraphSAGE layer (Hamilton et al., 2017) in MetaLink ( W ( l ) , U ( l ) are trainable):
[Equation]
Next, we discuss the special GNN design used in MetaLink which has shown to be successful.
Special GNN designs in MetaLink.
We make three extensions in the formulation in Equation 1.
First, since there are two types of nodes in our knowledge graph, we deﬁne different message passing functions for different message types, i.e. , the message from data nodes to task nodes, and the message from task nodes to data nodes.
Second, we include edge features in the message computation.
This is especially important for our formulation, since the task label values y ( i ) j are included as edge features, and should be considered during GNN message passing.
Concretely, we extend Equation 2 into:
[Equation]
where 1 [ v ∈ V d , u ∈ V t ] indicates the message type (whether from task to data, or data to task), and O ( l ) is an additional trainable weight that allows task label y ( u ) v participates in message passing.
Finally, we let each GNN layer make a prediction and sum them up as the ﬁnal prediction; this way, the ﬁnal prediction is obtained from mixed information from different hops of node neighbors.
We observe this multi-layer ensemble technique can help MetaLink make robust predictions.
MetaLink for relational meta setting.
Here we provide a detailed description on how to apply MetaLink for the relational meta setting in Algorithm 1.
At training time, since most of the existing multi-label datasets are not designed for meta setting, we manually simulate such setting by sampling a mini batch with support S =
(cid:110)
( x ( i ) , { y ( i ) j } j ∈ T ( i ) s ) (cid:111) and query set (cid:110)
( x ( i ) , { y ( i ) j } j ∈ T ( i ) aux )
→ { y ( i ) j } j ∈ T ( i ) s \ T ( i ) aux (cid:111) .
We make sure the sampled meta tasks T ( i ) s and auxiliary knowledge tasks T ( i ) aux have no intersection.
We ﬁrst use the feature extractor f θ to get data embeddings and use the embeddings to initialize data nodes.
To initialize task nodes, we either: (1) use 1 to initialize if the task is a meta task; Or, (2) use trained weights otherwise.
We construct the edge set
E by connecting each data-task pair based on (cid:110) { y ( i ) j } j ∈ T ( i ) s (cid:111) in the support set and (cid:110) { y ( i ) j } j ∈ T ( i )
aux (cid:111)
(we eliminate the relations we want to predict) in query set.
Now that we have all the components of the knowledge graph, we can apply the predictive graph model f φ to learn expressive data and task node embeddings and make predictions.
At test time, we use the same pipeline as in training to construct the knowledge graph and run inference.
Please refer to Appendix A for the pipelines of relational or meta settings.
E XPERIMENTS
Here we experimentally show that our proposed MetaLink ﬂexibly handles different settings and leverages knowledge from auxiliary tasks.
We ﬁrst evaluate our algorithms on Tox21 (Huang et al., 2016), Sider (Kuhn et al., 2016), ToxCast (Richard et al., 2016), and MS-COCO (Lin et al.,
Require: Dataset D train =
{ ( x , y ) } .
A parameterized embedding function f θ .
Last layer weights for each task w j .
A parameterized heterogeneous GNN f φ .
Number of GNN layers L .
1: for each iteration do 2:
S, Q ← SampleMiniBatch ( D train )
(cid:46) Simulate meta setting in training 3: { z } ← f θ ( x ) for x ∈ ( S, Q ) 4:
V (0)
d
= { h (0) i ← z for z ∈ { z }} (cid:46)
Initialize data nodes 5:
V (0) t
= { h (0)
j ← 1 if meta else w j for each w j } (cid:46) Initialize task nodes 6:
E = { e ij ← ( x ( i ) , t j ) for y ( i ) j ∈ ( S, Q ) } (cid:46) Initialize edges 7: for l = 1 to L do 8:
V ( l ) d , V ( l ) t ← GraphConv ( V ( l − 1) d , V ( l − 1) t , E ) with f φ 9: logits ← EdgePred ( V ( L ) d , V ( L ) t ) with f φ 10:
Backward (cid:16) Criterion ( logits , {{ y ( i ) j } j ∈ T ( i ) s \ T ( i ) aux } ∈ Q ) (cid:17) 2014) datasets with various controllable settings on relational multi-task learning.
To additionally demonstrate the advantage of MetaLink, we also include experiments on a well-studied task: few-shot learning.
Our core algorithm is developed using PyTorch (Paszke et al., 2017).
We use one NVIDIA
RTX 8000 GPU for each experiment and the most time-consuming one (MS-COCO) takes less than 24 hours.
Please refer to Appendix A for additional low-level implementation details.
4. E VALUATING M ETA L INK ON R ELATIONAL M
ULTI - TASK S ETTINGS
Datasets.
We simulate four relational multi-task settings (Figure 2) using four widely-used multilabel datasets.
Tox21 (Huang et al., 2016) contains 12 different toxicological experiments for each sample with binary labels (active/inactive).
Sider (Kuhn et al., 2016) is a database of marketed drugs and adverse drug reactions (ADR), grouped into 27 tasks.
ToxCast (Richard et al., 2016) contains about 8K pairs of molecular graphs and corresponding 617-dimensional binary vectors that represent different experimental results.
Microsoft COCO (Common Objects in Context)
(Lin et al., 2014) is originally a large-scale object detection, segmentation dataset.
By counting whether each type of object exists in a scene as a single task, it also serves as the default large-scale dataset for benchmarking multi-label classiﬁcation in vision.
There are 80 binary classiﬁcation tasks with an average of 2.9 positive labels per image.
Experimental setting.
We evaluate MetaLink on all four settings described in Figure 2, and we summarize the detailed conﬁgurations as follows.
S TANDARD SETTING : standard supervised learning.
In order to have a fair comparison, we evaluate the same set of tasks with unknown labels as we sampled in relational settings.
R ELATIONAL SETTING : we assume each example has access to labels of 20% of all tasks in each dataset.
We evaluate on the rest of the tasks with unknown labels.
M ETA
SETTING : we hold out 20% of the tasks at the training time.
We only evaluate on the held-out tasks at the test time.
We use 256-shot setting, meaning at test time, we use 256 data points as a support set to initialize the prototypes of unseen tasks.
The reason why the number of shots is much larger than what is commonly used in few-shot learning (1-shot, 5-shot) is that positive labels are sometimes sparse in certain tasks.
R ELATIONAL META SETTING : we hold out the same 20% of the tasks at the training time as in the meta setting.
At test time, we assume each unseen task has access to 20% of the labels of seen tasks.
We will release our splits of meta setting to promote reproducibility.
Baselines.
Though the main motivation of our work is to utilize auxiliary task labels, we still include a few baselines under the standard supervised setting in order to benchmark state-of-the-art results.
The simplest one is (1) Empirical risk minimization (ERM): we train the network with cross-entropy loss under the standard supervised setting; (2) Various commonly used graph neural network architectures designed for molecules, MPNN (Gilmer et al., 2017), DMPNN (Yang et al., 2019), MGCN (Lu et al., 2019), AttentiveFP (Xiong et al., 2019); (3) GROVER
(Rong et al., 2020) integrates Message Passing
Networks into the Transformer-style architecture.
By leverage pretraining, it achieves state-of-the-art results on the aforementioned molecule datasets; (4) Baseline++ (Chen et al., 2018): since there is no prior work on addressing meta setting for multi-task learning, we adapt Baseline++ from few-shot [Table] learning to this setting.
We ﬁrst train a feature extractor on the training set.
At test time we use the support set to train a linear classiﬁcation layer.
Results on biochemical datasets.
Table 5 summarizes the results on Tox21, Sider, and ToxCast datasets.
We ﬁrst clarify that when performing MetaLink under the standard setting, where there is no auxiliary task label to leverage, the knowledge graph we built degenerates to a trivial set since the edge set is empty.
It has a negligible difference compared with training a vanilla network with cross-entropy loss.
Thus, it is reasonable that MetaLink under the standard setting has similar performance with the baselines.
GROVER has better performance because it leverages large-scale pretraining.
In terms of comparing relational setting with standard setting, we conclude that the proposed MetaLink successfully leverages the auxiliary knowledge for each sample.
Notably, with auxiliary labels from 20% of the tasks, MetaLink surpassed GROVER with large-scale pretraining by a large margin.
In addition, we observe similar improvements between meta and relational meta settings.
The empirical results for meta or relational meta setting are very close to the performance in standard setting, proving the potential of data-efﬁcient research in biochemical domains.
Results on MS-COCO.
Table 2 summarizes the results on MS-COCO.
Note that as also mentioned in the settings, we need a large support set to ensure each task in the meta testing stage has at least one positive label.
This blocks the feasibility of using a large model and input size as it is done in common benchmarks.
Thus, we use ResNet-50 with an input size of 224 for all the experiments.
We observe consistent improvements for both relational setting and relational meta setting.
MetaLink also outperforms Baseline++ in the meta setting.
Note that since label distribution for most of the tasks is highly skewed, mAP degrades more in the meta setting than for the biochemical datasets.
[Table]
4. A BLATION S TUDIES
Does MetaLink truly learn to utilize correlations between tasks?
To better understand the improvement of our algorithm [Table] we ﬁrst plot the Pearson correlation heat map on Sider (Figure 3).
We
[Table]
on the Sider dataset.
We
MetaLink gains more improvement on the high
[Table]
then ﬁnd the top 3 tasks with the highest and lowest average correlation with respect to the rest of the tasks.
We report the performance of MetaLink on these two subsets of tasks, respectively (Table 3).
We observe that MetaLink demonstrates larger improvement on tasks with higher correlations.
This experiment veriﬁes that MetaLink learns to utilize correlations between tasks as expected.
How does MetaLink perform with varying ratio of auxiliary task labels?
We vary the ratio of additional labels per test point and report the results in Figure 4.
We observe consistent improvement as we gradually increase the number of auxiliary labels in each dataset.
This experiment demonstrates that MetaLink can successfully utilize marginal information whenever auxiliary task labels are added to the dataset.
The improvement is usually prominent for the ﬁrst few added labels.
In Sider, MetaLink can outperform GROVER by up are
4. E VALUATING M ETA L INK ON S TANDARD F EW - SHOT L EARNING D ATASETS
Since the relational multi-task setting is novel, there is a very small number of baselines that we can compare against.
The major aim of this section is to show the advantage of MetaLink in a well-studied problem: few-shot learning.
Note that there is a slight difference between the meta setting above and few-shot learning setting here.
For few-shot learning, for all the link label prediction tasks related to an input, there will be only one positive link.
This inductive bias could be easily incorporated by modeling all the link label predictions together using cross-entropy loss.
Experimental setting.
We evaluate performance on two standard benchmarks: miniImageNet (Vinyals et al., 2016) and tiered-ImageNet (Ren et al., 2018).
We compare against
MatchNet (Vinyals et al., 2016), Baseline++ (Chen et al., 2018), MetaOptNet (Lee et al., 2019), and
Meta-Baseline (Chen et al., 2020b), who only assume the input is a vector.
[Table]
Results.
Table 4 shows that our MetaLink outperforms the standard few-shot learning benchmarks.
Note that if we set KG Layer = 0, the proposed MetaLink degenerates to Meta-Baseline.
The experiments clearly demonstrate the beneﬁts of building a knowledge graph on the last layer.
Furthermore, as an ablation study we manipulate the number of KG layers and ﬁnd that in the few-shot image recognition setting, there’s an improvement for stacking 2 KG layers instead of 1, meaning non-linearity is useful.
We do not observe further improvements for more than 3 layers.
R ELATED W ORK
Multi-task learning.
Multi-task learning is a learning paradigm that jointly optimizes a set of tasks with shared parameters.
It is generally believed that relations across different tasks can improve the overall performance.
Some works cast it as a multi-objective optimization problem and introduce multiple gradient-based methods to reduce negative transfer among tasks (Fliege & Vaz, 2016; Lin et al., 2019).
Other works assign (adaptive) weights for different tasks using certain heuristics (Kendall et al., 2018; Chen et al., 2018).
Our empirical study is also closely related to multi-label learning, where the problem is usually decomposed into multiple binary classiﬁcation tasks (Tsoumakas &
Katakis, 2007).
There is a line of work learning to leverage the relationship among tasks (Haller et al., 2021; Zamir et al., 2020).
Wang et al.
(2016) utilized recurrent neural networks to transform labels into embedded label vectors to learn the correlation among labels.
The most recent work is
ML-GCN (Chen et al., 2019), which uses a GCN to map label graph into a set of inter-dependent classiﬁers.
Although the major motivation of our work is also about leveraging correlation among tasks, our problem formulation is new and thus yields a novel algorithm.
Exploring graph structure for data and tasks.
There are prior works exploring graph structure for data points or tasks, and graph structure was proven to be effective in some tasks.
Satorras & Estrach (2018) explores graph neural representations over data points only for few-shot learning.
Besides, some works study how to transfer knowledge among tasks through constructing a graph over task nodes/classiﬁers only (Liu et al., 2019; Chen et al., 2020a).
Along this direction, instead of building a fully connected graph, recent works utilize auxiliary task structure/knowledge graph to build the graph (Chen et al., 2019; Lee et al., 2018).
In contrast, and orthogonal to the papers above, we focus on modeling data-task relationships with the reinterpretation of the last layer.
In addition, with data-task relationships solely, we are still able to capture data-data, task-task relationships implicitly through higher-order message passing.
C ONCLUSION
We introduced relational multi-task settings in which the methods are required to learn to leverage labels on auxiliary tasks to predict on the new task.
These settings are impactful in the biomedical domain where labels of different tasks are often scarcely available.
To address these settings, we propose MetaLink, which is general enough to allow us formulating the above settings in a single framework.
We demonstrated that MetaLink can successfully utilize the relations among tasks, outperforming the state-of-the-art methods under the proposed relational multi-task learning setting, with up to 27% improvement in ROC AUC.
We limit our focus to model data-task relationships since most benchmark datasets do not have information on data-data or task-task relationships, though
MetaLink is expressive enough to model such relationships.
We leave extending MetaLink to more complex relationships or tasks for future work.
A CKNOWLEDGEMENT
We want to give special thanks to Hamed Nilforoshan, Michihiro Yasunaga, Weihua Hu, Paridhi Maheshwari, Xuechen Li, Wengong Jin, Octavian Ganea, Tommi Jaakkola and Regina Barzilay for the thoughtful discussions. We gratefully acknowledge the support of DARPA under Nos. HR00112190039 (TAMI), N660011924033 (MCS); ARO under Nos.
W911NF-16-1-0342 (MURI), W911NF-16-1-0171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), IIS-2030477 (RAPID), NIH under No. R56LM013365; Stanford Data
Science Initiative, Wu Tsai Neurosciences Institute, Chan Zuckerberg Biohub, Amazon, JPMorgan
Chase, Docomo, Hitachi, Intel, KDDI, Toshiba, NEC and UnitedHealth Group.
Jiaxuan You is supported by JPMC PhD Fellowship and Baidu Scholarship
A A DDITIONAL I MPLEMENTATION D ETAILS
Implementation details for biomedical datasets.
We use an adapted version of GraphSAGE
Hamilton et al.
(2017) as the base neural architecture, and implement MetaLink based on the descriptions in the main paper.
We use Adam optimizer, with initial learning of 0 .
001 and cosine learning rate scheduler.
The model is trained with a batch size of 128 for 50 epochs.
We search over the number of layers of [2 , 3 , 4 , 5] , and report the test set performance when the best validation set performance is reached.
Implementation details for MS-COCO.
We use ResNet-50 with an input size of 224 × 224 as the base neural architecture, and implement MetaLink based on the descriptions in the main paper.
We perform standard data augmentation, including random crop and horizontal ﬂipping.
We use Adam optimizer with an initial learning rate of 0 . 002 , weight decay of 1 × 10 − 4 .
The model is trained with a batch size of 256 for 40 epochs.
We use the same set of hyper-parameters for all the four settings mentioned in the main paper.
Implementation details for few-shot learning.
We use the widely adopted backbones ResNetwith an input size of 84 × 84 (Chen et al., 2020b).
We perform standard data augmentation, including random crop and horizontal ﬂipping.
We use SGD optimizer with an initial learning rate of 0 .
001 , momentum of 0.9, weight decay of 5 × 10 − 4 .
The model is trained with a batch size of 4 for 20 epochs.
Notice one batch includes a support set (with 5 × 5 examples) and a query set(with 15 examples).
In a nutshell, we use the same set of hyperparameters as (Chen et al., 2020b), which is exactly the degenerated model if we set KG layer = 0.
Additional algorithms.
We summarize the steps of MetaLink for relational setting in Algorithm 2, for meta setting in Algorithm 3.
Code release.
We will make all the source code public at the time of publication.
Algorithm 2 MetaLink Training in Relational
Require: Dataset D train =
{ ( x , y )
1: for each iteration do 2: { ( x , { y ( i ) j ∈ T ( i ) aux } , { y ( i ) j ∈ T ( i ) test }} ← SampleMiniBatch
( D train ) 3: { z } ← f θ ( x ) for x ∈ { ( x , { y ( i ) j ∈ T ( i ) aux } , { y ( i ) j ∈ T ( i ) test }} 4:
V (0)
d
= { h (0) i ← z for z ∈ { z }} (cid:46)
Initialize data nodes 5:
V (0) t
= { h (0)
j ← w j for w j ∈ { w j } (cid:46) Initialize task nodes 6:
E = { e ij ← ( x ( i ) , t j ) for y ( i ) j ∈
y ( i ) j ∈ T aux } (cid:46) Initialize edges 7: for l = 1 to L do 8:
V ( l ) d , V ( l ) t ← GraphConv ( V ( l − 1) d , V ( l − 1) t , E ) with f φ 9: logits ← EdgePred ( V ( L ) d , V ( L ) t ) with f φ 10:
Backward (cid:16) Criterion ( logits , {{ y ( i ) j } j ∈ T ( i ) test ) (cid:17)
B A DDITIONAL D ISCUSSIONS
Limitations.
i)
In principle, MetaLink leverages task correlation to achieve data efﬁciency/gain better performance.
Thus, it is less effective when tasks are not correlated.
We believe it is fair to assume a multi-task learning system will work worse when tasks as not correlated.
ii)
Though our method is general to model diverse multi-task scenarios, we limit our empirical study to multi-label learning mostly due to the limitation of dataset availability.
We hope that the proposed MetaLink can inspire researchers to collect more multi-
Require: Dataset D train =
{ ( x , y ) } .
A parameterized embedding function f θ .
A parameterized heterogeneous GNN f φ .
Number of GNN layers L .
1: for each iteration do 2:
S, Q ← SampleMiniBatch ( D train )
(cid:46) Simulate meta setting in training 3: { z } ← f θ ( x ) for x ∈ ( S, Q ) 4:
V (0)
d
= { h (0) i ← z for z ∈ { z }} (cid:46)
Initialize data nodes 5:
V (0) t
= { h (0)
j ← 1 } (cid:46) Initialize task nodes 6:
E = { e ij ← ( x ( i ) , t j ) for y ( i ) j ∈ S } (cid:46) Initialize edges 7: for l = 1 to L do 8:
V ( l ) d , V ( l ) t ← GraphConv ( V ( l − 1) d , V ( l − 1) t , E ) with f φ 9: logits ← EdgePred ( V ( L ) d , V ( L ) t ) with f φ 10:
Backward (cid:16) Criterion ( logits , {{ y ( i ) j } j ∈ T s } ∈ Q ) (cid:17)
Social impact.
In practice, collecting high-quality datasets is often expensive and time-consuming.
In the biomedical domain, labeling requires domain expertise and hence is very resource-intensive.
In some applications, it requires long-term experiments to get the ground truth labels.
The approach that we proposed could reduce the high labeling cost by utilizing auxiliary task labels that are already available.
In principle, MetaLink attempts to leverage correlation among tasks.
Thus, when it comes to future deployment, we suggest being careful about deﬁning tasks.
Because it might come across similar issues related to fairness as in other supervised learning problems, where the algorithm learns biased correlations from the datasets that are not considered to be appropriate.
The datasets we used for experiments are among the most widely-used benchmarks, which should not contain any undesirable bias.
C A DDITIONAL R ESULTS
Advantages on using heterogeneous weights.
As discussed in Section 3, we proposed to learn from the knowledge graph constructed via a heterogeneous GNN.
Here we empirically demonstrate the improvement of using distinct weights for updating data nodes and task nodes.
Results are summarized in Table
[Table] [Table]