Paper: Gradual Optimization Learning for Conformational Energy Minimization (ICLR-2024)

Reviewer comment (Experiments):
Using GO to generate the baseline dataset may limit the scalability of the approach.

Author response:
While working on the topic of molecular optimization, we found that this area is severely under-researched. To the best of our knowledge, there are no papers that present the idea of enriching the training dataset with optimization trajectories. However, in this year Open Catalyst challenge [https://opencatalystproject.org/challenge.html], a similar dataset of optimization trajectories for adsorbate-catalyst pairs has been published. We will cite the OCP23 challenge in the revision.

Reviewer comment (Experiments):
It is unclear why the authors do not use the molecular databases they mention in the “Introduction” section to extract the baseline dataset.

Author response:
First, the limited scalability of this approach is stated as the main motivation for our proposed GOLF framework. Second, we would like to note that we are not aware of any molecular databases that contain optimization trajectories. In our paper, we identify enriching the training dataset $\mathcal{D}_0$ with optimization trajectories as a crucial step to overcome the distribution shift. In Figure 1, we illustrate that enriching the training dataset with optimization trajectories helps alliviate the distribution shift. We agree, that in theory, datasets containing both near-optimal and sub-optimal conformations can be used to improve the optimization quality of NNPs. To test that, we trained a baseline model $\operatorname{SPICE}-f^{\text{baseline}}$ on a subset of SPICE dataset and showed that it suffers from the similar issues (see answer to “Experiments” section).

Reviewer comment (Writing):
Some of the discussion is somewhat cryptic and can benefit from some additional discussion or graphics.

Author response:
We appreciate your suggestions on improving the writing of our paper! Following various ideas proposed by the reviewers we have reworked the manuscript and uploaded a corrected version as a revision. If you still have any concerns or suggestions, please let us know. We will be happy to improve the writing further.

Reviewer comment (Experiments):
The results seem anecdotal, tied to the selected dataset and molecules, and thus may not generalize particularly well.

Author response:
In our experiments, we use a randomly selected subset $\mathcal{D}_0$ of the nablaDFT dataset. The conformations in the nabla DFT dataset are generated based on SMILES representations of molecules from the MOSES dataset. The MOSES dataset consists of a large number of druglike molecules and is considered to be representative.  To ensure reasonable computational cost of DFT computations, the subset  $\mathcal{D}_0$ is limited to molecules with number of atoms < 36. However, apart from the increased computational cost we do not forsee any issues with using GOLF with bigger molecules.

Therefore we significantly extended the main evaluation dataset (10x) by sampling more conformations with less than 35 atoms from the nablaDFT dataset. The extended evaluation dataset contains 20k conformations for X molecules. It contains  $\mathcal{D}\_{\text{test}}$ as its subset.
Moreover, we selected 1828 molecules (1 conformation per molecule) from the SPICE dataset, optimized them with $\mathcal{O}\_G$ to get the optimal energy. We evaluated all non-generative models from our paper on this new SPICE test set $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$.
We provide the results in the table below. We understand that the size of the $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$ may raise similar concerns, but we were unable to collect a larger dataset during the rebuttal due to the need for an expensive $\mathcal{O}\_G$ optimizations. We will extend the $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$ to ~20k conformations and provide results for generative models in the camera-ready version.

### SPICE test (1828 conformations)
|  | OpenFF | RDKit | $f^{baseline}$ | $f^{traj-10k}$ | $f^{traj-100k}$ | $f^{traj-500k}$ | $f^{GOLF-1k}$ | $f^{GOLF-10k}$ | $SPICE-f^{baseline}$ | $SPICE-f^{GOLF-10k}$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\overline{\operatorname{pct}}\_{100}$ | 61.14 | 77.58 | 72.71 | 73.22 | 77.06 | 80.25 | 80.31 | 82.41 | 79.76 | 86.59 |
| $\operatorname{pct}\_{\text{div}}$ | 5.38 | 8.15 | 27.73 | 17.45 | 16.02 | 12.80 | 13.56 | 12.41 | 23.08 | 15.09 |
| $\overline{E^{\text{res}}\_{100}}$ | 12.58 | 8.42 | 13.62 | 9.63 | 8.09 | 6.88 | 6.63 | 6.06 | 6.06 | 4.19 |
| $\operatorname{pct}\_{\text{solved}}$ | 0.05 | 10.72 | 6.35 | 5.17 | 15.90 | 25.65 | 17.53 | 27.73 | 8.91 | 19.41 |

The OpenFF is a non-neural force field used in the original SPICE paper [1]. $SPICE-f^{baseline}$ was trained on ~10000 conformations (2.5 conformations per molecule) from the SPICE dataset. $SPICE-f^{GOLF-10k}$ was trained with the same hyperparameters as $f^{GOLF-10k}$ and used $SPICE-f^{baseline}$ to initialize the NNP. These preliminary results show that models trained on nablaDFT and SPICE datasets demonstrate comparable performance on $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$. 

### Extended test (19477 conformations)

|  | RDKit | $f^{baseline}$ | $f^{rdkit}$ | $f^{traj-10k}$ | $f^{traj-100k}$ | $f^{traj-500k}$ | $f^{GOLF-10k}$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $\overline{\operatorname{pct}}\_{100}$ | 85.47 | 77.88 | 93.04 | 95.08 | 96.15 | 98.75 | 98.78 |
| $\operatorname{pct}\_{\text{div}}$ | 0.62 | 7.46 | 4.43 | 4.50 | 2.77 | 1.99 | 2.98 |
| $\overline{E^{\text{res}}\_{100}}$ | 5.50 | 8.58 | 2.77 | 1.96 | 1.50 | 0.53 | 0.51 |
|  $\operatorname{pct}\_{\text{solved}}$ | 4.05 | 8.18 | 35.42 | 37.01 | 52.71 | 73.41 | 77.26 |

These results show that all conclusions made in the first version of the paper hold on an extended evaluation dataset.

[1] Eastman, P., Behara, P. K., Dotson, D. L., Galvelis, R., Herr, J. E., Horton, J. T., ... & Markland, T. E. (2023). Spice, a dataset of drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1), 11.

