[Paper Review] MolDiff: Addressing the Atom-Bond Inconsistency Problem in 3D Molecule Diffusion Generation

Xingang Peng, Jiaqi Guan+2 more|arXiv (Cornell University)|May 11, 2023

Machine Learning in Materials Science11 citations

TL;DR

MolDiff jointly diffuses atoms and bonds in 3D molecules using an E(3)-equivariant framework with bond guidance, addressing atom-bond inconsistencies and achieving higher generation quality than prior models.

ABSTRACT

Deep generative models have recently achieved superior performance in 3D molecule generation. Most of them first generate atoms and then add chemical bonds based on the generated atoms in a post-processing manner. However, there might be no corresponding bond solution for the temporally generated atoms as their locations are generated without considering potential bonds. We define this problem as the atom-bond inconsistency problem and claim it is the main reason for current approaches to generating unrealistic 3D molecules. To overcome this problem, we propose a new diffusion model called MolDiff which can generate atoms and bonds simultaneously while still maintaining their consistency by explicitly modeling the dependence between their relationships. We evaluated the generation ability of our proposed model and the quality of the generated molecules using criteria related to both geometry and chemical properties. The empirical studies showed that our model outperforms previous approaches, achieving a three-fold improvement in success rate and generating molecules with significantly better quality.

Motivation & Objective

• Motivate addressing atom-bond inconsistency in 3D molecule diffusion generation.
• Propose a diffusion model that samples atoms and bonds simultaneously to maintain chemical consistency.
• Design an E(3)-equivariant graph neural network that updates atom and bond representations together.
• Introduce a bond-first noise schedule to stabilize training and generation.
• Develop evaluation metrics capturing geometry, drug-likeness, and structural rationality of generated molecules.

Proposed method

• Model a 3D molecule as M={A,R,B} with atom types A, coordinates R, and bonds B.
• Use a forward diffusion process that perturbs atom positions and discrete atom/bond types with absorbing-type noise.
• Employ a reverse diffusion with an E(3)-equivariant neural network to predict pθ(M^{t-1}|M^{t}) for atoms and bonds.
• Treat atom and bond types as discrete variables and separate bond and atom diffusion with a bond-first schedule.
• Introduce a bond predictor that guides atom position denoising via the gradient of log C, where C is bond-predictor confidence.
• Adopt an E(3)-equivariant message-passing scheme that updates vertex and edge representations to predict atom and bond types (Eq. 4).
• Train with a loss L^{t-1}=L_pos^{t-1}+λ1 L_atom^{t-1}+λ2 L_bond^{t-1} balancing position, atom-type, and bond-type denoising.

Experimental results

Research questions

• RQ1 Does joint diffusion of atoms and bonds improve realism and validity of generated 3D molecules compared to post-processing bond addition?
• RQ2 Can a bond-first diffusion schedule reduce atom-bond inconsistencies and improve chemical topology during generation?
• RQ3 How does an E(3)-equivariant architecture that updates both atoms and bonds affect generation quality?
• RQ4 What metrics best capture geometry, drug-likeness, and rationality of 3D molecules in diffusion-based generation?

Key findings

• MolDiff substantially improves generation performance over EDM, achieving higher validity, connectivity, and a near-perfect success rate.
• Discrete diffusion of atom and bond types outperforms continuous diffusion and post-processing bond assignment across multiple baselines.
• A bond-first diffusion schedule stabilizes training by letting bond types diffuse to a prior before atoms, improving bond consistency with later atom refinements.
• Incorporating a bond predictor-guided gradient during atom-position denoising enhances bond-length accuracy and overall molecular geometry.
• MolDiff yields superior 3D geometry metrics (lower RMSD, better bond-angle/dihedral distributions) and more realistic bond-type and ring statistics compared to EDM.
• MolDiff achieves a three-fold improvement in success rate over the baseline EDM and produces molecules with improved drug-likeness and structural quality.