PaperlayerView exact revision on arXiv

Source-linked AI summary

GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation

Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, Jian Tang

arXiv:2203.02923v1cs.LGq-bio.QM

TL;DR

Predicting stable molecular conformations from molecular graphs is difficult, while traditional simulation methods are expensive for large molecules. GeoDiff directly reverses a coordinate diffusion process with equivariant Markov kernels to enforce roto-translational invariance, and experiments show consistent state-of-the-art performance, especially on large molecules.

  • Problem

    Predicting stable molecular conformations from molecular graphs matters because molecular 3D structures determine biological and physical properties, while traditional MD and MCMC methods are computationally expensive for large molecules.

  • Method

    GeoDiff models atoms as particles and learns a reverse diffusion Markov chain directly over atomic coordinates, using equivariant transitions to produce roto-translationally invariant likelihoods.

  • Results

    GeoDiff consistently outperforms existing state-of-the-art machine-learning models across datasets and metrics, especially on challenging large molecules.

  • Takeaways & Limitations

    Direct coordinate generation bypasses intermediate geometric variables and supports accurate, diverse conformations without accumulated intermediate errors.

  • Takeaways & Limitations

    The diffusion process itself is not required to satisfy an invariance property, provided it efficiently draws noisy samples for training the generative process.

Abstract

from arXiv · show

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.

1 INTRODUCTION

GeoDiff directly generates molecular conformations by reversing a diffusion process over atomic coordinates while enforcing roto-translational invariance. Across multiple benchmarks, it consistently outperforms existing machine-learning approaches, especially for large molecules.

  • Stable molecular conformations are important because 3D structures determine molecular biological and physical properties, but traditional MD and MCMC methods are expensive for large molecules.
  • Existing approaches often model intermediate geometric variables, motivating direct coordinate modeling that also accounts for roto-translational invariance.
  • GeoDiff treats atoms as particles and learns a reverse Markov diffusion process that transforms noisy atomic positions into realistic conformations.
  • By operating directly on atomic coordinates, GeoDiff supports end-to-end training, avoids accumulated intermediate errors, and models multimodal conformational distributions with quality and diversity.
  • GeoDiff consistently outperforms existing state-of-the-art machine-learning approaches across multiple benchmarks, with especially large margins on challenging large molecules.

2 RELATED WORK

Prior conformation-generation methods commonly use invariant intermediate geometric elements, but errors in those elements can accumulate during coordinate reconstruction. Later score-based approaches model coordinate gradients while still relying on intermediate geometric information.

  • Many deep generative conformation models use atomic distances or torsion angles because these variables are roto-translationally invariant.
  • Distance-based VAE and flow models predict interatomic distances and then reconstruct coordinates with distance geometry.
  • CONFVAE makes this distance-to-coordinate pipeline end-to-end through bilevel optimization.
  • Noise in predicted distances can misguide coordinate searching, causing accumulated errors and inaccurate or erroneous structures.
  • CONFGF learns coordinate log-likelihood gradients but remains aided by intermediate geometric elements in practice.

3 PRELIMINARIES

The paper represents molecules as atom-bond graphs with three-dimensional coordinate conformations and formulates conformation generation as conditional modeling of stable structures. It also requires likelihoods to respect 3D rotation and translation transformations.

  • A molecule is represented as an undirected graph whose vertices are atoms and whose edges encode interatomic bonds and chemical types.
  • Each conformation is represented by an n×3 matrix of Cartesian coordinates, with one three-dimensional vector per atom.
  • The task is to learn pθ(C|G) from conformation samples so generated structures approximate the underlying Boltzmann distribution.
  • Equivariance requires transformed inputs and outputs to correspond under group transformations, and the paper considers the SE(3) group of 3D rotations and translations.
  • The estimated conformation likelihood must remain unaffected by translational and rotational transformations.

4 GEODIFF METHOD

GEODIFF formulates conformation generation as an equivariant diffusion framework that progressively corrupts and then restores molecular coordinates. It combines invariant initialization, equivariant Markov kernels, and equivariant neural networks to generate conformations directly from molecular graphs.

  • 4.1 FORMULATION: The method models forward diffusion as a fixed Markov chain that gradually injects noise into conformations until they approach isotropic Gaussian noise.The variance schedule is fixed rather than trainable, and arbitrary forward states can be sampled in closed form.
  • 4.1 FORMULATION: GEODIFF learns a reverse conditional Markov chain that starts from Gaussian noise and iteratively refines coordinates into conformations conditioned on graph G.Each reverse transition uses pθ(Ct−1|G, Ct), with neural networks estimating transition means.
  • 4.2 EQUIVARIANT REVERSE GENERATIVE PROCESS: The model addresses roto-translational invariance by combining an invariant initial density with SE(3)-equivariant Markov transitions.The paper proves that an invariant initial density and equivariant transitions induce an invariant generated density.
  • 4.2 EQUIVARIANT REVERSE GENERATIVE PROCESS: The initial noise distribution is made invariant by centering sampled coordinates at zero center of mass before evaluating or sampling the isotropic Gaussian density.This CoM-free construction ensures translational invariance while preserving rotational invariance of the isotropic Gaussian.
  • 4.2 EQUIVARIANT REVERSE GENERATIVE PROCESS: Graph field networks produce an SE(3)-equivariant noise vector field by updating coordinates with relative directions weighted by invariant features.The construction uses node embeddings, coordinate embeddings, interatomic distances, and neighborhood messages, including radius-based long-range interactions.

5 EXPERIMENT

