Paper: Stabilizing Policy Gradients for Stochastic Differential Equations by enforcing Consistency with Perturbation Process (ICLR-2024)

Reviewer comment (Presentation):
The authors need to make the equations and their associated explanations easy to read for general readers. The organization of the paper can be improved to increase clarity.

Author response:
A3: We will carefully polish the manuscript according to your suggestions and add a more detailed discussion for RL, SDEs, and diffusion models in appendix in the next revision.
 We here provide a more clear background and better notations of SDEs as follows:  


- (Equation 5) $dx = f(x,t) dt + g(t) dw$,
 which is the forward SDE that corresponds to the perturbation process (i.e., the forward process of diffusion models).  $dw$ is a Wiener process.  $f(\cdot,t):\mathbb{R}^d\to\mathbb{R}^d$ is a vector-valued function called the drift coefficient of $x(t)$.  $g(t)$ is a scalar function of time and known as diffusion coefficient of the underlying dynamic.
 - Following [3], applying Kolmogorov’s forward equation (Fokker-Planck equation) to an SDE, we can derive a probability ODE which describes how the marginal distribution ${p\_t}(x)$ evolves through time $t$. 
- (Equation 6) According to [4], an SDE as Equation 5 has a backward SDE: $dx= (f(x,t)- g^2(t)\nabla\_x\log p\_t(x))dt+g(t) d\bar{w}$, which shares the same marginal distribution $p_t(x)$ at time $t$. $d\bar{w}$ is the reverse Wiener process. 
- (Equation 1) $dx\_t=\pi\_\theta(x\_t,\theta)dt+g(t)d\bar{w}$, which is the approximated backward SDE parameterized by $\theta$. Conventionally, we use $\epsilon\_\theta(x\_t,t)$ to approximate the score $\nabla\_x\log p\_t(x)$. So we can let $\pi\_\theta(x\_t, \theta):=f(x\_t, t) - g^2(t)\epsilon\_\theta(x\_t, \theta)$. The approximated backward SDE corresponds to the generative process of diffusion models. Specifically, $x\_T$ is sampled from the prior distribution, evolves following this SDE, and arrives at $x\_0$. $x\_0$ is the generated sample.
- $p\_0$ and $q\_0$ in Definition 1: You can find their definitions in section 2.2. Specifically, We denote the marginal distribution at time $t$ of Equation 5 and 1 as $p\_t$ and $q\_t$, respectively. Thus, $p\_0$ (resp. $q\_0$) is $p\_t$ (resp. $q\_t$) at time $t=0$.

Reviewer comment (Writing):
In equation 1, what is sigma_t?

Author response:
A3: We will carefully polish the manuscript according to your suggestions and add a more detailed discussion for RL, SDEs, and diffusion models in appendix in the next revision.
 We here provide a more clear background and better notations of SDEs as follows:  


- (Equation 5) $dx = f(x,t) dt + g(t) dw$,
 which is the forward SDE that corresponds to the perturbation process (i.e., the forward process of diffusion models).  $dw$ is a Wiener process.  $f(\cdot,t):\mathbb{R}^d\to\mathbb{R}^d$ is a vector-valued function called the drift coefficient of $x(t)$.  $g(t)$ is a scalar function of time and known as diffusion coefficient of the underlying dynamic.
 - Following [3], applying Kolmogorov’s forward equation (Fokker-Planck equation) to an SDE, we can derive a probability ODE which describes how the marginal distribution ${p\_t}(x)$ evolves through time $t$. 
