Graph Diffusion Policy Optimization

Authors: yijing  liu
chao  du
tianyu  pang
chongxuan  li
min  lin
wei  chen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 21:33:24

Abstract: Recent research has made significant progress in optimizing diffusion models for downstream objectives, which is an important pursuit in fields such as graph generation for drug design. However, directly applying these models to graph presents challenges, resulting in suboptimal performance. This paper introduces graph diffusion policy optimization (GDPO), a novel approach to optimize graph diffusion models for arbitrary (e.g., non-differentiable) objectives using reinforcement learning. GDPO is based on an eager policy gradient tailored for graph diffusion models, developed through meticulous analysis and promising improved performance. Experimental results show that GDPO achieves state-of-the-art performance in various graph generation tasks with complex and diverse objectives. Code is available at https://github.com/sail-sg/GDPO.

Main paper: 
 introduction: 
 Graph generation, a key facet of graph learning, has applications in a variety of domains, including drug and material design [54] , code completion [8] , social network analysis [20] , and neural architecture search [64] . Numerous studies have shown significant progress in graph generation with deep generative models [34, 62, 69, 21] . one of The most notable advances in the field is the introduction of graph diffusion probabilistic models (DPMs) [61, 31] . These methods can learn the underlying distribution from graph data samples and produce high-quality novel graph structures.In many use cases of graph generation, the primary focus is on achieving specific objectives, such as high drug efficacy [60] or creating novel graphs with special discrete properties [22] . These objectives are often expressed as specific reward signals, such as binding affinity [10] and synthetic accessibility [7] , rather than a set of training graph samples. Therefore, a more pertinent goal in such scenarios is to train graph generative models to meet these predefined objectives directly, rather than learning to match a distribution over training data [72] .A major challenge in this context is that most signals are non-differentiable w.r.t. graph representations, making it difficult to apply many optimization algorithms. To address this, methods based on property predictors [29, 37] learn parametric models to predict the reward signals, providing gradient guidance for graph generation. However, since reward signals can be highly complex (e.g., results from physical simulations), these predictors often struggle to provide accurate guidance [44] . An alternative direction is to learn graph generative models as policies through reinforcement learning (RL) [72] , which enables the integration of exact reward signals into the optimization. However, existing work primarily explores earlier graph generative models and has yet to leverage the superior performance of graph DPMs [9, 68] . On the other hand, several pioneer works have seen significant these works are conducted on images. To the best of our knowledge, our work is the first effective method for aligning graph DPMs, filling a notable gap in the literature. 
preliminaries: 
 In this section, we briefly introduce the background of graph DPMs and policy gradient methods.Following Vignac et al. [61] , we consider graphs with categorical node and edge attributes, allowing representation of diverse structured data like molecules. Let X and E be the space of categories for nodes and edges, respectively, with cardinalities a = |X | and b = |E|. For a graph with n nodes, we denote the attribute of node i by a one-hot encoding vector x (i) ∈ R a . Similarly, the attribute of the edge1 from node i to node j is represented as e (ij) ∈ R b . By grouping these one-hot vectors, the graph can then be represented as a tuple G ≜ (X, E), where X ∈ R n×a and E ∈ R n×n×b .Graph diffusion probabilistic models (DPMs) [61] involve a forward diffusion process q(G 1:T |G 0 ) = T t=1 q(G t |G t-1 ), which gradually corrupts a data distribution q(G 0 ) into a simple noise distribution q(G T ) over a specified number of diffusion steps, denoted as T . The transition distribution q(G t |G t-1 ) can be factorized into a product of categorical distributions for individual nodes and edges, i.e., q(x (i) t |x (i) t-1 ) and q(e(ij) t |e (ij)t-1 ). For simplicity, superscripts are omitted when no ambiguity is caused in the following. The transition distribution for each node is defined as q(x t |x t-1 ) = Cat(x t ; x t-1 Q t ), where the transition matrix is chosen asQ t ≜ α t I + (1 -α t )(1 a 1 ⊤a )/a, with α t transitioning from 1 to 0 as t increases [2] . It then follows that q(x t |x 0 ) = Cat(x t ; x 0 Qt ) and q(x t-1 |x t , x 0 ) = Cat(x t-1 ;xtQ ⊤ t ⊙ x0 Qt-1 x0 Qtx ⊤ t ), where Qt ≜ Q 1 Q 2 • • • Q t and⊙ denotes elementwise product. The design choice of Q t ensures that q(x T |x 0 ) ≈ Cat(x T ;1 a /a), i.e., a uniform distribution over X . The transition distribution for edges is defined similarly, and we omit it for brevity.Given the forward diffusion process, a parametric reverse denoising process p θ (G 0:T ) = p(G T ) T t=1 p θ (G t-1 |G t ) is then learned to recover the data distribution from p(G T ) ≈ q(G T ) (an approximate uniform distribution). The reverse transition p θ (G t-1 |G t ) is a product of categorical distributions over nodes and edges, denoted as p θ (x t-1 |G t ) and p θ (e t-1 |G t ). Notably, in line with the x 0 -parameterization used in continuous DPMs [25, 32] , p θ (x t-1 |G t ) is modeled as:EQUATIONwhere p θ ( x 0 |G t ) is a neural network predicting the posterior probability of x 0 given a noisy graph G t . For edges, each definition is analogous and thus omitted.The model is learned with a graph dataset D by maximizing the following objective [61] :EQUATIONwhere G 0 and t follow uniform distributions over D and [[1, T ]], respectively. After learning, graph samples can then be generated by first sampling G T from p(G T ) and subsequently sampling G t from p θ (G t-1 |G t ), resulting in a generation trajectory (G T , G T -1 , . . . , G 0 ).Markov decision processes (MDPs) are commonly used to model sequential decision-making problems [17] . An MDP is formally defined by a quintuple (S, A, P, r, ρ 0 ), where S is the state space containing all possible environment states, A is the action space comprising all available potential actions, P is the transition function determining the probabilities of state transitions, r is the reward signal, and ρ 0 gives the distribution of the initial state.In the context of an MDP, an agent engages with the environment across multiple steps. At each step t, the agent observes a state s t ∈ S and selects an action a t ∈ A based on its policy distribution π θ (a t |s t ). Subsequently, the agent receives a reward r(s t , a t ) and transitions to a new state s t+1 following the transition function P (s t+1 |s t , a t ). As the agent interacts in the MDP (starting from an initial state s 0 ∼ ρ 0 ), it generates a trajectory (i.e., a sequence of states and actions) denoted as τ = (s 0 , a 0 , s 1 , a 1 , . . . , s T , a T ). The cumulative reward over a trajectory τ is given by R(τ ) = T t=0 r(s t , a t ). In most scenarios, the objective is to maximize the following expectation:EQUATIONPolicy gradient methods aim to estimate ∇ θ J RL (θ) and thus solve the problem by gradient descent. An important result is the policy gradient theorem [19] , which estimates ∇ θ J RL (θ) as follows:EQUATIONThe REINFORCE algorithm [58] provides a simple method for estimating the above policy gradient using Monte-Carlo simulation, which will be adopted and discussed in the following section. 
