Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

Summary

In this paper, the authors propose to integrate classical Krylov subspace methods, widely used in large linear inverse problems, with state-of-the-art (SOTA) diffusion models applied to medical image reconstruction. An efficient diffusion sampling strategy is proposed that, if the tangent space at a denoised sample forms a Krylov subspace, allows the Conjugate Gradient (CG) method to maintain data consistency updates within this subspace. This eliminates the need for the computationally expensive manifold-constrained gradient (MCG) reducing computation time. The approach achieves SOTA reconstruction quality in large scale inverse imaging problems, such as multi-coil MRI and 3D CT reconstruction.

Strengths

• The idea of applying Krylov subspace methods for efficient sampling strategies in stable diffusion models is novel and significant since practical applications such as medical image reconstructions are usually large inverse problems. • The organization of the paper is clear with useful introductions to the relevant topics such as diffusion models and Krylov subspace methods. • Comparison with various SOTA methods is provided showing that the proposed method achieves higher quality results with significant fewer number of iterations.

Weaknesses

Main weaknesses: • The comparison provided primarily centers on state-of-the-art (SOTA) diffusion models for reconstructing images from undersampled data and only a comparison with a the naïve Total Variation (TV) is included. However, I believe it lacks a comparison with the classical Compressed Sensing (CS) approach, which relies on sparse image representations and has shown success, especially in MRI applications (as in the classical reference [*]). It's unclear whether SOTA diffusion models can compete with traditional CS algorithms. Perhaps referencing recent studies that compare diffusion models with classical CS methods would offer a more comprehensive perspective on the field's current status. • It is not clear why the Krylov method, originally developed for square matrix A can be successfully applied to the image reconstruction problem when A has fewer rows than columns. See my question number 1 below. • The proposed method has a limitation since it only considers linear forward models with the form y=Ax (see my question number 2 below). • The methods requires that the tangent space on the clean manifold at a denoised sample forms a Krylov subspace and experimental results suggests that such assumption is met. However, an intuition about why the tangent space forms a Krylov subspace is missing. It would be useful to have identified conditions for such assumption become true (see my question number 3 below). • Experiments: Quality results are expressed in terms of PSNR and SSIM values, but it is not clear if these are averaged values and the number of instances considered in the corresponding average. The experiments should also report variability of the results, for example showing standard deviation for a more comprehensive comparison of different methods (See my question number 4). Minor comments: - Fig. 2 caption should indicate what are the numbers at the bottom of each subimage, which I think they are PSNR/SSIM - Acronym NFE is not defined in the paper, please introduce its definition (Number of Forward Evaluations?) - Just before eq. (3): “two subspace K and L” -> “two subspaces K and L” - Just after eq. (7): “for given the optimization problem” -> “for a given the optimization problem” - Page 6: “guranteed” -> “guaranteed” - Problem setting considers only the noiseless case (equation (29)), however experiments in Fig. 2 shows some results in the noisy case. Please include the noisy case in the problem setting. - In the problem setting, s^{(1)}, s^{(2)}, …, s^{(c)} are not defined. For single-coil case s^{(1)}=I but it is not clear how they are defined in the multi-coil case. Please, clarify. - PC sampler is not defined. [*] Lustig, M., Donoho, D. and Pauly, J.M. (2007), Sparse MRI: The application of compressed sensing for rapid MR imaging. Magn. Reson. Med., 58: 1182-1195. https://doi.org/10.1002/mrm.21391

Questions

1. In section 2, it is stated that in the forward model (1): y=Ax, we can assume A square without losing generality because, if not, we can consider the modified equation (2): A*y=A*Ax, because its solution is also a solution of problem (1). However, equation (1) has infinite number of solutions when A has fewer rows than columns as in the considered in the image reconstruction scenarios, and the solution of eq. (2) is unique if A is full rank matrix. It is not clear why the unique solution of eq. (2) is the most preferable solution over all possible solutions of eq. (1). Could you please clarify about this? 2. The proposed method has a limitation since it only considers linear forward models with the form y=Ax and uses Krylov method to minimize l(x) = ||y - Ax||^2. However, just after eq. (6) it is mentioned that Krylov method can be extended to non-linear problems and the example l(x) = 0.5||y - Ax||^2 is given, which is confusing since it is the same optimization problem as in the linear model y=Ax. Can you clarify what do you mean by extending Krylov method to nonlinear problems? Can thus you apply your method to an inverse problem having a nonlinear forward model? 3. The methods requires that the tangent space on the clean manifold at a denoised sample forms a Krylov subspace and experimental results suggests that such assumption is met. What is the intuition behind this assumption? Can you provide some conditions for such assumption become true? 4. Regarding the reported PSNR and SSIM values: Are these average values? If so, how many samples have you consider computing the average? Can you also report the variability of these measures to derive conclusions about the comparison of methods? Are the differences between methods statistically significative in some sense?

Rating

