Meta-Guided Diffusion Models for Zero-Shot Medical Imaging Inverse Problems

We thank the reviewer for their constructive comments. Below are our point-by-point responses to Weaknesses and Questions:

Weakness 1: We acknowledge that “Robust Compressed Sensing MRI” is one of the first pioneer papers in the field, cited in our introduction section. Our decision not to test this method was driven by two primary reasons. Firstly, according to Chung et al. (2022) [1], the approach described in “Robust Compressed Sensing MRI” addresses linear inverse problems by employing an approximation $\nabla_{x_t} \log p_t(y|x) \approx \frac{A^H(y - Ax)}{\sigma^2}$, which is valid when ${n}$ is considered Gaussian noise with variance $\sigma^2$. However, this approximation only holds $t=0$ and fails at other noise levels relevant to the generative process. Secondly, while state-of-the-art (SOTA) methods did not benchmark against this approach, Song's method (Score-Med) [2] did include it as a baseline and outperformed it. Moreover, our results have surpassed those of Song's method.

Regarding the inclusion of supervised methods, we faced uncertainty. The challenge lies in the lack of fair comparative studies, as supervised methods have different architectures and hyper-parameters. Despite that, we have now reported the quantitative comparisons with SOTA supervised methods in Appendix Table 4. Consistent with other zero-shot inverse problem solvers, such as those discussed in Chung et al. (2022) [1], Wang et al. (2022) [3], and Kawar et al. (2022) [4], our MGDM model demonstrates superiority over existing supervised methods.

Weakness 2: We recognize your concerns about the conjugate symmetry in k-space when applying Fourier Transform to its real-valued images like BraTS, which may affect the effective undersampling factor, potentially halving the apparent ACR. I addressed this issue in my work and included a citation to the paper you have recommended, which discusses the biases that can arise from the misuse of public data. While I continue to utilize the BraTS dataset for its merits, I ensure to mention this potential pitfall for clarity and accuracy.

Quoted from our limitation part: “Additionally, the BraTS dataset, employed both in our study and by the baseline methods, has been indicated in a recent paper (Shimron et al., 2022) to have an overestimated undersampling factor, which arises from the conjugate symmetry of k-space inherent in real-valued images.”

Additional results for the FastMRI dataset, including undersampling patterns and detailed descriptions, are provided in the main paper and Appendix.

Weakness 3: We first want to acknowledge that other baseline methods have also reported the capability to reconstruct MRI BraTs datasets under x24 acceleration factor. For example, in Song’s paper [4], they reported a similar PSNR of 29.42; ours is slightly higher, with a PSNR of 30.04. We believe that learning informative priors can significantly aid in signal recovery. The training and testing datasets in Brats come from the same distribution, so we expect to achieve good results and effectively recover brain tumors. However, if tested on a dataset outside this distribution, our method may not perform as well, which would require additional adaptations. Our reasoning is further supported by our robust data consistency step.

Minor Problem 1: We apologize for the oversight and thank you for bringing it to my attention. Now, all figures, including Figure 5, are provided in high resolution without compression artifacts in the revised manuscript. Additionally, the orientation of the knee images is corrected to the standard presentation.

Minor Problem 2: The citation for ESPIRiT is now corrected, and a thorough review of all other references is conducted to ensure their accuracy. Thank you for pointing out this discrepancy.

Question 1: Please refer to our response to Weakness 3.

Question 2: We have now added Figure 6 to show how the reconstruction results look like with different ablations.

[1] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.14687, 2022a.

[2] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint arXiv:2111.08005, 2021.

[3] Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. arXiv preprint arXiv:2212.00490, 2022.

[4] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in Neural Information Processing Systems, 35:23593–23606, 2022.

We thank the reviewer for their constructive comments. We highly invite the Reviewer to read our paper again, we have significantly improved our representation of the work, including typos, notations, concepts, and formulations.

Below are our point-by-point responses to Weaknesses and Questions:

Weakness 1 (Presentation): We have now reworked our presentation of the paper with major changes highlighted in blue.

• Diffusion Models: We had initially aimed to make the paper as comprehensive as possible. Following your suggestion, we have now updated that section by removing the continuous part and rendering the discrete section more concise.

• Typos: Thank you for your feedback. We have now thoroughly checked grammar and polished our writing. If you notice any further discrepancies, please know they will be addressed in our continuous efforts to improve.

