Solving Inverse Problems via Diffusion Optimal Control

Authors: henry  li
marcus  pereira
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:22:31

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.

Main paper: 
 introduction: 
 Diffusion models Song and Ermon [2019] , Ho et al. [2020] have been shown to be remarkably adept at conditional generation tasks Dhariwal and Nichol [2021] , Ho and Salimans [2022] , in part due to their iterative sampling algorithm, which allows the dynamics of an uncontrolled prior score function ∇ x log p t (x) to be directed towards an arbitrary posterior distribution by introducing an additive guidance term u. When this guidance term is the conditional score ∇ x log p t (y|x), the resulting sample is provably drawn from the desired conditional distribution p(x|y) Song et al. [2020a] .A central obstacle to this framework is the general difficulty of obtaining the conditional score function ∇ x log p t (y|x t ) due to its dependence on the noisy diffusion variate x t rather than just the final sample x 0 Chung et al. [2023a] . In large-scale conditional generation tasks such as class-or text-conditional sampling the computational overhead of training a time-dependent conditional score function from scratch is deemed acceptable, and is indeed the approach taken by Rombach et al. [2022] , Saharia et 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 a zero-shot capacity for an arbitrary forward model. This bottleneck has spawned a flurry of recent research dedicated to approximating the conditional score ∇ x log p t (y|x t ) as a simple function of the noiseless likelihood log p(y|x 0 ) Choi et al. [2021] , Chung et al. [2022] , Rout et al. [2024] , Chung et al. [2023a] , Kawar et al. [2022] , Chung et al. [2023b] . However, as we will demonstrate in this work, these approximations impose a significant cost to the performance of the resulting algorithm.Figure 1 : Conceptual illustration comparing a probabilistic posterior sampler to our proposed optimal control-based sampler. In a probabilistic sampler, the model relies on an approximation x0 ≈ x 0 to guide each step (left). We are able to compute x 0 exactly on each step, resulting in much higher quality gradients ∇ log p(y|x 0 ) and an improved trajectory update (right).To address these issues, we propose a 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 Song et al. [2020a] , 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 of 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 diffusionbased inverse problem solvers. In particular, the idealized posterior sampling score Song et al. [2020b] -approximated by existing methods -can be recovered exactly as a 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 × 256 dataset. 
background: 
 Notation We use lowercase letters for denoting scalars a ∈ R, lowercase bold letters for vectors a ∈ R n and uppercase bold letters for matrices A ∈ R m×n . Subscripts indicate Jacobians and Hessians of scalar functions, e.g. l x ∈ R n and l xx ∈ R n×n for l(x) : R n → R, respectively. We overload notation for time-dependent variables, where subscripts imply dependence rather than derivatives w.r.t. time, e.g., x t = x(t). Furthermore, V (x t ) and Q(x t , u t ) are scalar functions despite being uppercase, in line with existing optimal control literature Betts [1998] .The diffusion modeling literature uses the following reverse-time Itö SDE to generate samples Song et al. [2020b] ,EQUATIONwhere x t ∈ R n is the state vector, f : R n → R n and g : R → R are drift and diffusion terms that can take different functional forms (e.g., Variance-Preserving SDEs (VPSDEs) and Variance-Exploding SDEs (VESDEs) in Song et al. [2020b] ), ∇ xt log p t (x t ) is the score-function and w t ∈ R n is a vector of mutually independent Brownian motions. The above SDE has an associated ODE called the Figure 2 : Predicted x 0 used in a probabilistic framework (above) compared to ours (below) for a general diffusion trajectory. The full forward rollout in our proposed framework allows for the predicted x 0 (and therefore ∇ xt log p(y|x 0 )) to be efficiently computed for all t = 0, . . . , T .probability-flow (PF) ODE given byEQUATIONwith the same marginals p t (x t ) as the SDE, which allow for likelihood computation [Song et al., 2020b , Li et al., 2024] . All practical implementations of diffusion samplers require a timediscretization of the PF-ODE. One such discretization is the well-known Euler-discretization which gives,EQUATIONwhere, ∆t 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 Euler-discretization and any high-order discretization techniques can also be employed. More concisely, we have,x t-1 = h(x t ), where h : R n → R n (4)which describes the general non-linear dynamics of the corresponding discrete-time diffusion sampler.Inverse problems are a general class of problems where an unknown signal is reconstructed from observations obtained by a forward measurement process Ongie et al. [2020] . The forward process is usually lossy, resulting in an ill-posed signal recovery task where a unique solution does not exist. The forward model can generally be written asEQUATIONwhere A : R n → R d is the forward operator, y ∈ R d the measured signal, x 0 ∈ R n the unknown signal to be recovered, and η ∼ N (0, σI d ) the noise (with variance σ 2 ) in the measurement process.Given the forward model Eq. ( 5) and a measurement y, sampling from the posterior distribution p θ (x|y) can then be performed by solving the corresponding conditional Itö SDEEQUATIONwhere, invoking Bayes rule,EQUATIONAs with the unconditional dynamics, Eq. ( 6) has a corresponding ODEEQUATIONwhich has an approximate solution obtained by the Euler discretizationEQUATION)Figure 3 : Inverse problem solution as a function of total diffusion timesteps T for the 4× 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 .Optimal control is the structured and principled approach to the guidance of dynamical systems over time. Many methods have been developed in the optimal control literature and are popularly referred to as trajectory optimization algorithms Betts [1998] . Perhaps the most well-known is the Iterative Linear Quadratic Regulator (iLQR) algorithm which uses a first-order approximation of the dynamics and second-order approximations of the value-function Li and Todorov [2004] .Formally, let us define an arbitrary user-defined global cost functionEQUATIONcomposed of a sum over scalar-valued running and terminal cost functions ℓ t and ℓ 0 . Optimal control theory dictates that the value function V (x t , t) := min {un} n=1 n=t J t satisfies the following recursive relation also known as Bellman's Principle of OptimalityEQUATION)The iLQR algorithm centers around approximating the state-action value function, Q(x t , u t ) := ℓ t (x t , u t ) + V (x t-1 , t -1), (12) from which the value function can be recovered as V (x t , t) = min ut Q(x t , u t ).Then given a state transition function x t = h(x t+1 , u t+1 ) where we crucially note that we have defined time to flow backwards from t = T, . . . , 0, the iLQR algorithm has feedforward and feedback gainsk = -Q -1 uu Q u and K = -Q -1 uu Q ux (13)The update equations can be written asEQUATION) Given the feedforward and feedback gains {(K t , k t )} T t=0 and x0 := x 0 , we can recursively obtain the locally optimal control at time t as a function of the present states x t and controls u t asEQUATION) u * t = u t + λk + K( xt -x t ).(16) For a more detailed treatment of iLQR as well as a derivation of the equations, please see Appendix B. 