method: 
 In this section, we study the problem of learning graph DPMs from arbitrary reward signals. We first present an MDP formulation of the problem and conduct an analysis on the failure of a direct application of REINFORCE. Based on the analysis, we introduce a substitute termed eager policy gradient, which forms the core of our method Graph Diffusion Policy Optimization (GDPO).A graph DPM defines a sample distribution p θ (G 0 ) through its reverse denoising process p θ (G 0:T ).Given a reward signal r(•) for G 0 , we aim to maximize the expected reward (ER) over p θ (G 0 ):EQUATIONHowever, directly optimizing J ER (θ) is challenging since the likelihood p θ (G 0 ) is unavailable [25] and r(•) is black-box, hindering the use of typical RL algorithms [6] . Following Fan et al. [16] , we We generate connected graphs with increasing number of nodes. Node categories are disregarded, and the edge categories are binary, indicating whether two nodes are linked. The graph DPM is initialized randomly as a one-layer graph transformer from DiGress [61] . The diffusion step T is set to 50, and the reward signal r(G 0 ) is defined as 1 if G 0 is connected and 0 otherwise. We use 256 trajectories for gradient estimation in each update. The learning curve illustrates the diminishing performance of DDPO as the number of nodes increases, while GDPO consistently performs well. formulate the denoising process as a T -step MDP and obtain an equivalent objective. Using notations in Sec. 3, we define the MDP of graph DPMs as follows:EQUATIONwhere the initial state s 0 corresponds to the initial noisy graph G T and the policy corresponds to the reverse transition distribution. As a result, the graph generation trajectory (G T , G T -1 , . . . , G 0 ) can be considered as a state-action trajectory τ produced by an agent acting in the MDP. It then follows that p(τ |π θ ) = p θ (G 0:T ). 2 Moreover, we have R(τ ) = T t=0 r(s t , a t ) = r(G 0 ). Therefore, the expected cumulative reward of the agentJ RL (θ) = E p(τ |π θ ) [R(τ )] = E p θ (G 0:T ) [r(G 0 )] is equivalent to J ER (θ), and thus J ER (θ) can also be optimized with the policy gradient ∇ θ J RL (θ):EQUATIONwhere the generation trajectory τ follows the parametric reverse process p θ (G 0:T ).The policy gradient ∇ θ J RL (θ) in Eq. ( 7) is generally intractable and an efficient estimation is necessary. In a related setting centered on continuous-variable DPMs for image generation, DDPO [6] estimates the policy gradient ∇ θ J RL (θ) with REINFORCE and achieves great results. This motivates us to also try REINFORCE on graph DPMs, i.e., to approximate Eq. ( 7) with a Monte Carlo estimation:EQUATIONwhere{G (k) 0:T } K k=1 are K trajectories sampled from p θ (G 0:T ) and {T k ⊂ [[1, T ]]} K k=1are uniformly random subsets of timesteps (which avoid summing over all timesteps and accelerate the estimation).However, we empirically observe that it rarely converges on graph DPMs. To investigate this, we design a toy experiment, where the reward signal is whether G 0 is connected. The graph DPMs are randomly initialized and optimized using Eq. (8) . We refer to this setting as DDPO. Fig. 2 depicts the learning curves, where the horizontal axis represents the number of queries to the reward signal and the vertical axis represents the average reward. The results demonstrate that DDPO fails to converge to a high reward signal area when generating graphs with more than 4 nodes. Furthermore, as the number of nodes increases, the fluctuation of the learning curves grows significantly. This implies that DDPO is essentially unable to optimize properly on randomly initialized models. We conjecture that the failure is due to the vast space constituted by discrete graph trajectories and the well-known high variance issue of REINFORCE [58] . A straightforward method to reduce variance is to sample more trajectories. However, this is typically expensive in DPMs, as each trajectory requires multiple rounds of model inference. Moreover, evaluating the reward signals of additional trajectories also incurs high computational costs, such as drug simulation [48] . This prompts us to delve deeper at a micro level. Since the policy gradient estimation in Eq. ( 8) is a weighted summation of gradients, we first inspect each summand term ∇ θ log p θ (G t-1 |G t ). With the parameterization Eq. ( 1) described in Sec. 3.1, it has the following form:EQUATIONwhere we can view the "weight" term as a weight assigned to the gradient ∇ θ p θ ( G 0 |G t ), and thus ∇ θ log p θ (G t-1 |G t ) as a weighted sum of such gradients, with G 0 taken over all possible graphs. Intuitively, the gradient ∇ θ p θ ( G 0 |G t ) promotes the probability of predicting G 0 from G t . Note, however, that the weight q(G t-1 |G t , G 0 ) is completely independent of r( G 0 ) and could assign large weight for G 0 that has low reward. Since the weighted sum in Eq. ( 9) can be dominated by gradient terms with large q(G t-1 |G t , G 0 ), given a particular sampled trajectory, it is fairly possible that ∇ θ log p θ (G t-1 |G t ) increases the probabilities of predicting undesired G 0 with low rewards from G t . This explains why Eq. ( 8) tends to produce fluctuating and unreliable policy gradient estimates when the number of Monte Carlo samples (i.e., K and |T k |) is limited. To further analyze why DDPO does not yield satisfactory results, we present additional findings in Appendix A.5. Besides, we discuss the impact of importance sampling techniques in the same section.To address the above issues, we suggest a slight modification to Eq. ( 8) and obtain a new policy gradient denoted as g(θ):EQUATIONwhich we refer to as the eager policy gradient. Intuitively, although the number of possible graph trajectories is tremendous, if we partition them into different equivalence classes according to G 0 , where trajectories with the same G 0 are considered equivalent, then the number of these equivalence classes will be much smaller than the number of graph trajectories. The optimization over these equivalence classes will be much easier than optimizing in the entire trajectory space.Technically, by replacing the summand gradient term ∇ θ log p θ (G t-1 |G t ) with ∇ θ log p θ (G 0 |G t ) in Eq. ( 8), we skip the weighted sum in Eq. ( 9) and directly promotes the probability of predicting G 0 which has higher reward from G t at all timestep t. As a result, our estimation does not focus on how G t changes to G t-1 within the trajectory; instead, it aims to force the model's generated results to be close to the desired G 0 , which can be seen as optimizing in equivalence classes. While being a biased estimator of the policy gradient ∇ θ J RL (θ), the eager policy gradient consistently leads to more stable learning and better performance than DDPO, as demonstrated in Fig. 2 . We present the resulting method in Fig. 1 and Algorithm 1, naming it Graph Diffusion Policy Optimization (GDPO). 
