Graph Generation with Diffusion Mixture

TL;DR

This work tackles the challenge of generating graphs with accurate topology using diffusion models. It introduces GruM, a diffusion framework that explicitly models graph topology by predicting the final graph through a graph mixture, implemented as a weighted mean of endpoints from endpoint-conditioned OU bridge processes. The training objective derives from a likelihood bound via the Girsanov theorem, enabling simulation-free learning of the graph mixture with an efficient objective that maximizes data likelihood. GruM demonstrates strong performance across general graphs and 2D/3D molecular generation tasks, outperforming previous diffusion-based and graph-generative baselines in topology accuracy and molecule stability. The approach supports both discrete and continuous features, achieves fast convergence of the graph mixture, and offers practical benefits such as early stopping to reduce generation time, with potential impact on drug design and related domains.

Abstract

Generation of graphs is a major challenge for real-world tasks that require understanding the complex nature of their non-Euclidean structures. Although diffusion models have achieved notable success in graph generation recently, they are ill-suited for modeling the topological properties of graphs since learning to denoise the noisy samples does not explicitly learn the graph structures to be generated. To tackle this limitation, we propose a generative framework that models the topology of graphs by explicitly learning the final graph structures of the diffusion process. Specifically, we design the generative process as a mixture of endpoint-conditioned diffusion processes which is driven toward the predicted graph that results in rapid convergence. We further introduce a simple parameterization of the mixture process and develop an objective for learning the final graph structure, which enables maximum likelihood training. Through extensive experimental validation on general graph and 2D/3D molecule generation tasks, we show that our method outperforms previous generative models, generating graphs with correct topology with both continuous (e.g. 3D coordinates) and discrete (e.g. atom types) features. Our code is available at https://github.com/harryjo97/GruM.

Figures (19)

• Figure 1: Illustration of the graph generative process. (a: Denoising diffusion model, b: GruM (ours), c: Graph mixture) For GruM, we design the generative process as a mixture of endpoint-conditioned diffusion processes (Eq. \ref{['eq:ou_bridge']}), namely the OU bridge mixture (Eq. \ref{['eq:mixture']}), which is driven toward the graph mixture (green) by its drift (Eq. \ref{['eq:drift_mixture']}). Our GruM in (b) successfully generates graphs with valid topology by predicting the final result via learning the graph mixture as a weighted mean of data (Eq. \ref{['eq:predicted_graph']}). The predicted graph of GruM converges in an early stage to the correct topology as visualized in (c). In contrast, previous denoising diffusion models in (a) often fail to capture the correct topology as they learn the score or noise for denoising (red), without explicit knowledge of final graph structure.
• Figure 2: (Left) Topology analysis. We compare Spec. MMD and V.U.N of the graph mixture from GruM against the implicit prediction computed from GDSS, ConGress, and DiGress which we provide details in Appendix \ref{['sec:app:exp:general_implementation']}. (Middle) MMD between the test set and the graph mixture of GruM through the generative process. (Right) The complexity of GruM with and without using the inductive bias, measured by the Frobenius norm of the Jacobian of the models.
• Figure 3: (Left) Generation results on the 3D molecule datasets. Best results are highlighted in bold which is the average of 3 different runs. The baseline results are taken from hoogeboom22edm and wu22bridge. (Right) Convergence of the generative process. We compare the convergence of the graph mixture from GruM and the implicit prediction computed from EDM. We measure the convergence (L2 distance) and report the molecule stability of the predictions.
• Figure 4: (Left) Generation results on the Planar dataset. Best results are highlighted in bold, where smaller MMD and larger V.U.N. indicate better results. (Right) Generated graphs by learning the drift. Visualized graphs are randomly sampled without curation.
• Figure 7: The experimental results for the variant of EDM where it aims to predict the final result (EDM-Var.). (Left) Generation results on the 3D molecule QM9 datasets. Best results are highlighted in bold where the higher stability indicates better results. (Right) Convergence of the generative process. We compare the convergence of the graph mixture from GruM, the implicit prediction computed from the estimated noise of EDM, and the predicted result of EDM-Var. We measure the convergence with L2 distance and further visualize the molecule stability of the predictions through the generative process.