diffusion optimal control: 
 We motivate our framework by observing that the reverse diffusion process Eq. ( 1) is an uncontrolled non-linear dynamical system that evolves from some initial state (at time t = T ) to some terminal state (at time t = 0). By injecting control vectors u t into this system we can influence its behavior and hence its terminal state (i.e., the generated data) to sample from a desired p(x|y). There are two obvious ways to inject control into this process:Input: λ, T, y, x T Initialize u t , k t , K t as 0 for t = 1 . . . T , {x ′ t } T t=0 as uncontrolled dynamics for iter = 1 to num_iters doV x , V xx ← ∇ x0 log p(y|x 0 ), ∇ 2 x0 log p(y|x 0 ) ▷ Initialize derivatives of V (x t , t) for t = 1 to T do Compute k t , K t , V x , V xx ▷ See Eqs. (13), (14) end for for t = T to 1 do x t-1 ← h(x t , λk t + K t (x t -x ′ t )) ▷ Update x t-1 with new u t x ′t ← x t end for end for 1. In input perturbation control, we apply the u t before the diffusion step:EQUATION2. In output perturbation control, u t is applied after the diffusion step:EQUATION)Observe that iLQR is 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 at time t = 0i.e., ℓ t (x t , u t ) does not depend on x t .In the case of input perturbation control, we observe from Eq. ( 17) that h x = h u , whereas output perturbation control implies that h u = I, resulting in the left and right equations, respectively:EQUATION)EQUATION)EQUATIONEQUATION)EQUATIONThe derivatives of V can then be backpropagated using the following equations:EQUATIONIn high dimensional systems such as Eq. 3, matrices may be singular. Therefore, a Tikhonov regularized variant of iLQR is often employed, where matrix inverses are regularized by a diagonal matrix αI Tassa et al. [2014] .Compared to the dynamics in traditional application areas of optimal control, those we 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 x , h u and the second-order derivative matricesV xx , Q xx , Q ux , Q xu ,and Q uu are particularly expensive to compute, store, and perform downstream operations against. For example, in a three-channel 256 × 256 image, these matrices naively contain (256 × 256 × 3) 2 ≈ 39B parameters. In Appendix D.1 we propose and analyze three modifications to the standard iLQR 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 on inverse problem solving tasks. 
