Paper: Meta-Guided Diffusion Models for Zero-Shot Medical Imaging Inverse Problems (ICLR-2024)

Reviewer comment (Experiments):
Miss the baseline comparisons with supervised method and [Robust compressed sensing MRI with deep generative priors]. I think the current benchmark comparisons are fair, but would curious on how MGDM compare to supervised methods [like MODL: https://arxiv.org/abs/1712.02862] and one of the earlier generative model for MRI recon using Langevin dynamics: [Robust compressed sensing MRI with deep generative priors].

Author response:
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.

Reviewer comment (Experiments):
I would appreciate more results on FastMRI dataset, I only see one visual example - Figure 5, without the undersampling pattern, and much descriptions.

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

Reviewer comment (Evaluation):
From results in Figure 2, I am amazed but also confused on the third column ACR=24, despite the effective undersampling rate, with only a few lines, I don't expect the model accurately predicting the tumor. Could you please elaborate on this results? I would like to learn more on what you think?

Author response:
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.

Reviewer comment (Presentation):
The figures are not of high-resolution and some of them are compressed. For example, Figure 5, I can actually see jpeg artifacts in all images, which is not acceptable for imaging related papers. Please fix this problem. Also the knee orientation is up-side-down, please fix it.

Author response:
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.

Reviewer comment (Theory):
Please provide some comprehensive explanation on why MGDM can recover the brain tumor given such a high undersample rate.

Author response:
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.

Reviewer comment (Presentation):
Could you add some visual results for your ablation studies? I would want to see how the reconstruction results look like with different ablations.

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

Reviewer comment (Theory):
**Method**: the presentation of the method is very handwavy: there is no derivation of why this sampling is supposed to work even with very strong assumptions. The only "theoretical" grounding is proposition 3.1 which to me amounts a bit to mathiness given its simplicity. Another example is the introduction of the discrepancy gradient which is not discussed.

Author response:
**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.

Reviewer comment (Evaluation):
**Results**: given how good the results are, it's important to question why the improvement is so big. If I focus on knee MRI reconstruction, there is a +4dB improvement in PSNR: this is absolutely huge but it isn't discussed. In particular, it should be viewed in comparison with fully supervised methods like the ones presented in the 2019 fastmri challenge which are nowhere near the performance reported here.

Author response:
**Weakness 3 (Results):**

We appreciate your comments about our results. After thorough checking, we noticed that we had reported the unifrom1D instead of the Gaussain1D from the Score-MRI paper. We apologize for this insight and have now corrected their values in Table 2. 

| Method                                   | 4x ACR PSNR↑     | 4x ACR SSIM↑    | 8x ACR PSNR↑     | 8x ACR SSIM↑    |
|------------------------------------------|------------------|-----------------|------------------|-----------------|
| DPS (Chung et al., 2022a)                | 22.41±3.33       | 0.650±0.080     | 21.87±2.91       | 0.607±0.076     |
| DDNM (Wang et al., 2022)                 | 35.87±2.68       | 0.873±0.065     | 34.04±2.70       | 0.847±0.071     |
| Score-MRI (Chung & Ye, 2022)             | 33.96±1.27       | 0.858±0.028     | 30.82±1.37       | 0.762±0.034     |
| MGDM (ours)                              | **36.94±2.70**   | **0.888±0.062** | **34.98±2.66**   | **0.856±0.070** |
|------------------------------------------|------------------|-----------------|------------------|-----------------|

Additionally, in Table 4 of the Appendix, we present supervised results. Consistent with other zero-shot inverse problem solvers, such as those discussed in Chung et al. (2022), Wang et al. (2022), and Kawar et al. (2022), our MGDM model demonstrates superiority over existing supervised methods.

Reviewer comment (Reproducibility):
**Code**: while some code is provided, the README has not been updated from the Wang et al. repo, which makes it difficult to know how to look at the relevant code, i.e. where the new sampling method is introduced.

Author response:
**Weakness 3 (Code):**
We would like to clarify that the code supporting the findings of our study is based on three GitHub repositories, each corresponding to a dataset: DDNM by Wang et al., Score-Med by Song et al., and Score-MRI by Chung et al. We need time to enhance the documentation and improve the code's readability. We recognize the importance of reproducibility in scientific research and plan to make the code publicly available once these improvements are made. In the meantime, I have added a simple README file that explains how to run the code.

in our paper, we do not claim to have derived our algorithm directly from Bayesian inference. While we initially discussed Bayesian inference, this aspect has since been refined and is not the basis for our algorithm's derivation.

To address your concerns regarding the gap, we have incorporated additional lines in our revised manuscript.

**Question 1:** In our paper, we do not claim to have derived our algorithm from Bayesian inference. While we initially discussed Bayesian inference, this aspect has since been refined and is not the basis for our algorithm's derivation.

**Question 2:** To address your concerns regarding the gap, we have incorporated additional lines in our revised manuscript. 
  > A central challenge in zero-shot inverse problem solving is how to guide an unconditional prediction to conform to the measurement information. Existing methods rely on deficient projection or inefficient posterior score approximation guidance, which often leads to suboptimal results. In this paper, we propose a Meta-Guided Diffusion Model (MGDM) that tackles this challenge through a _bi-level_ guidance strategy, where the _outer level_ solves a proximal optimization problem to impose measurement consistency and the _inner level_ approximates the measurement-conditioned posterior mean as the initial prediction.

**Question 3:** We have clarified these points in the text for better understanding. The proximal optimization process is outlined as Step 10 in Algorithm 2. Additionally, the 'discrepancy gradient step' is synonymous with what we refer to as the 'refinement step’.

1. The exclusion of **proximal optimization** (step **9**) and **refinement** (step **11**) termed `no-pr`.
2. The omission of **initial prediction** (step **8**) and **refinement** (step **11**) designated as `no-ir`.
3. The absence of **initial prediction** alone (step **8**) noted as `no-i`.
4. The removal of **refinement** alone (step **11**) referred to as `no-r`.

Reviewer comment (Theory):
**Question**: How can the new sampling be derived using a bayesian formulation?

Author response:
**Question 1:** In our paper, we do not claim to have derived our algorithm from Bayesian inference. While we initially discussed Bayesian inference, this aspect has since been refined and is not the basis for our algorithm's derivation.

Reviewer comment (Novelty):
**Question**: How can you explain the gap between this work and previous works?

Author response:
**Question 2:** To address your concerns regarding the gap, we have incorporated additional lines in our revised manuscript. 
  > A central challenge in zero-shot inverse problem solving is how to guide an unconditional prediction to conform to the measurement information. Existing methods rely on deficient projection or inefficient posterior score approximation guidance, which often leads to suboptimal results. In this paper, we propose a Meta-Guided Diffusion Model (MGDM) that tackles this challenge through a _bi-level_ guidance strategy, where the _outer level_ solves a proximal optimization problem to impose measurement consistency and the _inner level_ approximates the measurement-conditioned posterior mean as the initial prediction.

Reviewer comment (Writing):
**Question**: What is (pr) i.e. proximal optimization and refinement?

Author response:
**Question 3:** We have clarified these points in the text for better understanding. The proximal optimization process is outlined as Step 10 in Algorithm 2. Additionally, the 'discrepancy gradient step' is synonymous with what we refer to as the 'refinement step’.

1. The exclusion of **proximal optimization** (step **9**) and **refinement** (step **11**) termed `no-pr`.
2. The omission of **initial prediction** (step **8**) and **refinement** (step **11**) designated as `no-ir`.
3. The absence of **initial prediction** alone (step **8**) noted as `no-i`.
4. The removal of **refinement** alone (step **11**) referred to as `no-r`.

Reviewer comment (Novelty):
Seems the method is a combination of DPS (Chung et al., 2022a) and DDNM (Wang et al., 2022), which somehow limits the contribution of this paper.

Author response:
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}}}
$$, representing a direct analytical approach, which utilizes a straightforward formula involving matrix operations. The solution is a constructive combination of the least-squares fit and the null space component, aiming to be as faithful to the data measurement as possible while respecting the initial prediction. This is a deterministic and direct method, which clearly follows the path dictated by the data and initial conditions.

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.

