Two-Stage Pretraining for Molecular Property Prediction in the Wild

Authors: Kevin Tirta Wijaya, Minghao Guo, Michael Sun, Hans-peter Seidel, Wojciech Matusik, Vahid Babaei
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: large-scale pretraining, molecular property prediction

Abstract:
Accurate property prediction is crucial for accelerating the discovery of new compounds. Although deep learning models have achieved remarkable success, their performance often relies on large amounts of labeled data that are expensive and time-consuming to obtain. Thus, there is a growing need for models that can perform well with limited experimentally-validated data. In this work, we introduce MoleVers, a versatile pretrained model designed for various types of molecular property prediction *in the wild*, i.e., where experimentally-validated molecular property labels are scarce. MoleVers adopts a two-stage pretraining strategy. In the first stage, the model learns molecular representations from large unlabeled datasets via masked atom prediction and *dynamic denoising*, a novel task enabled by a new branching encoder architecture. In the second stage, MoleVers is further pretrained using auxiliary labels obtained with inexpensive computational methods, enabling supervised learning without the need for costly experimental data. This two-stage framework allows MoleVers to learn representations that generalize effectively across various downstream datasets. We evaluate MoleVers on a new benchmark comprising 22 molecular datasets with diverse types of properties, the majority of which contain 50 or fewer training labels reflecting real-world conditions. MoleVers achieves state-of-the-art results on 20 out of the 22 datasets, and ranks second among the remaining two, highlighting its ability to bridge the gap between data-hungry models and real-world conditions where practically-useful labels are scarce.

Decision: reject
Reviewer scores: 5, 5, 5, 3

Reviews and discussion:
<thread>
creation date:
1730363106939

thread_type:
official_review

<thread_object>
creator: reviewer_stin

<content>
summary:
This paper introduces MoleVers, a two-stage pre-training framework designed for molecular property prediction. The first stage of pre-training includes a masking atom prediction strategy alongside a dynamic denoising approach. The second stage involves supervised learning tasks, leveraging labels generated by DFT methods.

strengths:
1. The paper is well-written and easy to follow.
2. The dynamic scaling of the denoising process is an innovative and interesting approach.

weaknesses:
1. Overall, the proposed pre-training strategy lacks novelty. Both the masking and denoising strategies have been widely adopted in previous pre-training methods, such as Uni-Mol[1] and 3D-denoising. Additionally, the second stage of pre-training, which involves supervised task learning with DFT-generated properties, is also used by Transformer-M[2]. The motivation for splitting the pre-training into two distinct stages needs further clarification.
[1] Uni-Mol: A Universal 3D Molecular Representation Learning Framework
[2] One Transformer Can Understand Both 2D & 3D Molecular Data

2. The DFT process is computationally expensive, especially for large molecules. Relying on DFT-generated labels for pre-training tasks in MRL may not be a feasible or efficient choice.
3. The rationale for decoupling the denoising and masking atom tasks is unclear, as these two distributions operate on different modalities—coordinates and atom types, respectively.
4. Results for each subtask on QM9 and MoleculeNet should be reported to demonstrate the model's superiority effectively.

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730685003900

thread_type:
official_review

<thread_object>
creator: reviewer_blw3

<content>
summary:
In real-world scenarios where molecular property labels are scarce, there is a need to develop property prediction methods that can perform well with limited labeled data. To address this, the authors propose a two-stage pretraining approach, MoleVers. The first stage employs atom masking and dynamic denoising, where a branching encoder is used to mitigate interference between denoising pretraining with larger noise scales and the masking task. To evaluate MoleVers' fine-tuning performance on limited labeled data, the authors introduce the MPPW benchmark, comprising 22 biological and chemical property subsets, each containing between 38 and 123 molecules. MoleVers outperforms four baseline molecular pretraining methods and also shows superior performance on QM property prediction datasets with abundant data. The ablation studies demonstrate the effectiveness of both pretraining stages, as well as the impact of noise scale and pretraining data on performance.