improved posterior sampling: 
 We demonstrate that our optimal control-based sampler overcomes several practical obstacles that plague existing diffusion-based methods for inverse problem solvers.In a probabilistic framework, solutions to inverse problems incur a discretization error from the numerical solution of Eq. ( 8) that decays poorly with the total diffusion steps T of the diffusion process. While much research has been conducted on the acceleration of unconditional diffusion processes Song et al. [2020a] , Jolicoeur-Martineau et al. [2021] , Karras et al. [2022] , Meng et al. [2023] , sample quality appears to decay much more aggressively in diffusion-based inverse problem solvers (Figure 3 ).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 ∇ x log p(y|x t ) decays quickly with time (Figure 2 ), and so fewer timesteps lead to fewer chances at low t to correct errors made at high t. On the other hand, since optimal control directly casts the discretized process as an end-to-end control episode, it produces a feasible solution for any number of discretization steps T .Intractability of ∇ xt log p(y|x t ) When the forward model A is known and η comes from a simple distribution, the conditional likelihood p(y|x t ) can be derived in closed form for t = 0. On the other hand, the dependence of y on x t for t > 0 is generally not known without explicitly computing x 0 , which requires sampling from the diffusion process. Ultimately, obtaining the conditional score term ∇ xt log p(y|x t ) is a highly nontrivial task Song et al. [2020b] .To sidestep this issue, many works Meng and Kabashima [2022] , Song et al. [2022] , Chung et al. [2023a] factorize this term as the integralp(y|x t ) = p(y|x 0 )p(x 0 |x t )dx 0 (26)and then apply a series of approximations to recover a computationally feasible estimate of the conditional score. First, the marginal p(x 0 |x t ) is replaced by the marginal conditioned on x 0 , i.e.p(x 0 |x t , x 0 ) = N (x 0 , σ 2 I) Kim and Ye [2021] . Next, the x 0 -centered marginal is replaced by the posterior mean E[x 0 |x t ] given by Tweedie's formula Efron [2011] . Finally, the true score is replaced by the learned score network.While these approximations are necessary in a probabilistic framework, we show that they are not required in our method. Intuitively, this is because the linear quadratic regulator backpropagates the control cost log p(y|x) through a forward trajectory rollout, which naturally computes the true conditional score at each time t. Moreover, our model always estimates x 0 |x t exactly (up to the discretization error induced by solving Eq. 3), rather than forming an approximation x0 ≈ x 0 (Figure 2 ). We formalize this observation with the following statement.Theorem 4.1. Let Eq. 3 be the discretized sampling equation for the diffusion model with output perturbation mode control (Eq. 18). Moreover, let the terminal costℓ 0 (x 0 ) = -log p(y|x 0 ) (27)be twice-differentiable and the running costsEQUATION)Then the iterative linear quadratic regulator with Tikhonov regularizer α produces the controlEQUATIONIn other words, by framing the inverse problem as an unconditional diffusion process with controls u t , our proposed method produces controls that coincide precisely with the desired conditional scores ∇ xt log p(y|x 0 ).Let us further assume that log p(y|x t ) = log p(y|x 0 ), i.e., x t contains no additional information about y than x 0 . This assumption results in the posterior mean approximation in Chung et al. [2023a] under stochastic dynamics (Eq. 1), where we additionally obtain exact computation of x 0 , rather than x0 ≈ x 0 via Tweedie's formula Kim and Ye [2021] . Under the deterministic ODE dynamics (Eq. 2), we recover the true posterior sampler under appropriate choice of Tikhonov regularization constant α.Lemma 4.2. Under the deterministic sampler with output perturbation mode control, α = 1 g(t) 2 ∆t recovers posterior sampling (Eq. 9).We demonstrate a similar result with input mode perturbation. Theorem 4.3. Let Eq. 3 be the discretized sampling equation for the diffusion model with input perturbation mode control (Eq. 17). Moreover, letEQUATIONand the running costsEQUATION)Then the iterative linear quadratic regulator with Tikhonov regularizer α = 1 g(t) 2 ∆t produces the dynamical sytemEQUATIONwhere x t := x t + u t .Observe that Eq. ( 32) can be understood as a predictor-corrector sampling method, where the predictor produces an unconditional reverse diffusion update and the corrector produces a conditional correction step on the intermediary variablex t = xt -u t .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. Certainly, improved approximation of the score terms result in a 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(y|x 0 ) when the terms composing the posterior sampling equation (Eq. ( 8)) are well approximated (Figure 6 ). Modeling errors can occur even in foundation models. There exist several methods that seek to alleviate the errors incurred by Tweedie's formula (being a 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. 
