Molecule Relaxation by Reverse Diffusion with Time Step Prediction

Khaled Kahouli, Stefaan Simon Pierre Hessmann, Klaus Robert Muller, Shinichi Nakajima, Niklas Wolf Andreas Gebauer

Primary Area: generative models

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: diffusion time prediction, diffusion models, generative modeling, quantum chemistry, molecule relaxation

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: In this work, we introduce MoreRed, a diffusion model with time step prediction for molecule relaxation, which can denoise unstable molecules while offering potential enhancements for unconditional data generation.

Abstract: Molecule relaxation---finding the stable state of an unstable configuration---is an important subtask for exploring the chemical compound space, for instance, to identify novel drugs or catalysts. Existing methods rely on local energy minimization with the gradients (i.e., force field) estimated through computationally intensive ab initio methods or approximated by a neural network trained on large expensive datasets encompassing \emph{labeled stable and unstable} molecules. In this work, we propose molecule relaxation by reverse diffusion (MoreRed), a novel purely statistical approach where unstable molecules are seen as \emph{noisy} samples to be denoised by a diffusion model equipped with a time step predictor to handle arbitrarily noisy inputs. Notably, MoreRed learns a simpler pseudo energy surface instead of the complex physical energy surface and is trained on a significantly smaller dataset consisting of solely \emph{unlabeled stable} molecules, which is considerably less expensive to generate. Nevertheless, our experiments demonstrate its competitive performance to the state-of-the-art baseline in terms of the quality of the relaxed molecules inferred. Furthermore, we identify the high potential that time step prediction has to enhance the performance of data generation, where our findings are promising both in molecular structure and image generation.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 7978