Reviewer comment (Experiments):
GOLF requires running the high-fidelity simulation, GO, to produce the anti-gradients initial training data. In a different section of the paper, the claim is that for a sufficiently complex system, the physical simulation may not succeed in completing even a single iteration in a reasonable amount of time. Taken together, these two statements seem to suggest that GOLF is atomic complexity scale limited. How do the authors claim to address this apparent limitation?

Author response:
Yes, the GOLF indeed requires a high-fidelity simulator to collect the training data which may may fail on a sufficiently complex system. However, any deep learning model that aims to match the quality of the physical simulator is subject to the same limitation. Motivated by this limitation, we tested the generalization of NNPs trained with GOLF to larger molecules (see section 6.4).

Reviewer comment (Experiments):
The additional data generated for Active Learning are selected based only on errors. Without some sort of approach to balance the introduction of new data, does the approach bias the dataset distribution, causing the learned distribution to fail to generalize to other molecules?

Author response:
First of all, we would like to note that we test the NNPs on a subset from the same distribution than the train dataset. The train dataset $\mathcal{D}\_0$ and the evaluation dataset $\mathcal{D}\_{\text{test}}$ share no molecules. Our evaluation results described in Section 6.1 and 6.2 show that the NNP successfully generalizes to unseen molecules from the same distribution. Considering the balancing of the new data, we use the following technique: 10% of each mini-batch is sampled from the initial training dataset $\mathcal{D}\_0$. We have added this detail to the appendix of the revised manuscript.

Reviewer comment (Novelty):
How different is the GOLF approach from the active learning technique presented by Kulichecnko et al. (2023)?

Author response:
The method presented in the Kulichenko et al. paper is an interesting and promising approach for the exploration of the configuration space of a single molecule. Unfortunately, it requires training a new NNP ensemble for each new molecule, which is a completely different problem setup compared with ours. Moreover, the focus of the method is to enrich the MD trajectories with the conformations for which the ensemble is uncertain, which is shown to improve the energy prediction. In theory, this approach may also improve the atomic forces prediction, but it remains to be studied.

Reviewer comment (Writing):
The writing is not very clear, especially in the introduction section. It takes me a while to understand simply enriching the training dataset is actually a preliminary baseline method the authors want to compare with. Also, lots of experiment details are mixed with the method, which makes the paper not very easy to read.