• Zero-Shot: Zero-shot learning is one of the several learning paradigms aimed at out-of-distribution generalization, where the algorithm is trained to categorize objects or concepts that it has not been exposed to during training. In this paradigm, the set of classes during training, denoted as $\mathcal{Y}^{train}$, is distinct and non-overlapping with the set of classes during testing, denoted as $\mathcal{Y}^{test}$. The challenge is for the model to generalize from the training classes to the test classes, a significant leap as it must make accurate predictions for classes without having any direct prior examples to learn from. For inverse problems, where the distribution of measurements $\mathcal{Y}$ can change based on the undersampling pattern, the term 'zero-shot solver' is sometimes applied. It also follows the precedent set by prior works that have employed this terminology for plug-and-play approaches, such as DDNM and SSD [1]. This usage is intended to reflect the solver’s ability to adapt to different undersampling patterns without retraining, similar to how zero-shot learning generalizes to new classes without prior exposure.

• Confusing notation: Regarding your concerns on the confusing notations, we've reworked our notations in all mathematical equations and algorithms to make them consistent and easy to follow. Also, we would happily address any, if you could kindly point us to specific places where confusion might arise.

Weakness 2 (Method):

• Derivation: We understand that although we did not explicitly state that our approach is derived from Bayesian principles, some terminology used in the previous introductory presentation could have led to confusion that our approach is based on Bayesian inference. Therefore, we have decided to revise the presentation accordingly. The modified texts are now highlighted in blue in the revised manuscript and as following:

  "To mitigate the ill-posedness, it is essential to incorporate an additional assumption based on prior knowledge to constrain the space of possible solutions. In this manner, the inverse problem then can be addressed by optimizing or sampling a function that integrates this prior or regularization term with a data consistency or likelihood term {ongie2020deep}. A prevalent approach for prior imposition is to employ pre-trained deep generative models {bora2017compressed, jalal2021robust}."

• Theory: We respectfully disagree with the reviewer’s sentiment of the theoretical foundation in our work and we find their characterization of our results to be, rather, unfair. We have designed an algorithm that is to a great extend theoretically motivated. Proposition 3.1, as the name suggest, only small piece of our work. Another important part of our work is the practical side, which clearly shows that our findings and analysis are on the right track.

• Discrepancy Gradient step: Regarding, discrepancy gradient step, we added some information in the paper and our reasoning. Discrepancy gradient guides $\mathbf{x}_{t-1}$ towards an equilibrium between two values x_hat_0|t and x_tilde_0|t. This step serves to add additional corrections based on the specific characteristics of the problem at hand, such as incorporating more detailed information from the measurement $\mathbf{y}$ or adjusting for errors of approximations made in previous steps. This adjustment can mitigate the effects of any errors introduced in earlier steps, thus ensuring a more consistent and robust convergence toward an optimal solution. Empirical evidence suggests that this step can moderately enhance the reconstruction results. We have now explained this clearly in the text.

[1] Gongye Liu et al. Accelerating Diffusion Models for Inverse Problems through Shortcut Sampling. 2023.

Dear Reviewer,

We appreciate your feedback and would like to first provide some response about your concern on comparison with supervised approach:

Sampling Mask Clarification: We would like to clarify that the term 'Gaussian 1D random sampling mask' actually refers to a 1D Cartesian mask. We want to make sure the Reviewer has not confused this with 'Gaussian 2D or Possoin Disk mask', which are commonly used for accelerating 3D MRI. The current Gaussian 1D random sampling mask features a full-acquisition in the readout dimension and pseudo-randomly skipped sampling steps in the phase encoding dimension. Importantly, this approach aligns with clinical practice, as many 2D MRI acquisitions employ similar 1D randomly-distributed sampling masks for compressed-sensing-based image acceleration. It's worth noting that this 'random' 1D Cartesian sampling mask pattern is applied deterministically during the pulse sequence design. We want to highlight that our last author possesses a substantial background in MRI, encompassing pulse sequence design and image reconstruction, ensuring the clinical relevance of our work.

