[Paper Review] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation

Minkai Xu, Lantao Yu+4 more|arXiv (Cornell University)|Mar 6, 2022

Machine Learning in Materials Science173 citations

TL;DR

GeoDiff directly models atomic coordinates with an SE(3)-equivariant diffusion framework to generate molecular conformations, achieving state-of-the-art results on conformation generation, especially for large molecules.

ABSTRACT

Predicting molecular conformations from molecular graphs is a fundamental problem in cheminformatics and drug discovery. Recently, significant progress has been achieved with machine learning approaches, especially with deep generative models. Inspired by the diffusion process in classical non-equilibrium thermodynamics where heated particles will diffuse from original states to a noise distribution, in this paper, we propose a novel generative model named GeoDiff for molecular conformation prediction. GeoDiff treats each atom as a particle and learns to directly reverse the diffusion process (i.e., transforming from a noise distribution to stable conformations) as a Markov chain. Modeling such a generation process is however very challenging as the likelihood of conformations should be roto-translational invariant. We theoretically show that Markov chains evolving with equivariant Markov kernels can induce an invariant distribution by design, and further propose building blocks for the Markov kernels to preserve the desirable equivariance property. The whole framework can be efficiently trained in an end-to-end fashion by optimizing a weighted variational lower bound to the (conditional) likelihood. Experiments on multiple benchmarks show that GeoDiff is superior or comparable to existing state-of-the-art approaches, especially on large molecules.

Motivation & Objective

• Motivate and develop a diffusion-based generative model that directly outputs 3D atomic coordinates for molecules.
• Ensure roto-translation invariance via SE(3) equivariant design and a CoM-free invariant initialization.
• Provide end-to-end trainable framework with a tailored ELBO objective for conditional generation given molecular graphs.

Proposed method

• Formulate forward diffusion of 3D coordinates C^t with fixed variance schedule to C^T ~ N(0, I).
• Define a conditional reverse Markov chain p_θ(C^{t-1}|G, C^t) to progressively denoise while conditioning on the molecular graph G.
• Parameterize the reverse mean μ_θ using an SE(3) equivariant neural network (Graph Field Network) ensuring roto-translation equivariance.
• Use a CoM-free invariant initial density and CoM-free Markov kernels to guarantee translational invariance; enforce rotation equivariance through μ_θ and the network architecture.
• Train via a weighted ELBO that reduces to a weighted L2 loss between ε and ε_θ, with two strategies to enforce equivariance (alignment-based and chain-rule approaches).
• Provide a sampling algorithm that starts from C^T ~ N(0, I) and iteratively samples C^{t-1} using p_θ(C^{t-1}|G, C^t).

Experimental results

Research questions

• RQ1 Can a diffusion model operating directly on 3D coordinates achieve competitive or superior conformer generation performance compared to models relying on invariant geometric intermediates?
• RQ2 How can roto-translation invariance be guaranteed in both training and inference for geometry generation?
• RQ3 What are effective parameterizations and architectures (e.g., SE(3)-equivariant networks) to model the reverse diffusion step for molecular conformations?
• RQ4 Does end-to-end coordinate-based diffusion generation scale to larger, drug-like molecules and provide diverse, high-quality conformations?

Key findings

• GeoDiff consistently outperforms state-of-the-art ML baselines on GEOM-Drugs, with notable gains in COV-R and MAT-R metrics.
• GeoDiff achieves higher COV-R and lower MAT-R than baselines for large molecules, demonstrating strong coverage and accuracy.
• GeoDiff-C generally yields slightly better optimization than GeoDiff-A, suggesting chain-rule based training better guides learning.
• The method yields substantially improved diversity and quality of generated conformations, particularly on the Drugs dataset where large molecules are tested.
• Qualitative samples show GeoDiff captures both local and global structural features effectively.