Stochastic Optimal Control Sampling for Diffusion Inverse Problems
For practitioners solving image inverse problems with diffusion models, SOCS offers a computationally efficient and high-quality method that outperforms prior SOC-based approaches.
SOCS models diffusion denoising as a dynamical system and injects closed-form control signals at each step, achieving accurate measurement-aligned reconstructions with improved visual fidelity and stronger quantitative performance across diverse image inverse tasks.
Benefiting from the strong ability to capture data distributions, diffusion models have become powerful tools for solving image inverse problems. The key is to controllably steer the sampling trajectory toward the measurements while respecting the diffusion prior. In this work, we introduce Stochastic Optimal Control Sampling (SOCS), which models the denoising process as a dynamical system and injects control signals via SOC. Previous SOC-based approach addresses inverse problems by optimizing over the entire trajectory, which is computationally expensive. In contrast, we derive a closed-form control update and apply it at each sampling step, pulling the measurement-consistent clean prediction back onto the denoising flow. In SOCS, we can readily modulate the control strength to align with the diffusion model's native capabilities and thereby enhance perceptual quality. Our method is compatible with a variety of linear stochastic differential equation backbones. Extensive experiments across a broad spectrum of image inverse tasks demonstrate that SOCS achieves accurate measurement-aligned reconstructions with improved visual fidelity and stronger quantitative performance.