PaperlayerView exact revision on arXiv

Source-linked AI summary

MDM: Molecular Diffusion Model for 3D Molecule Generation

Lei Huang, Hengtong Zhang, Tingyang Xu, Ka-Chun Wong

arXiv:2209.05710v1cs.LGq-bio.BM

TL;DR

3D molecule generation can be limited by difficulty representing local interatomic relations and by insufficient diversity. MDM adds force-aware equivariant encoders and per-step latent controls, and reports improved benchmark and targeted-generation outcomes over prior methods.

  • Problem

    3D molecule generators have difficulty representing local interatomic constraints, and their generated molecules can lack diversity.

    Source: Introduction · p. 2

  • Method

    MDM augments interatomic linkages, uses separate equivariant encoders for forces of different strengths, and introduces a distributional latent control at each diffusion and reverse step.

    Source: Introduction · p. 2

  • Results

    MDM outperforms EDM on QM9 and GEOM; uniqueness and novelty improve by 4.1% to 31.4%, and targeted generation retains validity and stability.

    Source: Introduction · p. 2

  • Takeaways & Limitations

    The authors conclude that MDM exceeds previous state-of-the-art models and can generate 3D molecules with desired properties.

    Source: Conclusion · p. 7

    The Markov-kernel design takes molecular geometries' roto-translation invariance into account.

    Source: Dual Equivariant Score Neural Networks · p. 3

Abstract

from arXiv · show

Molecule generation, especially generating 3D molecular geometries from scratch (i.e., 3D \textit{de novo} generation), has become a fundamental task in drug designs. Existing diffusion-based 3D molecule generation methods could suffer from unsatisfactory performances, especially when generating large molecules. At the same time, the generated molecules lack enough diversity. This paper proposes a novel diffusion model to address those two challenges. First, interatomic relations are not in molecules' 3D point cloud representations. Thus, it is difficult for existing generative models to capture the potential interatomic forces and abundant local constraints. To tackle this challenge, we propose to augment the potential interatomic forces and further involve dual equivariant encoders to encode interatomic forces of different strengths. Second, existing diffusion-based models essentially shift elements in geometry along the gradient of data density. Such a process lacks enough exploration in the intermediate steps of the Langevin dynamics. To address this issue, we introduce a distributional controlling variable in each diffusion/reverse step to enforce thorough explorations and further improve generation diversity. Extensive experiments on multiple benchmarks demonstrate that the proposed model significantly outperforms existing methods for both unconditional and conditional generation tasks. We also conduct case studies to help understand the physicochemical properties of the generated molecules.

Introduction

MDM addresses difficulty modeling local interatomic relations and limited generation diversity in 3D molecule generation. It combines distinct equivariant encoders with per-step latent controls, and reports improvements over EDM on benchmark and targeted-generation tasks.

Introduction

  • Generating valid, unique, property-restricted molecules manually is costly, motivating de novo generation for drug discovery.The paper describes molecule generation as a crucial task in drug discovery.Source: Introduction · p. 1
  • MDM augments adjacent-atom linkages and uses separate equivariant networks to model interatomic forces of different strengths.Nearby atoms are treated as covalently bonded, while more distant pairs are associated with van der Waals forces.Source: Introduction · p. 2
  • Per-step latent variables condition diffusion and reverse steps on distributional representations that can be sampled to explore diverse 3D geometries.The paper presents this sampling as a way to encourage thorough exploration during generation.Source: Introduction · p. 2
  • MDM outperforms EDM on QM9 and GEOM; uniqueness and novelty improve by 4.1% to 31.4%, and targeted generation retains validity and stability.GEOM molecules average 46 atoms, compared with 18 in QM9.Source: Introduction · p. 2

Related work

Prior work spans 2D molecule generation and 3D geometry generation, while conformation prediction differs in its required inputs and outputs.

Related work

  • 2D generation methods produce SMILES strings or molecular graphs, but ignore 3D structure information that is crucial for molecular properties.The passage names sequential models, VAEs, GANs, and normalizing flows among the approaches.Source: Related work · p. 2
  • 3D methods include autoregressive G-Schnet, one-shot E-NF, and diffusion-based generation; Hoogeboom et al. (2022) use a fully connected adjacency matrix that ignores molecular graph topology.The passage distinguishes these approaches by their generation frameworks and treatment of molecular structure.Source: Related work · p. 2
  • Molecule generation outputs a complete molecule, whereas conformation prediction takes a molecular graph as input and outputs atom coordinates.Source: Related work · p. 2

