Ambient Diffusion Posterior Sampling: Solving Inverse Problems with Diffusion Models Trained on Corrupted Data

Feedback: I am confused by the bolding in Table 2. For methods trained with clean data, no values are bolded. For the R=2 and R=4 sections, A-DPS has equal performance to A-ALD in some metrics. The A-ALD metrics that are equally performant should be bolded for fairness. Please revise this table so that it is consistent.

Response: We thank the Reviewer for pointing out the inconsistencies in Table 2. We have revised the table to ensure consistent bolding throughout. Initially, metrics for methods trained with clean data ($R=1$) were not bolded as our primary focus was on comparisons between self-supervised approaches. However, based on your feedback, we have now highlighted the best performance for each inference $R$ within each training scenario, including for methods trained with clean data. Additionally, the bolding issue with metrics for our proposed methods, A-ALD and A-DPS has been corrected to ensure fairness and consistency.

See weaknesses. Although, I do have one actual question:

Question: Did you consider training your MRI diffusion model (ambient or other) with the multi-coil data? To be clear: I am not asking you to do this, as it is a fundamental departure from what you have here. That said, I wonder if you tried this. I know that in, e.g., fastMRI, the number of coils is inconsistent, but you can get around this issue with virtual coils [1]. I suspect that the overall performance could be further improved with a multi-coil diffusion process, rather than a coil-combined one.

[1] T. Zhang, J. M. Pauly, S. S. Vasanawala, and M. Lustig, “Coil compression for accelerated imaging with Cartesian sampling,” 2013

Response: We thank the Reviewer for this insightful question. Currently, we define A_train according to Eq. 2.7. This is convenient because it lets us use an image-to-image network architecture such as that used by EDM. However, an alternative could be to use the A operator directly as in Eq. 2.5, in which case the network would go from multi-coil k-space to an image. This network could be preceded by an IFFT and would effectively be a multi-coil-to-image network, and could be standardized using coil compression [1]. We agree that such a network architecture could be interesting and potentially more expressive, as we currently collapse the multi-coil information through zero-filling adjoint. Another possibility is to define A_train as the pseudo-inverse applied to A. We appreciate the suggestion and view this as an interesting direction for future research. We have incorporated this discussion in the future directions section of our revised manuscript.

We thank a lot the Reviewer fo their valuable Review. We are glad that the Reviewer appreciated the novelty of our proposed algorithm and the overall presentation of our findings.

