Bi-level Contrastive Learning for Knowledge Enhanced Molecule Representations

Thank you for acknowledging the clarity, coherence, and comprehensive experimental scope of our paper. We address your concerns and answer your questions below. We also uploaded a revision and used blue to mark the new changes.

W1 (Novelty and Performance):

We appreciate your insight on the novelty of our work. For a comprehensive understanding, we recommend revisiting our response to reviewer C94D - Parts 1 & 2, where we elaborate on Gode's unique contributions and novelty.

To clarify, KANO and Gode are very different frameworks where KANO can only deal with Chemical Element KG (which contains element-wise relationships) while Gode can only deal with Molecule-Centric Biomedical KG (which contains sophisticated biomedical relationships for molecules).It's crucial to note that the types of KGs employed by these two frameworks are not interchangeable.

Achieving performance comparable to, or even surpassing, KANO on 10 out of 11 tasks - a benchmark method in our field - using a completely different knowledge graph should not be seen as a weakness. Instead, this accomplishment highlights the effectiveness of our framework.

Below is a table delineating the key differences between the two KGs

We are confident that the extensive data richness, superior scalability, and frequent updates inherent to the Molecule-Centric KG sub-graph from MolKG, along with its subsequent iterations, will establish it as an essential and highly effective modality for knowledge-enhanced molecular representation in the field.

W2 (Ablation Study Results in Figure 2):

We appreciate the opportunity to clarify the implications of our ablation study presented in Figure 2. It is crucial to emphasize that the variants we discuss in this figure are not minor adjustments but integral components of our framework. Each variant plays a significant role, and their removal can significantly impact performance, as we will detail:

1. KGE (K-GNN embedding initialization): This is a vital process in the pre-training of K-GNN, as further evidenced by Figure 4. Its exclusion leads to a substantial performance decrease, a finding consistent with the standard application of KGE in KANO. The impact of removing KGE is clearly demonstrated when comparing Case 1 and Case 3 in Figure 2.
2. $\kappa$-hop: The difference between a 3-hop and a 2-hop subgraph is significant, with the former having nearly 20 times more nodes and edges. This increase in complexity introduces noise, making the graph more challenging to learn from. This aspect is unique to our framework, as KANO's Chemical Element KG does not require exploration of this parameter. The performance variation between these configurations is evident when comparing Case 3 and Case 5 in Figure 2.
3. Pretraining: The self-supervised pretraining of K-GNN is crucial for aggregating local information within the molecule-centric KG subgraph into the central molecule node. The importance of this step is highlighted in Table 7 (Appendix G.1) of our revised paper.
4. $\mathbf{h}_{\mathrm{f}}$: The functional group feature, extracted using RDKit, is a commonly used attribute in property prediction, as seen in frameworks like GROVER and KANO. Omitting this feature, as shown by comparing Case 9 and Case 3, leads to a noticeable performance decline. This finding aligns with KANO's observations in their Extended Data Fig. 1-a in their paper [1].

In summary, each component in our framework plays a substantial role in ensuring optimal performance. Our ablation study methodically examines these components, thereby solidifying the effectiveness of our approach.

Thank you for your valuable feedback.

Regarding the deviations in experimental results you highlighted:

MolCLR: The results of its default setting were sourced from the KANO paper, as detailed in Appendix D.2. We independently reproduced its pre-training using the GTransformer backbone and reported it as $\text{MolCLR}_{\text{GTrans}}$

GROVER: The deviations you’ve observed are also documented in other studies, as well as in KANO's findings. We addressed this in our response to Reviewer ar3U (W2&Q2). The variance is evident in both scenarios: using the provided pre-trained models or after re-training them ourselves, which is corroborated by the discussions in their issues repository.

KANO: To our knowledge, we are the first to compare our work directly with KANO. This necessitated that we reproduce their pre-training process rather than rely on their available models. We adhered to their default setting, pre-training KANO over 50 epochs, which required approximately 21 hours on a single A6000 GPU.

We assure you that our reported results are authentic and the outcome of a meticulous and repeatable process. Acknowledging the potential for variability due to non-deterministic factors (like how GROVER explained their reproducibility issue in their repo-readme) we took the additional step of running the pre-training process three times to confirm the stability of our findings.

