[Paper Review] MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation

Clément Vignac, Nagham Osman+2 more|arXiv (Cornell University)|Feb 17, 2023

Machine Learning in Materials Science10 citations

TL;DR

MiDi jointly generates molecular graphs and their 3D conformations via end-to-end diffusion, achieving high rates of stable, valid molecules on GEOM-DRUGS compared to prior approaches.

ABSTRACT

This work introduces MiDi, a novel diffusion model for jointly generating molecular graphs and their corresponding 3D arrangement of atoms. Unlike existing methods that rely on predefined rules to determine molecular bonds based on the 3D conformation, MiDi offers an end-to-end differentiable approach that streamlines the molecule generation process. Our experimental results demonstrate the effectiveness of this approach. On the challenging GEOM-DRUGS dataset, MiDi generates 92% of stable molecules, against 6% for the previous EDM model that uses interatomic distances for bond prediction, and 40% using EDM followed by an algorithm that directly optimize bond orders for validity. Our code is available at github.com/cvignac/MiDi.

Motivation & Objective

• Motivate end-to-end generation of molecules that jointly capture graph structure and 3D conformation.
• Propose a diffusion-based model that denoises both graph and coordinates simultaneously.
• Enhance 3D-equivariant processing with a novel rEGNN-enhanced Transformer architecture.
• Demonstrate superior unconditional generation performance on GEOM-DRUGS and QM9 datasets.

Proposed method

• Represent molecules as graphs embedded in 3D with node features (atom types, charges) and edge features (bond types).
• Use a mixed Gaussian (coordinates) and discrete (atom types, charges, bonds) diffusion process with component-specific noise schedules.
• Introduce an adaptive noise schedule to prioritize coordinates and bond types before atom types and charges.
• Develop a Relaxed Equivariant Graph Neural Network (rEGNN) layer that extends EGNN with non-translation-invariant features.
• Incorporate the rEGNN into a Transformer-based denoising network that maintains SE(3) equivariance.
• Train with a loss combining coordinate regression and cross-entropy on discrete features.

Figure 1 : Samples from our model. MiDi generates simultaneously a 2D graph structure and a 3D conformation that is consistent with this structure.

Experimental results

Research questions

• RQ1 Can a diffusion model jointly learn molecular graphs and 3D coordinates in an end-to-end differentiable manner?
• RQ2 Does adaptive component-wise noise scheduling improve the quality and stability of generated molecules?
• RQ3 Does a relaxed SE(3)-equivariant architecture (rEGNN) enhance graph+conformer generation over standard EGNNs?
• RQ4 How does MiDi perform on unconditional molecule generation compared to OpenBabel-based post-processing and prior diffusion-based methods on QM9 and GEOM-DRUGS?

Key findings

• MiDi achieves high rates of stable molecules on GEOM-DRUGS, outperforming EDM-based approaches and Open Babel post-processing. (Adaptive MiDi: Mol stable 91.6%, At stable 99.8%, Validity 77.8%, Connected 100.0%).
• On QM9, MiDi outperforms the base EDM model on graph-based metrics, though Open Babel optimization remains strong for simple datasets.
• MiDi delivers end-to-end differentiable generation of both 2D graph structures and 3D conformers, avoiding a separate bond-prediction step.
• The adaptive noise schedule improves sample quality, guiding the model to prioritize coordinates and bond types before atom types and charges.
• MiDi scales to large drug-like datasets (GEOM-DRUGS) and demonstrates stronger end-to-end performance than methods relying on post-hoc bond prediction.

Figure 2 : The noise schedule is tuned separately for each component. Atom coordinates and bond types are denoised earlier during sampling, while atom types and formal charges are updated later in the process. Experimentally, the adaptive schedule allows to obtain better 3D conformers and more stabl