Controllable Graph Generation with Diffusion Models via Inference-Time Tree Search Guidance

AI summary

Overview

Research area: Graph generation with diffusion models, controllable generation, inference-time guidance, and Monte Carlo Tree Search (MCTS).

Technical level: Advanced. The paper assumes familiarity with diffusion models, denoising score matching, latent-space versus discrete graph representations, and the four-phase MCTS loop (selection, expansion, simulation, backpropagation).

Scope: The paper proposes TreeDiff, an MCTS-guided dual-space diffusion framework that improves controllable graph generation at inference time without retraining, and evaluates it on 2D and 3D molecular generation benchmarks under conditional and unconditional settings.

What This Paper Is About

Diffusion models can generate graphs by iteratively denoising from noise, but they offer weak control over desired properties and can produce unstable, invalid outputs, especially since small structural changes in a graph can drastically change its meaning. Existing inference-time guidance methods steer sampling at test time without retraining, but they make greedy, step-level decisions with no lookahead, so they cannot plan over long denoising trajectories and fail to improve when given more compute. TreeDiff's goal is to replace these local heuristics with structured, tree-search-based planning that yields more controllable, stable, and scalable graph generation.

Key Contributions

1. MCTS as an inference-time controller for graph diffusion. TreeDiff reframes the denoising trajectory as a sequential decision process and applies Monte Carlo Tree Search over it, enabling lookahead planning and integration of multiple objectives instead of greedy step-wise correction. It is described as a plug-and-play inference-time method that requires no retraining.

2. Macro-step expansion. Rather than expanding one denoising step per tree node (which would create hundreds or thousands of near-duplicate layers), TreeDiff groups multiple denoising updates into a single transition, expanding from $\mathcal{G}t$ directly to $\mathcal{G}{t-k}$. The horizon $k$ is sampled adaptively as $k \sim \mathcal{N}(k_{\text{base}}, (\sigma_k k_{\text{base}})^2)$ with $k_{\text{base}} = T_{\text{rem}}/D_{\text{rem}}$ and $\sigma_k = 0.1$, subject to $1 \leq k \leq T_{\text{rem}}$.

3. Dual-space denoising. Each expansion alternates $n$ continuous latent denoising steps with $m \ll n$ lightweight categorical refinement steps in graph space. The reprojected latent is treated as a directional guidance vector $g_{t-n}$ (the gradient of an implicit potential) that softly reshapes subsequent reverse steps, so MCTS selects and backs up values in the latent domain while trajectories remain anchored to structurally valid graphs.

4. Dual-space verifier plus trajectory distillation. A learned verifier $V_\phi$ predicts terminal rewards $\hat{R}_t \approx R(\mathcal{G}_0)$ directly from partially denoised states, removing the need for expensive full rollouts. The encoder, decoder, and discrete denoiser are trained by distilling 500 full denoising trajectories produced by a pretrained diffusion model, rather than being trained jointly with the diffusion process.

Main Findings

• Best average rank on conditional 2D molecule generation: On Synth. & BACE, Synth. & BBBP, and Synth. & HIV, TreeDiff achieves average rank (A.R.) 1.9, the best among all compared methods. For example, on Synth. & HIV it reaches MAE 0.302 and Accuracy 97.90%, and on Synth. & BBBP it reaches MAE 0.342 and Accuracy 94.81%.

• Avoids the diversity-similarity trade-off: TreeDiff reports both high diversity and high similarity rather than trading one for the other. On Synth. & BACE it obtains Div. 89.30 (versus 82.38 for Graph DiT) and Sim. 87.82 (versus 87.52 for Graph DiT).

• Best conditional 3D QM9 results in both ID and OOD regimes: TreeDiff attains A.R. 1.4 across six SE(3)-invariant quantum properties under both in-distribution and out-of-distribution settings. Under in-distribution conditions it reports MAE of 14.92 (polarizability α, Bohr³), 95.21 (gap Δε, meV), 61.55 (HOMO ε_H, meV), 60.19 (LUMO ε_L, meV), 6.07 (dipole μ, D), and 13.41 (heat capacity C_v). Under out-of-distribution shift the errors rise for all methods, but TreeDiff reports the lowest values: 28.55, 341.76, 99.53, 170.11, 14.24, and 23.51 on the same six targets.

• Strong unconditional 2D generation: On QM9, TreeDiff reports validity 99.89, FCD 0.091, NSPDK 0.0004, and scaffold similarity 0.932; on ZINC250k it reports validity 98.95, FCD 2.135, NSPDK 0.0011, and scaffold similarity 0.492, with A.R. 1.8 overall.

• Strong unconditional 3D generation: On QM9, TreeDiff reports validity & uniqueness 94.8, atom-level stability 98.8, and molecule-level stability 88.6; on Drugs it reports atom-level stability 86.5, with A.R. 1.8.

• Favorable inference-time scaling: Across four generation benchmarks, TreeDiff continues to improve as additional inference computation is allocated, whereas standard diffusion, Best-of-N, and a state-of-the-art inference-time guidance method saturate quickly (and can even degrade) under increased compute.

• Consistent improvement over inference-time baselines: The paper reports that inference-time guidance consistently improves diffusion baselines, and that TreeDiff delivers the most stable and substantial gains, with baseline inference-time methods sharing GraphDiT as the backbone for the 2D conditional tasks and LDM-3DG as the backbone for the 3D conditional tasks at matched compute budgets.