- (Equation 6) According to [4], an SDE as Equation 5 has a backward SDE: $dx= (f(x,t)- g^2(t)\nabla\_x\log p\_t(x))dt+g(t) d\bar{w}$, which shares the same marginal distribution $p_t(x)$ at time $t$. $d\bar{w}$ is the reverse Wiener process. 
- (Equation 1) $dx\_t=\pi\_\theta(x\_t,\theta)dt+g(t)d\bar{w}$, which is the approximated backward SDE parameterized by $\theta$. Conventionally, we use $\epsilon\_\theta(x\_t,t)$ to approximate the score $\nabla\_x\log p\_t(x)$. So we can let $\pi\_\theta(x\_t, \theta):=f(x\_t, t) - g^2(t)\epsilon\_\theta(x\_t, \theta)$. The approximated backward SDE corresponds to the generative process of diffusion models. Specifically, $x\_T$ is sampled from the prior distribution, evolves following this SDE, and arrives at $x\_0$. $x\_0$ is the generated sample.
- $p\_0$ and $q\_0$ in Definition 1: You can find their definitions in section 2.2. Specifically, We denote the marginal distribution at time $t$ of Equation 5 and 1 as $p\_t$ and $q\_t$, respectively. Thus, $p\_0$ (resp. $q\_0$) is $p\_t$ (resp. $q\_t$) at time $t=0$.

Reviewer comment (Theory):
Why is it natural to train SDEs with RL? What if the training data only consist of samples with desired characteristics? Do we still need RL in that case? The authors cited relevant work, but the basic idea/counterexample should be mentioned here (in a sentence).

Author response:
A2: Thanks for your comments. To address your concerns, let's take the molecular optimization problem as an example, where we want to find molecules with high binding affinity given a protein pocket.  Firstly, in many practical scenarios, the training set does not consist of data with desired characteristics. For the molecular optimization problem, TargetDiff is trained on protein-molecule pairs with Vina scores mainly ranging from -5 to -7. And with RL, we can significantly improve the average Vina score of generated molecules to -9.07. Secondly, we want to argue that even if the training data only consist of samples with desired properties, we still need RL.  Suppose we're testing TargetDiff on new proteins, the performance may be unsatisfactory (please check Figure 4), especially when the structure of the test data is different from those in the training set. In such cases, we need to exploit RL to further optimize the structure of the molecules.

Reviewer comment (Writing):
In equation 5, where is p_t as mentioned afterwards?

Author response:
A3: We will carefully polish the manuscript according to your suggestions and add a more detailed discussion for RL, SDEs, and diffusion models in appendix in the next revision.
 We here provide a more clear background and better notations of SDEs as follows:  


- (Equation 5) $dx = f(x,t) dt + g(t) dw$,
 which is the forward SDE that corresponds to the perturbation process (i.e., the forward process of diffusion models).  $dw$ is a Wiener process.  $f(\cdot,t):\mathbb{R}^d\to\mathbb{R}^d$ is a vector-valued function called the drift coefficient of $x(t)$.  $g(t)$ is a scalar function of time and known as diffusion coefficient of the underlying dynamic.
 - Following [3], applying Kolmogorov’s forward equation (Fokker-Planck equation) to an SDE, we can derive a probability ODE which describes how the marginal distribution ${p\_t}(x)$ evolves through time $t$. 
- (Equation 6) According to [4], an SDE as Equation 5 has a backward SDE: $dx= (f(x,t)- g^2(t)\nabla\_x\log p\_t(x))dt+g(t) d\bar{w}$, which shares the same marginal distribution $p_t(x)$ at time $t$. $d\bar{w}$ is the reverse Wiener process. 
- (Equation 1) $dx\_t=\pi\_\theta(x\_t,\theta)dt+g(t)d\bar{w}$, which is the approximated backward SDE parameterized by $\theta$. Conventionally, we use $\epsilon\_\theta(x\_t,t)$ to approximate the score $\nabla\_x\log p\_t(x)$. So we can let $\pi\_\theta(x\_t, \theta):=f(x\_t, t) - g^2(t)\epsilon\_\theta(x\_t, \theta)$. The approximated backward SDE corresponds to the generative process of diffusion models. Specifically, $x\_T$ is sampled from the prior distribution, evolves following this SDE, and arrives at $x\_0$. $x\_0$ is the generated sample.
- $p\_0$ and $q\_0$ in Definition 1: You can find their definitions in section 2.2. Specifically, We denote the marginal distribution at time $t$ of Equation 5 and 1 as $p\_t$ and $q\_t$, respectively. Thus, $p\_0$ (resp. $q\_0$) is $p\_t$ (resp. $q\_t$) at time $t=0$.

Reviewer comment (Presentation):
Figure 1 and 2 need to be further up in the paper, as they are the motivating examples for the authors’ rationale of the paper.