Comparison with Fully-Supervised Methods: We acknowledge your point about the potential disadvantage of the 'random' sampling mask approach in the context of training the supervised methods. Since training data containing a mix of different sampling mask effects can be challenging for supervised approaches. Although we believe that the Gaussian 1D distribution underpinning these masks may mitigate the mixture effect to some extent, as these masks are very similar to some degree. However, conducting a comprehensive comparison against fully-supervised methods falls beyond the scope of our current study. Implementing various state-of-the-art supervised methods within our revision timeline is a complex undertaking. Hence, we opted to refer to previous publications and quote relevant numbers. We have made efforts to adjust and clarify these comparisons in Appendix Table 4 to highlight the disadvantages that the 'Gaussian 1D random sampling mask' presents to supervised learning approaches. If you strongly believe that the supervised performance reported in the Chung and Ye paper is biased, we are open to removing Table 4 from our Appendix. After reevaluating the tables in the Chung and Ye paper, we can confirm that we are reporting the DuDorNet numbers, not those of E2E-VarNet. We kindly request your review of this clarification.

As mentioned in our response to 'Weakness 3 (Results),' we have rectified the numbers we previously misreported. The corrected results indicate an improvement of approximately 1dB over the state-of-the-art, which we believe represents a reasonable advancement. We apologize for any confusion arising from our previous use of a different sampling mask.

Thank you for your continued feedback and consideration. We are currently working on writing down the Equations deriving E. 13 from Eq 12, which will be posted soon.

We are grateful for the Reviewer's constructive feedback regarding the presentation and clarity of our equations. We acknowledge that some of the derivation steps may not have been immediately clear, especially to readers who are not intimately familiar with the foundational works of DPS (Chung et al., 2022a) [1] and Ravula et al. (2023) [2].

In response to your critique, we have refined these derivations to enhance their readability and accessibility. This includes incorporating relevant assumptions directly following each equation in the Appendix 4.2, rather than grouping them together in the main document prior to Equation 13. We believe this approach has provided a clearer and more intuitive understanding of our methodology.

The only assumption made here is measurement noise is Gaussian, which is a standard assumption in most MRI simulation studies. All other \simeq are approximations derived in DPS (Chung et al., 2022a) [1] and Ravula et al. (2023) [2]. These are now clearly explained together with the Equations both in the main document and in Appendix.

We thank the reviewer for their constructive comments. Below are our point-by-point responses to Weaknesses and Questions:

Weakness 1:

We respectfully disagree with the comment that our proposed method “is a combination of DPS and DDNM”. DDNM is given by the expression $$ \hat{\mathbf{x}} = {\mathcal{A}^{\dagger}\mathbf{y}}+{(\mathbf{I}-\mathcal{A}^{\dagger}\mathcal{A})}{\bar{\mathbf{x}}}

$$On the other hand, MGDM has a closed-form solution derived from balancing two terms in the objective function.$$

\hat{\mathbf{x}}_{t} = \underset{\mathbf{x}}{\mathrm{argmin}} \frac{1}{2} ||\mathbf{y} -\mathcal{A}\mathbf{x}||_2^2 + \frac{\lambda}{2} ||\mathbf{x}-\mathbf{x}_0t||_2^2

$$MGDM comes from the closed-form of the optimization problem, reflecting an optimization-based approach. It aims to find the best route that respects closeness to the path. The closed-form solution embodies a compromise that includes fitting the model to the data and adhering to a prior reference x0|t, which is introduced through the regularization term λ. The solution represents an equilibrium between minimizing the least-squares error and the regularization term, even when expressed in closed form. This method represents an iterative approach to balance, between fitting the data and incorporating prior knowledge, reflecting a philosophical stance even in the presence of a closed-form solution. Secondly, our MGDM method differs from the DPS method's purpose and application of gradient-based updates. While both methods involve gradient-based updates, they serve distinct roles. In Algorithm 2 (step 8) of our MGDM, the gradient-descent update is used to refine our clean image initial prediction with the measurement y as a conditioning factor (i.e., approximated measurement-conditioned posterior mean). This contrasts DPS, where gradient-based updates are applied to the noisy image after each reverse sampling process. In response to your valuable feedback, we have revised our text to ensure that these distinctions between our method, DDNM, and DPS are clearer. **Weakness 2:** We appreciate the Reviewer's inquiry and understand the confusion regarding the differences in acquisition modes, particularly from a non-medical-imaging perspective. For instance, in MRI, acquiring only a few 2D slices with anatomical gaps between them is common practice to save scan time. In such cases, stacking them into a continuous 3D volume is not feasible. In our approach, we focus on reconstructing 2D images from individual 2D slices obtained from multi-slice 2D acquisitions. Each slice is treated as an independent entity, and our primary goal is to reconstruct the best possible 2D image from the raw 2D slice measurements. This is consistent with other SOTA methods we are comparing against in our paper. However, suppose one decides to acquire multi-slice 2D images covering the 3D slab without gaps between the slices, as the Reviewer suspected. In that case, potential discontinuities may indeed arise when stacking these 2D images to form a 3D volume in other views, such as coronal or sagittal. It's important to note that these discontinuities are inherent limitations of multi-slice 2D acquisitions and are not specific to our reconstruction methods. To address this limitation and ensure smoother transitions in the 3D volume, one viable solution is presented in a recent paper [1]. This approach introduces Total Variation regularization in the Z-dimension, effectively enforcing a smooth transition between slices in the 3D volume. We highly value this valuable feedback and have revised our text to emphasize the distinction between multi-slice 2D acquisitions and full-slab 3D acquisitions. Additionally, we have referenced the mentioned paper to provide readers with further insights into addressing the challenges associated with 3D reconstruction from 2D diffusion models. [1] Lee, Suhyeon, et al. "Improving 3D Imaging with Pre-Trained Perpendicular 2D Diffusion Models." arXiv preprint arXiv:2303.08440 (2023).$$

