Learning to Group Auxiliary Datasets for Molecule
Authors: Tinglin Huang
Ziniu Hu
Zhitao Ying
Venue: NeurIPS 2023 poster (2023)
Keywords: molecule
routing mechanism
meta gradient

Abstract: The limited availability of annotations in small molecule datasets presents a challenge to machine learning models. To address this, one common strategy is to collaborate with additional auxiliary datasets. However, having more data does not always guarantee improvements. Negative transfer can occur when the knowledge in the target dataset differs or contradicts that of the auxiliary molecule datasets. In light of this, identifying the auxiliary molecule datasets that can benefit the target dataset when jointly trained remains a critical and unresolved problem. Through an empirical analysis, we observe that combining graph structure similarity and task similarity can serve as a more reliable indicator for identifying high-affinity auxiliary datasets. Motivated by this insight, we propose MolGroup, which separates the dataset affinity into task and structure affinity to predict the potential benefits of each auxiliary molecule dataset. MolGroup achieves this by utilizing a routing mechanism optimized through a bi-level optimization framework. Empowered by the meta gradient, the routing mechanism is optimized toward maximizing the target dataset's performance and quantifies the affinity as the gating score. As a result, MolGroup is capable of predicting the optimal combination of auxiliary datasets for each target dataset. Our extensive experiments demonstrate the efficiency and effectiveness of MolGroup, showing an average improvement of 4.41%/3.47% for GIN/Graphormer trained with the group of molecule datasets selected by MolGroup on 11 target molecule datasets.