Reviewer comment (Writing):
What do you mean by 'a novel class of "fully" probabilistic Deep Learning Models (DLMs)'?

Author response:
**Question 1**: By fully probabilistic deep models we in fact meant deep generative models. Generative models can indeed be considered a subset of fully probabilistic deep learning models. Recall that fully probabilistic models capture the entire probability distribution of the dataset.  They can be used for a wide range of tasks, including but not limited to prediction, regression, and generation.  On the other hand, generative models are specifically designed to sample and generate new data points that are like the ones they were trained on. To do this, they must learn the probability distribution of the dataset. Therefore, they inherently deal with the probabilistic nature of the data. In essence, generative models are one application of fully probabilistic deep learning models, with the specific task of data generation. The term "fully" in the context of fully probabilistic deep learning models is meant to convey completeness in terms of probabilistic modeling.  

However, to avoid confusion, we have now replaced "fully probablistic deep learning models" with "deep generative models".

Reviewer comment (Presentation):
Seems in section 2.1 that there is no need to include both Discrete-time formulation and Continuous-time formulation, just introduce the one related to this paper.

Author response:
**Question 2**: 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.

Reviewer comment (Theory):
The proximity term penalizes deviations from the initial estimate, which seems like the data consistency in MRI acceleration?

Author response:
**Question 3**: As explained in the introduction and method sections, the proximity term in Eq [10] ensures the solution remains close to its initial prediction estimated from the denoiser pre-trained by diffusion models. This is not about the data consistency in MRI. In fact, the fidelity term in Eq [10] enforces that the reconstructed image is consistent with the acquired measurements in the transformed domains. We have now made this clear in the text.

