Solving Inverse Problems via Diffusion Optimal Control

TL;DR

This paper introduces a novel diffusion-based optimal control framework for inverse problems, overcoming key limitations of probabilistic sampling methods and enabling versatile, high-quality signal reconstruction across various tasks.

Contribution

The authors propose a diffusion optimal control approach inspired by iLQR, capable of handling any differentiable measurement operator and improving inverse problem solving.

Findings

Achieves superior image reconstruction quality.

Handles complex nonlinear inverse problems.

Establishes a new baseline in inverse problem solving.

Abstract

Existing approaches to diffusion-based inverse problem solvers frame the signal recovery task as a probabilistic sampling episode, where the solution is drawn from the desired posterior distribution. This framework suffers from several critical drawbacks, including the intractability of the conditional likelihood function, strict dependence on the score network approximation, and poor $\mathbf{x}_0$ prediction quality. We demonstrate that these limitations can be sidestepped by reframing the generative process as a discrete optimal control episode. We derive a diffusion-based optimal controller inspired by the iterative Linear Quadratic Regulator (iLQR) algorithm. This framework is fully general and able to handle any differentiable forward measurement operator, including super-resolution, inpainting, Gaussian deblurring, nonlinear deblurring, and even highly nonlinear neural classifiers. Furthermore, we show that the idealized posterior sampling equation can be recovered as a special case of our algorithm. We then evaluate our method against a selection of neural inverse problem solvers, and establish a new baseline in image reconstruction with inverse problems.