An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

We thank the reviewer for the constructive suggestions. We are happy that the reviewer recognizes our idea as "simple" and "elegant", our method "can have a substantial practical impact", and our paper "is well structured". We have provided explanations and clarifications regarding all the concerns, as shown below:

1. Qualitative results with artifacts

  Thanks for your careful review! In Figure 3, we used the officially released Ambient Diffusion checkpoint for image posterior sampling rather than training the SOTA models ourselves. We double-checked our code and found no issues.

  However, training clean diffusion models from corrupted images is extremely challenging, and the missing information can cause qualitative results on some datasets to appear artefactual. The good performance in Figure 7 of the Ambient Diffusion paper leverages AFHQ, whereas our results report on CIFAR10. CIFAR10 images contain more high-frequency features than AFHQ images, making inverse problems more challenging and leading to more reconstruction artifacts. The original Ambient Diffusion paper also noted similar issues, stating that learning clean diffusion models on CIFAR10 is harder compared to CelebA-HQ and AFHQ (Sec. 5.1 in their paper) :

“for CelebA-HQ and AFHQ, we manage to maintain a decent FID score even with 90% of the pixels deleted. For CIFAR-10, the performance degrades faster, potentially because of the lower resolution of the training images”.

2. Dependence on clean data for initialization

  Our method relies on a small amount of clean data as starting point, which we acknowledge as a limitation, as discussed in the conclusion. However, we have minimized this requirement to as few as 10 clean images (Fig. 5(a)). Additionally, we have explored alternative initialization strategies, including using out-of-distribution (OOD) clean images or classical pre-processing methods to initialize EM iterations (see Fig.5 and Tbl. 4 in the main paper). Our preliminary findings suggest that our EM approach can be initialized using low-frequency image statistics (e.g., smoothness) learned from OOD or pre-processed images. Investigating the elimination of the clean data requirement is an important research direction that we leave for future work. We will discuss this further in the revision.

3. Gaussian noise in forward models

  All examples in the paper involve Gaussian noise in their degradation models, including inpainting and deblurring, as practical imaging problems always include measurement noise. However, the noise added to masked and blurry images is minimal (std=0.0001 and 0.02, respectively) and may not be visible. We also support different noise models, as discussed below.

4. Extending to other types of measurement noise

  Our method can be extended to more realistic noise models. As discussed in the original DPS paper [1], it can handle various noise statistics, such as Gaussian and Poisson. Our EM framework has flexibility in replacing our E-step with DPS under new noise assumptions or even more advanced posterior sampling algorithms such as plug-and-play Monte Carlo (PMC)[2].

5. Learning parameters of degradation models

  Thank you for suggesting this important topic! Learning degradation models complements our method, which focuses on understanding the relationship between observations and underlying images rather than just image priors. This challenge can also be addressed using the EM framework, where the M-step iteratively estimates the degradation model’s parameters by incorporating the gradient of the data fidelity term with respect to these parameters. We will discuss it in the paper and continue to explore this exciting direction in the future.

We thank the reviewer once again for the insightful suggestions. We hope our explanations have clarified all the concerns. We are happy to address any further questions or discussions the reviewer may have. Thank you!

Reference:

[1] Hyungjin Chung, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, Jong Chul Ye. Diffusion Posterior Sampling for General Noisy Inverse Problems, ICLR 2023.

[2] Yu Sun, Zihui Wu, Yifan Chen, Berthy Feng, and Katherine L. Bouman. Provable Probabilistic Imaging using Score-based Generative Priors, arXiv 2023.

I don't have any additional questions.

We thank the reviewer for the insightful suggestions. We are happy that the reviewer recognizes our approach "addresses a significant challenge in the field", and our experiments are "extensive" and "show the effectiveness of the proposed method". We will release all the code and the checkpoint to the community to stimulate future research. Below we address each question of the reviewer.

1. Source code

  We are happy to provide a private repository of our code to address your concerns. Following NeurIPS author guidelines,

If you were asked by the reviewers to provide code, please send an anonymized link to the AC in a separate comment, make sure the code itself and all related files and file names are also completely anonymized.

  we have sent a link to the AC for anonymity verification and will share it with you once approved.

2. Impact of number of clean images on model performance

  Preliminary experiments in Section 5.4 and Figure 5(a) show that using 10 to 500 clean images results in similar posterior sampling outcomes after the first EM stage, indicating that the diffusion model successfully learns important image statistics with limited data. Since clean images initialize the diffusion model by providing low-frequency statistics, the model performance will not be significantly affected by their number.