Author response:
We agree that some additional context regarding the use of $f^{\text{traj-*}}$ as a baseline is needed. We added clarifying sentences in the revision of the manuscript. Unfortunately, we do not see a way to separate the method from some experiment details, as the GOLF is mainly motivated by the empirical results obtained in section 4. We believe that introducing our framework without a preliminary discussion of experimental details and results will be more confusing for the reader. We will be happy to consider any suggestions on improving the writing of the manuscript.

Reviewer comment (Evaluation):
The calculation of COV and MAT looks problematic. It seems the authors optimize **one conformation per molecule** and take them of the entire test set as the generation set. However, in the conformation generation setting, models generate **multiple conformations per molecule** to construct the generation set, and then COV and MAT are calculated per molecule, and finally the average / median of them on the entire test set is reported.

Author response:
We agree that such metrics are indeed not ideal for the conformation optimization task due to the problem you described above. However, we decided to compute these metrics to provide a fair comparison with ConfOpt and Torsional Diffusion, as the main results in these papers are reported in terms of COV and MAT.

Reviewer comment (Experiments):
Are the conformations in nablaDFT dataset equilibrium ones or the intermediate state sampled from the optimization process? How large is the training set D0? More description about this dataset is needed.

Author response:
The conformations in the nablaDFT are generated using the following procedure:

1. First, several thousand different conformations per molecule are sampled using Rdkit's ETKG [1] based on a SMILES representation of the molecule.
2. Second, these conformations are clustered using Butina [2].
3. The authors then identify the smallest number of clusters needed to cover 95% of generated conformations.
4. Centroids of these identified clusters are chosen and evaluated with a DFT-based oracle.

The training dataset contains 4000 molecules and 10000 conformations (~2.5 conformations per molecule). We apologize for not providing this information in the first version of the paper. We added this information in the revision. We also added a brief description of the nablaDFT dataset in the appendix of the revised manuscript.

[1] [Wang, S., Witek, J., Landrum, G. A., & Riniker, S. (2020). Improving conformer generation for small rings and macrocycles based on distance geometry and experimental torsional-angle preferences. *Journal of chemical information and modeling*, *60*(4), 2044-2058.] 

[2] Barnard, J. M., & Downs, G. M. (1992). Clustering of chemical structures based on two-dimensional similarity measures. *Journal of chemical information and computer sciences*, *32*(6), 644-649.

Reviewer comment (Evaluation):
ConfOpt and TorsionDiff are designed to generate equilibrium low-energy conformers, and not guaranteed to achieve a lower energy by repeatedly applied. Thus, I think it's unfair to compare these models with GOLF in terms of pct.

Author response:
We would like to clarify that in our experiments, these models are trained to generate a single optimal conformation from an initial random conformation. We use these models in the intended way and do not repeatedly apply any of them. The training dataset for these models consists of optimal conformations obtained with $\mathcal{O}_G$ and, thus, we expect them to generate optimal conformations, altough the inference of these models is different to $f^{\text{GOLF-*}}$, we think it is still valuable to compare the generative approach (ConfOpt, TorsionDiff) with an iterative optimization approach.

Reviewer comment (Writing):
The statement of "We hypothesize that in the case of ConfOpt, the main problems are the choice of the architecture and the fact that the model generates optimal conformations from SMILES and does not use initial geometries." doesn't make sense to me. ConfOpt takes the 2D molecular graph as input and also utilizes initial 3D conformations.

