Controllable Graph Generation with Diffusion Models via Inference-Time Tree Search Guidance
Year: 2025
Type: preprint
Abstract: Graph generation is a fundamental problem in graph learning with broad applications across Web-scale systems, knowledge graphs, and scientific domains such as drug and material discovery. Recent approaches leverage diffusion models for step-by-step generation, yet unconditional diffusion offers little control over desired properties, often leading to unstable quality and difficulty in incorporating new objectives. Inference-time guidance methods mitigate these issues by adjusting the sampling process without retraining, but they remain inherently local, heuristic, and limited in controllability. To overcome these limitations, we propose TreeDiff, a Monte Carlo Tree Search (MCTS) guided dual-space diffusion framework for controllable graph generation. TreeDiff is a plug-and-play inference-time method that expands the search space while keeping computation tractable. Specifically, TreeDiff introduces three key designs to make it practical and scalable: (1) a macro-step expansion strategy that groups multiple denoising updates into a single transition, reducing tree depth and enabling long-horizon exploration; (2) a dual-space denoising mechanism that couples efficient latent-space denoising with lightweight discrete correction in graph space, ensuring both scalability and structural fidelity; and (3) a dual-space verifier that predicts long-term rewards from partially denoised graphs, enabling early value estimation and removing the need for full rollouts. Extensive experiments on 2D and 3D molecular generation benchmarks, under both unconditional and conditional settings, demonstrate that TreeDiff achieves state-of-the-art performance. Notably, TreeDiff exhibits favorable inference-time scaling: it continues to improve with additional computation, while existing inference-time methods plateau early under limited resources.
Source: arXiv (Cornell University)
Authors: Zhao, Jiachi, Wang, Zehong, Liao, Yamei, Zhang, Chuxu, Ye, Yanfang
Institutions: University of Notre Dame, University of Connecticut
Language: English
Subfield: Computational Theory and Mathematics
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: III: Small: A New Machine Learning Paradigm Towards Effective yet Efficient Foundation Graph Learning Models, SCC-IRG Track 1: A New AI-driven Paradigm to Promote Community Resilience for Teenagers and Young Adults in Preventing Opioid Misuse and Addiction, Collaborative Research: III: Small: Data-driven and AI-powered Innovations for Tackling the Nexus of Poverty-driven Food Insecurity and the Opioid Crisis, III: Medium: A Data-driven and AI-augmented Framework for Collaborative Decision Making to Combat Infectious Disease Outbreaks, D-ISN: An AI-augmented Framework to Detect, Disrupt, and Dismantle Opioid Trafficking Networks
OpenAlex ID: W7091186262
Cites:
Cited by:
Related to: