PaperlayerView exact revision on arXiv

Source-linked AI summary

MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation

Clement Vignac, Nagham Osman, Laura Toni, Pascal Frossard

arXiv:2302.09048v2cs.LG

TL;DR

Existing molecule generators often separate graph generation from 3D conformer generation, leaving connectivity to a later prediction step. MiDi jointly denoises molecular graphs and coordinates, and on GEOM-DRUGS it generates stable molecules at a higher rate than EDM with either distance-based or Open Babel bond prediction.

  • Problem

    Existing molecule generators typically produce either molecular graphs or 3D conformers, with conformer models relying on a later bond-prediction step that prevents end-to-end differentiability.

    Source: §1 · p. 2

  • Method

    MiDi jointly generates a molecular graph and its 3D coordinates by denoising both components, using a component-specific noise schedule and a Transformer with rEGNN layers.

    Source: §1 · p. 2

  • Results

    91.6% of MiDi-generated molecules are stable on GEOM-DRUGS, versus 5.5% for EDM and 40.3% for EDM with Open Babel bond prediction.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors report that MiDi also applies to drug-discovery tasks beyond unconditional generation, indicating versatility and potential for drug-discovery pipelines.

    Source: §1 · p. 3

    For practical applications, explicit hydrogen atoms can be removed because they can be inferred from heavy atoms; methods trained only on heavy atoms usually perform better because they use smaller graphs.

    Source: §5.1 · p. 11

Abstract

from arXiv · show

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.

1 Introduction

MiDi jointly generates molecular graphs and 3D conformations, addressing the separation between graph generation and conformer generation in prior models. On GEOM-DRUGS, it produces stable molecules at a higher rate than EDM with either distance-based or Open Babel bond prediction.

1 Introduction

  • MiDi jointly generates molecular graphs and corresponding 3D coordinates, denoising both together to produce stable graphs consistent with generated conformers.It represents atom types, formal charges, and bond types alongside coordinates.Source: §1 · p. 2
  • MiDi combines Gaussian noise for 3D coordinates with discrete diffusion for other components, adding noise to atom types and charges faster than to coordinates and bonds.This schedule encourages initial focus on conformation and bond types before refining atom types and formal charges.Source: §1 · p. 2
  • The proposed Transformer incorporates rEGNN layers that use non-translation-invariant features yet remain translation- and rotation-equivariant with Gaussian noise in the zero center-of-mass subspace.The authors identify this equivariance as crucial for high performance.Source: §1 · p. 2
  • 91.6% of MiDi-generated molecules are stable on GEOM-DRUGS, versus 5.5% for EDM and 40.3% for EDM with Open Babel bond prediction.The passage presents this comparison as evidence for MiDi's end-to-end differentiable approach.Source: §1 · p. 2
  • MiDi is also applied to drug-discovery tasks beyond unconditional generation, which the authors describe as evidence of its versatility and potential for drug-discovery pipelines.The passage does not specify those additional tasks.Source: §1 · p. 3

2 Related Work

Prior molecule-generation work either generates 3D conformers without learning connectivity or generates graphs without associated 3D coordinates. Related approaches combine these modalities in different ways, while protein-generation methods exploit chain-specific structure.

Concurrent work

  • Concurrent 2D+3D diffusion models [17] and [36] use Gaussian coordinate diffusion and discrete graph diffusion; MiDi uniquely improves on standard EGNN layers, handles formal charges, and reports explicit-hydrogen results on GEOM-DRUGS.[17] proposes richer positional encodings, while [36] introduces guidance for predicting bond lengths.Source: Concurrent work · p. 3

Molecule generation in 3D

  • 3D molecule models such as EDM generate conformers but not connectivity, which must be inferred afterward from distances or specialized software, limiting complex-molecule performance and end-to-end differentiability.These models have also been conditioned on properties or another point cloud, including for protein binding and linker generation.Source: Molecule generation in 3D · p. 3

Graph Generation

  • Graph-then-conformer pipelines exploit known-graph symmetries, but combining the stages breaks end-to-end differentiability and makes unconditional generation more difficult.Fragment-based methods instead combine chemically relevant substructures, but are harder to adapt to 3D.Source: Graph Generation · pp. 3–4