8: accept, good paper

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

This paper proposes a novel and efficient diffusion sampling strategy that synergistically combine the diffusion sampling and Krylov subspace methods. Specifically, it prove that if the tangent space at a denoised sample by Tweedie’s formula forms a Krylov subspace, then the CG initialized with the denoised data ensures the data consistency update to remain in the tangent space. This eliminates the necessity of computing the manifold-constrained gradient (MCG), leading to a more efficient diffusion sampling method. The proposed method achieves state-of-the-art reconstruction quality on challenging real-world medical inverse imaging problems, and more than 80 times faster inference time than the previous state-of-the-art method.

Strengths

A novel and efficient diffusion sampling strategy is proposed, synergistically combining diffusion sampling and Krylov subspace methods. It certifies that if the tangent space at a denoised sample by Tweedie’s formula forms a Krylov subspace, then the Conjugate Gradient initialized with the denoised data ensures that the data consistency update remains in the tangent space. This also eliminates the need to compute the manifold-constrained gradient.

Weaknesses

1.The detailed process of M-step CG update is lacking. 2.In page 7, I think clarifying Equation (29) will help the reader understand the problem setting. 3.An important point of this paper is the negation of the necessity to compute the manifold-constrained gradient, but there is a lack of elaboration on using Tweedie's formula to obtain the tangent space for gradient calculation.

Questions

1.In page 3, " the m-th order Krylov subspace is defined by equation (2)", this "m-th" is not shown in equation (2). 2.For high-dimensional data, the computational complexity of the Conjugate Gradient (CG) algorithm is significant, and each step of inverse sampling requires solving using the CG algorithm. Does this affect the model's efficiency? 3.Which has a higher computational complexity, this method or the classical accelerated diffusion models(e.g., DDIM)?

Rating

6: marginally above the acceptance threshold

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

In this paper, the authors aim to solve the inverse problem through diffusion model and propose the conjugate gradient (CG) method based on its connection to Krylov subspace. Specifically, in each time step, they update the noisy estimation based on the CG update and finally obtain the inverse results. Experiments show that the proposed method works well.

Strengths

- Using CG method in each time step is reasonable and can mitigate the difficulty of solving the original problem. - The results seem promissing.

Weaknesses

- How to initialize $x_0$, from Gaussian noise? - In equation (19), what is $A_t$ and $y_t$, are they the same with $A$ and $y$? If so, will this cause some bias or problem? - For different steps, will $l$ in equation (22) be fixed? If so, how to get an optimal $l$? - In table 1, what's the meaning of mask pattern? - In table 2, it seems that with more steps, the performance drops. Any explanation?

Questions

Please see above

Rating

6: marginally above the acceptance threshold

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

This work proposes the Decomposed Diffusion Sampling (DDS) method as a Diffusion model-based Inverse problem Solvers (DIS) for inverse problems in medical imaging. The work is based on the observation that a diffusion posterior sampling (DPS) with the manifold constrained gradient (MCG) is equivalent to one-step projected gradient on the tangent space at the “denoised" data by Tweedie’s formula, this work provides multi-step update scheme on the tangent space using Krylov subspace methods. The experiments show that performing numerical optimization schemes on the denoised representation is superior to the previous methods of imposing DC. Further, the work devises a fast sampler based on DDIM that works well for both VE/VP settings. With extensive experiments on multi-coil MRI reconstruction and 3D CT reconstruction, it was shown that DDS achieves superior quality while being ≥ ×80 faster than the previous DIS.

Strengths

The main strength of the work is the novel theoretic insight: the geometric interpretation that the diffusion posterior sampling (DPS) with the manifold constrained gradient (MCG) is one-step projection to the tangent space of the clean data manifold, and using Conjugate Gradient type method, one can achieve multiple steps projection within the tangent space. All the theoretic proofs are given in details, all the lemmas are well formulated and clearly explained. The experimental results are convincing.

Weaknesses

The work involves both Krylov space theory and diffusion model, it will be more helpful for general audience to give brief overview of both theories in the appendix. Especially. the key observation : DPS with MCG is a one-step projection to the tangent space of the clean data manifold. Furthermore, it needs more explanation for the problem : why Krylov space method can guarantee to stay in the tangent space.

Questions

1. Do we need the assumption that the clean data manifold is an affine subspace ? How about general curved manifold ? 2. Why Krylov space method guarantee to stay in the tangent space ? 3. Please elaborate the geometric view of diffusion model more. Is the diffusion process on the manifold or in the ambient Euclidean space ?

Rating

6: marginally above the acceptance threshold

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

General response

We would like to thank the reviewers for their constructive and thorough reviews. We are encouraged that the reviewers think that **the work offers novel theoretic insight** (Q4Mc), **method is reasonable and results are promising** (wm9K), **proposes a novel and efficient strategy** (tJRQ), and **the idea is novel, significant, and clear** (QCsq). For point-to-point response, please refer to below.