Two-Stage Pretraining for Molecular Property Prediction in the Wild

The proposed pre-training strategy is not novel enough as the masking and denoising strategies have been widely adopted in previous pre-training methods, such as Uni-Mol [1] and 3D-denoising.

We are in full agreement that masked atom prediction (MAP) and denoising tasks, both used in our first pretraining stage, have been widely adopted in various works. However, the technical novelty of our first pretraining stage lies in the dynamic denoising and the branching encoder, which decouples the denoising and MAP tasks. This novel design fixes a significant problem in denoising pretraining, where it was believed that large noise scales (e.g., those leading to larger than 1 Å atom displacement) typically hurt the downstream performance [1-4]. With our model, we demonstrate that using larger noise scales (up to 10 Å) actually improves the downstream performance—a new insight for denoising-based molecular pretraining.

We also note that our technical contributions go beyond the pretraining strategy. Specifically, we:

• Propose a new benchmark (Section 4) that reflects the conditions in the real-world. Here, the datasets contain 50 or less training samples with experimentally-validated properties. We believe our benchmark is more realistic compared with QMs or MoleculeNet, which use datasets with hundreds to hundreds of thousand training samples. For reference, only around 0.37% of assays in the ChemBL database have 100 or more labeled molecules.
• Standardize the pretraining dataset in our benchmark to limit the influence of pretraining dataset quality to the downstream performance. The standardization is important to ensure fairness when comparing multiple pretraining methods, as pretraining dataset quality can significantly affect the downstream performance (Section 5.6)
• Provide extensive discussions on the effects of (1) noise scales, (2) pretraining dataset size, (3) pretraining dataset diversity.

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.

We are happy to expand and improve the rationale for the decoupling, which is already discussed in Section 3.1.3. (line 142 - 152). The reviewer is right that these tasks operate on different modalities (coordinates and atom types). However, their joint learning via a shared encoder is suboptimal as the network capacity would be disproportionately allocated to solving the more complex pretraining task, denoising, at the expense of the MAP performance. By introducing a branching structure and connecting the branches through an aggregation module, we allow each task-specific encoder to learn independently and focus on its respective objective. This separation reduces interference between the tasks and enables more efficient optimization, ultimately leading to better downstream performance.

Supervised task learning with DFT-generated properties is also used by Transformer-M[5]. The motivation for splitting the pre-training into two distinct stages needs further clarification.

A single-stage approach such as used by Transformer-M [5] would require either (a) labeling all molecules in the dataset, or (b) a bespoke model whose components can be dynamically activated or deactivated according to label availability. In contrast, our two-stage pretraining approach offers a practical advantage: we can use a large, unlabeled dataset in the first stage and a smaller, labeled dataset in the second stage in a plug and play manner. Any models that can learn via self-supervision can be trained with our approach with minimal modifications (Appendix A2).

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.

Calculating the HOMO, LUMO, and dipole moment of 100K molecules available in GDB17 only takes less than 48 hours. It is a one-time computation that can improve the prediction performance at a negligible cost compared with collecting new labels through wet lab-based experiments (Line 352-355).

Thanks for authors' responses. Some of my concerns have been addressed. However, I still feel that the pre-training strategies lack novelty, and the motivation behind decoupling denoising and MAP remains unconvincing to me. I will discuss with the other reviewers before reaching a final decision.

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.

All experimental results on our benchmark in the Table (1-6, 8) are the average of three different train/test splits (Section 5.3). Unfortunately, we could not include the standard deviations without removing one baseline due to the space limitation. We provide a snippet of the prediction performance on the individual splits in Table A below.

Could the authors conduct experiments to compare dynamic and static denoising pretraining, as well as examine the impact of branching on the results?

Thank you for the suggestion! We have conducted an additional ablation study focusing on the effects of the branching encoder and dynamic noise. The results, presented in Table B below, show a general trend where downstream performance degrades with higher noise levels when a single encoder with static noise approach is used. In most cases, utilizing the branching encoder resolves this issue. Ultimately, combining the branching encoder with dynamic noise yields the best performance across all four assays. We will include this in the revised manuscript.

What is the network structure of the Denoising and MAP heads? What is the "pooling with multihead attention" (PMA)

