Chemistry-Inspired Diffusion with Non-Differentiable Guidance

Thank you for the detailed suggestions on the presentation of our paper! We address the weaknesses and questions as follows.

1. Justification of SPSA & Comparison to other methods

• SPSA justification: we have added more descriptions and our motivation for SPSA in Appendix F Motivation for Gradient Estimation with SPSA. We previously considered using finite difference to calculate the gradient as $\frac{F([z_{x, t}+\zeta 1_{N \times 3}, z_{h, t}])-F([z_{x, t}-\zeta 1_{N \times 3}, z_{h, t}])}{2\zeta}$, however, it violates the zero-mean requirements for 3D diffusion models to be equivariant (i.e., $z_{x, t}\pm\zeta 1_{N \times 3}$ shifts the mean by $\pm\zeta$), which motivates the choice of our SPSA based estimation (i.e., the perturbation is sampled from the normal distribution).

• Alternative methods: we have conducted new experiments and added the results for evolutionary algorithms in the (global) official comments above. (please see Appendix I in our current revised paper). It turns out that ChemGuide outperforms the evolutionary algorithm in terms of force RMS and energy about ground state.

2. Quality of the guidance gradients

We employed a neural network potential AIMNet2 [1] to calculate the force, and compared its backpropagated gradient with ChemGuide measured in cosine similarity in Appendix H3 ChemGuide vs. Neural Guidance on Force. We can observe that the cosine similarity over time oscillates around 0, suggesting that the guidance from ChemGuide and neural networks are different, hence the discrepancy in performance.

We report the results of neural network guidance on force in the (global) official comments above. (Appendix H3 in our current revised paper). We observe that using neural networks for force guidance underperforms ChemGuide, and significantly increases memory demands, which makes guidance on large molecules extremely difficult (e.g., a 48 GiB GPU can only generate 5 molecules at a time on GEOM, making sampling a large amount of molecules almost linear in time).

3. Compatibility between guidance methods and property

From the chemistry definitions, there is no direct evidence to suggest that one property would be harder or easier than another for the neural regressor to learn. However, we can observe from Figure 7 on the generalization analysis of neural regressor, that on property $\alpha$ the regressor is more robust in terms of perturbations compared to other properties, meanwhile, in Table 4, noisy guidance performs better than clean guidance on $\alpha$, and vice the versa on other properties. We hypothesize that this is because noisy guidance requires backpropagation through the diffusion model, while clean guidance translates the optimized representation $x_0$ in the clean space back to $x_t$ with pre-defined schedules. Thus, with a robust regressor (e.g., on $\alpha$), noisy guidance might be preferred over clean guidance as it directly backpropagates to $x_t$; given a less precise regressor on out-of-distribution molecules, clean guidance would be favored as the errors are more likely to be broadcast through both the VAE decoder and the latent diffusion model for noisy guidance.

However, our analyses are primitive and do not include many factors (e.g., different architectures of the diffusion models), and further investigation in this direction would be very valuable.

Thank you for providing new experiments and revising your manuscript to incorporate the changes requested.

Overall, while I agree with the other reviewers about the existence of multiple non-differentiable guidance methods, I find the simplicity of the current method (using zeroth-order optimization to approximate the gradients of the oracle) appealing, especially in light of the discussion the authors provided when comparing to evolutionary algorithms and a property predictor neural network.

I will thus maintain my score.

Thank you for taking the time and we appreciate your input to help make our paper a better work! We address the weaknesses and questions as follows.

1. Novelty and motivation

We would like to re-emphasize that our focus is on, for molecule optimization, how to achieve controllable generation at inference time from quantum chemistry to bypass the requirement to train a neural property predictor or a conditional diffusion model, which require a large amount of labels that are hard to acquire. We choose stability (defined only on 3D) to explore this motivation, and hence experiment on 3D diffusion models.

We are not proposing a new guidance method (e.g., as in [2]) nor are claiming to be the first to introduce derivative-free guidance to diffusion models. In fact, our method builds on previous proposed guided diffusion methods [3, 4], and the use of zeroth-order optimization techniques [5, 6] for gradient estimation is to overcome the difficulty that we can not backpropagate the chemistry software.

The evolution algorithm [1] relies on stochastic mutations and pruning, which requires a large search space to be effective and do not directly change the underlying conditional distribution from which the diffusion model samples. Gradients provide informative guidance by indicating the correct direction for optimizing stability. Although the evolutionary algorithm selects the best variant at each evolutionary step, it doesn’t rely on the gradient, so optimization direction is essentially random, making the process less controllable.

