Molecule Relaxation by Reverse Diffusion with Time Step Prediction

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.