Research map / Reinforcement learning / Policy gradient and actor-critic
Stabilizing Policy Gradients for Stochastic Differential Equations via Consistency with Perturbation Process
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DiffAC, an actor-critic policy gradient method for SDE/diffusion policies that stabilizes training by enforcing consistency between the SDE policy and its forward perturbation process, evaluated on structure-based drug design (Vina score) and text-to-image fine-tuning.
Problem
Policy gradients for SDE-parameterized generative models are estimated from a finite set of trajectories, making them ill-defined in data-scarce regions and causing unstable training and poor sample complexity.
Approach
Constrain the backward SDE policy to be consistent with its forward perturbation process via score matching, so critic and policy-gradient samples can be generated cheaply by perturbing initial samples D_0 across the whole space; add a KL regularizer toward a score-matched reference policy to keep consistency during RL updates.
Key results
- DiffAC achieves the best average Vina score of -9.07 (best run) on CrossDocked2020, versus -6.53 for the best EEGSDE variant without online training and -8.58 for Online EEGSDE-0.1.
- On text-to-image with the prompt 'A green colored rabbit', DiffAC improves ImageReward scores and outperforms DPOK by a large margin (per curves; no numbers given in the text).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DiffAC instead of Online EEGSDE-0.1 | TargetDiff-based SBDD generation · CrossDocked2020 test set | Vina score (Best Run Avg.): -8.58 -> -9.07 | Helped |
| Use DiffAC instead of Online EEGSDE-0.1 | TargetDiff-based SBDD generation · CrossDocked2020 test set | Vina score (Last Run Avg.): -7.78 -> -8.50 | Helped |
| Add online critic training to EEGSDE | EEGSDE · CrossDocked2020 test set | Vina score (Best Run Avg.): -6.53 -> -8.58 | Helped |
| Use energy guidance coefficient 1 instead of 0.1 | EEGSDE · CrossDocked2020 test set | Vina score (Best Run Avg.): -6.53 -> -3.30 | Hurt |
| Use a too-small regularization coefficient (ablation) | DiffAC · CrossDocked2020, three protein targets | Vina score | Mixed |
| Use a too-large regularization coefficient (ablation) | DiffAC · CrossDocked2020, three protein targets | Vina score | Hurt |
| Use a proper regularization coefficient (ablation) | DiffAC · CrossDocked2020, three protein targets | Vina score | Helped |
| Estimate policy gradients using the forward perturbation process rather than real backward-process trajectories | DiffAC text-to-image fine-tuning · Text-to-image, prompt “A green colored rabbit” | ImageReward score | Helped |
Surprising findings
- Too small regularization coefficients perform well initially but fail after a few optimization steps, showing the consistency gap accumulates during RL updates.
- Online EEGSDE with the largest guidance coefficient (1) gives worse last-checkpoint Vina scores than the best run, indicating instability.
Limitations
- Consistency is only approximately enforced in practice via KL regularization, so the gradient estimate is biased; the bound depends on KL terms that are not measured end-to-end.
- Evaluation is limited to a single docking-based oracle (AutoDock Vina) and a single text prompt for text-to-image.
Experimental setup
Datasets: CrossDocked2020
Benchmarks: CrossDocked2020
Metrics: Vina score (mean) Vina score (median) ImageReward score
Baselines: AR Pocket2Mol TargetDiff EEGSDE Online EEGSDE DPOK Stable Diffusion
Builds on: TargetDiff DDPG REINFORCE DPOK DDPM
Closest papers by meaning
- An Efficient High-dimensional Gradient Estimator for Stochastic Differential Equations (NeurIPS 2024 · similarity 0.86)
- Flow Matching Policy Optimization with Mirror Descent and Entropy Constraints (NeurIPS 2026 · similarity 0.85)
- Learning Optimal Deterministic Policies with Stochastic Policy Gradients (ICML 2024 · similarity 0.84)
- PRECISE: SDE-Consistent Stochastic Sampling for RL Post-Training of Flow-Matching Models (NeurIPS 2026 · similarity 0.84)
- Convergence Analysis of Policy Gradient Methods with Dynamic Stochasticity (ICML 2025 · similarity 0.84)
- Efficient Training of Neural Stochastic Differential Equations by Matching Finite Dimensional Distributions (ICLR 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.