3. Completely eliminating the dependence on clean data

  This is indeed an important research direction. In this paper, we have explored strategies to eliminate the dependence on clean data, such as initializing EM iterations with classical pre-processing methods, as detailed in Appendix B. Additionally, leveraging a generative foundation model like Stable Diffusion as the initial model could be a potential alternative solution, as suggested by other reviewers (Consistent Diffusion Meets Tweedie, ICML 2024). We will continue to explore this direction in future work.

We thank the reviewer once again for the insightful suggestions. We greatly appreciate the recognition and are more than willing to address any further questions or discussions the reviewer may have. Thank you!

We thank the reviewer for the thoughtful comments. We believe there may have been some misunderstandings regarding our work, which we will clarify. Below we carefully address all the questions from the reviewer.

1. Quantifying the number of required clean images

  We totally agree that quantifying the amount of required clean images is important and we have carefully investigated this problem in Section 5.4, Figure 5(a), and Table 4 of the main paper. Our experiments show that diffusion models trained on as few as 10 clean images provide adequate initialization. Moreover, to further reduce the need for clean data, we also explored using out-of-distribution (OOD) or pre-processed images for initialization, which achieved comparable results. We will highlight these findings more clearly in the revision.

2. Convergence claims & conditions for the forward operator

  In our paper, convergence refers to reaching a local minimum of the optimization problem, as established by the properties of the classical EM algorithm and supported by previous theoretical studies [1]. We discuss this briefly in Section 3.3, lines 124-126, of the original paper. The forward operator and noise process do not affect this type of convergence but do influence the convergence speed by determining the amount of information in the corrupted observations. We will clarify the sentences and definitions about convergence in the final version.

3. Why sparse clean data provide “good” initialization

  This question has been partially addressed in Section 4.2, lines 160-165, of the original paper. Sparse clean data helps the initial diffusion model learn common structures in natural images, such as continuity, smoothness, and object profiles. In this context, “good” means capturing important low-frequency statistics of natural images. From the EM algorithm perspective, “good” initialization means defining a local minimum solution close to the global minimum.

4. Limitation of requiring clean images

  Considering the minimal number of clean images required for initialization and the alternative strategies we've validated for our EM framework (as explained in answer (1)), we respectfully disagree that this is a major limitation of our algorithm. Additionally, in many scientific applications, accessing some clean images is reasonable, as the lack of clean images is often due to cost rather than physical restrictions. For instance, in fluorescent microscopy, while high-SNR images are limited due to the need for long exposure times and the risk of phototoxicity, it is still feasible to obtain some high-SNR images to support the training of clean diffusion models from extensive low-SNR microscopy data.

We thank the reviewer once again for the valuable comments. We hope our explanations have clarified all the questions. We are happy to address any further concerns or discussions the reviewer may have. We kindly hope the reviewer could consider raising the score if satisfied with our responses. Thank you!

Reference:

[1] Dempster, A.P.; Laird, N.M.; Rubin, D.B. Maximum Likelihood from Incomplete Data via the EM Algorithm. Journal of the Royal Statistical Society, 1977

We thank the reviewer for the additional comments. We are pleased that many of your concerns have been addressed. Below, we provide additional clarification regarding the convergence claim and the conditions of the forward operator.

1. Convergence Claim

  As explained in our rebuttal, the convergence we refer to is the local convergence inherited from the EM framework. We explicitly mentioned this in our original paper (lines 124-126). This claim is further supported by our extensive experimental results across various imaging tasks (see Appendix C). We promise to rephrase our abstract and add further clarifications on this point in the revised manuscript to avoid any confusion.

2. Forward Operator

  In our previous responses, we misunderstood the reviewer's comments as asking us to discuss the effects of various corruption types (e.g., inpainting, denoising, deblurring) on the method’s effectiveness. We fully agree that there are certain cases where recovering a clean distribution from corrupted data is information-theoretically impossible, such as when the forward operator matrix is entirely composed of zeros. We will include additional discussions on these scenarios in the revised manuscript.

We hope our explanations address all the concerns of the reviewer. We would greatly appreciate it if the reviewer could reconsider the evaluation of our paper. Thank you!

We thank the reviewer for the invaluable feedback. As noted by the reviewer, our idea is "fresh" and "elegant," our algorithm is "agnostic to the type of corruption," and our results "outperform the prior state-of-the-art." Below we provide more explanations and additional results to address each concern from the reviewer.

