Diffuse, Sample, Project: Plug-and-play Controllable Graph Generation
[Paper]    [GitHub]

Kartik Sharma1, Srijan Kumar1 and Rakshit S Trivedi2,
1Georgia Institute of Technology, 2Massachusetts Institute of Technology




TL;DR: We present a novel method to enable controllable graph generation of pre-trained graph diffusion models.



Technical Abstract

Diffusion models lend transformative capabilities to the graph generation task, yet controlling the properties of the generated graphs remains challenging. Recent approaches augment support for controlling soft, differentiable properties but they fail to handle user-specified hard constraints that are non-differentiable. This often results in vague control, unsuitable for applications like drug discovery that demand satisfaction of precise constraints, e.g., the maximum number of bonds. To address this, we formalize the problem of controlled graph generation and introduce PRODIGY (PROjected DIffusion for controlled Graph Generation), an innovative plug-and-play approach enabling the generation of graphs with precise control, from any pre-trained diffusion model. PRODIGY employs a novel operator to project the samples at each diffusion step onto the specified constrained space. For a large class of practical constraints and a variety of graphs, our extensive experiments demonstrate that PRODIGY empowers state-of-the-art continuous and discrete diffusion models to produce graphs meeting specific, hard constraints. Our approach achieves up to 100% constraint satisfaction for non-attributed and molecular graphs, under a variety of constraints, marking a significant step forward in precise, interpretable graph generation.


Model





Results

PRODIGY enables us to generate realistic graphs while following constraints that are only satisfied by 10% of test graphs.


PRODIGY is also able to generate synthetically valid graphs while satisfying arbitrary constraints.


Different molecular constraints can be satisfied across different diffusion models.


Examples of molecular graphs with only Carbon and Oxygen atoms.



Try on your data!

Demo coming soon! Stay tuned.

Paper and Bibtex

Webpage: https://https://prodigy-diffusion.github.io
Code: https://prodigy-diffusion.github.io/code
Paper: https://openreview.net/forum?id=ia0Z8d1DbY
Arxiv: https://arxiv.org/abs/2402.NNNNN


Bibtex:

@inproceedings{sharma2024diffuse,
title={Diffuse, Sample, Project: Plug-And-Play Controllable Graph Generation},
author={Sharma, Kartik and Kumar, Srijan and Trivedi, Rakshit},
booktitle={Forty-first International Conference on Machine Learning},
year={2024},
url={https://openreview.net/forum?id=ia0Z8d1DbY}
}



The template is built on top of the one build by Phillip Isola and Richard Zhang.