Author response:
A5: Thanks for your valuable suggestions. Because time is limited, we have not yet had time to make adjustments. In the next revision, we will advance these two pictures as you suggested.

Reviewer comment (Presentation):
In section 2.3, it is better to state that the MDP is mapping to Equation 5 rather than Equation 1.

Author response:
A4: Thanks for the suggestions. However, the forward process in Equation 5 is not an MDP as there is no action (please refer to Q4 for reviewer 4LjP for a more intuitive explanation of forward process). In Equation 1, we view $dx\_{t}=\pi\_\theta(x\_t,\theta)dt+g(t)d\bar{w}$  as a MDP by taking $x\_t$ as state, $\pi\_\theta(x\_t,\theta)$ as the selected action, $x\_{t-dt}$ is the next state and the equation 1 itself stands for the transition function.

Reviewer comment (Writing):
In definition 1, what is p0, q0? There are no such symbols in Equation 1.

Author response:
A3: We will carefully polish the manuscript according to your suggestions and add a more detailed discussion for RL, SDEs, and diffusion models in appendix in the next revision.
 We here provide a more clear background and better notations of SDEs as follows:  


- (Equation 5) $dx = f(x,t) dt + g(t) dw$,
 which is the forward SDE that corresponds to the perturbation process (i.e., the forward process of diffusion models).  $dw$ is a Wiener process.  $f(\cdot,t):\mathbb{R}^d\to\mathbb{R}^d$ is a vector-valued function called the drift coefficient of $x(t)$.  $g(t)$ is a scalar function of time and known as diffusion coefficient of the underlying dynamic.
 - Following [3], applying Kolmogorov’s forward equation (Fokker-Planck equation) to an SDE, we can derive a probability ODE which describes how the marginal distribution ${p\_t}(x)$ evolves through time $t$. 
- (Equation 6) According to [4], an SDE as Equation 5 has a backward SDE: $dx= (f(x,t)- g^2(t)\nabla\_x\log p\_t(x))dt+g(t) d\bar{w}$, which shares the same marginal distribution $p_t(x)$ at time $t$. $d\bar{w}$ is the reverse Wiener process. 
- (Equation 1) $dx\_t=\pi\_\theta(x\_t,\theta)dt+g(t)d\bar{w}$, which is the approximated backward SDE parameterized by $\theta$. Conventionally, we use $\epsilon\_\theta(x\_t,t)$ to approximate the score $\nabla\_x\log p\_t(x)$. So we can let $\pi\_\theta(x\_t, \theta):=f(x\_t, t) - g^2(t)\epsilon\_\theta(x\_t, \theta)$. The approximated backward SDE corresponds to the generative process of diffusion models. Specifically, $x\_T$ is sampled from the prior distribution, evolves following this SDE, and arrives at $x\_0$. $x\_0$ is the generated sample.
- $p\_0$ and $q\_0$ in Definition 1: You can find their definitions in section 2.2. Specifically, We denote the marginal distribution at time $t$ of Equation 5 and 1 as $p\_t$ and $q\_t$, respectively. Thus, $p\_0$ (resp. $q\_0$) is $p\_t$ (resp. $q\_t$) at time $t=0$.

Reviewer comment (Theory):
It was not clear why reinforcement learning can effectively respect the boundary conditions of the underlying problem while supervised learning approach fail to do so beyond the numerical results.

Author response:
A1: Thanks for your suggestions. We want to clarify that we have already included the desired properties in the objective, as the goal of reinforcement learning is to efficiently and reliably optimize properties through interactions. And policy gradient-based methods are the SOTA of this research field in various applications, please refer to [1] for a comprehensive survey about RL and policy gradients. Therefore, the algorithm that encodes both score-matching and the desired properties as you mentioned, exactly falls into the class of reinforcement learning. Our contribution is a novel policy gradient-estimator for SDE-based policy. With our policy gradient estimator, the optimization process becomes more stable and more efficient.

Reviewer comment (Novelty):
How would an alternative method fail, for example, one that encodes both model matching score and the desired properties into the objective function and then solve the corresponding supervised learning problem, was not articulated.