For tSMC[2], it proposes a novel guidance method and requires $p_{y|x^0}(y|x^0)$, which is a classifier that is trained on labels (refer to Section 5.2 and 6.2 in [2] on the choice of such likelihood $p$) and thus not directly comparable to our works, further, they experiment on protein binding, which is intrinsically different from our task in molecule optimization.

We have excluded the sentence “To the best of our knowledge, this is the first work of its kind.” in the introduction to avoid further confusion in our current revised paper.

2. Related works

For the broader discussion of guided diffusion beyond the scope of our work on molecule optimization, we have added the papers you suggested to enlighten further discussions.

3. Additional experiments

As stated above, our motivation is not derivative-free guidance but to bypass the need to train a neural predictor or a conditional diffusion model to achieve controllable generation.

• Evolutionary algorithm: we added the results using the evolutionary algorithm [1,7] to the global comment above, please also see Appendix I for detailed results and analysis. We can observe ChemGuide outperforms the evolutionary algorithm in terms of force RMS and energy above the ground state.
• Neural network guidance: we replaced the non-differentiable oracle with a neural network potential AIMNet2, the results are reported in the global comment above, please also see Appendix H3 for detailed results and analysis. It turns out that using AIMNet2 for guidance negatively impacted performance.

References

[1] Schneuing, A., Du, Y., Harris, C., Jamasb, A., Igashov, I., Du, W., Blundell, T., Lió, P., Gomes, C., Welling, M. and Bronstein, M., 2022. Structure-based drug design with equivariant diffusion models. arXiv preprint arXiv:2210.13695.

[2] Wu, L., Trippe, B., Naesseth, C., Blei, D. and Cunningham, J.P., 2024. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36.

[3] Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis, 2021.

[4] Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209.14734, 2022.

[5] James C. Spall. Multivariate stochastic approximation using a simultaneous perturbation gradient approximation. IEEE Transactions on Automatic Control, 37:332–341, 1992. URL https://api.semanticscholar.org/CorpusID:122365276.

[6] Yurii Nesterov and Vladimir Spokoiny. Random gradient-free minimization of convex functions. Foundations of Computational Mathematics, 17(2):527–566, 2017

[7] Huang, Y., Ghatare, A., Liu, Y., Hu, Z., Zhang, Q., Sastry, C.S., Gururani, S., Oore, S. and Yue, Y., Symbolic Music Generation with Non-Differentiable Rule Guided Diffusion. In Forty-first International Conference on Machine Learning, 2024.

Thank you for taking the time to review! We appreciate your feedback and address the weaknesses and questions as follows.

1. Relaxation with xTB

• Comparing to GeoLDM with xTB relaxation We have included this baseline in tables 1-3 by reporting “energy above ground state”, which is defined in Appendix B as “ the difference between the energy before and after using xTB or DFT (specified in the context) to optimize the geometry of the generated molecules”. We can observe that ChemGuide consistently achieves the smallest energy difference before and after quantum chemistry optimization at different levels (xTB, DFT) across datasets (QM9, GEOM).

• Relaxation inside xTB The number of geometric optimization steps is decided by xTB given the specific input molecule, which is a built-in and required process that can not be reduced or relaxed to achieve speed-up.

• Relaxation on the guidance steps of ChemGuide There are two ways to accelerate ChemGuide, 1. add guidance for every $n$ step; 2. add less consecutive guidance steps.

  • Every $n$ step guidance, the results are in Appendix H1 (previously Appendix F1), where we choose $n=1, 3, 5$. We can observe that a. $n=1$ (no relaxation) performs the best, suggesting a trade-off between time and performance, and b. ChemGuide with $n=3, 5$ skip steps performs better than the baselines with smaller forces and energy above the ground state, demonstrating the effectiveness of our method.

  • Less consecutive guidance steps In the global official comment above, we conduct experiments that apply the guidance for 100, 200, and 300 steps to generate 500 molecules, and report the results (updated in Appendix H2 in our current revised submission). Similar trends can be observed that reducing guidance steps lowers performance.

• Relaxation on replacing xTB with neural networks In the global official comment in Table 3-8 above, we also relax the dependency on xTB and replace xTB with a neural network AIMNet2 [1], and report the results (detailed results and analysis are in H.3 of our current revised submission). We can observe that using a neural network as guidance hurts the performance, and further, significantly increases GPU memory demands.

• Hardware level relaxation, we can increase the number of CPU cores assigned to xTB when calculating force, and this functionality is implemented in our code.

2. Order of Section 3.1

