SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Complexity. Given  $M$  candidate auxiliary datasets, we combine them with the target dataset and use MolGroup for grouping. At each training step, we sample data from each dataset with equal probability to form a mini-batch, ensuring that each dataset contributes equally to the training process. The total number of training instances is set as  $N \times M$  to cover all instances in the target dataset, where  $N$  is 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.
To show the effectiveness of MolGroup, we apply it to 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 experimental results can be 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 performance on the validation set (P), and the combination of the performance and the difference of fingerprint distribution  $(\mathrm{P} + \mathrm{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, we 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 a subset of datasets through the task discriminator. Furthermore, we present the results of Pretrain-Finetune 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 ChemBL [33], which can be categorized into three groups: medication, quantum mechanics, and chemical analysis. Our focus is solely on the small molecule datasets that have less than 10,000
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 a different decoder for each one. The overall training loss is calculated as the unweighted mean of the losses for all included tasks.
Method	BBBP(↑)	ClinTox(↑)	Tox21(↑)	BACE(↑)	FreeSolv(↓)	qm8(↓)
Only-target	66.620.028	56.450.023	74.230.005	75.020.026	3.8421.579	0.03850.001
Beam search(P)	66.020.015	57.860.068	74.710.004	67.340.039	3.2710.479	0.05530.002
Beam search(P+S)	67.690.034	57.630.036	74.360.003	69.740.056	3.3310.287	0.04940.000
TAG	60.660.014	57.980.028	70.320.006	70.020.076	3.9220.748	0.06370.001
Task2vec	68.180.011	47.300.031	68.000.005	74.710.032	3.3830.766	0.06350.001
MTDNN	66.560.021	52.900.039	71.870.003	69.910.026	3.4280.733	0.05230.002
UA	60.410.008	51.990.078	68.160.004	61.750.018	4.0950.334	0.06250.001
Gradnorm	61.210.007	53.080.070	59.400.041	64.830.028	4.3560.589	0.06570.006
Pretrain-Finetune	56.590.026	56.000.037	50.640.015	64.800.052	4.3910.043	0.06370.001
MolGroup	68.360.016	59.770.027	75.660.004	77.330.015	3.1160.279	0.03850.001
Due to the space limit, here we present the comparison results of 6 datasets in Table 1. Our results show that MolGroup 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. We observed that MolGroup assigned low-affinity scores to all candidate auxiliary datasets for qm8 and qm9, resulting in no 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 is constrained by the limited exploration space, and task grouping methods neglect the data structure difference, leading to suboptimal performance.
Table 1: Performance comparison of GIN on target molecule datasets, with  $\uparrow$  indicating higher is better and  $\downarrow$  indicating lower is better.
Method	BBBP(↑)	ClinTox(↑)	Tox21(↑)	BACE(↑)	FreeSolv(↓)	qm8(↓)
Only-target	66.400.019	77.590.028	75.970.009	78.910.023	2.0040.088	0.03720.002
Beam search(P)	67.530.010	76.770.117	76.610.008	81.580.038	2.1290.215	0.04730.001
Beam search(P+S)	67.150.030	77.700.076	76.830.007	80.680.024	2.3120.191	0.04580.001
TAG	69.620.008	78.080.036	76.570.006	80.100.016	2.0380.178	0.05370.001
Task2vec	66.930.022	72.920.049	75.150.007	80.840.017	1.9410.123	0.05300.001
MTDNN	66.710.039	71.970.078	76.840.008	79.790.001	2.1730.321	0.05120.001
UW	65.370.000	80.520.027	72.160.009	75.640.000	2.4050.415	0.05690.000
Gradnorm	60.400.045	43.460.083	55.060.002	48.000.005	2.5520.595	0.16940.014
MolGroup	69.660.009	81.210.040	77.340.006	82.940.017	1.8790.032	0.03720.002
Table 2: Performance comparison using Graphormer on 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 it on the specific downstream datasets with the selected auxiliary datasets. As shown in Table 2, the extensive parameter space and pretraining offer a 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 an average relative improvement of  $3.35\%$  across all the target datasets, demonstrating the generalizability of the auxiliary dataset groupings exploited by the lightweight surrogate model GIN.
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 or 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 to 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.
Ablation study. We conduct an ablation study to 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 as the target dataset, and average the gating scores  $\alpha_{m}$  every epoch, which is illustrated in Fig. 5(a). The routing mechanism tends to assign a low weight to the auxiliary-specific parameters (<0.15 in most cases) to increase the dataset-specific capacity, failing to differentiate the
affinity 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 a distinguishable affinity score distribution. Similar phenomena can be found in the other datasets.
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 running environment (see Appendix D). We pick BBBP and Tox21 as the target datasets, and the results are 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 on the number of training instances used. Although we initially input all candidate auxiliary datasets to MolGroup, we only train for a few epochs at each round, and the majority of them are filtered out after the first or 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 task≠benefit. In Section 3, we conclude that task similarity measured by task embedding is correlated with relative improvement. But Task2vec [1] fails to achieve better performance in our experiments since it doesn't incorporate structure information. It is likely that the datasets with similar tasks can serve as data augmentation and are able to 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. We combine 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.
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 to a 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 a 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, it fails to select high-affinity auxiliary datasets in our cases. Here we plot the inter-task affinity scores measured by TAG for BBBP in Fig. 7(c). Compared with MolGroup's results shown in Fig. 3(a), we observe that most datasets' affinity scores produced by TAG are homogeneous, centered around 0. Even the dataset with the highest affinity score (FreeSolv) is 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 can be easily influenced by the other dataset when they are trained together in a shared architecture. In contrast, the specific parameters assigned to each dataset in MolGroup can alleviate this issue and result in more stable training.
PCBA is an effective booster. One interesting finding is 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 a broad range of biological activities, making it a potent booster for small molecule property prediction tasks that can benefit from both structure and task.
BBBP(↑)	ClinTox(↑)	ToxCast(↑)	Tox21(↑)	ESOL(↓)	FreeSolv(↓)	Lipo(↓)
Only-target	66.620.028	56.450.023	60.690.010	74.230.005	1.5630.040	3.8421.579	0.80630.015
+PCBA	67.110.023	57.770.028	62.050.007	74.810.006	1.4630.020	3.5630.989	0.80210.009
Table 3: Cases whose performance is improved by PCBA.
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 the dataset level and does not consider subdataset-level information. As a result, it cannot fully exploit the potential benefits obtained from specific subdatasets within a large dataset. This limitation could impact the performance of MolGroup on certain datasets, such as qm8 and qm9, where MolGroup fails to group high-affinity auxiliary datasets.
[61] Biao Zhang, Ankur Bapna, Rico Sennrich, and Orhan First. Share or not? learning to schedule language-specific capacity for multilingual translation. 2021.
[62] Biao Zhang, Ivan Titov, and Rico Sennrich. On sparsifying encoder outputs in sequence-to-sequence 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: a universal 3d molecular representation learning framework. 2023.
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 on their applications or source domains: medication (HIV, SIDER, Lipo, ClinTox), 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.e., Tox21, BACE, BBBP, ClinTox, SIDER, ToxCast, ESOL, FreeSolv, Lipo, qm8 and qm9. All the involved datasets can be accessed and downloaded from OGB<sup>4</sup> or MoleculeNet repository<sup>5</sup>. The overall statistics are 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	1	128	12	1	1	2	17
split	scaffold	scaffold	scaffold	scaffold	scaffold	scaffold	scaffold
metric	ROAUC	AP	ROAUC	ROAUC	ROAUC	ROAUC	AP
Table 4: Dataset statistics(1).
SIDER	ToxCast	ChemBL	ESOL	FreeSolv	Lipo	qm8	qm9
#graphs	1,427	8,576	456,309	1,128	642	4,200	3,000	3,000
#tasks	27	617	1,310	1	1	1	12	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 a brief description of each dataset's task in this study:
- HIV: Test the ability of compounds to 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 a collection of inhibitors targeting human  $\beta$ -secretase 1.
- 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 on 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 in 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.
<sup>4</sup>https://ogb.stanford.edu/
$^{5}$ https://moleculenet.org/
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 a 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 Graphomer [57], we fix the batch size as 128 and train the model for 30 epochs. AdamW [31] with a 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 is set as 128 and 5 respectively. The size of position embedding, in-degree embedding, and out-degree embedding are fixed as 512. All the results are reported after 5 different random seeds.
Grouping method settings. As for MolGroup, we apply GIN as 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  $\lambda$  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  $\mathbf{e}^{\text{task}}$  to make sure that the dot product between the target dataset' embedding and auxiliary dataset's embedding starts with 0. The task embedding size is set as 16, and the number of the processing steps in Set2Set is set as 2. The pseudo-code is presented in Algo.1.
As for beam search, we apply GIN as 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 a 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<sup>6</sup>. First, we 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<sup>7</sup>. 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 a 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.
For the preliminary analysis shown in Fig.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, we 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 on 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
$^{6}$ https://github.com/awslabs/aws-cv-task2vec
<sup>7</sup>https://github.com/google-research/google-research/tree/master/tag
best-performing model selected based on its performance on the validation set. Cross-entropy loss is used for classification tasks and mean squared error 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 5 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 6 and Table 7. It can be observed that pretrained Graphormer outperforms GIN significantly, consistent with the previous studies. In addition, MolGroup achieves the best performance in most cases, with an 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 a promising candidate by directly comparing the performance of different groupings. It is worth noting that overall, our proposed MolGroup still achieves better performance compared to the other baseline methods.
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 to learn a shared representation that can effectively capture the characteristics of all the datasets.
Method	SIDER(↑)	ToxCast(↑)	ESOL(↓)	Lipo(↓)	qm9(↓)
Only-target	59.350.010	60.690.010	1.5630.040	0.80630.015	0.03030.001
Beam search(P)	54.250.024	63.370.004	1.4310.058	0.80730.011	0.04560.001
Beam search(P+S)	56.580.038	61.010.010	1.4760.061	0.81470.014	0.02580.000
TAG	55.640.005	58.080.003	1.4170.066	0.81700.017	0.04530.000
Task2vec	55.740.008	57.450.004	1.4360.050	0.80780.010	0.04680.000
MTDNN	55.950.015	59.020.007	1.4990.027	0.81550.017	0.04760.001
UA	56.260.009	59.930.003	1.4800.060	0.91300.017	0.04940.000
Gradnorm	55.690.005	53.310.010	1.5410.091	1.07550.001	0.05740.005
Pretrain-Finetune	51.430.009	51.060.005	1.4800.099	1.07760.066	0.04980.010
MolGroup	59.210.013	63.910.005	1.4020.037	0.79960.007	0.03030.001