# Stochastic Optimal Control Sampling for Diffusion Inverse Problems

arXiv: https://arxiv.org/abs/2606.28785 · HTML: https://mlsanity.com/paper/2606.28785/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Inverse problems and image restoration, one of 14 directions in Diffusion models.
- Builds on 32 of its 51 references that are in the map.
- First posted on arXiv 27 Jun 2026.

## Placed in

- Diffusion models / [Inverse problems and image restoration](https://mlsanity.com/diffusion/inverse-problems-and-image-restoration/)

## Builds on (most cited first, 30 of 32)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Image Super-Resolution via Iterative Refinement](https://mlsanity.com/paper/2104.07636/) (2021)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Palette: Image-to-Image Diffusion Models](https://mlsanity.com/paper/2111.05826/) (2021)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [Denoising Diffusion Restoration Models](https://mlsanity.com/paper/2201.11793/) (2022)
- [Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model](https://mlsanity.com/paper/2212.00490/) (2022)
- [Improving Diffusion Models for Inverse Problems using Manifold Constraints](https://mlsanity.com/paper/2206.00941/) (2022)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [Denoising Diffusion Models for Plug-and-Play Image Restoration](https://mlsanity.com/paper/2305.08995/) (2023)
- [DiffIR: Efficient Diffusion Model for Image Restoration](https://mlsanity.com/paper/2303.09472/) (2023)
- [A Variational Perspective on Solving Inverse Problems with Diffusion Models](https://mlsanity.com/paper/2305.04391/) (2023)
- [Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing](https://mlsanity.com/paper/2407.01521/) (2024)
- [Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models](https://mlsanity.com/paper/2307.00619/) (2023)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors](https://mlsanity.com/paper/2405.18782/) (2024)
- [FlowDPS: Flow-Driven Posterior Sampling for Inverse Problems](https://mlsanity.com/paper/2503.08136/) (2025)
- [Explore Image Deblurring via Blur Kernel Space](https://mlsanity.com/paper/2104.00317/) (2021)
- [RB-Modulation: Training-Free Personalization of Diffusion Models using Stochastic Optimal Control](https://mlsanity.com/paper/2405.17401/) (2024)
- [Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance](https://mlsanity.com/paper/2402.02149/) (2024)
- [Steering Rectified Flow Models in the Vector Field for Controlled Image Generation](https://mlsanity.com/paper/2412.00100/) (2024)
- [Generative Modeling with Phase Stochastic Bridges](https://mlsanity.com/paper/2310.07805/) (2023)
- [UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control](https://mlsanity.com/paper/2502.05749/) (2025)

## Abstract (arXiv)

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.