PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariant Diffusion for Molecule Generation in 3D

Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, Max Welling

arXiv:2203.17003v2cs.LGq-bio.QMstat.ML

TL;DR

Molecule generation in 3D must account for Euclidean symmetries, while prior non-autoregressive approaches have faced scaling or training-cost limitations. EDM learns equivariant diffusion over coordinates and atom types, with likelihood computation, and outperforms previous molecular generation models in log-likelihood and stability.

  • Problem

    3D molecule generation requires models that leverage Euclidean symmetries, while prior approaches have faced sampling or training scalability limitations.

  • Method

    EDM defines a diffusion process over atom positions and features and learns its denoising process with an E(3)-equivariant neural network, while enabling likelihood computation.

  • Results

    EDM outperforms previous molecule generation models in log-likelihood and molecule stability.

  • Takeaways & Limitations

    EDM scales beyond very small molecules, generates valid conformations while explicitly modeling hydrogens, and is evaluated on the larger GEOM-DRUGS dataset.

  • Takeaways & Limitations

    A Gaussian model for categorical atom-type features would introduce an undesired bias.

Abstract

from arXiv · show

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.

1. Introduction

Molecular generation benefits from modeling 3D Euclidean symmetries, but existing approaches face ordering or training-scalability limitations. EDM addresses these issues with equivariant diffusion over coordinates and atom types, supporting likelihood computation and improved generation results.

  • Motivation: 3D molecular data obeys translation, rotation, and reflection symmetries represented by the Euclidean group E(3).Leveraging these symmetries is described as important for generalization.
  • Limitations of prior methods: Autoregressive models impose an atom order and are difficult to scale during sampling, while normalizing flows are expensive to train because they integrate differential equations.These limitations motivate a diffusion-based alternative.
  • EDM: EDM denoises a diffusion process jointly over continuous atom coordinates and categorical atom types without requiring a particular atom ordering.The model is introduced as an E(3)-equivariant diffusion model for direct 3D molecule generation.
  • Results: On QM9, EDMs generate up to 16 times more stable molecules than E-NFs while requiring half the training time.The reported scaling behavior also permits training on the larger GEOM-Drugs dataset.
  • Contributions: EDM provides a probabilistic analysis that enables molecule likelihood computation and supports both continuous and categorical features.The paper reports improvements over previous molecule-generation models in log-likelihood and molecule stability.

2. Background

Diffusion models learn generation by reversing a noise-adding process, while equivariant modeling preserves the relevant transformations of molecular coordinates and features. EDM combines these ideas through equivariant transition distributions parameterized by an EGNN.

  • 2.1. Diffusion Models: A diffusion process progressively adds noise to data, and the generative model learns the reverse denoising process.The process can be expressed through Markov transition distributions and a variational lower bound on log-likelihood.
  • 2.1. Diffusion Models: In a diffusion parameterization z_t = α_t x + σ_t ε, α_t controls retained signal and σ_t controls added noise.The schedule transitions from approximately clean data at t = 0 toward noise at t = T.
  • 2.1. Diffusion Models: The denoising model replaces the unknown data point with a neural-network estimate and defines reverse transitions from that estimate.The network may predict Gaussian noise rather than the clean data because optimization is easier in that parameterization.
  • 2.2. Equivariance: Equivariance requires transforming the input to produce the corresponding transformation of the output, including rotations, reflections, and translations in E(3).For molecular point clouds, coordinates transform while associated features remain invariant.
  • 2.2. Equivariance: EDM uses an equivariant neural network to parameterize equivariant reverse transitions, so an invariant initial distribution yields an invariant generated marginal.The model uses EGNNs, which operate on fully connected graphs whose nodes carry coordinates and features.

3. EDM: E(3) Equivariant Diffusion Model

EDM defines an E(3)-equivariant diffusion process over molecular coordinates and features, trains an equivariant denoiser, and derives likelihood terms for mixed continuous and categorical data. Sampling iteratively reverses the noise process from Gaussian variables.

  • 3. EDM: E(3) Equivariant Diffusion Model: EDM adds noise to atom positions and features, learns the generative denoising process with an equivariant neural network, and derives log-likelihood equations.The diffusion model treats coordinates and atom features together while respecting E(3) structure.
  • Optimization: Training samples a timestep and Gaussian noise, constructs z_t = α_t[x, h] + σ_t ε, and minimizes the squared error between ε and the network prediction.Coordinate noise is centered so it lies in the zero-center-of-gravity subspace.
  • Denoising network: The denoiser predicts noise from z_t and t, then converts that prediction into estimates of the clean coordinates and features.An EGNN receives the noised coordinates and appends t/T to node features; coordinate outputs are projected to zero center of gravity.
  • Equivariance: The coordinate denoising distribution is equivariant because an equivariant noise prediction rotates with the input and isotropic noise preserves the transition structure.This establishes the required rotational and reflection equivariance of the denoising process.
  • Sampling: Sampling starts from z_T ∼ N(0, I) and iteratively draws z_{t−1} from p(z_{t−1}|z_t) before producing x and h from p(x, h|z_0).The reverse process follows the learned denoising transitions from high noise to molecular variables.
  • Optimization objective: Training with w(t) = 1 outperformed training with the variational weighting w(t) = 1 − SNR(t−1)/SNR(t).This held even when evaluating the probabilistic variational objective.
  • Categorical features: For categorical atom types, EDM uses one-hot representations because integer representations are unnatural and introduce bias.The categorical likelihood models probabilities from the noised one-hot representation rather than treating categories as ordered integers.