reward functions for graph generation: 
 In this work, we study both general graph and molecule reward signals that are crucial in real-world tasks. Below, we elaborate on how we formulate diverse reward signals as numerical functions.Validity. For graph generation, a common objective is to generate a specific type of graph. For instance, one might be interested in graphs that can be drawn without edges crossing each other [43] . For such objectives, the reward function r val (•) is then formulated as binary, with r val (G 0 ) ≜ 1 indicating that the generated graph G 0 conforms to the specified type; otherwise, r val (G 0 ) ≜ 0.Similarity. In certain scenarios, the objective is to generate graphs that resemble a known set of graphs D at the distribution level, based on a pre-defined distance metric d(•, •) between sets of graphs. As an example, the Deg(G, D) [38] measures the maximum mean discrepancy (MMD) [18] between the degree distributions of a set G of generated graphs and the given graphs D. Since our method requires a reward for each single generated graph G 0 , we simply adopt Deg({G 0 }, D) as the signal. As the magnitude of reward is critical for policy gradients [58] , we definer deg (G 0 ) ≜ exp -Deg({G 0 }, D) 2 /σ 2, where the σ controls the reward distribution, ensuring that the reward lies within the range of 0 to 1. The other two similar distance metrics are Clus(G, D) and Orb(G, D), which respectively measure the distances between two sets of graphs in terms of the distribution of clustering coefficients [55] and the distribution of substructures [1] . Based on the two metrics, we define two reward signals analogous to r deg , namely, r clus and r orb .Novelty. A primary objective of molecular graph generation is to discover novel drugs with desired therapeutic potentials. Due to drug patent restrictions, the novelty of generated molecules has paramount importance. The Tanimoto similarity [4] , denoted as J(•, •), measures the chemical similarity between two molecules, defined by the Jaccard index of molecule fingerprint bits. Specifically, J ∈ [0, 1], and J(G 0 , G ′ 0 ) = 1 indicates that two molecules G 0 and G ′ 0 have identical fingerprints. Following Xie et al. [65] , we define the novelty of a generated graph G 0 asNOV(G 0 ) ≜ 1 -max G ′ 0 ∈D J(G 0 , G ′ 0 ), i.e., the similarity gap between G 0 and its nearest neighbor in the training dataset D, and further define r NOV (G 0 ) ≜ NOV(G 0 ). Drug-Likeness. Regarding the efficacy of molecular graph generation in drug design, a critical indicator is the binding affinity between the generated drug candidate and a target protein. The docking score [10] , denoted as DS(•), estimates the binding energy (in kcal/mol) between the ligand and the target protein through physical simulations in 3D space. Following Lee et al. [37] , we clip the docking score in the range [-20, 0] and define the reward function asr DS (G 0 ) ≜ -DS(G 0 )/20.Another metric is the quantitative estimate of drug-likeness QED(•), which measures the chemical properties to gauge the likelihood of a molecule being a successful drug [5] . As it takes values in the range[0, 1], we adopt r QED (G 0 ) ≜ I[QED(G 0 ) > 0.5].Synthetic Accessibility. The synthetic accessibility [7] SA(•) evaluates the inherent difficulty in synthesizing a chemical compound, with values in the range from 1 to 10. We follow Lee et al. [37] and use a normalized version as the reward function:r SA (G 0 ) ≜ (10 -SA(G 0 ))/9. 