We thank the Reviewer for their constructive comments. Below are our point-by-point responses to Weaknesses and Questions:

Weakness 1: We appreciate the Reviewer's observation regarding introducing two additional parameters, ζ and ρ, in our algorithm (MGDM sampling) compared to DDNM. In response to this concern, we have conducted additional experiments to investigate the impact of these parameters on the algorithm's performance. Our findings indicate the following: (a) Once the parameters ζ and ρ are appropriately tuned within specific ranges, the performance of our MGDM becomes stable, making the parameter tuning process easier. (b) Moreover, we have observed that the parameters once tuned on a small number of specific datasets, exhibit the ability to generalize well to unseen datasets. This suggests that it is not always necessary to perform parameter tuning for each subject, as the tuned values can provide satisfactory performance across different datasets. It's worth noting that hyperparameter tuning is a universal challenge in machine learning and not unique to our method. Our revised manuscript provides detailed insights into our parameter selection process, allowing readers to replicate it in their experiments. Moreover, our ablation study of Table 3 showed that even without the two additional steps involving ζ and ρ (i.e., Ours_{no_ir}), our new proximal projection step alone can already marginally outperform DDNM.

Weakness 2: We appreciate the Reviewer's keen attention to detail. We acknowledge the error in our citation of the 'ESPIRIT' paper and apologize for this oversight. We have promptly replaced it with the correct reference, the paper by Zbontar et al. Furthermore, we have thoroughly reviewed all other references to ensure their accuracy.

Weakness 3: We appreciate the Reviewer's feedback regarding the limitations of our 2D digital simulation compared to practical 3D cone-beam or helical CT data acquisitions. It's important to clarify that our choice of 2D simulation using parallel-beam geometry aligns with the methodology adopted by several existing methods, including diffusion models such as MCG [1], SIN-4c-PRN [2], and ScoreMed [3]. This allows for a direct comparison of algorithm performances.

We acknowledge the limitations of this simplified 2D simulation and the absence of practical effects like beam hardening. To address these concerns more comprehensively, we have referenced the paper 'Combining ordered subsets and momentum for accelerated X-ray CT image reconstruction' for further discussion, as suggested by the Reviewer. In light of the short timeframe for revision, we understand that a 3D realistic CT phantom simulation is a valuable avenue for future work, and we appreciate the Reviewer's input in this regard.

Quoted from our limitation part: “It should be noted that our CT simulation adheres to the 2D parallel beam geometry assumption, aligning with the baseline models used in other studies for direct comparison. This differs from the more complex and realistic 3D cone-beam CT or helical CT simulations (Kim et al., 2014).”

Weakness 4: Thank you for pointing out the typographical errors in our manuscript. We have carefully reviewed the document and corrected the typo "fidility" to "fidelity" at the bottom of page 5, along with a thorough check for any additional typos throughout the text.

[1] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35:25683–25696, 2022b.

[2] Haoyu Wei, Florian Schiffers, Tobias Wu ̈rfl, Daming Shen, Daniel Kim, Aggelos K Katsaggelos, and Oliver Cossairt. 2-step sparse-view ct reconstruction with a domain-specific perceptual net- work. arXiv preprint arXiv:2012.04743, 2020.

[3] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint arXiv:2111.08005, 2021.