Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation

TL;DR

This paper introduces JODO, a joint 2D and 3D diffusion model for complete molecule generation, improving the quality and versatility of generated molecules by modeling both molecular graphs and geometries simultaneously.

Contribution

The paper proposes a novel Diffusion Graph Transformer that jointly models 2D and 3D molecular data, enabling more accurate and versatile molecule generation.

Findings

JODO outperforms baselines on QM9 and GEOM-Drugs datasets.

The model achieves fast few-step sampling.

It effectively supports inverse molecular design.

Abstract

Designing new molecules is essential for drug discovery and material science. Recently, deep generative models that aim to model molecule distribution have made promising progress in narrowing down the chemical research space and generating high-fidelity molecules. However, current generative models only focus on modeling either 2D bonding graphs or 3D geometries, which are two complementary descriptors for molecules. The lack of ability to jointly model both limits the improvement of generation quality and further downstream applications. In this paper, we propose a new joint 2D and 3D diffusion model (JODO) that generates complete molecules with atom types, formal charges, bond information, and 3D coordinates. To capture the correlation between molecular graphs and geometries in the diffusion process, we develop a Diffusion Graph Transformer to parameterize the data prediction model that recovers the original data from noisy data. The Diffusion Graph Transformer interacts node and edge representations based on our relational attention mechanism, while simultaneously propagating and updating scalar features and geometric vectors. Our model can also be extended for inverse molecular design targeting single or multiple quantum properties. In our comprehensive evaluation pipeline for unconditional joint generation, the results of the experiment show that JODO remarkably outperforms the baselines on the QM9 and GEOM-Drugs datasets. Furthermore, our model excels in few-step fast sampling, as well as in inverse molecule design and molecular graph generation. Our code is provided in https://github.com/GRAPH-0/JODO.