Solving Inverse Problems via Diffusion Optimal Control
Authors: Henry Li
Marcus Aloysius Pereira
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion models
inverse problems
optimal control
Primary area: diffusion_based_models

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
Henry Li Yale University henry.1i@yall. 1ii@yale...... edu
Mareus Pereira Bosch Center for Artificial Intelligence marcus pereira@us bosch. com
Abstract
Existing approaches to Liffusion-basee inverse problem solvers frame the signal recovery task as 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 Xo prediction quality. We demonstrate that these limitations can sidestepped by reframing the gener- ative process discrete optimal control episode. We derive diffusion-baat bbsed optimal controller inspired by the iterative Linear Quadratic Regulator (iLQR) algo- rithm. 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 selection of neural inverse problem solvers, and establish a new baseline in image reconstruction with inverse problems
Introduction
Diffusion models Song and Ermon 22019) Ho et al. (2020) have been shown to be remarkably adept conditional generation tasks Dhariwal and Nichol [202T] Ho and Salimans 2022], part due their iterative sampling algorithm, which allows the dynamics of an uncontrolled prior score function Vx log Pr (x) be directed towards an arbitrary posterior distribution by introducing an additive guidance term u. When this guidance term is the conditional score Vx log (V x), the resulting sample provably drawn from the desired conditional distribution p(x x yy) Song al. (2020).
central obstacle to this framework is the general difficulty of obtaining the conditional score function Vx p, (x,) due its dependence the noisy diffusion variate x, rather than just the final sample Xo [Chung et al. [2023a]. In large-scale conditional generation tasks such as class- or text-conditional sampling the computational overhead of training time-dependent conditional score function from scratch is deemed acceptable, and is indeed the approach taken by Rombach et al. [202]] Saharia al. [2022]. and many others. However, this solution is not acceptable in inverse problems where the goal is to design a generalized solver that will work in zero-shot capacity for an arbitrary forward model.
This bottleneck has spawned a flurry of recent research dedicated to approximating the conditional score Vx log Pa (y) x,)
simple function of the noiseless likelihood log p(y xn) Choi et al. 2021
Chung et al. (2022] Rout et al. 20244), [Chung et al. [2023a] Kawar et al. [2022] Chung et al. 2023b
However, we will demonstrate in this work, these approximations impose significant cost the
performance of the resulting algorithm.
'Wor partially completed during intersshi at Bosch AI. com/
CCode is available at htttt::///// http///haaarraa//aaffuuooooooooo control
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

S1187-1
Aix.)
Axx.)
Key
Naj N1
XT-1
Conceptual illustration comparing a probabilistic posterior sampler to our proposed optimal control-based sampler. In a probabilistic sampler, the model relies approximation Xo & X to guide each step (left). We are able to compute X exactly on each step, resulting in much higher quality gradients V log ply (Ro) and an improved trajectory update (right).
Figure 1:
To address these issues, we propose novel framework built from optimal control theory where such approximations are no longer necessary. By framing the reverse diffusion process as an optimal control episode, we are able to detach the inverse problem solver from the strict requirements of the conditional sampling equation given by [Sng al. [2200), while still leveraging the exceptionally powerful prior of the unconditional diffusion process. Moreover, we find that the desired score function directly arises as the Jacobian the value function.
We summarize our contributions as follows:
We present diffusion optimal control, a framework for solving inverse problems via the lens of optimal control theory, using pretrained unconditional off-the-shelf diffusion models.
We show that this perspective overcomes many core obstacles present in existing diffusion- based inverse problem solvers. particular, the idealized posterior sampling score Song et al. [202T] approximated by existing methods can be recovered exactly as specific case of our method.
We showcase the advantages of our model empirically with quantitative experiments and qualitative examples, and demonstrate state-of-the-art performance on the FFHQ 256 x 256 dataset.
2
Background
Notation We use lowercase letters for denoting scalars E R, lowercase bold letters for vectors IR" and uppercase bold letters for matrices Hessians of scalar functions, e.g. Ix E RR" and 1xx Raxn for 1(x) : IR - R, respectively. We overload notation for time-dependent variables, where subscripts imply dependence rather than derivatives w.r.t. time, e.g., x, x(t). Furthermore, V(x)) and Q(x,, u:) are scalar functions
mm Subscripts indicate Jacobians and
despite being uppercase, in line with existing optimal control literature Betts 1998
2.1 Diffusion Models
The diffusion modeling literature uses the following reverse-time Itö SDE to generate samples Song et al. [002])
dx, [((x) V (x: (()
where X, E RD the state vector, f IR - RD and R + are drift and diffusion terms that can take different functional forms (e.g., Variance-Pressevvnn SDEs (VPSDEs) and Variance-Exploding SDEs (VESDEs) in Song al. [2021). Vx, log pt (x;) is the score-function and w, E R a vector of mutually independent Brownian motions. The above SDE has associated ODE called the

---

Predicted
True
X- D
(Vin Tamalle' Sermii))
o (N(:)) Imm forvrrd toller)
7
Xg
Figure 2: Predicted Xo used in probabilistic framework (above) compared to ours (helow) for general diffusion trajectory. The full forward rollout in our proposed framework allows for the predicted Xy (and therefore Vx, loy ) to be efficiently computed for all f = O.
probability-flow (PF) ODE given by
dx, = dx. + [f(x,) - _0(1)Vxx, log log):(x:))d.
(2)
with the same marginals (xx) as the SDE, which allow for likelihood computation (Song et al. 2021 Li al. 2024]. All practical implementations of diffusion samplers require dineeiiscetizzation of the PF-ODE. One such discretization is the well-known Euler-discretization which gives,
x = x, [f(x,) -2 [g(()>>x (x:):((;)]AA
where, At is the length of the discretization interval and we have reversed the time evolution by changing the sign of the drift. We are not restricted to only using the Fuler-discretization and any high-order discretization techniques can also be employed. More concisely, we have,
X h(xx); where h R" R"
which describes the general non-linear dynamics of the corresponding discrete-time diffusion sampler.
2.2 Posterior Sampling for Inverse Problems
Inverse problems are general class of problems where an unknown signal reconstructed from observations obtained by forward measurement process Ongie et al. [2020). The forward process is usually lossy, resulting an ill-posed signal recovery task where unique solution does not exist. The forward model can generally be written as
= A(xo) + 17,
(5)
where A : T + IRd is the forward operator, y E TR. the measured signal, Xp E R the unknown signal to be recovered, and 17 (O, Lla the noise (with variance o²) in the measurement process.
Given the forward model Eq. 5 and measurement y, sampling from the posterior distribution py(x can then be performed by solving the corresponding conditional Itö SDE
dx f(x) V log lo g(())d
(6)
where, invoking Bayes rule,
Vx log (xl)) = Vx (x) + Vx logp:(y)) xx).
As with the unconditional dynamics, Eq. (6) has a corresponding ODE
dx = [x(x) 2 g() (x)))()
(8)
which has an approximate solution obtained by the Euler discretization
X x, + (f(x,)
-:000-20000 g(t))Vv, (xg))xx
(9)
3

---

x y
T
T 16
T 32
T 64
7
128
T - 256
Figure 3: Inverse problem solution as a function of total diffusion timesteps T for the 4x 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 .Ax = y, and more stable across T.
2.3 Optimal Control
Optimal control the structured and principled approach the guidance of dynamical systems over time. Many methods have been developed in the optimal control literature and are popularly referred as trajectory optimization algorithms Betts [1998]. Perhaps the most well-known is the Iterative Linear Quadratic Regulator (iLQR) algorithm which uses first-order approximation of the dynamics and second-order approximations of the value-function and Todorov (2004).
Formally, let us define an arbitrary user-defined global cost function
JT Σ (x,,, u,) + lo(xo): 1=T
(10)
composed a sum over scalar-valued running and terminal cost functions l and Ep. Optimal control theory dictates that the value function V Vxxt, t) := min [u. Jr satisfies the following recursive relation also known as Bellman's 's Principle of Optimality
V(x,, t) = min (xx,,, u,) + (X
(11)
The iL.QR algorithm centers around approximating the state-action value function, Q(x;, u,) = (x,,, u,) V(x,11 1-11-, 1), from which the value function can be recovered as V(x,, t) min,, Q(x,. u.).
(12)
Then given a state transition function x, h(x++1, Up+11 ) where erucially note that have defined time to flow backwards from / = T..... 0, the iLQR algorithm has feedforward and feedback gains
k = - Qu
and
K = =-Qui00 Qux
(13)
The update equations can be written as
Vx Qx - Qunk and
Vxx Qxx K" Quuk.
(14)
Given the feedforward and feedback gains {(K, k,) k)) and 8g == Xo, we can recursively obtain the locally optimal control time as function of the present states x, and controls
x, h(x 11
(15)
u u, + Ak xx).
(16)
For more detailed treatment of iLQR as well as derivation of the equations, please Appendix
Diffusion Optimal Control
We motivate our framework by observing that the reverse diffusion process Eq. 11 is an uncontrolled non-linear dynamical system that evolves from some initial state (at time = ) to some terminal state (at time / 0). By injecting control vectors into this system we can influence its behavior and hence terminal state (i.e., the generated data) to sample from desired pxx(y). There are two obvious ways to inject control into this process:

---

Algorithm Diffusion Optimal Control
Input: X, T, X. Initialize u, kt, K, as 0 for t = ......,, (x1 as uncontrolled dynamics
for iter to num_ -imer do
Vx, Vxx f- Vx, logp(y/Xon (xo), V2 logply (x)))
> Initialize derivatives of ((x,, ()
for 1
T do
Compute k,, K,, Vx, Vxx
D. See Eqs. 13
end for
for f = T to do
X h(xe, Ak; K2(X, x())
Update X - with new u
X + Xx
end for end for
In input perturbation control, we apply the u, before the diffusion step:
x (x, u,) [f(x+ + uz) 1 v0( log Px u.)] AA.
(17)
2. In output perturbation control, u, is applied after the diffusion step:
X
[f()) 2 20) lo (x:) At + us-
(18)
Observe that iLQR formulated for general discrete-time dynamic processes. When applied specifically to the reverse diffusion dynamics of diffusion models, we are able to make several simplifications. First, we assume that we do not have access to any guidance except time i.e., (x,,, Ue) does not depend on X,.
In the case of input perturbation control, we observe from Eq. (7) that hy hu, whereas output perturbation control implies that h, = I, resulting in the left and right equations, respectively:
Qx h V
Qx h V
(19)
Qn bu + h V
Qu eu + V
(20)
Qxx = h Vi hx
Qxx h V hx
(21)
Qux Qxu = h V hx
Qux
Q V h&
(22)
Quu luu h VV Vx
Quu lau + Vi
(23)
The derivatives of V can then be backpropagated using the following equations:
Vx Qx -K Quuku = Qxx - -
= Qxx + Q QuaQu
(24)
Vxx Qxx KT KTQuuk
Qxx
(25)
In high dimensional systems such as Eq. B matrices may be singular. Therefore, a Tikhonov regularized variant of iLQ is often employed, where matrix inverses are regularized by diagonal matrix al Tassa et al. 2014))
3.1 High Dimensional Control
Compared to the dynamics in traditional application areas of optimal control, those consider in Eqs. (17 18) are much higher dimensional in the state x and control u variates. Therefore, iLQR faces several unique computational bottlenecks when applied to such control problems.
In particular, the Jacobian matrices h h, and the second-order derivative matrices Vxx, Qxx, Qux, Qxu: and Qun are particularly expensive to compute, store, and perform down- stream operations against. For example, in three-channel 256 x 256 image, these matrices naively contain (256 x 256 3)² ze 39B parameters.
5

---

Ground Truth
Mlassremeent
Ground Trush
Mensurement
Ours
Rane Inpoir ting
Super Resouttinn
Figure 4: Examples from inverse problem tasks on FFHQ 256 x 256. From left to right each column contains ground truth, measurement, Diffusion Posterior Sampling (DPS), and ours.
In Appendix D.I we propose and analyze three modifications the standard iL.QR algorithm: randomized low rank approximations, matrix-free evaluations, and action updates via an adaptive optimizer. that significantly reduce runtime and memory constraints while introducing minimal deterioration to performance inverse problem solving tasks.
4 Improved Posterior Sampling
We demonstrate that our optimal control-based sampler overcomes several practical obstacles that plague existing diffusion-basse -based methods for inverse problem solvers.
Brittleness to Discretization In probabilistic framework, solutions inverse problems incur discretization error from the numerical solution of Eq. (8) that decays poorly with the total diffusion steps 7 of the diffusion process. While much research has been conducted the acceleration of unconditional diffusion processes Song et al. [2020), Besiiceuur-Martinne et al.) [2021]. Karras eet al. [2022]. Meng al. 20033, sample quality appears decay much more aggressively in diffusion- based inverse problem solvers (Figure
We theorize that this is due to two reasons: 1) the posterior sampler Eq. 9 is only correct in the limit of infinitely small time steps, and 2) the quality of the approximated conditional score term V, x log piyppxxxx (xx) decays quickly with time (Figure 2). and so fewer timesteps lead to fewer chances at low t to correct errors made at high f. On the other hand, since optimal control directly casts the discretized process
end-to-end control episode, it produces feasible solution for any number of discretization steps T'.
Intractability of Vx. log x(y)()));) When the forward model A is known and comes from simple distribution, the conditional likelihood p(y (x.) can be derived in closed form for = 0. On the other hand, the dependence of y on x, for > 0 is generally not known without explicitly computing Xo, which requires sampling from the diffusion process, Ultimately, obtaining the conditional score term
Xx, log p(y) xx) is highly nontrivial task Song et al. [2021).
To sidestep this issue, many works Meng and Kabashima 2022 Song et al. [2022 Chung et al. 2023al factorize this term the integral
p(ylx;) = f ply
(26)
and then apply a series of approximations to recover a computationally feasible estimate of the conditional score. First, the marginal p(xo/x;) ) is replaced by the marginal conditioned Xo, i...