Should you wish, we invite you to reproduce the KANO pre-training process for further validation or to review additional papers that have attempted a similar reproduction. Considering the bias introduced by pre-training, we will also pre-train KANO as many times as possible before the rebuttal deadline, and report the standard deviation of the performance. Please see our results of extensive experiments for KANO below.

We are confident in our results and remain open to discussing our methodology and findings in more detail.

For Additional Experiments

Thanks for reminding us again about this, but there are a few to note:

In your original feedback, you mentioned you would like to see the direct comparison between KANO and Gode regarding experiments in Figures 3, 4, and 5.

We responded that Figure 3 is not feasible to conduct during this rebuttal window, as KANO has already done this. We would need to pre-train KANO for $5 (\text{KG ratios}) \times 3 (\text{independent runs}) \times 21 (\text{hours/epoch} = 315 (\text{hours})$ to replicate this experiment.

The experiment in Figure 4 is unique to Gode, as the tasks mentioned therein are defined by us. It is not applicable to KANO.

For Figure 5, we have already included the embedding visualization of KANO.

Regarding training time, KANO takes 21 hours and Gode takes 9 hours when $\kappa=2$, and 71 hours when $\kappa=3$. As we reported in Figure 2, the performance is better when $\kappa=2$. Thus, the pre-training of Gode is also more efficient than KANO.

Regarding KANO's performance we reported

Dear Reviewer,

In the past three days, we have pre-trained 7 KANO models from scratch and tested each of them on downstream tasks. Among these, KANO-1, KANO-2, and KANO-3 were trained on three A6000 GPUs, while the remaining four were trained on A100 GPUs separately.

We have uploaded all training logs for both pretraining and fine-tuning here:

<download link (training logs)>

Additionally, we provide these 7 pretrained KANO models here:

<download link (pretrained models)>

We summarize the results below. We also discovered that the fine-tuning code provided by KANO's repository is not randomly seeded but is specifically set to 51, 52, and 53 by the author. Therefore, we tested each model with a randomly seeded setup. This is crucial to ensure consistency with the experimental setup described in their paper.

All experimental settings, except for the random seeds, are exactly the same as those in KANO's repo.

Classification Tasks:

Regression Tasks:

We are still fine-tuning models on QM8, as it is time-consuming. We will also include it soon. (done)

We will update KANO's performance in Tables 2 and 3 based on those new results in our final revision.

We sincerely hope that the new results, along with the transparent training logs and models, can address your concerns and doubts about KANO's performance reported by us.

Please let us know if you have any other concerns.

Thanks for recognizing the comprehension and insightful research direction of our work. We address your concerns and answer your questions below. We also uploaded a revision and used blue to mark the new changes.

For Weaknesses:

1. We acknowledge that molecule graph pre-training strategies have been well-studied, as in the paper by Sun et al. [1]. However, the focus of our work is to investigate how biomedical knowledge graphs can enhance molecule representations, rather than proposing new strategies for molecule graph pre-training.

2. The paper by Xia et al. [2] does not discuss pre-training strategies for bio-knowledge graphs. Instead, it mentions only the Chemical Element KG, as proposed by KCL [3], which is significantly different from our constructed biomedical knowledge graph, MolKG:

   • Chemical Element KG is sourced from the Periodic Table of Elements and focuses on the attributes of chemical elements.
   • MolKG is a molecule-centric biomedical KG, integrating molecule-tailored biomedical information from PubChemRDF (the largest KG for chemistry) and PrimeKG (a precision medicine-oriented KG).

3. Ye et al.'s work [4] does not construct its own biomedical knowledge graph but uses existing datasets (Luo's dataset, Hetionet, Yamanishi_08's dataset, and BioKG). None of these datasets contains the abundant molecules found in our downstream datasets. In contrast, our MolKG construction can scale to any molecules contained by PubChem or PrimeKG, with their relevant KG triples extracted. We detailed our construction steps in Appendix B, along with the processing scripts.

4. We would like to clarify some key differences in our approach compared to Fang et al, 2022 [3].

   In both papers by KCL [3] and KANO [5], contrastive learning is conducted between (1) the molecule graph and (2) the augmented molecule graph, where (2) is obtained by linking element-wise KG triples from the Chemical Element KG directly to each element of the molecule.

   In our work, differently, we conduct contrastive learning between (1) the molecule graph, and (2) the central molecule node in (1)'s corresponding KG sub-graph. Unlike KCL or KANO, we do not augment molecule graphs explicitly but treat the molecule as a node in the KG sub-graph. We also designed a self-supervised pre-training strategy (Eq. 3) for this KG sub-graph, the effectiveness of which is demonstrated in Appendix G.1.

   Please also feel free to read Appendix F for more details regarding our comparison to similar works.

5. As explained in points 1 to 4, our work is not a simple combination of the mentioned research but introduces major innovations:

   1. A tailored molecule-centric KG providing comprehensive biomedical knowledge.
   2. A novel integration of molecular graphs with biomedical domain knowledge, treating molecules as nodes in KG sub-graphs.

   KANO [5] is an improved version of KCL [3] in both simplicity and performance stability. Thus, our comparison with KANO, as shown in Tables 2 and 3, underscores our unique contribution.

Dear Reviewer,

Thank you for your questions. We are pleased to provide prompt responses as follows:

(a) Different KGs?

Yes, we have explored this aspect:

• Figure 3: Shows the results of applying three different KGs in varying ratios (20%, 50%, 80%) to the original MolKG. KANO also did such a study in their "Extended Data Fig. 1".
• Author Response to Reviewer afSU - Part 2 - Q1: We present new results with four different KGs, each with unique relation types. We will include a more comprehensive version of this study in our final revision.

Regarding testing on KGs other than MolKG, the answer is no. Our method is specifically designed for molecule-centric KGs, such as MolKG, which is currently the only KG of such type. Similarly, the KANO framework is tailored for the Chemical Element KG they proposed. The KG types are not interchangeable between different frameworks like Gode and KANO. KANO also only studies the Chemical Element KG.

(b) Different CL Method?

No, we did not perform an ablation study on different contrastive learning methods. This decision aligns with the standard practices in the field, as evidenced by similar methodologies in notable works like MolCLR and KANO. These studies, too, did not explore alternative contrastive learning methods. Nevertheless, we would love to conduct this study and hopefully integrate it into our final revision.

We hope these clarifications enhance your understanding of our proposed method.

Please remind me if ablation experiments (or sth alike) were conducted to reveal the combinatorial effects of (a) using a different KG, and (b) using a different CL method.

Thank you for appreciating our paper's clear motivation, innovative use of knowledge graphs in molecular representation, and the accessible structure of our method. We address your concerns and answer your questions below. We also uploaded a revision and used blue to mark the new changes.

W1: Experiments should ideally be conducted under the same settings.

Thank you for highlighting the importance of uniform experimental settings across different models. In our study, we paid careful attention to this aspect to ensure a fair and meaningful comparison of our GODE method with the baseline models.

To address your concern regarding the standardization of experimental settings, we acknowledge that our original description in Appendix D may have caused some confusion. To clarify, this meant that while we utilized the original model architectures as provided by the authors, we also meticulously standardized all other experimental settings for a fair comparison. Our goal was to ensure that each method was evaluated on a consistent and comparable basis.

Taking your feedback into consideration, we have further refined our experimental design for an even more equitable comparison. Specifically, we have now reproduced $\text{MolCLR}\_{\text{GTrans}}$ using the GTransformer backbone and $\text{KANO}\_\text{CMPNN}$ with its default CMPNN backbone as specified in the KANO paper. It's crucial to note that the results for GROVER and KANO presented in our initial submission were both based on the GTransformer backbone.

In our revised submission, we have added a new section, Appendix D.2 Baseline Models, which provides detailed insights into the experimental settings used for each baseline model. This addition aims to offer a transparent and comprehensive view on how we ensured uniformity in our comparative analysis.

The results reflecting these changes also have been updated in Table 2 and Table 3 of our revised paper (in blue for easy identification). We believe these enhancements and clarifications in our methodology significantly strengthen the validity of our findings and address your concerns regarding the experimental setup's consistency.

W2&Q2: The experimental results for GROVER appear to exhibit significant deviations from what is reported in the original paper.

We appreciate your concern regarding the observed deviations in GROVER's performance in our study compared to the results reported in its original paper [1]. We would like to emphasize that such discrepancies are not unique to our research. Similar deviations have been reported in other significant studies, including MGSSL [2], GEM [3], Mole-bert [4], and KANO [5], all of which were published in highly respected venues.

Our experimental settings strictly adhered to the guidelines and configurations outlined in the original GROVER paper.

W3: Recommendation for Additional Ablation Studies

Thank you for your insightful recommendation to further investigate the necessity of the various steps in our methodology. In response, we have conducted detailed ablation studies, the results of which are presented in Appendix G.1 - Table 7 of our revised manuscript.

These studies decisively highlight the integral roles of both M-GNN and K-GNN pre-training within our framework. Each pre-training phase is shown to uniquely and significantly contribute to the overall enhancement of our model's performance. This evidence robustly counters the notion of redundancy in our methodological approach.

We hope that these findings will satisfactorily address your concerns about the steps previously perceived as redundant in our methodology.

References

[1] Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W. and Huang, J.. Self-supervised graph transformer on large-scale molecular data. In NeurIPS 2020.

[2] Zhang, Z., Liu, Q., Wang, H., Lu, C. and Lee, C.K. Motif-based graph self-supervised learning for molecular property prediction. In NeurIPS 2021.

[3] Fang, X., Liu, L., Lei, J., He, D., Zhang, S., Zhou, J., Wang, F., Wu, H. and Wang, H., 2022. Geometry-enhanced molecular representation learning for property prediction. Nature Machine Intelligence, 4(2), pp.127-134.

[4] Xia, J., Zhao, C., Hu, B., Gao, Z., Tan, C., Liu, Y., Li, S. and Li, S.Z. Mole-bert: Rethinking pre-training graph neural networks for molecules. In ICLR 2023.

[5] Fang, Y., Zhang, Q., Zhang, N., Chen, Z., Zhuang, X., Shao, X., ... & Chen, H. (2023). Knowledge graph-enhanced molecular contrastive learning with functional prompt. Nature Machine Intelligence, 1-12.

Dear Reviewer,

As the rebuttal deadline for our submission is rapidly approaching, we wanted to reach out to ensure that all your concerns have been addressed satisfactorily. Should there be any further questions or points of clarification needed, please let us know. We are ready and willing to provide additional information or engage in further discussion to address any issues.

We are happy to know that several of your concerns have been addressed by us. We answer your new questions and address your new concerns below.

(1) "GROVER has provided fine-tuned models"

As previous works (e.g., GEM, KANO) fine-tuned GROVER by themselves and reported the performance, we keep consistent with such a setting. We believe It is not rigorous to use the provided fine-tuned models as there is a possibility that they have been trained with the test data. Furthermore, GROVER provides an efficient framework to fine-tune the pre-trained model on downstream tasks. It only takes a few minutes to fine-tune a model on a task.

Nevertheless, we will provide you with the results of those provided fine-tuned models soon, in our later response to you. Please see the results below.

(2) "Why select GTransformer, MPNN, and GIN as the backbones to test?"

We believe GTransformer is a sophisticated backbone to engage both edge and node information mutually, whereas MPNN and GIN are simpler but show good performance in Tables 2 and 3. Note that most previous works like KANO did not study the backbone, and took the best-performing GNN (e.g., CMPNN for KANO) as their backbone.

(3) "High runtime overhead due to different encoders"

We mentioned "During fine-tuning, embeddings from K-GNN remain fixed, updating only the parameters of M-GNN" in Section 4.1 - Implementation, which means the embedding from K-GNN is pre-computed and simply concatenated to the M-GNN embedding as an additional feature (together with the extracted functional group feature) during fine-tuning or prediction. This means the fine-tuning/prediction is almost as efficient as the fine-tuning/prediction with M-GNN (e.g., GTransformer) itself.

Here is an example for your information: fine-tuning one epoch on the BACE task takes 10.2 seconds for Gode while taking 9.6 seconds for GROVER (GTransformer).

To thoroughly validate the results obtained from GROVER, we undertook an additional fine-tuning of the pre-trained GROVER model (specifically the large version available in their repository [1]) using a 10-fold cross-validation approach, as opposed to the 3-fold setting originally described in the paper. The detailed results of this rigorous re-evaluation are presented below.

For complete transparency and to facilitate further analysis, we have uploaded both the short and long versions of the training logs. These can be accessed via the following link:

<10-fold GROVER fine-tuning logs>

Classification Tasks:

Regression Tasks:

The results from the 10-fold validation align closely with the 3-fold results we reported in the paper, thereby reinforcing the validity of the GROVER performance that we initially presented.

References

[1] https://github.com/tencent-ailab/grover

Q1: When the knowledge graph information is useful or even harmful for a particular molecule property prediction task?

This is an insightful question. In our work using GINE for K-GNN, both node (entity) and edge (relation) information in the knowledge graph can influence molecular representation. We recognize the importance of identifying which information in the KG is beneficial or detrimental and plan to conduct a systematic study on this in the future. To this end, we conducted a preliminary experiment to understand the impact of different relations on molecular property prediction tasks:

For the BBBP task, removing "xlogp3" and "xlogp3-aa" led to a significant decrease in performance, suggesting their vital role in the model. In contrast, omitting "neighbor_2d", "neighbor_3d", and "has_same_connectivity" resulted in a slight improvement, indicating that these relations may not be crucial or could be harmful.

In the BACE task, the impact of removing "tautomer_count" and "covalent_unit_count" was the most pronounced (a decrease of 0.9%), highlighting their importance in this specific prediction task. The minimal impact of removing "indication" (a decrease of 0.2%) suggests that this relation might not be as critical for BACE as for BBBP.

These results provide initial insights into the task-specific importance of different relations in the knowledge graph. We plan to extend this preliminary analysis into a more comprehensive and systematic study in the future, which will allow us to more accurately determine the role and impact of various knowledge graph relations (or entities) on different molecule property prediction tasks.

Q2: "Visual grouping of the methods ..."

Thank you for your suggestion. We have split the non-KG-based and KG-based methods in Table 2 and Table 3 in our revised paper.

Q3: What affects the steepness of the improvement as knowledge graph size is increased?

The observations in Figure 3 are derived from a single run, intended to illustrate the impact of KG size on model performance across different tasks. Due to the high computational cost involved — requiring re-training of the K-GNN for each ratio per run and subsequent contrastive learning — the experiment was not scaled to multiple runs. Therefore, the steepness observed in the current Figure 3 may contain an element of randomness, and drawing definitive conclusions from it is challenging. However, we certainly plan to prioritize this aspect as a key area of future research, potentially integrating it with the systematic study outlined in response to Q1.

Q4: Why wasn’t the Ye et al, 2021 model benchmarked?

Thanks for this important note! We have benchmarked it (named KGE_NFM) in Table 2 and Table 3 in our revised paper.

Thank you for your insightful and constructive feedback on our submission. We greatly appreciate your recognition of the strengths of our work. We address your concerns and answer your questions below. We also uploaded a revision and used blue to mark the new changes.

W1: Number of seeds in the experiment

Thanks for pointing out that the number of seeds (3) is low in experiments. While we recognize the value in your suggestion that using a greater number of seeds could yield more robust and convincing results, this experimental setup was chosen to maintain consistency with previous works [1, 2, 3], as suggested by MoleculeNet's paper [4]. We value your suggestion and will explore the use of additional seeds in our future research.

W2: Missing some of the larger benchmark datasets. (e.g., HIV, PCBA).

We appreciate your observation regarding the omission of larger datasets like HIV and PCBA in our study. Our primary focus was to explore how biochemical knowledge graphs can enhance molecule representation for property prediction. To maintain a consistent and comparable framework, we chose to validate our method on datasets used by our backbone - GROVER [1], a popular model that leverages molecule information only.

We acknowledge the potential value of extending our approach to larger datasets, as this could further validate and strengthen our findings. However, constructing a more extensive MolKG that includes molecules from these datasets, along with conducting K-GNN pre-training and contrastive learning with the new KG, requires substantial time beyond this author's response window. Thus, we decided to include those results in our final revision.

References

[1] Rong, Y., Bian, Y., Xu, T., Xie, W., Wei, Y., Huang, W. and Huang, J.. Self-supervised graph transformer on large-scale molecular data. In NeurIPS 2020.

[2] Zhang, Z., Liu, Q., Wang, H., Lu, C. and Lee, C.K. Motif-based graph self-supervised learning for molecular property prediction. In NeurIPS 2021.

[3] Fang, Y., Zhang, Q., Zhang, N., Chen, Z., Zhuang, X., Shao, X., ... & Chen, H. (2023). Knowledge graph-enhanced molecular contrastive learning with functional prompt. Nature Machine Intelligence, 1-12.

[4] Wu, Z., Ramsundar, B., Feinberg, E.N., Gomes, J., Geniesse, C., Pappu, A.S., Leswing, K. and Pande, V., 2018. MoleculeNet: a benchmark for molecular machine learning. Chemical science, 9(2), pp.513-530.