SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tinglin Huang
Ziniu Hu2
$^{1}$ Yale University,  $^{2}$ University of California, Los Angeles
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<sup>1</sup>, 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.
Predicting and understanding molecular properties is a fundamental task in biomedical and chemical fields [51, 14, 23, 55], such as evaluating the toxicity of new clinical drugs, characterizing the binding results for the inhibitors of human  $\beta$ -secretase, and predicting the thermodynamic 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 it difficult to collect a large and diverse labeled dataset for training. This limited availability of labeled data poses a challenge for traditional machine learning methods, which struggle to effectively generalize on small
Figure 1: Relative improvement of using the combination of target dataset and auxiliary dataset over only using target dataset:  $(\mathrm{perf}(a,b) - \mathrm{perf}(a)) / \mathrm{perf}(a)$ , where  $a$  is target dataset and  $b$  is auxiliary dataset.
Source code is available at https://github.com/Graph-and-Geometric-Learning/MolGroup.
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 a 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 as 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 a shared representation [60, 45, 20]. 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<sup>2</sup>. The results demonstrate how auxiliary datasets positively or negatively affect the target dataset. For instance, ToxCast greatly improves the performance of ClinTox, 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<sup>3</sup>. In light of this, how can we measure the affinity of auxiliary datasets to a target dataset?
Measured by both structure and task similarity — a clear answer is well supported by our 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.
Present work. In this paper, we focus on designing a dataset grouping method for molecules. Here we propose MolGroup, a routing-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 to route between auxiliary-specific or target-specific parameters through a bi-level optimization framework. The routing mechanism is optimized toward maximizing the target dataset's performance, and its learned gating score serves as 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:
- 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, a 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.
2More experimental details can be found in Appendix C.
3In 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.
Molecule representation pretraining. Recently, the pretrain-finetune paradigm has been extensively applied to learn the representation of molecules [57, 19, 16, 37, 52, 55]. These methods involve pretraining a graph neural network on a large-scale molecule dataset through contrastive [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 auxiliary 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] to 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 as an integer program problem and solve it through optimization. Recent works [12, 44] learn the underlying affinity between different tasks by quantifying the effect of the task'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 incorporating both structure and task information, our method can achieve better performance.
While MolGroup and MTG-Net [44] are both related to meta-learning approaches for solving grouping problems, the specific objectives differ. MTG-Net formulates the grouping problem as a meta-learning problem and learns the relationship between tasks with a few meta-training 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 to determine the fusion significance of two datasets' parameters, which is used as the affinity score.
3 Understanding Relationship between Molecule Datasets
Molecule datasets comprise two critical aspects of information: the structural characteristics of the molecules and the associated predictive tasks. Using this insight, we analyze the relationship between molecule datasets by dividing them into these two dimensions. We conduct an empirical study that links the discrepancies in structure and task between datasets to changes in performance when they train together. Molecule structure and predictive task are quantified as the distribution of fingerprint [5, 7] and task embedding [1]:
- Fingerprint distribution: Fingerprint is a multi-hop vector used to character a 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.
- **Task embedding:** We apply Task2vec [1] to quantify task representation which uses Fisher Information Matrix associated with a 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 KL divergence as 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
Figure 3: (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, as 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 Fig.3(a). This combined similarity demonstrates a stronger correlation with the performance compared to 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), we can observe that the structure and task correlation exhibit a compensatory relationship. For instance, a dataset may present a low structural correlation but a high correlation in task similarity, as 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 a 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 a 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 a positive correlation, indicating that the auxiliary datasets with similar structures and predicted tasks can effectively augment the target dataset and improve its performance.
There are also negatively related cases, such as 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 on out-of-distribution generalization [58, 54] which demonstrate 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.
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-based 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, a molecule dataset grouping method. MolGroup includes two modules: (1) routing mechanism used to quantify the affinity between two datasets, and (2)
Figure 4: Overview of routing mechanism and bi-level optimization framework.
bi-level optimization framework used to 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.1 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  $\mathcal{D}_T$  and a set of auxiliary datasets  $\{\mathcal{D}_{Aux}\}_M$ . To this end, we propose to use the routing mechanism that dynamically allocates the impact of each auxiliary dataset on the target dataset across the network's sub-layers. Intuitively, this routing mechanism allows us to 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 a GNN at  $l$ -th layer as:
$$
\mathbf {z} ^ {(l + 1)} = f _ {\theta} ^ {(l)} \left(\mathbf {z} ^ {(l)}\right), \tag {1}
$$
where  $\mathbf{z}^{(l + 1)}$  is the representation of the input batch at  $l$ -th layer, and  $f_{\theta}^{(l)}(\cdot)$  denotes the convolution function with parameter  $\theta$  at  $l$ -th layer. As illustrated in Fig. 4, in MolGroup, every dataset is assigned a specific parameter  $\theta_T, \theta_1, \dots, \theta_M$ , 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  $\mathcal{B}_T$  and  $m$ -th auxiliary dataset  $\mathcal{B}_m$ , the convolution with routing mechanism for  $m$ -th auxiliary dataset at  $l$ -th layer is calculated as:
$$
\mathbf {z} _ {m} ^ {(l + 1)} = \alpha_ {m} f _ {\theta_ {T}} ^ {(l)} \left(\mathbf {z} _ {m} ^ {(l)}\right) + (1 - \alpha_ {m}) f _ {\theta_ {m}} ^ {(l)} \left(\mathbf {z} _ {m} ^ {(l)}\right), \tag {2}
$$
$$
\text {w i t h} \quad \alpha_ {m} = g _ {m} \left(\mathcal {B} _ {T}, \mathcal {B} _ {m}\right), \tag {3}
$$
where  $\mathbf{z}_m^{(l)}$  is the representation of  $m$ -th auxiliary dataset at  $l$ -th layer, and  $\alpha_{m}$  is the gating score generated by the routing function  $g_{m}(\cdot)$ . 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 a minor impact on  $\theta_T$ , particularly when  $\alpha_{m} = 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  $\theta_T$ , particularly when  $\alpha_{m} = 1$ , indicating a hard-parameter sharing architecture.
The calculation of  $g(\cdot)$  is parameterized using two modules to capture task affinity and structure affinity. For task affinity, we assign learnable embeddings  $\mathbf{e}_T^{\mathrm{task}}$  and  $\mathbf{e}_m^{\mathrm{task}}$  for the target dataset and auxiliary dataset respectively, which updated through optimization. For structure affinity, we first extract the fingerprint feature  $x^{\mathrm{fp}}$  for each molecule instance in the input batch. Then, we embed it using a weight matrix and apply Set2Set [48] to obtain a single embedding  $\mathbf{e}^{\mathrm{struct}}$ :
$$
\mathbf {e} ^ {\text {s t r u c t .}} = \operatorname {S e t 2 S e t} \left(\left\{\mathbf {W} x ^ {\mathrm {f p}} \right\} _ {n}\right), \tag {4}
$$
where  $\mathbf{W}$  is a learnable weight matrix, and  $n$  is the number of the instance in the batch. Task affinity score and structure affinity score are computed as the cosine similarity and fused as the dataset affinity
Algorithm 1: Iterative filtering process with MolGroup
Input: Target dataset  $\mathcal{D}_T$  with  $N$  training instances; Candidate auxiliary datasets  $\{\mathcal{D}_A\} _M$  GNN model with specific parameters  $\theta_T,\{\theta_m\}_M$  for each dataset and routing mechanism  $g(\cdot)$  Number of iterations  $R$  ;Number of epochs  $E$  ; Learning rate  $lr$  ; Batch size  $B$    
// Filtering round   
for  $r\gets 1,\dots ,R$  do Random initialize  $\theta_T,\theta_1,\dots ,\theta_M$  and  $g(\cdot)$ $\phi_1,\dots ,\phi_M\gets 0.$  // Affinity scores  $I\gets NM / / B$  // Number of iteration in each epoch   
// Training epoch   
for  $e\gets 1,\dots ,E$  do // Training step for iter  $\leftarrow 1,\dots ,I$  do Sample mini-batch  $\{\mathcal{B}_T,\mathcal{B}_1,\dots ,\mathcal{B}_M\}$  from current datasets. Obtain losses  $l_{T},\{l_{m}\}_{M}$  and affinity scores  $\{\alpha_{m}\}_{M}$  by feeding mini-batch to GNN. // Bi-level optimization framework 1)  $\theta_T(\{\alpha_m\}_M)\gets \theta_T - lr\cdot \nabla_{\theta_T}\sum_m^M l_m$  2) Update  $g(\cdot)$  through target dataset's loss function with  $\mathcal{B}_T$  and  $\theta_T(\{\alpha_m\})_M$  Update all the parameters  $\theta_T,\theta_1,\dots ,\theta_M$  through  $l_{T},l_{1},\dots ,l_{M}$  if  $e = E$  then // Average affinity scores in final epoch for  $m\gets 1,\dots ,M$  do end end end // Remove datasets according to threshold  $\{\mathcal{D}_A\}_{M'}\gets \{\mathcal{D}_{A_m}|\phi_m\geq 0.6\} .$ $M\gets M^{\prime}$
score, which can be formulated as:
$$
\alpha_ {m} = g _ {m} \left(\mathcal {B} _ {T}, \mathcal {B} _ {m}\right) = \sigma \left(\lambda \mathbf {e} _ {T} ^ {\text {t a s k}} \cdot \mathbf {e} _ {m} ^ {\text {t a s k}} + (1 - \lambda) \mathbf {e} _ {T} ^ {\text {s t r u c t .}} \cdot \mathbf {e} _ {m} ^ {\text {s t r u c t .}}\right), \tag {5}
$$
where  $\sigma(\cdot)$  is the logistic-sigmoid function, and  $\lambda$  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  $\lambda$  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 in a target dataset-aware way is challenging since it is only used during the forward pass of the auxiliary dataset. To address this, we propose to use a bi-level [34] optimization framework to incorporate the guidance of the target dataset. The framework utilizes the target dataset's performance, using the parameters updated by the auxiliary dataset, as a signal to guide the learning of the routing mechanism.
Equipped with the routing mechanism, the overall optimization can be formulated as:
$$
\min  _ {\theta_ {T}} L _ {T} \left(\mathcal {B} _ {T}; \theta_ {T}\right), \tag {6}
$$
$$
\min  _ {\theta_ {T}, \theta_ {m}} L _ {m} \left(\mathcal {B} _ {m}; \theta_ {T}, \theta_ {m}, \alpha_ {m}\right), \tag {7}
$$
where  $L_{T}$  and  $L_{m}$  denote the loss function of the target dataset and auxiliary dataset. It can be found that the optimization of  $\theta_{T}$  is closely linked to the auxiliary task and its relevance is determined by the routing function. We explicitly represent such dependency as  $\theta_{T}(\alpha_{m})$ . With the goal of obtaining an optimal gating score  $\alpha_{m}$  parameterized by  $g_{m}(\cdot)$ , given  $M$  auxiliary datasets, we further formulate the optimization as a bi-level framework:
$$
\min  _ {\alpha_ {1}, \dots , \alpha_ {M}} L _ {T} \left(\mathcal {B} _ {T}; \theta_ {T} \left(\left\{\alpha_ {m} \right\} _ {M}\right)\right), \tag {8}
$$
$$
\text {w h e r e} \quad \theta_ {T} \left(\left\{\alpha_ {m} \right\} _ {M}\right) = \underset {\theta_ {T}} {\arg \min } \sum_ {m} ^ {M} L _ {m} \left(\mathcal {B} _ {m}; \theta_ {T}, \theta_ {m}, \alpha_ {m}\right), \tag {9}
$$
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 an affinity score below the predefined threshold are then filtered out. This process is repeated for  $t$  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.1.