Learning to Group Auxiliary Datasets for Molecule
Tinglin Huang'
Ziniu Hu?
Rex Ying'
Yale University, University of California, Los Angeles
Abstract
The limited availability of annotations in small molecule datasets presents chal- lenge to machine learning models. To address this, one common strategy is to collaborate with additional auxiliary datasets. However, having more data does not always guarantee improvements. Negative transfer can occur when the knowledge in the target dataset differs contradicts that of the auxiliary molecule datasets. In light of this, identifying the auxiliary molecule datasets that can benefit the target dataset when jointly trained remains critical and unresolved problem. Through an empirical analysis, we observe that combining graph structure similarity and task similarity can serve more reliable indicator for identifying high-affinity auxiliary datasets. Motivated by this insight, we propose MolGroup', which sep- arates the dataset affinity into task and structure affinity to predict the potential benefits of each auxiliary molecule dataset. MolGroup achieves this by utiliz- ing a routing mechanism optimized through bi-level optimization framework. Empowered by the meta gradient, the routing mechanism is optimized toward maximizing the target dataset's performance and quantifies the affinity as the gating score. As result, MolGrroup is capable of predicting the optimal combination of auxiliary datasets for each target dataset. Our extensive experiments demonstrate the efficiency and effectiveness of MolGroup, showing average improvement 4.41 %6/3.47% for GIN/Graphormer trained with the group of molecule datasets selected by MolGroup on 11 target molecule datasets.
1 Introduction
0.100
Painwise relative improvement
0.075
Predicting and understanding molecular prop- erties fundamental task in biomedical and chemical fields [51, 14, 23, 55], such as eval- uating the toxicity of new clinical drugs, char- acterizing the binding results for the inhibitors of human --secretase, and predicting the ther- modynamic property of organic molecules. In recent years, machine learning-based methods have shown promise in this domain [19, 57, 30, 63, 50]. However, labeling molecules requires expensive real-world clinical trials and expert knowledge, making difficult to collect large and diverse labeled dataset for training. This limited availability of labeled data poses chal- lenge for traditional machine learning methods, which struggle to effectively generalize on small
Tox21
0.050
BACE
Dataset
BB8P
0.025
Target Target
SIDER
ToxCast
ESOL
FreeSolv
-0.050
Tox21
BACE
anop
Chintas
SIDER
lorCast
ESOL
TramSola
amm HI
PCBA
MUV
Chemat
-0.075
-0.10
Auxiliary Dataset
Figure 1: Relative improvement of using the com- bination of target dataset and auxiliary dataset over only using target dataset: (perf(a, b) perf(a))/per((a), perf(a), where a is target dataset and b is auxiliary dataset.
Source code is available
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Measuring
Banking
combination training
datasen
Novel
Dra
Task
Predict ciinical trial toxicity on drugs
Input target dataset rgt
Todisity prediction
Structure
High-afTinily
audilirry dataset
Daur sizun
Task
Predict drug side-effects on marketed drugs
Affenity
datasets
Existing
Structure
MolGroup
0.58
0.79
Daucal mar
Task
Predict barrier permeebility on blood-brain dataset
Threshold
>0.4
0.13
Structure
Figure 2: Overview of our proposed MolGroup to identify the auxiliary datasets with high affinity. (1) Measure the affinity between the given new molecule dataset and the existing molecule datasets in terms of task and structure. (2) Rank the auxiliary datasets and filter them according to predetermined threshold. (3) Train the model with the combination of the selected auxiliary datasets and the target dataset.
datasets. To alleviate this, many methods resort to incorporating additional sources of data to the learning process [19, 6, 29, 41, 57], such augmenting the current drug data with other known drugs of toxicity. This approach can significantly improve performance by introducing useful out-of-distribution knowledge [58, 54] and emphasizing attention on relevant features [12, 38).
Unfortunately, negative transfer can occur when the included data competes for model capacity or when the model fails to learn shared representation [60, 45, 20J. This phenomenon is also evident in our empirical study, as illustrated in Figure 1. We pair each target dataset with every other dataset and measure the relative improvement achieved by the combination'. The results demonstrate how auxiliary datasets positively or negatively affect the target dataset. For instance, ToxCast greatly improves the performance of Clin'Tox, while FreeSolv consistently degrades all other datasets' performance. These findings highlight the presence of underlying affinity between different datasets, where auxiliary datasets with high affinity yield greater benefits". In light of this, how can we measure the alfinity of auxiliary datasets to a larget dataset?
Measured by both structure and task similarity preliminary analysis (Section 3). We compare the similarity between molecule datasets based on task and graph structure separately and examine their correlation with the relative improvement. Our results confirm that combining both task and structure similarity leads to a stronger correlation with relative improvement. However, existing works on task grouping [12, 1, 60, 45] primarily model task similarities. They rely on exhaustively searching or examining the effect of one task on another. Without further incorporating structure affinity, these methods fail to achieve better performance.
clear answer is well supported by our
Present work. In this paper, we focus on designing a dataset grouping method for molecules. Here propose MolGroup, a routing-basee based molecule grouping method. As shown in Fig. 2, MolGroup involves calculating the affinity scores of each auxiliary dataset based on the graph structure and task information, and selecting the auxiliary datasets with high affinity. The selected datasets are then combined and fed into the downstream model. The key idea is to adopt a routing mechanism in the network and learn route between auxiliary-ppecffic optimization frumework. The routing mechanism is optimized toward maximizing the target dataset's performance, and its learned gating score serves the affinity score. This score is calculated by measuring and combining the task and structure affinity between datasets. Our main contributions can be summarized below:
target-specific parameters through bi-level
We provide an empirical study to investigate how different molecule datasets affect each other's learning, considering both task and structure aspects.
We propose MolGroup, molecule dataset grouping method, which outperforms other methods on 11 molecule datasets and achieves an average improvement of 4.41 %/3.47% for GIN/Graphormer.
Mooe experimental details can be found in Appendix C.
'In this paper. "affinity" refers to the impact of one dataset on the performance of another dataset. High- affinity datasets benefit the target dataset, while low-affinity datasets negatively affect the target dataset. It is distinct from "similarity" in this context.

---

2 Related Work
Molecule representation pretraining. Recently, the pretrain-finetume paradigm has been ex- lensively applied to learn the representation of molecules [57, 19, 16, 37, 52, 55]. These meth- ods involve pretraining graph neural network on large-scale molecule dataset through con- trastive [35, 19, 47, 49, 56, 49] or generative [16, 19, 63] proxy tasks, and finetuning the model on specific downstream task. In our experiments, we demonstrate that incorporating high-affinity auxil- iary datasets can significantly improve the performance of pretrained models on downstream target datasets. In addition, our proposed MolGroup can seamlessly integrate with the pretrained models, and the grouping results obtained by the lightweight surrogate model can effectively generalize.
Task grouping. Many applications use multi-task learning [29, 8, 59, 28, 41, 20] reduce the inference time required to conduct multiple tasks individually or enhance the performance of the tasks with limited labels. To avoid negative transfer, previous methods have been proposed to separate tasks with high affinity. Early studies [60, 45, 22] formulate this an integer program problem and solve through optimization. Recent works [12, 44) learn the underlying affinity between different tasks by quantifying the effect of the task's s gradient or active learning strategy. In this paper, we focus on grouping the auxiliary molecule datasets that can maximize the performance of the target dataset. By method can achieve better performance.
incorporating both structure and task information,
While MolGroup and MTG-Net [44] are both related meta-learning approaches for solving grouping problems, the specific objectives differ. MTG-Net formulates the grouping problem as meta-learning problem and learns the relationship between tasks with a few meta- tetta-trinnng instances, while our approach uses the meta gradient to optimize the routing mechanism instead of applying the meta-learning paradigm.
Conditional computation. A plethora of studies applies conditional computation for scaling up model capacity or adapting computation to input [4, 39, 15, 26, 42, 62, 3]. For instance, Mixture-of- Experts layer [42, 26, 36] introduces the routing mechanism to route each token of a sequence to a subset of experts, which increases the number of parameters without incurring too much computational cost. Some methods apply conditional computation to limit the attention operation for speeding up the inference [10, 3, 62, 46). Another line of work in multilingual modeling [61] includes the routing mechanism in each layer to enforce the model to learn language-specific behavior, which can be considered as a simplified neural architecture search framework. Different from these existing works, the routing mechanism in MolGroup aims to endow the model with the ability determine the fusion significance of two datasets' parameters, which is used as the affinity score.
Understanding Relationship between Molecule Datasets
Molecule datasets comprise two critical aspects of information: the structural characteristics of the molccules and the associated predictive tasks. Using this insight, we analyze the relationship between molecule datasets by dividing them into these two dimensions. We conduct study that links the discrepancies in structure and task between datasets to changes in performance when they train together. Molecule structure and predictive task quantified as the distribution of fingerprint [5, 7] and task embedding [1]:
empirical
Fingerprint is multi-hop vector used character molecule's structure, with each bit representing a certain chemical substructure. We first convert all the instances into MACC [11] fingerprint, then obtain the fingerprint distribution of the dataset by computing the occurrence frequency of each bit.
Fingerprint distribution:
Task embedding: We apply Task2vec [1] to quantify task representation which uses Fisher Information Matrix associated with pretrained probe network. GIN [53] is applied as the probe network with extracted atom and bond features as input [18], and the Monte Carlo algorithm is employed to approximate the Fisher Information Matrix, following the suggested implementation.
Considering the asymmetric nature of the impact between molecule datasets, we use asymmetric KI divergence the similarity metric [9, 25]. Structural similarity between datasets is measured by fingerprint distribution similarity, while task similarity is measured by task embedding similarity. We plot the regression curve and calculate the Pearson correlation between relative improvement and
3

---

Pearson coefficient (()
[a]
0.16
0.06
0.36
(b)
Structure coreeation
[c]
Task correlation
0.0
Poarson confficien
0.4
0.6
-0.1
0.4
Relative
0.7
-0.2
0.0
-0.3
0.2
-0.4
Ror22
Bicg
Lito
oreg
Nar21
Becg
Clor
sinER
ESSO
Preestali
Liyeo
-0.4
Freed
Struct. similarity
Task similarity
Struct.+Task
Target dotaset
Drgect Datasct
Figure (a) Regression curves between relative improvement and the measures of structure similarity, task similarity, and their mixing. (b,c) Pearson correlation between relative improvement and similarity of fingerprint distribution/task embedding of each molecule dataset individually.
structural/task similarity over all the combination pairs in Fig.3(a). Additionally, we compute the structural and task similarity between each target dataset and the other 14 datasets individually. We
then calculate the Pearson correlation between the similarity scores and the corresponding relative improvement for each target dataset, shown in Fig.3(b) and (c).
Combination of task and structure leads to better performance. We combine the structure similarity and task similarity for each dataset pair by adding them together, and calculate the Pearson coefficient between the combined similarity and relative improvement in Fi).3(a). This combined similarity demonstrates stronger correlation with the performance compared using the two similarities separately, demonstrating the effectiveness of incorporating both structure and task information. Despite this. the existing methods of multi-task grouping primarily focus on leveraging the underlying interdependence among tasks, without explicitly incorporating the structure feature.
Structure and task are compensatory. Based on Fig.3(b) and (c), can observe that the structure and task correlation exhibit a compensatory relationship. For instance, dataset may present a low structural correlation but high correlation in task similarity, observed in Tox21, where the structural correlation is 0.04, while the task correlation is 0.66. This suggests that these two sources of information contribute to the performance gain in complementary manner, such that the inadequacy of one source does not preclude performance improvement facilitated by the other. This observation also demonstrates that the affinity of an auxiliary dataset to target dataset should be determined by the discrepancies in both structure and task.
Both similar and dissimilar structures and tasks can benefit target dataset. According to Fig.3(b) and (c), 8 out of 11 cases in the structural analysis and 7 out of 11 cases in the task analysis show positive correlation, indicating that the auxiliary datasets with similar structures and predicted tasks can effectively augment the target dataset and improve its performance.
There also negatively related cases, such SIDER in Fig.3(b) and BACE in Fig.3(c), indicating that larger discrepancies in graph structure and task can potentially lead to greater improvement. This observation aligns with prior studies that dissimilar graph structures and tasks can improve the performance in small datasets with limited annotations. These findings confirm that the additional information required from the other sources of data varies across different target datasets, and both similar and dissimilar structures and tasks from the auxiliary datasets can potentially benefit the target dataset.
-f--iiiiibittiion generalization [58, 54] which demonstrate
Based on the above observations, we argue that an ideal criterion for selecting auxiliary datasets should incorporate measures of both structure and task similarity to the target dataset, and balance these two features. In this study, we propose MolGroup, a routing- grouping method that allows us to capture the task and graph structure affinity between two molecule datasets.
4 Grouping Molecule Datasets with MolGroup
In this section, we introduce MolGroup, molecule dataset grouping method. MolGroup includes two modules: (1) routing mechanism used to quantify the affinity between two datasets, and (2)

---

Routing mechanism across layers
Bi-level optimization framework
2] Meta gradient:
Decodert
Structure feature (elter Gaaz
2l route
Penas figge (r rg)
Convolution (45]
Task feature leigh club)
Routting by danoset affintly (frauee)
Transformation
Auxilary speciiic param Bause
Audilary dataset
Apgragation
Aucllary-peellic w params
Tonget-peccffcc paraers
LAXX
Target-pecific param Frut
Encoder
Bryt
Target
1) Ir V Pegcheu
Tarpetopeatt parm
Orge
Zrge
Logtt
Ootaset1
DotasetM
Figure 4: Overview of routing mechanism and li-level optimization framework.
bi-level optimization framework used update the routing mechanism through meta gradient. Unlike previous methods, the proposed routing function can comprehensively measure the affinity from two perspectives: task and graph structure. We formally define these two modules in Section 4. and Section 4.2, and explain our selection process in Section 4.3. More details can be found in Appendix B.
4.1 Routing Mechanism
The objective of MolGroup is to determine the affinity between a target dataset DD and a set of auxiliary datasets {D (uz M- To this end, we propose allocates the impact of each auxiliary dataset on the target dataset across the network's 's sub-layers. Intuitively, this routing mechanism allows us control the contribution of the auxiliary datasets to the final predictions, ensuring the incorporation of the high-affinity auxiliary datasets while mitigating any interference. We first formulate the graph convolution process of GNN at 1-th layer as:
the routing mechanism that dynamically
(1+1)
(1) (z)
(1)
where (1++) is the representation the input batch (-th layer, and ( (.) denotes the convolution function with parameter A at I-th layer. As illustrated in Fig. 4, in MolGroup, every dataset is assigned a specific parameter 0y, ... A.. and we apply the routing mechanism to the auxiliary dataset's convolution process while keeping the convolution process of the target dataset unchanged. Specifically, given an input batch of the target dataset By and m-th auxiliary dataset Bm, the convolution with routing mechanism for m-th auxiliary dataset 1-th layer is calculated as:
((+1) Z m
am Or (2!!)
p(0) (2()) + (1 - (1-aam) fm (z!!))
(())
(2)
with
Olm = Jm (BT, B.n)
(3)
where 2!! is the representation of m-th auxiliary dataset at 1-th layer, and Om is the gating score generated by the routing function (m).). The routing function endows the model with the capability of determining the impact of the auxiliary dataset on the target dataset's parameter. A closed gate suggests that the backward gradient has minor impact on 0r, particularly when am 0. indicating no interaction between the auxiliary dataset and the target dataset. Conversely, an open gate encourages the gradient from the auxiliary dataset to operate on 0r, particularly when a. indicating a hard-parameter sharing architecture.
The calculation of g(.) is parameterized using two modules to capture task affinity and structure affinity. For task affinity, we assign learnable embeddings e and e for the target dataset and auxiliary dataset respectively, which updated through optimization. For structure affinity, we first extract the fingerprint feature 'P for each molecule instance in the input batch. Then, we embed using a weight matrix and apply Set2Set [48] to obtain single embedding snict.
struct. = Set2Set ((W=*)))
(4)
where Win a learnable weight matrix, and n the number of the instance the batch. Task affinity score and structure affinity score are computed as the cosine similarity and fused the dataset affinity
5

---

Algorithm 1: Iterative filtering process with MolGroup
Input: Target dataset DDr with N training instances; Candidate auxiliary datasets {DA)M: GNN model with specific parameters 0r, {fm] M for each dataset and routing mechanism g(-); Number of iterations R; Number of epochs E; Learning rate ir; Batch size B.
// Filtering round
for r + 1, do
Random initialize Pr,01,.... Oar and g(-).
01,... ΦM &- Affinity scores
I N NM// // Number of iteration in each epoch
I1 Training epoch
for e < 1,...
do I/ Training step
for iter t- 1. . do
Sample
BI. BB from current datasets. Obtain losses /T, {m J and affinity scores {am by feeding mini-batch to GNN.
// Bi-level optimization framework
1) Or (famm))as + OT - Ir. Ver Σm Im-
2) Update g(-) through target dataset's loss function with By and 0r (famm })M- Update all the parameters 0r, 01, AM through IT. 41, 1 ".M.
ife == E then
1/ Average affinity scores in final epoch
for + 1.... M do 6m t pm + am/L
end
end
end end
// Remove datasets according to threshold
(DAIM* (DAm. 1m 3 2 0.6).
M M+M'.. M'.
end
Output: Auxiliary datasets with high affinity {DA)mr M.
score, which can be formulated as:
Om 9m (Br, Bm)
(Aelek cjask (1 A)e
Cy
struct.
where (-) is the eegistic-sigmodd function, and is a hyperparameter used to balance these two affinity scores. It can be found that the task affinity score is determined globally during training, while the structure affinity score is computed at a per-step level of granularity. To ensure stable learning, a high value for À is suggested. Further analysis can be found in Appendix E. We average the dataset affinity score over a continuous subset of training steps to obtain a final affinity score.
4.2 Bi-level Optimization Framework
The learning of the routing mechanism depends on how the auxiliary dataset's gradient affects the target dataset's performance. However, optimizing this routing mechanism target datasct-aware way challenging since it only used during the forward pass the auxiliary dataset. To address this, we propose to bi-level [34] optimization framework to incorporate the guidance of the target dataset. The framework utilizes the target dataset's 's performance, using the parameters updated by the auxiliary dataset, signal guide the learning of the routing mechanism.
Equipped with the routing mechanism, the overall optimization can formulated as:
min Ly (Br; 0r), 0r
(6)
min Lm (Bmi Br, Bm, asm). By 8
(7)

---

where LT and Lm denote the loss function of the target dataset and auxiliary dataset. It can be found that the optimization of 0r is closely linked to the auxiliary task and its relevance is determined by the routing function. We explicitly represent such dependency as @r(mm). With the goal of obtaining an optimal gating score Omm parameterized by 9m((), given M auxiliary datasets, we further formulate the optimization as bi-level framework:
min LT(BT; OT (Ham)ar))
(8)
M
where 0r((am3m)) = arg min (B Or, On Om),
Σ m
(9)
By
As illustrated in Fig.4, the update process of the routing function is split into two steps: first, the model parameters except the routing function are updated using the gradient from the auxiliary tasks; second, we reuse this computation graph and calculate the meta gradient of the routing mechanism.
4.3 Iterative Filtering and Selection
We adopt an iterative filtering process to perform the selection of auxiliary datasets. Initially, all datasets are trained together in a model equipped with the routing mechanism for several epochs. At the final epoch, we average the generated affinity score across all training steps and take it as the final dataset affinity. Datasets with affinity score below the predefined threshold are then filtered out. This process is repeated for é iterations, and in the last round of filtering, the top-k remaining auxiliary datasets are selected based on their affinity scores. Finally, these selected datasets are combined with the target dataset and used for training the model. The pseudo-code is presented in Algo.
Complexity. Given M candidate auxiliary datasets, combine them with the target dataset and use MolGroup for grouping. At each training step. we sample data from each dataset with equal probability form a mini-batch, ensuring that each dataset contributes equally the training process. The total number of training instances is set as N x M to cover all instances in the target dataset, where N the target dataset size. Therefore, the training complexity of each round is O(NM). This approach is efficient as the number of auxiliary datasets is typically filtered, and further analysis is provided in Section 5.2.
Experiments
To show the effectiveness of MolGroup, we apply it 15 molecule datasets with varying sizes and compare it with 8 baseline models. Moreover, we conduct some further studies including an ablation study and efficiency comparison, and conclude some findings regarding the grouping results. The statistics of the involved datasets are summarized in Appendix A, the experimental setting is detailed in Appendix C, and the extensive experimenta results can found in Appendix D.
5.1 Dataset Grouping Evaluation
Baselines. We compare MolGroup with three classes of approaches: 1) search-based methods, 2) grouping-based methods, and 3) the methods that train on all the datasets. Regarding the search- based methods, we apply beam search [32, 13] and consider two kinds of criteria, i.e., the target dataset's 's performance on the validation set (P), and the combination of the performance and the difference of fingerprint distribution (P+S). We select the candidates with the highest criterion value and train the model for a few epochs at each selection step to speed up the search process. As for the grouping-based methods, apply TAG [12] and Task2vec [1], which calculates the pairwise affinity using the gradient-based strategy. For the last class of methods, we consider Unweighted Averages(UA), Gradnorm [8], and MTDNN [27), where MTDNN selects subset of datasets through the task discriminator. Furthermore, we present the results of Pretrain-Finetuu strategy [19, 57], where the model is first trained on PCQM4Mv2 [17] and then finetuned on the downstream dataset.
Dataset. Our study utilizes 15 molecule datasets of varying sizes obtained from MoleculeNet [51, 18] and ChemBI. [33], which can categorized into three groups: medication, quantum mechanies, and chemical analysis. Our focus is solely on the small molecule datasets that have less than 10,000
7

---

instances in their training sets, totaling 11 target datasets. We follow the original split setting, where qm8 and qm9 are randomly split, and scaffold splitting is used for the others.
Architectures. All the model architectures include a standard encoder-decoder, where GIN [53] is used as the encoder in MolGroup and other baseline methods that explicitly group datasets. We
evaluate the selected auxiliary datasets using GIN and the SOTA model Graphormer [57]. To
accommodate the different tasks, we use different decoder for each one. The overall training loss is calculated as the unweighted mean of the losses for all included tasks.
Method
BBBP(1)
ClinTox(1)
Tox21(1)
BACE(1)
FreeSolv(L)
qm8(L)
Only-target
66.62p.ud 028
56.456.023
74.23p.m 0o5
75.02g.ul 0206
3.8421.572
0.0385p.uur
Beam search(P)
66.02..om 015
37.86g.068
74.71p.00 004
67 67.34g. 0.09
3.271o.479
0.0553g.00
Beam search(P+S)
67. 67.69p.oo 034
57. 57.630.030
74.36p. 003
69. 09.740.0 7.006
3.331 0.287
0.0494..om
TAG
60.66p.014
57. 57.980.028
70.02..or 70.32u.u06
3.922b.74s
0.0637p.um
Task2vec
68. 18p.8p.011
47 47-300.031
68.00p.005
74.71 0.002
3.3830.766
0.0635o.om
MTDNN
66.56p.021
52.90g.039
71.87
69.9 19.91g.026
3.428p.733
0.0523
UA
60.41g.00s
51.990.078
68. 16.006
61.75g.018
4.095p.334
0.0625g.001
Gradnorm
61.21p.007
53. 53.080.070
59. 59.40p.041
64.83g.028
4.356p.589
0.0657 006
Pretrain-Finetune
56.59p.om 026
56.00p.037
50.64 50.64g.015
64.80
4.391g.048
0.0637
MolGroup
68.360.016 59.770.027
75.66g.004
77 03.01
3.1166.279 0.0385p.oor
3. 3.1166
Table 1:
Performance comparison of GIN on target molecule datasets, with 1 indicating higher is
better and + indicating lower is better.
Due to the space limit, here we present the comparison results of datasets in Table 1. Our results show that MolGroup outperforms all the baseline methods and consistently improves the performance of the backbone model, with average relative improvement of 6.47% across all datasets. We observed that MolGroup assigned low-affinity scores to all candidate auxiliary datasets for qm8 and qm9, resulting in selection of auxiliary datasets for those two target datasets. We attribute this phenomenon to the significant difference between quantum chemistry and other domains. We also observe that UA, Gradnorm, and Pretrain-Finetune methods perform even worse than the model trained only on the target dataset due to the significant distribution shift among these datasets. In addition, the search-based methods' performance constrained by the limited exploration space, and task grouping methods neglect the data structure difference, leading to suboptimal performance.
Method
BBBP(1)
ClinTox(t)
Tox21(1)
BACE(1)
FreeSolv(L)
qm8(l)
Only-target
66.40p.019 019
77.590.028
75.970.009
78.91..022 02.023
2.004o.088
0.0372
Beam search(P)
67.53p.010
76.770.117
76.61
81.58g.03s
2.129
0.0473g.001
Beam search(P+S)
67. 67.15p.0 030
77.700.076
76.83..00
80.68o.0 0211
2.3120.191
0.0458..00
TAG
69.62n. 008
78.080.036
76.57p.006
80.10p.or
2.038p.178
0.0537
Task2vec
66.93p.um
72.92o.019
75.150.007
80.84..o
1.941..129
0.0530p.um
MTDNN
71 71.970.078 66.71g.039
76.84p.008
79.79o.001
2.1730.321
0.0512o.nm1
UW
65.37p.000
80.52..027
72.16p.009
75.64g.000
2.405p.415
0.0569
Gradnorm
60.40g.045
43. 43.460.083
55. 55.06p.002
48.00g.005
2.552p.505
0.16944.014
MolGroup
69.66 009 81.21o.010
77.3 77.340.006
82.94o.017
1.879..03. 03.
0.0372
Table 2: Performance comparison using Graphormer target molecule datasets.
Additionally, we apply the resulting groups to evaluate the impact of dataset combination using Graphormer as the encoder [57]. We first pretrain Graphormer on PCQM4Mv2 [17], which is the largest graph-level prediction dataset, and then finetune on the specific downstream datasets with the selected auxiliary datasets. As shown in Table 2, the extensive parameter space and pretraining offer significant improvement to Graphormer, consistent with the previous studies [57, 43). Moreover, the dataset combinations found by MolGroup outperform all the baseline methods and provide average relative improvement 3.35% across all the target datasets, demonstrating the generalizability of the auxiliary dataset groupings exploited by the lightweight surrogate model GIN.
8

---

1004
TAG
HTDNN
KaGrup
Task2vnc
Beaa search
(a]
Metaa
(b)
0.8
0.3
0.7
12500
Aftitty
0-2
9.6
2020
10000
7500
0.1
04
5000
0.0
3000
2500
Training Epoch
Traninng
0
Grouping Method
Grouping Method
Figure 5: (a) Comparison between MolGroup with and without meta gradient on BBBP. (b) Efficiency comparison between MolGroup and the other grouping methods.
5.2 Further Analysis
Overall selected results. In Fig. 6, we visualize the auxiliary datasets with top-3 affinity scores to the target dataset measured by MolGroup. PCBA is selected by most of the datasets due to their diverse structure useful out-of-distribution information. This phenomenon is supported by Fig. 1, where PCBA demonstrates significant benefits across datasets. We also observe that Tox21 benefits ClinTox and ToxCast, particularly for toxicity-related tasks. Additionally, despite belonging distinct domains, some datasets exhibit a great affinity to the target dataset, such as qm8 for BBBP and ESOL for Lipo. These findings demonstrate that MolGroup can effectively reveal the underlying affinity between molecule datasets.
MUV
ToxCast
qma
Lipo
BBBP
Tox21
Tox21
Auxiliany dataset
FreeSolv
BACE
Target dataset
ToxCast
ESOL
PCRA
ClinTox
BEeP
Ablation study. We conduct an ablation study verify the impact of the bi-level optimization framework. Specifically, the routing mechanism is solely updated by the auxiliary dataset's task without the guidance of the target dataset. We pick BBBP the target dataset, and average the gating scores a'm every epoch, which illustrated in Fig. 5(a). The routing mechanism tends to assign low weight to the auxiliary-specific parameters (<0.15 15 most cases) to increase the dataset-specific capacity, failing to differentiate the affimity of the auxiliary dataset to the target dataset. Conversely, the meta gradient can optimize the routing mechanism toward maximizing the performance of the target dataset, which learns distinguishable affinity score distribution. Similar phenomena can be found the other datasets.
FreeSolv
HIV
SIDER
SIDER
Lipo
ESOL
Figure 6: Selected grouping for each target dataset.
Efficiency Comparison. We compare the wall-clock time performance of MolGroup with the other grouping methods using the same rumning environment (see Appendix D). We pick BBBP and Tox21 the target datasets, and the results presented in Fig. 5(b). Beam search explores combinations in a greedy manner, but its computational complexity increases exponentially as the search space grows. TAG efficiently computes inter-task affinity by averaging the lookahead loss proposed during training. MTDNN includes all auxiliary datasets and utilizes a task discriminator to select instances, which limits its computational performance based the number of training instances used. Although we initially input all candidate auxiliary datasets to MolGroup, we only train for few epochs at each round, and the majority of them are filtered out after the first second round of training. Specifically, only 6 and 5 out of 14 candidates remain after the first round of filtering for BBBP and Tox21 respectively.
Comparison with Task2vec: similar benelt. In Section 3, we conclude that task similarity measured by task embedding correlated with relative improvement. But Task2vec [1] fails to achieve better performance in our experiments since it doesn't 't incorporate structure information. It likely that the datasets with similar tasks can serve as data augmentation and are able benefit each other during training. We choose Tox21, ToxCast, and ClinTox as examples since they share similar tasks of predicting the absence of toxicity or qualitative toxicity measurement. them pairwise and observe that half of the cases result in negative transfers. To delve deeper into this phenomenon, we average the structure affinity score measured by MolGroup over every epoch and plot the results for positive and negative transfer cases in Fig. 7(a) and (b) respectively.
combine

---

(lock TeCant)
FbaCan Positive
- Kalrbe RoCast
(b)
Ted:CCilo
Negative
Mescae Ginke!
(Cirfii bo23)
(c)
TaCall
0.60
0.44
0.42
0.40
50 N W de un 0.
0.50
0.38
Affiiity scores
0.45
0.36
0.40
0.34
0.35
0.32
0.30
0.30
0.28
1
2
3
6
7
8
9
5
6
7
Training Eooch
Training Epoch
Treining Epoch
Figure 7: (a.b) Structure affinity score of positive and negative cases measured by MolGroup. (c)
Inter-task affinity score measured by TAG on BBBP.
Our result reveals that the structure affinity consistently converged low score during training on negative transfer cases, suggesting a fundamental mismatch between the target task and the auxiliary dataset. The results also demonstrate that the combination of task and structure affinity can lead to more comprehensive measurement, which is consistent with the insights presented in Sec. 3.
Comparison with TAG: differential affinity scores perform better. TAG [12] is the SOTA method in task grouping; however, fails to select high-affinity auxiliary datasets in our cases. Here plot the inter-task affinity scores measured by TAG for BBBP in Fig. 7(c). Compared with MolGroup's 's results shown in Fig. 3(a), we observe that most datasets' affinity scores produced by TAG are homogeneous, centered around O. Even the dataset with the highest affinity score (FreeSolv) the negative auxiliary dataset to BBBP, as demonstrated in Fig. 1. One reason is that the lookahead loss of every dataset pair used in TAG trained together a shared architecture. In contrast, the specific parameters assigned to each dataset in MolGroup alleviate this issue and result in more stable training.
be easily influenced by the other dataset when they
PCBA is an effective booster. One interesting finding that dataset PCBA [17] can boost the performance of most of the small molecule datasets, as shown in Table 3. This dataset offers both a diverse range of chemical compounds with unique scaffold structures, comprising over 350,000 training instances, and an extensive collection of 128 bioassay annotations that represent broad range of biological activities, making a potent booster for small molecule property prediction tasks that can benefit from both structure and task.
BBBP(t) ClinTox(f) ToxCast(1) Tox21(t) ESOL(L)
FreeSolv(L)
Lipo(l)
Only-target 66. 62..222 56.45o.023 60.69o.010 74.234 74.23p.00s
1.563p.040
3.8421.579
0.8063p.015
+PCBA 67. 1o.025 57.77o.028
62.05p.007 74.81p.00e 1.463p.020
3.563p.989
0.8021o.009
Table 3: Cases whose performance is improved by PCBA.
6 Conclusion
A common strategy to improve the performance of small molecule datasets is to incorporate additional data during training. In this paper, we conduct an empirical study to investigate the relationship between molecule datasets and conclude that an ideal criterion should combine graph structure and task discrepancies. Motivated by this, we propose MolGroup which quantifies the affinity between datasets by a routing mechanism that is optimized by the bi-level optimization framework. We evaluate our method on 15 molecule datasets and compare it with 8 baseline models, and the results demonstrate the effectiveness of MolGroup.
Limitation. MolGroup operates at dataset level and does not consider subdataset-leww information. As result, it cannot fully exploit the potential benefits obtained from specific subdatasets within large dataset. This limitation could impact the performance of MolGroup on certain datasets, such as qm8 and qm9, where MolGroup fails to group high- auxiliary datasets.
10

---

References
[1] Alessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran, Subhransu Maji. Charless C Fowikes, Stefano Soatto, and Pictro Perona. Task2vec: Task embedding for meta-lcaming. In Proceedings of the IEEE/CVF international conference on computer vision, pages 6430-6439, 2019.
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.0644 2016.
[3] Ankur Bapna, Naveen Arivazhagan, and Orhan Firat. Controlling computation versus quality for neural sequence models. arXiv preprint arXiv:2002.07 106, 2020.
[4] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: arXiv:/.00..3433 2013.
[5] Alice Capecchi, Daniel Probst, and Jean-Louis Reymond. One molecular fingerprint to rule them all: drugs, biomolecules, and the metabolome. Journal of cheminformatics, 12(1):1-15, 2020.
[6] Fabio Capela, Vincent Nouchi, Ruud Van Deursen, Igor V Tetko, and Guillaume Godin. Multitask learning on graph neural networks applied to molecular property predictions. arXiv preprint 1910.11110.11/// 2019.
[7] Adrià Streto-Massagué, María José Ojeda, Cristina Valls, Miquel Mulero, Santiago Garcia-Vallvé, and Gerard Pujadas. Molecular fingerprint similarity search in virtual screening. Methods, 71:58 63, 2015.
[8] Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normal- ization for adaptive loss balancing in deep multitask networks. In International conference on machine learning, pages 794 803. PMLR, 2018.
[9] Carl Doersch. Tutorial variational autoencoders. arXiv preprint arXiv:1606.05908 2016.
[10] Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 190911990..111 1/56, 2019.
[11] Eli Fernández-de Gortari, César R Garefa-Jacas, Karina Martinez-Mayorge and José L. Medina-Frannce Database fingerprint (dfp): an approach to represent molecular databases. Journal of cheminformatiis, 9(1):1-9, 2017.
[12] Chris Fifty, Ehsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, and Chelsea Finn. Etficiently identifying task groupings for multi-task learning. Advances in Neural Information Processing Systems, 34:27503-2755 2021.
[13] Markus Freitag and Yaser Al-Onaizan. Beam search strategics for neural machine translation. arXiv preprint arXiv:7702.018000 00066 2017.
[14] Anna Gaulton, Louisa J Bellis, A Patricia Bento, Jon Chambers, Mark Davies, Anne Hersey, Yvonne Light, Shaun McGlinchey, David Michalovich, Bissan Al-Lazikani, et al. Chembl: a large-scale bioactivity database for drug discovery. Nucleic acids rexearch, (OODDI:DI100-DIIUI
2012.
[15] Felix A Gers, Jürgen Schmidhuber, and Fred Cummins. Learning to forget: Continual prediction with Istm. Neural computation, (4100)2441124711 2000.
[16] Zhenyu Hou, Xiao Liu, Yukuo Cen, Yuxiao Dong, Hongxia Yang, Chunjie Wang, and Jie Tang. Graphmae: Self-supervised masked graph autoencoders. In Proceedings the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 594,604, 2022.
[17] Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Isc: large-scale challenge for machine learning on graphs. arXiv preprint 2rXiv:2103.0944 103.09430, 2021.
[18] Weihua Hu, Matthias Fey, Marinka Zimik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118-22131 18-22133, 2020.
[19] Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitmik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint arXiv: 1905. 12265, 2019.
[20] Ziniu Hu, Zhe Zhao, Xinyang Yi, Tiansheng Yao, Lichan Hong. Yizhou Sun, and Ed Chi. Improving multi-task generalization via regularizing spurious correlation. Advances in Neural Information Processing Systems, 35:114 1146, 2022.
11

---

[21] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference machine learning, pages 448-456. pmir, 2015.
[22] Zhuoliang Kang, Kristen Grauman, and Fei Sha. T.carning with whom to share in multi-task fcature learning. In Proceedings of the 28th International Conference Machine Learning (ICMI,1 I), pages 521-528, 2011.
[23] Sunghwan Kim, Paul A Thiessen, Evan E Bolton, Jie Chen, Gang Fu, Asta Gindulyte, Lianyi Han, Jane He, Siqian He, Benjamin A Shoemaker, et al. Pubchem substance and compound databases. Nucleic acids research, 44(D1):D1202 D1213, 2016.
[24] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412::441.66980 2014.
[25] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv: 1312.61 14, 2013.
[26] Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh-Thang Luong. and Orhan Firat. Beyond distillation: Task-level mixture-of -f-eeertss for efficient inference. arXiv preprint arXiv:2110003442. 2021.
[27] Po-Nien Kung, Sheng-Siang Yin, Yi-Cheng Chen, Tse-Hsuan Yang, and Yun-Nung Chen. Efficient multi-task auxiliary learning: selecting auxiliary data by feature similarity. In Proceedings of the 2021 Conference on Empirical Methods Natural Language Processing, pages 416-428, 2021.
[28] Xi Lin, Hui-Ling Zhen, Zhenhua Li, Qing-Fu Fu Zhang, and Sam Kwong. Pareto multi-task learning. Advances in neural information processing systems, 32, 2019.
[29] Shengchao Liu, Meng Qu, Zuobai Zhang, Huiyu Cai, and Jian Tang. Structured multi-task learning for molecular property prediction. In International Conference Artificial Intelligence and Statistics, pages 8906-8920. PMLR, 2022.
[30] Shengchao I Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre-training molecular graph representation with 3d geometry. arXiv preprint arXiv:21 2021.
[31] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization.
arXiv preprint
arXiv:1711.05100 2017.
[32] Clara Meister, Tim Vieira, and Ryan Cotterell. beam search is the answer, what was question? arXiv preprint 2tti:200000065500 2020.
[33] David Mendez, Anna Gaulton, Patrícia Bento, Jon Chambers. Marleen De Veij, Eloy Félix, María Paula Magariños, Juan F Mosquera, Prudence Mutowo, Michat Nowotka, et al. Chembl: towards direct deposition of bioassay data. Nucleic acids research, (7(D1):D930-DD)
2019.
[34] Hieu Pham. Zihang Dai, Qizhe Xie, and Quoe V Le. Meta pseudo labels. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11557-11568, 2021.
[35] Jiezhong Qiu, Qibin Chen, Yuxiao Dong, Jing Zhang, Hongxia Yang. Ming Ding, Kuansan Wang, and Jie Tang. Gee: Graph contrastive coding for graph neural network pre-training. Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pages 1150-1 160, 2020.
[36] Maksim Riabinin and Anton Gusev, Learning@ home: Crowdsoureed training of large neural networks using decentralized mixture-of-experts arXiv preprins arXiv:/2002.04013 2020.
[37] Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenhing Huang, and Junzhou Huang. Self- supervised graph transformer on large-scale molecular data. Advances Neural Information Processing Systems, 33:12559- 2020.
[38] Sebastian Ruder. An overview of multi-task learning in deep neural networks. arXiv preprint (:XXvv/706.00000 2017.
[39] Hasim Sak, Andrew W Senior, and Françoise Beaufays. Long short-term memory recurrent neural network architectures for large scale acoustic modeling. 2014.
[40] Andrew M Saxe, James I. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6110, 2013.
[41] Ozan Sener and Vladlen Koltun. Musti-task learning as multi- optimization. Advances in neural information processing systems, 31, 2018
12

---

[42] Noam Shazeer, Azalia Mirhoseini. Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.0653 2017.
[43] Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie-Yan Liu. Benchmarking graphormer on large-scale molecular modeling datasets. arXiv preprint axXvv2203.008810, 2022.
[44] Xiaozhuang Song, Shun Zheng, Wei Cao, James Yu, and Jiang Bian. Efficient and effective multi-task grouping via meta learning on task combinations. In Advances in Neural Information Processing Systems, 2022.
[45] Trevor Standley, Amir 7amir, Dawn Chen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi-task learning? In International Conference on Machine Learning. pages 9120-9132. PMLR, 2020.
[46] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin, Adaptive attention span in transformers. arXiv preprint 1XXi::199000799 2019.
[47] Ruoxi Sun. Hanjun Dai, and Adams Wei Yu. Does gnn pretraining help molecular representation? Advances in Neural Information Processing Systems, 35: 12:1609--, 2022.
[48] Oriol Vinyals, Samy Bengio, and Manjunath Kudlur. Order matters: Sequence to sequence for sets, arXiv preprint 1xxv:1511.06391,
[49] Yuyang Wang, Jianren Wang. Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence, 4(3):299-877, 2022.
[50] Oliver Wieder, Stefan Kohlbacher, Mélaine Kuenemann, Arthur Garon, Pierre Ducrot, Thomas Seidel, and Thierry Langer. A compact review of molecular property prediction with graph neural networks. Drug Discovery Today: Technologies, 37:1-12, 2020.
[51] Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu. Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 9(2):513 530, 2018.
[52] Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li, Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molccular representation for drug discovery with the graph attention mechanism. Journal of medicinal chemistry, 63166)8779--776, 2019.
[53] Keyulu Xu. Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint arXiv: 1810.0086. 2018
[54] Jingkang Yang, Kaiyang Zhou, Yixuan Li, and Ziwei Liu. Generalized out-of-distribution detection: A survey. arXiv preprint
2021.
[55] Kevin Yang. Kyle Swanson. Wengong Jin, Connor Coley, Philipp Eiden. Hua Gao, Angel Guzman-Perrs Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling, 5968)3337--3388, 2019.
[56] Zhen Yang, Tinglin Huang, Ming Ding, Yuxiao Dong, Rex Ying, Yukuo Cen, Yangliao Geng, and Jie Tang. Batchsampler: Sampling mini-batches for contrastive learning in vision, language, and graphs. arXiv preprint arXiv:2306.03355 2023.
[57] Chengxuan Ying. Tianle Cai,, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877-28888, 2021.
[58] Junchi Yo, Jian Liang, and Ran He. Mind the label shift of augmentation--assd graph ood generalization. arXiv preprint arXiv:2303. 14859, 2023.
[59] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsca Finn. Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems, 33:5824-5836. 2020.
[60] Amir R Zamir, Alexander Sax, William Shen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 372227722,222 2018
13

---

[61] Biao Zhang, Ankur Bapna, Rico Sennrich, and Orhan Firat. Share not? learning to schedule language- specific capacity for multilingual translation. 2021.
[62] Biao Zhang, Ivan Titov, and Rico Sennrich. On sparsifying encoder outputs in sequeence-to-sequeeee models. arXiv preprint arXiv:2004. 11854, 2020.
[63] Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: universal 3d molecular representation learning framework. 2023.
14

---

A
Dataset Details
Our study utilizes 15 molecule datasets of varying sizes obtained from MoleculeNet [51, 18] and ChemBL [33], which can be categorized into three groups based their applications source domains: medication (HIV, SIDER, Lipo, Clin'Tox), quantum mechanics (qm8, qm9), and chemical analysis (ESOL. FreeSolv, ChemBL, MUV, BACE, BBBP, ToxCast, Tox21, PCBA). As for qm8 and qm9, we randomly sample 3,000 graphs to construct the datasets. We use the original split setting, where qm8 and qm9 are randomly split, and scaffold splitting is used for the others. The small molecule datasets with less than 10,000 instances in the training sets are selected as target molecule datasets, i.c., Tox21, BACE, BBBP, ClinTox, SIDER, ToxCast, ESOL, FrecSolv, Lipo, qm& and qm9. All the involved datasets can be accessed and downloaded from OGB4 or MoleculeNet repository". The overall statistics summarized as follows:
HIV
PCBA
Tox21
BACE
BBBP
ClinTox
MUV
#graphs
41,127
437,929
7,831
1,513
2,039
1,477
93,087
#tasks
128
12
2
17
split
scaffold
scaffold scaffold scaffold scaffold scaffold scaffold
metric
ROAUC
AP
ROAUC ROAUC ROAUC ROAUC
AP
Table 4: Dataset statistics(I).
SIDER
ToxCast ChemBL
ESOL
FreeSolv
Lipo
qm8
qm9
#graphs #tasks
1,427
8,576
456,309
1,128
642
4,200
3,000
3,000
617
1,310
12
split
scaffold
scaffold scaffold scaffold scaffold scaffold random random
metric
ROAUC ROAUC ROAUC RMSE RMSE
RMSE
MAE
MAE
Table 5: Dataset statistics(2).
Here we list brief description of each dataset's task in this study:
HIV: Test the ability of compounds inhibit HIV replication, i.e., estimate whether the molecule is inactive or active.
PCBA: Predict 128 molecular properties of each compound, where all the properties are cast as binary labels.
Tox21: Predict the toxicity of each compound on 12 different targets.
BACE: Predict the binding results (either binding or non-binding) for collection of inhibitors targeting human 8-secretase
BBBP: Predict the permeability properties of drugs across the barrier between circulating blood and the brain's extracellular fluid.
ClinTox: Classify a collection of drug compounds by predicting their (1) clinical trial toxicity (or absence thereof) and (2) FDA approval status.
MUV: Evaluate 17 molecule properties for around 90 thousand compounds specifically designed to validate virtual screening methodologies.
SIDER: Predict the drug side effects 27 system organ classes.
ToxCast: Predict the toxicity of the compound on over 600 targets.
ChemBL: Predict the molecular properties of a set of bioactive molecules.
ESOL: Estimate the water solubility of each molecule.
FreeSolv: Estimate the hydration-free energy the water of each molecule.
Lipo: Evaluate the lipophilicity of the drug molecule.
qm8/qm9: Estimate the energetic, electronic, and thermodynamic properties of the molecule.
hhttpa://ogb.stanttte stanford edu/ 'nttps://moleculeenee org/
15

---

B Implementation Details
Backbone model settings. As for GIN [53], we fix the batch size as 128 and train the model for 50 epochs. We use Adam [24] with learning rate of 0.001 for optimization. The hidden size and number of layers are set as 300 and 5 respectively. We set the dropout rate as 0.5 and apply batchnorm [21] in each layer. All the results are reported after 5 different random seeds.
As for Graphormer [57], we fix the batch size as 128 and train the model for 30 epochs. Adam [31] with learning rate of 0.0001 is used as the optimizer. The hidden size, number of layers, and number of attention heads are set as 512, 5, and 8 respectively. We set the dropout rate and attention dropout rate as 0.1 and 0.3. Layernorm [2] is applied across layers. The maximum number of nodes and the distance between nodes in the sampled graph set as 128 and respectively. The size of position embedding, in-degree embedding, and out-degree embedding are fixed as 512. All the results are reported after different random seeds.
Grouping method settings. As for MolGroup, we apply GIN the encoder. During the training,
we randomly select datasets with equal probability and sample data from them to ensure that each
dataset contributes an equal number of samples to the constructed mini-batch. We use MACC [11] fingerprint to calculate the structure affinity score. We fix the number of the filtering iterations as 3 and the balance coefficient is set as 0.9. During each round, the auxiliary datasets with affinity scores below 0.6 will be filtered out. If none of the datasets fall below this threshold, we filter out the dataset with the lowest affinity score. The orthogonal initialization [40] is applied to initialize the learnable task embedding eaack to make sure that the dot product between the target dataset' embedding and auxiliary dataset's 's embedding starts with O. The task embedding size set as 16, and the number of the processing steps in Set2Set is set as 2. The pseudo-code is presented in Algo.
As for beam search, we apply GIN the encoder. The beam width and search depth are both set to 3. During the search process, we train the model for 3 epochs using each candidate dataset and evaluate its performance using the validation set loss. Additionally, we consider criterion based on the combination of the difference in fingerprint distribution and performance. Specifically, we average and normalize these two metrics to determine the criterion.
As for Task2vec [1], we employ GIN as the probe network and follow the official implementation First, fix the encoder and train the decoder for 10 epochs. Then we apply the Monte Carlo algorithm to compute the Fisher Information Matrix.
As for TAG [12], we use GIN as the encoder and follow the official implementation?. TAG involves training all the datasets and computing the lookahead loss between target dataset and auxiliary datasets. The lookahead loss is accumulated over multiple epochs and used as the affinity scores.
As for MTDNN [27], we train all the datasets together for each target dataset and apply an additional task discriminator to classify the source of the dataset for the input instances. We train 50 epochs for GIN and 30 epochs for Graphormer, and the instances in auxiliary datasets that have probability greater than 0.6 of being classified as the target dataset will be selected.
As for Gradnorm [8], we train all the datasets together and update the weights of the loss based on the gradients of the last shared GNN layer.
C Experimental Details
For the preliminary analysis shown in 1, we conduct the study using a set of 15 molecule datasets. Among these datasets, 9 datasets with less than 10,000 instances in the training set are selected as target datasets. We pair the target datasets with every other dataset and measure the relative improvement the combination achieves. To mitigate the issue of varying dataset sizes, upsample or downsample the training sets of all datasets to ensure an equal number of training instances, specifically 5,000 instances. All the reported results are based 5 different random seeds.
For the dataset grouping evaluation, we train the model using the combined datasets and assess its performance on the target datasets. We then report the model's performance on the test set using the
'nttps://github /ttrpar//gihau...oomamsss//////s task2vec
'https://githin httetthhh.o.omeeeeeeeeeeeeeeeeee google- rewearrheeeecceeeeeeeeeer wowooll ctweeerreeeeeeeeeeeeeeeeeeaeeeee htwpr//rerrrrrrerrr////////ee/ee
16

---

best best-performing model selected based on its performance on the validation set. Cross-entropy loss used for classification tasks and mean squared eтror loss is used for regression tasks. The overall training loss is calculated as the unweighted mean of the losses for all included tasks. All the reported results are based on different random seeds.
D Overall Experimental Results
Running environment. The experiments are conducted on a single Linux server with The Intel Xeon Gold 6240 36-Core Processor, 361G RAM, and 4 NVIDIA A100-40GB. Our method is implemented on PyTorch 1.10.0 and Python 3.9.13.
D.1 Dataset Grouping Evaluation
Here we present the performance comparison over the other 5 target datasets in Table and Table It can be observed that pretrained Graphormer outperforms GIN significantly, consistent with the previous studies. In addition, MolGroup achieves the best performance most cases, with average relative improvement of 3.52% and 3.10% for GIN and Graphormer. However, a notable exception occurs with qm9 where our proposed method is unable to surpass beam search. In this instance, MolGroup assigns low-affinity scores to each auxiliary dataset due to the significant disparity between quantum chemistry and the other domains, as shown in Section D.3. Nonetheless, despite the limited efficiency, beam search is capable of identifying promising candidate by directly comparing the performance of different groupings. achieves better performance compared to the other baseline methods.
worth noting that overall, our proposed MolGroup still
Additionally, we attribute the poor performance or even worse results of the Unweighted Average, Gradnorm, and Pretrain-Finetune methods to the significant distribution gap between different datasets. These methods struggle learn a shared representation that can effectively capture the characteristics of all the datasets.
Method
SIDER(1)
ToxCast(t)
ESOL(L)
Lipo(J)
qm9(()
Only-target
59.35p.010
60.69g.010
1.563p.040
0.8063
0.0303g.00
Beam search(P)
54.25g.024
63.37o.004
1.431
0.8073g.011
0.0456g.001
Beam search(P+S)
56.58
61.01o.010
1.476p.061
0.8147
0.0258o.om
TAG
55.64g.0m
58.080.003
1.417p.066
0.81 181777
0.0453g.com
Task2vec
55.74g
57.450.00
1.436p.050
0.8078
0.0468g.o0 .000
MTDNN
55.95p.07
59.020.007
1.499p.u27
0.8155g.017
0.04766.00
UA
56.26g.m 009
59.930.000
1.480p.oo 060
0.9130g.017
0.0494o.co 000
Gradnorm
55.69o.00s
53.31o.010
1.541..091
1.0755g.001
0.0574g.000 0000
Pretrain-Finetune
51.43g.009
51.066.005
1.480 099
1.07760.086
0.0498
MolGroup
59. 59.21o.or 0133
63.9 63.91o.005
1.402 037
0. 0.7996..oom
0.03033.o0
Table 6: Performance comparison GIN on target molecule datasets, with + indicating higher
better and J indicating lower is better.
D.2 Takeaway
Grouping more high-affinity datasets improves performance. Previous studies on task grouping have assumed that low-order relationships can be an effective indicator of high-order ones. It suggests that combining multiple source datasets, which can benefit the target dataset, leads to better performance when learned together. Our experimental results also confirm this phenomenon. We take eight datasets as examples and add them one by one with the top three datasets having the highest affinity score as measured by MolGroup. Results are presented in Table 8, where Top/ a, b. b.. denotes the combination of auxiliary datasets with top-performing ones. We can find that combining more auxiliary datasets leads to better performance in most cases. Besides, training with high-affinity datasets can significantly reduce the variant of FreeSolv (1.579 + 0.279), indicating more robust representation learned from the auxiliary datasets.
17

---

Method
SIDER(T)
ToxCast(1)
ES
Lipo(!)
qm9(L)
Only-target
62.050.021
66. 66.16p.004
1.054g.0m 0553
0.7432o.usz
0.02730.u01
Beam sParch(P)
60.80g.005
67.36p
0.989
0.7640p.om
0.0399
Beam search(P+S)
62. 62.670.009
66.00 003
1.026g.00 024
0.7524p.01r
0.0262p.000
TAG
63.57o.004
65.40p.um
1.0150.045
0.7507p.014
0.0432o.ure
Task2vec
60.699.019
63.28p
0.997 0.28
0.75620.0z
0.0429..oo
MTDNN
61 61.430.024
65.05p.022
1.045g.003
1p.7160.032
0.0420p.000
UW
58. 58.240.014
62. 01.00
1.120g.084 0884
0.7887 0.067
0.05080.000
Gradnorm
51.65g.018
52. 52.40p.007
1.400g.099
1.0482
0.21544.
MolGroup
63.75 63.750.018
68.68p.oo
0.978p.047
0. 0.73044.018
0.02730.001
Table 7: Performance comparison using Graphormer target molecule datasets.
Combinations
ClinTox
Tox21
FredSolv
BBBP
BACE
ToxCaet
ESOL.
Lipo
Only Target
56. 050.
74. 23,000
3. 3.89
66 06.62g.0 028
75.02o. 028
60. 0100
1.5630.040
0.80630.015
+Top!1]
74.81o.006
3. 3.563p.989
67.36p 023
71.277 045
62. 680 002
1.502..045
0. 0.8096p.014
1,p[1.2
57. 480 032
75. 220-03
3. 977
68. 68.64n
71.13g
6d. 180
1.5246.077
0.8078p.011
-Top(1.2.3)
59. 770.027
75. 66g 004
3.1 3.116u 279
68 08.36g.018
77. 77.33p.0 015
63. 63.91p.ous
1.4022.010
0.7996p.oo ous
Table 8: Performance of top-k combinations.
MolGroup is the hard version of routing mechanism. Here we investigate the use the routing mechanism in the training of the final model. Specifically, we train the encoder with the routing mechanism on all the auxiliary datasets and adjust the influence between datasets instead of iterative filtering. as shown in Table. 9. It can be observed that the model trained without the negative auxiliary datasets (i.e., MolGroup) outperforms significantly the model trained with negative datasets (i.c., Final model with 9(-)), although when it is equipped with routing mechanism. Besides, MolGroup can be considered hard version of the final model including the routing mechanism. Rather than incrementally eliminating negative training signals, it directly filters out negative auxiliary datasets.
BBBP( ToxCast(t) Tox21(t)
ESOL(1)
FreeSolv(L)
Only-target
66. 62...2.. 60.69o.om
74.230.005
1.5630.040
3. 3.8421.579
MolGroup
68.36p.016
63.91o.00s
75.666.004
1.402a.037
3. 3.1166.2790
Final Model with 9()
66.83p.a30
57.06p.o11
69.37o.o18
3.040b.1384
5.483).3647
Table 9: Comparison with final model trained on all auxiliary datasets.
D.3
Case Study
To give an intuition of the learning process of MolGroup, we show the learning curves of the affinity scores in different filtering rounds Fig. 8. In each round, auxiliary datasets with affinity scores below 0.6 are removed, and if none of the datasets fall below this threshold, the dataset with the lowest affinity score is filtered out. It can be observed that a significant number of auxiliary datasets are removed in the first or second round. Furthermore, the learning curves tend to converge after epochs in most cases. The dataset PCBA remains in the final round most cases, indicating its general benefit to the other target datasets. We also notice that the majority of auxiliary datasets are assigned high-affinity scores for FreeSolv, as demonstrated in Fig. 1, suggesting that all auxiliary datasets contribute positively to its performance.
Additionally, we plot the learning curves for qm8 and qm9, which don't have suggested auxiliary datasets. As shown in Fig. 9, all the auxiliary datasets are assigned affinity scores lower than the threshold for both qm8 and qm9, resulting in the removal of all datasets after the first round. We attribute this to the large discrepancy in the structure and task between quantum chemistry and the other domains such as medication. Molecules in quantum chemistry have diverse structures, including both natural and hypothetical compounds. They are studied to explore the behaviors of various functional groups. On the other hand, molecules in medication are more focused on structure. Their study revolves around predicting and optimizing molecular properties relevant to drug design.
18

---

BAC4
Chanis
Frdda
Uge
R0A
BBBP 1sz Round
2nd
B8BP Round
Tax21 Round
Tox21 2nd
Tox21
0.9
0.8
0.9
1.0
0.9
0.9
Aifinity
0.5
0.0
0.8
0:7
0.7
0.5
0.6
0.0
0.66
0.7
0.55
0.
0.6
0.7
0.4
0.3
0.5
0.3
0.3
0.2
0..4
6
8
6
8
8
6
8
6
8
Taining Epoch
Training Epoch
Training Epoch
Training Epoch
Training Epoch
Triiinn Epoch
EX0L
BACE
hrddon
and
Upe
PCSA
2001
uxe
PL4 Lip
Round
ClinTox Round
Lipo
2nd
0.9
0.9
10.9
1..9
0.8
0.8
0.8
88 0.7
0.7
30 0.5
0.6
0.6
0.5
0.5
0.4
0.3
0.2
6
8
0.3
6
6
8
6 Training Esoch
6 Training Spoch
6
Taining Epoch
Traniin Epoch
Training Epoch
Taaninn Epach
May
Beza
and
Charil
Tocae
2337
Caribe
Use
PCBA
FreeSolv 1st Round
FreeSolv 2nd Round
FreeSolv 3rd Round
ToxCast 1st Round
ToxCast 2nd Round
ToxCast 3rd Round
1.0
0.8
0.8
0.8
0.8
0.4
0.6
0.7
0.5
0.3
0.6
0.4
0.5
0.6
0.5
0.7
0.5
0.6
0.7
0.4
0.3
0.3
0.2
Training Epoch
Training Epach
Taining Epoch
Training Epoch
Training Epach
Figure 8: Learning curves of affinity scores where red dashed line represents threshold.
Tox21 qm8 qm9
Upp
ToxCast
MUV BACE
FreeSolvv
ChemBL PCBA
SIDER
Clinfox
B8BP
qm8 1st Round
qm9 1st Round
0.60
0.65
0.55
0.60
Affinity SCOPE
0.50
0.55
0.45
0.50
0.40
0.45
0.35
0.40
0.30
0.35
0.25
0.30
0.20
0.25
1 2
6 7 8
9
1
6 Training Epoch
7
Training Epoch
Figure 9: Learning curves for qm& and gm9,
E Parameter Analysis
F.1 Analysis on balance coefficient A
To analyze the impact of the balance coefficient X, we vary À in the range of {0.9, 0.,, and pick BBBP, Tox21, and Lipo as examples. The learning curves of the affinity scores corresponding to different values of À are plotted in 10. From the results, we observe that decreasing led to lower affinity scores for the auxiliary datasets, which fails to effectively discriminate the affinity of individual auxiliary datasets. One reason for this is the instability introduced by parameter initialization and the per-step level computation of structure affinity scores. As a result, lower values of X cause the MolGroup to assign lower affinity scores to stabilize the training of the target dataset. In light of this, we suggest setting X to high value, i.e., 0.9.
E.2 Analysis number of filtering rounds R
As shown in Fig.8, the auxiliary datasets with the highest affinity scores change as we progress through the different filtering rounds. To investigate the impact of the filtering rounds, we test the performance of the top-3 auxiliary datasets selected in each rounds. The comparison results are shown in Table 10 and the selected datasets are illustrated Fig. 11. We observe that the auxiliary datasets selected in the 3rd round exhibit the best performance. This is attributed to the filtering
19

---

HEeSeN
000
Chaeel
JB3
8B6P A-0.9
BBBP A-0.7
B88P A-0.5
BBBP A-0.3
0.6
0.7
0.7
0.9
0.5
0.5
Wfinity
0.6
0.5
0.3
0.2
0.3
0.2
0.3
0.1
0.1
0.0
0.0
0.0
4
Training
Training Epoch
Training Epoch
Training
Lipo
Am0.5
1.0
0.8
0.9
0.8
0.6
0.5
0.7
0.4
0.6
0.3
0.5
0.4
0.3
0.2
0.2
0.2
0.1
0.0
0.0
2
4
6
6
6
Training Epoch
Training Epoch
Training Epech
Training Epoch
Tox21 Am0.9
Tox21 A=0.7
Tex21 Am0.5
Tox21 A=0.3
0.9
1.0
0.8
0.6
0.5
0.6
0.4
0.4
0.4
0.3
0.3
0.1
0.0
6 8
Training Epoch
Training Epoch
Training Epach
Training Epach
Figure 10: Learning curves of affinity scores with different A.
process in the 1st and 2nd rounds, which removes negative datasets and helps alleviate interference.
As result, MolGroup can estimate the afñinit of each auxiliary dataset more accurately, leading to improved performance on the target dataset.
BBBP
Tox21
ClinTox
Lipo
FreeSolv
ToxCast
R 67. 67.94..ms
75.666.004
57. 77.666.o2z
0.8013p.on
3.2288p.543
58.82p.007
R 67. 018
75.660.001
57. 17.666.ozz 027
0. 0.8020o.co
3.2280o.cr 627
63.47p.om
88.366.016
75.66% 004
59. 00.777..22 027
0. 0.79966.om 007
3.1160.277
03..91.005
Table 10: Performance with different number of filtering rounds R.
R
R =
MUV
MUV
Tox21
FOBA
Tox21
Tex
Tox21
ToxCast
PCBA
SIBCR
Clin Tox
HIV
HM
ClinTox
Tox21
FreeSelh
ToxCast
BBBP
H800
Liço
Tox
FreeSolv
ESOL
BACE
Lipa
Lipo
GinTex
BACE
Figure 11: Auxiliary datasets with top-3 affinity scores with different R where highlight the different edges these two rounds.
20

---

Broader impact
Impact on machine learning research. We propose a novel strategy that combines the routing mechanism with the meta gradient to quantify the impact of one dataset another. Previously, the routing mechanism was used increase the model capacity. Our proposed framework can inspire various extensive applications in machine learning. including neural architecture searching (NAS) and data-centric AI. Specifically, our framework enables the network to determine the optimal routing path through the meta gradient. This empowers the network to control and modify the layerwise architecture, leading to improved performance. Besides, has potential to enhance data-centric Al approaches by providing tool to analyze and understand the relationship between different datasets or sub-datasets.
Impact on biological research. We investigate the relationship between molecule datasets, and, focusing on both the structural and task dimensions. our findings shed light on how different molecule datasets impact each other. The analytical approach we employ be extended to other biological data, such as protein and RNA. Furthermore, our method has the potential to be utilized for data instance filtering in the biological domain. This is particularly important given that biological data often contains various types of noise. By filtering out data instances with negative effects on improve the quality and reliability of biological data used in research.
downstream tasks,
Impact on the society. Our study has significant implications for society, particularly in the field of biomedieine and drug discovery. By understanding the relationship between molecule datasets and their impact on each other, we can gain insights into how different molecules interact and influence biological processes. It can aid in the development of more effective drugs and therapies by identifying molecules that have a positive impact on specific biological targets or diseases.
21


Reviews:
[{"review_id": "zGGeawGTBk", "reviewer": "Reviewer_B1V2", "summary": "This paper investigates how different molecule datasets affect each other’s learning, considering both task and structure aspects. It proposes a routing-based molecule grouping method to calculate the affinity scores of each auxiliary dataset based on the graph structure and task information, and select the auxiliary datasets with high affinity. The selected datasets are then  combined and fed into the downstream model. Experiments show a large improvement for GIN/Graphormer trained with the selected group of molecule datasets.", "questions": "If the proposed method is model-agnostic?", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. This paper studies how to select additional auxiliary molecule datasets to improve the prediction performance on the target dataset. Annotated molecule dataset is difficult to obtain and naively introduce auxiliary molecule datasets may result in negative transfer. This paper investigates an interesting problem how to properly select auxiliary molecule datasets. The authors design a dataset grouping method for molecules, which considers both task and structure aspects.\n\n2. Design a routing mechanism to quantify the affinity between two datasets and bi-level optimization framework used to update the routing mechanism through meta gradient. The proposed routing function can comprehensively measure the affinity from two perspectives: task and graph structure. The bi-level optimization uses the parameters updated by the auxiliary dataset as a signal to guide the learning of the routing mechanism.\n\n3. Experimental results how strong improvement of the baseline methods. The proposed method outperforms all the baseline methods and consistently improves the performance of the backbone model, with an average relative improvement of 6.47% across all datasets. Extra analysis of the proposed method is presented.", "weaknesses": "1. The authors did not study how auxiliary molecule datasets impact different models. If the proposed method is model-agnostic?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable comments on our paper! Here we additionally present the performance of GCN (#layer=2, hidden dim=300, dropout rate=0.5) with different grouping methods in the following table, where a consistent improvement can be observed. Such a phenomenon also aligns with the performance of GIN and Graphormer shown in Tables 1 and 2 in our paper, demonstrating that the grouping can be transferred to different models. We appreciate your suggestion and will conduct a more comprehensive evaluation in our future work.\n\n| GCN    | BBBP(↑) | ToxCast(↑) | Tox21(↑) | ESOL(↓)  | FreeSolv (↓)|\n| -------- | ------- |  ------- | ------- | ------- | ------- |\n| Only-target   |   0.6309(±0.0136)|\t0.6253(±0.0086)|\t0.7427(±0.0061)|\t1.4720(±0.0277)|\t3.3941(±0.2063)|\n| TAG   |   0.6159(±0.0123)|\t0.6164(±0.0037)|\t0.6996(±0.0058)|\t1.5028(±0.0423)|\t2.7482(±0.2939)|\n| MolGroup  |   0.6369(±0.0084)|\t0.6290(±0.0062)|\t0.7527(±0.0054)|\t1.3745(±0.0297)|\t2.5496(±0.2700)|"}, {"review_id": "luoITIgacG", "reviewer": "Reviewer_VvkN", "summary": "The paper proposed MolGroup, a dataset grouping method designed to aid molecule property prediction. Motivated by preliminary empirical analysis, MolGroup separates the dataset affinity into task and structure affinity, and uses a routing mechanism to quantify the affinity between a pair of datasets. The routing mechanism is optimized through a bi-level optimization framework.  Experiments on 11 target molecule datasets show that MolGroup yields a ~4% increase across two architectures.", "questions": "1. Could you elaborate/reiterate on how you train the final model? Why do you not use the routing mechanism in your final model (I suppose it is due to performance issues)?\n2. Please address the quality and significance issue stated above.", "limitations": "I do not see any significant, unreported negative societal impact.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "**Originality**\n\nThe separation of task and structure affinity is novel. The application of bi-level optimization for quantifying dataset affinity is novel.\n\n**Clarity**\n\nThe paper is well-structured and easy to follow.", "weaknesses": "**Quality**\n\nIn Fig. 3a, I don't think we can draw the conclusion that structure and task affinities are compensatory. If we remove the outliers, the points seem randomly distributed, which means that the affinities and the relative improvements are not correlated. This phenomenon might significantly undermine the subsequent arguments.\n\n**Significance**\n\nI agree that molecular property prediction is a very important task with limited data. But if putting this much effort and data result in a mere 4% performance increase, how do you convince the research community that we shall continue on this direction?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for the insightful comments! We will address your concerns point by point.\n\n```\nQ1: In Fig. 3a, I don't think we can draw the conclusion that structure and task affinities are compensatory. If we remove the outliers, the points seem randomly distributed, which means that the affinities and the relative improvements are not correlated. This phenomenon might significantly undermine the subsequent arguments.\n```\n\n**Response**: Thanks for your suggestions! The outliers are the cases of dataset pair with significant improvements or degradations, such as the FreeSolv dataset with others. Here we exclude all the datasets exhibiting huge performance changes, i.e., FreeSolv, qm8, and qm9, and the regression curves are included in the attachment. It can be observed that the combination of structure and task information still exhibits a stronger correlation compared to the use of the individual one, which is consistent with our analysis shown in the paper. \n\nBesides, the advantage of applying both task and structure affinity can be further demonstrated by the case studies in Section 5.3, where the structure affinity score explains why datasets with similar tasks cannot benefit from each other. We will add more analysis in the revised version.\n\n```\nQ2: I agree that molecular property prediction is a very important task with limited data. But if putting this much effort and data result in a mere 4% performance increase, how do you convince the research community that we shall continue on this direction?\n```\n\n**Response**: Actually improving the small molecule datasets with a 4% gain is non-trivial and challenging, which can also be observed in the previous state-of-the-art methods[1,2]. These previous methods employ large parameter space and additional information (i.e., 3D structure and large pretrained dataset) to achieve a relative improvement of 4%-6%, comparable to ours.\n\nMoreover, our proposed method has the additional advantage of efficiency compared with search-based methods, as illustrated in Figure 5 (b). The grouping results are model-agnostic and hold for various backbone models, as demonstrated by Tables 1 and 2. We would like to highlight that the model trained with the grouping dataset can still benefit from other techniques like pre-training.\n\n[1] Zhou, G., Gao, Z., Ding, Q., Zheng, H., Xu, H., Wei, Z., ... & Ke, G. (2023). Uni-Mol: a universal 3D molecular representation learning framework.  \n[2] Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W., & Huang, J. (2020). Self-supervised graph transformer on large-scale molecular data. *Advances in Neural Information Processing Systems*, *33*, 12559-12571.\n\n```\nQ3: Could you elaborate/reiterate on how you train the final model? Why do you not use the routing mechanism in your final model (I suppose it is due to performance issues)?\n```\n\n**Response**: We first merge all the selected auxiliary datasets together and train an initial model on this merged dataset. Note that the initial model, i.e., the final model, doesn’t employ the routing mechanism which is only utilized for auxiliary grouping datasets. The data is sampled from each dataset with equal probability during the training.\n\nOur proposed routing mechanism indeed can be used to train the final model on all the datasets and adjust the influence between datasets in an end-to-end manner. But such a method suffers from two limitations:\n\n1. Training a model with more datasets will consume significantly more computational resources and time. \n2. Our empirical study shows that the model trained without the negative auxiliary datasets outperforms significantly the model trained with negative datasets, although when it is equipped with a routing mechanism. It can be demonstrated in the following table:\n\n| GIN | BBBP(↑) | toxcast(↑) | tox21(↑) | esol(↓) | freesolv(↓) |\n| --- | --- | --- | --- | --- | --- |\n| Only-target | 0.6662(±0.0284) | 0.6069(±0.0102) | 0.7423(±0.0057) | 1.5635(±0.0408) | 3.8421(±1.5796) |\n| Molgroup | 0.6836(±0.0163) | 0.6391(±0.0058) | 0.7566(±0.0044) | 1.4028(±0.0372) | 3.1166(±0.2790) |\n| Final Model with routing mechanism | 0.6683(±0.0307) | 0.5706(±0.0114) | 0.6937(±0.0180) | 3.0400(±0.1383) | 5.4833(±1.3547) |\n\nMolGroup can be considered a **hard version** of the final model including the routing mechanism. Rather than incrementally eliminating negative training signals, it directly filters out negative datasets. Furthermore, a final model trained with the routing mechanism fails to yield grouping results that can be transferred to other models."}, {"review_id": "pS0GGPltdb", "reviewer": "Reviewer_NGKd", "summary": "The paper addresses the challenge of limited annotations in small molecule datasets and proposes a method called MolGroup to identify auxiliary datasets that can benefit the target dataset when jointly trained. MolGroup utilizes a routing mechanism optimized through a bi-level optimization framework to separate dataset affinity into task and structure affinity. The proposed method demonstrates its effectiveness in predicting the optimal combination of auxiliary datasets for each target dataset and outperforms existing baselines.", "questions": "Please see the weaknesses", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides a clear motivation for the problem of limited annotations in small molecule datasets and the challenges associated with incorporating auxiliary datasets. This highlights the practical relevance of the research.\n2. The proposed MolGroup method is well-explained and builds on the insights obtained through empirical analysis. Particularly, the preliminary study on the relative improvement and the similarities is interesting.\n3. The extensive experiments demonstrate the efficiency and effectiveness of MolGroup.", "weaknesses": "1. The Pearson coefficients presented in Figure 3 are relatively low, all below 0.5. This raises doubts about the claim that the combination of task and structure leads to better performance. There is a potential risk of negative transfer that could negatively impact the main task's performance. It would be better to provide further insights or explanations to address this concern.\n2. Given that the authors propose to use meta learning to strengthen the learning process, it is suggested that an ablation study be conducted to demonstrate the effectiveness of this approach. Comparing the performance with and without meta learning would provide a clearer understanding of its contribution to the proposed MolGroup method.\n3. In line 187, the authors propose to assign learnable embeddings for the tasks. Why not use Task2vec to generate the task embedding, as is done in Section 3? The reasoning behind this choice is not adequately explained. It would be beneficial for the authors to provide a justification for this decision and discuss any potential implications or advantages of using learnable embeddings over Task2vec.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your feedback on our work! We will explain your concerns point by point.\n\n```\nQ1: The Pearson coefficients presented in Figure 3 are relatively low, all below 0.5. This raises doubts about the claim that the combination of task and structure leads to better performance. There is a potential risk of negative transfer that could negatively impact the main task's performance. It would be better to provide further insights or explanations to address this concern.\n```\n\n**Response**: The single metric(fingerprint/task similarity) indeed doesn’t correlate strongly with the performance gain. However, **the primary purpose** of this analysis is to demonstrate that integrating both structural and task information yields a stronger correlation with performance gains rather than proposing a novel metric to predict affinity.\n\nBesides, compared with these single metrics, the proposed MolGroup can better capture the affinity between datasets in terms of task and structure, as demonstrated in Table1,2 and Fig.7(a). Combining both sides of information can lead to a more comprehensive understanding of the affinity between datasets. We will add more analysis in the revised version.\n\n```\nQ2: Given that the authors propose to use meta learning to strengthen the learning process, it is suggested that an ablation study be conducted to demonstrate the effectiveness of this approach. Comparing the performance with and without meta learning would provide a clearer understanding of its contribution to the proposed MolGroup method.\n```\n\n**Response**: Thanks for your suggestion! We have presented the learning curves of the affinity scores without utilizing the bi-level framework in Fig.5(a). The result shows a homogeneous distribution among the auxiliary datasets, indicating that the model struggles to distinguish the affinity of the auxiliary dataset without bi-level training. We will conduct more analysis to provide a more comprehensive understanding in future work.\n\n```\nQ3: In line 187, the authors propose to assign learnable embeddings for the tasks. Why not use Task2vec to generate the task embedding, as is done in Section 3? The reasoning behind this choice is not adequately explained. It would be beneficial for the authors to provide a justification for this decision and discuss any potential implications or advantages of using learnable embeddings over Task2vec.\n```\n\n**Response**: As demonstrated in Fig. 3(a) and Tables 1 and 2, Task2vec exhibits a weak correlation with the performance gain (0.05 Pearson coefficient). Additionally, the performance of the model trained with the auxiliary datasets selected by Task2vec is poor. We will add more discussions in the revised version."}, {"review_id": "xc2lHgFyol", "reviewer": "Reviewer_ohh4", "summary": "The authors address the problem of, in a transfer learning, meta-learning, or few-shot learning setting involving molecules, which datasets might be most useful in providing positive auxiliary information that does not damage model performance on the task of interest. \n\nThe paper proposes a method, MolGroup, to identify the datasets which, if provided as auxiliaries to a model, most increase the score of a target dataset. The method makes use of a routing mechanism to select the optimal auxiliaries, and demonstrates improvement of both GIN and Graphormer models on some of 11 target datasets. ", "questions": "A number of questions are raised in the \"weaknesses\" section above. ", "limitations": "The authors do not discuss limitations. ", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "\nThe method is novel and well-motivated, with routing mechanism and bilevel optimization which has not previously been applied in this setting. \n\nThe method does find auxiliary datasets which provide a relative improvement on the target datasets, where the gains are small but nonetheless present across the board. In addition the method is not computationally infeasible, and the mention of efficiency of the method as measured by wall clock time is valuable. \n\nThe experiments are comprehensive and carefully performed, taking into account SOTA modelling techniques when comparing final performance. ", "weaknesses": "It is rather unclear initially whether the authors propose to calculate an affinity score based upon calculated task embedding and fingerprint distribution differences between datasets, or learn an affinity score as the gating score. Line 65-68 in particular are unclear on this point. Clarification would be very useful at this point in the manuscript. \n\nThe assertion that fingerprint and task embedding similarity are strictly different measures. In practice, task embeddings are highly dependent on the distribution of input features, regardless of labels, and therefore are rather similar to fingerprint embedding similarities. An earlier justification for this assertion in the manuscript would be very helpful. \n\nWhile Figure 3 demonstrates that these correlation measures between auxiliary and target dataset are themselves correlated with relative improvement of performance, it is not especially convincing that the two correlation measures are distinct. For instance, some tasks show positive correlation for one measure and negative for another, others do not. It is not clear how to interpret this information. Expansion around these plots would be valuable. \n\nIt would be useful to see whether the method works on another domain outside of molecular property prediction -- have the authors considered this? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your feedback on our work. We will address your main concern point by point.\n\n```\nQ1: It is rather unclear initially whether the authors propose to calculate an affinity score based upon calculated task embedding and fingerprint distribution differences between datasets, or learn an affinity score as the gating score. Line 65-68 in particular are unclear on this point. Clarification would be very useful at this point in the manuscript.\n```\n\n**Response**: We apologize for any confusion and will make a revision in the next version. In our work, we propose to learn the gate scores with a routing mechanism to quantify the affinity between datasets. Thank you for bringing this to our attention.\n\n```\nQ2: The assertion that fingerprint and task embedding similarity are strictly different measures. In practice, task embeddings are highly dependent on the distribution of input features, regardless of labels, and therefore are rather similar to fingerprint embedding similarities. An earlier justification for this assertion in the manuscript would be very helpful.\n```\n\n**Response**: Thank you for the suggestion! Actually, the task embedding is obtained using the GIN which applies extracted node and edge features rather than directly utilizing fingerprint features. It can inherently minimize the correlation between fingerprint and task embedding similarity, a fact evidenced by the substantial difference in their Pearson coefficients (0.16 vs 0.06). We will include this discussion in the next version to prevent any confusion.\n\n```\nQ3: While Figure 3 demonstrates that these correlation measures between auxiliary and target dataset are themselves correlated with relative improvement of performance, it is not especially convincing that the two correlation measures are distinct. For instance, some tasks show positive correlation for one measure and negative for another, others do not. It is not clear how to interpret this information. Expansion around these plots would be valuable.\n```\n\n**Response**: Thanks for your suggestions! Here are some examples using asymmetric KL divergence as our similarity measurement:\n\n1. For the datasets Esol and Freesolv, the fingerprint similarity is 0.0212, while the task similarity is 0.4084. Their tasks are both related to water solubility but structural distributions vary substantially.\n2. For the datasets Tox21 and SIDER, the fingerprint similarity is 0.8852, while the task similarity is 0.0062. The result reveals that the two datasets share similar structure distribution, yet their tasks differ significantly.\n\nWe will incorporate more discussion in the revised version. Besides, we have included some case studies in Section 5.3, which shows the learned structure affinity scores among the datasets with the toxicity-related task (Tox21, ToxCast, and ClinTox). \n\n```\nQ4: It would be useful to see whether the method works on another domain outside of molecular property prediction -- have the authors considered this?\n```\n\n**Response**: Thanks for your suggestion! We believe that MolGroup can be potentially extended into other biomedical domains, such as protein and single-cell data, where data distribution varies and label annotations are costly. We will keep exploring it in our future work.\n\n```\nLimitations: The authors do not discuss limitations.\n```\n\n**Response**: The limitation is included in the Conclusion section. We apologize for any oversight and will ensure greater clarity in the revised version."}, {"review_id": "31X6SkJMXy", "reviewer": "Reviewer_sjEd", "summary": "This paper studies an interesting and meaningful problem, that is, how to make the best use of available molecule data for achieving superior performance on the prediction of molecule properties. This paper first conducts an extensive study on many widely used benchmark datasets and discovers some interesting patterns. Then the authors propose a routing-based method that considers both task similarity and molecule similarity to predict the best combination of auxiliary datasets for each target dataset.", "questions": "1. What does $\\mathbb{z}$ in Eqn. (1) mean? Is it the representation for a batch or the whole dataset?\n2. Why the learnable paratmeter of $g_m(\\cdot,\\cdot)$ is not in Eqn. 7?\n3. Is the proposed method sensitive to the threshold in Sec. 4.3 ?\n4. In addition to task similarity and structural similarity, what are the other possible factors determining whether an auxiliary task is helpful or not? I hope the authors can discuss this.\n5. What is the difference between search-based methods and group-based methods (line 237)?\n6. As shown in Fig. 6, FreeSolv and ESOL are not quite helpful to each other. This seems strange --- FreeSolv is about Hydration Free Energy (HFE) and ESOL is about water solubility. Theoretically, these two properties should be strongly correlated as HFE is an important factor in water solubility. Do you have any comments on this result?", "limitations": "N.A.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The studied problem is interesting and meaningful because it is a very often case in the drug discovery pipeline. \n2. The paper is clearly written and easy to follow.\n3. The empirical study is well-designed and conducted, and the proposed method is well-motivated.", "weaknesses": "1. Computation cost will be high if there are many tasks involved.\n2. Some datasets studied are essentially formed by multiple properties (e.g. Tox21, Clintox). It would be better if these properties are treated as one single task, instead of considered together.\n3. This paper studies many datasets, and therefore it would be better if the authors can give a brief introduction to each of the datasets (e.g. what property is the dataset about?). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your valuable comments on our paper. We will explain your concerns point by point.\n\n```\nQ1: Computation cost will be high if there are many tasks involved.\n```\n\n**Response**: It is indeed true that handling numerous datasets, especially when including large auxiliary sets, can consume significant computational resources. However, our proposed method can alleviate this by automatically grouping the datasets with the routing mechanism in a single model. This approach leads to improved efficiency compared with the other baseline models, as demonstrated by Fig.5(b).\n\n```\nQ2: Some datasets studied are essentially formed by multiple properties (e.g. Tox21, Clintox). It would be better if these properties are treated as one single task, instead of considered together.\n```\n\n**Response**: Thanks for your great suggestions! The current model only considers dataset-level information so it cannot fully exploit the sub dataset-level information in terms of task or structure. We will explore it in our future work.\n\n```\nQ3: This paper studies many datasets, and therefore it would be better if the authors can give a brief introduction to each of the datasets (e.g. what property is the dataset about?).\n```\n\n**Response**: Good suggestion! The introduction can be found in [MoleculeNet](https://moleculenet.org/). We will include it in the next version. \n\n```\nQ4: What does 𝕫 in Eqn. (1) mean? Is it the representation for a batch or the whole dataset?\n```\n\n**Response**: It represents the representations of the input batch. We apologize for any confusion and will make a revision in the next version.\n\n```\nQ5: Why the learnable parameter of $g_m(\\cdot,\\cdot)$ is not in Eqn. 7?\n```\n\n**Response**: Sorry for any confusion. $\\alpha_m$ is produced by $g_m(\\cdot,\\cdot)$ and we only present the former in Eqn.7 for simplifying the expression. We will clarify the relationship between $\\alpha_m$ and $g(\\cdot, \\cdot)$ in the revised version of the paper.\n\n```\nQ6: Is the proposed method sensitive to the threshold in Sec. 4.3 ?\n```\n\n**Response**: In our experiments, we found that setting a high threshold, e.g., 0.8, may lead to unstable selection results, as it potentially filters out most auxiliary datasets in the first or second round. A low threshold, like 0.2, requires more iterations for a more refined grouping. We have shown the learning curves of each iteration in Appendix D.3 with the threshold of 0.6, from which a stable selection process can be observed.\n\n```\nQ7: In addition to task similarity and structural similarity, what are the other possible factors determining whether an auxiliary task is helpful or not? I hope the authors can discuss this.\n```\n\n**Response**: Thanks for the suggestions! Besides the task and the structure, one significant factor is the strategy employed in training the model using multiple auxiliary datasets. Here we simply merge the datasets and sample them with equal possibility, which can potentially be extended. For example, selectively training different portions of the model parameters with specific datasets might lead to improved performance. We will thoroughly discuss this in the revised version of our manuscript.\n\n```\nQ8: What is the difference between search-based methods and group-based methods (line 237)?\n```\n\n**Response**: The distinction between search-based methods and grouping-based methods lies in two aspects:\n\n1. Criterion: Grouping-based methods utilize a learnable metric to evaluate the affinity between datasets, such as the task embedding similarity of Task2vec. Conversely, search-based methods assess affinity by directly testing performance or by integrating fundamental features, such as similarities in fingerprint features.\n2. Selection process: Grouping-based methods select the top-K auxiliary datasets based on the defined criterion. Search-based methods explore the auxiliary datasets through a breadth-first approach, retaining a candidate set at each level of exploration.\n\n```\nQ9: As shown in Fig. 6, FreeSolv and ESOL are not quite helpful to each other. This seems strange --- FreeSolv is about Hydration Free Energy (HFE) and ESOL is about water solubility. Theoretically, these two properties should be strongly correlated as HFE is an important factor in water solubility. Do you have any comments on this result?\n```\n\n**Response**: Although these two datasets share similar tasks, the difference in molecule structure distribution can still result in performance degradation, as demonstrated in Fig.7(a)(b). Besides, a similar task doesn't always lead to better performance, and actually, the structure distribution of ESOL dataset doesn’t exhibit a high correlation to the performance gain, as shown in Fig.3(b)."}]

Author rebuttal:
We thank the reviewers for noting that we propose a novel method (VvkN,ohh4) to address a meaningful problem (B1V2,VvkN,sjEd) with a clear motivation (VvkN,NGKd,ohh4,sjEd), and the paper is well-written and easy to follow (VvkN,NGKd,ohh4,sjEd). We further summarize our key contributions as follows:

1. We study a new angle to improve the performance of the molecule dataset with limited annotations by utilizing the auxiliary datasets with high affinity. The strategy is compatible with the other training strategy such as pertaining. 
2. We conduct a study to analyze how different molecule datasets affect each other in terms of task and structure similarity. This investigation has led to some interesting findings, such as the compensatory relationship between structure and task, and the performance gains achieved by integrating both aspects.
3. A routing-based molecule grouping method optimized by a bi-level learning framework is proposed which achieves the SOTA performance across 11 molecule datasets. 

Besides, we include the regression curves between relative improvement and the task/structure similarity measurement excluding the datasets with huge performance changes in the attachment (reviewer VvkN).

Meta-review: The paper addresses the limited data problem in molecule property prediction by proposing a novel method, MolGroup, which effectively leverages auxiliary datasets to improve performance on target datasets. The proposed method includes a routing mechanism and bi-level optimization, which is well-motivated and clearly presented. Extensive experiments are also comprehensive and carefully conducted, demonstrating the effectiveness of the proposed method. While there are some concerns raised by the reviewers, such as the interpretation of correlation and the ablation study, the overall contributions and significance of the paper outweigh these concerns. The authors are encouraged to revise the paper with respect to the comments in the final version.
Decision: Accept (poster)