Protein Generation

  • Protein diffusion methods scale to larger structures by exploiting chain adjacency and predicting successive Cα angles, but those improvements are specific to chain graphs.The passage says these strategies cannot be used for arbitrary molecules.Source: Protein Generation · p. 4

3 Background

Diffusion models corrupt data through a Markovian noise process and train a denoising network to recover clean inputs. For molecular 3D data, equivariance requirements govern the noise, prior, network, and loss.

Denoising Diffusion Models

  • A diffusion noise model creates a Markovian trajectory of increasingly corrupted data, and the denoising network learns to predict the clean input rather than a noisy intermediate target.Sampling reverses the diffusion sequence using network predictions.Source: Denoising Diffusion Models · p. 4
  • Gaussian noise suits continuous data and discrete state-space diffusion suits categorical data; both frameworks make key diffusion training and sampling quantities tractable.Discrete diffusion represents class transitions with categorical distributions and transition matrices.Source: Denoising Diffusion Models · p. 5

SE(3)-Equivariance with Diffusion Models

  • SE(3)-equivariant diffusion requires equivariant noise, an invariant inference prior, an equivariant denoising network, and a compatible loss; EGNNs use rotation-invariant messages and relative-coordinate updates.The normalization term in the EGNN coordinate update supports numerical stability when concatenating many layers.Source: Denoising Diffusion Models · p. 5, SE(3)-Equivariance with Diffusion Models · p. 6

4 Proposed Model

MiDi represents molecules as graphs embedded in 3D and jointly denoises their coordinates and graph features. Its component-specific noise schedule and equivariant Transformer network are designed to generate molecular structure across these modalities.

4 Proposed Model

  • MiDi represents each molecule with atom types, formal charges, 3D coordinates, and bond types, and uses dense adjacency tensors with absence of a bond as a bond type.The graph is G = (x, c, R, Y).Source: §4 · p. 6

4.1 Noise Model

  • Positions receive Gaussian noise in the zero center-of-mass subspace, while atom types, formal charges, and bond types use discrete diffusion with categorical transitions.The position-noise subspace has dimension 3(n−1).Source: §4.1 · p. 6

4.2 Adaptive Noise Schedule

  • The adaptive schedule denoises atom coordinates and bond types earlier than atom types and formal charges, and is reported to improve conformers and molecular stability.The schedule parameters vary by component; for QM9, νr = 2.5, νy = 1.5, and νx = νc = 1, while for GEOM-DRUGS, νr = 2.Source: Figure 2 · p. 7, §4.2 · pp. 7–8

4.3 Denoising Network

  • The denoising network uses a Transformer to predict the clean graph from a noisy one, processing graph-level features, coordinates, atom features, and edge features.Coordinates are handled separately from other node features to maintain SE(3) equivariance.Source: §4.3 · p. 8

Relaxed Equivariant Graph Neural Networks (rEGNNs)

  • The rEGNN extends EGNN processing beyond pairwise distances with rotation-invariant descriptors relative to the molecule’s center of mass; these features facilitate generating more connected molecules.The additional descriptors include ||ri||2, ||rj||2, and cos(ri, rj).Source: Relaxed Equivariant Graph Neural… · pp. 8–9

Update Block

  • The update block updates edge features from coordinate differences and node and global features, then updates node, global, and coordinate features using the other representations.Pairwise features are pooled into node representations with PNA layers.Source: Update Block · pp. 9–10

Integration into a Transformer Architecture

  • The feed-forward network applies MLPs to nodes and edges in parallel, while the coordinate update projects outputs onto the zero center-of-mass subspace.The text also describes equivariant normalization choices for the Transformer architecture.Source: Integration into a Transformer… · p. 10

4.4 Training Objective

  • Training predicts coordinates with mean-squared error and categorical features with cross-entropy, combining the loss components in a weighted sum.The final weights are λr = 3, λx = 0.4, λc = 1, and λy = 2.Source: §4.4 · pp. 10–11

5 Experiments

Experiments compare unconditional generation on QM9 and GEOM-DRUGS, evaluating both molecular graphs and conformers. MiDi’s reported comparisons differ across these datasets.

