Year: 2021

Abstract:
A key assumption in multi-task learning is that test data points have no access to labels of some auxiliary tasks during inference. This presents an opportunity to extend multi-task learning to utilize available auxiliary task labels, 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 novel 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 auxiliary labels. 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 flexibility 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.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 RELATIONAL MULTI-TASK LEARNING SETTINGS
text: 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
}

heading: 3 METALINK FRAMEWORK
text: 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.

heading: 3.1 BUILD A KNOWLEDGE GRAPH ON TASK HEADS
text: 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.

heading: 3.2 LEARN FROM THE KNOWLEDGE GRAPH VIA A HETEROGENEOUS GNN
text: 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 wj , h (0) j = wj . 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 hv based on an iterative
1Without loss of generality, we assume a task is a binary classification or 1-D regression
aggregation of local network neighborhoods. The l-th iteration of GraphConv(·), or the l-th layer, can be written as:
h(l)v = AGG (l) ( {MSG(l)(h(l−1)u ), u ∈ NG(v)},h(l)v ) (1)
where h(l)v is the node embedding after l iterations, h (0) v have been initialized as explained above, andNG(v) denotes the direct neighbors of v. AGG is the acronym for aggregation function and MSG 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 ( CONCAT(h(L)i ,h (L) j ) ) .
In general, MetaLink should work with any GNN architecture that follows the definition in Equation 1. We use the GraphSAGE layer (Hamilton et al., 2017) in MetaLink (W(l), U(l) are trainable):
h(l)v = U (l)CONCAT
( MEAN ( {RELU(W(l)h(l−1)u ), u ∈ N (v)} ) ,h(l−1)v ) (2)
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 define 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:
h(l)v = U (l)CONCAT
( MEAN ( {RELU(W(l)1[v∈Vd,u∈Vt]h (l−1) u + O (l)y(u)v ), u ∈ N (v)} ) ,h(l−1)v ) (3)
where 1[v ∈ Vd, u ∈ Vt] 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 final prediction; this way, the final 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 = { (x(i), {y(i)j }j∈T (i)s ) } and query set{
(x(i), {y(i)j }j∈T (i)aux )→ {y (i) j }j∈T (i)s \T (i)aux
} . We make sure the sampled meta tasks T (i)s and auxiliary
knowledge tasks T (i)aux have no intersection. We first 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 { {y(i)j }j∈T (i)s } in the support set and { {y(i)j }j∈T (i)aux
} (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.

heading: 4 EXPERIMENTS
text: Here we experimentally show that our proposed MetaLink flexibly handles different settings and leverages knowledge from auxiliary tasks. We first 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., 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
Algorithm 1 MetaLink Training in Relational Meta Setting Require: Dataset Dtrain = {(x, y)}. A parameterized embedding function fθ. Last layer weights for
each task wj . A parameterized heterogeneous GNN fφ. Number of GNN layers L. 1: for each iteration do 2: S,Q← SampleMiniBatch(Dtrain) . Simulate meta setting in training 3: {z} ← fθ(x) for x ∈ (S,Q) 4: V (0)d = {h (0) i ← z for z ∈ {z}} . Initialize data nodes 5: V (0)t = {h (0) j ← 1 if meta else wj for each wj} . Initialize task nodes 6: E = {eij ← (x(i), tj) for y(i)j ∈ (S,Q)} . 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 ( Criterion(logits, {{y(i)j }j∈T (i)s \T (i)aux } ∈ Q) )
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.

heading: 4.1 EVALUATING METALINK ON RELATIONAL MULTI-TASK SETTINGS
text: 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 classification in vision. There are 80 binary classification 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 configurations as follows. STANDARD 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. RELATIONAL 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. META 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. RELATIONAL 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 learning to this setting. We first train a feature extractor on the training set. At test time we use the support set to train a linear classification layer.
Results on biochemical datasets. Table 5 summarizes the results on Tox21, Sider, and ToxCast datasets. We first 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-efficient 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.

heading: 4.2 ABLATION STUDIES
text: Does MetaLink truly learn to utilize correlations between tasks? To better understand the improvement of our algorithm, we first plot the Pearson correlation heat map on Sider (Figure 3). We
1https://github.com/Megvii-Nanjing/ML-GCN
then find 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 verifies 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 first few added labels.

heading: 4.3 EVALUATING METALINK ON STANDARD FEW-SHOT LEARNING DATASETS
text: 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.
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 benefits of building a knowledge graph on the last layer. Furthermore, as an ablation study we manipulate the number of KG layers and find 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.

heading: 5 RELATED WORK
text: 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 classification 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 classifiers. 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/classifiers 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.

heading: 6 CONCLUSION
text: 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.
REPRODUCIBILITY STATEMENT
The pseudo-code is provided in the paper. We will make code and data splits publicly available to ensure reproducibility.

heading: A ADDITIONAL IMPLEMENTATION DETAILS
text: 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 flipping. 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 ResNet-12 with an input size of 84×84 (Chen et al., 2020b). We perform standard data augmentation, including random crop and horizontal flipping. 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 Setting Require: Dataset Dtrain = {(x, y)}. A parameterized embedding function fθ. Last layer weights for
each task {wj}. A parameterized heterogeneous GNN fφ. Number of GNN layers L. 1: for each iteration do 2: {(x, {y(i)
j∈T (i)aux }, {y(i) j∈T (i)test }} ← SampleMiniBatch(Dtrain)
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}} . Initialize data nodes 5: V (0)t = {h (0) j ← wj for wj ∈ {wj} . Initialize task nodes 6: E = {eij ← (x(i), tj) for y(i)j ∈ {y (i) j∈Taux} . 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 ( Criterion(logits, {{y(i)j }j∈T (i)test ) )

heading: B ADDITIONAL DISCUSSIONS
text: Limitations. i) In principle, MetaLink leverages task correlation to achieve data efficiency/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-task learning datasets with diverse use cases.
Algorithm 3 MetaLink Training in Meta Setting Require: Dataset Dtrain = {(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(Dtrain) . Simulate meta setting in training 3: {z} ← fθ(x) for x ∈ (S,Q) 4: V (0)d = {h (0) i ← z for z ∈ {z}} . Initialize data nodes 5: V (0)t = {h (0) j ← 1} . Initialize task nodes 6: E = {eij ← (x(i), tj) for y(i)j ∈ S} . 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 ( Criterion(logits, {{y(i)j }j∈Ts} ∈ Q) )
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 defining 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.

heading: C ADDITIONAL RESULTS
text: 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 5.

Review summaries:
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. 

In this paper, they introduce relational multi-task learning in which they construct a graph with the data points and task and labels as edges. They represent the data points as the embedding of NN model and task as the last layer of NN model for that task. Then they solve the link label prediction problem between each node and task by using GNN and heterogeneous message passing method to predict the label of test data point on a new task at inference time.