Reviewer comment (Experiments):
Why did not test on the brain data from fastMRI? Can we share the same model for BraTS and brain from fastMRI?

Author response:
**Question 5**: The choice of dataset for model validation is crucial for demonstrating the model's generalizability and robustness. In our current study, we focused on the BraTS dataset due to its diverse image contrasts and tumor pathologies, as also adopted by a baseline method we are comparing against [1]. 

To address your point on model sharing between BraTS and fastMRI, it is possible to apply the same model to both datasets. However, adaptations for out-of-distribution (ODD) are necessary to account for the different characteristics and objectives of each dataset. We acknowledge the importance of evaluating our model on multiple datasets to establish its generalizability, and we did this on different modalities and anatomical choices. We also are considering fastMRI data in future work to validate our model's performance for ODD analysis.

Reviewer comment (Writing):
I am not quite sure what the mentioned 'zero-shot setting' is? For MRI accereration, we do have the fully sampled k-sapce data as the reference to train the neural network. What is the 'zero-shot setting' here?

Author response:
**Question 6**: We want to make sure this is clear to the Reviewer that different from existing supervised methods using the fully sampled k-space as the reference (label) for training, our approach does not involve any paired data (undersampled and fully-sampled k-space) into our unsupervised training process. The undersampled k-space is only used for sampling/reconstructing the corresponding MRI image during inference time. 

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

Reviewer comment (Experiments):
Can we share the model for different acceleration rates?

Author response:
**Question 7**: Yes, as explained in response to question 6, the nature of zero-shot image reconstruction allows us to adapt the model for various acceleration rates. All the results reported for each test dataset on different acceleration factors share the same model for that dataset.

Reviewer comment (Experiments):
Two extra parameters $\zeta$ and $\rho$ are introduced in the proposed Algorithm 2 (MGDM sampling) in comparison to Algorithm 1 (DDNM sampling). It is clear not how much of the improvements over DDNM were obtained by painstakingly tuning these two new parameters.

Author response:
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.

Reviewer comment (Writing):
While a slice-wise 2D imaging simulation is appropriate for demonstrating the practical efficacy for MR reconstruction and the MRI simulation, as in the fastMRI paper, appears to be realistic, the 1989 ESPIRIT reference for estimate the parallel multi-coil sensitivities appears to be incorrect. The cited 1989 paper doesn't consider any special considerations for the parallel MR imaging problem and the correct reference appears to be: ESPIRiT — An Eigenvalue Approach to Autocalibrating Parallel MRI: Where SENSE meets GRAPPA, by Uecker, et al., 2014. This is the reference from the fastMRI paper by Zbontar et al.

Author response:
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.

Reviewer comment (Experiments):
In the case of CT imaging, 3D cone-beam CT or helical CT are the common practical data acquisition techniques and 2-D simulations are not particularly convincing. Please refer to the following paper for a reasonable 3-D simulation as well as results on real sinogram data: Kim, Donghwan, Sathish Ramani, and Jeffrey A. Fessler. "Combining ordered subsets and momentum for accelerated X-ray CT image reconstruction." IEEE transactions on medical imaging 34, no. 1 (2014): 167-178. It may be acceptable that the authors do not simulate practically important effects such as beam hardening, but a 2-D simulation and an FBP baseline can be misleading. It is also common practice in CT reconstruction to use an anthropomorphic digital phantom during simulation in order to uncover reconstruction artifacts that might get hidden when using a reconstructed CT image, itself full of noise and other artifacts.

Author response:
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).”

Reviewer comment (Writing):
Some typos exist in the paper, e.g. "fidility" at the bottom of page 5.

Author response:
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.