Solving Inverse Problems via Diffusion Optimal Control

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.

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

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?

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

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.

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.

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. * The writing is generally good (see more comments regarding writing in questions). It is concise and to the point with good intuitions provided. * Efforts (e.g. low-rank approximations) have been made to bring down the computational cost in iLQR for the high-dimensional image setting. * 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. * 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

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

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

There's not much discussion about the limitations in the paper. There is no potential negative societal impact.

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.

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.

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.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.

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.

Strengths

1. This paper is well-written, with a good summary of previous methods and their shortcomings. 2. 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. 2. The experiments are limited. Quantitative results on different datasets and nonlinear inverse problems are lacking.

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

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes. They mentioned limitations and impacts in the appendix, which looks good to me.

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.

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

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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors has adequately addressed the limitations.

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.

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): * "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

See "weaknesses" above.

Rating

4

Confidence

3

Soundness

1

Presentation

2

Contribution

2

Limitations

See "weaknesses" above.

Please review the authors' rebuttal and respond

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

Response

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.