Research map / Probabilistic and scientific modeling / Variational and generative sampling
Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference
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 trust-region constrained iterative diffusion optimization for stochastic optimal control, which yields a geometric annealing between prior and target path measures with equidistant Fisher-Rao steps, and applies it to diffusion-based sampling, transition path sampling, and diffusion model fine-tuning.
Problem
Iterative diffusion optimization for stochastic optimal control struggles with high variance and mode collapse when the target path measure differs substantially from the uncontrolled prior.
Approach
Solve a sequence of KL-trust-region constrained SOC problems around the previous control iterate; the Lagrangian becomes another SOC problem and the optimal multiplier gives a geometric annealing whose steps are equispaced in Fisher-Rao distance, with buffer-based off-policy losses (log-variance, SOC matching).
Key results
- On Alanine Dipeptide, TR-LV reaches RMSD 0.29 ± 0.03 Å and THP 61.25 ± 4.05%, versus 0.47 Å and 39.58% for TPS-DPS.
- On Chignolin, TR-LV improves THP to 43.95 ± 5.64% from 25.00 ± 10.69% for TPS-DPS.
- On GMM targets, non-trust-region methods suffer mode collapse above d=10, while trust region methods remain robust up to d≈150; TR-SOCM matches adjoint matching fine-tuning performance at a fraction of the cost.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add trust region updates to the log-variance loss | TPS-DPS · Alanine Dipeptide | RMSD: 0.47 ± 0.18 -> 0.29 ± 0.03 | Helped |
| Add trust region updates to the log-variance loss | TPS-DPS · Alanine Dipeptide | THP: 39.58 ± 28.13 -> 61.25 ± 4.05 | Helped |
| Add trust region updates to the log-variance loss | TPS-DPS · Chignolin | RMSD: 1.06 ± 0.08 -> 0.90 ± 0.01 | Helped |
| Add trust region updates to the log-variance loss | TPS-DPS · Chignolin | THP: 25.00 ± 10.69 -> 43.95 ± 5.64 | Helped |
| Use trust region updates instead of unconstrained updates | SOC methods on GMM sampling · GMM, dimensions beyond d = 10 | control error | Helped |
| Use trust region updates | Diffusion-based sampling methods · Many Well, 32 modes | absolute log-normalization error | Helped |
| Add replay-buffer training passes | TR-SOCM fine-tuning · Stable Diffusion 1.5 reward fine-tuning | trajectories required for a fixed number of model updates: three times fewer trajectories | Helped |
| Add trust region optimization rather than using a fixed, tuned Lagrangian multiplier (ablation) | TR-LV on GMM sampling · GMM | control error | Helped |
| Add trust region optimization rather than using a fixed, tuned Lagrangian multiplier (ablation) | TR-LV on GMM sampling · GMM | target evaluations: several orders of magnitude | Helped |
Surprising findings
- Trust region optimization reduces control error and cuts the number of target evaluations by several orders of magnitude compared to Langevin preconditioning on GMM targets.
Limitations
- Replay buffer stores full trajectories, requiring CPU-GPU transfers in large-scale settings such as text-to-image fine-tuning.
- Analysis restricted to KL trust regions; other divergences left for future work.
- Performance degrades only at high dimensions (d >= 150) on GMM, where mode collapse begins to appear.
Experimental setup
Datasets: Gaussian Mixture Model targets Many Well target Alanine Dipeptide MD Chignolin MD text prompts for Stable Diffusion fine-tuning
Benchmarks: GMM sampling Many Well Alanine Dipeptide transition path sampling Chignolin transition path sampling ImageReward CLIP-Score Human Preference Score DreamSim diversity
Metrics: L2 control error Sinkhorn distance |Delta log Z| total variation RMSD transition hit percentage energy of transition state ImageReward CLIP-Score HPS
Baselines: relative entropy cross entropy log-variance loss adjoint matching SOCM UMD SMD PIPS TPS-DPS
Builds on: Denoising Diffusion Sampler adjoint matching log-variance loss stochastic optimal control matching TPS-DPS
Closest papers by meaning
- Stochastic Optimal Control Matching (NeurIPS 2024 · similarity 0.86)
- Regional Expected Improvement for Efficient Trust Region Selection in High-Dimensional Bayesian Optimization (AAAI 2025 · similarity 0.86)
- Stochastic Optimal Control for Diffusion Bridges in Function Spaces (NeurIPS 2024 · similarity 0.85)
- Stochastic Optimization in Semi-Discrete Optimal Transport: Convergence Analysis and Minimax Rate (NeurIPS 2025 · similarity 0.85)
- Constrained Diffusion with Trust Sampling (NeurIPS 2024 · similarity 0.85)
- Trust Region Methods for Nonconvex Stochastic Optimization beyond Lipschitz Smoothness (AAAI 2024 · similarity 0.85)
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.