GEODIFF is evaluated on conformation generation and property prediction using GEOM-QM9 and GEOM-Drugs, with metrics designed to assess quality and diversity. It consistently outperforms state-of-the-art machine-learning baselines, especially for large molecules, while force-field optimization improves accuracy without sacrificing diversity.

  • Experiment setup: Experiments cover equilibrium conformation generation for small and drug-like molecules on GEOM-QM9 and GEOM-Drugs.Each training split contains 40,000 molecules with five conformations per molecule.
  • Experiment setup: The study compares GEODIFF with five machine-learning baselines and the RDKIT conformation-generation method.The machine-learning baselines are CVGAE, GRAPHDG, CGCF, CONFVAE, and CONFGF.
  • Evaluation: Coverage and Matching metrics use aligned-coordinate RMSD to assess generated-conformation quality and diversity.Higher COV and lower MAT indicate more realistic conformations; thresholds are 0.5Å for QM9 and 1.25Å for Drugs.
  • Conformation generation: GEODIFF consistently outperforms state-of-the-art machine-learning models across datasets and metrics, with the largest margins on the challenging Drugs dataset.GEODIFF-C generally performs slightly better than GEODIFF-A and is used in subsequent comparisons.
  • Conformation generation: GEODIFF +FF preserves superior Recall-based diversity while significantly improving Precision-based accuracy on GEOM-Drugs.The comparison reflects the differing tendencies of machine-learning models and RDKIT under force-field optimization.
  • Property prediction: GEODIFF’s property-prediction performance indicates more accurate conformations across molecules because molecular properties are highly geometry-sensitive.Results are reported as mean absolute errors of predicted ensemble properties in eV.

6 CONCLUSION

The conclusion presents GEODIFF as a probabilistic molecular-conformation model that combines denoising diffusion with geometric representations. Its invariant likelihood objective and equivariant Markov kernels support competitive performance across multiple tasks, while future work targets efficiency and harder structures.

  • Conclusion: GEODIFF combines denoising diffusion models with geometric representations for molecular-conformation generation.The reverse generative dynamics are parameterized as a Markov chain.
  • Conclusion: Equivariant Markov kernels impose roto-translational invariance on the generated-conformation density.The framework derives a tractable invariant objective from a variational lower bound to optimize likelihood.
  • Conclusion: Experiments across multiple tasks show that GEODIFF is competitive with existing state-of-the-art models.
  • Conclusion: Future work includes improving or accelerating diffusion models and extending the method to challenging structures such as proteins.

A PROOFS

The proofs establish tractable marginal and posterior distributions for the forward diffusion process. These properties provide the closed-form quantities needed to derive the model’s tractable training objective and explain convergence toward whitened noise.

  • Forward diffusion: The forward process uses a variance schedule with α_t = 1 − β_t and cumulative product ᾱ_t = ∏_{s=1}^t α_s.These quantities are used in deriving the tractable objective.
  • Forward diffusion: Independent Gaussian noise makes the marginal q(C_t|C_0) Gaussian with mean √ᾱ_t C_0 and variance (1 − ᾱ_t)I.This gives a closed-form way to evaluate C_t given C_0.
  • Forward diffusion: As the number of diffusion steps grows, the cumulative signal factor approaches zero and the process converges to a whitened noisy distribution.
  • Forward diffusion: The appendix identifies a tractable posterior q(C_{t−1}|C_t,C_0) for the forward process.This posterior is another key property used to derive the final tractable objective.

A.2 PROOF OF PROPOSITION 1

The proof establishes that equivariant Markov transitions preserve the desired roto-translational invariance, while the model layers maintain the corresponding geometric properties. It also derives the ELBO weighting and formalizes the center-of-mass-free state space used by the diffusion process.

  • SE(3)-equivariant Markov transitions induce an SE(3)-invariant output density when the initial density is invariant.
  • A GFN layer propagates invariant hidden features and rotationally equivariant, translationally invariant coordinates through distance-based messages and coordinate updates.
  • Composing L GFN layers preserves the same equivariance property by induction.
  • The ELBO objective omits a constant KL term and combines the terminal likelihood term with the remaining KL-divergence terms.
  • Because the forward and reverse conditional distributions share covariance ˜β_tI, their KL divergence reduces to a weighted squared ℓ2 distance between means.
  • Center-of-mass removal is represented by a symmetric projection Q, and equivariant reverse kernels keep trajectories within the resulting center-of-mass-free subspace.

C EXPERIMENT DETAILS

The experimental implementation parameterizes diffusion means with invariant MPNNs and equivariant GFNs. The setup also specifies additional GEODIFF hyperparameters in a dedicated table.

  • Diffusion means ϵθ are parameterized as compositions of invariant MPNNs and the proposed equivariant GFNs.
  • The default MPNNs use 4 layers and hidden embeddings of dimension 128 before producing invariant atom embeddings.
  • Table 4 reports additional hyperparameters for GEODIFF.

D.1 RESULTS FOR GEOM-QM9

This section identifies the GEOM-QM9 results as a reported benchmark evaluation, while also noting an additional Drugs experiment with fewer diffusion steps. The reduced-step model remains competitive with existing baselines.

  • The GEOM-QM9 results are reported in Table 5.
  • Table 5 reports GEOM-QM9 results without force-field optimization.
  • With T = 1000 diffusion steps, GEODIFF is slightly weaker than with 5000 steps but outperforms all existing baselines on the Drugs dataset.

E MORE VISUALIZATIONS

The paper provides visualizations of GEODIFF-generated molecular structures. The displayed drug-like molecules are selected from the GEOM-Drugs test split.

  • The additional generated-structure visualizations use molecules selected from the GEOM-Drugs test split.
  • Figure 3 visualizes drug-like conformations generated by GEODIFF.
  • The visual material concerns generated molecular structures rather than numerical benchmark results.
Loading 2203.02923v1…