strengths:
1. The writing is clear and easy to understand.
2. Experimental settings are well-detailed, with fair comparisons based on a shared pretraining dataset, making the results convincing.
3. Addressing the challenge of limited labeled data for biological and chemical property prediction is crucial, and the two-stage approach provides an effective solution that could inspire further research in this area.

weaknesses:
1. Given that each task in the proposed MPPW benchmark includes only 38–123 molecules, there is a risk that the results are not robust and variable. However, I did not see any mention of repeated experiments or the reporting of standard deviations. I suggest the authors perform multiple runs and report the mean and standard deviation to make the results in Table 1 more robust and convincing.
2. The authors propose a dynamic denoising pretraining strategy and a branching network design, but no ablations are provided to test the effectiveness of these design choices. Could the authors conduct experiments to compare dynamic and static denoising pretraining, as well as examine the impact of branching on the results?
3. The model descriptions could be more thorough to aid reader comprehension. For example, what is the network structure of the Denoising and MAP heads? What is the "pooling with multihead attention" (PMA) module, and what does the pair representation entail? A detailed description of the encoder architecture would also be helpful, perhaps included in the main text or the appendix.

questions:
1. In Section 3.1, you mention that "the denoising tasks are equivalent to learning a molecular force field" and that this explanation is based on the assumption that the input data is at equilibrium. Do the RDKit-generated conformations used here meet this "equilibrium" assumption?
2. In Section 3.1.2, you mention that increasing the noise scale can cover more unseen molecules. This issue was also addressed in [1], where they argued that the noise scale of an isotropic Gaussian correlates with the accuracy of force fields, showing that excessively small or large noise scales negatively impact force field accuracy (extended data Fig. 3 [1]). In contrast, this paper suggests that larger noise scales are beneficial. How do you think of this discrepancy?
[1] Pre-training with fractional denoising to enhance molecular property prediction, Nature Machine Intelligence, 2024.
3. Additionally, in the dynamic denoising pretraining, which is compared to the denoising diffusion model, does the dynamic noise scale $$\sigma$$ serve as an input to the network as in diffusion models?
4. In Section 3.2, you selected HOMO, LUMO, and dipole moment prediction as the second-stage pretraining tasks. Why were these tasks chosen? Are they related to downstream properties in MPPW? Regarding your assumption that "representations useful for predicting one type of property could also help in predicting others," does this require that the second-stage properties are correlated with downstream properties? After all, property prediction tasks are different from general pre-training tasks and are already learning high-level features.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730707106011

thread_type:
official_review

<thread_object>
creator: reviewer_kqgy

<content>
summary:
The paper introduces MoleVers, a two-stage pretrained model aimed at improving molecular property prediction in low-data scenarios. The model first undergoes self-supervised training through masked atom prediction (MAP) and dynamic denoising. In the second stage, MoleVers utilizes auxiliary labels generated by cost-effective computational methods to further refine molecular representations. The model is evaluated on a newly curated benchmark, which includes 22 datasets with 50 or fewer labeled samples each. Experimental results indicate that MoleVers achieves competitive performance on them, suggesting its robustness in both low-data settings.

strengths:
- **Valuable Research Question**: The research question addressed in this paper focuses on an area in the field that has lacked sufficient attention and remains underexplored. The authors provide a feasible solution for effective property prediction in low-data scenarios.

- **Benchmark Contribution**: The MPPW benchmark is a valuable addition to the field, as it reflects real-world scenarios where labeled molecular data is often scarce.

- **Strong Empirical Results**: MoleVers achieves consistently high performance across various assays, demonstrating its generalizability and effectiveness in low-data regimes.

weaknesses:
**Weaknesses:**

- My primary concern with this paper is the lack of technical novelty. Although the authors propose a two-stage pretraining approach, each pretraining stage employs widely used methods, such as MAP and dynamic denoising. Thus, in terms of technical contribution, this work may lean more toward engineering improvements.

