Paper: Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems (ICLR-2024)

Reviewer comment (Experiments):
The comparison provided primarily centers on state-of-the-art (SOTA) diffusion models for reconstructing images from undersampled data and only a comparison with the naïve Total Variation (TV) is included. However, 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.

Author response:
**W1. Are DIS better than classic CS methods? References?**

**A.** Thank you for the suggestion. Please see modified section 4.1., where we cite two works on using diffusion model-based inverse problem solvers, which show that the proposed methods largely outperform the classic CS algorithms such as L1-wavelet [1] and L1-ESPiRiT [2]. While we do not directly compare ours against these CS baselines, we do compare against [3], implying superiority.

Reviewer comment (Theory):
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.

Author response:
**W2, Q1. Relation between the solution to $\mathbf{y} = \mathbf{A}\mathbf{x}$ and the modified equation by left multiplying $\mathbf{A}^\top$**

**A.** Thanks for your careful observation. We agree that even after using $\mathbf{A}^*\mathbf{y} = \mathbf{A}^*\mathbf{A}\mathbf{x}$ to obtain a square matrix for Krylov method, the solution is not unique. This is why the notion of Krylov subspace is important. By enforcing that the solution lies in a specific lower-dimensional Krylov subspace whose dimension is less than equal to $\text{rank}(\mathbf{A}^*\mathbf{A})$, the ill-posedness of the zero-finding in $\mathbf{A}^*\mathbf{y} = \mathbf{A}^*\mathbf{A}\mathbf{x}$ can be overcome.

Reviewer comment (Theory):
The proposed method has a limitation since it only considers linear forward models with the form y=Ax.

Author response:
**W3, Q2. The proposed method can be applied to linear models. Can you clarify what do you mean by extending Krylov method to nonlinear problems?**

**A.** Apologies for the confusion. When solving problems with a nonlinear operator $\mathcal{A}(\cdot)$ we can change the objective to $\|\|\mathbf{y} - \mathcal{A}(\mathbf{x})\|\|_2^2/2$. Then, one can use methods such as Newton-Krylov method to linearize the given problem with Taylor approximation, and apply standard Krylov method in that approximate regime. We have modified our manuscript to elaborate without experiments as this is beyond the scope of this work.

Reviewer comment (Theory):
The method requires that the tangent space on the clean manifold at a denoised sample forms a Krylov subspace and experimental results suggest 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 an assumption to become true.

Author response:
**W4, Q3. What is the intuition behind the assumption that the tangent of the data manifold forms a Krylov subspace?**

**A.** In practical use of Krylov subspace in linear inverse solver, the approach becomes powerful and require less CG optimization step especially when the initialization is good. This is because the solution may lie in lower dimensional affine space around the initialization point. 

Our diffusion inverse solvers combined with Tweedie’s denoising step gives quite a good approximation of the true solution especially when the reverse sampling progresses, so that CG may achieve a good solution with a small number of iteration. Our intuition is to find an iterative solver to satisfy the data consistency with the smallest number of iterations in order to avoid the results from falling off the correction diffusion manifold. Empirically we found that CG was one of the most powerful solvers with the initialization with diffusion solver, which led us to believe that the local tangent space may be in the form of the Krylov subspace.

We believe that the empirical study in Section F.1. does indeed show that our assumption is close to reality. That said, it would be definitely interesting to further investigate on when the assumptions will perfectly hold as a future direction of research.

Reviewer comment (Evaluation):
In the 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 is unspecified. The experiments should also report variability of the results, for example showing standard deviation for a more comprehensive comparison of different methods.

Author response:
**W5, Q4. Experiments: how many slices were considered? Variability of the measures and their statistical significance?**

**A.** For CS-MRI, the total number of test slices is 292. For CT, it is 448, 256, 256 slices for axial, coronal, and sagittal slices, respectively. The results reported in the Tables are the average values. For completeness, we also included the std values in Tab. 1. Due to limited space, we included the std values for Tab. 4 in Appendix Tab. 8. We do not perform statistical tests as it would require pairwise comparison of DDS against all the baseline methods.

Reviewer comment (Presentation):
Fig. 2 caption should indicate what are the numbers at the bottom of each subimage, which are presumably PSNR/SSIM.

Author response:
**MC. Fig 2 caption should indicate what the numbers mean in the image**

**A.** Thank you for the careful reading. Fixed.

Reviewer comment (Reproducibility):
The detailed process of M-step CG update is lacking.

Author response:
**A.** For completeness, we included Algorithm 1 which is the precise algorithm used throughout the work.

Reviewer comment (Writing):
In page 7, I think clarifying Equation (29) will help the reader understand the problem setting.

Author response:
**A.** We included Fig. 3 as an illustration of the imaging physics, and added a pointer in the main text.

Reviewer comment (Theory):
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.

Author response:
**A.** Thank you for the suggestion. In the modified Appendix B.1., where we cover the preliminaries on diffusion models, we have also elaborated on the use of Tweedie’s formula and how it relates to diffusion models and our work.

Reviewer comment (Writing):
In page 3, "the m-th order Krylov subspace is defined by equation (2)", this "m-th" is not shown in equation (2).

Author response:
**A.** Thank you for spotting the typo. Fixed.

Reviewer comment (Evaluation):
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?

Author response:
**A.** The computation cost for diffusion sampling is approximately linear to the neural function evaluation (NFE) count. This is the computation time needed for neural network forward pass dominates and far exceeds the time needed for applying CG, even for the large-scale inverse problems that we consider. 

For instance, removing the CG steps, DDS would be equivalent to unconditional DDIM sampling. For the case of CS-MRI where the dimension is (15 $\times$ 320 $\times$ 320), DDIM vs DDS sampling time is 13.50 vs 13.01 [s], showing that there exists only a small difference.

