[Paper Review] Equivariant Diffusion for Molecule Generation in 3D

Emiel Hoogeboom, Victor García Satorras+2 more|arXiv (Cornell University)|Mar 31, 2022

Machine Learning in Materials Science118 citations

TL;DR

Introduces an E(3) equivariant diffusion model (EDM) that jointly denoises 3D coordinates and atom-type features to generate molecules, with a probabilistic likelihood framework and strong empirical gains over prior 3D molecular generators.

ABSTRACT

This work introduces a diffusion model for molecule generation in 3D that is equivariant to Euclidean transformations. Our E(3) Equivariant Diffusion Model (EDM) learns to denoise a diffusion process with an equivariant network that jointly operates on both continuous (atom coordinates) and categorical features (atom types). In addition, we provide a probabilistic analysis which admits likelihood computation of molecules using our model. Experimentally, the proposed method significantly outperforms previous 3D molecular generative methods regarding the quality of generated samples and efficiency at training time.

Motivation & Objective

• Motivate the need for 3D molecular generation that respects Euclidean symmetries (E(3)).
• Propose an EDM that jointly models continuous coordinates and discrete atom features in a diffusion framework.
• Enable tractable likelihood computation for generated molecules under the proposed model.
• Demonstrate empirical improvements over prior 3D molecular generation approaches in quality and efficiency.

Proposed method

• Define an equivariant denoising diffusion model operating on coordinates x and features h.
• Use an EGNN-based dynamics phi to predict denoised noise in an equivariant way.
• Formulate a likelihood objective with L_t terms and a practical noise parametrization predicting epsilon.
• Handle discrete atom types via one-hot encoding and appropriate likelihood terms for x and h.
• Incorporate a zero-centered gravity constraint for x and scale features to stabilize training.
• Allow conditional generation by incorporating properties c into the denoising network input.

Experimental results

Research questions

• RQ1 Can an E(3) equivariant diffusion model generate 3D molecular conformations and atom types without fixed atom ordering?
• RQ2 Does EDM achieve better log-likelihood and molecule stability compared to prior equivariant methods (e.g., E-NF, G-Schnet) on QM9?
• RQ3 Can the model efficiently train and scale to larger drug-like datasets while providing a tractable likelihood?
• RQ4 How do discrete atom features (types and charges) integrate into a joint diffusion process with continuous coordinates?
• RQ5 What is the impact of conditioning on molecular properties for targeted molecule generation?

Key findings

• EDM achieves a more favorable negative log-likelihood (-110.7 ± 1.5) than baselines (E-NF -59.7, G-Schnet N.A., GDM -94.7, GDM-aug -92.5].
• Atom stability is higher for EDM (98.7 ± 0.1%) than baselines (85.0%, 95.7%, 97.0%, 97.6%).
• Molecule stability is higher for EDM (82.0 ± 0.4%) than baselines (4.9%, 68.1%, 63.2%, 71.6%).
• EDM generates more valid and unique molecules in 3D than competing models (Table 2 results: 97.5 ± 0.2% valid, 94.3 ± 0.2% valid and unique with H included).
• EDM can train more efficiently and scale to larger datasets like GEOM-Drugs, outperforming prior 3D diffusion and normalizing flow approaches.