Sliced Denoising: A Physics-Informed Molecular Pre-Training Method

We greatly appreciate your expert feedback and your insightful concerns. Your concerns and questions are responded as follows:

Weakness 1: The claim above Eq. 6 is not correct and the authors should rephrase the sentence. The two citations should be removed.

We appreciate the valuable suggestion. Indeed, electrostatic and van der Waals interactions are important in molecular simulation to maintain the long-range structure of molecules in classical simulations. However, we do not expect the self-supervised pretraining force target to be as accurate as precise computational methods, which would greatly increase the burden of the pretraining task. We strike a balance between complexity and effectiveness, and raise the force field accuray by 42%(as shown in experiment in section 4.1) without substentially increasing the time complexity of pre-training(as shown in responses to ft62 question 1).

As for the explaination for the approximation, we would like to rephase is as follows. "Secondly, we drop the last two terms in order to get a quadratic form of energy function in equation 6. Despite these long-range electrostatic and van der Waals interactions are important in classical simulations, we find the accuracy of the approximated energy function is much higher than existing self-supervised pre-training methods in section 4.1." If you have a better explanation or further suggestions, please feel free to respond us.

Weakness 2: I think more details and examples of those parameter vectors can substantially help the reader in understanding the noise design.

Thanks for the helpful suggestion. Here we provide some examples of BAT noise in the figure https://anonymous.4open.science/r/SliDeRebuttal-EEC9/STDBat32.jpg. (a)(b) Standard deviations of BAT noise for acetaldehyde. (c) Standard deviations of BAT noise for 2-methylpyridine. We can see that the noise scale can distinguish different bond types and atom types, resulting in better molecular modeling. We have added the examples in appendix C.1 in the new version.

Weakness 3: Missing parenthesis in the second exponential term of Eq. 8.

Thank you for the kind reminder, and we have revised it in the new version.

Weakness 4: The superiority of the SliDe can be strengthened by more downstream experiments, especially energy prediction, for example, on the ANI-1x dataset.

We add an experiment on ANI-1x energy prediction as below. The result is compared with baseline nonequilibrium denoising (Noneq) [1] that performs coordinate denoising as Coord (Zaidi et al., 2022), while Noneq is pre-trained on ANI-1 and ANI-1x that contain nonequilibrium structures. Noneq also adopts TorchMD-NET as the backbone model. SliDe performs better in both pre-train improvement and without pre-train setting, indicating the superiority of SliDe pre-training method over nonequilibrim coordiante denoising as well as the improved backbone model GET over TorchMD-NET.

[1] Denoise Pretraining on Nonequilibrium Molecules for Accurate and Transferable Neural Potentials. J. Chem. Theory Comput. 2023.

Question 1.1: "What does DFT label supervised mean?" Isn’t “Training from Scratch” also supervised? The authors should elaborate.

"DFT label supervised" means pre-training by learning the DFT force labels. Besides, your understanding of “Training from Scratch”, "Coord" and "Frad" is exactly right. We add more explanations in appendix B.2 in the revised version.

Question 1.2:The unit of the prediction MAE should also be included in the table.

Thank you for the useful suggestion, and we have added the unit in the table.

We have carefully considered your valuable suggestions and have made necessary revisions. The revised content is shown in blue text in the latest version of the paper, which can be found in: https://anonymous.4open.science/r/SliDeRebuttal-EEC9/SLIDE_ICLR2024.pdf

If you have any further questions and feedbacks, please feel free to contact us. If our answer has successfully resolved your issue, we kindly request that you consider raising the rating. Thanks!

I would like to raise the question of if ICLR is the correct venue for this submission. I would assume that a broader audience will have trouble understanding the paper due to the amount of domain knowledge involved.

Response: We believe our paper is very suitable for ICLR for the following reasons.