- Some experimental setups and procedures lack sufficient motivation and explanation:

   1. In Section 3.2, the authors mention selecting HOMO, LUMO, and Dipole Moment as auxiliary tasks. However, as quantum properties, these may be less beneficial for ADMET tasks, raising doubts about the validity of the assumption in line 226.
   
   2. In line 298, the authors use a subset of the stage 1 pretraining dataset for auxiliary task labeling and further pretraining in stage 2. This implementation detail lacks motivation—why was a subset of the stage 1 pretraining dataset chosen for stage 2?
   
   3. In Section 5.4, the authors point out that the benefit from stage 2 pretraining might depend on the nature of the downstream task, which heightens concerns over the choice of task properties in stage 2.

   4. In line 425 of Section 5.6, the authors explore the impact of pretraining dataset diversity by filtering molecules based on certain elements. However, do these filtered datasets differ in size? This could also contribute to performance differences, so it would be preferable if the authors controlled for dataset size in this analysis.

- The authors’ evaluation in high-data settings is insufficiently comprehensive. They selected only three datasets (QM7, QM8, QM9), which are closely related to the auxiliary tasks. I recommend a more diverse downstream task evaluation on MoleculeNet or TDC datasets.

questions:
Please refer to the questions in Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712640405

thread_type:
official_review

<thread_object>
creator: reviewer_yqcy

<content>
summary:
This paper designs a two-stage pre-training framework for molecular property prediction in the wild. The authors claim that previous methods are not good at generalizing to the testing cases with few labeled samples provided. The first stage of the proposed method additionally equips denoising objectives as branching encoder learning compared to the previous method. The second stage is training the model via weakly supervised learning with coarse labels (also through computation tools). Experimental results on a new benchmark dataset demonstrate the proposed method's effectiveness.

strengths:
(1) The presentation of this paper is quite clear. The major idea is quite straightforward. It can be expected that better denoising frameworks and weakly supervised learning can be effective for enhancing few-shot molecular property prediction performance. 

(2) Comprehensive experiments have been covered in this work. And new large-scale benchmark dataset has been introduced for evaluation in this work.

weaknesses:
(1) The algorithmic contribution appears somewhat limited. Firstly, both the denoising pretraining strategy and masking strategy have been explored in previous works. Therefore, this study primarily concentrates on the interaction between these two modules during the pretraining stage, which lacks novelty. In the second stage, the effectiveness of weakly supervised learning is not surprising, provided the computed coarse labels are essentially accurate. Overall, the efficacy of the proposed strategies has been largely validated by previous works, which somewhat diminishes the significance of this study.

(2) It appears that the research works related to few-shot molecular property prediction have been overlooked in the discussion of related works. Interestingly, the experiments conducted in this study primarily focused on the few-shot molecular property prediction task, yet this line of research has not been addressed in the related works section.

(3) Given that the benchmark dataset introduced in this work is novel, it is not entirely convincing that the proposed method consistently outperforms previous approaches. The number of baseline methods covered is relatively small compared to the total number of molecular pretraining methods. It seems that some molecular denoising pretraining frameworks, which are crucial for validating the effectiveness and novelty of the proposed method, have not been included in the comparison.

questions:
(1) Could the authors elaborate more on the novelty of the first stage of pre-training? What are the key novel points that distinguish the proposed method from previous approaches?

(2) Could authors have more discussions over related works about the few-shot molecular property prediction? There is also a series of works working on this problem. Considering this work is actually tackling the same problem, it is appropriate to have a few more discussions and even experimental comparisons over the previous approaches.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731993975518

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their insightful comments. We are pleased that all reviewers recognize the value of our research question of property prediction in the small-data regime and appreciate the comprehensiveness of our experimental evaluation (yqCY, kQgY). We are also encouraged by the positive reception of the proposed two-stage approach and dynamic denoising, which were noted as innovative and inspiring (StiN, BLw3). Here, we address the recurring comments regarding our contributions.