We changed the order accordingly and moved Section 3.1 on noisy and clean guidance into Section 2 Preliminary. Thank you for your constructive suggestion to present our paper clearly!

3. Eq. 15 typo

We fixed the typo accordingly. Thank you for your constructive suggestion to present our paper clearly!

Reference

[1] Dylan Anstine, Roman Zubatyuk, and Olexandr Isayev. Aimnet2: a neural network potential to meet your neutral, charged, organic, and elemental-organic needs. 2024.

We hope the above information helps address your concerns and questions. We are happy to answer any future questions you may have. Thank you!

We included additional results and summaries in https://openreview.net/forum?id=4dAgG8ma3B&noteId=UrUPt5fBhg.

2. Comparing ChemGuide against DFT

The results against DFT are provided in Table 3 - for every generated 3D geometry by our model and the baseline models, we used DFT/B3LYP/6-31G(2df,p), the same level of theory as the QM9 dataset, to get its optimized geometry (ground state) and compared their energy difference. We showed that ChemGuide-generated 3D geometries are more stable and closer to the ground-state geometries than the baseline models while also being more successful in generating valid molecules.

In other words, we were able to achieve improved performance even in DFT accuracy by leveraging a lower-level GFN2-xTB theory with a reasonable cost.

3. Handle certain molecular configurations

Our choice of using GFN2-xTB as the non-differentiable oracle would help address the potential difficulties in handling certain molecular configurations such as large molecules. It is a semi-empirical density functional-based tight-binding method with reasonable accuracy and computational cost, and unlike neural networks, does not suffer from out-of-distribution generalization problems (see Figure 7 in Appendix H.4). When a decoded 3D geometry from step $t$ is invalid during the generative process, GFN2-xTB will catch such undesired situations. Since the estimated guidance is not informative in such cases and calculating properties for invalid conformations is theoretically not correct, we add no guidance for the invalid molecules in such cases at step $t$ and allow more steps for the model to improve by itself.

4. Automated scale scheduling

We noticed the need for automated scale scheduling. However, since our paper mainly focuses on introducing quantum chemistry to the generative process for molecules, we mentioned this in our future work section on page 10 line 530:

“Additionally, we observe that the guidance scale plays a critical role in the generative process of guided diffusion for molecules, but it is generally difficult to determine in advance. This opens the possibility for research into more effective methods for selecting guidance scales, including automated scale scheduling that optimizes guidance strength.”

We expect that automated scale scheduling will continue to improve ChemGuide from the results shown in Tables 1 and 2. This would require extra designing, tuning, and computational costs and is beyond the scope of this work.

5. Code availability

We provided access to our code in Appendix A: https://anonymous.4open.science/r/ChemGuide in our initial submission, and added the link in the abstract in our current revised version.

References

[1] De Cao, Nicola, and Thomas Kipf. "MolGAN: An implicit generative model for small molecular graphs." arXiv preprint arXiv:1805.11973 (2018).

[2] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.

[3] Xu, Minkai, et al. "Geometric latent diffusion models for 3d molecule generation." International Conference on Machine Learning. PMLR, 2023.

Thank you for taking the time to review and provide feedback! We address the weaknesses and questions as follows.

1. More baselines, guidance methods & datasets

- Additional baselines

Our ChemGuide model aims to generate stable 3D molecules (atoms + coordinates) with geometries close to their ground states. In our model, a generated 3D molecule is represented by its atoms and their coordinates. In comparison, MolGAN [1] was designed to generate valid 2D molecules with correct valency requirements. In MolGAN, a molecule is represented by a graph, with atoms as the vertices and bonds as the edges. For molecule stability, it can be only defined on 3D structures over 2D graphs.

While MolGAN shows high potential in generating valid 2D molecules, generating valid 3D molecules is inherently more challenging. On the other hand, we are also targeting generating stable 3D molecules. To start, a valid 3D molecule must also meet the requirements for stable 2D molecules - correct valencies. In the 3D scenario, a bond is inferred from the generated coordinates by quantum mechanics. The model needs to learn the underlying quantum chemistry rules to generate a valid 3D molecule, such as the correct distance for desired bonds, the correct dihedral angles for correct configurations, etc.

We designed our model to generate 3D molecules that are chemically correct but also thermodynamically stable with close-to-ground-state geometries. Additionally, we tasked our model to perform conditional generation of 3D molecules with desired quantum chemical properties such as HOMO-LUMO gap and polarizability - these properties are also more intricate than the properties used in MolGAN, such as drug likeliness and solubility.

