Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems

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.

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.

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.

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.

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.

MC. Fig 2 caption should indicate what the numbers mean in the image

A. Thank you for the careful reading. Fixed.

References

[1] 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

[2] Uecker, Martin, et al. "ESPIRiT—an eigenvalue approach to autocalibrating parallel MRI: where SENSE meets GRAPPA." Magnetic resonance in medicine 71.3 (2014): 990-1001.

[3] Jalal, Ajil, et al. "Robust compressed sensing MRI with deep generative priors." NeurIPS 2021.

**W1. The detailed process of $M$-step CG update is lacking.

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

W2. In page 7, I think clarifying Eq. (29) will help the reader understand the problem setting.

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

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

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.

Q1. In page 3, $m$-th is not shown in equation (2).

A. Thank you for spotting the typo. Fixed.

Q2,3. For high-dim data, the compute complexity of CG is significant. Does this affect model's efficiency? How does it compare to DDIM?

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.

W1, Q2. Brief overview on diffusion and Krylov space theory is needed. Why do Krylov subspace methods guarantee the iterations to stay in the tangent space?

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

Q1. Do we need the assumption that the clean data manifold is an affine subspace? How about general curved manifold?

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]

Q3. Please elaborate the geometric view of diffusion model more. Is the diffusion process on the manifold or in the ambient Euclidean space?

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.

References

[1] Chen, Minshuo, et al. "Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data." ICML 2023.

[2] Rout, Litu, et al. "Solving linear inverse problems provably via posterior sampling with latent diffusion models." NeurIPS 2023.

[3] Rout, Litu, et al. "A theoretical justification for image inpainting using denoising diffusion probabilistic models." arXiv preprint arXiv:2302.01217 (2023).

[4] Chung, Hyungjin, et al. "Improving diffusion models for inverse problems using manifold constraints." NeurIPS 2022.

W1. How do we initialize $\mathbf{x}_0$ from Gaussian noise?

A. We sample from random Gaussian noise, i.e. sampling $\mathbf{x}_T \sim \mathcal{N}(0, \mathbf{I})$ as in standard practice of diffusion models.

W2. In equation (19), what is $A_t$ and $y_t$? Are they same with $A$ and $y$? If so, will this cause some bias or problem?

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.

**W3. For different steps, will $l$ in equation (22) be fixed? If so, how to get an optimal $l$?

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

W4. In table 1, what's the meaning of mask pattern?

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.

W5. In table 2, it seems that with more steps, the performance drops. Any explanation?

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.