The contributions of this work can be summarized as follows:
- **The branching encoder and dynamic denoising**. While masked atom prediction and denoising tasks have been employed in various works, denoising with larger noise scales (e.g., larger than **1 Å**) is previously thought to be detrimental to the downstream prediction performance [1-4]. With our novel design, we demonstrated the contrary that using larger noise levels (up to **10 Å**) in the first pretraining stage can actually improve the downstream performance.

- **The two-stage pretraining approach**. We demonstrate that performing a supervised pretraining following the self-supervised pretraining helps with prediction performance, even if the pretraining and downstream properties are not strongly-correlated. Moreover, splitting the pretraining into two stages allows our strategy to be used in a plug and play manner. Thus, we can use any unlabeled datasets (stage1) and labeled datasets (stage 2) to train any models (Appendix A2). 

- **A new benchmark for molecular property prediction that reflects real-world conditions**. Existing benchmarks, such as [6], usually provide 600 to 130K labeled molecules. This number of labeled samples is unrealistic; only around 0.37% of assays in the ChemBL database have 100 or more labels. The proposed MPPW benchmark addresses this gap by using experimentally-validated datasets that mostly contain 50 or less training samples.


Reference

[1] Zhou, Gengmo, et al. "Uni-mol: A universal 3d molecular representation learning framework." (2023).

[2] Yang, Junwei, et al. "MOL-AE: Auto-Encoder Based Molecular Representation Learning With 3D Cloze Test Objective." (2024).

[3] Zaidi, Sheheryar, et al. "Pre-training via denoising for molecular property prediction." (2022).

[4] Ni, Yuyan, et al. "Pre-training with fractional denoising to enhance molecular property prediction." (2024).

[5] Luo, Shengjie, et al. "One transformer can understand both 2d & 3d molecular data." (2023).

[6] Wu, Zhenqin, et al. "MoleculeNet: a benchmark for molecular machine learning." (2018).

[7] Ju, Wei, et al. "Few-shot molecular property prediction via hierarchically structured learning on relation graphs." (2023)

[8] Guo, Zhichun, et al. "Few-shot graph learning for molecular property prediction." (2021)

[9] Wang, Yaqing, et al. "Property-aware relation networks for few-shot molecular property prediction." (2021)

[10] Lee, Juho, et al. "Set transformer: A framework for attention-based permutation-invariant neural networks." International conference on machine learning. (2019)
</content>
</thread_object>

</thread>


Metareview:
metareview:
(a) Summary: The paper proposes MoleVers, a two-stage pretrained model for molecular property prediction. In the first stage, it uses masked atom prediction and dynamic denoising with a branching encoder on large unlabeled datasets. The second stage employs supervised pretraining with auxiliary labels from inexpensive computational methods. The model is evaluated on a new benchmark (MPPW) with 22 datasets having mostly 50 or fewer training labels. 

(b) Strengths:

   - The research addresses an important problem of property prediction with limited experimentally-validated data.
   - The introduction of the MPPW benchmark is valuable as it reflects real-world conditions with scarce labeled data better than existing benchmarks.
  - The paper presents comprehensive experimental evaluations.

(c) Weaknesses:
- The technical novelty of the pretraining strategy is questioned, as masked atom prediction and denoising tasks have been used before. The key innovation of dynamic denoising and branching encoder needs further justification.
- Some experimental details lack sufficient motivation and explanation, such as the choice of auxiliary tasks and the use of a subset of the stage 1 dataset in stage 2.
- The paper could have more in-depth discussions on related works, especially in the area of few-shot molecular property prediction.

(d) Rejection Reasons:
- The lack of clear and significant technical novelty in the proposed pretraining approach is a major concern. Although the authors attempt to highlight the dynamic denoising and branching encoder, the overall contribution seems incremental.
- The insufficient motivation and explanation for certain experimental setups make it difficult to fully assess the validity and significance of the proposed method.
- The limited discussion on related works and the potential overlap with existing research directions reduce the impact and uniqueness of the paper.