---

p(xo|Xe,XX) Xb) = (x Kim and Ye [2021). J. Next, the Xo-centered marginal is replaced by the posterior mean E Xg x given by Tweedie's 's formula Efron 2011]. Finally, the true score is replaced by the learned score network.
While these approximations are necessary in probabilistic framework, show that they required in our method. Intuitively, this is because the linear quadratic regulator backpropagates the control cost logpl, p(y) x) through forward trajectory rollout, which naturally computes the true conditional score at each time t. Moreover, our model always estimates Xo Xt exactly (up to the discretization error induced by solving Eq. 3). rather than forming an approximation Xy = Xu (Figure 2). We formalize this observation with the following statement.
Theorem 4.1. Let Eq. 3 be the discretized sumpling equation for the diffusion model with output perturbation mode control (Eq. 18.. Moreover, let the terminal cost
(x(xo) log p
(27)
be twice-differentiabl and the running costs
(x u,) 0.
(28)
Then the iterative linear quadratic regulator with Tikhonov regularizer a produces the control
ur a (oy)(y)xx))
(99
In other words, by framing the inverse problem as an unconditional diffusion process with controls ue, our proposed method produces controls that coincide precisely with the desired conditional scores Vx, logp(y/(x))
Let us further assume that logp(y)x;) = log Ply Xg). i.e.. x, contains no additional information about y than Xp- This assumption results in the posterior mean approximation Chung et al. [2023a under stochastic dynamics (Eq. [T]. where we additionally obtain exact computation of Xo, rather than So R Xy via Tweedie's 's formula Kim and Ye [202T] Under the deterministic ODE dynamics (Eq. 2
we recover the true posterior sampler under appropriate choice of Tikhonov regularization constant a.
Lemma 4.2. Under the deterministic sampler with output perturbation mode control, 901
recovers posterior sampling (Eq. 9
We demonstrate similar result with input mode perturbation.
Theorem 4.3. Let Eq. be the discretized sampling equation for the diffusion model with input perturbation mode control (Eq. 17.. Moreover, let
lo(xo)
(30)
and the running costs
lo(x,, u.) =.
(31)
Then the iterative linear quadratic regulator with Tikhonov regularizer 9(0)-t produces the dynamical sytem
X, xx,+((((E))
(0(() (X.) Vx (ogp)((//))))
(32)
where X, X4 U4-
Observe that Eq. 32 can be understood as predictor-corrector sampling method, where the predictor produces an unconditional reverse diffusion update and the corrector produces a conditional correction step the intermediary variable x, X, Ut
Ultimately, these results demonstrate that our proposed method is able to recover the idealized sampling procedure under mild assumptions on the diffusion optimal control algorithm.
7

---

Inpainting
Inpainting
Gaussian Deblurring
Motion Deblurring
FID
LPIPS
FID
LPIPS
FID
LPIPS
FID
LPIPS
FID
LPIPS
Ours (NFE= 2500)
32.47
0,171
15.93
0.053
20.22
0.122
31.80
0.189
39.40
0.217
Ours (NFE 1000)
37.53
0.189
20.75
0.108
23.88
0.164
35.24
0.191
45,99
0.233
PSLD INFE = 1000)
34.28
0.201
21.34
0.096
43.11
0.167
41.53
0.221
Flab-Diflusion' * (NF1 varies)
53.95
0.195
65.35
0.280
64.57
0.267
DDNM (NFE- 1000)
68.94
0.328
105.3
0.802
72.28
0.483
126.0
0.995
DPS (NFE 1000)
39.35
0.214
33.12
0.168
21.19
0.212
44.05
0.257
39,92
0.242
DDRM (NFB= 1000)
62.15
0.294
42.93
0.204
69.71
0.587
74,92
0.332
MCG 1000)
87.64
0.520
40.11
0,309
29.26
0.286
101.2
0.340
310.5
0,702
PNP-ADMM
66.52
01,353
0.406
0.692
0.441
89.08
0,405
Score-SDE (NFE = 1000)
96.72
0.56.
60.06
0.331
76.54
0.612
109.0
0.403
292.2
0.657
ADMM-TV
110.6
0.322
0.463
186.7
0.507
152.3
0.508
Table 1: Quantitative evaluation (FID, LPIPS) of model performance inverse problems on the FFHQ 256x256-1K dataset.
Dependence on the Approximate Score While our theoretical results require that the learned score function se (x,, t) approximates the true data score log Pr (x(,, , t). we emphasize that the performance of our method does not necessitate this condition. In fact, we find that reconstruction performance is theoretically and empirically robust to the accuracy of the approximated prior score sy (x, t) Rd Xx, log pt (xx) conditional score Vx, log ( xx) Vx, log pe ()(y)),) (x,) terms. This is because the optimal control-basee solution is formulated for the optimization of generalized dynamical systems, and thus agnostic to the diffusion sampling process.
Certainly, improved approximation of the score terms result better-informed prior and usually higher sample quality. However, we demonstrate that our sampler produces remarkably reasonable solutions even in the case of randomly initialized diffusion models. Conversely, probabilistic posterior samplers can only sample from p when the terms composing the posterior sampling equation (Eq. (8) are well approximated (Figure 6). Modeling errors can occur even in foundation models. For example, this scenario may arise models trained on regions where there underrepresenteed examples in the data. When these arise from existing social or ethical biases, they can further perpetuate amplify biases to the resulting model if left caaddresseeBoluubaar et al. [2016). Birhane jtt al. [021], Srivastava et al. [2022).
There exist several methods that seek to alleviate the errors incurred by Tweedie's formula (being mean approximation of the diffusion process), including Song et al. (2024) which imposes a hard data consistency optimization loop at various points in the diffusion process, and Rout et al. [2023) which includes a stochastic averaging loop in each step of the diffusion process. However, these methods still rely on Tweedie's ' formula for the error reduction scheme, which assumes access to a ground truth score function. Ultimately, the aforementioned problems in the present section are exacerbated in existing samplers, and relatively less consequential in our solver.
5 Related Work
The recent success of diffusion models in image generation [Snng and Ermon [2019] Ho et al. [2020.
Song et al. [2021) Rombach et al. (2022) has spawned a surge of research in deep sarring--bssed solvers to inverse problems. (Song et al. [020] demonstrated a strategy for provably sampling from the solution set p(xly) of general inverse problem y A(x) using only an unconditional prior score model Vx log Pr (x) and a forward probabilistic model log p(y|x;), (x)). However, a crucial problem arises in the intractability of forward probabilistic model, which depends on the noisy x, rather than the final Xg. This has resulted in a series of approximation algorithms Choi et al. [2021), Kawar et al.. 2022
Chung et al. 2022 2023a bl, Kawar et al. 2023 for the true conditional diffusion dynamics.
Topics in control theory have been applied to deep learning al. (2020), Pereira al. (2020) as well as diffusion modeling Berner et al. [022). Optimal control can also be connected to diffusion processes via wewaard-bacwwar SDEs Chen et al. (202T) However, these ideas have not been applied to guided conditional diffusion processes solely at inference time, nor for guided conditional sampling. Our proposed optimal control-based algorithm is, to our knowledge, the first such framework for deep inverse problem solvers.
8

---

0
Ours
DPS
2
4
S
5
S
S
6
O
7
7 6
7
7
F
D
R
8
2
Figure 6: Robustness to approximation quality of the score function. We consider the 4x super-resolution task with a randomly initialized diffusion model. Since the reverse diffusion process is longer well approxi- mated, DPS cannot produce a feasible solu- tion, while our method still can.
9
5
9
9 9
Figure Examples from the class- conditional inverse problem. DPS (left) is compared against ours (right). Each row is a different target MNIST class.
6 Experiments
Following previous work Chung et al. 2023a Meng and Kabashimal [2022] Kawar et al. 2022] we consider five inverse problems. 1) In 4x image super-resolution, operator. 2) In randomized inpainting, we uniformly omit 92% of all pixels (across all channels). 3) In box inpainting, from each side of the image, as in Chung et al. [2022]. 4) In Gaussian deblurring, we use a kernel of size 61 61 and standard deviation 3.0. In motion deblurring,
the bicubic downsampling
mask out 128 x 128 block uniformly sampled from a 16 pixel margin
generate images according to a library of point spread functions with kernel size 61 x 61 and intensity 0.5. Following the experimental design in Chung et al. (2023a). we apply Gaussian noise with standard deviation 0.05 to all measurements of the forward model.
We compare against generalized diffusion inverse sampler (Score-SDFF) proposed in Song et al. 2022]. Diffusion Posterior Sampling (DPS) Chung et al. [2023a], Denoising Diffusion Restoration Models Kawar et al. [2022], Manifold Constrained Gradients (MCG) Chung et al. [2022] 2. well as two recent latent diffusion-bern based methods Fabian et al. [2023) Saftt--ifftsion"" ) and Rout et al. 2024] (PSLD). For non-diffusion baselines, compare against Plug-and aan--lld Play Alternating Direction Method of Multipliers (PnP-ADMM) with neural proximal maps Chan et al. (2016), Zhang et al.
[017]. and iital-variation based alternating direction method of multipliers (TV-ADMM) baseline proposed in Chung et al. 2023a]
We validate our results on the high resolution human face dataset FFHQ 256 x 256 Karras et al. (2099). Several methods are model agnostic (DPS, DDRM, MCG, and thus evaluated with the same pre-trained diffusion models. To fairly compare between all models, all methods use the model weights from Chung et al. [2023a]. which are trained on 49K FFHQ images, with IK images left as a held-out set for evaluation. We compare algorithm against competing frameworks on these last 1K images. We report our results on FFHQ 256 x 256 in Table and demonstrate improvements on all tasks against previous methods. Finally. demonstrate the performance of our algorithm the nonlinear inverse problem of class-conditional generation. Namely, let A(x) = classifier( and p(y|x x) be its associated probability. We compare method DPS the inverse task of generating an MNIST digit given a label y. Compared to images generated by DPS, images from our method exhibit more pronounced class alignment and higher overall sample quality (Figure
7 Conclusion
In this paper we presented a novel perspective on tackling inverse problems with diffusion models framing the discretized reverse diffusion process discrete time optimal control episode. We demonstrate that this framework alleviates several core problems in probabilistic solvers: its depen- dence on the approximation quality of the underlying terms in the diffusion process, its sensitivity to the temporal discretization scheme, its inherent inaccuracy due to the intractability of the conditional score function. We also show that the diffusion posterior sampler can be seen as a specific case of
tetttps://itthuu wwtp://eevoorrroeoooooooooiooii.
FFor Gaussian blur and random inpainting, Flash- Diffusion uses randomly sampled, but less severe degrada- tion operators than our experimental setup.

---

our optimal control- basee sampler. Finally, leveraging the improvements granted by our solver, we validate the performance of our algorithm on several inverse problem tasks across several datasets, and demonstrate highly competitive results.
References
Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:221.003364 2022.
John T Betts. Survey numerical methods for trajectory optimization. Journal of guidance, control, and dynamics, 21(2):193-207, 1998.
Abeba Birhane, Vinay Uday Prabhu, and Emmanuel Kahembwe. Multimodal datasets: misogyny. pornography, and malignant stereotypes. arXiv preprint arXiv:2110.01663 2021.
Tolga Bolukbasi, Kai- Rei Chang, James Y Zou, Venkatesh Saligrama, and Adam T Kalai. Man is to computer programmer as woman to homemaker? debiasing word embeddings. Advances in neural information processing systems, 29, 2016.
Stanley Chan, Xiran Wang, and Omar Elgendy. Plug-and-plas admm for image restoration: Fixed-point convergence and applications. IEEE Transactions on Computational Imaging, 3(1): 84-98, 2016.
Tianrong Chen, Guan-Horng Liu, and Evangelos A Theodorou. Likelihood training of schr) odinger bridge using
sdes theory. arXiv preprint arXiv:21 12211, 2021.
Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Condi- tioning method for denoising diffusion probabilistic models. arXiv preprint arXiv:2/08.02938 2021.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35:25683-25666, 2022.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Mare I. Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. International Conference on Learning Representations, 2023a.
Hyungjin Chung, Jeongsol Kim, and Jong Chul Ye. Direct diffusion bridge using data consistency for inverse problems. arXiv preprint arXiv:2305. 19809, 2023b.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 2011.
Zalan Fabian, Berk Tinaz, and Mahdi Soltanolkotabi. Adapt and diffuse: Sample-adaptive recon- struction via latent diffusion models. arXiv preprint arXiv:2309.0666 2023.
Nathan Halko, Per-Gunnar Martinsson, and Joel A Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review, 53(2): 217-288, 2011.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598. 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Matthew D Houghton, Alexander B Oshin, Michael J Acheson, Evangelos A Theodorou, and Irene M Gregory. Path planning: Differential dynamic programming and model predictive path integral control on vtol aircraft. In AIAA SCITECH 2022 Forum, page 0624, 2022.
10

---

David H Jacobson. New second-order and first-order algorithms for determining optimal control: A differential dynamic programming approach. Journal of Optimization Theory and Applications, 411-440, 1968.
Alexia Martinaa,,, Ke Li, Rémi Piché-Taillefer Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint S.Ki:22105.140880 2021.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401-4410, 2019.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 25:26565-26677, 2022.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in Neural Information Processing Systems, 35:23593-23606, 2022.
Bahjat Kawar, Noam Elata, Tomer Michaeli, and Michael Elad. Gsure-hased diffusion model training with corrupted data. arXiv preprint arXiv:2305. 13/28, 2023.
Kwanyoung Kim and Jong Chul Ye. Noise2score: tweedic's ' approach to self-supervised image denoising without clean images. Advances in Neural Information Processing Systems, 34:864-874 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 44126698,, 2014.
Dana A Knoll and David E Keyes. Jacobian-free newtonnkkrylee methods: a survey of approaches and applications. Journal of Computational Physies, 193(2):357-397, 2004.
Henry Li, Ronen Basri, and Yuval Kluger. Likelihood training of cascaded diffusion models via hierarchical volume- -luuee-reserring maps. In The Twelfth International Conference on Learning Representations, 2024. URL (ttpsttts:: http://pppnneeeeeeeeeeeeorooroee
Weiwei Li and Emanuel Todorov. Iterative linear quadratic regulator design for nonlinear biological movement systems. In First International Conference on Informatics in Control, Automation and Robotics, volume 2, pages 222-229. SciTePress, 2004.
Guan-Horg Liu, Tianrong Chen, and Evangelos A Theodorou. Ddpnopt: Differential dynamic programming neural optimizer. arXiv preprint arXiv:2002.08809 2020.
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 14297-14306, 2023.
Xiangming Meng and Yoshiyuki Kabashima. Diffusion model based posterior sampling for noisy linear inverse problems. arXiv preprint arXiv:2211.123333 2022.
Gregory Ongie, Ajil Jalal, Christopher A Metzler, Richard G Baraniuk, Alexandros G Dimakis, and Rebecca Willett. Deep learning techniques for inverse problems in imaging. IEEE Journal on Selected Areas in Information Theory, 1(1):39-56, 2020.
Samet Oymak, Zalan Fabian, Mingchen Li, and Mahdi Soltanolkotabi. Generalization guaran- tees for neural networks via harnessing the low-rank structure the jacobian. arXiv preprint (ttp://906.53922. 2019.
Mareus Pereira, Ziyi Wang, Tianrong Chen, Emily Reed, and Evangelos Theodorou. Feynman neural network architectures for stochastic control using second-order fbsde theory. In Learning for Dynamics and Control, pages 728-738. PMLR, 2020.
Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):100, 2008

---

Robin Rombach, Andreas Blattmann, Dominik Lorenz. Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and puttern recognition, pages 10684-10695, 2022.
Litu Rout, Yujia Chen, Abhishek Kumar, Constantine Caramanis, Sanjay Shakkottai, and Wen-Sheng Chu. Beyond first-order tweedie: Solving inverse problems using latent diffusion. arXiv preprint 1:X.::2312.00852,, 2023.
Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alex Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint 1rXiv:1706.0444/ 2017.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, et Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494, 2022.
Tomohiro Sasaki, Koki Ho, and E Glenn Lightsey. Nonlinear spacecraft formation flying using constrained differential dynamic programming. In Proceedings of AAS/AIAA Astrodynamics Specialist Conference, 2022.
Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen. Solving inverse problems with latent diffusion models via hard data consistency. arXiv preprint arXiv:2307.08123. 2024.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse--uuided diffusion models for inverse problems. In International Conference Learning Representations, 2022.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differcntiaa equations. In International Conference on Learning Representations, 2021. URL httpttt://peeeeeiiieeee
PXTIG12RRHS
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonso et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint 04615, 046.:,200004461,, 2022.
arXiv:2206.00
Yuval Tassa, Tom Erez, and William Smart. Receding horizon differential dynamic programming. Advances in neural information processing systems, 20, 2007.
Yuval Tassa, Nicolas Mansard, and Emo Todorov. Control-limited differential dynamic programming. In 2014 IEEE International Conference on Robotics and Automation (ICRA), pages 1168-1175. IEEE, 2014.
Emanuel Todorov and Weiwei Li. A generalized iterative lgg method for locally-optimal feedback control of constrained nonlinear stochastic systems. Proceedings of the 2005, American Control Conference, 2005., pages 300-306. IEEF, 2005.
Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng. and Lei Zhang. Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising. IEEE transactions on image processing, 26(7): 31 3142-3155 2017.
12

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We demonstrate results through rigorous analysis of our algorithm and extensive experiments on multiple inverse problem settings over several datasets.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Yes, the paper discusses the runtime cost of the work, and provides equivalent budget analysis, where still demonstrates competitive performance on each benchmark.
Justification:
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The paper provides full proofs for all theory in the appendix.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper discloses all hyperparameters and implementation details in the appendix.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described supplemental material?
Answer: [Yes]
The paper provides open access to the data, which publicly available. The authors will release code upon acceptance.
Justification:
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper provides all details in the appendix.
7. Experiment Statistical Significance
13

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: Experiments for other works do not provide error bars, therefore error bars would not benefit the analysis in this paper.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Experiments can be run on any GPU A4000 or later.
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http://rrrrr..cccouubuui//cccaca
Answer: [Yes]
Justification: We confirm the NeurIPS Code of Ethics in every respect.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The paper discusses this in the appendix.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The results in this paper paper do not have high risk for misuse.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We credit all creators and original owners of assets.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets introduced.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No research is performed with human subjects.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
14

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No research is performed with human subjects.
15

---

A Impact Statement
This paper builds on large body of existing work and presents an improved technique for solving generic nonlinear inverse problems, which can be seen generalization of guided diffusion modeling. Controlling the diffusion process in a generative model has many societal applications, and thus a broad range of downstream impacts. We believe that understanding the capabilities and limitations of such models in public forum and open community is essential for practical and responsible integration of these technologies with society. However, the ideas presented in this work, as well as any other work in this field, must be deployed with caution to the inherent dangers of these technologies.
B Deriving the Iterative Linear Quadratic Regulator (iLQR)
Differential Dynamic Programming (DDP) is very popular trajectory optimization algorithm that has a rich history of theoretical results Jacobson 1968) as well as successful practical applications in roboties Tassa et al. (2007] 2014], aerospace Houghton al.) [2022), (Sasaki al.) [202] and biomechanics Todorov and Li [2005]. It falls under the class of indirect methods for trajectory
optimization, wherein Bellman's principle of optimality defines the so-called optimal value-function
which in turn can be used to determine the optimal control. This is in contrast to so-called direct methods which cast the problem at hand into nonlinear constrained optimization problem.
To formulate an optimal control algorithm we first define the state transition function of a dynamical system
X
h Uz).
(33)
The next ingredient that we need for our optimal control approach is a cost function J(t,, uz) E IR. This is used to define a performance criterion that iLQR can optimize with respect to the set of controls {u, != (i.e., the control trajectory going backwards from time t T to t 1). The cost-function is defined as follows:
Jr Σ (xx, us) + lo(xo),
(34)
where, & and lo are scalar-valued functions which are commonly referred as the running cost- function and the terminal cost-function respectively.
To obtain the sequence of optimal controls, we employ the dynamic programming principle. To do so, we first introduce the notion the Value-function defined as follows:
V(x;, t)
min Je = min
I (x un) + (o(xo)]
(35)
(u,, ) (u,, (mmm)
Intuitively, the Value-function resembles the optimal cost-to-go starting from time step and state x; until the end of the time horizon (i.c., t = 0). Using this definition, one can easily derive the following recursive relation also known Bellman's Principle of Optimality:
Vxx,, () min [6(x,, un) + (Xx-1,
( 11)]
(36)
often useful defintion used the derivation of the iLQR Riccati equations that of the State-Action Value- Salue-Function Q(x,, u,) given by.
Q(xi, uz) (x,,, uz) + VX t 1-1) 1) Therefore, V(x;, t) min Qxx,, u,)
(37)
(38)
A sketch of the derivation of the Riccati equations is as follows: we take second-order Taylor expansions of both Q(x,, u.) and V(x,, t) around nominal state and action trajectories of [x,]==)
and {ū, [=] respectively. Next, substitute these into Eq. and equate the first- and second-
16

---

order terms to yield the following relations between the derivatives of Q, e and V:
Qx lx Vx
(39)
Qu la + h
(40)
Qxx = 4x + h Vx
(41)
Qxu = xx + h V hu
(42)
Qux lux + h V hx
(43)
Qun Lun + huVhh
(44)
where hx, and hu, are the Jacobians of the dynamics function h(x,, we). evaluated at time step t, w.r. the state and the control vectors respectively. For ease of notation, we have dropped the subscript I and therefore all derivatives above should be considered to be evaluated at time step 4, while V and Vxx above to indicate the gradient and hessian of the Value-function evaluated at the next time step (i.e., at time step 1).
Next, we substitute for the second-orde approximation of Q(x, uz) into Eq. (38) and note that u, can be written in terms of the nominal control as follows:
U + du,
This results in a quadratic objective w.r.t Sú, and the minimization in Eq. 38 can be performed exactly resulting in the following optimal perturbation from the nominal control trajectory:
gug = kG + K28x;
(45)
where, the feedforward and feedback gains are given by the following expressions:
QuiQu
(46)
K Qud Qux
(47)
Finally, by substituting for the optimal duž back into (3. we can drop the min operator and equate the first- and second-order terms on both sides. This results the following Riccati equations:
Vx Qx K Quuk
(48)
Vxx Qxx K" Quk.u
(49)
This concludes the sketch derivation of the Riccati equations. The algorithm roughly proceeds as follows:
1. We start with an initial guess of the the nominal control trajectory [ū.] and generate the corresponding nominal state trajectory (xt 10 using x = h(x++1, U1+1).
2. By noticing from Eq. (35) that V(xo, 0) ((xo) we can obtain expressions for Vx and Vxx evaluated at &g-
3. Next, we compute the derivatives of Q given by equations. (x,)
444 using (ūr) and
4. Using the derivatives of Q, we can compute the feedforward and feedback gains using equations (46 (47
5. Finally, using the Riccati equations 48 backwards in time.
49
we can propagate both 1/x and Vxx one step
6. We then repeat the steps 3, 4 and 5 until we backpropagate the derivatives of V to time step T.
7. This completes one iteration of iLQR. At the end of each iteration the gains are used to produce the updated nominal control trajectory as follows:
u, = ú, + ak + K(x: - -x.)
(50)
where, x, is the state obtained by unrolling the dynamics subject to the updated controls:
X; = h(x++1 11+
8. The new nominal control trajectory of is used to produce new nominal state trajectory x and the algorithm is repeated from step onwards until convergence or a fixed number of iterations.
17

---

C Proofs
Theorem 4.1. Let Eq.. be the discretized sampling equation for the diffussion model with output perturbation mode control (Eq. 18.. Moreover, let the terminal cost
(o(xo) - -lg ((yp((())
(27)
be twice-differentiable and the running costs
(x(x;, U.) 0.
(28)
Then the ilerative linear quadratic regulator with Tikhonov regularizer a produces the control
u, = aVx. logp(y)xo).
(29)
Proof. We demonstrate the result via induction for t = I.....T T.
Since assume that luu 0, Vxx vanishes:
Vxx Qxx Qux QuiQux
(51)
li Vixlax V/ Vocx(VXx xxx -
Vxxhx
(52)
O.
(53)
Similarly, Vx also greatly simplifies
Vx Qx + Q Qui
(54)
N + lix V Vxx (xxx) V'
(55)
xx LV x-
(56)
Turning the Tikhonov regularized feedforward term,
k - Qud Qu
(57)
(h T' Vxx hx -t
(58)
+ aI) a
(59)
1 Vz
(60)
Finally, the feedback term disappears due to the vanishing Vxx
K 2ux
(61)
0.
(62)
Explicitly denoting the dependence of V and V on f, we can rewrite Eq. 56 as
(( h (t-1)
ox--1
Ox. 0x,-1
Combining this observation with the fact that to - log (xo), we can conclude that
V(0
Vx. log p(y|
(63)
where depends on x₃ via the state transition function h (Eq. 18 Therefore, we have that
V
a
Vx
log p(y(xx)
K = 0.
Finally, given our action update (Eq.
15
we can conclude our desired result
u,
V x, logp(y|Xo).
(64)
a
18

---

Lemma C.1.
Under the deterministic sampler with output perturbation mode control. 9(1)= recovers posterior sampling (Eq.
Proof. Substituting in a 9(0)=A7 to Eq. we observe that Eq. 18 can now be written as
1 X [f(x)) 2 (1) (x.) + Vx, (yg
(65)
Under the determinstic sampler, we can conclude that log (y xx)) log ( xt), since each X has unique path through the sample space. Therefore, we conclude that Eq. 65 resembles the ideal posterior sampler equation We conclude our proof.
Theorem 4.3. Let Eq. be the discretized sampling equation for the diffusion model with input perturbation mode control (Eq. 17 Moreover, let
(o(xo) - llog p(yp((()))
(30)
and the running costs
l((x;, u.) =0
(31)
Then the iterative linear quadratic regulator with Tikhonov regularizer a(1/at produces the dynamical sytem
X = X + (t)) (Xt.)
Vx px
(32)
where X, X4 + Up.
Proof. We similarly demonstrate the result via induction for = I...., T.
Again, assuming that fuu Vxx vanishes:
Vxx = Qxx - QuxQuuQun
(66)
= Qxx - Qxx luu +Qxx)= -1xxx
(67)
=0, 0,
(68)
whereas Vx greatly simplifies
Vx = Qx TTx Qu
(69)
= =
(70)
Turning the feedforward and feedback terms, have
k Qu
(71)
- (h Vxx hx al) QQu
(72)
-0 + (al) -t.
(73)
1 11 1/
(74)
and
K = -Qii
(75)
0.
(76)
We observe that
(2(0 = h
(t-1))
19

---

SR x4
Random Inpainting
Box Inpainting
Gaussian Deblurring
Motion Deblurring
T
50
50
50
50
50
num iters
50
100
100
100
step_size
1c -
3
le 3
le 3
1c
1c 3
(o(xo)
[A(x)) yll
||A(xx)) - yll
[|A(x)) - yll
[|A(x)) - yll le
[[((x)) - yll
Ic 4
le 4
le
4
le 4
4xx,, u.)
al|uel|
|ulll
||ulll
al uell
al|u,
&
control_mode
input mode
input mode
input mode
input mode input mode
Table 2: Hyperparameters for FFHQ experiments.
Therefore, noting that v2 log p we have
k V((1)
(h (1(0)) (((t1))
Xx, log p())o))))
a
xx, logp(y/xo(xx))
a
Applying the feedforward terms to the diffusion sampling process, we have
- (x: + u,) [f(x; + ue)
[1) log + At
We define the intermediary variable which has dynamics
x; = X, u1
(77)
x = x (f(x)) _g(t)² Vx log (x:)) -u.
(78)
We now can see that, letting 2lg we obtain
1 g())) ((t)) (V + Vx (x)
X = Xi (f(e))
D
Implementation
For all experiments, use publicly available datasets and pre-trained model weights. For the FFHQ 256 x 256 experiments, we use the last 1K images of the dataset for evaluation. For MNIST, we do use images directly in the inverse classification task. The images were only used for training the pretrained diffusion model.
For models, we used the pretrained weights from Chung et al. 2023a for FFHQ 256 256 tasks, and the Hugging Face laurent/ hwtreettmmnitt222 diffusion model for MNIST experiments. No further training is performed any models. Further hyperparameters ply x) in MNIST class-guided classification, we use a simple convolutional neural network with two convolutional layers and two MLP layers, trained the entire MNIST dataset.
be found in Table [ For the classifier
D.1 High Dimensional Control
To speed up our proposed method, we leverage the following three modifications to the standard iLQR algorithm.
20

---

SR x-1
Random Inperning
Box Iepainring
Gurasian Dubluring
Mition Dublariing
PSNK
SSIM T
MSG:
PSNK 7
XS0M
M15
PSNK
SSIM
MbN
PSNK
SSIM
MSL-1
PSNK
55IM 1 MSa
OZ 550
27.45
0.792
117.0
31.84
0.382
42.57
25.33
0.803
198.6
24.99
0,284
206.1
25.05
0.721
201.9
(T= 1000
1950
DORM 1000)
INP-ADMM
36.55
0.865
11:65
9337
0.512
24.65
222.9
ScIE-SDEE
0.617
13.52
2891
12620
14291
ADMM-TT
21.36
0.803
267.4
1781
0.314
1926
22.00
0.788
407.5
22.37
0.301
326.8
21.96
0758
4754
Table 3: Quantitative evaluation (PSNR, SSIM, MSE) performance on inverse problems on the
FFHQ 256x256-1K dataset.
SR x4
Random Inpainting
LPIPS 4
PSNR
SSIM 1
MSE 4
LPIPS +
PSNR 1
SSIM 1
MSE 4
k: 0
0.254
24.00
0.691
141.2
0.121
28.33
0.755
56.74
k: =
0.171
27.45
0.792
117.0
0.053
31.84
0.882
42.57
k 4
0.171
27.47
0.794
116.4
0.052
31.99
0.883
41.12
k = 16 0.170
27.43
0.799
117.5
0.050
32.12
0.891
39.90
Table 4: Ablative study on the effect of rank in the low rank and matrix-frt -free approximations on performance (LPIPS, PSNR, SSIM, NMSE) of our proposed model on the FFHQ 256x256- K dataset dataset.
Randomized Low-Rank Approximation The first and second order terms Eqs. 19 25 are
corresponding Taylor expansions of deep neural functions. Even with the of automatic differen- tiation libraries, the formation of these matrices is incredibly expensive, requiring at least dim (x) backpropagation passes (where dim(x) 39B in some experiments). To reduce the cost of com- puting these matrices, we utilize their known low rank structure Sagun et al. [017)] Oymak et al. 2019).
Leveraging advanced techniques in randomized numerical linear algebra, we estimate Eqs. (19 25 using randomized SVD Halko et al. [2011]. For any matrix A E _ this is four step process. 1) We sample a random matrix Q NO, InxR). 2) We obtain AΩ Y E R m k 3) We form basis over the columns of Y. e.g. by taking the Q matrix in QR factorization QR Y. 4) We approximate A == Q Q..
Notably, we observe that when A is a Jacobian (or Hessian) matrix, it can be approximated purely vector and vector-Jacobian (Hessian-vector and vector-Hessian, resp.) products without ever materializing A itself. Moreover, a key result in randomized linear algebra is that this algorithm can approximate up to accuracy (mmnkrk+1) 1) (Theorem 1.1 in Halko et al. [011]
through lacotian-vector
Notably, if A has low rank structure where 3K such that the k + 1th singular value dx+1 the approximation is exact.
0, then
Matrix-Fas FFee Evaluation Inspired by matrix- techniques in numerical optimization Knoll and
Keyes [2004), we demonstrate a strategy for forming the action update (15) without materializing
the costly dim(x) x dim(x) matrices the iLQR algorithm 19 25 which we shall denote as an
SR X4
Random Inpainting
LPIPS 4
PSNR 1
SSIM 1
MSE +
LPIPS +
PSNR
SSIM
MSE +
a 0
le
7
0.173
27.49
0.794
115.9
0.050
31.80
0.879
42.96
le
0.171
27.45
0.792
117.0
0.053
31.84
0.882
42.57
&
0.172
27.43
0.799
117.5
0.050
31.85
0.891
42,47
from Lemma 4.2
0.170
27.44
0.788
117.3
0.051
31.86
0.880
42.44
Table 5: Ablative study on the effect of the Tikhonov regularization coefficient a on performance (LPIPS, PSNR, SSIM, NMSE) of our proposed model the FFHQ 256x256-18 KK dataset dataset. No results are reported for a as the algorithm encountered numerical precision errors during matrix inversion.
21

---

(85) SR X4
Random Inpainting
LPIPS 4
PSNR
SSIM
MSE +
LPIPS +
PSNR t
SSIM t
MSE 4
T = 10
0.198
27.48
0.783
125.6
0.168
27.46
0.771
123.7
T 20
0.1923
31.79
0.859
117.0
0.108
34.41
0.910
42.57
T 50
0.171
27.45
0.792
90.79
0.053
31.84
0.882
40.56
= 200 0.155
28.55
0.811
43.05
0.048
32.05
0.899
23.17
Table 6: Ablative study on the effect of ' on performance (LPIPS, PSNR, SSIM, NMSE) of our proposed model on the FFHQ 256x256-1 dataset dataset.
indexed set of matrices {A.). We do this by forming projections of each A; against corresponding set dim(x) x column-orthogome matrices {Q₂ ). which we denote as B, = Q A, These matrices then be stored reduced cost as (Q., B,) pairs
Matrix multiplications between any A,A, can then be approximated up to rank with respect the projected matrix, Q; Q... i.e.
A, = Q,
(79)
However, to prevent materialization the full size of any matrices, we drop the leading Q obtaining new projected matrix pair (Qx. Bk), where Qk Q
Adam Optimizer
Finally, we precondition gradients via the Adam optimizer Kingma and Ba
2014 before applying the feedback gains, rather than applying a backtracking line search Tassa et al. 2014 resulting in the action update
Pk: + K.(X: x(),
(80)
where P is the preconditioning matrix produced by the Adam optimizer. This reduces the overall run- time of the algorithm while still accounting for second-order information that respects the nonlinearity of the optimization landscape.
D.2 Computational Complexity Analysis
Incorporating all three modifications, we can provide a realistic runtime and space complexity analysis of our presented algorithm with respect to the rank k, the data dimension d, diffusion steps m, and number of iL.QR iterations n.
Combining both the low rank and matrix-free approximations, we obtain the updated equations for input mode perturbation (where projection matrices are written P to avoid overloading the Q function notation):
Qx = h V
(81)
Qu lu + h2
(82)
PQxxPT PQ.xx PT = PQ PT Ph V PT
(83)
PQ..pP' = Pl P
Ph V h P?,
(84)
To simplify notation, each projection matrix is the same in reality, this need not be the case.
Note that Q and Q2 are simply of size d and therefore image-sized. For all our datasets, these each take 0.2 MB to store and are therefore negligible, and we do not project these variables. When lyu is diagonal (as it is in our case), we can obtain the projected inverse for Quu as
PQ PC-!p" + -P..IP"(C- ( -PP
where C -Phh VV2 hy P
via a direct application of the Woodbury matrix inversion formula Petersen et al. (2008) which has cost O(k3 + kdd). Finally, compute the projected updates Vxx, K. well the full-precision
22

---

Vx, k terms via
k = -P"'QQ'I'"'
(86)
Vx = Qx - PPPK
Pk
(87)
PKP -PQUPP PT PVXXP" PQxxP" - PK'P"PQmPP"PKK
(88)
(89)
Where applicable, we leverage vector- Jector-Jacobinn products from standard automatic differentiation li- braries (e.g- torch.. func. vip) which have runtime complexity (11). Computing the Vx, Vxx K
terms in Eqs. 46)- (9) costs O kd²) FLOPs in terms of matrix multiplications (dominated by the matrix inverse of l x matrix q)) Crucially, it incurs O(k) neural function evaluations (NFEs), which dominates the runtime of the algorithm. Since this computation is performed for each diffusion step and iLQR iteration, the total runtime complexity our algorithm (vummk kxl))) matrix multiplication FLOPs and (nnnk) NFEs, with Omm2² - d) space complexity. In terms of time complexity, the NFEs are dominating cost, accounting for 97% of computation time.
D.3 Sensitivity Hyperparameters
6 we investigate the effect of the rank of the low rank approximation and matrix -reee projections, the Tikhonov regularization coefficient a, and the diffusion time T on the performance of our method on the FFHQ 256x256 dataset. We evaluate performance on the super-resolution and random inpainting tasks, with the same setup as in Section
In Tables
Low-Rank and Matrix-Free Rank From Table it is clear that there is a significant performance gain from even a rank one approximation of the first- and second-order matrices. The gains from subsequent increases in the rank approximation diminish quickly. This is because increasing the rank of the approximation only improves the approximation of the second-order terms. The first order Vx, Qx. Qu terms are always modeled exactly in ((1) time per iteration due to their amenability to vector-Jacobian products. From Theorems 4. |4. we see that even when the second order terms are zero (i.e.. the result of assumption & 0), exactly recover the true posterior sampler. Therefore, the second-order terms are less important, though still useful for imposing a quadratic trust-region regularization to the algorithm. Therefore, ultimately choose k 1 for three reasons:
1. the rank only affects the quadratic approximation of the iL.QR algorithm (and does not affect theoretical results in Theorems 4.1 4.3)
2. k 1 already allows second-order propagation of the quadratic trust-region regularization, and
3. subsequent increases in & have a minimal effect on the performance the algorithm.
Tikhonov Regularizer Table 5 demonstrates that our algorithm is relatively robust to the Tikhonov regularization parameter, except when a O. Under this condition, any ill-conditioning of Quu results in division by zero errors, resulting in the failure of the algorithm. Therefore, we simply choose let le 4, since the effect of Tikhonov regularizer minimal.
Diffusion Steps Finally, we observe in Table that increasing the diffusion time results in higher quality samples- - though the cost of increased computation time. Therefore, choice of T requires balancing computational cost and sample quality, and is ultimately highly user-dependent. When the computational and latency budget is relatively high, large can used to improve sample quality. Conversely, when this budget is low, we find that even T 20 provides reasonable samples.
23


Reviews:
[{"review_id": "Rrw304IrPq", "reviewer": "Reviewer_VcWe", "summary": "The paper uses tools from optimal control to introduce a novel approach for solving inverse problems with diffusion models. The authors propose reframing the generative process of diffusion models as a discrete optimal control problem allowing to leverage the iterative Linear Quadratic Regulator (iLQR) algorithm. Tackling limitations of existing probabilistic sampling methods, the resulting method demonstrates promising performance for inverse problems on FFHQ, such as super-resolution, inpainting, and deblurring.", "questions": "See \"weaknesses\" above.", "limitations": "See \"weaknesses\" above.", "rating": 4, "confidence": 3, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "While many connections between optimal control and diffusion models have been established, the proposed algorithm leverages variants of iLQR to provide a fresh perspective on training-free posterior sampling with diffusion models. The paper provides additional theoretical guarantees as well as multiple modifications (randomized low-rank approximations, matrix-free evaluations, and adaptive optimizers) to reduce computational costs. Finally, several ablations are presented for the proposed method.", "weaknesses": "1. The claims of the paper are not sufficiently supported by experiments and/or theory (the first statements in the following are just examples---similar statements can be found throughout the paper):\n\t* \"dependence on the approximation quality of the underlying terms in the diffusion process\": only a result for a single image is provided (Fig. 6). The current paper also does not seem to provide theoretical results for such robustness as claimed in \"reconstruction performance is theoretically and empirically robust to the accuracy of the approximated prior score\".\n\t* \"its sensitivity to the temporal discretization scheme\": for the baselines, again only a result for a single image is provided (Fig. 3). Moreover, the number of steps is typically reduced to accelerate the algorithm. Accordingly, methods should be compared in terms of performance vs. runtime/flops and not the number of diffusion steps. It seems that the proposed method is significantly more expensive than competing methods (in particular, since `num_iters>=50` full simulations are used).\n\t* \"its inherent inaccuracy due to the intractability of the conditional score function\": The conditional score function remains intractable, one just obtains an approximation via iLQR, since the obtained $x_0$'s obtained from the iLQR iterations only *approximately* converge to the posterior distribution *in the limit*. Statements like \"Moreover, our model always estimates $x_0$ exactly, rather than forming an approximation $\\hat{x}_0 \\approx x_0$\" sound misleading. Using iLQRs, we simulate \"nominal\" trajectories and thus iteratively obtain an approximate candidate for $x_0$ which will be used for the refinement of the control. In a similar (however, useless) fashion one could also use, e.g., DPS to obtain an estimate of $x_0$ and then run a probability flow ODE simulation where the scores are conditioned on $x_0$ (instead of $x_t$) to have a \"method [that] produces controls that coincide precisely with the desired conditional scores\". However, the advantage of DPS lies in the fact that only a single simulation is needed.\n\t* \"on several inverse problem tasks across several datasets\": Apart from a single figure on MNIST (without metrics and for only a single baseline and task), results are only provided for FFHQ.\n\n2. Moreover, several of the mentioned limitations have been already tackled by alternative approaches to posterior sampling with diffusion models, e.g., variational approaches (https://arxiv.org/abs/2305.04391) or resampling strategies (https://arxiv.org/abs/2307.08123).\n3. Finally, the appendix could provide further details on\n\t* hyperparameter choices and optimization for the baselines.\n\t* precise assumptions for the theorems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful discussion and insightful critique. We hope to clarify some of the points of the paper, and alleviate the reviewer's concerns below in a point by point response.\n\n**Only a single a single image is provided in Figs. 3 and 6.**\nIn the rebuttal PDF Figures 9 and 10, we provide further experiments that demonstrate the concepts in Figure 3 (accuracy of the predicted score) and Figure 6 (robustness to the approximated score), respectively.\n\n**No theoretical results for such robustness [to the score approximation].**\nThis theoretical statement comes from simple fact that DPS (and related algorithms) sample from p(x|y) by solving the conditional reverse diffusion process (Eq. 1), where $\\nabla \\log p_t(x_t)$ is approximated by the neural score function $s_\\theta$ (equivalently $\\epsilon_\\theta$, $v_\\theta$, etc. in noise- and v- prediction models, respectively, up to a simple reparameterization). Therefore, the performance of DPS (and related algorithms) depends on the approximation $s_\\theta \\approx \\nabla \\log p_t(x_t)$. Our algorithm makes no such assumption, and is thus robust to this approximation error. We further illustrate this point with more examples in the rebuttal Figure [x].\n\n**Methods should be compared in terms of performance vs runtime / flops and no the number of diffusion steps.**\nWe do run an experiment under an equivalent runtime/flops budget in the Ours (T=20) case in Table 1, where we restrict our method to 1000 NFEs, we still show significant gains over DPS, and comparable performance against the most recent state-of-the-art methods on FFHQ256-1K. Empirically, we find that our method has a similar runtime to DPS on an A6000 GPU (130s vs 125s). That said, we do agree that listing only the number of diffusion steps can misrepresent our method. Therefore, we have added NFEs to Table 1 (just like Table 7 in the rebuttal PDF).\n\n**The proposed method is significantly more expensive than competing methods.**\nWe maintain that our method is not significantly more expensive than competing methods. For details, see computational complexity analysis in Section D.3, where we show that iteration-for-iteration, our method has similar costs to DPS. Indeed, under the $T=50$ configuration, our method will run for more iterations. \n\n**The conditional score function remains intractable in the proposed method.**\nWe respectfully disagree. Let us first follow DPS (Chung et. al, 2022) and define the conditional likelihood as $p(y | x_0)$. Therefore, the conditional score function at time t is $\\nabla_{x_t} \\log p(y | x_t)$. Unlike DPS, we consider both deterministic and stochastic dynamics in our theory. \nUnder deterministic dynamics, we first note that the conditional **likelihood** is tractable, since $p(y | x_t) = p(y | x_0)$ — i.e., $x_0$ is determined by $x_t$, by definition — and $p(y | x_0)$ can be obtained via the forward rollout of Algorithm 1.\nNow, we have the distinct advantage over DPS where $\\nabla_{x_t} \\log p(y | x_0)$ **is** exactly backpropagated to $x_t$ (i.e., the backward pass of Algorithm 1, see Theorems 4.1-4.3). For this reason, we also maintain that the conditional **score** is tractable. This claim cannot be made by DPS, because 1) DPS is derived from stochastic dynamics and 2) DPS relies on Tweedie’s formula at each step to “approximate” the gradient of the conditional likelihood.\n\n**The statement that our model estimates $x_0$ exactly, rather than forming an approximation sounds misleading.**\nWe also respectfully disagree here. We do estimate $x_0$ exactly. Note that $x_0$ as defined in DPS (as well as our work) is $x_0 | x_t$. In other words, the denoised image **given the current noisy image $x_t$** at time $t$. Observe that we *do* obtain this term exactly in each iLQR step, up to the discretization error of the diffusion solver.\nThat said, we have added the qualification “up to the discretization error of the numerical solver” where applicable. We note that we do emphasize the discrete nature of our method in the abstract and throughout the paper, such as lines 67-68, 115, and 265.\n\n**With DPS one could also (uselessly) obtain an ODE estimate of $x_0$ and then compute the scores from there.**\nIndeed, this is possible --- but like the reviewer says, uselessly slow. The impracticality of this strategy illustrates the advantage of our method. Using iLQR, we are able to combine the prediction of $x_0$ in the forward rollout with the score computations used in the feedback step in one single sweep (Algorithm 1). To see this, observe that the scores can be computed for all timesteps with a single forward solve (and backprop) of the ODE. This is not possible for DPS, since each conditional score *must be* computed sequentially, **AFTER** the previous score has already been applied to the trajectory. If DPS were to use the reviewer’s suggested strategy, it would require on average $n / 2$ extra NFEs per update, resulting in a sum total of 500K NFEs under the current parameterization. By comparison, our method achieves superior performance to DPS using as little as 1K NFEs.\n\n**Results are only provided for MNIST and FFHQ.**\nWe additionally provide experiments on ImageNet-1K, and demonstrate that our performance holds, with similar trends to Table 1.\n\n**Please discuss alternative approaches to posterior sampling with diffusion models, e.g., variational approaches (https://arxiv.org/abs/2305.04391) or resampling strategies (https://arxiv.org/abs/2307.08123).**\nWe thank the authors for bringing up these concurrent works. We have added them to the manuscript. We note that ReSample (Song et. al, 2024) suffers from the same approximation error as DPS when estimating $x_0$ via Tweedie’s. Conversely, RED-Diff (Mardani et al., 2024) indeed does not require Tweedie’s, but requires approximations to their variational framework (e.g., the stop-gradient) when computing their proposed variational loss, thus also yielding an approximate solution.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gqiErJAYCq", "reviewer": "Reviewer_d6zE", "summary": "This paper tackles inverse problem via the perspective of optimal control. By treating the diffusion process (ODE) as a non-linear dynamic system and the extra guidance term as control signal, the authors manage to optimize the diffusion trajectory via the iterative Linear Quadratic Regulator (iLQR) algorithm. Several techniques are used to make the iLQR algorithm more efficient. This paper show good results on FFHQ dataset.", "questions": "1. Following my first point in weakness, can the author provide a comparison in sampling time (e.g. second, or NFE) of different methods. Only comparing diffusion timesteps are not very fair considering the proposed method needs to iteratively update over the whole trajectory for many times.\n\n2. Under different initializations, can the proposed algorithm always be able to find a good solution? And will the optimized results look same or different?", "limitations": "The authors has adequately addressed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea is interesting and reasonable. Using optimal control to solve the inverse problem enables us to optimize the whole sampling trajectory and avoid the error for estimating $x_0$ via Tweedie's formula. And the results on FFHQ dataset looks good.", "weaknesses": "1. High computation cost: Despite the advantages mentioned above, one obvious drawback of this method is the potential high computation cost. This includes: \n\n      a. Computing and storing the Jacobian matrices, which can be of very high dimension, can be very costly. Although the authors \n      further propose some techniques to reducing the cost, these methods might also bring extra approximation error as well as more hyper-parameters to tune; \n\n     b. Optimizing the the whole trajectory requires evaluating the whole trajectory for many times and do iterative updates. This requires more computation. Thus, though in Table 1, the authors denoted their methods as $T=50$ and $T=20$, considering the iterative update nature over the whole trajectory, this might not be directly comparable (and might actually need more computation) to other methods, which are denoted as $T=1000$. And the authors might have to greatly reduce the timesteps to make the whole algorithm affordable, this might also bring extra approximation error.\n\n\n2. Lack of more complex dataset: Though the authors achieve good performance on FFHQ dataset, considering the human face data is relatively easy (aligned, not very multimodal), it is still not very clear to me how the proposed method can work on more complex dataset, for example, on ImageNet. From my own experience, the ImageNet data can be much harder than the FFHQ human face data in the inverse problem. And considering the approximation error introduced in iLQR algorithm, computing the Jacobian matrices as well as using less timesteps, it might raise concerning regarding whether the proposed algorithm can work well on more complex dataset.\n\n3. Minor suggestion: I think it might be better for the authors to add more introduction for the optimal control part in the main paper. Or at least give more clear introduction for the notation used in 2.3. Currently, I find it not very clear to people without much background in optimal control.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude that the reviewer appreciates both the theoretical and empirical results of our work, and for bringing several insightful shortcomings of our work to our attention. Below, we address the reviewers concerns on a point-by-point basis.\n\n**High computational cost:**\n\n**a) Jacobian matrix computation can be costly, and low rank approximations bring additional error.**\nIndeed, Hessian matrices in iLQR are costly at full rank. Moreover, we agree that low rank approximations generally introduce some error to the computation. However, we maintain that this error is negligible. In our ablation study, we see that even a rank-one approximation of these matrices is sufficient, with further increases in rank providing quickly diminishing returns (see Table 4 in Appendix). Using modern randomized linear algebra libraries (e.g. randomized SVD), this brings the computation of the Hessian matrices to cost $\\mathcal{O}(d)$, where $d$ is the data dimension. Therefore, the cost of our algorithm is on par with competing methods, e.g. DPS. (For more details, please find a thorough complexity analysis in Appendix D.2.) We also find that there are few hyperparameters introduced by these approximations, and the sensitivity of the algorithm to these hyperparameters is low (Tables 4, 5, and 6).\n\n**b) The proposed algorithm is iterative over the simulated trajectory. Reported $T$ in Table 1 may not be directly comparable.**\nThe reviewer is correct, our Algorithm 1 requires multiple passes through the diffusion trajectory. Our intention with displaying the time steps $T$ in Table 1 is to show the stability of our algorithm even at very low time steps. This is a nice property to have, and can be useful in certain cases. For example, DPM-Solver (Lu et. al, 2022) is known to be unstable at >100 timesteps, and is usually used for T ~ 20. Methods like DPS will struggle in this regime.\nThat being said, we do agree with the reviewer that in terms of computation time, the number of neural function evaluations (NFEs) is also important to consider, and have added it to Table 1, as well as the new Table 7 in the rebuttal PDF. Overall, we find that our model performs favorably against other methods under an equivalent computational budget (the T=20 case in Table 1 — and 1000 NFEs), and establishes a new baseline at T=50 and 2500 NFEs.\n\n**More complex dataset than FFHQ, e.g., ImageNet.**\nWe agree that further experiments can demonstrate the generalizability of our results. To this end, we evaluate our model on the ImageNet-1K dataset replicating the experimental setup in (Chung et. al, 2022), and on some nonlinear forward operators. Please see the general rebuttal and PDF.\n\n**Minor suggestion: More introduction for the optimal control part in the main paper.**\nWe fully support this suggestion. As the reviewer has noticed, we moved a large amount of the introduction to the appendix due to page constraints. However, following this suggestion we have included a slightly more intuitive discussion on iLQR and its mechanism in Section 2.3.\n\n**Only comparing diffusion timesteps is unfair. Please provide a second form of comparison for sampling time.**\nWe agree, and again note that the “Ours (T=20)” entry in Table 1 is run for 50 trajectory iterations, yielding 1000 NFEs. This results in roughly the same sampling time (e.g., on an A6000, DPS and our method both take ~120s to run). We also maintain the same 1000 NFE budget in the new Table 7 showcasing results on ImageNet.\n\n**Robustness to different initializations.**\nWe currently initialize our actions as identically 0, and our states as i.i.d. normal (as prescribed by the DDPM algorithm). In the rebuttal Figure X we show the result of the algorithm with different state and action initializations. We currently do not fix the seed of the state initializations, and our algorithm obtain similar results on each run. Therefore we find our method to be relatively robust to different initializations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t31CKNl1by", "reviewer": "Reviewer_KZ9Z", "summary": "The paper uses the optimal control theory to solve the diffusion posterior sampling problem by iterative Linear Quadratic Regulator (iLQR) algorithm. The method could be utilized to solve both linear and nonlinear inverse problems.  Experiments on MNIST and FFHQ demonstrate the outperformance of the proposed method.", "questions": "1. As shown in Algorithm 1, the method's time complexity is  $O(T)$ in each iteration. Although the $T$ is relatively small $(=50)$ in the experiments,  num_iters $\\times T$ would be large, e.g. $50\\times 50 = 2500$ as shown in Table 2 in the appendix.  Also, the initialization of $\\{x_T'\\}$ requires $T$ NFEs (number of function evaluations). Is this correct?  How about the computational efficiency of the proposed method? I would like to see a more detailed comparison of the method with other baselines like DPS in terms of time.\n\n2. More baselines need to be compared such as [1], [2], [3] and [4]. The settings in these works might be a bit different. Some settings might be different. Can you clarify the proposed method's advantage over these baselines?\n\n3. Previous methods like DPS have done extensive experiments on both linear and nonlinear inverse problems across both FFHQ and ImageNet datasets. However. the experiments in the paper seem to be somewhat limited. I have two questions about the experiments. 1) Since there are only quantitative results for linear inverse problems (note that the results in Table 1 are all linear), can you clarify the proposed method's advantages in nonlinear problems such as phase retrieval, nonlinear deblurring, and so on? 2) Can you show more results on ImageNet, which is a broader dataset that contains more than one domain, such as faces in FFHQ?\n\n[1] Zehao Dou, and Yang Song. Diffusion Posterior Sampling for Linear Inverse Problem Solving: A Filtering Perspective, ICLR 2024\n\n[2] Morteza Mardani, Jiaming Song, Jan Kautz, Arash Vahdat. A Variational Perspective on Solving Inverse Problems with Diffusion Models. ICLR 2024\n\n[3] Jiaming Song, Arash Vahdat, Morteza Mardani, Jan Kautz. Pseudoinverse-Guided Diffusion Models for Inverse Problems. ICLR 2023\n\n[4] Zhu, Y., Zhang, K., Liang, J., Cao, J., Wen, B., Timofte, R., and Van Gool, L. Denoising diffusion models for plug-and-play image restoration. CVPR 2023", "limitations": "Yes. They mentioned limitations and impacts in the appendix, which looks good to me.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well-written, with a good summary of previous methods and their shortcomings. \n\n2. The proposed method that interprets the reverse diffusion process as an uncontrolled non-linear dynamical system is novel. Theoretical support is provided to verify the algorithm.", "weaknesses": "1. The method is well-backed but might be computationally exhausting.\n\n2. The experiments are limited. Quantitative results on different datasets and nonlinear inverse problems are lacking.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their appreciation of the optimal control perspective, and for their insightful discussion and bringing many recent works to our attention. We respond to comments in detail, on a point-by-point basis below.\n\n**The method is well-backed but might be computationally exhausting.**\nWe maintain that our method is not significantly more expensive than competing methods. For details, see computational complexity analysis in Section D.3, where we show that iteration-for-iteration, our method has similar costs to DPS. Indeed, under the $T=50$ configuration, our method will run for more iterations. Moreover, under the $T=20$ configuration in Table 1, our method has a very similar runtime as DPS (130s vs 125s on an NVIDIA A6000 GPU).\n\n**Quantitative results on different datasets and nonlinear inverse problems are lacking.**\nThank you for this suggestion. We have added experiments on ImageNet 256 x 256-1K and the phase retrieval and nonlinear blurring inverse problems, and find that our algorithm compares very favorably against existing methods.\n\n**Algorithm 1 takes 2500 NFEs in the $T=50$ case in Table 1.**\nIndeed, the reviewer is correct. We see that reporting only $T$ may not show the full picture, and have updated Table 1 to show NFEs. The new Table 7 (see main rebuttal PDF) displaying ImageNet results already shows $NFEs$.\n\n**What is the computational efficiency of the proposed method?**\nWhile the full matrix computations in equations 19-25 would be expensive to compute in full, we find that low rank approximations are very effective and result in minimal deterioration of the algorithm performance (for more details, see Appendix D.3 and specifically Table 4). Moreover, we find that the algorithm has the same computational complexity as most existing algorithms (e.g., DPS), and is similarly dominated by neural network calls (i.e., 97% of the wall-clock time is dominated by NFEs). Therefore, we also use NFEs to measure runtime.\n\n**I would like to see a more detailed comparison of the method with other baselines like DPS in terms of time.**\nUnder the $Ours (T=20)$ configuration in Table 1, we take the number of iterations to be $50$, resulting in $1000$ NFEs. Under this equivalent computational budget to other methods in Table 1, we show that we still outperform prominent methods such as DPS and DDRM, while comparing favorably to state-of-the-art algorithms. Finally, we verify that the computational budget is indeed equivalent by comparing our method to DPS and MCG, letting $T = 1000$. On an NVIDIA A6000 GPU, all three methods take ~120s to run.\n\n**Please discuss advantages over [1], [2], [3] and [4].**\nWe thank the reviewer for bringing these works to our attention. After careful study, we note that [1, 3, 4] all rely on Tweedie’s formula to predict an estimated $x_0$, given $x_t$, to formulate the correction step at each time of the reverse diffusion process. Therefore, these methods all suffer from the same pitfall of DPS (and MCG, and other algorithms of the same vein), where $\\nabla log p(y | x_0)$ must be approximated.\n\nWe do observe that [2] stands out from the other three in that does not require this approximation. However, the variational framework that is proposed is itself intractable in their inner optimization loop, and they require a stop-gradient on the score network to maintain tractability. Overall, these related works are important to discuss, and we have added these references to our manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cDxOP39xPC", "reviewer": "Reviewer_GqZF", "summary": "This paper proposes a new approach to conditional generation tasks through score-based diffusion models, with a focus on inverse problems.   As an alternative to using the likelihood $p(y | x_t)$ to guide the time-reversed SDE towards the posterior distribution, the authors reformulate this as an optimal control problem.    Starting from the ODE-based time-reversed flow for the unconditioned prior, the authors derive a controller based on the iLQR algorithm to guide the particles towards high posterior probability regions.    The authors provide theory to demonstrate that the optimal guidance coincides precisely with the desired conditional scores.  They demonstrate the method on a number of benchmarks including image inpainting and other inverse problems.", "questions": "Can the authors provide some insight on how to choose the key parameters $(m, n, k)$?  -- does the optimal control method allow substantially smaller $m$?   When does one approach become more computationally effective than the other?   I can imagine, when $m$ is sufficiently small, that PSLD, DPS will start to outperform this method with comparable computational cost.    \n\nMinor comments:  the metrics LPIP, PSNR, SSIM are reported, but at no point are these explained, or are references provided.  These are well known in some communities, but not for the wider readership.   Small typos around references were found.", "limitations": "The main limitation is the large computational cost of this methodology.   This has been identified and acknowledged.  The authors have claimed to do an equivalent budget analysis, but this maybe needed a bit more details (wall-clock time, etc).   Potential societal impacts are addressed in the Appendix.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and very clear.   The method appears novel and addresses a legitimate challenge in conditional diffusion models.  As the authors acknowledge: optimal control formulation of diffusions exist, but not (to my knowledge) in the context of guiding conditional diffusion models.     The theoretical results provide a sound justification of the validity of the approach.   The numerical results demonstrate that it is competitive in terms of accuracy compared to baseline, established methodology.", "weaknesses": "The main weakness is the sheer computational cost of the algorithm,  the need to compute very expensive hessians drastically limits the practical use of this method.    The authors suggest a number of low rank approximations to mitigate this, but it is unclear how much is lost by introducing them.      One point of question is the interplay between the number of diffusion steps $m$ and $T$.   As $m\\rightarrow \\infty$, for $T$ fixed and large we expect that the baseline conditional diffusion model will improve significantly in accuracy.  Generally, I feel that the configuration of the baseline has not been explored (or if it has, it has not been reported carefully).   Similarly, the authors claim that they have done equivalent budget analysis in the experiments -- I could not find the details of this: is it the case that the computational cost is the equivalent?  Have the author really explored the hyper-parameter space for these methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive assessment of our work, and their insightful critique. Below we provide a point-by-point response to the reviewer's discussion points.\n\n**Main weakness is the computational cost of the algorithm, e.g., computing Hessians. Approximation error is not well understood.**\nIndeed, the Hessians would be very difficult to compute in their entirety. However, we want to mention that we do investigate the computational trade-off that occurs when approximating the Hessian $V_{xx}$. Our primary method technique is the randomized low rank SVD (Halko et. al, 2011), which provides very strong approximation guarantees when there does exist low rank structure in the matrix. Indeed, the Hessians of overparameterized neural functions have been well known to be low rank (Sagun et. al, 2017), and we also verify this empirically in Table 4 in the Appendix, where we see that the algorithm performance does not deteriorate with sparser representations, in terms of rank. Moreover, even letting the rank $k = 1$ obtains very strong performance (see also Table 1). We generally found in our experiments that a low-rank approximation is sufficient to impose the quadratic trust-region optimization to our model, and does not meaningfully deteriorate performance.\n\n**Understanding the interplay between the number of diffusion steps $𝑚$ and $𝑇$.**\nWe generally find that, fixing a computational budget (e.g., 1000 NFEs), the number of diffusion steps and number of iLQR iterations can quite significantly affect the final outcome of the model. The key observation is that the iLQR objective optimizes strictly the inverse constraint, whereas the diffusion model acts as a regularizer. With a fixed budget, increasing one will come at the cost of the other. Therefore, if the measurement is very noisy, or the solution is known to be well-supported by the diffusion prior distribution, then one should take a large number of diffusion steps, and fewer number of iLQR iterations. Conversely, if the measurement is noise free, or the prior is not very informative, then the reverse should be done.\n\n**Equivalent budget analysis. Is the computational cost equivalent?**\nWe perform equivalent budget analysis in two ways. First, we note that theoretically, our Algorithm 1 costs  $\\mathcal{O}(nm(k^3 + kd^2))$ in terms of FLOPS and $\\mathcal{O}(mnk)$ in terms of neural function evaluations (NFEs). Perhaps interestingly, we found that NFEs dominate the wall-clock time of the algorithm, accounting for 97%. We let $k=1$ in all reported experiments outside of our ablation study in Table 4. More details can be found in Appendix D.2.\n\nEmpirically, the results of our equivalent budget analysis is summarized in the Ours (T=20) entry Table 1. As discussed, neural network calls take up ~97% of the wall-clock runtime of our solver. Therefore, we measure computational budget by the number of neural function evaluations (NFEs) of the method. DPS, MCG, DDNM, and DDRM allow a computational budget of 1000 NFEs, and most other methods in Table 1 follow a similar guideline. Therefore, the entry “Ours (T=20)” runs with 50 iLQR iterations and T=20, satisfies an equivalent budget to these models. We see that under this restricted budget our method beats out several prominent methods, including DPS, while comparing favorably against other state-of-the-art methods.\n\n\n**Have the author really explored the hyper-parameter space for these methods. Can the authors provide some insight on how to choose the key parameters (𝑚,𝑛,𝑘)?**\nPlease see Tables 4, 5, and 6 for an in-depth exploration of the main hyperparameters of our method.\n\nWe assume the reviewer is using the notation from Appendix D.2 for $(m, n, k)$. Fixing a computational budget $\\mathcal{N}$, there is a clear trade-off between the number of diffusion steps $m$ and iLQR iterations $n$, since they are inversely related. More diffusion steps bias the model towards the diffusion prior, whereas more iLQR iterations bias the model towards the observation $y$.\n\n**When $m$ is small, DPS and PSLD should outperform this method with comparable cost.**\nWe in fact observe the reverse phenomenon with $m$! In other words, our model actually outperforms DPS at low $m$. We conduct the simple comparison below on the super-resolution task on FFHQ. We fix $n = 50$, and report LPIPS.\n| Algorithm / $m$| 1 | 5 | 10 | 25 | 50 | 100 |\n|-|-|-|-|-|-|-|\n| Ours | 0.643 | 0.472 | 0.277 | 0.185 | 0.171 | 0.169 |\n| DPS | 0.927 | 0.799 | 0.485 | 0.395 |  0.354 |  0.331 |\n\n**The metrics (e.g., LPIPS, PSNR, SSIM) are mentioned but not described.**\nThank you for pointing this out. We agree that this is can be confusing and have added references to the LPIPS, PSNR, and SSIM metrics.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h6DcTUtfHv", "reviewer": "Reviewer_u6BH", "summary": "This paper proposes diffusion optimal control that solves inverse problems via posterior sampling by combining the power of a pre-trained unconditional diffusion model and the iterative Linear Quadratic Regulator algorithm to produce optimal controls that steer the reverse diffusion process to correctly recover the original signal. \nThe framework is general and able to handle any differentiable forward measurement operator and establishes a new baseline in image reconstruction.", "questions": "* What is the rough runtime of each method in Table 1?\n* Notations such as $p_t(x|y)$ are confusing to me. The subscript $t$ on $p$ suggests a family of distributions but I think that's not the case. My understanding of the randomness is the following. First, the random variable $x_0$ is drawn from the clean image distribution. Then $y = A(x_0) + \\eta$. In parallel, we also have random variables $x_t$ obtained deterministically from $x_0$ by evolving along the ODE of the forward diffusion. A better notation in my opinion is to put the subscript $t$ on $x$, like $p(x_t|y)$.\n* In (29), what is the meaning of $\\nabla_{x_t} \\log p(y|x_0)$?\n* Line 149, what do you mean by \"produces a feasible solution\"? What's the meaning of being feasible?\n* The texts around line 160 are hard to parse. What is the notation $p(x_0|x_t,x_0)$? What is the $x_0$-centered marginal?\n* Line 176, why is $\\log p(y|x_t) = \\log p(y|x_0)$ an assumption, not a consequence?\n* In the paragraph of Line 204, I'm confused about why the diffusion model can be taken to have randomly initialized weights. This does not seem to result in any meaningful application, since there is no information about the image distribution. For instance, in Figure 6, the produced result looks even worse than the input $y$.\n\nMinor comments\n* Line 82, what is $\\theta$?\n* In Section 2.3, it would be good to include the dimension of each variable. In (13), it would be good to clarify that $k, K$ all depend on $t$.\n* Algorithm 1 appears unreferenced in the main text. In the input, what is $x_T$? Is it just drawn from a Gaussian? What is the output?\n* Line 117, $\\ell_t(x_t,u_t) = 0$, not just not depend on $x_t$. It would also be good to say here what exactly is $\\ell_0$ for input/output perturbation controls instead of defining it later. \n* Theorem 4.1 is missing transitions for presenting the conclusion (29).\n* The values of $\\alpha$ in the two theorems appear to be missing a factor of $2$.\n* In (32), it should be $\\widetilde{x}_{t-1}$ on the left side of the equation.\n* Line 253, there is a parenthesis not closed.", "limitations": "There's not much discussion about the limitations in the paper. There is no potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* The idea of augmenting the reverse diffusion sampling process with a perturbation control is quite novel and general for arbitrary cost functions, although the paper focuses specifically on the cost for posterior sampling.\n* The writing is generally good (see more comments regarding writing in questions). It is concise and to the point with good intuitions provided.\n* Efforts (e.g. low-rank approximations) have been made to bring down the computational cost in iLQR for the high-dimensional image setting.\n* The empirical performance of the proposed method is strong and establishes new state-of-the-art results.", "weaknesses": "* The runtime of the proposed method seems high and is not much discussed. iLQR is a global-in-time iterative method that could potentially require a large number of iterations to converge (and all nice things discussed in Section 4. rely on this convergence). On top of that, for each iteration, there needs to be $\\Omega(T)$ matrix solves which can be quite slow given the dimension of the images (even with techniques like low-rank approximation).  It would be interesting to see ablation studies on the effect of num_iter in Algorithm 1. It would be also more convincing to report the runtime of each method in Table 1.\n* There is no analysis of the approximation error of iLQR (the first and second-order Taylor approximations) in the studied setting. Specifically, it seems to me that a lot of heavy lifting is done when the control $u_t$ is designed to be only a perturbation of the reverse diffusion step. For instance, does this imply that the value function is smoother (hence the Taylor approximation is more accurate) when parameterized by $u_t$?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer’s positive assessment of theoretical contributions, for their in-depth reading of our manuscript, and for their suggested improvements. We respond to comments in detail below.\n\n**The runtime of Algorithm 1 seems high.**\nWithout taking any approximations of the Hessian and Jacobian matrices in Eqs. 19-25, Algorithm 1 will indeed be relatively costly, compared to the other algorithms we compare against in Table 1. However, we maintain that our method is not significantly more expensive than competing methods. For details, see computational complexity analysis in Section D.3. Indeed, under the $T=50$ configuration, our method is slightly more expensive than competing methods. That being said, we run an equivalent runtime/flops budget in the Ours (T=20) case in Table 1, where we restrict our method to 1000 NFEs, which take up the majority (~97%) of the runtime of our algorithm. We still show significant gains over DPS, and comparable performance against the most recent state-of-the-art methods on FFHQ256-1K. Empirically, we find that our method has a similar runtime to DPS on an A6000 GPU (130s vs 125s). \n\n\n**iLQR could require a large number of iterations to converge.**\nIndeed, iLQR could take many iterations to converge. However, we observe near convergence for nearly all of the settings we consider in our work. Moreover, true convergence is not necessarily desirable, since the forward operator measurements $y$ are noisy. Therefore, a fully converged iLQR method would overfit to the noisy, and produce inferior solutions.\n\n\n**...(and all nice things discussed in Section 4. rely on this convergence).**\nWe respectfully disagree. A central claim in our paper, and in Theorems 4.1-4.3 in Section 4 is the ability to compute the conditional score function $\\nabla_{x_t} \\log p(y | x_0)$. This claim does not rely at all on the convergence of the iLQR algorithm. In fact, this quantity is a fundamental property of the backward pass of the iLQR algorithm (the second loop in Algorithm 1). Therefore, we obtain the true $\\nabla_{x_t} \\log p(y | x_0)$ in every pass of our the iLQR algorithm, including the first step.\n\n**Furthermore, each iteration needs to be $\\Omega(T)$ matrix solves which can be quite slow.**\nThis is true. We agree with the reviewer's statement. However, the constant swallowed by the $\\Omega$ notation is important, and rather small here. After an ablation study, we found (somewhat surprisingly) that even a rank-1 approximation of the relevant matrices is sufficient to obtain competitive results on our benchmarks (Table 4). Under a rank-1 approximation, many of the matrix operations are simply $\\mathcal{O}(d)$, where $d$ is the size of the image.\n\n**It would be interesting to see ablation studies on the effect of num_iter in Algorithm 1.**\nThis is a good point. We provide the ablation study on the FFHQ256-1K dataset with the super-resolution task, letting $T=50$ and $k=1$. Performance is evaluated with the LPIPS metric.\n| num_iter | 1 | 5 | 10 | 25 | 50 |\n|-|-|-|-|-|-|\n| | 0.491 | 0.411 | 0.322 | 0.236 | 0.171 |\n\n**No analysis of the approximation error of iLQR (the first and second-order Taylor approximations) in the studied setting.**\nWe respectfully disagree. We do analyze the approximation quality of our proposed iLQR Algorithm 1 in Theorems 4.1 and 4.3. Here, we show that, in the deterministic setting, the solution $\\mu_t$ is precisely the desired conditional score $\\nabla \\log p(y | x_0)$, which is the quantity desired in DPS-based solvers at each step. Therefore, in the deterministic setting, the approximation error arises entirely from the discretization error of the numerical solver for $x_0$ itself. In the stochastic setting, the approximation error will also come from the randomness of the diffusion process.\n\n**What is the rough runtime of each method in Table 1?**\nBelow we report the approximate runtimes for each method on an NVIDIA A6000 GPU. The first two are ours with different choices of $T$.\n| (T = 50) | (T=20) | DPS | MCG | PSLD | DDNM | DDRM | Score-SDE |\n|-|-|-|-|-|-|-|-|\n| 259s | 130s | 125s | 123s | 251s | 122s | 35s | 50s|\n\n**Elaboration on $p_t(x|y)$ notation.**\nIn general, $p_t(x|y)$ and $p_t(x)$ *do* come from a family of distributions, and refer to the marginal distribution of the conditional diffusion process at time $t$. Each $x_t \\sim p_t(x|y)$ can be obtained, given $x_0$, via $x_t = \\alpha x_0 + \\sqrt{1 - \\alpha} \\epsilon$, where $x_0 \\sim p(x | y)$ (the distribution of our solution set), and $\\epsilon \\sim \\mathcal{N}(0, \\mathbf{I})$.\n\n**In (29), what is the meaning of $\\nabla_{x_t} \\log⁡ 𝑝(𝑦|𝑥_0)$?**\nWe apologize for the confusing notation. This is the derivative of the conditional log likelihood $\\log p(y | x_0)$, which depends on $x_0$, which in the reverse diffusion process depends on $x_t$. We thank you for this question and have clarified this notation in the text immediately following (29).\n\n**What is a feasible solution in Line 149?**\nWe mean feasibility in the constrained optimization sense – a solution is feasible when it satisfies all the constraints of the problem, i.e., $y \\approxeq f(x)$. Since our proposed algorithm is closed loop (i.e., changes at each step $t$ depend on the final error at $t=0$), the algorithm can make global modifications to the diffusion trajectory at each $t$. To our knowledge, this is not possible for general inverse problems with competing methods like DPS, MCG, etc.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uzBjyGYbx6", "reviewer": "Reviewer_L2Ru", "summary": "The paper addresses the limitations of existing diffusion-based inverse problem solvers, which typically frame signal recovery as a probabilistic sampling task. The authors propose a novel approach that redefines the generative process as a discrete optimal control task. Inspired by the iterative Linear Quadratic Regulator  algorithm, this new framework named diffusion optimal control, can handle various differentiable forward measurement operators, including super-resolution, inpainting, and deblurring.", "questions": "1. What is the purpose of injecting control vectors u_t into the reverse diffusion process, and how do they influence the terminal state of the system?\n\n2. How are the gradients V_x and Hessians V_xx of the value function used within the optimal control framework, and what is their significance?\n\n3. Please indicate what is output of Algorithm 1.\n\n4. Can you give some high level description about: How does using an adaptive optimizer for action updates improve the iLQR algorithm, and what impact does it have on the performance and efficiency of solving inverse problem tasks?", "limitations": "The method requires the forward measurement operator to be differentiable. In cases where this is not possible or practical, the applicability of the proposed framework may be limited.\n\nThe performance evaluations rely on specific pretrained models. The method’s robustness and performance with other pretrained models, or models trained on different data distributions, would need further investigation.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper introduces a novel framework based on optimal control theory to solve diffusion-based inverse problems, moving away from the traditional probabilistic sampling approaches. This is a significant theoretical advancement.\n2. The framework addresses critical drawbacks of current methods, such as the intractability of the conditional likelihood function and dependence on score network approximations. This leads to more robust and potentially more accurate solutions.", "weaknesses": "1. The method involves complex mathematical formulations and optimal control theory, which may pose challenges for implementation and understanding by practitioners who are not familiar with these concepts. The need to compute Jacobian and Hessian matrices, as well as the regularized inverses, may lead to significant computational demands, particularly in high-dimensional settings.\n\n2. Lacking of enough experiments, such as MRI reconstruction or other medical images. Including more diverse datasets and additional baseline methods would provide a more comprehensive evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We highlight our gratitude that the reviewer appreciates both the theoretical and empirical results of our work, and for bringing several insightful shortcomings of our work to our attention. Below we respond to the reviewers concerns in a point-by-point basis.\n\n**Mathematical formulations are complex and potentially challenging to implement.**\nWe agree that the mathematical framework behind optimal control theory is sophisticated, and differs from that of standard diffusion models, which comes from nonequilibrium thermodynamics. However, our approach has several distinct advantages. First, we can treat the generative process as a black box dynamical system, which vastly abstracts the pre-existing mathematics in the reverse diffusion process. Practitioners may thus choose between standard diffusion modeling and optimal control-based modeling based on their expertise. Second, we can leverage extensive research from the optimal control community to improve diffusion-based inverse solvers. Third, we are able to sidestep the intractability of the conditional score function, as discussed in Section 4.\n\n**Computing Jacobian and Hessian matrices and their inverse may have significant computational demands.**\nWe note that even a rank-one approximation of these matrices is sufficient, with further increases in rank providing quickly diminishing returns (see Table 4 in Appendix). Using modern randomized linear algebra libraries (e.g. randomized SVD), this brings the computation of the Hessian matrices to cost $\\mathcal{O}(d)$, where $d$ is the data dimension. Therefore, the cost of our algorithm is equivalent to the cost of, e.g., performing a DPS step. (For more details, see a thorough complexity analysis in Appendix D.2.)\n\n**Including more diverse datasets would provide a more comprehensive evaluation.**\nThank you for this insight. We agree, and additionally include results on ImageNet, and more nonlinear settings. Please see the main rebuttal and PDF.\n\n**Purpose and mechanism of injected control vectors in the reverse diffusion process.**\nThe control vectors u_t are simply perturbations to the original unconditional diffusion process, and are widely used in conditional sampling with diffusion models. Analogues include the conditional score term in DPS (Chung et al., 2022), the classifier gradient in classifier guidance (Dhariwal and Nichol, 2021), or the classifier-free guidance term (Ho and Salimans, 2022).\n\n**How are $V_x$ and $V_{xx}$ used in the optimal control framework?**\nIn optimal control, $V(x)$ is the value function, which intuitively represents \"desirability\" of each state x in the dynamical system to the user. In this case, $x$ is the reverse diffusion iterate. $V_x$ and $V_{xx}$ are simply derivatives of this value function. We can think of $V(x)$ as the negative loss of the optimal control system. Therefore we use the derivatives $V_x$ and $V_{xx}$ much like they are used in an optimization framework to guide a solver towards (locally) optimal solutions.\n\n**Please indicate the output of Algorithm 1.**\nThe output of Algorithm 1 is the perturbed dynamics ${x’_t}_{t=1}^T$. The final iterate of which is the solution to the inverse problem. Thank you for this comment, we have clarified this point in the manuscript.\n\n**High level intuition on the adaptive optimizer for action updates.**\nNote that the default choice of optimizer for iLQR implementations is the backtracking line search. We instead use an adaptive optimizer, and motivate this design decision with three main observations.\n\nFirst, recall that given the update direction $v$ the backtracking line search finds the optimal step size $\\lambda$ to update the actions by applying $u’ = u + \\lambda v$ to our reverse diffusion process. Evaluating each $\\lambda$ requires re-computing the inverse problem loss under the new proposed $u’$, we would have to re-run the reverse diffusion process, which is quite expensive.\n\nSecond, the backtracking line search is used with generally smooth landscapes, and not as amenable to highly nonconvex settings. Since each marginal density of the diffusion process is essentially the data distribution convolved with a Gaussian distribution, it is highly nonconvex (especially at the low noise regime) and thus ill-suited for our setting.\n\nThird and last, since the diffusion process occurs on images which can be very high dimensional, our optimization space is also very high-dimensional. Adaptive optimizers such as Adam, Adagrad, RMSprop, etc. are specially designed for deep learning settings where there is a similarly high dimensional optimization space. Therefore, it is a natural choice for our problem setting.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their thoroughness and diligence in reading our manuscript. We received a lot of sound, constructive criticism and positive feedback. This guided our revisions and further experiments in this rebuttal period, and we believe that the paper has meaningfully improved as a result. Below, we summarize some common assessments (good and bad) shared by multiple reviewers. We then address each reviewer's concerns on a point-by-point basis below.

**Strengths:**

All reviewers noted the novelty of our application of the well-known iLQR algorithm in optimal control theory to diffusion modeling. Reviewers VcWe, KZ9Z GqZF, and L2Ru found the theoretical contributions sound and compelling. KZ9Z, GqZF, and u6BH found the writing generally clear, and all reviewers provided suggestions that further improved the readability of our work.

**Weaknesses:**

**Lack of further empirical experiments.**
Many reviewers felt that, while performance is very strong on FFHQ, it is important to validate the performance of our algorithm on other datasets and tasks so as to understand the generalizability of the approach. We considered this feedback and extend the empirical results of our algorithm on two fronts: **ImageNet experiments** and **further nonlinear tasks**. In both settings, we constrain the number of function evaluations (NFEs) of our algorithm to 1000, in line with other works in these settings.

In Table 7 and Figure 7 in the attached rebuttal PDF, we demonstrate very favorable performance on ImageNet compared to existing works. We note that some works (e.g. PSLD) cannot be applied to ImageNet, since there do not exist unconditional latent ImageNet models, and are therefore absent from our benchmark. 

In Table 8, we compare our model to DPS on two nonlinear inverse problems, phase retrieval and nonlinear deblurring following the setup in DPS (Chung et. al, 2022), and find that we improve on the baselines set by DPS on the phase retrieval task, and are on par with DPS on the nonlinear deblurring task.

Finally, we further corroborate our claim that our proposed algorithm is more robust to approximation errors in the neural score function. We demonstrate an out-of-domain inverse problem setting, where the model is tasked to recover an image that is out-of-distribution to the training distribution of the pretrained model. In Figure 10 in the rebuttal PDF, the right hand block shows an inverse problem setting for reconstructing ImageNet images, where the diffusion model is initialized with a pretrained model on the FFHQ dataset. As expected, our proposed method outperforms DPS on this out-of-domain inverse problem setting.

Overall, these results round out the empirical evaluation of our work, demonstrating that our proposed algorithm is capable of obtaining high quality solutions to inverse problems on a bevy of datasets and problem settings.

**Concerns on computational cost.**
Reviewers were also concerned by the computational cost of the algorithm. Indeed, optimal control was designed for generally simpler, low-dimensional systems, and many computations in a vanilla iLQR algorithm are  the Hessians would be very difficult to compute in their entirety. However, we want to mention that we do investigate the computational trade-off that occurs when approximating the Hessian $V_{xx}$. Our primary method technique is the randomized low rank SVD (Halko et. al, 2011), which provides very strong approximation guarantees when there does exist low rank structure in the matrix. Indeed, the Hessians of overparameterized neural functions have been well known to be low rank (Sagun et. al, 2017), and we also verify this empirically in Table 4 in the Appendix, where we see that the algorithm performance does not deteriorate with sparser representations, in terms of rank. Moreover, even letting the rank $k = 1$ obtains very strong performance (see also Table 1). We generally found in our experiments that a low-rank approximation is sufficient to impose the quadratic trust-region optimization to our model, and does not meaningfully deteriorate performance.

Many also noted that we only report $T$, not $NFEs$ in Table 1, which can be misleading about the runtime of our algorithm. Our original intention with displaying the time steps $T$ in Table 1 is to show the stability of our algorithm even at very low time steps. This is a nice property to have, and can be useful in certain scenarios. For example, consistency (Song et. al, 2023), progressively distilled (Salimans et. al, 2022) or adversarially finetuned diffusion models (e.g., Imagine-Flash, Kohler et. al), or even regular diffusion models sampled with certain solvers (e.g. DPM-Solver, Lu et. al, 2022) are known to be unstable at larger $T$. Methods like DPS, PSLD, MCG, etc., will thus struggle in this regime, and fail to find good solutions to the inverse problem at hand. That being said, we do agree that listing only $T$ is a somewhat unfair representation of our algorithm, and have adjusted the notation in Table 1 in line with Table 7 in our rebuttal PDF. In particular, Ours ($T=20$) now reads Ours (NFE = 1000) and Ours ($T = 50$) now reads Ours ($T = 2500$), and all other methods in Table 1 also list NFEs, where applicable.

Meta-review: After reading the reviews comments and the original paper,

Main strengths: 
1. Novel idea. This manuscript proposed an optimal control perspective on solving linear inverse problems with diffusion models, which is an interesting and novel idea. The use of iLQR is very interesting and well-motivated. 
2. Good empirical results. The empirical results show that the proposed method can achieve impressive reconstruction performances in terms of FID, LPIPS, as shown in Table 1. 

Main weaknesses: 
1. Theoretical arguments not sound and rigorous. 
2. Insufficient comparisons. More baseline methods and data should be incorporate to validate the method. Moreover, while FID and LPIPS are good in Table 1, the performances in terms of SSIM are much worse compared to other methods. 

Overall, given that the idea of using optimal control is novel and the performance is relatively good, I recommend acceptance but the authors should address the concerns raised in the rebuttal and discussions.
Decision: Accept (poster)