We present the first diffusion-based framework that can learn an unknown distribution using only highly-corrupted samples. This problem arises in scientific applications where access to uncorrupted samples is impossible or expensive to acquire. Another benefit of our approach is the ability to train generative models that are less likely to memorize individual training samples since they never observe clean training data. Our main idea is to introduce additional measurement distortion during the diffusion process and require the model to predict the original corrupted image from the further corrupted image. We prove that our method leads to models that learn the conditional expectation of the full uncorrupted image given this additional measurement corruption. This holds for any corruption process that satisfies some technical conditions (and in particular includes inpainting and compressed sensing). We train models on standard benchmarks (CelebA, CIFAR-10 and AFHQ) and show that we can learn the distribution even when all the training samples have $90\%$ of their pixels missing. We also show that we can finetune foundation models on small corrupted datasets (e.g. MRI scans with block corruptions) and learn the clean distribution without memorizing the training set.
Paper
Similar papers
Peer review
Summary
This paper proposes to train diffusion models that can recover corrupted data without training on clean data. The key idea is, given a corruption matrix $A$ one can further sample a corruption matrix $\tilde{A}$ given $A$ and the model learns to predict all the existing pixels. It is empirically shown that this trick ensures robustness against higher corruption levels and can restore data with better performance than those that's trained on clean data.
Strengths
- The authors propose the first diffusion-based method that can restore data from corruption without training from clean data. - The method can reuse regular diffusion sampler without much modification. - It is guaranteed theoretically for the model to recover clean data with some rank assumptions. - The method alleviates the memorization issues of diffusion models by considering corruption schemes.
Weaknesses
- More exposition is needed on some details. What kind of $\tilde{A}$ is used during inference? Is it necessary to sample using the same $\delta$ during training? Do we fix a $\tilde{A}$ for all sampling steps or resample each step? Does taking the expectation over multiple samples of such $\tilde{A}$ work better? - The paper lacks implementation details. E.g. $h_\theta(\tilde{A}, \tilde{A}x_t, t)$ explicitly depends on $\tilde{A}$. How is this dependency implemented in practice? - Line 248 the authors forgot line break. - The authors only considered random masking corruption. However, there are many other types of linear corruption schemes. Does the method stay robust against other corruption such as Gaussian blur, etc.? - It would be great to analyze the effect of different sampling scheme for $\tilde{A}$. How does FID change w.r.t $\tilde{A}$ with increasing levels of further corruption? Is there a sweet spot for further corruption? - Why is there no comparison with AmbientGAN for Table 1, as it is an important baseline?
Questions
My questions and concerns are listed in the section above.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
- The authors have adequately discussed limitations of the model. Some further discussion on societal impact is encouraged as it is closely related to privacy issues of training data.
Summary
In summary, the authors propose a diffusion-based framework that can learn unknown distributions from highly-corrupted samples, allowing the training of generative models without relying on clean training data. Their approach introduces additional measurement distortion and successfully predicts original corrupted images. The method is applicable to various corruption processes and achieves promising results on benchmark datasets.
Strengths
- The problem formulation by itself is interesting, and the proposed method is novel. - The paper proposes a new and interesting domain of learning the image data distribution w/o access to the ground truth data samples. - The paper is well written and easy to follow. - The proposed training and sampling procedure is scalable and easy to incorporate into the current diffusion model framework. - The authors attach the theory for the effectiveness of the proposed method. - The proposed method only needs 1 NFE to produce comparable results. - The method paves a way to alleviate the memorize issue of diffusion model.
Weaknesses
- From my perspective, I do not see significant weaknesses in this paper. - One potential issue is the significant drop of FID when the model is trained with images corrupted with a large ratio of pixels. However, I think this should not be criticized.
Questions
- Can you provide the performance of the proposed method in Table 1 with more timesteps of sampling? Does the performance improve? If no, then can you please explain why?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The author lists the limitation of the proposed method in the paper which is reasonable. I appreciate the authors' adequate limitation summary and it sheds light on future improvements.
Summary
This paper describes a method to learn a denoising diffusion model only with corrupted data. This is an important problem in many areas of applied science where there is no access to ground truth. Another important potential benefit of this method is to overcome memorization of the training images. The main idea in this method is described as introducing additional corruption and training on this doubly corrupt data. The authors describe a derivation to estimate conditional expectation of the uncorrupted image only using the corrupted data. They show (Table 1) that their method outperforms other methods in solving random inpainting problem. Also, the method is used to fine tune Deepfloyd IF on smaller samples. This result is used to show that the method overcomes memorization.
Strengths
This work is motivated in two ways: learning the score of distribution of clean data with access to corrupted data, and avoiding memorization. These are both very important practical and theoretical topics which this work tries to address. The trick which is used to train the models (adding more corruption) is clever and the network achieves to learn to inpaint and denoise images (figure 7). As mentioned under weaknesses, I think the first goal is not fully achieved and the theoretical results seem incorrect and practical results seem to be constrained to specific types of corruption. However, the memorization results seem very impressive, although the memorization analysis is carries only for the fine tuned model.
Weaknesses
There are a number of fundamental technical and conceptual flaws which need to be considered and fixed: 1) The training setup was not clear in the text, but this is what I gathered: the target image during training is a corrupted image $Ax_0$, and the input is a noisy and more corrupted image $BAx_0 + \sigma \eta $. All throughout the corruptions are assumed to be random or block missing pixels. The network learns to remove the noise and inpaint pixels that are removed by B. At the test time, the network removes noise and also inpaint *all* pixels since it doesn't know which pixels are dropped due to A and which pixels are dropped due to B. If this is the training setup, please clarify in the text. If not, please describe what was the training setup. 2) Equation (3), the objective after additional corruption, is the distance between the doubly corrupted image and the $\textbf{clean}, x_0,$ image. The entire premise of the work is that clean image is not available so why is it assumed to be available during training? The objective should be the distance from $Ax_0$ instead. It needs to be clarified whether this is a typo or the authors actually used clean data, $x_0$, during training. 2) If that is a typo and the target image during training is $Ax_0$ then after training is completed the network works as a denoiser plus inpainter. As a result, the Tweedie equation is not a good description of the output of this network anymore. That is, the output is not $E[x_0|x_t]$, where $x_t = x_0 + \sigma_t \eta$. So, this output cannot be used directly as an estimate of score in the diffusion model . 3) To remedy the above mentioned problem, the authors propose eq 3.3 in which they claim to approximate the score, $E[x_0|x_t],$ with $E[x_0|\tilde{A}x_t, \tilde{A}]$. It is not at all clear what is the justification behind this approximation. Again, this expectation is a very different entity from the actual score, and it is not clear why direct use of it makes sense to estimate the score. This expectation is the solution for inverse problem given the forward measurement $A$, as apposed to the solution for mere denoising problem (i.e. the score). 4) Additionally, two more terms are added to the update line in eq 3.4. The justification for this is described from line 166 to 175. It starts with $\gamma_t$ going to zero when $t$ approaches zero. It is not clear as to why $\gamma_t$ goes to zero when $t$ goes to zero. Please clarify this assumption. The reasoning that follows this assumption is also not clear and sounds ad hoc. Did you add these terms because the update line of eq 3.3 did not work in practice? What is the intuition or theory behind this choice? Please clarify both eq 3.3. (why did you estimate one expectation with another?) and eq 3.4 (why did you add two terms and why do they make sense?). 5) The theory section needs at least a re-write because it is not clear what is the goal of this section. The section starts with the goal of proving that the optimal estimate of clean image given the corrupted image is equal to $E[x_0| Ax_t = y, A]$. It is not clear why the authors need to prove this, since this is a basic fact from Bayesian machine learning: the optimal estimation of corrupted data is the conditional mean of the predictive distribution (refer to textbooks like Bishop). The problem is that this expectation is not equal to the score, so theoretically it can't be used to estimate the score. Of course what is learned by this network is approximately the optimal reconstruction of $A(x+\sigma \eta)$, but as long as it is not the score, it is not theoretically valid to use it iteratively in a diffusion model. In a nutshell, the learned network is a denoiser/inpainter not a score estimator (i.e. pure denoiser). 6) On top of the above mentioned issues, this network learns to reconstruct + denoise images only if $E_{A|\tilde{A}}[A^TA]$ is full rank. This strong assumption requires a high level of randomness in the corruption which is not very common in many real-world applications. For example, if the ground truth images miss some information systematically (let's say they are all blurred) the network will not be able to reconstruct.
Questions
In addition to theoretical questions and comments I made under weaknesses: Why do you think you achieve optimal sampling performance in one iteration? How does the training time compares to uncorrupted models? Do you need to train on more samples? ________________________________________________________ Note: I am willing to raise the score if the questions asked under Weaknesses and here are addressed.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
2 fair
Presentation
2 fair
Contribution
3 good
Limitations
The main limitation of the work (aside from some jumps and ambiguity in the theoretical results) is the type of corruption this method can work with. The corruption must be linear and in addition $E_{A|\tilde{A}}[A^TA]$ must be full rank. This rules out many real world applications.
Summary
The paper focuses on learning clean distributions from corrupted data. In the training diffusion model, the training dataset contains only highly-corrupted examples. They propose a training algorithm of restoration model by introducing additional measurement distortion.They also provide sampling methods and theoretical analysis. Experimental results show that their superior performance.
Strengths
* The problem of handling corrupted datasets is significant even in generative model learning. This paper could be seen as one that explores this directions. * The paper is well-structured and easy to follow. * The fact that a model which is trained by corrupted data does not memorize the training dataset, as supported by Figure 1 and Figure 4, is very important. This might be gives significant implications for various applications.
Weaknesses
* Sampling Process * It would be helpful to provide an algorithm or a detailed explanation of the sampling process from scratch. * Equations 3.3 and 3.4 describe how the sample at time $t$ is generated from sample at $t-\Delta t$. It is necessary to explain how the evaluation of single networks allows the generation of images in the experiments. * Related work * It would be beneficial to include a discussion of previous research on dealing with incomplete datasets in generative models, such as [1, 2, 3, 4]. [1] Li, S. C. X., Jiang, B., & Marlin, B. (2018, September). MisGAN: Learning from Incomplete Data with Generative Adversarial Networks. In International Conference on Learning Representations. [2] Mattei, P. A., & Frellsen, J. (2019, May). MIWAE: Deep generative modelling and imputation of incomplete data sets. In International Conference on Machine Learning (pp. 4413-4423). PMLR. [3] Ipsen, N. B., Mattei, P. A., & Frellsen, J. (2020, September). not-MIWAE: Deep Generative Modelling with Missing not at Random Data. In International Conference on Learning Representations. [4] Richardson, T. W., Wu, W., Lin, L., Xu, B., & Bernal, E. A. (2020). Mcflow: Monte carlo flow models for data imputation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 14205-14214). * Experiments * It would be beneficial to provide the FID results for Figure 5 to demonstrate the differences. * In addition to Figure 5, it would be valuable to compare the proposed model with existing generative models such as AmbientGAN in various experiments. * Presentation * It would be helpful to have better spacing between subfigures in Figure 2 to improve the clarity of the captions. * Proper citations are needed in the background section. * Figure 3 would be better placed within a paragraph rather than between paragraphs.
Questions
Please see Weaknesses part.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
They provided in the last paragraph in the main paper.
Response to Authors
Thank you for your answer. I will keep the current score.
Thank you for responding to my comments and questions. The comment about eq 3.2 was indeed due to my misreading of it. Thanks for clarification. In my initial review, as I had mentioned, I had assumed the apparent missing A was a typo and not a conceptual error. So the rest of my initial review was based on the assumption that you had the correct form of the eq 3.2 (as you did). Regarding the theory section, I am convinced now that Theorem 4.1 is valuable beyond what I was evaluated initially. However, it relies on the strong assumption on full rank $E[A^TA]$, so a discussion on this assumption should be included in the text. This assumption excludes many real work corruptions. It should be mentioned earlier in the text (in the intro or even abstract). At its current state, the text makes the impression that this method magically learns the score from any corrupted data, until section 4. It can be mentioned earlier that this method works as long as the corruption is not systematic. For more discussion, refer to literature on Stein Risk Estimator (SURE) in which the minimizer is estimated without access to clean data for the case of Gaussian corruption. Your proof can be thought of as a generalization of this old idea. It is very interesting that the corruption does not require training on more data points. I think it is worth to include this in the main text. A minor note: do not assume that all readers are familiar with the notations you use. Explain in the text what $\gamma$ is. The paper should be self-sufficient in terms of notations and definitions. Regarding the final update line, eq 3.3 and eq 3.4, I still find the discussion unsatisfactory. This is at the core of the method, and I am not convinced why this crude approximation (in 3.3) is acceptable. Of course if the corruption is small, the approximation would be not too far, but that is not what this paper claims to do. Additionally the term added in 3.4 is not well motivated. Again, if there is a good motivation from the paper you're borrowing this from, it should be described and explained in the text. Why is that term a good addition to your update line? At it's current state, there is no clear connection between the theory section and this added term. I raise my score from 4 to 5 since the authors addressed some of my questions. There is still room for improvement regarding the above point (point 3 and 4 in my initial review).
Additional Discussion
We thank the Reviewer for engaging in the discussion and for raising their score. Their valuable time is deeply appreciated! We are glad that the issue that arose from the misreading has now been resolved, that the Reviewer appreciated our Theorem 4.1 and that the Reviewer found satisfying that we didn’t use more data points for the corrupted models. We will definitely include this in the main text, as recommended by the Reviewer. We will also make it more clear from the Introduction that our method does not work for arbitrary corruptions. In fact, there is no method that learns the true distribution for arbitrary corruptions. The assumption we make on the corruption process is necessary to learn from inpainted data – if we never observe certain pixel locations there is no way to identify the distribution in these locations unless we make assumptions on the distribution itself. As we pointed out in our rebuttal, previous methods such as AmbientGAN, make such assumptions (e.g. from AmbientGAN: “A critical assumption for our framework and theory to work is that the measurement process is known and satisfies certain technical conditions”). These assumptions are not very restrictive when we control the corruption process (e.g. to reduce memorization) but preclude (theoretically) applying our method for non-systematic corruption processes, as the Reviewer noted correctly. We will clarify this very early in the text, as recommended by the Reviewer. We also want to point out that our method could potentially achieve reasonable performance in practice for corruption processes that violate such assumptions, but in this case it does not come with any theoretical guarantees. We will also further highlight how our proof adds to the literature of Stein Risk Estimation (SURE) and generalizes this to the case of non Gaussian corruption. We thank the Reviewer for recommending this. Finally, we will add the explanation from our rebuttal to the main text regarding what $\gamma_t$ is to avoid potential confusions. Regarding the remaining concern of the Reviewer about the sampling, we want to start by acknowledging again that we are making an approximation there. We clearly state this in the Limitations Section of our work (“Further, in this work we only experimented with very simple approximation algorithms to estimate $\mathbb E[x_0|x_t]$ using our trained models”) and in the Method Section (“we approximate $\mathbb E[x_0|x_t]$ given the predictions of $\mathbb E[x_0|Ax_t, A]$”). Surprisingly, this simple approximation outperforms previous baselines and works reasonably well even for high corruptions (potentially because there a lot of redundancies in natural images). We are planning to explore in future work how to make this step exact and we might get a significant performance boost (especially in the high corruption regime) by fixing this. Regarding the term in the Equation 3.4, this is inspired by the video diffusion models literature. For video generation with diffusion models, the models are trained to denoise the current frame but are regularized so that denoised version is not too different compared to the prediction for neighboring frames. This heuristic is needed for smooth transitions in the video generation process. Similarly, in our setup, we only have access to models that predict given limited context and we want to combine the predictions such that there are no artifacts or inconsistencies. The fixed mask sampler ensures there are no artifacts by always masking the same pixels. In Eq. 3.4, we added this extra term to (heuristically) remedy that the values in the masked pixels in the fixed mask sampler only depended on the final prediction – we update all pixels at every prediction using many different masks and ensure consistency of the different predictions through this added term. Again, this is a heuristic step and future work might significantly improve performance by making this exact, thus we understand if the Reviewer finds this approximation unsatisfactory. We did not end up using this sampler for the experiments in the main paper due to the extra computation needed in return for marginal benefits, e.g. see Appendix E.3 and Table 5. Since the Reviewer did not appreciate this attempt to heuristically improve the sampling process, we can move this entirely to the Appendix (together with the discussion above) and make space for the rest of the requested changes that are more critical to the paper. We hope this discussion is useful and we want to thank again the Reviewer for helping us improve our paper!
Thanks for clarifying the point about the added term in eq 3.4. I believe it will significantly improve the clarity of the paper if, as you suggested, this is moved to the appendix. Or at least a similar description of the heuristic is offered in the text. In its current state, it seems very arbitrary. Although the approximation of $E[x_0|x_t]$ by $E[x_0|Ax_t,A]$ is unjustified, I think this paper overall contains significant contributions to the field. Thus I raise my score to 6.
As the discussion period approaches its end, we would like to kindly ask if the Reviewer had a chance to read our rebuttal. We thank again the Reviewer for their time and their constructive feedback in their initial review and we hope that our rebuttal addressed the Reviewer's questions.
Thank you for the rebuttal
The authors have resolved my concerns. I have raised my score.
Decision
Accept (poster)