Graph Diffusion Policy Optimization

1 Introduction

• Objective-driven graph generation targets predefined rewards such as drug efficacy or discrete graph properties, rather than matching a training-data distribution.Examples of reward signals include binding affinity and synthetic accessibility.Source: §1 · p. 1
• Because graph rewards are often non-differentiable, predictor-based guidance can be inaccurate for complex signals, while existing RL work had not leveraged graph diffusion models.The paper identifies exact reward integration through RL as an alternative.Source: §1 · p. 1
• Directly extending continuous-diffusion optimization approaches to discrete graph diffusion models produces substantial empirical failure.GDPO adapts the discrete diffusion process to an MDP and introduces an eager policy gradient tailored to graph diffusion.Source: §1 · p. 2
• GDPO reports a 41.64% to 81.97% average reduction in generation-test distance and a 1.03% to 19.31% improvement in effective-drug generation, using 1/25 of the training samples.The reported results cover various scenarios and objectives.Source: §1 · p. 2

Graph Generative Models

• Graph diffusion models are the focus of GDPO because earlier graph DPM work learned from datasets, whereas this paper targets arbitrary reward signals.The related work describes both continuous-diffusion graph models and DiGress's discrete transitions.Source: §2 · p. 2

Controllable Generation for Graphs

• Prior controllable-generation methods use conditional models or property-predictor guidance, but predictor accuracy can be poor for difficult properties such as drug efficacy.MOOD adds out-of-distribution control to predictor-based approaches.Source: §2 · p. 2

Graph Generation using RL

• Graph RL methods include policy-gradient molecular environments and Q-learning, while the paper characterizes existing methods as computationally costly with limited policy-model capabilities.The cited examples include REINVENT, GCPN, MolDQN, DGAPN, and FREED.Source: §2 · p. 2

Aligning DPMs

• DPOK [16] and DDPO [6] use policy gradients to align image diffusion models with black-box rewards, while AlignProp [50] and DRaFT [11] backpropagate through differentiable rewards.These works are conducted on images; the paper presents GDPO as an effective method for aligning graph DPMs.Source: §2 · pp. 2–3

3 Preliminaries

• Graphs with categorical node and edge attributes are represented as G = (X, E), with one-hot feature vectors grouped into node and edge tensors.The formulation supports structured data such as molecules.Source: §3 · p. 3

3.1 Graph Diffusion Probabilistic Models

• Graph DPMs gradually corrupt data toward an approximately uniform categorical noise distribution, then learn reverse transitions to recover the data distribution.The reverse process uses neural predictions of clean node attributes, with an analogous construction for edges.Source: §3.1 · p. 3

3.2 Markov Decision Process and Policy Gradient

• An MDP specifies states, actions, transition probabilities, rewards, and an initial-state distribution; policy-gradient methods estimate the gradient of expected return.The paper's overview depicts sampling trajectories, weighting accumulated gradients by reward, and averaging across trajectories.Source: §3.2 · pp. 3–4, Figure 1 · p. 4
• REINFORCE estimates policy gradients using Monte Carlo simulation, providing the policy-gradient baseline discussed before the paper's method.The general MDP formulation describes trajectories as sequences of states and actions with cumulative rewards.Source: §3.2 · p. 4

4.1 A Markov Decision Process Formulation

• A graph DPM’s reverse denoising process becomes a T-step MDP whose policy is the reverse transition distribution and whose cumulative reward equals the final graph’s reward.This equivalence lets the expected-reward objective be optimized with a policy gradient.Source: §4.1 · p. 5

4.2 Learning Graph DPMs with Policy Gradient