4. Related Work

Prior 3D molecule-generation methods include autoregressive models, grid-based VAEs, and equivariant normalizing flows, while graph-based methods omit conformer information. EDM instead denoises a diffusion process and extends coordinate-prediction diffusion methods to discrete atom features.

  • Autoregressive 3D generators sample atoms iteratively from order-dependent distributions, while graph-based generators do not provide conformer information.
  • Grid-based VAE methods represent atoms on fixed grids and use 3D convolutions for molecule generation.
  • E-NF uses an equivariant normalizing flow that integrates a differential equation, whereas diffusion denoising is described as scaling better during training.
  • Coordinate-prediction methods model molecular conformations, and recent score-based and diffusion approaches are effective when neural networks respect data symmetries.
  • EDM extends coordinate-prediction diffusion methods by incorporating discrete atom features and deriving equations for log-likelihood computation.

5. Experiments

Experiments evaluate EDM for unconditional 3D molecule generation, conditional property control, and larger-molecule generation, using stability, likelihood, validity, uniqueness, property-prediction, and energy-distribution metrics.

  • Molecule Generation — QM9: QM9 experiments train EDM to generate 3D coordinates, atom types, and integer-valued charges for molecules with up to nine heavy atoms.The dataset contains 130k molecules and uses 100K/18K/13K train/validation/test partitions.
  • Molecule Generation — QM9: EDM outperforms E-NF, G-Schnet, and its non-equivariant counterpart on all Table 1 metrics, including negative log-likelihood and stability.The lower negative log-likelihood indicates sharper peaks in the learned model distribution.
  • Molecule Generation — QM9: EDM generates high validity and uniqueness rates despite deriving bonds from 3D distances, and explicit hydrogen modeling causes little performance deterioration.The comparison includes GraphVAE, GraphTransformerVAE, and Set2GraphVAE, although results marked (*) are not directly comparable because those models do not use 3D coordinates to derive bonds.
  • Conditional Molecule Generation: Conditional EDM outperforms the Naive and #Atoms baselines for every conditioned property except µ, showing property incorporation beyond atom count for most properties.The remaining gap to the QM9 lower-bound indicates room for improvement.
  • Conditional Molecule Generation: Interpolating polarizability α with fixed noise produces less isometrically shaped molecules at larger α values, consistently across different runs.Polarizability measures a molecule’s tendency to acquire an electric dipole moment under an external electric field.
  • GEOM-Drugs: On GEOM-DRUGS, EDM outperforms non-equivariant counterparts on all reported metrics and captures the training-set energy distribution well.The dataset contains molecules with up to 181 atoms and 44.4 atoms on average; evaluation includes atom stability and Wasserstein distance between energy histograms.

6. Conclusions

EDM scales 3D molecule generation beyond very small molecules, explicitly models hydrogen atoms, and is evaluated on the larger GEOM-DRUGS dataset.

  • EDM scales better than previous non-autoregressive models, which mostly focused on molecules with up to 9 atoms.
  • EDM generates valid conformations while explicitly modeling hydrogen atoms.
  • Evaluation on GEOM-DRUGS extends the model to larger molecules and supports future drug-size 3D molecule generation research.

A. The zero center of gravity, normal distribution

The coordinate distribution is defined on the zero-center-of-gravity subspace, while likelihood and KL computations can be carried out consistently in the ambient space under shared-variance conditions.

  • The coordinate variable x is a point cloud whose center of gravity is zero.
  • A normal distribution is placed on the zero-center-of-gravity subspace, with its mean µ in the same subspace.
  • Samples can be drawn in the ambient space and projected by subtracting their center of gravity, exploiting isotropic variance.
  • For equal variances, KL divergence between distributions on the subspace can be computed consistently in the ambient space.
  • When denoising and posterior variances differ, the divergence depends on subspace dimensionality rather than ambient-space dimensionality.
  • The combined coordinate-and-feature KL divergence decomposes into independent terms and uses squared Euclidean distance after concatenation.

B. Additional Details for the Method