Author response:
We follow the original implementation [https://github.com/guanjq/confopt_official], which takes the 2D molecular graph, generates a random conformation with Rdkit’s ETKG, and optimizes it with MMFF. This implementation does not provide any means to supply the model with an initial conformation of our choice.

Reviewer comment (Experiments):
Does $f^{traj-10k /100k }$ keep the total number of updates equal to $5 \times 10^5$? If so, please provide more training details, otherwise, the comparison between them and $f^{GOLF}$ is unfair.

Author response:
Yes, the number of NN updates for all models except for $f^{\text{traj-500k}}$ is equal to $5 \times 10^5$. For $f^{\text{traj-500k}}$ the number of updates is equal to $10^6$ as $5 \times 10^5$ is not enough for the training to converge. We have added a table with hyperparameters along with an additional explanatory sentence in the appendix.

Reviewer comment (Experiments):
Dataset Limitation: The paper may be limited by the availability and diversity of datasets used for testing, potentially impacting the generalizability of the results.

Author response:
We agree that the experimental section could benefit from a more diverse evaluation dataset. Therefore we significantly extended the main evaluation dataset (10x) by sampling more conformations with less than 35 atoms from the nablaDFT dataset. The extended evaluation dataset contains 20k conformations for X molecules. It contains  $\mathcal{D}\_{\text{test}}$ as its subset.
Moreover, we selected 1828 molecules (1 conformation per molecule) from the SPICE dataset, optimized them with $\mathcal{O}\_G$ to get the optimal energy. We evaluated all non-generative models from our paper on this new SPICE test set $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$.
We provide the results in the table below. We understand that the size of the $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$ may raise similar concerns, but we were unable to collect a larger dataset during the rebuttal due to the need for an expensive $\mathcal{O}\_G$ optimizations. We will extend the $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$ to ~20k conformations and provide results for generative models in the camera-ready version.

### SPICE test (1828 conformations)
|  | OpenFF | RDKit | $f^{baseline}$ | $f^{traj-10k}$ | $f^{traj-100k}$ | $f^{traj-500k}$ | $f^{GOLF-1k}$ | $f^{GOLF-10k}$ | $SPICE-f^{baseline}$ | $SPICE-f^{GOLF-10k}$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\overline{\operatorname{pct}}\_{100}$ | 61.14 | 77.58 | 72.71 | 73.22 | 77.06 | 80.25 | 80.31 | 82.41 | 79.76 | 86.59 |
| $\operatorname{pct}\_{\text{div}}$ | 5.38 | 8.15 | 27.73 | 17.45 | 16.02 | 12.80 | 13.56 | 12.41 | 23.08 | 15.09 |
| $\overline{E^{\text{res}}\_{100}}$ | 12.58 | 8.42 | 13.62 | 9.63 | 8.09 | 6.88 | 6.63 | 6.06 | 6.06 | 4.19 |
| $\operatorname{pct}\_{\text{solved}}$ | 0.05 | 10.72 | 6.35 | 5.17 | 15.90 | 25.65 | 17.53 | 27.73 | 8.91 | 19.41 |

The OpenFF is a non-neural force field used in the original SPICE paper [1]. $SPICE-f^{baseline}$ was trained on ~10000 conformations (2.5 conformations per molecule) from the SPICE dataset. $SPICE-f^{GOLF-10k}$ was trained with the same hyperparameters as $f^{GOLF-10k}$ and used $SPICE-f^{baseline}$ to initialize the NNP. These preliminary results show that models trained on nablaDFT and SPICE datasets demonstrate comparable performance on $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$. 

### Extended test (19477 conformations)

|  | RDKit | $f^{baseline}$ | $f^{rdkit}$ | $f^{traj-10k}$ | $f^{traj-100k}$ | $f^{traj-500k}$ | $f^{GOLF-10k}$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $\overline{\operatorname{pct}}\_{100}$ | 85.47 | 77.88 | 93.04 | 95.08 | 96.15 | 98.75 | 98.78 |
| $\operatorname{pct}\_{\text{div}}$ | 0.62 | 7.46 | 4.43 | 4.50 | 2.77 | 1.99 | 2.98 |
| $\overline{E^{\text{res}}\_{100}}$ | 5.50 | 8.58 | 2.77 | 1.96 | 1.50 | 0.53 | 0.51 |
|  $\operatorname{pct}\_{\text{solved}}$ | 4.05 | 8.18 | 35.42 | 37.01 | 52.71 | 73.41 | 77.26 |

These results show that all conclusions made in the first version of the paper hold on an extended evaluation dataset.

Reviewer comment (Theory):
Complexity: It seems that the complexity of GOLF is not clearly discussed in the paper.

Author response:
Inference: GOLF uses NNPs that rely on Message Passing [2]. In the worst case, the computational complexity is quadratic in terms of the number of atoms in the molecule. We provide average inference time for baselines and NNPs and compare them with the average $\mathcal{O}_G$. All the measurements were made on the same machine with a 16-core Intel(R) Xeon(R) Gold 6278C and a single Tesla V100 GPU. The results are in the table below.

| Model | Iterative Optimization with NNP (100 steps) | ConfOpt (single step) | Torsional Diffusion (single step) | UniMol+ (single step) | Iterative Optimization with $\mathcal{O}_G$ (until convergence, ~25 steps on average) |
| --- | --- | --- | --- | --- | --- |
| Average Inference time (seconds) | 1.312 | 0.005 | 0.48 | 0.002 | 2634.787 |

Training: around 2/3 of the computational complexity during training comes from various NN operations. We use DFT-based oracle, which complexity is $O(N^4)$ [3]. For example, the training of the GOLF-10k on our machine with 120 workers for DFT calculations takes approximately 20 hours .

[1] Eastman, P., Behara, P. K., Dotson, D. L., Galvelis, R., Herr, J. E., Horton, J. T., ... & Markland, T. E. (2023). Spice, a dataset of drug-like molecules and peptides for training machine learning potentials. Scientific Data, 10(1), 11.

[2] Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., & Dahl, G. E. (2020). Message passing neural networks. *Machine learning meets quantum physics*, 199-214.

[3] Kohn, W., & Sham, L. J. (1965). Self-consistent equations including exchange and correlation effects. *Physical review*, *140*(4A), A1133.

Reviewer comment (Experiments):
Why not use more datasets besides nablaDFT?

Author response:
We have not considered datasets such as QM9 and GEOM as they contain equilibrium conformations. We have not considered MD17 and other similar datasets due to the low diversity of the molecules.

Additionally, to answer your first question, we trained the $f^{\text{baseline}}$ and the $f^{\text{GOLF-10k}}$ on a subset of SPICE dataset. For the SPICE training set $\mathcal{D}^{\text{SPICE}}\_{\text{train}}$, we collected 4000 molecules with roughly 2.5 conformations per molecule resulting in approximately 10k conformations (the same size as $\mathcal{D}\_{\text{train}}$). We evaluated these models both on the $\mathcal{D}\_{\text{test}}^{\text{SPICE}}$ and the extended $\mathcal{D}\_{\text{test}}$. The results are in the table below.

| NablaDFT | $f^{GOLF-10k}$ | $\operatorname{SPICE}-f^{GOLF-10k}$ |
| --- | --- | --- |
| $\overline{\operatorname{pct}}_{100}$ | 98.78 | 91.71 |
| $\operatorname{pct}_{\text{div}}$ | 2.98 | 12.69 |
| $\overline{E^{\text{res}}\_{100}}$ | 0.51 | 3.27 |
| $\operatorname{pct}\_{\text{solved}}$ | 77.26 | 17.27 |

| SPICE | $f^{GOLF-10k}$ | $\operatorname{SPICE}-f^{GOLF-10k}$ |
| --- | --- | --- |
| $\overline{\operatorname{pct}}_{100}$ | 82.41 | 86.59 |
| $\operatorname{pct}_{\text{div}}$ | 12.41 | 15.09 |
| $\overline{E^{\text{res}}\_{100}}$ | 6.06 | 4.19 |
| $\operatorname{pct}\_{\text{solved}}$ | 27.73 | 19.41 |

The main take-away from these tables is that the percentage of optimizations solved is larger for the $f^{\text{GOLF-10k}}$ trained on NablaDFT. This is likely due to the presence of near-optimal conformations in SPICE which can in theory result in collecting highly correlated training samples.

Reviewer comment (Theory):
What's the complexity of GOLF? It seems that you only show experiment results on subsets of nablaDFT. Is it because the computational complexity of the method is very high?

Author response:
The high computational complexity of the evaluation procedure is tied to the need to optimize every conformation with $\mathcal{O}_G$. Estimating the percentage of optimized energy and other reported metrics requires running a minimization procedure with a DFT-based oracle. For example, it takes ~593 CPU days to calculate optimal energies for 19477 conformations in the extended evaluation dataset. While the optimization of 19477 conformation with GOLF-10k takes ~7.1 hours on a single V100 GPU.