Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference

2025-08-17

AI Summary

In stochastic optimal control and inference, approximating path-space measures becomes challenging when the target and prior measures differ substantially, leading to gradient mismatch and poor convergence in conventional optimization. To address this, we propose a trust-region-based iterative constrained optimization framework—the first to incorporate trust-region mechanisms into path-space measure transport—enabling geometrically annealed, progressive approximation. Explicit feasibility constraints ensure stable updates at each iteration, while a systematic annealing step-size selection rule is introduced. Our method unifies gradient-based optimization, optimal transport theory, and robust constraint principles. Experiments on diffusion sampling, transition path simulation, and diffusion model fine-tuning demonstrate significant improvements in training stability and convergence speed. The approach exhibits strong generalization across diverse tasks, validating both its effectiveness and broad applicability.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Search and Optimization: Sampling/Simulation-based Search
• Constraint Satisfaction and Optimization: Constraint Learning and Acquisition

Application Category

• Security and Privacy: Large-scale security measurements
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Solving stochastic optimal control problems with quadratic control costs can be viewed as approximating a target path space measure, e.g. via gradient-based optimization. In practice, however, this optimization is challenging in particular if the target measure differs substantially from the prior. In this work, we therefore approach the problem by iteratively solving constrained problems incorporating trust regions that aim for approaching the target measure gradually in a systematic way. It turns out that this trust region based strategy can be understood as a geometric annealing from the prior to the target measure, where, however, the incorporated trust regions lead to a principled and educated way of choosing the time steps in the annealing path. We demonstrate in multiple optimal control applications that our novel method can improve performance significantly, including tasks in diffusion-based sampling, transition path sampling, and fine-tuning of diffusion models.

Problem

• Solving stochastic optimal control with trust regions
• Approaching target measure via geometric annealing
• Improving performance in diffusion-based tasks

Innovation

• Trust region constrained measure transport
• Geometric annealing with trust regions
• Gradient-based optimization in path space