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

Summary

The authors present a new method for training diffusion models from corrupted data. The presented method is an EM algorithm that in turn (i) uses the current model to sample clean images given the noisy input and then (ii) use these clean samples to train the diffusion model. The authors show the results for data corrupted by noise, blurring and missing pixels.

Strengths

* This is a conceptual paper that does not focus on empirically optimised architecture details, but instead shows the applicability of a simple elegant new idea. This is a plus. * I can think of many applications in e.g. in scientific imaging (e.g. microscopy or astronomy) that suffer heavily from noise and degradation and for which clean data training is usually not available or very limited. So, I believe such methods can have a substantial practical impact. * The paper is well structured, with each building block (Score based diffusion models, posterior sampling, EM algorithm) being carefully introduced.

Weaknesses

* Most importantly, I think the evaluation is not as convincing as it could be. All qualitative results seem to have artefacts to a degree. This also seems to be true for the baselines. E.g. the ambient diffusion results in Figure 3 suffer from some form of pixel artefacts. This is not the case when looking at the ambient diffusion paper (e.g. Figure 7), where not such artefacts are visible. Granted, this is a different dataset. However, I wonder if there is a problem with the training. * I think the fact that the method requires some (albeit little) clean data as a first training step is a weakness of the method. * My understanding is that the current system assumes Gaussian noise as part of the forward model (line 149). This can be limiting. * There are some typos and language problems.

Questions

At the moment the degradation model is known and seems to have to involve Gaussian noise (line 149). How does this work for the blurring operation, which seems to be deterministic? Or is the amount of added Gaussian noise just not visible. How does it work for the in-painting task, which also does not seem to involve Gaussian noise? Typos and language: Line 103-104 "we defines" Table 4: "The results show that EM-Score is insensitivity to initializations"

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Do the authors think that it could be possible to include other degradation models, such as a more realistic combination of Gaussian noise Poisson shot noise? Would it potentially be possible to learn the parameters of a degradation model. e.g. the blur kernel of the parameters of the noise distribution? A learnable degradation model would drastically increase the applicability in real life imaging applications.

Summary

This paper we proposes an expectation-maximization (EM) approach to train diffusion models from corrupted observations. The extensive experiments show the effectiveness of the proposed method.

Strengths

1. The experiments are extensive. 2. The proposed approach addresses a significant challenge in the field where large clean datasets are often unavailable.

Weaknesses

1. The source code is not provided yet, but the authors promised to make it public in the future.

Questions

While you mention that only a small number of clean images are needed to initialize the EM iterations, does the number of clean images affect the final model performance in applications? Is it possible to completely eliminate the dependence on clean data?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

The authors presented the limitations and proposed the future direction to addressed them.

Summary

Authors propose using the EM algorithm to train a diffusion model on corrupted data. To initialize the process, the proposed method requires access to a "limited number of clean samples." The authors claim the method allows convergence to the true data distribution.

Strengths

**Originality.** The idea of using EM to train a diffusion model from corrupted data is specifically original; however, training generative models without access to clean data is a well-explored problem, even before the diffusion model era. **Quality and clarity.** The paper is easy to follow, and the problem is well-motivated. **Significance.** I do not find the proposed method in its current format significant, as it assumes access to a "limited dataset of clean images" without quantifying how small this set can be. Additionally, the paper claims convergence to the true data distribution without providing conditions on the forward operator and noise process, which is an unsubstantiated claim given the limited set of inverse problems explored.

Weaknesses

* Claiming convergence to the true data distribution is a grand claim, and simply providing examples from relatively simple inverse problems does not substantiate this claim. Therefore, due to unsubstantiated claims and limited quantitative discussion on the required conditions for the forward operator and the number of clean data samples needed, I find this paper not truly usable in practice.

Questions

* How do authors quantify "minimal clean images"? The authors mention that sparse clean data provides a good initialization of the DM's manifold. Why is this the case? What does "good" mean precisely in this context?

Rating

2

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

* I disagree with the authors that needing access to a minimal number of clean images is a minor limitation. AmbientDiffusion and some other prior work discussed do not assume access to such a dataset, and hence the comparisons might not be entirely fair.

Summary

The authors introduce a new framework for training diffusion models from corrupted data. Prior work on this research topic is based on the Ambient Diffusion framework or Stein's Unbiased Risk Estimate (SURE) idea. The authors propose an alternative methodology based on the Expectation-Maximization algorithm. The algorithm is tested on standard datasets and leads to improved performance over the considered baselines.