experiments: 
 In this section, we first examine the performance of GDPO on both general graph generation tasks and molecular graph generation tasks. Then, we conduct several ablation studies to investigate the effectiveness of GDPO's design. Our code can be found in the supplementary material. Datasets and Baselines. Following DiGress [61] , we evaluate GDPO on two benchmark datasets: SBM (200 nodes) and Planar (64 nodes), each consisting of 200 graphs. We compare GDPO with GraphRNN [69] , SPECTRE [43] , GDSS [31] , MOOD [37] and DiGress. The first two models are based on RNN and GAN, respectively. The remaining methods are graph DPMs, and MOOD employs an additional property predictor. We also test DDPO [6] , i.e., graph DPMs optimized with Eq. (8). Metrics and Reward Functions. We consider four metrics: Deg(G, D test ), Clus(G, D test ), Orb(G, D test ), and the V.U.N metrics. V.U.N measures the proportion of generated graphs that are valid, unique, and novel. The reward function is defined as follows: GCPN [68] , REINVENT [47] , FREED [67] and MOOD [37] . GCPN, REINVENT and FREED are RL methods that search in the chemical environment. MOOD, based on graph DPMs, employs a property predictor for guided sampling. Similar to general graph generation, we also compare our method with DiGress and DDPO. Besides, we show the performance of DiGress with property predictors, termed as DiGress-guidance.EQUATIONImplementation. We set T = 500, |T | = 100, N = 100, and K = 256 for both datasets. We use the same model structure with DiGress. See more details in Appendix A.1.Metrics and Reward Functions. Following MOOD, we consider two metrics essential for real-world novel drug discovery: Novel hit ratio (%) and Novel top 5% docking score, denoted as Hit Ratio and DS (top 5%), respectively. Using the notations from Sec. 5.2, the Hit Ratio is the proportion of unique generated molecules that satisfy: DS < median DS of the known effective molecules, NOV > 0.6, QED > 0.5, and SA < 5. The DS (top 5%) is the average DS of the top 5% molecules (ranked by DS) that satisfy: NOV > 0.6, QED > 0.5, and SA < 5. Since calculating DS requires specifying a target protein, we set five different protein targets to fully test GDPO: parp1, fa7, 5ht1b, braf, and jak2. The reward function for molecule property optimization is defined as follows:EQUATION)We do not directly use Hit Ratio and DS (top 5%) as rewards in consideration of method generality. The reward weights are determined through several rounds of search, and we find that assigning a high weight to r NOV leads to training instability, which is discussed in Sec. 6.3. More details about the experiment settings are discussed in Appendix A.4.Results. In Table 2 , GDPO shows significant improvement on ZINC250k, especially in the Hit Ratio. A higher Hit Ratio means the model is more likely to generate valuable new drugs, and GDPO averagely improves the Hit Ratio by 5.72% in comparison with other SOTA methods. For DS (top 5%), GDPO also has a 1.48% improvement on average. Discovering new drugs on MOSES is much more challenging than on ZINC250k due to its vast training dataset. In Table 3 , GDPO also shows promising results on MOSES. Despite a less favorable Hit Ratio on 5ht1b, GDPO achieves an average improvement of 12.94% on the other four target proteins. For DS (top 5%), GDPO records an average improvement of 5.54% compared to MOOD, showing a big improvement in drug efficacy. To validate whether GDPO correctly optimizes the model, we test the performance of GDPO on metrics not used in the reward signal. In Table 4, we evaluate the performance on Spectral MMD [43] , where the GDPO is optimized by Eq. (11) . The results demonstrate that GDPO does not show overfitting; instead, it finds a more powerful model. The results presented in Appendix A.5 further support that GDPO can attain high sample novelty and diversity. We then investigate two crucial factors for GDPO: 1) the number of trajectories; 2) the selection of the reward signals. We test our method on ZINC250k and set the target proteins as 5ht1b. In Fig. 3 (a), the results indicate that GDPO exhibits good sampling efficiency, as it achieves a significant improvement in average reward by querying only 10k molecule reward signals, which is much less than the number of molecules contained in ZINC250k. Moreover, the sample efficiency can be further improved by reducing the number of trajectories, but this may lead to training instability. To achieve consistent results, we use 256 trajectories. In Fig. 3 (b), we illustrate a failure case of GDPO when assigning a high weight to r NOV . Generating novel samples is challenging. MOOD [37] addresses this challenge by controlling noise in the sampling process, whereas we achieve it by novelty optimization. However, assigning a large weight to r NOV can lead the model to rapidly degenerate. One potential solution is to gradually increase the weight and conduct multi-stage optimization. 