related work: 
 The recent success of diffusion models in image generation [2022], Chung et al. [2022 Chung et al. [ , 2023a,b],b] , Kawar et al. [2023] for the true conditional diffusion dynamics.Topics in control theory have been applied to deep learning Liu et al. [2020] , Pereira et al. [2020] as well as diffusion modeling Berner et al. [2022] . Optimal control can also be connected to diffusion processes via forward-backward SDEs Chen et al. [2021] . 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. 6 Experiments [2017], and a total-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 × 256 Karras et al. [2019] . 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 1K images left as a held-out set for evaluation. We compare our algorithm against competing frameworks on these last 1K images. We report our results on FFHQ 256 × 256 in Table 1 , and demonstrate improvements on all tasks against previous methods. Finally, we demonstrate the performance of our algorithm on the nonlinear inverse problem of class-conditional generation. Namely, let A(x) = classifier(x) and p(y|x) be its associated probability. We compare our method to DPS on 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 5 ). 
conclusion: 
 In this paper we presented a novel perspective on tackling inverse problems with diffusion models -framing the discretized reverse diffusion process as a discrete time optimal control episode. We demonstrate that this framework alleviates several core problems in probabilistic solvers: its dependence 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 our optimal control-based 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.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. 
the new nominal control trajectory ū *: 
 t is used to produce a new nominal state trajectory x * t and the algorithm is repeated from step 2 onwards until convergence or a fixed number of iterations. 
 
 Appendix: This paper builds on a large body of existing work and presents an improved technique for solving generic nonlinear inverse problems, which can be seen as a 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 a 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 a very popular trajectory optimization algorithm that has a rich history of theoretical results Jacobson [1968] as well as successful practical applications in robotics Tassa et al. [2007 Tassa et al. [ , 2014]] , aerospace Houghton et al. [2022] , Sasaki et al. [2022] 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 a nonlinear constrained optimization problem.To formulate an optimal control algorithm we first define the state transition function of a dynamical system asEQUATIONThe next ingredient that we need for our optimal control approach is a cost function J(x t , u t ) ∈ R. This is used to define a performance criterion that iLQR can optimize with respect to the set of controls {u t } t=1 t=T (i.e., the control trajectory going backwards from time t = T to t = 1). The cost-function is defined as follows:EQUATIONwhere, ℓ t and ℓ 0 are scalar-valued functions which are commonly referred to as the running costfunction 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 of the Value-function defined as follows:V (x t , t) = min {un} n=1 n=t J t = min {un} n=1 n=t 1 n=t ℓ n (x n , u n ) + ℓ 0 (x 0 ) (35)Intuitively, the Value-function resembles the optimal cost-to-go starting from time step t and state x t until the end of the time horizon (i.e., t = 0). Using this definition, one can easily derive the following recursive relation also known as Bellman's Principle of Optimality:EQUATION)A often useful defintion used in the derivation of the iLQR Riccati equations is that of the State-Action Value-Function Q(x t , u t ) given by,EQUATIONA sketch of the derivation of the Riccati equations is as follows: we take second-order Taylor expansions of both Q(x t , u t ) and V (x t , t) around nominal state and action trajectories of {x t } t=0 t=T and {ū t } t=1 t=T respectively. Next, we substitute these into Eq.( 37) and equate the first-and second-order terms to yield the following relations between the derivatives of Q, ℓ and V :EQUATIONEQUATIONEQUATIONEQUATIONwhere h xt and h ut are the Jacobians of the dynamics function h(x t , u t ), evaluated at time step t, w.r.t the state and the control vectors respectively. For ease of notation, we have dropped the subscript t and therefore all derivatives above should be considered to be evaluated at time step t, while we use V ′x and V ′ xx above to indicate the gradient and hessian of the Value-function evaluated at the next time step (i.e., at time step t -1).Next, we substitute for the second-order approximation of Q(x t , u t ) into Eq. ( 38) and note that u t can be written in terms of the nominal control as follows:u t = ūt + δu t .This results in a quadratic objective w.r.t δu t and the minimization in Eq. ( 38) can be performed exactly resulting in the following optimal perturbation from the nominal control trajectory:EQUATIONwhere, the feedforward and feedback gains are given by the following expressions:EQUATIONFinally, by substituting for the optimal δu * t back into Eq.( 38), we can drop the min operator and equate the first-and second-order terms on both sides. This results the following Riccati equations:EQUATIONThis 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 {ū t } 1 t=T and generate the corresponding nominal state trajectory {x t } 0 t=T using x t = h(x t+1 , u t+1 ). 2. By noticing from Eq. ( 35) that V (x 0 , 0) = ℓ(x 0 ) we can obtain expressions for V x and V xx evaluated at x0 . 3. Next, we compute the derivatives of Q given by equations. ( 39)-(44) using {ū t } 1 t=T and {x t } 1 t=T . 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)-( 49), we can propagate both V x and V xx one step backwards in time. 6. We then repeat the steps 3, 4 and 5 until we backpropagate the derivatives of V to time step t = 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:EQUATIONwhere, x t is the state obtained by unrolling the dynamics subject to the updated controls:x t = h(x t+1 , ū * t+1 ).Theorem 4.1. Let Eq. 3 be the discretized sampling equation for the diffusion model with output perturbation mode control (Eq. 18). Moreover, let the terminal cost ℓ 0 (x 0 ) = -log p(y|x 0 ) (27) be twice-differentiable and the running costs ℓ t (x t , u t ) = 0.(28) Then the iterative linear quadratic regulator with Tikhonov regularizer α produces the controlEQUATIONProof. We demonstrate the result via induction for t = 1, . . . , T .Since we assume that ℓ uu = 0, V xx vanishes:EQUATIONSimilarly, V x also greatly simplifies asEQUATIONTurning to the Tikhonov regularized feedforward term,EQUATIONFinally, the feedback term disappears due to the vanishingEQUATIONExplicitly denoting the dependence of V x and V ′ x on t, we can rewrite Eq. 56 asV (t) x = h T x V (t-1) x = ∂x t-1 ∂x t ∂ ∂x t-1 V.Combining this observation with the fact that ℓ 0 = -log p(y|x 0 ), we can conclude thatEQUATIONwhere x 0 depends on x t via the state transition function h (Eq. 18). Therefore, we have thatk = - 1 α V ′ x = 1 α ∇ xt log p(y|x 0 ) K = 0.Finally, given our action update (Eq. 15), we can conclude our desired resultEQUATION)Lemma C.1. Under the deterministic sampler with output perturbation mode control, α = 1 g(t) 2 ∆t recovers posterior sampling (Eq. 9).Proof. Substituting in α = 1 g(t) 2 ∆t to Eq. 29, we observe that Eq. 18 can now be written asEQUATION)Under the determinstic sampler, we can conclude that log p t (y|x 0 ) = log p t (y|x t ), since each x t has a unique path through the sample space. Therefore, we conclude that Eq. 65 resembles the ideal posterior sampler equation 9. We conclude our proof.Theorem 4.3. Let Eq. 3 be the discretized sampling equation for the diffusion model with input perturbation mode control (Eq. 17). Moreover, letEQUATIONand the running costs ℓ t (x t , u t ) = 0.(31) Then the iterative linear quadratic regulator with Tikhonov regularizer α = 1 g(t) 2 ∆t produces the dynamical sytemEQUATIONwhere x t := x t + u t .Proof. We similarly demonstrate the result via induction for t = 1, . . . , T .Again, assuming that ℓ uu = 0, V xx vanishes:EQUATIONwhereas V x greatly simplifies asEQUATIONTurning to the feedforward and feedback terms, we haveEQUATIONandEQUATION)We observe thatV (t) x = - 1 α h T x V (t-1)x . Therefore, noting that V (0) x = log p(y|x 0 ), we havek = -V (t) x = - 1 α (h (t) x ) T V (t-1) x = - 1 α ∇ xt log p(y|x 0 ) = - 1 α ∇ xt log p(y|x 0 (x t )).Applying the feedforward terms to the diffusion sampling process, we havex t-1 = (x t + u t ) + [f (x t + u t ) - 1 2 g(t) 2 ∇ x log p t (x t + u t )]∆t.We define the intermediary variable x t = x t + u t , (77) which has dynamicsEQUATIONWe now can see that, letting α = ∆tg(t) 2 , we obtainx t = x t + [f ( x t ) - 1 2 g(t) 2 (∇ x log p t ( x t ) + ∇ x log p t (y|x 0 ))]∆t..For all experiments, we use publicly available datasets and pre-trained model weights. For the FFHQ 256 × 256 experiments, we use the last 1K images of the dataset for evaluation. For MNIST, we do not 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] To speed up our proposed method, we leverage the following three modifications to the standard iLQR algorithm. Randomized Low-Rank Approximation The first and second order terms in Eqs. (19-25) are corresponding Taylor expansions of deep neural functions. Even with the use of automatic differentiation 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 computing these matrices, we utilize their known low rank structure Sagun et al. [2017] , 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 ∈ R m×n this is a four step process. 1) We sample a random matrix Ω ∼ N (0, I n×k ). 2) We obtain AΩ = Y ∈ R m×k . 3) We form a basis over the columns of Y, e.g. by taking the Q matrix in a QR factorization QR = Y. 4) We approximate A ≈ Q T QA.Notably, we observe that when A is a Jacobian (or Hessian) matrix, it can be approximated purely through Jacobian-vector and vector-Jacobian (Hessian-vector and vector-Hessian, resp.) productswithout ever materializing A itself. Moreover, a key result in randomized linear algebra is that this algorithm can approximate A up to accuracy O(mnkσ k+1 ) (Theorem 1.1 in Halko et al. [2011] ). Notably, if A has low rank structure where ∃k such that the k + 1th singular value σ k+1 = 0, then the approximation is exact. indexed set of matrices {A i }. We do this by forming projections of each A i against a corresponding set of dim(x) × ℓ column-orthogonal matrices {Q i }, which we denote as B i := Q T i A i . These matrices can then be stored at reduced cost as (Q i , B i ) pairs. Matrix multiplications between any A i A j can then be approximated up to rank ℓ with respect to the projected matrix, Q i A i,Qi , i.e.EQUATIONHowever, to prevent materialization of the full size of any matrices, we drop the leading Q i , obtaining a new projected-matrix pair (Q k , B k ), whereQ k = Q i .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 updateEQUATIONwhere P is the preconditioning matrix produced by the Adam optimizer. This reduces the overall runtime of the algorithm while still accounting for second-order information that respects the nonlinearity of the optimization landscape.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 iLQR 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 as P to avoid overloading the Q function notation):EQUATION)EQUATION)PQ xx P T = PQ ux P T = PQ xu P T = Ph T x V ′ xx h x P T (83)EQUATION)To simplify notation, each projection matrix P is the same -in reality, this need not be the case. Note that Q x and Q u 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 ℓ uu is diagonal (as it is in our case), we can obtain the projected inverse for Q uu as PQ -1 uu P T = Pℓ -1 uu P T + Pℓ -1 uu P T (C -1 + P T ℓ -1 uu P) -1 Pℓ -1 uu P T where C = Ph T x V ′ xx h x P (85) via a direct application of the Woodbury matrix inversion formula Petersen et al. [2008] , which has cost O(k 3 + kd 2 ). Finally, we compute the projected updates V xx , K as well as the full-precisionhttps://github.com/LeviBorodenko/motionblurFor Gaussian blur and random inpainting, Flash-Diffusion uses randomly sampled, but less severe degradation operators than our experimental setup. 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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):
	* "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".
	* "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).
	* "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.
	* "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.

