Graph Diffusion Policy Optimization

Graph diffusion probabilistic models (DPMs) [61] involve a forward diffusion process $q(\bm{G}_{1:T}|\bm{G}_{0})=\prod_{t=1}^{T}q(\bm{G}_{t}|\bm{G}_{t-1})$, which gradually corrupts a data distribution $q(\bm{G}_{0})$ into a simple noise distribution $q(\bm{G}_{T})$ over a specified number of diffusion steps, denoted as $T$. The transition distribution $q(\bm{G}_{t}|\bm{G}_{t-1})$ can be factorized into a product of categorical distributions for individual nodes and edges, i.e., $q(\bm{x}^{(i)}_{t}|\bm{x}^{(i)}_{t-1})$ and $q(\bm{e}^{(ij)}_{t}|\bm{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(\bm{x}_{t}|\bm{x}_{t-1})=\operatorname{Cat}(\bm{x}_{t};\bm{x}_{t-1}\bm{Q}_{t})$, where the transition matrix is chosen as $\bm{Q}_{t}\triangleq\alpha_{t}\bm{I}+(1-\alpha_{t})(\bm{1}_{a}\bm{1}_{a}^{\top})/a$, with $\alpha_{t}$ transitioning from $1$ to $0$ as $t$ increases [2]. It then follows that $q(\bm{x}_{t}|\bm{x}_{0})=\operatorname{Cat}(\bm{x}_{t};\bm{x}_{0}\bar{\bm{Q}}_{t})$ and $q(\bm{x}_{t-1}|\bm{x}_{t},\bm{x}_{0})\!=\!\operatorname{Cat}(\bm{x}_{t-1};\frac{\bm{x}_{t}\bm{Q}_{t}^{\top}\odot~{}\bm{x}_{0}\bar{\bm{Q}}_{t-1}}{\bm{x}_{0}\bar{\bm{Q}}_{t}\bm{x}_{t}^{\top}})$, where $\bar{\bm{Q}}_{t}\triangleq\bm{Q}_{1}\bm{Q}_{2}\cdots\bm{Q}_{t}$ and $\odot$ denotes element-wise product. The design choice of $\bm{Q}_{t}$ ensures that $q(\bm{x}_{T}|\bm{x}_{0})\approx\operatorname{Cat}(\bm{x}_{T};\!\bm{1}_{a}/a)$, i.e., a uniform distribution over $\mathcal{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_{\theta}(\bm{G}_{0:T})=p(\bm{G}_{T})\prod_{t=1}^{T}p_{\theta}(\bm{G}_{t-1}|\bm{G}_{t})$ is then learned to recover the data distribution from $p(\bm{G}_{T})\approx q(\bm{G}_{T})$ (an approximate uniform distribution). The reverse transition $p_{\theta}(\bm{G}_{t-1}|\bm{G}_{t})$ is a product of categorical distributions over nodes and edges, denoted as $p_{\theta}(\bm{x}_{t-1}|\bm{G}_{t})$ and $p_{\theta}(\bm{e}_{t-1}|\bm{G}_{t})$. Notably, in line with the $\bm{x}_{0}$-parameterization used in continuous DPMs [25, 32], $p_{\theta}(\bm{x}_{t-1}|\bm{G}_{t})$ is modeled as:

where $p_{\theta}(\widetilde{\bm{x}}_{0}|\bm{G}_{t})$ is a neural network predicting the posterior probability of $\bm{x}_{0}$ given a noisy graph $\bm{G}_{t}$. For edges, each definition is analogous and thus omitted.

The model is learned with a graph dataset $\mathcal{D}$ by maximizing the following objective [61]:

where $\bm{G}_{0}$ and $t$ follow uniform distributions over $\mathcal{D}$ and $[\![1,T]\!]$, respectively. After learning, graph samples can then be generated by first sampling $\bm{G}_{T}$ from $p(\bm{G}_{T})$ and subsequently sampling $\bm{G}_{t}$ from $p_{\theta}(\bm{G}_{t-1}|\bm{G}_{t})$, resulting in a generation trajectory $(\bm{G}_{T},\bm{G}_{T-1},\dots,\bm{G}_{0})$.

Markov decision processes (MDPs) are commonly used to model sequential decision-making problems [17]. An MDP is formally defined by a quintuple $(\mathcal{S},\mathcal{A},P,r,\rho_{0})$, where $\mathcal{S}$ is the state space containing all possible environment states, $\mathcal{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 $\rho_{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 $\bm{s}_{t}\in\mathcal{S}$ and selects an action $\bm{a}_{t}\in\mathcal{A}$ based on its policy distribution $\pi_{\theta}(\bm{a}_{t}|\bm{s}_{t})$. Subsequently, the agent receives a reward $r(\bm{s}_{t},\bm{a}_{t})$ and transitions to a new state $\bm{s}_{t+1}$ following the transition function $P(\bm{s}_{t+1}|\bm{s}_{t},\bm{a}_{t})$. As the agent interacts in the MDP (starting from an initial state $\bm{s}_{0}\sim\rho_{0}$), it generates a trajectory (i.e., a sequence of states and actions) denoted as $\bm{\tau}=(\bm{s}_{0},\bm{a}_{0},\bm{s}_{1},\bm{a}_{1},\dots,\bm{s}_{T},\bm{a}_{T})$. The cumulative reward over a trajectory $\bm{\tau}$ is given by $R(\bm{\tau})=\sum_{t=0}^{T}r(\bm{s}_{t},\bm{a}_{t})$. In most scenarios, the objective is to maximize the following expectation:

Policy gradient methods aim to estimate $\nabla_{\theta}\mathcal{J}_{\textsc{RL}}(\theta)$ and thus solve the problem by gradient descent. An important result is the policy gradient theorem [19], which estimates $\nabla_{\theta}\mathcal{J}_{\textsc{RL}}(\theta)$ as follows:

The 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.

A graph DPM defines a sample distribution $p_{\theta}(\bm{G}_{0})$ through its reverse denoising process $p_{\theta}(\bm{G}_{0:T})$. Given a reward signal $r(\cdot)$ for $\bm{G}_{0}$, we aim to maximize the expected reward (ER) over $p_{\theta}(\bm{G}_{0})$:

However, directly optimizing $\mathcal{J}_{\textsc{ER}}(\theta)$ is challenging since the likelihood $p_{\theta}(\bm{G}_{0})$ is unavailable [25] and $r(\cdot)$ is black-box, hindering the use of typical RL algorithms [6]. Following Fan et al. [16], we 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:

where the initial state $\bm{s}_{0}$ corresponds to the initial noisy graph $\bm{G}_{T}$ and the policy corresponds to the reverse transition distribution. As a result, the graph generation trajectory $(\bm{G}_{T},\bm{G}_{T-1},\dots,\bm{G}_{0})$ can be considered as a state-action trajectory $\bm{\tau}$ produced by an agent acting in the MDP. It then follows that $p(\bm{\tau}|\pi_{\theta})=p_{\theta}(\bm{G}_{0:T})$.²²2With a slight abuse of notation we will use $\bm{\tau}=\bm{G}_{0:T}$ and $\bm{\tau}=(\bm{s}_{0},\bm{a}_{0},\bm{s}_{1},\bm{a}_{1},\dots,\bm{s}_{T},\bm{a}_{T})$ interchangeably, which should not confuse as the MDP relates them with a bijection. Moreover, we have $R(\bm{\tau})=\sum_{t=0}^{T}r(\bm{s}_{t},\bm{a}_{t})=r(\bm{G}_{0})$. Therefore, the expected cumulative reward of the agent $\mathcal{J}_{\textsc{RL}}(\theta)=\mathbb{E}_{p(\bm{\tau}|\pi_{\theta})}[R(\bm{\tau})]=\mathbb{E}_{p_{\theta}(\bm{G}_{0:T})}[r(\bm{G}_{0})]$ is equivalent to $\mathcal{J}_{\textsc{ER}}(\theta)$, and thus $\mathcal{J}_{\textsc{ER}}(\theta)$ can also be optimized with the policy gradient $\nabla_{\theta}\mathcal{J}_{\textsc{RL}}(\theta)$:

where the generation trajectory $\bm{\tau}$ follows the parametric reverse process $p_{\theta}(\bm{G}_{0:T})$.

The policy gradient $\nabla_{\theta}\mathcal{J}_{\textsc{RL}}(\theta)$ 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 $\nabla_{\theta}\mathcal{J}_{\textsc{RL}}(\theta)$ 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:

where $\{\bm{G}_{0:T}^{(k)}\}_{k=1}^{K}$ are $K$ trajectories sampled from $p_{\theta}(\bm{G}_{0:T})$ and $\{\mathcal{T}_{k}\!\subset\![\![1,T]\!]\}_{k=1}^{K}$ are 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 $\bm{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 $\nabla_{\theta}\log p_{\theta}(\bm{G}_{t-1}|\bm{G}_{t})$. With the parameterization Eq. (1) described in Sec. 3.1, it has the following form:

where we can view the “weight” term as a weight assigned to the gradient $\nabla_{\theta}p_{\theta}(\widetilde{\bm{G}}_{0}|\bm{G}_{t})$, and thus $\nabla_{\theta}\log p_{\theta}(\bm{G}_{t-1}|\bm{G}_{t})$ as a weighted sum of such gradients, with $\widetilde{\bm{G}}_{0}$ taken over all possible graphs. Intuitively, the gradient $\nabla_{\theta}p_{\theta}(\widetilde{\bm{G}}_{0}|\bm{G}_{t})$ promotes the probability of predicting $\widetilde{\bm{G}}_{0}$ from $\bm{G}_{t}$. Note, however, that the weight $q(\bm{G}_{t-1}|\bm{G}_{t},\widetilde{\bm{G}}_{0})$ is completely independent of $r(\widetilde{\bm{G}}_{0})$ and could assign large weight for $\widetilde{\bm{G}}_{0}$ that has low reward. Since the weighted sum in Eq. (9) can be dominated by gradient terms with large $q(\bm{G}_{t-1}|\bm{G}_{t},\widetilde{\bm{G}}_{0})$, given a particular sampled trajectory, it is fairly possible that $\nabla_{\theta}\log p_{\theta}(\bm{G}_{t-1}|\bm{G}_{t})$ increases the probabilities of predicting undesired $\widetilde{\bm{G}}_{0}$ with low rewards from $\bm{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 $|\mathcal{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 $\bm{g}(\theta)$:

which 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 $\bm{G}_{0}$, where trajectories with the same $\bm{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 $\nabla_{\theta}\log p_{\theta}(\bm{G}_{t-1}|\bm{G}_{t})$ with $\nabla_{\theta}\log p_{\theta}(\bm{G}_{0}|\bm{G}_{t})$ in Eq. (8), we skip the weighted sum in Eq. (9) and directly promotes the probability of predicting $\bm{G}_{0}$ which has higher reward from $\bm{G}_{t}$ at all timestep $t$. As a result, our estimation does not focus on how $\bm{G}_{t}$ changes to $\bm{G}_{t-1}$ within the trajectory; instead, it aims to force the model’s generated results to be close to the desired $\bm{G}_{0}$, which can be seen as optimizing in equivalence classes. While being a biased estimator of the policy gradient $\nabla_{\theta}\mathcal{J}_{\textsc{RL}}(\theta)$, 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).

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_{\mathrm{val}}(\cdot)$ is then formulated as binary, with $r_{\mathrm{val}}(\bm{G}_{0})\triangleq 1$ indicating that the generated graph $\bm{G}_{0}$ conforms to the specified type; otherwise, $r_{\mathrm{val}}(\bm{G}_{0})\triangleq 0$.

Similarity. In certain scenarios, the objective is to generate graphs that resemble a known set of graphs $\mathcal{D}$ at the distribution level, based on a pre-defined distance metric $d(\cdot,\cdot)$ between sets of graphs. As an example, the $\textit{Deg}(\mathcal{G},\mathcal{D})$ [38] measures the maximum mean discrepancy (MMD) [18] between the degree distributions of a set $\mathcal{G}$ of generated graphs and the given graphs $\mathcal{D}$. Since our method requires a reward for each single generated graph $\bm{G}_{0}$, we simply adopt $\textit{Deg}(\{\bm{G}_{0}\},\mathcal{D})$ as the signal. As the magnitude of reward is critical for policy gradients [58], we define $r_{\mathrm{deg}}(\bm{G}_{0})\triangleq\exp(-\textit{Deg}(\{\bm{G}_{0}\},\mathcal{D})^{2}/\sigma^{2})$, where the $\sigma$ controls the reward distribution, ensuring that the reward lies within the range of $0$ to $1$. The other two similar distance metrics are $\textit{Clus}(\mathcal{G},\mathcal{D})$ and $\textit{Orb}(\mathcal{G},\mathcal{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_{\mathrm{deg}}$, namely, $r_{\mathrm{clus}}$ and $r_{\mathrm{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(\cdot,\cdot)$, measures the chemical similarity between two molecules, defined by the Jaccard index of molecule fingerprint bits. Specifically, $J\in[0,1]$, and $J(\bm{G}_{0},\bm{G}^{\prime}_{0})=1$ indicates that two molecules $\bm{G}_{0}$ and $\bm{G}^{\prime}_{0}$ have identical fingerprints. Following Xie et al. [65], we define the novelty of a generated graph $\bm{G}_{0}$ as $\textit{NOV}(\bm{G}_{0})\triangleq 1-\max_{\bm{G}^{\prime}_{0}\in\mathcal{D}}J(\bm{G}_{0},\bm{G}^{\prime}_{0})$, i.e., the similarity gap between $\bm{G}_{0}$ and its nearest neighbor in the training dataset $\mathcal{D}$, and further define $r_{\textsc{nov}}(\bm{G}_{0})\triangleq\textit{NOV}(\bm{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 $\textit{DS}(\cdot)$, 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 as $r_{\textsc{ds}}(\bm{G}_{0})\triangleq-\textit{DS}(\bm{G}_{0})/20$.

Another metric is the quantitative estimate of drug-likeness $\textit{QED}(\cdot)$, 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_{\textsc{qed}}(\bm{G}_{0})\triangleq\mathbb{I}[\textit{QED}(\bm{G}_{0})>0.5]$.

Synthetic Accessibility. The synthetic accessibility [7] $\textit{SA}(\cdot)$ 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_{\textsc{sa}}(\bm{G}_{0})\triangleq(10-\textit{SA}(\bm{G}_{0}))/9$.

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).

Implementation. We set $T=1000$, $|\mathcal{T}|=200$, and $N=100$. The number of trajectory samples $K$ is $64$ for SBM and $256$ for Planar. We use a DiGress model with $10$ layers. More implementation details can be found in Appendix A.1.

Metrics and Reward Functions. We consider four metrics: $\textit{Deg}(\mathcal{G},\mathcal{D}_{test})$, $\textit{Clus}(\mathcal{G},\mathcal{D}_{test})$, $\textit{Orb}(\mathcal{G},\mathcal{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:

where we do not explicitly incorporate uniqueness and novelty. All rewards are calculated on the training dataset if a reference graph set is required. All evaluation metrics are calculated on the test dataset. More details about baselines, reward signals, and metrics are in Appendix A.3.

Results. Table 1 summarizes GDPO’s superior performance in general graph generation, showing notable improvements in Deg and V.U.N across both SBM and Planar datasets. On the Planar dataset, GDPO significantly reduces distribution distance, achieving an $\mathbf{81.97\%}$ average decrease in metrics of Deg, Clus, and Orb compared to DiGress (the best baseline method). For the SBM dataset, GDPO has a $\mathbf{41.64\%}$ average improvement. The low distributional distances to the test dataset suggests that GDPO accurately captures the data distribution with well-designed rewards. Moreover, we observe that our method outperforms DDPO by a large margin, primarily because the graphs in Planar and SBM contain too many nodes, which aligns with the observation in Fig. 2.

Datasets and Baselines. Molecule property optimization aims to generate molecules with desired properties. We evaluate our method on two large molecule datasets: ZINC250k [27] and MOSES [49]. The ZINC250k dataset comprises 249,456 molecules, each containing 9 types of atoms, with a maximum node count of 38; the MOSES dataset consists of 1,584,663 molecules, with 8 types of atoms and a maximum node count of 30. We compare GDPO with several leading methods: 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.

Implementation. We set $T=500$, $|\mathcal{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:

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_{\textsc{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_{\textsc{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_{\textsc{NOV}}$ can lead the model to rapidly degenerate. One potential solution is to gradually increase the weight and conduct multi-stage optimization.