Molecular relaxation by reverse diffusion with time step prediction

TL;DR

MoreRed reframes molecular relaxation as a denoising task using reverse diffusion, learning a simple pseudo PES from unlabeled equilibrium structures to map non-equilibrium inputs back to equilibrium. A diffusion time-step predictor enables adaptive denoising from arbitrary noise levels, and three variants (ITP, AS, JT) balance speed and robustness. On QM7-X, MoreRed achieves near-equilibrium geometries with low RMSD and energies within chemical accuracy, demonstrating strong data efficiency compared to classical force fields and neural-network force fields. The approach is robust to input from different sources and holds promise for scalable relaxation in data-limited chemical spaces, with open-source code and data for replication.

Abstract

Molecular relaxation, finding the equilibrium state of a non-equilibrium structure, is an essential component of computational chemistry to understand reactivity. Classical force field (FF) methods often rely on insufficient local energy minimization, while neural network FF models require large labeled datasets encompassing both equilibrium and non-equilibrium structures. As a remedy, we propose MoreRed, molecular relaxation by reverse diffusion, a conceptually novel and purely statistical approach where non-equilibrium structures are treated as noisy instances of their corresponding equilibrium states. To enable the denoising of arbitrarily noisy inputs via a generative diffusion model, we further introduce a novel diffusion time step predictor. Notably, MoreRed learns a simpler pseudo potential energy surface (PES) instead of the complex physical PES. It is trained on a significantly smaller, and thus computationally cheaper, dataset consisting of solely unlabeled equilibrium structures, avoiding the computation of non-equilibrium structures altogether. We compare MoreRed to classical FFs, equivariant neural network FFs trained on a large dataset of equilibrium and non-equilibrium data, as well as a semi-empirical tight-binding model. To assess this quantitatively, we evaluate the root-mean-square deviation between the found equilibrium structures and the reference equilibrium structures as well as their energies.

Figures (19)

• Figure 1: Schematic depictions of the physical PES (left) and the pseudo PES (right), which emerge from removing Gaussian noise from distorted structures. Learning the physical energy requires equilibrium (circles) and many non-equilibrium (crosses) training structures, while learning the pseudo PES only requires equilibrium training structures.
• Figure 2: The diffusion model applied to the molecular structure of imidazole. It incorporates two stochastic processes: a forward process and a reverse process. The forward process involves a fixed diffusion kernel $q(\mathbf{x}_t|\mathbf{x}_{t-1})$ to transform the original sample $\mathbf{x}_0 \sim q_\mathrm{data}(\mathbf{x}_0)$ into complete noise sample, $\mathbf{x}_T \sim q_T(\mathbf{x}_T)$, usually isotropic Gaussian noise $\mathcal{N}(\boldsymbol{0}, \mathbf{I})$. The backward process is a learned model with parameters $\theta$, which reverses the forward process, i.e. $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}) \approx q(\mathbf{x}_{t-1}|\mathbf{x}_{t})$. It maps a noise sample $\mathbf{x}_T$ from the tractable prior distribution $p_T(\mathbf{x}_T)=q_T(\mathbf{x}_T)$ to the complex target distribution of equilibrium structures, $q_\mathrm{data}(\mathbf{x}_0)$.
• Figure 3: a: Density plots of the RMSD of 10 000 non-equilibrium test structures from QM7-X and their equilibrium structures vs. the initial diffusion time step, $\hat{t}$, predicted by MoreRed-JT (top) and MoreRed-ITP/-AS (bottom). A brighter color (yellow) indicates a higher density of scatter points, a darker color (purple) indicates a lower density. For both, we show the Pearson correlation coefficient, $\rho$. The three crosses (red, orange, green) mark the structures for which a time step trajectory is shown in plot b in matching colors. b: Comparison of time step trajectories for relaxation of three non-equilibrium test structures from QM7-X in red, orange, and green (matching the crosses in plot a). MoreRed-JT follows an adaptive schedule (the time step and noise are predicted at each denoising step; solid lines) and MoreRed-ITP follows a fixed schedule (the time step is only predicted at the initial denoising step; dashed lines). The dots show the predicted (but not utilized) time steps for MoreRed-ITP during denoising with the fixed schedule. All two methods start from a comparable initial time step for each non-equilibrium structure. The inset box shows the RMSDs between the reference equilibrium geometry and the non-equilibrium structure after and before relaxation, respectively.
• Figure 4: The RMSD ratios, i.e. the RMSD after relaxation divided by the RMSD of the non-equilibrium inputs, of structures relaxed with the baselines and the MoreRed variants, trained on the DFTB3+MBD equilibrium structures reported in QM7-X, for 20 000 non-equilibrium structures from the QM7-X test split (left) and 6500 non-equilibrium structures obtained by adding Gaussian noise over 250 forward diffusion steps to equilibrium structures from the QM7-X test split (right). The median values and the percentage of failure cases, i.e. the cases where the RMSD ratio $\geq 1$, are shown to the left of each box plot in bold. In Figure \ref{['fig:confusion_matrix']}, we provide examples of relaxed structures comparing both MoreRed and the MLFF model.
• Figure 5: A pairwise comparison of molecular structures from QM7-X that were relaxed with both MoreRed-AS (MR-AS) and MLFF. In each block, from left to right, the equilibrium structure is shown (eq.), followed by the non-equilibrium structure (non-eq.) and the structures relaxed by MoreRed-AS and MLFF, all labelled with their respective RMSD to the equilibrium structure in angstrom (Å). The first two structures follow the CPK colouring, i.e. white for hydrogen, gray for carbon, blue for nitrogen, and red for oxygen. In the subsequent two depictions, the equilibrium structure (CPK) is superimposed with MoreRed-AS and MLFF relaxed structure. Structures with an decreased RMSD value relative to the non-equilibrium structure are shown in green, while failure cases with a higher RMSD compared to the initial structure are highlighted in salmon