2. 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).
3. Finally, the appendix could provide further details on
	* hyperparameter choices and optimization for the baselines.
	* precise assumptions for the theorems.
questions: questions: See "weaknesses" above.
limitations: limitations: See "weaknesses" above.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. High computation cost: Despite the advantages mentioned above, one obvious drawback of this method is the potential high computation cost. This includes: 

      a. Computing and storing the Jacobian matrices, which can be of very high dimension, can be very costly. Although the authors 
      further propose some techniques to reducing the cost, these methods might also bring extra approximation error as well as more hyper-parameters to tune; 

     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.


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

3. 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.
questions: 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.

2. 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: limitations: The authors has adequately addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. The method is well-backed but might be computationally exhausting.

2. The experiments are limited. Quantitative results on different datasets and nonlinear inverse problems are lacking.
questions: 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.

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

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

[1] Zehao Dou, and Yang Song. Diffusion Posterior Sampling for Linear Inverse Problem Solving: A Filtering Perspective, ICLR 2024

[2] Morteza Mardani, Jiaming Song, Jan Kautz, Arash Vahdat. A Variational Perspective on Solving Inverse Problems with Diffusion Models. ICLR 2024

[3] Jiaming Song, Arash Vahdat, Morteza Mardani, Jan Kautz. Pseudoinverse-Guided Diffusion Models for Inverse Problems. ICLR 2023