conclusion: 
 We introduce GDPO, a novel policy gradient method for learning graph DPMs that effectively addresses the problem of graph generation under given objectives. Evaluation results on both general and molecular graphs indicate that GDPO is compatible with complex multi-objective optimization and achieves state-of-the-art performance on a series of representative graph generation tasks. We discuss some limitations of our work in Appendix A.2. Our future work will investigate the theoretical gap between GDPO and DDPO in order to obtain effective unbiased estimators. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: In the appendix, we discuss the potential impacts of this work in detail.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This work does not involve any relevant datasets or models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All datasets used in this paper are open-source, and there are no copyright issues involved.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: A.1 Implementation Details.For all experiments, we use the graph transformer proposed in DiGress [61] as the graph DPMs, and the models are pre-trained on the training dataset before applying GDPO or DDPO. During fine-tuning, we keep all layers fixed except for attention, set the learning rate to 0.00001, and utilize gradient clipping to limit the gradient norm to be less than or equal to 1. In addition, due to significant numerical fluctuations during reward normalization, we follow DDPO [6] in constraining the normalized reward to the range from [-5, 5] . This means that gradients resulting from rewards beyond this range will not contribute to model updates. When there is insufficient memory to generate enough trajectories, we use gradient accumulation to increase the number of trajectories used for gradient estimation. We conducted all experiments on a single A100 GPU with 40GB of VRAM and an AMD EPYC 7352 24-core Processor.Training time and efficiency. Training DiGress on the ZINC250k dataset using a single A100 GPU typically takes 48-72 hours, whereas fine-tuning with GDPO takes only 10 hours (excluding the time for reward function computation). This high efficiency is in line with the findings in the practice of DDPO, which is different from traditional RL methods. Additionally, as in Fig. 3 and Sec 6.3, GDPO effectively improves the average reward of the model using only 10,000 queries. This sample size is notably small compared to the 250,000 samples present in the ZINC250k dataset, showing the impressive sample efficiency of GDPO.A.2 Limitations and Broader Impact.Below we list some limitations of the current work:• Potential for overoptimization: As an RL-based approach, a recognized limitation is the risk of overoptimization, where the DPM distribution may collapse or diverge excessively from the original distribution. In Section 6.3, we demonstrated a failure case where, with a high weight on novelty in the reward function, GDPO encounters a sudden drop in reward after a period of optimization. Future research could explore the application of regularization techniques, similar to those utilized in recent works such as DPO [51] , to mitigate this risk.• Inherited limitations of DPMs: Our method inherits certain limitations inherent to diffusion models, particularly concerning their training and inference costs. As we do not modify the underlying model architecture, these constraints persist.• Scalability to large graphs: The scalability of GDPO to larger graphs (e.g., with 500 or more nodes) remains unexplored.For broader impact, this paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.Baselines. There are several baseline methods for general graph generation, we summarize them as follows:• GraphRNN: a deep autoregressive model designed to model and generate complex distributions over graphs. It addresses challenges like non-uniqueness and high dimensionality by decomposing the generation process into node and edge formations.• SPECTRE: a novel GAN for graph generation, approaches the problem spectrally by generating dominant parts of the graph Laplacian spectrum and matching them to eigenvalues and eigenvectors. This method allows for modeling global and local graph structures directly, overcoming issues like expressivity and mode collapse.• GDSS: A novel score-based generative model for graphs is introduced to tackle the task of capturing permutation invariance and intricate node-edge dependencies in graph data generation. This model employs a continuous-time framework incorporating a novel graph diffusion process, T k ∼ Uniform([[1, T ]]) // Sample timesteps r k ← r(G (k) 0 )// Get rewards // Estimate reward mean and variancer ← 1 K K k=1 r k std[r] ← K k=1 (r k -r) 2 K-1// Estimate the eager policy gradientg(θ) ← 1 K K k=1 T |T k | t∈T k ( r k -r std[r] )∇ θ log p θ (G (k) 0 |G (k) t ) // Update model parameter θ ← θ + η • g(θ)characterized by stochastic differential equations (SDEs), to simultaneously model distributions of nodes and edges. • DiGress: DiGress is a discrete denoising diffusion model designed for generating graphs with categorical attributes for nodes and edges. It employs a discrete diffusion process to iteratively modify graphs with noise, guided by a graph transformer network. By preserving the distribution of node and edge types and incorporating graph-theoretic features, DiGress achieves state-ofthe-art performance on various datasets. • MOOD: MOOD introduces Molecular Out-Of-distribution Diffusion, which employs out-ofdistribution control in the generative process without added costs. By incorporating gradients from a property predictor, MOOD guides the generation process towards molecules with desired properties, enabling the discovery of novel and valuable compounds surpassing existing methods.Metrics. The metrics of general graph generations are all taken from GraphRNN [38] . The reported metrics compare the discrepancy between the distribution of certain metrics on a test set and the distribution of the same metrics on a generated graph. The metrics measured include degree distributions, clustering coefficients, and orbit counts (which measure the distribution of all substructures of size 4). Following DiGress [61] , we do not report raw numbers but ratios computed as follows:EQUATIONBesides, we explain some metrics that are used in the general graph generation:• Clus: the clustering coefficient measures the tendency of nodes to form clusters in a network. Real-world networks, especially social networks, often exhibit tightly knit groups with more ties between nodes than expected by chance. There are two versions of this measure: global, which assesses overall clustering in the network, and local, which evaluates the clustering around individual nodes. • Orb: Graphlets are induced subgraph isomorphism classes in a graph, where occurrences are isomorphic or non-isomorphic. They differ from network motifs, which are over-or underrepresented graphlets compared to a random graph null model. Orb will count the occurrences of each type of graphlet in a graph. Generally, if two graphs have similar numbers of graphlets, they are considered to be relatively similar.Implementation Details. Following FREED [67] , we selected five proteins, PARP-1 (Poly [ADPribose] polymerase-1), FA7 (Coagulation factor VII), 5-HT1B (5-hydroxytryptamine receptor 1B), BRAF (Serine/threonine-protein kinase B-raf), and JAK2 (Tyrosine-protein kinase JAK2), which have the highest AUROC scores when the protein-ligand binding affinities for DUD-E ligands are approximated with AutoDock Vina [13] , as the target proteins for which the docking scores are calculated. QED and SA scores are computed using the RDKit library.Baselines. There are several baseline methods for molecular graph generation under the given objectives, they are diverse in methodology and performance, we summarize them as follows:• GCPN: Graph Convolutional Policy Network (GCPN) is a general graph convolutional networkbased model for goal-directed graph generation using reinforcement learning. The GCPN is trained to optimize domain-specific rewards and adversarial loss through policy gradient, operating within an environment that includes domain-specific rules. • REINVENT: This method enhances a sequence-based generative model for molecular design by incorporating augmented episodic likelihood, enabling the generation of structures with specified properties. It successfully performs tasks such as generating analogs to a reference molecule and predicting compounds active against a specific biological target. • HierVAE: a hierarchical graph encoder-decoder for drug discovery, overcoming limitations of previous approaches by using larger and more flexible graph motifs as building blocks. The encoder generates a multi-resolution representation of molecules, while the decoder adds motifs in a coarse-to-fine manner, effectively resolving attachments to the molecule. • FREED: a novel reinforcement learning (RL) framework for generating effective acceptable molecules with high docking scores, crucial for drug design. FREED addresses challenges in generating realistic molecules and optimizing docking scores through a fragment-based generation method and error-prioritized experience replay (PER). • MOOD: please refer to Appendix A.3.Metrics. There are several metrics for evaluating the molecule properties, we summarize the meaning of these metrics as follows:• Docking Score: Docking simulations aim to find the best binding mode based on scoring functions. Scoring functions in computational chemistry and molecular modeling predict binding affinity between molecules post-docking. They are commonly used for drug-protein interactions, but also for protein-protein or protein-DNA interactions. After defining the score function, we can optimize to find the optimal drug-protein matching positions and obtain the docking score. • QED: Drug-likeness evaluation in drug discovery often lacks nuance, leading to potential issues with compound quality. We introduce QED, a measure based on desirability, which considers the distribution of molecular properties and allows the ranking of compounds by relative merit. QED is intuitive, transparent, and applicable to various settings. We extend its use to assess molecular target druggability and suggest it may reflect aesthetic considerations in medicinal chemistry. • SA: a scoring method for rapid evaluation of synthetic accessibility, considering structural complexity, similarity to available starting materials, and strategic bond assessments. These components are combined using an additive scheme, with weights determined via linear regression analysis based on medicinal chemists' accessibility scores. The calculated synthetic accessibility values align well with chemists' assessments. Study of the Reward Signals. In Table . 5, we showcase the performance of GDPO on Planar under different configurations of reward weights. We keep the three weights related to distance the same and adjust the weight of validity while ensuring that the sum of weights is 1. The results indicate that GDPO is not very sensitive to the weights of several reward signals for general graph generation, even though these weight configurations vary significantly, they all achieve good performance. Additionally, we found that GDPO can easily increase V.U.N to above 80 while experiencing slight losses in the other three indicators. When applying GDPO in practice, one can make a tradeoff between them based on the specific application requirements. The Impact of Important Sampling. The importance sampling technique in DDPO, aims to facilitate multiple steps of optimization using the same batch of trajectories. This is achieved by weighting each item on the trajectory with an importance weight derived from the density ratio estimated using the model parameters from the previous step θ prev and the current step θ (referred to as DDPO-IS):EQUATIONOur eager policy gradient, independently motivated, aims to address the high variance issue of the policy gradient in each step of optimization, as elaborated in Sec. 4.2. Intuitively, the eager policy gradient can be viewed as a biased yet significantly less fluctuating gradient estimation.We conducted a series of experiments on ZINC250k to compare DDPO, DDPO-IS, and GDPO. The experimental setup remains consistent with the description in Section 6.2. Additionally, considering that the importance sampling technique in DDPO and our eager policy gradient appear to be orthogonal, we also explored combining them simultaneously (referred to as GDPO-IS):EQUATIONIn Table . 6, while importance sampling enhances the performance of DDPO, consistent with the results reported in the DDPO paper, it does not yield improvements for GDPO-IS over GDPO. We speculate that this discrepancy may be due to the biasness of the eager policy gradient, rendering it incompatible with the importance sampling technique. We intend to investigate the mechanism and address this in our future work. Nevertheless, it is noteworthy that the performance of DDPO-IS remains inferior to GDPO, indicating the superiority of our proposed GDPO method. Novelty and Diversity of GDPO. To provide further insight into the novelty and diversity of our approach, we introduce two additional metrics:• Intersection over Union (IoU): We compare two sets of molecules: 1) 500 molecules generated by GDPO (denoted as GDPO) and 2) top 500 molecules among 10,000 molecules generated by our base DPM before finetuning (denoted as TopPrior). We then compute IoU=100 × |GDPO∩TopPrior| |GDPO∪TopPrior| %. We report an average IoU of 5 independent runs.• Uniqueness in 10k samples (Uniq): We generate 10,000 molecules and compute the ratio of unique molecules Uniq = 100 × # unique molecules # all molecules %.In Table . 7, these results show that GDPO has not converged to a trivial solution, wherein it merely selects a subset of molecules generated by the prior diffusion model. Instead, GDPO has learned an effective and distinct denoising strategy from the prior diffusion model.The Gap between Image DPMs and Graph DPMs. GDPO is tackling the high variance issue inherent in utilizing policy gradients on graph DPMs, as stated and discussed in Sec. 4.2. To provide clarity on what GDPO tackles, we would like to elaborate more on the high variance issue of policy gradients on graph DPMs. Consider the generation trajectories in image and graph DPMs:In image DPMs, the generation process follows a (discretization of) continuous diffusion process(x t ) t∈[0,T ] .The consecutive steps x t-1 and x t are typically close due to the Gaussian reverse denoising distribution p(x t-1 |x t ) (typically with a small variance).In graph DPMs, the generation process follows a discrete diffusion process (G T , . . . , G 0 ), where each G t is a concrete sample (i.e., one-hot vectors) from categorical distributions. Therefore, consecutive steps G t-1 and G t can be very distant. This makes the trajectory of graph DPMs more fluctuating than images and thus leads to a high variance of the gradient ∇ θ log p(G t-1 |G t ) (and the ineffectiveness of DDPO) when evaluated with same number of trajectories as in DDPO.Regarding the "distance" between two consecutive steps G t and G t-1 , our intuition stems from the fact that graphs generation trajectories are inherently discontinuous. This means that each two consecutive steps can differ significantly, such as in the type/existence of edges. In contrast, the generation trajectories of images, governed by reverse SDEs, are continuous. This continuity implies that for fine-grained discretization (i.e., large T ), x t and x t-1 can be arbitrarily close to each other (in the limit case of T → ∞). In Fig. 4 , the results support the explanation of GDPO. While we acknowledge that graphs and images reside in different spaces and typically have different representations, we believe the comparison with L2 distance can provide valuable insights into the differences between graph and image DPMs.GDPO on the Synthetic Tree-like Dataset. We first generate a tree and then connect a clique to the nodes of the tree, performing a specified number of rewrite operations as suggested. Based on The model receives a reward of 1 when it generates a tree with a clique; otherwise, the reward is 0.We then ablate the following factors to test the performance of GDPO.Rewrite Steps: In Fig. 6 (a), we demonstrate GDPO's performance across different rewrite steps, with four curves representing steps ranging from 0 to 3. Despite a notable decrease in the initial reward as the number of rewrite steps increases, GDPO consistently optimizes the Graph DPMs effectively to generate the desired graph structure.Graph Size: In Fig. 6 (b), we gradually increase the number of nodes from 16 to 40. The results show that graph size affects the initial reward but does not impact GDPO's optimization performance.Clique Position: We experiment with inserting the clique at different levels of the tree but find no significant difference. We believe this is because the position of the clique does not affect the initial reward of the Graph DPMs, leading to similar optimization results with GDPO.Comparison with Baseline: In Fig. 6 (c), we compare GDPO with DDPO. The results, consistent with those in Figure 2 of the paper, reveal a clear distinction between GDPO and DDPO in handling challenging data generation tasks.Comparison with the x 0 -prediction Formulation. Indeed, our eager policy gradient in Eq. 10, compared to the policy gradient of REINFORCE in Eq. 8, resembles the idea of training a denoising network to predict the original uncorrupted graph rather than performing one-step denoising. However, we note that training a denoising network to predict the original data is fundamentally a matter of parametrization of one-step denoising. Specifically, the one-step denoising p θ (x t-1 |G t ) is parameterized as a weighted sum of x 0 -prediction, as described in Eq. Analysis on the Bias-variance Trade off. The main bias of GDPO arises from modifying the "weight" term in Eq. 9, which shifts the model's focus more towards the generated results rather than the intermediate process, thereby reducing potential noise. Due to the discrete nature of Graph DPMs, the x 0 -prediction and x t-1 -prediction formulations cannot be related through denoising objectives as in continuous DPMs. This issue also complicates the connection between DDPO and GDPO. Wehave not yet identified a relevant solution and are still working on it. In our empirical study, we do not observe significant performance variance and tradeoff for GDPO given the current scale of experiments. This may be due to the graph sizes we explored not being sufficiently large. In future implementations, we will incorporate support for sparse graphs to assess GDPO's performance on larger graph datasets and investigate the tradeoff more thoroughly.Answer Justification: In the appendix, we provide these contents.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).For convenience, "no edge" is treated as a special type of edge.With a slight abuse of notation we will use τ = G0:T and τ = (s0, a0, s1, a1, . . . , sT , aT ) interchangeably, which should not confuse as the MDP relates them with a bijection. 