Even when we consider 3D SV-CT reconstruction where we additionally solve the optimization problem with ADMM-CG, and the data dimension is very high (256 $\times$ 256 $\times$ 448), DDIM vs DDS sampling time is 21:07 vs. 16:01 [m]:[s]. While in this case, the difference is more pronounced than the multi-coil CS-MRI case, we still see that DDS only introduces a reasonable amount of computational overhead considering the scale of the problem that we are solving.

Reviewer comment (Evaluation):
Which has a higher computational complexity, this method or the classical accelerated diffusion models (e.g., DDIM)?

Author response:
**A.** The computation cost for diffusion sampling is approximately linear to the neural function evaluation (NFE) count. This is the computation time needed for neural network forward pass dominates and far exceeds the time needed for applying CG, even for the large-scale inverse problems that we consider. 

For instance, removing the CG steps, DDS would be equivalent to unconditional DDIM sampling. For the case of CS-MRI where the dimension is (15 $\times$ 320 $\times$ 320), DDIM vs DDS sampling time is 13.50 vs 13.01 [s], showing that there exists only a small difference.

Even when we consider 3D SV-CT reconstruction where we additionally solve the optimization problem with ADMM-CG, and the data dimension is very high (256 $\times$ 256 $\times$ 448), DDIM vs DDS sampling time is 21:07 vs. 16:01 [m]:[s]. While in this case, the difference is more pronounced than the multi-coil CS-MRI case, we still see that DDS only introduces a reasonable amount of computational overhead considering the scale of the problem that we are solving.

Reviewer comment (Writing):
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?

Author response:
**A.** There is no $A_t$ and $y_t$ in eq. (19). Please let us know which equation you are referring to. We would be happy to answer your question.

Reviewer comment (Experiments):
For different steps, will $l$ in equation (22) be fixed? If so, how to get an optimal $l$?

Author response:
**A.** Yes. For all steps $l$ is fixed in our experiments, although this does not necessarily have to be the case. We opted for simplicity and let $l$ fixed. $l$ is a hyper-parameter that one could tune according to the circumstances (e.g. see Table 2). In order to find $l$, one could rely on heuristics, but we note that in practice, we achieve stable results for a wide range of $l$.

Reviewer comment (Presentation):
In table 1, what's the meaning of mask pattern?

Author response:
**A.** The mask pattern refers to the under-sampling pattern of Fourier space in compressed sensing MRI. Please see the first column of Figure 6 and 7 for examples of these “masks”. The white pixels are the *sampled* k-space points, where as the black pixels are the *unknown* k-space points that need imputation.

Reviewer comment (Evaluation):
In table 2, it seems that with more steps, the performance drops. Any explanation?

Author response:
**A.** Our theoretical analysis builds on top of the assumption that the clean data manifold is an $l$-th order Krylov subspace. When the iteration $M$ exceeds this value of $l$, the sample might deviate from the tangent, and thereby incur errors.

Reviewer comment (Writing):
The work involves both Krylov space theory and diffusion model, and it will be more helpful for a general audience to give a brief overview of both theories in the appendix.

Author response:
**A.** Thank you for the suggestion. Per the reviewer's request, we have included the preliminaries section in Appendix A, where we review both diffusion models and Krylov subspace methods.

Please see the modified Appendix A.2. on the background of Krylov subspace methods. In essence, Krylov subspace methods can be thought of as repeated applications of the matrix-vector product with some matrix $A$, which is responsible for constructing the span of the Krylov subspace. For every increased iteration, we have an increased order of the subset. Due to the property that the $M$-th order Krylov space is a subspace of $L$-th order Krylov space, i.e. $\mathcal{K}_M \subset \mathcal{K}_L$ if $M \leq L$, Krylov subspace methods are guaranteed to stay on the tangent space if the iteration count $M$ does not exceed the value of $L$.

Reviewer comment (Theory):
Do we need the assumption that the clean data manifold is an affine subspace? How about a general curved manifold?

Author response:
**A.** This assumption is required for diffusion models to act as the projector. That said, it should be noted that natural data manifolds can be well-approximated as locally affine manifolds in most cases, which makes our assumption practical. Note that the affine subspace assumption is widely used in the diffusion literature and has been used when 1) studying the possibility of the score estimation and distribution recovery [1], 2) showing the possibility of signal recovery, and most relevantly [2,3], 3) showing the geometrical view of the clean and noisy manifolds. [4]

Reviewer comment (Theory):
Why does the Krylov space method guarantee staying in the tangent space?

Author response:
**A.** Thank you for the suggestion. Per the reviewer's request, we have included the preliminaries section in Appendix A, where we review both diffusion models and Krylov subspace methods.

Please see the modified Appendix A.2. on the background of Krylov subspace methods. In essence, Krylov subspace methods can be thought of as repeated applications of the matrix-vector product with some matrix $A$, which is responsible for constructing the span of the Krylov subspace. For every increased iteration, we have an increased order of the subset. Due to the property that the $M$-th order Krylov space is a subspace of $L$-th order Krylov space, i.e. $\mathcal{K}_M \subset \mathcal{K}_L$ if $M \leq L$, Krylov subspace methods are guaranteed to stay on the tangent space if the iteration count $M$ does not exceed the value of $L$.

Reviewer comment (Theory):
Please elaborate more on the geometric view of the diffusion model. Is the diffusion process on the manifold or in the ambient Euclidean space?

Author response:
**A.** We apologize for the confusion. The diffusion process is on the ambient Euclidean space, but we make assumptions that the natural data manifold exists on a manifold, which is a subset of the Euclidean ambient space.