The MAP and denoising heads are implemented with multilayer perceptrons, while the aggregator module is implemented with pooling by multihead attention (PMA), a cross attention-based module introduced in [10]. The PMA module uses a query of size 1 x 512, and takes the molecule features F as key and value. We have added this explanation to the revised manuscript.

Do the RDKit-generated conformations used here meet this "equilibrium" assumption? RDKit-generated conformations are not necessarily at equilibrium and are less accurate than DFT calculations. However, they are sufficiently accurate for the denoising pretraining task, as also noted by [4].

[4] argued that excessively small or large noise scales negatively impact force field accuracy (extended data Fig. 3 [4]). In contrast, this paper suggests that larger noise scales are beneficial. How do you think of this discrepancy?

Thank you for bringing this work to our attention! We missed this work as it was published shortly before the submission deadline. We believe that both conclusions can be true, since the network architectures are different. The hypothesis in [4] is that “larger noise scales yield more irrational noisy samples” that could “degenerate the (downstream) performance notably”. In [4], the pretraining utilizes static noise scales and does not provide the noise scale information to the encoder. In our work, the denoising encoder is conditioned on the dynamic noise scale, i.e., the network is now aware of how noisy the input is, and is learning to denoise a wide variety of noisy inputs, from trivial to challenging noises. Therefore, the noise scale information and the trivial noisy samples can help the network “rationalize” the more irrational, high-noise samples.

Additionally, in the dynamic denoising pretraining, which is compared to the denoising diffusion model, does the dynamic noise scale σ serve as an input to the network as in diffusion models?

Yes, the noise scale serves as an input to the network. We will fix Figure 1 and Eq. 3 to better reflect this. Thank you for raising this question.

In Section 3.2, you selected HOMO, LUMO, and dipole moment prediction as the second-stage pretraining tasks. Why were these tasks chosen? 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?

Our hypothesis is based on the fact that all molecular properties are a function of their 3D structures (Line 232). Therefore, the pretraining properties do not have to be strongly-correlated with the downstream properties for the supervised pretraining to be beneficial. This is demonstrated in Table 1, as our model achieves good performance across the diverse type of properties in MPPW. We choose HOMO, LUMO, and dipole moment for the pretraining properties because they can be calculated within reasonable time and accuracy using DFT.

Each pretraining stage employs widely used methods such as MAP and denoising. Thus, in terms of technical contribution, this work may lean more toward engineering improvements.

We are in full agreement that masked atom prediction (MAP) and denoising tasks, both used in our first pretraining stage, have been widely adopted in various works. However, the technical novelty of our first pretraining stage lies in the dynamic denoising and the branching encoder, which decouples the denoising and MAP tasks. This novel design fixes a significant problem in denoising pretraining, where it was believed that large noise scales (e.g., those leading to larger than 1 Å atom displacement) typically hurt the downstream performance [1-4]. With our model, we demonstrate that using larger noise scales (up to 10 Å) actually improves the downstream performance—a new insight for denoising-based molecular pretraining.

We also note that our technical contributions go beyond the pretraining strategy. Specifically, we:

• Propose a new benchmark (Section 4) that reflects the conditions in the real-world. Here, the datasets contain 50 or less training samples with experimentally-validated properties. We believe our benchmark is more realistic compared with QMs or MoleculeNet, which use datasets with hundreds to hundreds of thousand training samples. For reference, only around 0.37% of assays in the ChemBL database have 100 or more labeled molecules.
• Standardize the pretraining dataset in our benchmark to limit the influence of pretraining dataset quality to the downstream performance. The standardization is important to ensure fairness when comparing multiple pretraining methods, as pretraining dataset quality can significantly affect the downstream performance (Section 5.6)
• Provide extensive discussions on the effects of (1) noise scales, (2) pretraining dataset size, (3) pretraining dataset diversity.n the effects of (1) noise scales, (2) pretraining dataset size, (3) pretraining dataset diversity.

Quantum properties may be less beneficial for ADMET tasks, raising doubts about the validity of the assumption in line 226, especially with the results in Section 5.4.

Our hypothesis is based on the fact that all molecular properties are a function of their 3D structures (Line 232). Therefore, the pretraining properties do not have to be strongly-correlated with the downstream properties for the supervised pretraining to be beneficial. This is demonstrated in Table 1, as our model achieves good performance across the diverse type of properties in MPPW. Note that the benefits of stage 2 might be reduced, but do not entirely vanish, on properties with weaker correlation to the pretraining properties (Section 5.4.). This is supported by the results in Table 2 row 3, which still consistently outperform the results in Table 2 row 1.

Why was a subset of the stage 1 pretraining dataset chosen for stage 2?

One advantage of splitting the pretraining into two stages is that we can use any unlabeled dataset in stage 1 and any labeled dataset in stage 2. Using a subset of the dataset used in stage 1 pretraining is a practical choice as we already have the data ready in the pipeline.

Do these filtered datasets in Section 5.6 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.

Thank you for pointing this out! The dataset size could indeed affect the downstream performance. That is why we fixed the training set size to 100K for each row in Table 6. We have added this detail to the revised paper.

The authors’ evaluation in high-data settings is insufficiently comprehensive. I recommend a more diverse downstream task evaluation on MoleculeNet or TDC datasets.

Thank you for the suggestion! We have conducted additional experiments on regression tasks in the MoleculeNet benchmark (high-data regime), and the results are shown in the table below. Combined with results in Table 7, we can see that our method outperforms state-of-the-art models in high-data regimes. Note that we have trained the baselines from scratch using the GDB17 dataset. This is to ensure that any performance gains are the results of the pretraining strategy, not the pretraining dataset quality.

It is important to note that our work primarily focuses on real-world scenarios where labeled training data often consists of 50 molecules or fewer. Therefore, the results on our MPPW benchmark (small-data regime) should be considered as the core evaluation, while results from QMs and MoleculeNet are included only for supplementary comparisons.

Thank you for the authors' response. Most of my concerns have been addressed; however, I still have reservations regarding the novelty of the work. I will make a final decision on whether to adjust my score after further discussion with the other reviewers.

Both the denoising pretraining strategy and masking strategy have been explored in previous works. The efficacy of the proposed strategies has been largely validated by previous works.

We are in full agreement that masked atom prediction (MAP) and denoising tasks, both used in our first pretraining stage, have been widely adopted in various works. However, the technical novelty of our first pretraining stage lies in the dynamic denoising and the branching encoder, which decouples the denoising and MAP tasks. This novel design fixes a significant problem in denoising pretraining, where it was believed that large noise scales (e.g., those leading to larger than 1 Å atom displacement) typically hurt the downstream performance [1-4]. With our model, we demonstrate that using larger noise scales (up to 10 Å) actually improves the downstream performance—a new insight for denoising-based molecular pretraining that has never been validated in previous works.

We also note that our technical contributions go beyond the pretraining strategy. Specifically, we:

• Propose a new benchmark (Section 4) that reflects the conditions in the real-world. Here, the datasets contain 50 or less training samples with experimentally-validated properties. We believe our benchmark is more realistic compared with QMs or MoleculeNet, which use datasets with hundreds to hundreds of thousand training samples. For reference, only around 0.37% of assays in the ChemBL database have 100 or more labeled molecules.
• Standardize the pretraining dataset in our benchmark to limit the influence of pretraining dataset quality to the downstream performance. The standardization is important to ensure fairness when comparing multiple pretraining methods, as pretraining dataset quality can significantly affect the downstream performance (Section 5.6)
• Provide extensive discussions on the effects of (1) noise scales, (2) pretraining dataset size, (3) pretraining dataset diversity.

Could authors have more discussions over related works about the few-shot molecular property prediction?

Thank you for the feedback! Although related, our work does not fall into the traditional few-shot learning framework. Prior studies in this area [7-9] typically formulate the few-shot prediction as an N-way K-shot classification problem, where N classes of molecules are sampled from a dataset, each with K examples [7]. This formulation is not applicable to regression tasks, which are the focus of the proposed MPPW benchmark. Therefore, we choose baseline models that follow a pretraining-finetuning paradigm. We have clarified this in the revised related work section.

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.

Due to space limitation, we choose to include UniMol and Mol-AE in Table 1 as they are the latest, best performing methods using MAP and denoising pretraining. We also include GNN-based methods, GraphMVP and Mole-BERT, as additional baselines. Does the reviewer have suggestions on which additional baselines are necessary to be added? We will try to add it during the discussion phase.