[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: limitations: Yes. They mentioned limitations and impacts in the appendix, which looks good to me.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: 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.    

Minor 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: 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.
paper_summary: paper_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. 
The framework is general and able to handle any differentiable forward measurement operator and establishes a new baseline in image reconstruction.
strength_weakness: strength_weakness: * 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.
* 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$?
questions: questions: * What is the rough runtime of each method in Table 1?
* 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)$.
* In (29), what is the meaning of $\nabla_{x_t} \log p(y|x_0)$?
* Line 149, what do you mean by "produces a feasible solution"? What's the meaning of being feasible?
* 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?
* Line 176, why is $\log p(y|x_t) = \log p(y|x_0)$ an assumption, not a consequence?
* 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$.

Minor comments
* Line 82, what is $\theta$?
* 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$.
* 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?
* 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. 
* Theorem 4.1 is missing transitions for presenting the conclusion (29).
* The values of $\alpha$ in the two theorems appear to be missing a factor of $2$.
* In (32), it should be $\widetilde{x}_{t-1}$ on the left side of the equation.
* Line 253, there is a parenthesis not closed.
limitations: limitations: There's not much discussion about the limitations in the paper. There is no potential negative societal impact.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.
questions: 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?

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

3. Please indicate what is output of Algorithm 1.

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

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

