Reverse Diffusion Refines Molecular Structures

By Dr Silpaja Chandrasekar, PhDReviewed by Susha Cheriyedath, M.Sc.Aug 13 2024

In a paper published in the journal Machine Learning: Science and Technology, researchers introduced MoreRed, a novel statistical method for molecular relaxation using reverse diffusion. Unlike traditional force field methods that relied on local energy minimization or neural network models that required extensive labeled datasets, MoreRed learned a simplified pseudo potential energy surface. It was trained on a smaller, unlabeled dataset of equilibrium structures, avoiding the need for non-equilibrium data, and showed improved accuracy in predicting equilibrium states.

Background

Past work in computational chemistry focused on methods like classical force fields, semiempirical models, and machine learning (ML)--based force fields to optimize geometry for understanding reactivity. These methods aimed to find equilibrium structures efficiently but often compromised accuracy. Emerging approaches have used generative models, particularly diffusion models, to generate molecular structures from equilibrium data. However, these models typically generate structures from scratch, limiting their ability to steer generation towards desired outcomes.

Advanced Molecular Relaxation

Diffusion models are latent-variable generative models designed to efficiently generate samples from complex data distributions, such as equilibrium molecular structures, where direct sampling is difficult. MoreRed reinterprets molecular relaxation as a denoising problem, employing a reverse diffusion process to restore non-equilibrium molecular structures to their equilibrium states.

Unlike traditional diffusion models that generate new samples from complete noise, MoreRed initiates the reverse process from an arbitrary noise level to reconstruct the nearest equilibrium structure. By learning a pseudo potential energy surface (PES) that depends on the diffusion time step, MoreRed requires only equilibrium structures for training, enhancing data efficiency compared to machine learning force fields (MLFFs).

 MoreRed introduces a time step predictor to accurately relax non-equilibrium structures and determine the appropriate starting point for the reverse diffusion process. This predictor estimates the noise level of the input structure, ensuring that the reverse process begins at the correct time step and leads to accurate molecular relaxation. The time step prediction is feasible because the latent distributions formed during the diffusion process become progressively smoother, making it easier to estimate the noise level and, thus, the diffusion time step.

MoreRed includes three variants for handling diffusion time step prediction. The first, MoreRed initial time prediction (MoreRed-ITP), predicts only the initial time step for the reverse process. The second, MoreRed adaptive scheduling (MoreRed-AS), predicts the time step before each denoising step, allowing corrections during the reverse process. The third, MoreRed joint training (MoreRed-JT), uses a single neural network with shared backbone representation for noise and time step predictions, optimizing them simultaneously for better performance. These variants offer flexibility and accuracy in molecular relaxation tasks.

Enhanced Relaxation Performance

The time step predictor in MoreRed estimated the deviation of non-equilibrium inputs from equilibrium structures, with its predictions correlating well with root mean square deviation (RMSD) despite the absence of Gaussian noise in the training data. MoreRed-JT's joint model for predicting time step and noise showed better correlation than the separate predictors in MoreRed-ITP/AS, and the adaptive time step schedule in MoreRed-JT proved more effective in achieving structures closer to equilibrium.

The relaxation performance of MoreRed, especially with adaptive scheduling, surpassed baseline methods, including ML force field (MLFF), Merck molecular force field 94 (MMFF94), and generalized free energy-based density functional tight-binding (xTB) (GFN2-xTB). MoreRed consistently achieved lower median RMSD ratios and fewer failures, with MoreRed-AS and MoreRed-JT showing the best results.

Testing on diffused equilibrium structures from QM7-X showed that MoreRed variants outperformed baseline methods, especially in handling less plausible inputs. However, a mismatch between RMSD results and density functional theory (DFT) energies, particularly in MoreRed-AS, was linked to a structure-label mismatch in the training data, corrected by retraining on MLFF-relaxed structures.

MoreRed demonstrated superior performance in molecular relaxation tasks, effectively bringing non-equilibrium structures closer to equilibrium and achieving lower energy levels when trained on appropriate datasets. The adaptive time step predictor proved particularly advantageous, and MoreRed's robustness across different input types was notable. It even outperformed more traditional methods like MMFF94 and semiempirical models like GFN2-xTB. This study underscores the importance of carefully considering the training data and evaluation metrics, particularly when dealing with different computational methods and data sources, to ensure accurate and meaningful comparisons in molecular relaxation tasks.

Conclusion

To sum up, the study introduced MoreRed, a novel, data-efficient molecular relaxation method utilizing reverse diffusion and time step prediction. MoreRed effectively mapped non-equilibrium structures to equilibrium ones without requiring non-equilibrium data. Integrating a time step, predictor allowed for robust relaxation across varying noise levels, with variants showing significant benefits in flexibility and accuracy. Compared to other methods, MoreRed demonstrated superior data efficiency and relaxation performance, though it faced challenges due to mismatched computational methods in the dataset.