Author response:
Thanks for your response. 

What we are claiming is the general-purpose algorithm which aims to maximize a reward through interactions falls into the class of RL. In addition to the survey [1] we have cited, here we directly quote the first sentence of Section 1.1 from the textbook [2]: "Reinforcement learning is learning what to do—how to map situations to actions—so as to maximize a numerical reward signal."  Therefore,  we thought you are proposing the objective of RL by mentioning "the desired properties into the objective function". 

However, we realized that you might be talking about a very specific algorithm about "algorithms with boundary constraints" and "supervised learning with boundary constraints". But to the best of our knowledge, we are not aware of such an algorithm in the research fields of RL, diffusion models, and SBDD. Therefore, in order to make sure we are on the same page, would you mind providing some references of the algorithm you are talking about? Only with specific literature, we can compare our method against them if they are relevant.

Moreover, we speculate that the "boundary constraint" that you mentioned might be some kind of regularization that penalizes the coordinates of the atoms of a molecule that has some physical conflicts or collisions with the protein. If so, we need to emphasize again that the physical conflicts or collisions are only one of the factors that account for a bad Vina score. The Vina score is complex and far beyond the simple spatial relation. The final goal is to maximize the Vina score instead of only avoiding such physical conflicts, so only considering such a penalty cannot achieve the goal. 

More generally, our goal is to maximize a reward, i.e., a scalar function of a sample and we do not assume any prior knowledge (e.g., boundary constraint is correlated to Vina scores) about this reward in learning. Therefore, our algorithm could also be applied to more general applications. For example, please refer to our response to Reviewer 4LjP's Q1 and Appendix D in the revision for the experiments on text-to-image generation. 


Besides, we want to clarify that, for structure-based drug design, TargetDiff [3] is the SOTA method, and, for text-to-image generation, Stable Diffusion [4] is the SOTA one. Both are diffusion models (i.e., SDE-based generative models). Since the final goal is to generate samples with a certain desired property, we apply RL to fine-tune the SOTA SDE-based generative models. And our main contribution is a much more stable and efficient policy gradient method dedicated to this scenario.


References:

[1] Arulkumaran, Kai, et al. Deep reinforcement learning: A brief survey. IEEE Signal Processing Magazine (2017).

[2] Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press (2018).

[3] Guan, Jiaqi, et al. 3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction. ICLR (2023)

[4] Rombach, Robin, et al. High-resolution image synthesis with latent diffusion models. CVPR (2022).

Reviewer comment (Evaluation):
Is the superiority shown in Figure 3 (and also Figure 5 in the appendix) a statistically consistent behavior across majority of examples? Are the numbers shown in the table in Figure 4 possible to have the standard deviation annotations next to the mean and medium values?

Author response:
A2: 
Yes. In Figures 5, 6, 7 in Appendix C, we plot the optimization curves of generated ligand molecules over 100 different protein pockets.  As obviously observed, in most cases, our method outperforms all the baselines. 

More specifically, our method outperforms TargetDiff in 100 out of 100 cases, outperforms baselines in 75 out of 100 cases, and is at least comparable with the best baseline in almost all cases. 

Moreover, as you suggested, we have reported the standard deviation in Table 1 in the revision.

Reviewer comment (Evaluation):
It is surprising that from Figure 4 it seems that TargetDiff does not respect the boundary at all, which seems to be contradicting one of the contributions claimed by the TargetDiff authors. Are there some structural properties associated with these scenarios? Is this a consistent behavior in the majority of the cases?

Author response:
A3: Thank you for your valuable feedback. We want to assure you that the results produced by TargetDiff are both reliable and reproducible. In many protein instances, TargetDiff respects the boundary conditions. However, as depicted in Figure 4, there are some proteins for which TargetDiff is unable to generate high-quality ligands. This could be due to the structural differences between these proteins and those present in the training set. Our results show that we can further significantly improve the pretrained diffusion model by introducing RL. 

Besides, notably, violation of the boundary condition is only one of the factors for bad Vina scores (i.e., the proxy for binding affinity, and also the reward in our method). We take the boundary condition as an example for illustration because it is easy to understand even without any background knowledge about biology and chemistry.