Strengths

1. The research topic of training diffusion from corrupted data is relevant and timely. 2. The authors propose a fresh idea for this interesting research problem. The proposed idea has several elegant properties: I) it is simple, ii) relatively easy to implement, and iii) the idea is oblivious to the type of corruption. 3. The fact that the proposed algorithm is agnostic to the type of corruption is important: as the authors correctly mention, prior work, such as Ambient Diffusion, is designed (or at least tested) only for inpainting. 4. The authors manage to outperform the prior state-of-the-art, Ambient Diffusion in the same setting (CIFAR-10 with 60% corruption). 5. The authors have a nice way of leveraging clean samples. Prior work assumes that we don't have any access to clean samples. This paper naturally integrates the available clean data by using them to initialize the EM algorithm. Overall, I believe this submission offers a promising alternative to Ambient Diffusion for learning diffusion models with corrupted data.

Weaknesses

1. The comparison to the baselines is somewhat limited/unfair. As the authors mention, Ambient Diffusion has been designed for inpainting. The authors use Ambient Diffusion for denoising and Deblurring. It is not clear how they do this since the original algorithm cannot be trivially extended to these settings. Even for the case of inpainting, the authors could report results for other levels of corruption as well (apart from $p=0.6$) and for other datasets apart from CIFAR-10. Regarding the comparisons with Ambient Diffusion, the only valid data point right now is for $p=0.6$. Providing a more comprehensive analysis would help the readers understand in what regime the proposed method works the best. 2. Specifically for the case of denoising, the authors might want to provide comparisons with the more recent work Consistent Diffusion Meets Tweedie which extends the Ambient Diffusion idea to the setting of noisy samples. 3. For some scientific applications, there is no available clean data at all. The proposed algorithm needs some clean data to initialize the EM and this could be limiting for such applications. 4. One issue with the proposed method is that it relies on solving inverse problems with diffusion models. This problem is computationally intractable and for that reason, there has been a large body of research papers proposing approximations to the conditional score. The authors rely on such an approximation (DPS) which is known that doesn't truly offer samples from the posterior unless the data distribution is very simple (e.g. Gaussian). 5. One issue with this method seems to be the training speed. Even for converge to a local optimum, EM typically requires many steps. Each step here is a new training of a diffusion model. Prior works on learning diffusion models from corrupted data only required a single training. 6. For some corruption types, it is impossible to ever reconstruct the true distribution from noisy observations. The authors do not discuss how the proposed algorithm would work in such cases. 7. Minor: There is a block that needs to be deleted from the NeurIPS checklist section.

Questions

See the Weaknesses Section above. I am willing to further increase my score if my concerns are properly addressed.

Rating

6

Confidence

5

Soundness

4

Presentation

2

Contribution

3

Limitations

The authors have adequately addressed the limitations of their method.

Rebuttal acknowledgement

I would like to thank the authors for their time and their efforts during the rebuttal. I believe that the experiments with Ambient Diffusion for tasks beyond inpainting should be completely removed. Extending this method to corruptions beyond random inpainting should be done in a way such that Theorem 4.1. of the paper is satisfied. The current mechanism for further corruption is not valid and hence the results are not valid. I would like to ask the authors to remove the results for such tasks to avoid confusing the readers. I would like to thank the authors for providing additional comparisons to Ambient Diffusion. I would urge the authors to provide more points and more datasets. This is crucial: different proposed methods should be evaluated in the same setting so that future work can be compared on a common benchmark. From the current results, it looks like the gap between the proposed method and Ambient Diffusion is closing for higher corruption. Results for higher corruption should be provided as well. In my opinion, it is absolutely fine if the proposed method even underperforms a baseline in some settings. The contribution is significant already. Regarding the ICML work, it provides a general algorithm for training diffusion models from noisy observations than can be applied to training from scratch as well. In any case, this is concurrent work and I agree with the authors that direct comparisons are not needed. Regarding discussion for other corruption types: this is not what I meant. I was referring to cases where it is information-theoretically impossible to learn the underlying distribution from measurements, e.g. you can't learn the distribution of human faces from a set of training images where the eyes are always missing. Overall, I believe this paper has some really nice aspects and hence I will keep my positive score of 6.