Stochastic Optimal Control Sampling for Diffusion Inverse Problems

2026-06-27

AI Summary

This work addresses the challenge of balancing sampling efficiency and reconstruction quality in diffusion-based image inverse problems by modeling the denoising process as a dynamical system. The authors propose a closed-form, stepwise guidance mechanism grounded in stochastic optimal control, which explicitly injects a control signal at each sampling step to steer the clean prediction toward consistency with observed measurements. Unlike trajectory-level optimization approaches, this method operates incrementally, substantially reducing computational overhead. By adaptively tuning the control strength to align with the diffusion model’s prior capabilities, the approach achieves high measurement consistency while enhancing perceptual quality. Extensive experiments demonstrate that the proposed method consistently attains superior visual fidelity and quantitative performance across a range of image inverse tasks.

Abstract

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.

Problem

• diffusion inverse problems
• stochastic optimal control
• image reconstruction
• measurement consistency
• sampling trajectory

Innovation

• Stochastic Optimal Control
• Diffusion Models
• Inverse Problems
• Closed-form Control
• Denoising Flow