• First of all, AI for drug discovery has become an important area in machine learning. To our knowledge, there are around 41 papers about molecules and proteins accepted by ICLR2023, accounting for 2.9% of the total accepted papers. It is in the same scale as federated learning, an important branch of the machine learning field, having 44 papers accepted. Besides, there are around 130 papers about molecules or proteins submitted to ICLR2024.
• Additionally, many of our pre-training and property prediction baselines are ICLR conference papers: Coord(ICLR2023), SE(3)-DDM(ICLR2023), Uni-mol(ICLR2023), Transformer-M(ICLR2023), SphereNet(ICLR2022), ET(ICLR2022), DimeNet(ICLR2020).
• Finally, our paper is closely relevant to the topics of "self-supervised representation learning" and "applications to physical sciences (physics, chemistry, biology, etc.)" among the Subject Areas of ICLR2024.

For the convenience of a broader audience, we provide a supplementary introduction to the involved background knowledge in Appendix E in the new version, which can be found at: https://anonymous.4open.science/r/SliDeRebuttal-EEC9/SLIDE_ICLR2024.pdf

If you have any further questions and feedbacks, please feel free to contact us. If our answer has successfully resolved your issue, we kindly request that you consider raising the rating. Thanks!

We greatly appreciate your expert feedback and your insightful concerns. Your concerns and questions are responded as follows:

Weakness 1: Gemnet (2021) and Nequip (2022) can outperform Slide on basically all of the MD17 and require no pre-training. The claim of new state of the arts is incorrect.

First of all, please note that the results of MD17 in Gemnet and Nequip use a different unit (meV/$\mathring{\textnormal{A}}$) than ours (kcal/mol/$\mathring{\textnormal{A}}$). After unifying the unit, SliDe's results in paper surpass Gemnet on 6/8 molecules and NequIP((l= 1) on all 8 molecules in MD17, as shown below.

Furthermore, we find the hyperparameter WoFE indicating weight of force over energy in loss functions is important for MD17 dataset[1]. For fair comparisons with Gemnet, we add a setting that uses the same WoFE as Gemnet that notably raises performance. (Due to the time limit, we have completed the test for five molecules so far, and it is enough to show SliDe's competence.)

We also notice that NequIP reports two versions of results on MD17: l=1 and l=3, where l is the order of spherical harmonics. Recent works[1][2] choose to only compare with NequIP(l=1) because NequIP with high order spherical harmonics is substantially slower[2]. For the same reason, we compare with NequIP(l=1).

Finally, our main contribution is the pre-training method which is applicable to many backbone models. Therefore, SliDe can actually apply to Gemnet, Nequip and other competitive backbone models.

[1] Spherical Message Passing for 3D molecular Graphs, ICLR2022

[2] TorchMD-NET: Equivariant Transformers for Neural Network based Molecular Potentials, ICLR2022

Weakness 2: I think that multiple random seeds should be reported.

All the results reported in paper use the same random seed 1. This follows the practice of previous literature[1][2][3], because the performance is quite stable for random seeds on QM9 dataset, as shown in [3][4][5]. However, to further validate the stability and to answer this question, we conduct experiments with three different seeds (seed=0,1,2) on homo and lumo tasks in QM9 and Aspirin force prediction in MD17. The results are given in the table below, showing that the performance is quite stable across these tasks.

[1]Shikun Feng, Yuyan Ni, Yanyan Lan, Zhi-Ming Ma, and Wei-Ying Ma. Fractional denoising for 3D molecular pre-training.

[2]Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3)-invariant denoising distance matching.

[3]Yi Liu, Limei Wang, Meng Liu, Yu-Ching Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spheri cal message passing for 3D molecular graphs.

[4]Kristof Sch ̈utt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra.

[5]Kristof T Sch ̈utt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R M ̈uller. Schnet – a deep learning architecture for molecules and materials.

Weakness 3: "Showing the best result over random seeds in table 1 is kind of strange."

As discussed in the response to weakness 2, we do not select the best result over random seeds. All the results in our submission are from seed=1.

Question 1: How does your slicing method related to sliced score matching?

The similarity is that they are both random slicing method which uses random projection to reduce computational cost. The word "slicing" both originates from Sliced Wasserstein distance. The difference is that the slicing technique is applied to different targets. In sliced score matching, the technique is used to reduce the computation of the gradient of the score matrix, while it is used to reduce the computation of the Jacobi matrix of the coordinate transformation function in our paper.