The method uses a numerically controlled diffusion schedule, an unbiased likelihood estimator, and an EGNN-based equivariant denoising process whose marginals remain invariant under orthogonal transformations.

  • Noise schedule: The noise schedule uses αt values that decrease from approximately 1 to 0, with a monotonically increasing negative log SNR curve γ(t).
  • Noise schedule: Clipping αt|t−1 from below by 0.001 bounds its reciprocal during sampling and avoids numerical instability.
  • Log-likelihood estimator: The log-likelihood estimator uses an additional forward pass for L0 because its large contribution would otherwise cause high estimator variance.
  • Log-likelihood estimator: The resulting likelihood estimate is unbiased, combining a sampled Lt term with L0 and the base KL term.
  • The dynamics: EDM dynamics are learned by an EGNN composed of equivariant graph convolutional layers with edge, inference, node, and coordinate operations.
  • Equivariant processes: If reverse transitions are equivariant and the terminal distribution is invariant, induction gives invariant marginals through p(z0).

C. Additional Details on Experiments

Experiments compare equivariant and non-equivariant graph models using distributional, structural, validity, novelty, and likelihood measures, while documenting important metric limitations.

  • Baselines and setup: The ablation baseline GDM uses a non-equivariant graph network that concatenates coordinates with other node features.
  • Baselines and setup: QM9 training takes approximately 7 days on one NVIDIA GeForce GTX 1080Ti GPU, while GEOM-DRUGS training takes approximately 5.5 days on three NVIDIA RTX A6000 GPUs.
  • Distributional evaluation: EDM captures the GEOM-DRUGS estimated-energy distribution well, whereas other methods produce too many low-energy compounds.
  • Metric limitations: RDKit-based validity can be artificially increased by reducing the number of predicted bonds, whereas hydrogen-aware stability metrics are harder to exploit.
  • Novelty: On QM9, novelty starts near 100% and decreases during training as the algorithm progressively captures the exhaustively enumerated data distribution.
  • Metric limitations: Log-likelihoods are unbounded for the paper’s effectively discrete conformational data and should therefore be interpreted alongside stability metrics.

D. Samples from our models

The EDM generates realistic 3D molecules, including large structures, but can produce disconnected components and long rings. Feature scaling affects stability, with atom types apparently learned after coordinates when appropriately scaled.

  • Samples: The QM9 and GEOM-DRUGS samples are uncurated random outputs from the trained EDM models.Their viewing angles can sometimes make structures difficult to inspect.
  • Samples: The GEOM-DRUGS model generates impressive large 3D structures but sometimes produces disconnected components.Disconnected components were observed only in QM9 models during early training stages.
  • Samples: GEOM-DRUGS samples exhibit two main failure cases: disconnected molecules and molecules containing long rings.The model has no regularization specifically preventing these phenomena.
  • Sampling process: Sampling chains show iterative denoising from random normal noise at t = T = 1000 toward t = 0, producing final coordinates and atom types.Atom types are visualized using the argmax of their categorical component.
  • Feature scaling: Feature-scaling ablations show a substantial effect on atom and molecule stability.The results suggest denoising is easier when atom types are decided after coordinates become relatively well defined.

E. Conditional generation

Conditional EDM extends the diffusion model to generate molecules conditioned on desired properties while leaving noise addition unchanged. The conditional setting achieves molecule stabilities similar to the non-conditional model across reported properties.

  • Conditional method: Conditional EDM adds a desired property c as input to the denoising neural network φ.The diffusion process that adds noise is unchanged.
  • Conditional method: The generative process samples molecule size and property jointly, then generates coordinates and atom types conditioned on both.The joint distribution is estimated on training data using a parametrized two-dimensional categorical distribution.
  • Implementation: The conditional experiment uses a 9-layer EGNN with 192 hidden features per layer, SiLU activations, Adam, learning rate 10^-4, and batch size 64.Only atom types and positions are modeled, not atom charges.
  • Results: 80.4% α, 81.73% ∆ε, 82.81% εHOMO, 83.6% εLUMO, 83.3% µ, and 81.03% Cv are the reported molecule stabilities over 10K generated samples.These conditional stabilities were similar to those obtained without conditioning.

QM9 Properties

Conditional EDM samples vary across molecular properties while holding the number of atoms fixed and, within each sweep, reusing the same reparametrization noise. Higher polarizability corresponds to less isotropic atom distributions across reported runs.

  • Properties: α denotes polarizability, the tendency of a molecule to acquire an electric dipole moment in an external electric field.The other reported properties include HOMO energy, LUMO energy, their energy gap, and heat capacity.
  • Conditional generation: The experiment sweeps 9 α values from 73.6 to 101.6 while keeping reparametrization noise fixed and using M = 19 nodes.Figure 9 shows 10 randomly selected sweeps with different reparametrization noises.
  • Conditional generation: For larger polarizability values, atoms are distributed less isotropically.The passage connects this pattern with encouraging larger dipole moments under an applied electric field.
Loading 2203.17003v2…