Paper: Sliced Denoising: A Physics-Informed Molecular Pre-Training Method (ICLR-2024)

Reviewer comment (Theory):
In section 3.1, the claim above Eq. 6 is not correct. Modeling long-range electrostatic (coulomb) interaction is critical in many areas including energy calculation and md simulation. For large system, long-range vdW outside of a cutoff distance may be neglected, but for small molecules in QM9 and MD17, it should not be neglected. When using GNN that takes 3D coordinate and atom type as input approximate energy function, the model should be able to learn those two terms, thus not affecting the approximation of Eq.6. However, the authors should rephrase the sentence. The two citations associated with the sentence does support the claim so they should be removed.

Author response:
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.

Reviewer comment (Reproducibility):
The noise design is very import in this work. In the BAT noise (Eq. 9), the parameter vectors are critical but there isn't detailed explaination of them. The authors briefly discussed in the section C.1, but I do think more details and example of those parameter vectors can substentially help reader in understanding the noise design.

Author response:
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.

Reviewer comment (Writing):
Missing parenthesis in the second exponential term of Eq. 8. $(\theta_{i} - \theta_{i, 0})^2$

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

Reviewer comment (Experiments):
I do think the superiority of the SliDe method can be strengthen by more downstream experiments, especially energy prediction. For example, the ANI-1x dataset (www.nature.com/articles/s41597-020-0473-z) is a excellent dataset for such task.

Author response:
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.
| 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%  |

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

Reviewer comment (Presentation):
Table 7 is confusing. My understanding is that “Training from Scratch” meaning no pretraining, and Coord and Frad meas pretrained with different method and then fine-tuned on MD17-Aspirin. What does DFT label supervised mean? Isn’t “Training from Scratch” also supervised? The authors should elaborate. The unit of the prediction MAE should also be included in the table.

Author response:
"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.

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

Reviewer comment (Experiments):
The baselines used in this paper are not up to date. In fact, Gemnet (2021) and Nequip (2022) can outperform Slide on basically all of the MD17 and require no pre-training. In addition, the authors claim they set a new state of the art on these benchmarks, which is incorrect.

Author response:
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.
| Force prediction MAE (kcal/mol/ $\mathring{\textnormal{A}}$) | WoFE | ASP | BEN | ETH | MAL | NAP | SAL | TOL | URA |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Gemnet-T | 1000 | 0.219  | **0.145** | **0.085**  | 0.155  | 0.055  | 0.127 | 0.060 | 0.097  |
| GemNet-Q | 1000 | $\underline{0.217}$  | **0.145** | $\underline{0.088}$  | 0.159  | 0.051  | 0.125  | 0.060 | 0.104  |
| NequIP((l= 1) | 1000  | 0.348 | 0.187 | 0.208 | 0.337 | 0.097 | 0.238 | 0.101 | 0.173 |
| SliDe(in paper) | 4 | **0.174**   | $\underline{0.169} $ | $\underline{0.088} $ | $\underline{0.154 }$ | $\underline{0.048}$ | $\underline{0.101}$ | $\underline{0.054}$ | $\underline{ 0.083 }$  |
| SliDe | 1000 | - | - | **0.085** | - | **0.043** | **0.095** | **0.049** | **0.079** |

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

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

Reviewer comment (Experiments):
For the downstream task only one random seed is used and the gains over other methods are relatively minor (i.e. compared to Coord and Frad). This makes me doubt that the results are really significant or if they are just due to tuning. I think that multiple random seeds should be reported.

Author response:
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.
| MAE | homo(meV) | lumo(meV) | ASP(kcal/mol/ $\mathring{\textnormal{A}}$) |
| --- | --- | --- | --- |
| seed=0 | 14.0 | 12.3 | 0.1714 |
| seed=1 | 13.6 | 12.3 | 0.1740 |
| seed=2 | 13.3 | 12.2 | 0.1744 |
| Mean(Standard deviation) | 13.63(0.35)  | 12.27(0.06) | 0.1733(0.0016) |



[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.

Reviewer comment (Evaluation):
Showing the best result over random seeds in table 1 is kind of strange. I think that the mean result should be shown.

Author response:
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.

Reviewer comment (Novelty):
How does your slicing method relate to sliced score matching [1]? [1] Song, Yang, et al. "Sliced score matching: A scalable approach to density and score estimation." Uncertainty in Artificial Intelligence. PMLR, 2020.

Author response:
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.

Reviewer comment (Experiments):
How does the scale of the pre-training data affect performance?

Author response:
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.
| Number of pre-training data | homo(MAE,meV) |
| --- | --- |
| w/o pretrain | 17.6 |
| 10W | 16.05 |
| 50W | 14.53 |
| 100W | 14.21 |
| ~300W(whole) | 13.6 |

Reviewer comment (Evaluation):
computational complexity: while the random slicing technique addresses computational challenges associated with jacobian matrix estimation, the overall computational demand and efficiency, especially in large-scale applications, are not comprehensively addressed​​​​​​.

Author response:
>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.

Reviewer comment (Experiments):
regarding computational efficiency: can the authors provide more details on the computational requirements of slide, especially when applied to large molecular datasets? how does its computational efficiency compare to existing methods?

Author response:
>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.

Reviewer comment (Experiments):
on generalizability and applicability: what are low-hanging fruits to test the generalizability of slide to other types of geometric data or applications beyond molecular science? how might the method need to be adapted for such scenarios?

Author response:
>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

Reviewer comment (Experiments):
empirical validation across diverse datasets: could the authors elaborate on potential plans to validate slide on a broader range of datasets, particularly those that may present different challenges than qm9 and md17, such as des15k or oc20 as in the coord paper https://arxiv.org/abs/2206.00133?

Author response:
>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.

| Dataset | Molecule size | Number of molecules  | Number of conformations |
| --- | --- | --- | --- |
| QM9    | ∼18 (3-29)  | 133,885 | 133,885 |
| MD17    | ∼13 (9-21)  | 8 | 3611115  |
| MD22 | ∼67 (42-370) | 7 | 223422 |
| ANI-1x | ∼ 15 (4-63) | 63865 | 5496771 |

- 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.

Reviewer comment (Theory):
dependence on equilibrium structures: the method's reliance on equilibrium structures, to be clear same as most other methods in this space, for training may limit its effectiveness in scenarios where such structures are not readily available or accurate. are there ways to advance molecular representation learning in such a setting?

Author response:
>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. 
| homo(MAE,meV) | PCQ(DFT) | PCQ(RDKit) |
| --- | --- | --- |
| 10W | 16.05 | 16.13 |
| 50W | 14.53 | 15.39 |

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