Question 2: How does the scale of the pre-training data affect performance?

We randomly extract 10W, 50W, 100W data from PCQM4Mv2 dataset. We pre-train on these subsets and finetune on homo(QM9). The results are provided below. We find more pre-train data contributes to better downstream performance. Meanwhile, SliDe is able to achieve competitive results when data is relatively scarce, such as 50W.

We have carefully considered your valuable suggestions and have made necessary revisions. The revised content is shown in blue text in the latest version of the paper, which can be found in: https://anonymous.4open.science/r/SliDeRebuttal-EEC9/SLIDE_ICLR2024.pdf

If you have any further questions and feedbacks, please feel free to contact us. If our answer has successfully resolved your issue, we kindly request that you consider raising the rating. Thanks!

Question 3: Elaborate on potential plans to validate slide on a broader range of datasets that present different challenges than qm9 and md17.

We plan to evaluate SliDe on larger number and size of molecules. To this end, we choose to test SliDe on MD22[2] and ANI-1x[3] datasets, whose comparison to MD17 and QM9 is shown below. ANI-1x provides large numbers of molecules with multiple equilibrium and nonequilibrium conformations. MD22 contains larger and diversified molecules than QM9 and MD17, including protein, lipids, carbohydrates, nucleic acids and supramolecules.

• Baseline: The result is compared with baseline nonequilibrium denoising (Noneq) [1] that performs coordinate denoising as Coord (Zaidi et al., 2022) in our paper, while Noneq is pre-trained on ANI-1 and ANI-1x that contain nonequilibrium structures. Noneq also adopts TorchMD-NET as the backbone model.

• ANI-1x: The result of experiment on ANI-1x energy prediction is shown below. SliDe performs better in both pre-train improvement and without pre-train setting, indicating the superiority of SliDe pre-training method over nonequilibrim coordiante denoising as well as the improved backbone model GET over TorchMD-NET. | ANI-1x energy prediction (MAE, kcal/mol) | Noneq | SliDe | | --- | --- | --- | | w/o pre-train | 1.50 | 1.362 | | pre-train | 1.01 | 0.786 | | pre-train improvement | 32.7% | 42.3% |

• MD22: Due to the time limit, we only test on nucleic acids, consisting of ATAT(60 atoms) and ATATCGCG(118 atoms). The data split follows [1] (train:validation:test=8:1:1). Since Noneq focus on energy prediction and does not report force prediction result, we reevaluate it with both energy and force tasks as a fair comparison. The results below demostrate SliDe is consistently effective for molecules with varied type and size, indicating the ability to apply to broader range of molecules. | MD22 MAE (kcal/mol/$\mathring{\textnormal{A}}$) | ATAT(60 atoms) | | ATATCGCG(118 atoms) | | | --- | --- | --- | --- | --- | | | Force | Energy | Force | Energy | | Noneq w/o pre-train | 0.2744 | 0.1098 | 0.7581 | 0.2993 | | Noneq w/ pre-train | 0.2194 | 0.0776 | 0.6574 | 0.2764 | | SliDe w/o pre-train | 0.0700 | 0.1021 | 0.0873 | 0.1049 | | SliDe w/ pre-train | 0.0444 | 0.0872 | 0.0596 | 0.0904 |

[1] Denoise Pretraining on Nonequilibrium Molecules for Accurate and Transferable Neural Potentials. J. Chem. Theory Comput. 2023.

[2] Accurate global machine learning force fields for molecules with hundreds of atoms, Science Advances, 2023.

[3] The ANI-1ccx and ANI-1x datasets, coupled-cluster and density functional theory properties for molecules. Scientific data, 2020.

Weakness 2 & Curiosity question: (1) Dependency on equilibrium structures

We find pre-training with inaccurate conformation also works. Due to the time limit, we extract subsets of 10W and 50W molecules from PCQM4Mv2 dataset and generate inaccurate conformations by RDKit (the conformation is optimized by MMFFOptimizeMoleculeConfs function provided by RDKit, but still not as accurate as DFT). We pre-train on the dataset we construct and test them on homo(QM9), as shown below.