Decision: The paper proposes a methodology to optimize graph diffusion probabilistic models for arbitrary reward signals. The methodology has several novel aspects including the Markov Decision Process formulation and the policy gradient considered. Despite the fact that the paper lacks strong theoretical support, the empirical results are quite strong.  The response of the authors to the different questions raised further clarified several aspects of the paper. Therefore, I recommend accepting the paper.

Reviews: review: paper_summary: paper_summary: This paper introduces Graph Diffusion Policy Optimization (GDPO), a novel approach to optimize graph diffusion models for arbitrary objectives using reinforcement learning. The key contributions are:

1. Formulating the denoising process of graph diffusion probabilistic models (DPMs) as a Markov decision process and proposing an "eager policy gradient" method tailored for graph DPMs to address high variance issues in standard policy gradient approaches.
2. Demonstrating state-of-the-art performance on various graph generation tasks, including general graph generation and molecular graph generation with complex objectives.

The authors show that GDPO significantly outperforms baseline methods, including other graph generation techniques and adaptations of diffusion model optimization approaches from the image domain (e.g., DDPO).
strength_weakness: strength_weakness: - Limited theoretical analysis: While the eager policy gradient is empirically effective, the paper lacks a rigorous theoretical treatment of its properties, particularly regarding the bias-variance trade-off.
-  The paper would benefit from a comparison to other RL-utilizing graph generation methods, particularly MolGAN https://arxiv.org/pdf/1805.11973 , which also applies RL techniques to molecular graph generation.
- Scalability concerns: The paper does not explore the method's performance on very large graphs (e.g., 500+ nodes), leaving questions about its scalability unanswered.
- Limited exploration of failure cases: While the authors provide a failure case related to novelty optimization, a more comprehensive exploration of scenarios where GDPO struggles would provide valuable insights into its limitations.
questions: questions: 1. How sensitive is GDPO to the choice of reward function? Are there certain types of rewards or objectives that are particularly challenging for the method (e.g., very sparse, high dynamic range, noisy...)?
2. I'd like to suggest an important ablation study on highly ambiguous graphs to better understand the limitations and robustness of GDPO. Specifically, I propose constructing a "noisy-tree-with-planted-motifs" dataset with the following process:

---

1. Generate a base tree structure with parameters:
   - Tree depth
   - Minimum fan-out factor
   - Maximum fan-out factor
   - Binned beta variation across the fan-out for a given node

2. Plant a tree of cliques as an expansion of  nodes at a chosen height in the base tree. This planted structure should have parameters:
   - Number of rings of cliques
   - Ring size and clique count

3. Apply a noisy rewiring process towards an Erdős-Rényi random graph with the same number of nodes and edges as the base graph. This process should take as parameters:
   - Number of rewiring steps
   - starting layer of the subgraph to be noised (starting at a given layer of the tree)
  
The noisy rewiring process would involve:
a) Randomly selecting an edge
b) Checking whether to remove it based on the probability of an edge existing in the target Erdős-Rényi graph
c) Selecting a random node pair to potentially add an edge
d) Repeating this process for the specified number of steps

I suggest ablating over:
1. The number of rewiring steps
2. The size of the subgraph to be noised
3. The overall graph size
4. The height at which the tree of cliques is planted
5. (optional) clique size and ring size

My hypothesis is that GDPO would struggle to maintain the planted tree of cliques structure and the overall tree structure as the noise level and graph size increase, as it would slowly fill up the capacity of the model with combinatorial structures. This ablation study would provide valuable insights into GDPO's performance on more complex and ambiguous graph structures.