Comments: title: Response
Response to authors
Please review the authors' rebuttal and respond
Rebuttal (Continued)
comment: We greatly appreciate the reviewer's continued discussion of our paper, and their more positive outlook. Below we respond to some additional points the reviewer raised.

**No theoretical results for such robustness [to the score approximation] -- see Lines 196-197.** Theoretically, the reasoning for such robustness is due to the simple fact that DPS (and other algorithms) differentiate through Tweedie's formula, which relies explicitly on the score. On the other hand, our algorithm differentiates through the true simulated forward dynamics, which can be any dynamical system --- this additional flexibility is granted by our generalized optimal control framework.

**Comparison to ICLR 2024 results.** Indeed, the ImageNet results from ICLR 2024 (and some other venues) are very relevant, and we shall include some of these concurrent results in our work, where applicable. We note that some works do not use the same experimental setting. For example, ReSample uses the much easier setting of $\sigma = 0.01$, whereas we keep the settings used by DPS $\sigma=0.05$. We still believe that our work provides a meaningful contribution to the field, and hope the reviewer agrees.

**The conditional score remains intractable.** We claim tractability of $\nabla_{x_t} \log p(y | x_t)$ in the deterministic case. We also claim tractability of $\nabla_{x_t} \log p(y | x_0)$ in the stochastic case, which DPS (and related models are still unable to compute). For details, see our response to *No theoretical results for such robustness [to the score approximation]*. We have made this clearer in the manuscript.

**Estimating $x_0$ requires approximation of the score.** There is a subtle distinction here. $x_0$ as defined by the reverse diffusion process requires the score, of course. $x_0$ as defined simply as the final state of the dynamical system solved by our system does NOT require the score. We will make this distinction more clear in our manuscript.

**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).** We thank the reviewer for bringing these papers to our attention. DiffPIR actually uses Tweedie's (See Line 3 in Algorithm 1). On the other hand, DCDP does not leverage Tweedie's. However, there is also no theoretical discussion of the correctness of the proposed approach.

**Assumptions for the theorems.** $x_t$ is defined as in Algorithm 1. We are missing a statement before Eq. 29, which should read: "Then the iterative linear quadratic regulator with Tikhonov regularizer $\alpha$ produces the control". We apologize for the confusion, and have added statements clarifying this in our manuscript.
I thank the authors for their hard work on the rebuttal and appreciate the additional results. However, I still have concerns about two points:

- Insufficient baselines: While I agree that DPS-like baselines share the same issues due to the use of Tweedie's formula, recent works like  FPS have shown significant performance improvements. Additionally, Resample (https://arxiv.org/abs/2307.08123, ICLR 2024) is an advanced method of PSLD, as both approaches are latent-based.

- Unsatisfactory results: In Table 1 of the paper, the proposed method (T = 20) does not outperform either PSLD or DPS and is not even competitive. I wouldn't expect a significant difference between the FFHQ (Table 1 in the paper) and ImageNet datasets(in the rebuttal pdf) to change this outcome.

Given these considerations and other reviewers' comments, I retain my score.
**No theoretical results for such robustness [to the score approximation]:** "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." It is clear (and also written in 196-197), that the theoretical results require such an approximation. Thus, it still remains unclear *why* the method is more robust. 

**Methods should be compared in terms of performance vs runtime / flops and no the number of diffusion steps. / The proposed method is significantly more expensive than competing methods. / Results are only provided for MNIST and FFHQ.** I thank the authors for the additional experiments which improve the empirical evaluation and resolve some of my concerns. A revised version of the paper should provide results for different NFEs to assess the scaling as well as provide more insights into runtime for all considered method (considering the maximal batchsize for each method for a given memory budget). However, when comparing against further baselines, such as FPS-SMC [1], DCDP [2], DiffPIR [3], RedDIFF [4], ReSample [5], the provided results for ImageNet, in particular for the deblurring tasks, seem not to be state-of-the-art anymore.

**The conditional score function remains intractable in the proposed method.** The conditional likelihood (also the one used for the ODE) is defined via $p(y|x_t) = \int p(y|x_0) p_{SDE}(x_0 | x_t) \mathrm{d}x_t$, where the conditional density of the SDE $p_{SDE}$ cannot just be replaced by the one for the ODE (which, I agree, would be a dirac, since for deterministic dynamics $x_0$ is exactly determined by $x_t$.).