Feedback: While the experimental performance of A-DPS is impressive, I think that the MRI results suffer from insufficient evaluation with all competitors. NRMSE being the only metric reported in Table 2 is a concern. It would be worthwhile to add some feature-based metrics like LPIPS/DISTS for all methods. You use VGG to compute FID, you can also use VGG for those metrics (since VGG agrees with radiologists' perceptions [1]). Alternatively, you could also use AlexNet as your feature extractor [1]. Either way, I think that it is important to see these results for your method and competitors. The lack of additional metrics holds this paper back. I understand the space concern given the size of Table 2, but even putting these results in an appendix would be useful to provide readers full context and make a convincing statement about A-DPS' superiority.

Response: We have expanded our evaluation to include LPIPS and DISTS metrics, which are now reported in Table 3. Additionally, Table 2 has been updated to include SSIM and PSNR alongside NRMSE. Regarding insufficient evaluation, we have implemented another self-supervised baseline stemming from SSDU [A]. Here are the results from this baseline. Similar to our previous results, we show that A-DPS outperforms other algorithms at high acceleration rates.

[A] Millard, Charles, and Mark Chiew. "A theoretical framework for self-supervised MR image reconstruction using sub-sampling via variable density Noisier2Noise." IEEE Transactions on Computational Imaging (2023).

Below we report NRMSE, LPIPS, and DISTS across 100 validation samples:

(cont')

We thank a lot the Reviewer for their time and their valuable Review. We are glad that the Reviewer appreciated the novelty of our proposed algorithm and the extensiveness of our evaluation strategy.

Feedback: Eq: 2.9: Are these equations correct?

Response: The Reviewer is absolutely right, there was a typo in the Equations. We fixed the typo and we updated the Equations 2.9 and A.3 in our revised submission. We thank the Reviewer for careful reading!

Feedback: What is the input to the diffusion model ? In Eq 2.10, it takes the additional subsampling matrix . In Eq 2.11, it takes the full forward operator

Response: The models from the Ambient Diffusion paper that were trained for natural images, take the matrix tilde $A$ as input. The Ambient Diffusion MRI models we trained for this paper, take $\tilde{P}$ as input. We updated the notation to make this point explicit. We thank the Reviewer for their useful feedback.

Feedback: An intuitive explanation of would be helpful for understanding the overall algorithm. Specifically, what does the training accomplish and why can this be directly plugged into Eq. 2.11?

Response: The training objective of Ambient Diffusion leads to a model that estimates $E[x_0 | \tilde{A} x_t, \tilde{A}]$ (or in the case of the MRI models $E[x_0 | \tilde{P} x_t, \tilde{P}])$. This is used in equation 2.11 to obtain the one-step prediction for $x_0$ from the current corrupted iterate $x_t$. Specifically, given a current iterate $x_t$, we create the linear measurements $Ax_t$ and we pass them through the model that is trained to obtain an estimation of $x_0$ given linear measurements with noise. We added this clarification in our revised submission.

Feedback: What is the reason or intuition for why an A-DPS trained on one corruption can generate posterior samples for a different corruption? A good hypothesis could better sell this property of the method.

Response: The result of the Ambient Diffusion training process is a generative model. Generative models can be used at inference time to solve a variety of different inverse problems. This idea was proposed for GANs in the CSGM paper and has been extended to diffusion algorithms and other classes of generative models. Our trained models with Ambient Diffusion leverage a DPS-type algorithm (A-DPS) to guide the generations towards solutions that explain the measurements. That’s intuitively what’s going on in this paper.

Minor Points:

Feedback: L93-94: This sentence is confusing because it sounds like the Ambient Diffusion framework is "the first framework to solve inverse problems with diffusion models learned from linearly corrupted data". I would reword this to avoid confusion.

Response: We rephrased this to: “We propose the first framework to solve inverse problems with Ambient Diffusion models. These are models that are trained using only access to linear measurements and they estimate the ambient score [...]”.

Feedback: Eq 2.4: It was not clear what represents. I would explain this matrix first before the equation, so the equation is more interpretable

Response: Good point, we updated the submission to include a short description of $\tilde{A}$ before the equation.

Feedback: L255: Typo with "In that regards the multiplicity..."

Response: Fixed, we thank the Reviewer for careful reading!

Feedback: L364: Typo with "Similar to the results in The authors of ..."

Response: Fixed, we thank the Reviewer for careful reading!

Feedback: The construction of Fig 3 is confusing with the labeling of "Recon-FS". Since it's vertically displaced, it is easy to misread. I would suggest changing this to "Error" or rotate it vertically so "Recon-FS" is on a single line

Response: We have updated the label to "Error (10x)" to improve clarity and ensure it is easier to interpret.

Questions

Question: Would it work better to train a single network on multiple acceleration rates rather than a separate network for each acceleration rate?

Response: Training a model at different corruption levels could lead to improved performance, e.g. this paper shows that a few good points can improve the performance of Ambient Diffusion models. For all our experiments, we assumed access to homogeneous data, but it is certainly interesting for future work to look at heterogeneous settings.

Question: Were the evaluations done on a single posterior sample for each method? Or do you generate many posterior samples and compute a posterior mean?

Response: Initially, we conducted a subset of experiments using multiple posteriors and observed that averaging posteriors led to slight improvements in metrics, but the benefits were consistent across all reconstruction algorithms. Due to computational and time constraints, we opted to evaluate single posterior samples.

Thank you to the authors for thoughtfully addressing the points raised in my review. After carefully reviewing the responses, I find that most of my concerns have been adequately resolved. As a result, I will maintain my overall score (6: marginally above the acceptance threshold) while increasing the soundness rating to 3 and presentation rating to 4.

Thank you! Overall, I find this to be a strong paper, particularly in the context of learning a diffusion model from undersampled MRI data, which has significant potential for practical applications. However, there are two key reasons why I did not raise my score to 8.

1. Is there a more formal mathematical explanation for why DM learned from undersampled or corrupted data can get better performance in solving certain inverse problems?

2. For MRI applications, validating the algorithm on real-world scenarios where fully sampled data is unattainable (e.g., dynamic liver MRI) would further strengthen the paper's impact and practical relevance.

We thank a lot the Reviewer for their time and their valuable Review. We are glad that the Reviewer appreciated the novelty of our proposed algorithm and the overall presentation of our findings.

Feedback: The theoretical approach in this work is compelling and well-suited for training with randomly inpainted images. However, in the MRI setting used here, different training samples may share the auto-calibration signal (ACS) region, potentially leading to an over-representation of corresponding k-space regions in the training data. A more solid approach might involve introducing a diagonal weighting matrix in the loss function to account for these oversampled regions. A similar work discussing this point for MRI reconstruction is [3].

Response: The auto-calibration signal (ACS) region is indeed shared across training samples, which could lead to an over-representation of certain k-space regions. While our current approach does not explicitly address this, we agree that introducing a diagonal weighting matrix in the loss function, as suggested, is a promising direction to mitigate this issue. We have included a discussion regarding this and reference [3] in the limitations and future directions section of our manuscript.

Furthermore, we have implemented the self-supervised baseline in [3] and have included these results under Table 2 (NRMSE, SSIM, PSNR) and Table 3 (LPIPS, DISTS). Similar to our previous results, we show that A-DPS outperforms other algorithms at high acceleration rates.

[3] C. Millard and M. Chiew, "A Theoretical Framework for Self-Supervised MR Image Reconstruction Using Sub-Sampling via Variable Density Noisier2Noise," in IEEE Transactions on Computational Imaging

Below we report NRMSE, LPIPS, and DISTS across 100 validation samples:

(cont')

Feedback: Further explanation to how the FID value is calculated for MRI data will be helpful.

Response: To quantify unconditional sample quality, we calculate FID scores using a pre-trained VGG network, following the methodology outlined in [4]. Specifically for MRI data, the k-space reconstructions are converted to magnitude images, and the FID is computed using embeddings from the convolutional layers of a VGG network trained on natural images. While VGG is not optimized for MRI data, it provides a perceptual metric that aligns with image sharpness and structural consistency, as validated in related works like [A].

We also compute the distribution metrics (Conditional FID) between samples from the posterior and the ground-truth distribution, using the same VGG network (included in Table 1). We have clarified this process in the revised manuscript.

[A] Bendel, Matthew, Rizwan Ahmad, and Philip Schniter. "A regularized conditional GAN for posterior sampling in image recovery problems." Advances in Neural Information Processing Systems 36 (2024).

Feedback: Some minor suggestions regarding the paper's presentation: (1) In Table 3, it would be beneficial to standardize the significant digits for the FID values to improve readability. (2) For Figures 7 and 8, adjusting the color bar could enhance visual clarity and make the figures easier to interpret.

Response:

1. We thank the Reviewer for their suggestion regarding Table 3. We have combined Tables 1 and 3 into a single, standardized Table 1 in the main paper. This table now presents both Unconditional and Posterior Conditional FID scores with all values standardized for improved readability. Additionally, we have highlighted the best-performing method for each training acceleration factor $R$ to enhance visibility.

2. Based on the Reviewer’s feedback, we have removed Figures 7 and 8. In their place, we have introduced an expanded Table 2, which extends the posterior reconstruction metrics table (Table 2) by including SSIM and PSNR metrics alongside NRMSE. We have also included a new Table 3 where we report LPIPS and DISTS for the accelerated MR reconstruction experiments. Significant values are highlighted in these tables to improve clarity and emphasize key results.

Thank you for providing the additional experiments on MRI data where no ground truth is available. These experiments addressed most of my concerns, so I am increasing my score to 8. I also appreciate the authors' efforts to make the paper more comprehensive overall.

Regarding the question: "Is there a more formal mathematical explanation for why diffusion models learned from undersampled or corrupted data can achieve better performance in solving certain inverse problems?"—I generally agree with your intuition. To the regularization you mentioned, I suggest considering the perspective I outlined in the first round (which is similar to but distinct from DRP [1]). I suggest including a discussion or future work section to elaborate on your hypothesis.

[1] Y. Hu, M. Delbracio, P. Milanfar, and U. S. Kamilov, “A Restoration Network as an Implicit Prior,” Proc. Int. Conf. Learn. Represent. (ICLR 2024) (Vienna, Austria, May 7-11).

We thank a lot the Reviewer for their time and their valuable Review. We are glad that the Reviewer appreciated the effectiveness of our proposed algorithm.

Feedback: Despite promising experiments, the novelty of the paper seems limited; the authors primarily extend ambient diffusion to MRI and adapt it within the DPS framework for solving inverse problems.

Response: Thank you for the feedback. While we agree that extending DPS to A-DPS is technically straightforward, we believe the novelty lies in the finding that models trained on corrupted data can outperform those trained on clean data, especially for severely ill-posed problems. Furthermore, we draw a parallel to the DPS algorithm, which uses a simple approximation to address the intractability of the likelihood term. Despite its simplicity, DPS has been widely adopted and demonstrated significant success across diverse applications. Similarly, while A-DPS represents an adaptation of DPS within the Ambient Diffusion framework, the experimental results showcasing its efficacy in solving inverse problems highlight its importance and utility. We hope the reviewer recognizes this contribution as a valuable advancement for practical inverse problem-solving.

Feedback: Below Eq. 2.11, the authors state, "We remark that similar to DPS, the proposed algorithm is an approximation to sampling from the true posterior distribution". However, the accuracy of this approximation should be rigorously compared to that of the original DPS, as it relies on estimating the score function from corrupted data.

Response: The accuracy of the DPS approximation itself is poorly understood – the authors of DPS only provide a rough characterization of the Jensen gap in Theorem 1 of their paper. Following the Reviewer’s recommendation, we included the same result for A-DPS in our paper in Equation 2.12 of the updated submission.

Feedback: The authors should consider including additional baselines [...]

Response: We thank the Reviewer for their suggestion. We implemented another recent self-supervised baseline stemming from the SSDU framework [1]. We have included these results in our revised submission under Table 2 (NRMSE, SSIM, PSNR) and Table 3 (LPIPS, DISTS). Similar to our previous results, we show that A-DPS outperforms other algorithms at high acceleration rates.

[1] Millard, Charles, and Mark Chiew. "A theoretical framework for self-supervised MR image reconstruction using sub-sampling via variable density Noisier2Noise." IEEE Transactions on Computational Imaging (2023).

Below we report Normalized Root Mean Squared Error (NRMSE) across 100 validation samples:

(cont' in the next response)

Below we report LPIPS and DISTS averaged across 100 validation samples:

Feedback: The paper needs to be further polished, for example, Eq 2.9 seems incorrect.

Response: We are grateful to the Reviewer for catching this typo. We updated Equation 2.9 and Equation A.3. in the proof with the correct notation.

We thank a lot the Reviewer for their time and their valuable Review. We are glad that the Reviewer appreciated the effectiveness of our proposed algorithm and the presentation of our work.

Feedback: Although this paper introduces the first self-supervised diffusion method for inverse problems, the overall novelty is somewhat limited. It appears to be a straightforward integration of existing ambient diffusion with well-known diffusion posterior sampling, with application to MRI reconstruction. While the investigation holds value from a research perspective, the new insight/concept might be limited.

Response: While we agree that extending DPS to A-DPS is technically straightforward, we believe the novelty lies in the finding that models trained on corrupted data can outperform those trained on clean data, especially for severely ill-posed problems. Furthermore, we draw a parallel to the DPS algorithm, which uses a simple approximation to address the intractability of the likelihood term. Despite its simplicity, DPS has been widely adopted and demonstrated significant success across diverse applications. Similarly, while A-DPS represents an adaptation of DPS within the Ambient Diffusion framework, the experimental results showcasing its efficacy in solving inverse problems highlight its importance and utility. We hope the reviewer recognizes this contribution as a valuable advancement for practical inverse problem-solving.

Feedback: In the training phase (as described in Equation 2.10), the network takes \tilde{A} as input, considering one type of imaging artifact. However, during testing (Equation 2.11), the network uses A as input, which features a different imaging artifact. This inconsistency raises concerns about potential suboptimal performance. Could the authors clarify this aspect?

Response: We thank the Reviewer for raising this important point. We noticed experimentally that the model is robust at inference time to matrices $A$, even when trained on matrices $\tilde{A}$. This is consistent with our experiments in Section 3.3 on pre-trained models for natural images. That said, in our codebase, we used $\tilde{A}$ for our experiments, so we updated Equation 2.11 to reflect this.

Feedback: There is insufficient discussion regarding why the proposed A-DPS method outperforms FS-DPS, especially at high undersampling rates. Given that FS-DPS uses a diffusion model trained on fully sampled data, the results seem unexpectedly favorable.

Response: Respectfully, we believe this observation is not a weakness but rather the main finding of our paper. The superior performance of A-DPS at high undersampling rates arises because models trained on corrupted data can be better suited to handle the ill-posed nature of inverse problems under severe corruption. By being exposed to corrupted training data, A-DPS can learn more robust priors that generalize better to high-acceleration regimes, unlike FS-DPS, which may overfit clean data and fail to adapt effectively to severe undersampling. We have validated this extensively across different data modalities and inverse problems in our experiments. Based on your feedback, we have incorporated this discussion in our revised submission.

Feedback: The manuscript lacks discussions with recent theoretical developments in self-supervised learning for MRI [1,2], which share similar proof strategy with Theorem 1 in this paper.```

Response: We thank the Reviewer for raising this point. We have included a discussion of these methods in our revised submission. Regarding [1], while it shares a similar proof strategy to Theorem 1 in our work, its primary focus is on learning a restoration model tailored to a specific corruption type. In comparison, Ambient Diffusion aims to learn a generative model, which provides greater flexibility at inference time to address a wider variety of corruption types and inverse problems. We believe this distinction highlights the broader applicability of our framework.

(cont')

Feedback: The numerical validation could be improved, with certain experimental setups are not fully explored (see specific questions below). Additionally, the baseline methods for MRI reconstruction cited are well-known but somewhat dated (2019, 2020). Given the extensive literature on MRI reconstruction, these baselines may not represent the current state of the field. Comparing recent alternatives, such as [1-3] and self-supervised or supervised methods in survey [3] and [4], respectively, could improve this study. There is also no comparison between A-OS and other end-to-end methods. Despite that, I fully understand the time constraints of the rebuttal period and do not intend to suggest that the authors include comparisons with new baselines.

Response: We thank the Reviewer for these valuable suggestions. Our primary focus with MRI evaluation was on self-supervised baselines, as these align more closely with our core contributions. We believe that state-of-the-art FS-DPS provides the best fully supervised comparison. Additionally, we include end-to-end MoDL, which, while slightly dated, serves as a close proxy for [3].

In regards to self-supervised baselines, we agree that including more recent baselines could enhance the study. We have implemented [1] and have included it in our revised submission under Table 2 (NRMSE, SSIM, PSNR) and Table 3 (LPIPS, DISTS). Similar to our previous results, we show that A-DPS outperforms other algorithms at high acceleration rates.

[1] Millard, Charles, and Mark Chiew. "A theoretical framework for self-supervised MR image reconstruction using sub-sampling via variable density Noisier2Noise." IEEE transactions on computational imaging (2023).

Below we report NRMSE, LPIPS, and DISTS across 100 validation samples:

(cont')

Question: The further corruption ratio is simply R+1. How does this further corruption ratio impact the performance, says R+2, R+3, etc?

Response: The further corruption during training should be as minimal as possible to avoid throwing away training data. We have not studied the effect of increased additional corruption at training time, but we believe that this would lead to performance deterioration.

Question: Does the supervised DPS share the same network architecture with the proposed method?

Response: Both models are based on the EDM codebase. The EDM model trained on the fully-sampled (R=1) dataset, as well as the four EDM models trained on L1-Wavelet reconstructions at each acceleration factor R = 2, 4, 6, 8, were all trained with 65 million parameters. The Ambient Diffusion models on the other hand were trained with 36 million parameters, for faster training. There were no other differences in the network architecture.

Question: How does the coil sensitivity map in this training dataset and the L1-wavelet baseline be estimated?

Response: We estimate the sensitivity maps using ESPIRiT [A] from a fully sampled auto calibration signal (ACS).

[A] Uecker, Martin, Peng Lai, Mark J. Murphy, Patrick Virtue, Michael Elad, John M. Pauly, Shreyas S. Vasanawala, and Michael Lustig. "ESPIRiT—an eigenvalue approach to autocalibrating parallel MRI: where SENSE meets GRAPPA." Magnetic resonance in medicine 71, no. 3 (2014): 990-1001.

Question: What is the time complexity of the methods?

Response: A-OS requires 1 function evaluation. The number of function evaluations required by A-DPS is the same as DPS and depends on the discretization used in the sampler.

Question: What is the actual under sampled ratio considering the sampling mask for different R has the same 20 lines of the auto calibration signal?

Response: The acceleration factor accounted for sampling of the central lines, i.e. it was defined as the matrix size divided by the number of sampled points.