1. Comparison to baselines

  Ambient Diffusion pioneered in learning clean diffusion models from masked images, achieving current STOA performance. While this insightful idea was originally designed for inpainting task, we feel it is more comprehensive to test it on denoising and deblurring tasks, using its original implementation with a naïve further corruption technique (randomly masking pixels for all tasks). We acknowledge that the original algorithm cannot be trivially extended to different settings. To avoid misinterpretation, we are open to labeling Ambient Diffusion's denoising and deblurring results as "not applicable," as these tasks fall outside its original design scope. In addition, we also provided SURE-Score in our paper as a versatile baseline capable of handling various corruptions.

  We agree that validation on different levels of corruption and other datasets would strengthen the understanding of our method. As suggested by the reviewer, we conducted additional experiments comparing the EM approach with Ambient Diffusion across different corruption levels (masking ratios of 0.4, 0.6, 0.8). These new results, presented in the table below and the attached PDF, demonstrate our method’s robustness and superior performance across diverse conditions. We will incorporate these findings as well as more experiments on new datasets (e.g., CelebA) in the revision.

2. Recent work "Consistent Diffusion Meets Tweedie"

  We thank the reviewer for pointing out this relevant ICML 2024 work, which cleverly finetunes Stable Diffusion (SD) to leverage the pre-trained knowledge in denoising tasks. We will add it in related work. However, as our method focuses on training diffusion models from scratch using corrupted observations, we feel the two works belong to different categories. A fairer comparison with the ICML work would involve applying an EM framework to finetune SD with noisy observations, which presents intriguing future research. We will explore it in future work.

3. The limitation of clean data initialization

  Our method relies on a small amount of clean data as a starting point, which we do acknowledge as a limitation, as discussed in the conclusion. However, we have minimized the requirement of the clean data to as few as 10 or 50 clean images (Fig. 5(a)), which we believe is a reasonable assumption for many scientific applications. For instance, in fluorescent microscopy imaging, high SNR images are limited due to long exposure times and potential phototoxicity rather than physical restrictions. Therefore, acquiring a small number of high-SNR images to aid training is feasible.

  On top of that, we have explored alternative initialization strategies, including using out-of-distribution (OOD) clean images or classical pre-processing methods to initialize EM iterations (see Fig. 5 and Tbl. 4 in the paper). Our preliminary findings suggest that our EM approach can be initialized using low-frequency image statistics (e.g., smoothness) learned from OOD or pre-processed images. Investigating ways to eliminate the need for clean data would be an important future research direction. We will discuss more in the revision.

4.DPS does not offer true posterior samples

  We acknowledge that DPS-based methods may not always accurately sample the true posterior due to their simplified assumptions. Our key contribution is the EM approach itself, which provides flexibility in choosing the posterior sampling method in its E-step. For example, it can replace DPS with more advanced posterior sampling techniques with stronger theoretical guarantees, such as plug-and-play Monte Carlo (PMC)[1], which ensures non-asymptotic stationarity. We will discuss it in the revision.

5. Training efficiency

  Although our method involves multiple training steps, as discussed in Sec. 4.3, we avoid training diffusion models from scratch at each M-step. Instead, we fine-tune the model in most iterations, only training from scratch in the final 1-3 steps to prevent memorizing poor samples from the initial stages. This ensures rapid convergence in each iteration and significantly reduces training time. In practice, our method has a similar training time as Ambient Diffusion (ours is 2 days with 4 NVIDIA A800 GPUs).

6. Discussion on corruption types that are impossible to recover

  We agree that for extremely corrupted images (e.g., >99% pixels masked), reconstructing the true distribution becomes challenging or even impossible due to the loss of spatial relationships among pixels. This is a common issue across all baselines. A potential solution could involve incorporating more prior knowledge, such as fine-tuning SD instead of training from scratch, similar to ICML 2024 work. We will discuss it in the revision.

7.Redundancy in the checklist section

  Nice catch! We will correct it.

We thank the reviewer again for the valuable insights and suggestions. We hope our explanations have clarified all the concerns. We are happy to address any further questions or discussions the reviewer may have. Thank you!

Reference:

[1] Yu Sun, Zihui Wu, Yifan Chen, Berthy Feng, and Katherine L. Bouman. Provable Probabilistic Imaging using Score-based Generative Priors, arXiv 2023.