5.1 Settings

  • Evaluation combines graph metrics and property-distribution comparisons with bond-length and bond-angle metrics to assess both graph structure and conformer quality.Bond-length distances are weighted by bond type; bond-angle comparisons are restricted to atoms with at least two neighbors.Source: §5.1 · pp. 11–12
  • QM9 contains molecules with up to nine heavy atoms and is split into 100,000 training, 20,000 validation, and 13,000 test molecules; benchmarks include explicit hydrogens for comparison.The authors note that hydrogens can be inferred from heavy atoms in most practical applications, and heavy-atom-only methods usually perform better.Source: §5.1 · p. 11, §5.2 · p. 12
  • On QM9, MiDi outperforms EDM on 2D metrics and obtains similar 3D metrics; Open Babel performs well on this simple dataset.The authors note that QM9 molecules have up to nine heavy atoms and their bonds can be determined easily from conformations.Source: §5.2 · pp. 12–13, Table 2 · p. 13

5.3 GEOM-DRUGS

  • GEOM-DRUGS contains 430,000 drug-sized molecules averaging 44 atoms, with up to 181 atoms; the adaptive MiDi model was fine-tuned from a uniform-schedule checkpoint.The split is 80% training, 10% validation, and 10% testing; each molecule contributes its five lowest-energy conformations.Source: §5.3 · p. 13
  • On GEOM-DRUGS, MiDi significantly outperforms Open Babel on most metrics, while EDM produces few valid and stable molecules and MiDi predicts more realistic bond angles.The table reports 95% confidence intervals from five samplings of the same checkpoint.Source: §5.3 · pp. 13–14, Table 3 · p. 14

6 Conclusions

MiDi jointly generates molecular graphs and conformations using a diffusion model that combines Gaussian and discrete noise. On unconditional generation with complex datasets, it outperforms prior 3D methods that predict bonds from conformations using predefined rules.

6 Conclusions

  • MiDi jointly generates a molecular graph and its corresponding 3D conformation with a denoising diffusion model.The model represents molecules as graphs embedded in 3D and trains a neural network to denoise both components.Source: §6 · p. 14
  • MiDi combines Gaussian and discrete diffusion, adapts the noise schedule to different components, and uses a graph Transformer with rEGNN layers.The schedule first forms rough estimates of conformation and graph structure, then tunes atom types and charges; centered inputs yield a translation- and rotation-equivariant network.Source: §6 · p. 14
  • On unconditional generation with complex datasets, MiDi outperforms prior 3D molecule-generation methods that infer bonds from conformations using predefined rules.The authors state that the model is currently evaluated on unconditional generation and suggest pocket-conditioned generation as a potential downstream application.Source: §6 · p. 14

A.1 QM9 with Implicit Hydrogens

On QM9 with implicit hydrogens, all methods perform well overall, but EDM’s lookup table sometimes predicts incorrect edge types, producing invalid molecules. The adaptive noise schedule does not help on this simpler dataset, and MiDi’s bond-length predictions are less precise than EDM’s.

A.1 QM9 with Implicit Hydrogens

  • The adaptive noise schedule that improved GEOM-DRUGS is not effective on QM9, where the uniform schedule performs better.The passage does not establish why this difference occurs.Source: §A.1 · p. 19
  • Bond-length predictions are good for all methods, but MiDi is less precise than EDM.The passage attributes EDM’s precision to its use of learning-rate decay and an exponential moving average.Source: §A.1 · p. 19
  • EDM’s lookup table sometimes predicts incorrect edge types, resulting in invalid molecules.Table 4 reports results from five samplings of the same checkpoint.Source: §A.1 · p. 19, Table 4 · p. 20

B Samples from our model

The figures show non-curated molecular samples across QM9 and GEOM-drugs, using implicit or explicit hydrogens.

B Samples from our model

  • Figure 4 shows non-curated QM9 samples with implicit hydrogens.Source: Figure 4 · p. 21
  • Figure 5 shows non-curated QM9 samples with explicit hydrogens.Source: Figure 5 · p. 21
  • Figure 6 shows non-curated GEOM-drugs samples with implicit hydrogens.Source: Figure 6 · p. 22
  • Figure 7 shows non-curated QEOM-drugs samples with explicit hydrogens.Source: Figure 7 · p. 22

Related papers

Loading 2302.09048v2…