Questions I'd like answered through this ablation study:

1. How does GDPO's performance change as the noise level (number of rewiring steps) increases?
2. What is the impact of graph size on GDPO's ability to maintain the planted tree of cliques and the overall tree structure?
3. Is there a critical point in the noise level or graph size where GDPO's performance significantly degrades?
4. How does the height at which the tree of cliques is planted affect GDPO's ability to maintain this structure? (i.e., how well does it maintain long-range dependeny)
5. How does GDPO compare to baseline methods on these more challenging graph structures?
6. Can you identify specific types of motifs or structures within the planted tree of cliques that GDPO struggles to maintain or generate in these noisy environments?
limitations: limitations: I think the limitations in the appendix are acceptable
paper_summary: paper_summary: This paper introduces graph diffusion policy optimization (GDPO), a policy gradient method for optimizing graph diffusion probabilistic models with respect to non-differentiable reward signals. By establishing the connection between a T-step denoising process and a T-step Markov Decision Process (MDP), policy gradient methods can be applied to graph diffusion models. While a previous work (DDPO) report s competitive generation quality for image diffusion models trained with the classical REINFORCE algorithm, the authors empirically observe a convergence issue when applying REINFORCE to graph diffusion models. This issue is possibly due to the increasingly vast space constituted by discrete graph trajectories as the number of nodes in the graph increases. To address this issue, the authors propose GDPO, a modified policy optimization objective. Empirical studies demonstrate the effectiveness of the proposed modified objective.
strength_weakness: strength_weakness: **W1.** Some experiment settings are not completely clear from the description. See the questions below.

**W2.** The proposed modification coincides with the idea of training a denoising network to predict the original uncorrupted graph rather than perform one-step denoising. Some discussions are expected.

**W3.** Classifier-based and classifier-free guidance are two popular approaches for training conditional diffusion models and have been previously explored for graph diffusion models. Some discussions on the potential pros and cons of the RL approach against them are expected.
questions: questions: **Q1.** For the graph diffusion baselines considered, are they trained using classifier-free or classifier-based guidance?

**Q2.** A key question in understanding the pros and cons of the different conditional generation approaches is sample efficiency. If ground truth reward functions are employed for GDPO training, then similarly we can label the samples generated by non-RL conditional graph diffusion models as extra training samples. The key question is then how the performance of the different models changes against the number of reward function evaluations.
limitations: limitations: The paper considers discrete time graph diffusion models. Whether the proposed approach is effective for continuous time graph diffusion models remains unexplored.
paper_summary: paper_summary: The paper studies the problem of learning graph diffusion generative models on arbitrary non-differentiable objectives using policy gradients. Authors argue that the recently proposed DDPO technique doesn't work well on the discrete, graph-related learning tasks and consider a modified objective and a corresponding gradient estimate which they refer to as Graph Diffusion Policy Optimization (GDPO). They show that GDPO performs significantly better on a number of reward functions and datasets.
strength_weakness: strength_weakness: ## Underexplored potential of modern RL techniques

While the proposed GDPO is valueable regardless, it is unclear to me that the biased objective / gradient estimate of eager gradients is necessary. Neither DDPO nor this paper explores even the simplest techniques such as actor-critic, PPO / TRPO or more sophisticated versions of importance sampling. DDPO showed good performance on image tasks without such advances and it could well be that it would still work well on graph tasks.

## Lack of bias-variance analysis

Authors acknowledge that eager gradients is a biased version of the standard policy gradients but don't provider either theoretical or empirical analysis of the bias-variance tradeoff between the two methods. 

In the appendix (line 660) authors a strange statement that importance sampling is used to reduce variance. It is not obvious to me all. Importance sampling can both reduce or increase variance depending on the proposal or the policy generating the trajectory. What it achieves is that it allows to train on experiences generated by policies other than the current policy being optimized.
questions: questions: My main question is what is the nature of the bias the GDPO introduces? Is there a reasonable objective that is followed go eager gradients? Can we still interpret them as policy gradients in some kind of modified MDP?

Figure 4: what is $D_G$ and $D_I$ concretely? Does it make sense to assess $L_2$ distances on discrete objects with categorical features?

What happens if you apply GDPO to image diffusion? Does it still work or work better?
limitations: limitations: See above, I believe the paper needs a clearer discussion of the bias introduced by policy gradients. At the moment, it is not clear what is the connection of GDPO to the reverse diffusion MDP.

Comments: title: Thank you for your feedback
Thank you for your support and raising the score
comment: We appreciate your valuable feedback. We will further polish the paper and incorporate the rebuttal discussions into the final revision. Thank you!
I thank authors for their response which I found helpful but ultimately not convincing enough for me to raise my score.
Thank you for your timely response and for acknowledging our efforts. We appreciate your positive feedback, especially regarding the strengthened contribution of the ablation study. We will further explore deeper/wider graphs and incorporate these findings, along with the rebuttal discussions, into the final revision. Thank you once again!
Thank your for performing the ablation study and addressing my other points. I am surprised and intrigued by the results of the ablation study and am now curious on how it would look like when investing a larger compute amount to deeper/wider graphs, but I fully understand this is not feasible in this round of investigation.  Since I already had the highest score across reviewers,  I will retain my score as is for now, but I  think this and the other additions have  made the paper stronger and I'm now quite comfortable with my score.
We greatly appreciate your valuable feedback and the score improvement. We will further polish the paper and incorporate the rebuttal discussions into the final revision. Thank you!
Thank you for your response, which has addressed most of my questions. I've increased the "Soundness score" from 2 to 3 and the rating from 5 to 6.