Preliminaries

The paper represents a molecular geometry through atom features and coordinates, then describes diffusion as forward noising followed by a learned reverse process.

Preliminaries

  • A molecular geometry G=(A,R) comprises atom features A, including atom types and charges, and atom coordinates R.Source: Preliminaries · p. 2

Diffusion Model

  • The forward diffusion process adds Gaussian noise over steps 1 through T, with the variance schedule controlling the mixture.As the step increases, the distribution approaches standard Gaussian noise.Source: Diffusion Process · pp. 2–3
  • The reverse process learns Gaussian transitions to approximate the reverse distribution and generate molecules from standard Gaussian noise.The transition mean is parameterized by a neural network, while the variance is user-defined.Source: Reverse Process · p. 3
  • The reverse process can be viewed as moving from low- to high-density regions under a learned logarithmic gradient, resembling Langevin dynamics.The negative predicted noise is interpreted as the data score.Source: Reverse Process · p. 3

MDM: Molecular Diffusion Model

MDM combines local and global equivariant encoders with variational noise in reverse diffusion, alongside training and sampling procedures designed to preserve molecular-geometry symmetries and support diverse generation.

MDM: Molecular Diffusion Model

  • MDM uses dual equivariant score networks for local and global edges, and a VAE module supplies conditional noise intended to improve generation diversity.The local and global networks model intramolecular and distant interactions, respectively.Source: MDM: Molecular Diffusion Model · p. 3

Dual Equivariant Score Neural Networks

  • Local edges within a 2 Å radius model covalent interactions, while global edges capture distant interactions such as van der Waals forces.The two edge types are processed by separate equivariant encoders.Source: Edge Construction · p. 4
  • The equivariant Markov kernels treat atom features as invariant to coordinate transformations and coordinates as equivariant to them.The design uses an equivariant Markov kernel to preserve an invariant molecular distribution across diffusion steps.Source: The Equivariant Markov Kernels · p. 4, Dual Equivariant Score Neural Networks · p. 4

Enhanced Diversity via Variational Noising

  • Variational noising adds a learned noise variable to the reverse process; the authors report significantly improved performance when sampling it from U(−1, +1).The encoder produces a mean and standard deviation used to obtain the added noise by reparameterization.Source: Enhanced Diversity via Variational… · pp. 4–5

Training

  • Training samples diffused coordinates in pairwise-distance space to maintain equivariance, and zero-center-of-mass noise supports rotational and translational invariance.The authors report better performance when omitting γ from the training objective and using the simplified objective.Source: Training · pp. 5–6

Sampling

  • Sampling starts from a Gaussian-noise molecular state and repeatedly generates less chaotic states through reverse transitions until the final molecule is obtained.The process runs for T steps using the learned reverse Markov kernels.Source: Sampling · p. 6

Experiments

Experiments evaluate MDM on QM9 and GEOM for molecular geometry generation and test its ability to generate molecules with targeted properties. MDM outperforms baseline models on both geometry benchmarks and performs better than EDM on most conditional properties.

Molecular Geometry Generation

  • Evaluation uses QM9 and GEOM, 10,000 generated samples, and four metrics: Validity, Uniqueness, Novelty, and Stability.For GEOM, the experiments build single bonds because pairwise distances make bond construction difficult for its larger, complex molecules; all bond types are reported in the Appendix.Source: Dataset · p. 6, Baselines and Setup · p. 6, Metrics · p. 6

Results and Analysis

  • On GEOM, MDM exceeds EDM by 30.9% in Validity, 30.4% in Uniqueness, 30.4% in Novelty, and 48.5% in Stability.GEOM molecules average 46 atoms, compared with 18 in QM9.Source: Results and Analysis · p. 6
  • On QM9, MDM exceeds EDM by 6.9% in Validity, 4.1% in Uniqueness, and 31.1% in Novelty.Source: Results and Analysis · p. 6
  • MDM uses two equivariant graph networks to distinguish chemical-bond and van der Waals forces, capturing structural patterns at different atomic spacings.The paper relates this design to improved validity and fewer generated ions compared with EDM.Source: Results and Analysis · p. 6
  • MDM improves Uniqueness and Novelty over MDM-NV, which the authors attribute to the controlling variable enabling exploration during intermediate Langevin-dynamics steps.Source: Results and Analysis · p. 6
  • MDM and MDM-NV outperform all baselines on QM9 and GEOM, with larger performance gaps on GEOM.The comparison covers 10,000 generated molecules on each dataset.Source: Results and Analysis · p. 6, Table 1 · p. 7