From the above considerations, MolGAN cannot be used as a 3D benchmark, but we are happy to consider other 3D benchmarks if more related models can be suggested. For a more robust evaluation of ChemGuide, we selected EDM [2] and GeoEDM [3] as the baseline models - they are most relevant to our work in generating valid 3D molecules (with conditioned properties) and their results were the recent state of the art. Our stability criteria are a proper superset of the stability definitions in the baselines and thus more robust and strict (see tables 1-3). We also conducted a comprehensive comparison of conditional 3D molecule generation (see tables 4-6).

- Other existing guidance methods

In the global official comment above, we provide more guidance results for stability optimization with 1. neural networks for force prediction (Appendix H3) and 2. evolutionary algorithms (Appendix I). We can observe that ChemGuide performs the best compared to different guidance methods, with mostly decreased energy above ground state percentages, which demonstrates the effectiveness of our method.

We would also like to note that ChemGuide is to achieve guidance in the form of gradients directly from quantum chemistry, there is no equivalent counterpart to strictly compare with, since neural networks are surrogate models for force prediction (trained on the labels produced by quantum chemistry) and evolutionary algorithms incur uncontrollable randomness during the evolution process.

- Additional datasets

In our results, we presented the results with the QM9 dataset and the GEOM dataset. They are two of the most widely recognized large datasets in the community and they are large: QM9 has 134K ground-state molecules and GEOM has 37M conformational molecules. Considering the scale of the dataset and the level of quantum chemical theory, QM9 and GEOM are two of the best benchmark datasets in the field, on which existing ML algorithms often struggle. Considering the computational cost, we focused more in-depth on the current two datasets to provide more comprehensive results on them rather than going in breadth across a multitude of datasets.

If the reviewer has recommendations for additional datasets, we would be happy to look into them.

- Evolutionary algorithm

We explore using evolutionary algorithms as optimization. Details can be found in Appendix I. We follow [1], and the parameters of the algorithm are variant size $k$ and evolution interval $E$. During the diffusion process, the population is preserved, and evolution is performed at fixed intervals. At each evolution step, $k-1$ noise perturbations are added to the population, resulting in $k$ variants (we treat the non-perturbed molecule as a variant). The best variant is then selected as the new population based on evaluations from the non-differentiable oracle (i.e. xTB). Specifically, the oracle calculates the force RMS for each variant, and the variant with the lowest force RMS is selected. We explore several combinations of variant size and evolution interval in the below table. ChemGuide significantly outperforms it in terms of force RMS and energy above the ground state.

Table 3: Evolutionary algorithm results of 500 molecules on QM9

Reference

[1] Huang, Y., Ghatare, A., Liu, Y., Hu, Z., Zhang, Q., Sastry, C.S., Gururani, S., Oore, S. and Yue, Y., Symbolic Music Generation with Non-Differentiable Rule Guided Diffusion. In Forty-first International Conference on Machine Learning, 2024.

We address common weaknesses and questions as follows.

1. More baselines

- AIMNet2

we replaced xTB with the neural network potential AIMNet2 for guiding force optimization. Detailed results and analysis are presented in Appendix H.3. We evaluated its performance by sampling 500 molecules from QM9 and 50 molecules from GEOM. We found two significant drawbacks of using neural networks as guidance:

Performance: Using AIMNet2 for guidance negatively impacted performance. Specifically, for QM9, ChemGuide consistently outperforms AIMNet2 across all metrics. For GEOM, while AIMNet2 achieves similar results to ChemGuide in terms of force, validity, and stability, it struggles with the energy above the ground state. This indicates the optimization process is prone to being trapped in local minima, far from the global minima.

To understand the difference, we analyzed the cosine similarity between the guidance provided by ChemGuide and AIMNet2. The results showed that their gradients are nearly orthogonal, which accounts for the observed performance differences.

Memory Constraints: AIMNet2 outputs forces in the shape [N, 3]. Backpropagating through these forces to compute guidance requires the Hessian matrix, resulting in a memory complexity of [N, 3, N, 3]. This complexity increases by a factor of $3N$ which significantly limits scalability. For GEOM, we could only sample 5 molecules at a time on a 48GiB GPU. In contrast, ChemGuide allowed us to sample 100 molecules per batch on the same device.

We report the performances below. * and bold denote the overall best result and our best result, respectively. Percentage changes between our results and GeoLDM are shown in parentheses.

Table 1: AIMNet neural force guidance on 500 molecules sampled on QM9

Table 2: AIMNet neural force guidance on 50 molecules sampled on GEOM