• DDPO’s learning curves diminish as graph size increases and fail to reach high reward beyond four nodes, while GDPO performs consistently well.The toy task rewards connected graphs; the authors report that DDPO’s fluctuations also grow with node count.Source: §4.2 · p. 5, Figure 2 · p. 5
• Direct REINFORCE estimates rarely converge on graph DPMs; the authors conjecture high variance from the vast discrete trajectory space, while additional trajectories are costly.Their analysis says transition-gradient sums can emphasize candidate clean graphs independently of reward, including low-reward predictions, making estimates fluctuate when samples are limited.Source: §4.2 · pp. 5–6

4.3 Graph Diffusion Policy Optimization

• GDPO replaces transition log-gradient terms with gradients of predicting the final graph, directly favoring higher-reward outputs across denoising steps.This eager policy gradient optimizes over equivalence classes of trajectories sharing the same final graph, and is biased but reported to learn more stably than DDPO.Source: §4.3 · p. 6

5.1 Reward Functions for General Graph Generation

• For general graph generation, validity is a binary reward for whether a graph meets a specified type, while similarity rewards compare generated and reference graph distributions.The similarity signals use degree, clustering-coefficient, or substructure distributions; each generated graph receives a reward based on its distance to the reference set.Source: §5.1 · pp. 6–7

5.2 Reward Functions for Molecular Graph Generation

• Molecular novelty is the fingerprint-similarity gap between a generated molecule and its nearest training-set neighbor.The paper uses Tanimoto similarity and sets the novelty reward to this gap.Source: §5.2 · p. 7
• Molecular drug-related rewards encode docking score and QED, with QED rewarded only above 0.5.Docking scores are clipped to [−20, 0] and converted to reward as −DS(G0)/20.Source: §5.2 · p. 7
• Synthetic accessibility is rewarded using the normalized inverse of the compound’s 1-to-10 synthesis-difficulty score.The reward is (10 − SA(G0))/9.Source: §5.2 · p. 7

6.1 General Graph Generation

• General-graph rewards weight degree, clustering, and orbit terms by 0.1 each and validity by 0.7, without directly rewarding uniqueness or novelty.Source: §6.1 · p. 8
• GDPO reduces Planar distribution distances by 81.97% on average versus DiGress and reports a 41.64% average improvement on SBM, with gains in Deg and V.U.N on both datasets.Source: §6.1 · p. 8

6.2 Molecule Property Optimization

• Molecule optimization evaluates Hit Ratio and DS (top 5%) across five proteins, but uses QED, SA, novelty, and docking score as rewards instead.The reward weights are 0.1 for QED and SA, 0.3 for novelty, and 0.5 for docking score.Source: §6.2 · p. 9
• On ZINC250k, GDPO averages +5.72% Hit Ratio and +1.48% DS (top 5%); on MOSES, gains are +12.94% Hit Ratio across four targets and +5.54% DS versus MOOD.The MOSES Hit Ratio is less favorable on 5ht1b.Source: Table 2 · p. 8, §6.2 · p. 9, Table 3 · p. 9

6.3 Generalizability, Sample Efficiency, and A Failure Case

• On Spectral MMD, GDPO shows no overfitting and instead finds a more powerful model when optimized using Eq. (11).Source: §6.3 · pp. 9–10, Table 4 · p. 9
• With 10k reward queries, GDPO significantly improves average reward; fewer trajectories may improve sample efficiency but risk instability, while high rNOV can rapidly degenerate the model.Gradually increasing rNOV weight through multi-stage optimization is suggested as a potential solution.Source: §6.3 · p. 10, Figure 3 · p. 10

• GDPO addresses graph generation under specified objectives and achieves state-of-the-art performance across representative general and molecular graph-generation tasks.Future work will investigate the theoretical gap between GDPO and DDPO to obtain effective unbiased estimators.Source: §7 · p. 10

• Fine-tuning uses pretrained DiGress graph transformers, updates only attention layers, clips gradient norms to at most 1, and constrains normalized rewards to [−5, 5].Source: §A.1 · p. 15
• DiGress training on ZINC250k takes 48–72 hours on one A100 GPU, while GDPO fine-tuning takes 10 hours excluding reward computation.Source: §A.1 · p. 15
• GDPO improves average reward using 10,000 queries, a sample count notably below ZINC250k’s 250,000 samples.Source: §A.1 · p. 15

