Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Solving Inverse Problems via Diffusion Optimal Control

Henry Li, Marcus Pereira

TL;DR

This work addresses the brittleness and intractability of diffusion-based inverse problem solvers that rely on conditional posterior sampling. By reframing the reverse diffusion as a discrete-time optimal-control episode and deriving a diffusion-based iLQR controller, the authors create a general solver that accommodates arbitrary differentiable forward operators and can recover the ideal posterior sampling equation as a special case. They introduce high-dimensional adaptations (randomized low-rank, matrix-free, and Adam-preconditioned updates) to tackle scalability, and provide theoretical connections between the value-function Jacobian and the true conditional scores. Empirically, the method achieves state-of-the-art performance on FFHQ $256 imes256$ and demonstrates robustness to score quality and discretization, including a class-guided MNIST example, highlighting the practical impact for nonlinear inverse problems with diffusion priors.

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.

Solving Inverse Problems via Diffusion Optimal Control

TL;DR

This work addresses the brittleness and intractability of diffusion-based inverse problem solvers that rely on conditional posterior sampling. By reframing the reverse diffusion as a discrete-time optimal-control episode and deriving a diffusion-based iLQR controller, the authors create a general solver that accommodates arbitrary differentiable forward operators and can recover the ideal posterior sampling equation as a special case. They introduce high-dimensional adaptations (randomized low-rank, matrix-free, and Adam-preconditioned updates) to tackle scalability, and provide theoretical connections between the value-function Jacobian and the true conditional scores. Empirically, the method achieves state-of-the-art performance on FFHQ and demonstrates robustness to score quality and discretization, including a class-guided MNIST example, highlighting the practical impact for nonlinear inverse problems with diffusion priors.

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 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.

Paper Structure

This paper contains 28 sections, 6 theorems, 64 equations, 5 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Notation
  4. Diffusion Models
  5. Posterior Sampling for Inverse Problems
  6. Optimal Control
  7. Diffusion Optimal Control
  8. High Dimensional Control
  9. Improved Posterior Sampling
  10. Brittleness to Discretization
  11. Intractability of $\nabla_{\mathbf{x}_t} \log p(\mathbf{y}|\mathbf{x}_t)$
  12. Dependence on the Approximate Score
  13. Related Work
  14. Experiments
  15. Conclusion
  16. ...and 13 more sections

Key Result

Theorem 4.1

Let Eq. eq:discretized_reverse_ode be the discretized sampling equation for the diffusion model with output perturbation mode control (Eq. eq:output_perturbation). Moreover, let the terminal cost be twice-differentiable and the running costs Then the iterative linear quadratic regulator with Tikhonov regularizer $\alpha$ produces the control

Figures (5)

  • Figure 1: Predicted $\mathbf{x}_0$ used in a probabilistic framework (above) compared to ours (below) for a general diffusion trajectory. The full forward rollout in our proposed framework allows for the predicted $\mathbf{x}_0$ (and therefore $\nabla_{\mathbf{x}_t} \log p(\mathbf{y} | \mathbf{x}_0)$) to be efficiently computed for all $t = 0, \dots, T$.
  • Figure 2: Inverse problem solution as a function of total diffusion timesteps $T$ for the $4 \times$ super-resolution task. Compared to DPS (top row), our method (bottom row) produces solutions that are higher quality, in greater agreement with the inverse problem contraint $\mathcal{A}\mathbf{x} = \mathbf{y}$, and more stable across $T$.
  • Figure 3: Examples from inverse problem tasks on FFHQ $256 \times 256$. From left to right each column contains ground truth, measurement, Diffusion Posterior Sampling (DPS), and ours.
  • Figure 4: Robustness to approximation quality of the score function. We consider the $4 \times$ super-resolution task with a randomly initialized diffusion model. Since the reverse diffusion process is no longer well approximated, DPS cannot produce a feasible solution, while our method still can.
  • Figure :

Theorems & Definitions (9)

  • Theorem 4.1
  • Lemma 4.1
  • Theorem 4.2
  • Theorem C.1
  • proof
  • Lemma C.0
  • proof
  • Theorem C.1
  • proof