Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Molecular Graph Generation with Flow Matching and Optimal Transport

Xiaoyang Hou, Tian Zhu, Milong Ren, Dongbo Bu, Xin Gao, Chunming Zhang, Shiwei Sun

TL;DR

GGFlow is proposed, a discrete flow matching generative model incorporating optimal transport for molecular graphs and it incorporates an edge-augmented graph transformer to enable the direct communications among chemical bounds.

Abstract

Generating molecular graphs is crucial in drug design and discovery but remains challenging due to the complex interdependencies between nodes and edges. While diffusion models have demonstrated their potentiality in molecular graph design, they often suffer from unstable training and inefficient sampling. To enhance generation performance and training stability, we propose GGFlow, a discrete flow matching generative model incorporating optimal transport for molecular graphs and it incorporates an edge-augmented graph transformer to enable the direct communications among chemical bounds. Additionally, GGFlow introduces a novel goal-guided generation framework to control the generative trajectory of our model, aiming to design novel molecular structures with the desired properties. GGFlow demonstrates superior performance on both unconditional and conditional molecule generation tasks, outperforming existing baselines and underscoring its effectiveness and potential for wider application.

Improving Molecular Graph Generation with Flow Matching and Optimal Transport

TL;DR

GGFlow is proposed, a discrete flow matching generative model incorporating optimal transport for molecular graphs and it incorporates an edge-augmented graph transformer to enable the direct communications among chemical bounds.

Abstract

Generating molecular graphs is crucial in drug design and discovery but remains challenging due to the complex interdependencies between nodes and edges. While diffusion models have demonstrated their potentiality in molecular graph design, they often suffer from unstable training and inefficient sampling. To enhance generation performance and training stability, we propose GGFlow, a discrete flow matching generative model incorporating optimal transport for molecular graphs and it incorporates an edge-augmented graph transformer to enable the direct communications among chemical bounds. Additionally, GGFlow introduces a novel goal-guided generation framework to control the generative trajectory of our model, aiming to design novel molecular structures with the desired properties. GGFlow demonstrates superior performance on both unconditional and conditional molecule generation tasks, outperforming existing baselines and underscoring its effectiveness and potential for wider application.

Paper Structure

This paper contains 35 sections, 4 theorems, 28 equations, 4 figures, 5 tables, 5 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Flow Matching and Diffusion Models
  4. Molecular Graph Generative Models
  5. Methods
  6. Discrete Flow Matching for Molecular Graph Generation
  7. Optimal transport for graph flow matching
  8. GraphEvo: Edge-augmented Graph Transformer
  9. Permutation Property Analysis
  10. Goal-Guided Framework For Conditional Molecule Generation
  11. Experiment
  12. Moleuclar Graph Generation
  13. Generic Graph Generation
  14. Conditional Molecule Generation
  15. Conclusion
  16. ...and 20 more sections

Key Result

Theorem 1

If the distributions $p(G^0)$ and $p(G^1)$ are permutation invariant and the cost function maintains this invariance, then the optimal transport map $\phi$ also respects this property, i.e., $\phi(G^0, G^1) = \phi(\pi G^0,\pi G^1)$, where $\pi$ is a permutation operator.

Figures (4)

  • Figure 1: Illustration of generative trajectories using different methods. The generative trajectories are learned by the diffusion model (left), flow matching model (center), and flow matching model with optimal transport (right).
  • Figure 2: Visualization of generated samples of different models in different molecular datasets
  • Figure S1: (a-d) Data distribution of the flow matching model, $\pi_0$ is the original distribution (orange), $\pi_1$ is the target data distribution (blue), and the red dots are the data distribution generated by the model. (e-h) In reinforcement learning, the flow matching model conducts exploration/sampling trajectories
  • Figure S2: Visualization of generated samples of our model in different datasets

Theorems & Definitions (8)

  • Theorem 1
  • Proposition 1
  • Proposition 2
  • Theorem 2: Optimal prior distribution
  • proof
  • proof
  • proof
  • proof