Exploring Molecule Generation Using Latent Space Graph Diffusion

TL;DR

This work investigates latent-space diffusion for molecular graph generation, contrasting Gaussian diffusion in latent space, heat-diffusion, and flow matching while comparing GNN and EGNN backbones. Molecules are encoded as atom point clouds in a latent space, then decoded into labeled molecular graphs with an edge-type predictor, enabling de novo generation via iterative restoration from noise. Experiments on QM9 reveal clear trade-offs: 2D latent spaces yield high validity with reasonable diversity, while higher dimensions increase uniqueness at the cost of validity and computation; heat-diffusion offers high validity with lower diversity, and flow matching provides a balanced performance. The results highlight that representation choices (latent vs input space) and diffusion strategy critically influence generation quality and computational efficiency, informing future design of diffusion-based molecular generators.

Abstract

Generating molecular graphs is a challenging task due to their discrete nature and the competitive objectives involved. Diffusion models have emerged as SOTA approaches in data generation across various modalities. For molecular graphs, graph neural networks (GNNs) as a diffusion backbone have achieved impressive results. Latent space diffusion, where diffusion occurs in a low-dimensional space via an autoencoder, has demonstrated computational efficiency. However, the literature on latent space diffusion for molecular graphs is scarce, and no commonly accepted best practices exist. In this work, we explore different approaches and hyperparameters, contrasting generative flow models (denoising diffusion, flow matching, heat dissipation) and architectures (GNNs and E(3)-equivariant GNNs). Our experiments reveal a high sensitivity to the choice of approach and design decisions. Code is made available at github.com/Prashanth-Pombala/Molecule-Generation-using-Latent-Space-Graph-Diffusion.

Figures (5)

• Figure 1: Schematic overview of our training process. (1) An autoencoder is trained to embed atoms into a low-dimensional latent space. Molecule reconstruction involves two steps: predicting a feature vector for each atom, and predicting molecular bonds for each pair of atoms. (E)GNN means that either a vanilla GNN or an EGNN is used. (2) Information in the latent space is removed using a degradation operator (e.g., adding noise), and an (E)GNN is trained to restore the original data. (3) For de novo molecule generation, we start with random noise (or the endpoint of the degradation process) and iteratively add information using the (E)GNN until we obtain a final point cloud, which is decoded into a molecule.
• Figure 2: This scatter plot shows the relationship between the number of trainable parameters and the percentage of valid, unique, and novel molecules generated. The graph highlights that different model configurations show varying performance and model performance is not always dependent on the number of parameters.
• Figure 3: This scatter plot shows the relationship between validity and uniqueness across different model configurations.
• Figure 4: This scatter plot illustrates the relationship between autoencoder training time and diffusion model training time across different model configurations.
• Figure 5: Comparison of molecular structures from the QM9 dataset (top) with generated molecules (bottom). The top half displays 50 randomly selected molecules from the QM9 dataset, while the bottom half showcases 50 randomly selected valid, unique, and novel molecules generated using a GNN-based diffusion model.