Stabilizing Policy Gradients for Stochastic Differential Equations via Consistency with Perturbation Process
This work addresses stability issues in training SDEs with policy gradients, which is important for applications like drug design, though it is incremental as it builds on existing methods with a novel constraint.
The paper tackles the instability of policy gradients when training stochastic differential equations (SDEs) for generative modeling by constraining the SDE to be consistent with its perturbation process, which improves sample complexity and control. It achieves a Vina score of -9.07 on the CrossDocked2020 dataset for drug design, outperforming other methods.
Considering generating samples with high rewards, we focus on optimizing deep neural networks parameterized stochastic differential equations (SDEs), the advanced generative models with high expressiveness, with policy gradient, the leading algorithm in reinforcement learning. Nevertheless, when applying policy gradients to SDEs, since the policy gradient is estimated on a finite set of trajectories, it can be ill-defined, and the policy behavior in data-scarce regions may be uncontrolled. This challenge compromises the stability of policy gradients and negatively impacts sample complexity. To address these issues, we propose constraining the SDE to be consistent with its associated perturbation process. Since the perturbation process covers the entire space and is easy to sample, we can mitigate the aforementioned problems. Our framework offers a general approach allowing for a versatile selection of policy gradient methods to effectively and efficiently train SDEs. We evaluate our algorithm on the task of structure-based drug design and optimize the binding affinity of generated ligand molecules. Our method achieves the best Vina score -9.07 on the CrossDocked2020 dataset.