A.2 Limitations and Broader Impact.

• High novelty weighting can lead GDPO to a sudden reward drop after a period of optimization.The authors identify overoptimization as a risk in which the DPM distribution may collapse or diverge excessively from its original distribution.Source: §A.2 · p. 15
• GDPO inherits diffusion models’ training and inference costs because it leaves the underlying architecture unchanged.Source: §A.2 · p. 15
• Scalability to graphs with 500 or more nodes remains unexplored.Source: §A.2 · p. 15
• The authors report no societal consequence that they feel must be specifically highlighted.Source: §A.2 · p. 15

A.3 General Graph Generation

• The baselines span GraphRNN, SPECTRE, GDSS, DiGress, and MOOD, covering autoregressive, spectral GAN, score-based, discrete diffusion, and property-guided approaches.Source: §A.3 · pp. 15–16
• Algorithm 1 takes a graph DPM, reward signal, diffusion-step and trajectory-sample counts, learning rate, and training-step count, and outputs a final graph DPM.Its shown steps sample trajectories and timesteps and obtain rewards.Source: §A.3 · p. 16
• General graph metrics compare test-set and generated-graph distributions for degree, clustering-coefficient, and orbit counts, reporting ratios rather than raw numbers.Source: §A.3 · p. 16
• Clustering coefficients measure node clustering, while orbit counts capture graphlet occurrences and treat similar graphlet counts as indicating relatively similar graphs.Source: §A.3 · p. 16

A.4 Molecule Property Optimization

• Molecule optimization used five selected protein targets, with QED and SA computed using RDKit; the appendix also describes GCPN, REINVENT, HierVAE, FREED, and molecule-property metrics.The targets were selected for the highest AUROC scores when DUD-E ligand binding affinities were approximated with AutoDock Vina.Source: §A.4 · pp. 16–17

A.5 Additional Results of the GDPO

• GDPO is not very sensitive to several reward-weight configurations, can raise V.U.N above 80 with slight losses in other indicators, and outperforms DDPO-IS.Importance sampling improves DDPO but does not improve GDPO-IS over GDPO.Source: §A.5 · p. 18, Table 6 · p. 18
• GDPO learned a distinct denoising strategy rather than merely selecting a subset of molecules generated by the prior diffusion model.The passage refers to Table 7’s novelty and diversity results.Source: Table 7 · p. 18, §A.5 · p. 19
• Graph DPM trajectories can fluctuate more than image DPM trajectories because consecutive discrete graph samples can be far apart, contributing to high-variance policy gradients.The comparison uses normalized per-dimension L2 distances for models with 1,000 steps and 512 sampled trajectories.Source: §A.5 · p. 19, Figure 4 · p. 19
• GDPO effectively optimized tree-with-clique generation across 0–3 rewrite steps and 16–40 nodes, despite lower initial rewards as these factors increased.Graph size affected initial reward but not GDPO’s optimization performance.Source: Figure 6 · p. 20, §A.5 · p. 20
• Clique position produced no significant difference, while GDPO and DDPO showed a clear distinction on challenging tree-like generation tasks.The authors attribute similar optimization results across clique positions to similar initial rewards.Source: Figure 6 · p. 20, §A.5 · p. 21

A.6 Discussions

• RL approaches such as GDPO accommodate arbitrary rewards without intermediate approximations and collect labeled data during training from the current model distribution.The authors also report a significant sample-efficiency gap relative to classifier-based and classifier-free guidance.Source: §A.6 · p. 21
• GDPO’s bias comes from modifying the weight term to emphasize generated results over intermediate steps; the explored graph sizes may be too small to reveal the tradeoff.The authors plan to investigate it further on larger graph datasets with sparse-graph support.Source: §A.6 · p. 21