Conditional Molecular Generation

  • For conditional generation, MDM surpasses Naive, #Atoms, and EDM on almost all properties except µ and Cv.Its MAE for HOMO and the HOMO-LUMO gap is slightly below the QM9 lower-bound baseline; the authors suggest a slight distribution difference between De and Dt may explain this.Source: Results · p. 7, Table 2 · p. 7
  • With 19 atoms and fixed initial noise, generated molecules across interpolated Polarizability values show the expected trend toward less isometric shapes at higher values.The paper uses 19 atoms because that is the most frequent molecular size in QM9.Source: Results · p. 7, Figure 4 · p. 7

Conclusion

The conclusion presents MDM as a 3D molecule generator that combines augmented interatomic linkages, separate equivariant encoders for forces of different strengths, and a controlling variable for diversity.

Conclusion

  • MDM generates 3D molecules from scratch using augmented interatomic linkages, separate equivariant encoders for different force strengths, and a controlling variable to improve diversity.The paper reports that MDM exceeds previous state-of-the-art models and generates molecules with desired properties.Source: Conclusion · p. 7

Notations used in this paper

The paper introduces a table of notations to aid reading.

Notations used in this paper

Proof of the diffusion model

The appendix supplies proofs for diffusion-model properties and outlines the marginal distribution of the diffusion process.

Proof of the diffusion model

  • The appendix provides proofs for derivations of several diffusion-model properties and refers readers to Ho, Jain, and Abbeel (2020) for detailed discussion.Source: Proof of the diffusion model · p. 10

Marginal distribution of the diffusion process

  • The marginal distribution passage recalls q(Gt | G0), invokes Gaussian noise addition, and states that αt = 1 − βt with independent standard Gaussian samples.It says Gt can be obtained using the reparameterization trick.Source: Marginal distribution of the diffusion… · p. 10

The parameterized mean µθ

The section describes parameterizing the reverse-process mean, using pairwise distances to represent atomic coordinates, and deriving an ELBO objective for optimization.

The parameterized mean µθ

  • The learned mean µθ uses the same formulation as ˜µ but parameterizes ϵ.The reverse transition is described as a learned Gaussian approximation to q(Gt−1 | Gt).Source: The parameterized mean µθ · pp. 10–11

Decompose atomic coordinates to pairwise distances

  • Pairwise distances represent atomic coordinates to achieve equivariance, with a function mapping coordinates to interatomic distances for density estimation.The neural network estimates molecular log density from those distances.Source: Decompose atomic coordinates to… · p. 11

The ELBO objective

  • Because the data log likelihood is hard to calculate directly, the paper derives an ELBO objective for optimization.The section then further derives the ELBO objective.Source: The ELBO objective · p. 12

Experiments details

The experiments cover QM9 and GEOM, describe bond construction and six conditional-generation properties, and report MDM's performance in an all-bond GEOM comparison.

Dataset and implementation

  • QM9 contains over 130K small molecules averaging 18 atoms; GEOM molecules average 46 atoms including hydrogens, with 290K lowest-energy conformations used for MDM training.For GEOM, MDM is trained with Adam for 200K iterations.Source: Dataset and implementation · p. 12, Geom · p. 13

Bond prediction

  • Generated bonds are assigned from atomic pairwise distances using chemical-bond lookup-table ranges, with a single-bond threshold interval given as an example.The stated example assigns a single bond when Dd,ij ≤ d_ij ≤ Ds,ij.Source: Bond prediction · p. 13

Selected properties

  • The conditional-generation evaluation selects six molecular properties: polarizability, HOMO, LUMO, HOMO-LUMO gap, dipole moment, and heat capacity.Heat capacity is specified at 298.15K.Source: Selected properties · p. 13

Additional experiment results on Geom

  • On GEOM, MDM surpasses all baseline models in the all-bond setting; Table 4 compares 10,000 generated molecules, with higher values indicating better performance.The passage reports that EDM drops significantly on all metrics in this setting.Source: Table 4 · p. 13, Additional experiment results on Geom · p. 13

Samples generated by MDM

The paper presents generated-sample visualizations for QM9 and GEOM.

Samples generated by MDM

Loading 2209.05710v1…