**The statement that our model estimates $x_0$ exactly, rather than forming an approximation sounds misleading.**
Indeed one can just simulate the SDE, however, the distribution of $x_0$ depends on the approximation of the score (see also the first comment).

**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).** Note that resampling-based methods, such as DiffPIR and DCDP, use a simulation instead of Tweedie's formula to obtain an estimate of $x_0$.

**Assumptions for the theorems.** It should be defined what exactly is meant by $x_t$ (does it originate from the approximate score, the ground-truth score, the forward/backward simulation), on which equation (29) depends. Also, equation (29) comes without any statement.


Given the additional empirical evidence, I adjusted my score. However, I still think that the presentation should be improved, both theoretically (unclear statements, see above) as well as empirically (compare to further baselines, see above).

---

[1] Dou, Zehao and Song, Yang. "Diffusion Posterior Sampling for Linear Inverse Problem Solving: A Filtering Perspective.", ICLR 2024.

[2] Li, Xiang, et al. "Decoupled data consistency with diffusion purification for image restoration.", arXiv preprint arXiv:2403.06054.

[3] Zhu, Yuanzhi, et al. "Denoising Diffusion Models for Plug-and-Play Image Restoration.", CVPR workshop NTIRE 2023.

[4] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. arXiv preprint arXiv:2305.04391, 2023.

[5] 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. ICLR, 2023
Thanks for providing detailed explanation. I changed my score to be 5.
I would like to thank the authors for their responses. After reviewing the rebuttal pdf and comments, my concerns are resolved and I would like to raise my rating.
I thank the authors for their detailed response, and clarifications which were insightful.

Taking into account the new information provided, I will increase my score.
Dear authors, thank you for the detailed response and the additional experiments. This clarifies most of my questions. I would like to keep my positive score unchanged.
Dear Reviewers,

Thank you once again for your time and effort in reviewing for NeurIPS 2024.

The authors have already submitted their rebuttal to your comments, and the discussion phase has now begun. For those who have not yet responded, please review the authors' responses carefully at your earliest convenience. Your timely feedback on whether your previous concerns have been adequately addressed is crucial to the reviewing process.

Thank you very much!

Best regards,
Area Chair
**Line 176: Why is $\log⁡ 𝑝(𝑦|𝑥_0) = \log⁡ 𝑝(𝑦|𝑥_t)$ an assumption, not a consequence?**
Indeed, $\log p(y | x_t) = \log p(y | x_0)$ is a consequence under the deterministic ODE dynamics. However, we also consider the setting of DPS, where the process is stochastic. Under this setting, there may be more than one $y$ (and thus $x_0$) associated with a given $x_t$, since $p(x_t|x_0)$ is Gaussian and supported everywhere. In this case, the equality is no longer guaranteed, but is assumed in many prior works (e.g., DPS).

**Line 204: What is the purpose this experiment with randomly initialized weights?**
Indeed, this application is itself not meaningful. We meant to demonstrate DPS samplers *require* a well-approximated score model to function properly, whereas we do not. Of course, the performance of our model also improves with a better diffusion prior, but the deterioration when the score is not well approximated is more graceful than DPS (i.e., in less represented parts of the dataset). This can be useful when using a general inverse solver (e.g. ImageNet or foundation model) and solving for an image from an unknown or poorly approximated distribution. 

To further illustrate our point, we also demonstrate that the same trend occurs even when the diffusion model weights are initialized from a pretrained model, but from a different distribution than the source image. In Figure 10 in the rebuttal PDF, the right hand block shows an inverse problem setting for reconstructing ImageNet images, using an FFHQ model. As expected, our proposed method outperforms DPS on this out-of-domain inverse problem setting.
**Hyperparameter choices for the baselines.**
For baselines from competing works in Table 1, we directly replicate the reported hyperparameters from the respective papers. For our baselines, please refer to Table 2 for the hyperparameter values. For the T = 20 result, we let T = 20 and let the number of iterations be 50. For insight into how we selected these values, please see Section D.3 and Tables 4, 5, and 6.

**Assumptions for the theorems.**
The assumptions required for Theorem 4.1 are summarized in Equations 27 and 28. Similarly, Theorem 4.2 requires Equations 30 and 31. We have additionally clarified that we require $\alpha > 0$, and $\log p(y | x)$ to be twice-differentiable (which holds when $p(y | x)$ is Gaussian, as is the case in our settings). We take our theoretical results very seriously and would be happy to clarify any other ambiguities in the theorems.