• Not reported in the available content: The paper does not report wall-clock runtime, compute-cost tables, hyperparameter sensitivity, or ablation results in the truncated text. The text states that implementation details appear in Appendix C, which is not included. Dataset sizes for QM9, ZINC250k, BACE, BBBP, HIV, and Drugs are also not stated in the provided content.

Methodology in Plain English

The researchers treat the diffusion model's denoising procedure as a sequence of decisions rather than a fixed script. Instead of taking one small denoising step at a time and accepting whatever comes out, they build a search tree over possible denoising trajectories and use Monte Carlo Tree Search to decide which branches to explore.

Three practical problems had to be solved for this to work. First, a denoising trajectory is very long, so making one tree node per step is hopeless. They compress the trajectory by having each node represent several denoising steps at once (a "macro-step"), with the size of that jump chosen adaptively based on how many steps and how much tree depth remain. Second, it is unclear what a tree node should actually contain: working purely in a continuous latent embedding is fast but can decode into invalid graphs, while working directly on graph structure is reliable but slow. They therefore alternate between the two — a few latent denoising steps, a decode to graph space, a few cheap discrete refinement steps, then an encode back to latent space, where the difference between the refined and original latent acts as a gentle directional nudge for the remaining steps. Third, evaluating how good a partial trajectory will turn out requires completing the whole denoising run and computing properties, which is far too expensive inside a search loop. They instead train a small verifier that looks at both the latent state and the decoded graph and predicts the final reward, so the search can estimate value immediately and prune poor branches early.

To train these auxiliary pieces, they first run a pretrained diffusion model to produce 500 complete denoising trajectories and use them as teacher demonstrations: an encoder–decoder pair is trained as a time-conditioned variational autoencoder to reconstruct intermediate graphs along those trajectories, and a graph-space denoiser is trained on consecutive teacher states to predict the next denoised structure. The verifier is trained as a regression model on the same trajectories, labeled with the terminal reward, and is made more robust by perturbing latent states before decoding.

Evaluation covers conditional 2D molecular generation on BACE, BBBP, and HIV (additionally conditioned on synthetic accessibility and synthetic complexity, using 6:2:2 train/validation/test splits and 10,000 generated molecules per task, averaged over five runs), conditional 3D generation on QM9 with six quantum properties under in-distribution and out-of-distribution targets, and unconditional generation on 2D QM9 and ZINC250k plus 3D QM9 and Drugs.

Why This Matters

Impact on research. The paper argues that the standard frame for inference-time guidance in graph diffusion — greedy, step-level, local correction — has a structural ceiling, which shows up empirically as an inability to benefit from more inference compute. By demonstrating that a search-based controller can produce monotone gains with additional computation on four benchmarks, it suggests a different design principle for controllable generation: treat sampling as planning, and invest compute in lookahead rather than in more candidate samples.

Real-world applications (all mentioned in the paper's motivation):

• Drug and molecular discovery: generating candidate molecules that satisfy specific chemical or physical properties, such as high QED or low SAS, under synthetic accessibility and synthetic complexity constraints.
• Materials discovery: proposing new structured candidates that meet physical property targets.
• Knowledge graphs and knowledge graph completion: generating and manipulating linked entity structure under domain constraints.
• Web-scale systems: recommender networks, link prediction, and simulation of evolving online interaction graphs.

Industry relevance. Because TreeDiff is presented as a plug-and-play inference-time method requiring no retraining, an organization with an already-trained graph diffusion model could in principle add it as a sampling-time controller. The verifier's reward is defined as a weighted sum of arbitrary objective functions, so practitioners can combine multiple property targets at inference time without changing the underlying generative model.

Future Directions

• Extending from molecules to Web-scale graphs. The motivation section discusses knowledge graphs, recommender networks, and online interaction graphs, but all reported experiments are on molecular benchmarks; whether the approach scales to large, heterogeneously-typed Web graphs is an open question.
• Reducing dependence on pre-collected teacher trajectories. The auxiliary encoder, decoder, denoiser, and verifier are all distilled from 500 pretrained diffusion trajectories. How sensitive the method is to the number and quality of these trajectories, and whether they can be replaced by cheaper supervision, is not addressed in the available content.
• Characterizing the cost-accuracy frontier. The paper shows favorable inference-time scaling but reports no explicit runtime or compute-cost accounting in the available content, so the practical trade-off between additional search and wall-clock budget remains unquantified.
• Understanding the verifier under distribution shift. TreeDiff performs best relative to baselines in the out-of-distribution QM9 setting. Whether the verifier's reward prediction stays reliable when targets are pushed further outside the training range, or when rewards are composed from many competing objectives, is a natural next question.

Target Audience

Researchers and graduate students working on generative models for graphs, diffusion model sampling and guidance, and molecule or material design, who are comfortable with MCTS, latent-variable models, and denoising formulations. Practitioners applying pretrained graph diffusion models who need controllable, multi-objective sampling without retraining will also benefit, particularly those working in computational chemistry and drug discovery. Readers looking for a beginner-level introduction to graph generation would find the paper demanding, and readers seeking full ablation studies or cost analyses will need the appendix, which is not part of the provided content.

Authors’ 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.