Compared with training from scratch MAE=17.6 meV, pre-training with RDKit data is notably effective. Compared with pre-training with accurate conformations, less accurate conformations compromise the performance. Overall, SliDe is robust to inaccurate conformations, revealing the potential of SliDe in larger scale of pre-training.

(2) Are there ways to advance molecular representation learning in such a setting?

Although this issue is beyond the scope of this article, it is an academic issue that we believe is worth exploring. We deem the force and energy labels of the nonequilibrium conformation may be helpful to model the local energy function landscape and design a reasonable pre-training task[4][5]. But the specific usage of labels and task design still needs more research.

[4] May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular Conformations, Neurips 2023

[5] Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields, submitted to ICLR2024

We greatly appreciate your expert feedback and your insightful concerns. Your concerns and questions are responded as follows:

Weakness 1 & question 1: (1)Can the authors provide more details on the computational requirements of slide, especially when applied to large molecular datasets?

Please note that all the terms in the loss function (eq.17) except for $GNN_\theta(x)$, are calculated in data preprocessing. As the deep learning package (specifically PyTorch in this case) generates multiple workers in the DataLoader through parallel processes, each responsible for loading and processing individual data samples before adding the processed batches to a queue, this approach ensures that the time-consuming training process remains unimpeded, as long as the queue is consistently filled with batches. Therefore, it will not become the bottleneck of the training process, as $GNN_\theta(x)$ is what takes most of the time.

The computational complexity of the regression target for one molecule is $O(N_v\times\max$ { N,M,F }) for SliDe, where N is the number of atoms, M is the dimension of relative coordinates (total number of bond lengths, bond angles, and bond torsions), F is the complexity of coordinate transformation between Cartesian coordinates and relative coordinates.

Two factors affecting the complexity is $N_v$ and $F$. As for coordinate transformation, the complexity $F=O(\max$ { $N,M$ })[1], which is executed efficiently by RDKit. As for the $N_v$ times sampling procedure, in fact, for every $i\in\{1,\cdots,N_v\}$, the regression target can be calculated in parallel. Therefore the time complexity above can be further reduced to $O(\max${N,M,F}). This makes pre-training on large molecular datasets possible.

[1]On Updating Torsion Angles of Molecular Conformations, Journal of Chemical Information and Modeling 2005

(2)How does its computational efficiency compare to existing methods?

As for comparison to existing method, SliDe and Frad are theoretically in the same scale. The computational complexity of the regression target for one molecule is $O(\max${ $M_{rb},N,F_{rb}$ }) for Frad, where $M_{rb}$ is the number of rotatable bonds and $F_{rb}$ is the coordinate transformation between Cartesian coordiantes and diheral angles of rotatable bonds. It is in the same scale to SliDe's $O(\max${N,M,F}).

In practise, although we do not process all $N_v$ targets in parallel and incorporate edge update in network architecture, the training time of Frad and SliDe does not vary significantly. Frad pre-training takes 1d1h14m on 8 NVIDIA A100 GPU, and SliDe pre-training takes 1d17h1m on 8 Tesla V100 GPU.

Question2: Generalizability and applicability to other types of data.

Except for small molecules, our SliDe can potentially be used in pre-training for protein and material. To be specific, SliDe can replace coordinate denoising in existing protein pre-training method[1][2] and replace force label prediction in material pre-training method[3]. The parameters in the energy function should be adapted for these scenarios and they can be obtained from AMBER Force Field[4]. As a consequence, SliDe can expand its applications in a series of downstream tasks such as binding affinity prediction, structure-based virtual screening, and material property prediction.

[1] Uni-mol: A universal 3D molecular representation learning framework, ICLR2023

[2] General-purpose Pre-trained Model Towards Cross-domain Molecule Learning, submitted to ICLR2024

[3]From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction, submitted to ICLR2024

[4]https://ambermd.org/AmberModels.php

We have carefully considered your valuable suggestions and have made necessary revisions. The revised content is shown in blue text in the latest version of the paper, which can be found in: https://anonymous.4open.science/r/SliDeRebuttal-EEC9/SLIDE_ICLR2024.pdf

If you have any further questions and feedbacks, please feel free to contact us. Thanks!