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

arXiv: https://arxiv.org/abs/2305.12347 · HTML: https://mlsanity.com/paper/2305.12347/

## Key facts

- Cited by 27 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 64 of its 82 references that are in the map.
- First posted on arXiv 21 May 2023.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 64)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models](https://mlsanity.com/paper/2211.01095/) (2022)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation](https://mlsanity.com/paper/2107.03502/) (2021)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [MolGAN: An implicit generative model for small molecular graphs](https://mlsanity.com/paper/1805.11973/) (2018)
- [GraphVAE: Towards Generation of Small Graphs Using Variational Autoencoders](https://mlsanity.com/paper/1802.03480/) (2018)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation](https://mlsanity.com/paper/1806.02473/) (2018)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Do Transformers Really Perform Bad for Graph Representation?](https://mlsanity.com/paper/2106.05234/) (2021)
- [Learning from Protein Structure with Geometric Vector Perceptrons](https://mlsanity.com/paper/2009.01411/) (2020)

## Built on by (most cited first, 27 of 27)

- [Fisher Flow Matching for Generative Modeling over Discrete Data](https://mlsanity.com/paper/2405.14664/) (2024)
- [Mixed Continuous and Categorical Flow Matching for 3D De Novo Molecule Generation](https://mlsanity.com/paper/2404.19739/) (2024)
- [A Survey of Generative AI for de novo Drug Design: New Frontiers in Molecule and Protein Generation](https://mlsanity.com/paper/2402.08703/) (2024)
- [Diffusion Models in textitDe Novo Drug Design](https://mlsanity.com/paper/2406.08511/) (2024)
- [FlowMol3: Flow Matching for 3D De Novo Small-Molecule Generation](https://mlsanity.com/paper/2508.12629/) (2025)
- [Exploring Discrete Flow Matching for 3D De Novo Molecule Generation](https://mlsanity.com/paper/2411.16644/) (2024)
- [Geometry Informed Tokenization of Molecules for Language Model Generation](https://mlsanity.com/paper/2408.10120/) (2024)
- [Unified Guidance for Geometry-Conditioned Molecular Generation](https://mlsanity.com/paper/2501.02526/) (2025)
- [Controlled Generation with Equivariant Variational Flow Matching](https://mlsanity.com/paper/2506.18340/) (2025)
- [Straight-Line Diffusion Model for Efficient 3D Molecular Generation](https://mlsanity.com/paper/2503.02918/) (2025)
- [Rethinking Diffusion Models with Symmetries through Canonicalization with Applications to Molecular Graph Generation](https://mlsanity.com/paper/2602.15022/) (2026)
- [TFG-Flow: Training-free Guidance in Multimodal Generative Flow](https://mlsanity.com/paper/2501.14216/) (2025)
- [SynCoGen: Synthesizable 3D Molecule Generation via Joint Reaction and Coordinate Modeling](https://mlsanity.com/paper/2507.11818/) (2025)
- [Diffusion Twigs with Loop Guidance for Conditional Graph Generation](https://mlsanity.com/paper/2410.24012/) (2024)
- [Clifford Group Equivariant Diffusion Models for 3D Molecular Generation](https://mlsanity.com/paper/2504.15773/) (2025)
- [Diffusion-based Graph Generative Methods](https://mlsanity.com/paper/2401.15617/) (2024)
- [Pharmacophore-Conditioned Diffusion Model for Ligand-Based De Novo Drug Design](https://mlsanity.com/paper/2505.10545/) (2025)
- [Predictive Feature Caching for Training-free Acceleration of Molecular Geometry Generation](https://mlsanity.com/paper/2510.04646/) (2025)
- [Toward Better Geometric Representations for Molecule Generative Models](https://mlsanity.com/paper/2605.07693/) (2026)
- [MAST: Motif-Augmented Diffusion with Search Tree for Spectroscopic Molecular Structure Elucidation](https://mlsanity.com/paper/2610.12067/) (2026)
- [Learning Disentangled Equivariant Representation for Explicitly Controllable 3D Molecule Generation](https://mlsanity.com/paper/2412.15086/) (2024)
- [Energy Loss Functions for Physical Systems](https://mlsanity.com/paper/2511.02087/) (2025)
- [VEDA: 3D Molecular Generation via Variance-Exploding Diffusion with Annealing](https://mlsanity.com/paper/2511.09568/) (2025)
- [Generating Physically Consistent Molecules with Energy-Based Models](https://mlsanity.com/paper/2605.18381/) (2026)
- [Smoothing Dark Areas in Molecular Latent Diffusion](https://mlsanity.com/paper/2606.13955/) (2026)
- [S1-Omni: A Unified Multimodal Reasoning Model for Scientific Understanding, Prediction, and Generation](https://mlsanity.com/paper/2607.15686/) (2026)
- [TreeRef-BFN: Equivariance-Free De Novo Molecule Generation based on 2D Topology and Internal 3D Geometry](https://mlsanity.com/paper/2609.32502/) (2026)

## Abstract (arXiv)

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.