Title: Ambient Diffusion: Learning Clean Distributions from Corrupted Data
Venue: NeurIPS 2023
Accepted: yes
Keywords: corrupted data

generative models

ambient gan

inverse problems

learning from measurements
Abstract: 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 any individual training sample, 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.

Ambient Diffusion:

Abstract

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 any individual training sample, 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.

1 Introduction

Diffusion generative models [54, 23, 57] are emerging as versatile and powerful frameworks for learning high-dimensional distributions and solving inverse problems [34, 11, 36, 28]. Numerous recent developments [58, 30] have led to text conditional foundation models like Dalle-2 [45], Latent Diffusion [50] and Imagen [52] with an incredible performance in general image domains. Training these models requires access to high-quality datasets which may be expensive or impossible to obtain. For example, direct images of black holes cannot be observed [12, 19] and high quality MRI images require long scanning times, causing patient discomfort and motion artifacts [28].

Recently, Carlini et al. [8], Somepalli et al. [55], and Jagielski et al. [27] showed that diffusion models can memorize examples from their training set. Further, an adversary can extract dataset samples given only query access to the model, leading to privacy, security and copyright concerns. For many applications, we may want to learn the distribution but not individual training images e.g. we might want to learn the distribution of X-ray scans but not memorize images of specific patient scans from the dataset. Hence, we may want to introduce corruption as a design choice. We show that it is possible to train diffusions that learn a distribution of clean data by only observing highly corrupted samples.

Training

lmages

Images

Generated samples

Generated samples

Nearest neighbors

Prior work in supervised learning from corrupted data. The traditional approach to solving such problems involves training a restoration model using supervised learning to predict the clean image based on the measurements [47, 48, 63, 41]. The seminal Noise2Noise [39] work introduced a practical algorithm for learning how to denoise in the absence of any non-noisy images. This framework and its generalizations [5, 38, 59] have found applications in electron microscopy [16], tomographic image reconstruction [62], fluorescence image reconstruction [65], blind inverse problems [20, 5], monocular depth estimation and proteomics [6]. Another related line of work uses Stein's Unbiased Risk Estimate (SURE) to optimize an unbiased estimator of the denoising objective without access to non-noisy data [18]. We stress that the aforementioned research works study the problem of restoration , whereas are interested in the problem of sampling from the clean distribution. Restoration algorithms based on supervised learning are only effective when the corruption level is relatively low [15]. However, it might be either not possible or not desirable to reconstruct individual samples. Instead, the desired goal may be to learn to generate fresh and completely unseen samples from the distribution of the uncorrupted data but without reconstructing individual training samples .

Indeed, for certain corruption processes, it is theoretically possible to perfectly learn a distribution only from highly corrupted samples (such as just random one-dimensional projections), even though individual sample denoising is usually impossible in such settings.

AmbientGAN [7] showed that general d dimensional distributions can be learned from scalar observations, by observing only projections on one-dimensional random Gaussian vectors, in the infinite training data limit. The theory requires an infinitely powerful discriminator and hence does not apply to diffusion models. MCFlow [49] is a framework, based on a variant of the EM algorithm, that can be used to train normalizing flow models from missing data. Finally, MIWAE [42] and Not-MIWAE [26] are frameworks to learn deep latent models (e.g. VAEs) from missing data when the corruption process is known or unknown respectively.

Our contributions. We present the first diffusion-based framework to learn an unknown distribution D when the training set only contains highly-corrupted examples drawn from D . Specifically, we consider the problem of learning to sample from the target distribution p 0 ( x 0 ) given corrupted

Random

in-

painting

Further corrup-

tion

samples A x 0 where A ∼ p ( A ) is a random corruption matrix (with known realizations and prior distribution) and x 0 ∼ p 0 ( x 0 ) . 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 provide an algorithm that provably learns E [ x 0 | ˜ A ( x 0 + σ t η ) , ˜ A ] , for all noise levels t and for ˜ A ∼ p ( ˜ A | A ) being a further corrupted version of A . The result holds for a general family of corruption processes A ∼ p ( A ) . For various corruption processes, we show that the further degradation introduced by ˜ A can be very small.

We use our algorithm to train diffusion models on standard benchmarks (CelebA, CIFAR-10 and AFHQ) with training data at different levels of corruption.

Given the learned conditional expectations we provide an approximate sampler for the target distribution p 0 ( x 0 ) .

We show that for up to 90% missing pixels, we can learn reasonably well the distribution of uncorrupted images. We outperform the previous state-of-the-art AmbientGAN [7] and natural baselines.

We show that our models perform on par or even outperform state-of-the-art diffusion models for solving certain inverse problems even without ever seeing a clean image during training. Our models do so with a single prediction step while our baselines require hundreds of diffusion steps.

We use our algorithm to finetune foundational pretrained diffusion models. Our finetuning can be done in a few hours on a single GPU and we can use it to learn distributions with a few corrupted samples.

We show that models trained on sufficiently corrupted data do not memorize their training set. We measure the tradeoff between the amount of corruption (that controls the degree of memorization), the amount of training data and the quality of the learned generator.

We open-source our code and models: https://github.com/giannisdaras/ambient-diffusion.

2 Background

Training a diffusion model involves two steps. First, we design a corruption process that transforms the data distribution gradually into a distribution that we can sample from [58, 14]. Typically, this corruption process is described by an Ito SDE of the form: d x = f ( x , t )d t + g ( t )d w , where w is the standard Wiener process. Such corruption processes are reversible and the reverse process is also described by an Ito SDE [3]: d x = ( f ( x , t ) -g 2 ( t ) ∇ x log p t ( x ) ) d t + g ( t )d w . The designer of the diffusion model is usually free to choose the drift function f ( · , · ) and the diffusion function g ( · ) . Typical choices are setting f ( x , t ) = 0 , g ( t ) = √ d σ 2 t d t (Variance Exploding SDE) or setting f ( x, t ) = -β ( t ) x , g ( t ) = √ β ( t ) (Variance Preserving SDE). Both of these choices lead to a Gaussian terminal distribution and are equivalent to a linear transformation in the input. The goal of diffusion model training is to learn the function ∇ x log p t ( x ) , which is known as the score function.

To simplify the presentation of the paper, we will focus on the Variance Exploding SDE that leads to conditional distributions x t = x 0 + σ t η .

Vincent [61] showed that we can learn the score function at level t by optimizing for the scorematching objective:

Specifically, the score function can be written in terms of the minimizer of this objective as:

This result reveals a fundamental connection between the score-function and the best restoration model of x 0 given x t , known as Tweedie's Formula [17]. Specifically, the optimal h θ ∗ ( x t , t ) is given by E [ x 0 | x t ] , which means that

Inspired by this restoration interpretation of diffusion models, the Soft/Cold Diffusion works [14, 4] generalized diffusion models to look at non-Markovian corruption processes: x t = C t x 0 + σ t η . Specifically, Soft Diffusion proposes the Soft Score Matching objective:

and shows that it is sufficient to recover the score function via a generalized Tweedie's Formula:

For these generalized models, the matrix C t is a design choice (similar to how we could choose the functions f , g ). Most importantly, for t = 0 , the matrix C t becomes the identity matrix and the noise σ t becomes zero, i.e. we observe samples from the true distribution.

3 Method

As explained in the introduction, in many cases we do not observe uncorrupted images x 0 , either by design (to avoid memorization and leaking of sensitive data) or because it is impossible to obtain clean data. Here we study the case where a learner only has access to linear measurements of the clean data, i.e. y 0 = A x 0 , and the corruption matrices A : R m × n . We note that we are interested in non-invertible corruption matrices. We ask two questions:

Is it possible to learn E [ x 0 | A ( x 0 + σ t η ) , A ] for all noise levels t , given only access to corrupted samples ( y 0 = A x 0 , A ) ?

If so, is it possible to use this restoration model E [ x 0 | A ( x 0 + σ t η ) , A ] to recover E [ x 0 | x t ] for any noise level t , and thus sample from the true distribution through the score function as given by Tweedie's formula (Eq. 2.3)?

We investigate these questions in the rest of the paper. For the first, the answer is affirmative but only after introducing additional corruptions, as we explain below. For the second, at every time step t , we approximate E [ x 0 | x t ] directly using E [ x 0 | A x t , A ] (for a chosen A ) and substitute it into Eq. 2.3. Empirically, we observe that the resulting approximate sampler yields good results.

3.1 Training

For the sake of clarity, we first consider the case of random inpainting. If the image x 0 is viewed as a vector, we can think of the matrix A as a diagonal matrix with ones in the entries that correspond to the preserved pixels and zeros in the erased pixels. We assume that p ( A ) samples a matrix where each entry in the diagonal is sampled i.i.d. with a probability 1 -p to be 1 and p to be zero.

We would like to train a function h θ which receives a corruption matrix A and a noisy version of a corrupted image, y t = A ( x 0 + σ t η ) ︸ ︷︷ ︸ where η ∼ N (0 , I ) , and produces an estimate for the

x

t

conditional expectation. The simplest idea would be to simply ignore the missing pixels and optimize for:

Despite the similarities with Soft Score Matching (Eq 2.4), this objective will not learn the conditional expectation. The reason is that the learner is never penalized for performing arbitrarily poorly in the missing pixels. Formally, any function h θ ′ satisfying A h θ ′ ( A, y t , t ) = A E [ x 0 | A x t , A ] is a minimizer.

Instead, we propose to further corrupt the samples before feeding them to the model, and ask the model to predict the original corrupted sample from the further corrupted image.

Concretely, we randomly corrupt A to obtain ˜ A = BA for some matrix B that is selected randomly given A . In our example of missing pixels, ˜ A is obtained from A by randomly erasing an additional fraction δ of the pixels that survive after the corruption A . Here, B will be diagonal where each element is 1 with probability 1 -δ and 0 w.p. δ . We will penalize the model on recovering all the pixels that are visible in the sample A x 0 : this includes both the pixels that survive in ˜ A x 0 and those that are erased by ˜ A . The formal training objective is given by minimizing the following loss:

The key idea behind our algorithm is as follows: the learner does not know if a missing pixel is missing because we never had it (and hence do not know the ground truth) or because it was deliberately erased as part of the further corruption (in which case we do know the ground truth). Thus, the best learner cannot be inaccurate in the unobserved pixels because with non-zero probability it might be evaluated on some of them. Notice that the trained model behaves as a denoiser in the observed pixels and as an inpainter in the missing pixels. We also want to emphasize that the probability δ of further corruption can be arbitrarily small as long as it stays positive.

The idea of further corruption can be generalized from the case of random inpainting to a much broader family of corruption processes. For example, if A is a random Gaussian matrix with m rows, we can form ˜ A by deleting one row from A at random. If A is a block inpainting matrix (i.e. a random block of fixed size is missing from all of the training images), we can create ˜ A by corrupting further with one more non-overlapping missing block. Examples of our further corruption are shown in Figure 2. In our Theory Section, we prove conditions under which it is possible to recover E [ x 0 | ˜ A x t , ˜ A ] using our algorithm and samples ( y 0 = A x 0 , A ) . Our goal is to satisfy this condition while adding minimal further corruption, i.e. while keeping ˜ A close to A .

3.2 Sampling

To sample from p 0 ( x 0 ) using the standard diffusion formulation, we need access to ∇ x t log p t ( x t ) , which is equivalent to having access to E [ x 0 | x t ] (see Eq. 2.3). Instead, our model is trained to predict E [ x 0 | ˜ A x t , ˜ A ] for all matrices A in the support of p ( A ) .

We note that for random inpainting, the identity matrix is technically in the support of p ( A ) . However, if the corruption probability p is at least a constant, the probability of seeing the identity matrix is exponentially small in the dimension of x t . Hence, we should not expect our model to give good estimates of E [ x 0 | ˜ A x t , ˜ A ] for corruption matrices A that belong to the tails of the distribution p ( A ) .

The simplest idea is to sample a mask ˜ A ∼ p ( ˜ A ) and approximate E [ x 0 | x t ] with E [ x 0 | ˜ A x t , ˜ A ] . Under this approximation, the discretized sampling rule becomes:

This idea works surprisingly well. Unless mentioned otherwise, we use it for all the experiments in the main paper and we show that we can generate samples that are reasonably close to the true distribution (as shown by metrics such as FID and Inception) even with 90% of the pixels missing.

We include pseudocode for our sampling in Algorithm 1. In the Appendix (Section E.4), we ablate an alternative choice for sampling that can lead to slight improvements at the cost of increased function evaluations.

Algorithm 1 Ambient Diffusion Fixed Mask Sampling (for VE SDE)

We underline that the Fixed Mask Sampler does not come with a theoretical guarantee of sampling from p 0 ( x 0 ) . In the Appendix, we prove that whenever it is possible to reconstruct p 0 ( x 0 ) from corrupted samples, it is also possible to reconstruct it using access to E [ x 0 | A x t , A ] (Lemma A.3). However, as stated in the Limitations section, we were not able to find any practical algorithm to do so. Nevertheless, experimentally the Fixed Mask Sampler has strong performance and beats previously proposed baselines for different classes of generative models.

4 Theory

As elaborated in Section 3, one of our key goals is to learn the best restoration model for the measurements at all noise levels, i.e., the function h ( A, y t , t ) = E [ x 0 | y t , A ] . We now show that under a certain assumption on the distribution of A and ˜ A , the true population minimizer of Eq. 3.2 is indeed essentially of the form above. This assumption formalizes the notion that even conditional on ˜ A , A has considerable variability, and the latter ensures that the best way to predict A x 0 as a function of ˜ A x t and ˜ A is to optimally predict x 0 itself. All proofs are deferred to the Appendix.

Theorem 4.1. Assume a joint distribution of corruption matrices A and further corruption ˜ A . If for all ˜ A in the support it holds that E A | ˜ A [ A T A ] is full-rank, then the unique minimizer of the objective in equation 3.2 is given by

Two simple examples that fit into this framework (see Corollaries A.1 and A.2 in the Appendix) are:

Inpainting: A ∈ R n × n is a diagonal matrix where each entry A ii ∼ Ber(1 -p ) for some p > 0 (independently for each i ), and the additional noise is generated by drawing ˜ A | A such that ˜ A ii = A ii · Ber(1 -δ ) for some small δ > 0 (again independently for each i ). 1

Gaussian measurements: A ∈ R m × n consists of m rows drawn independently from N (0 , I n ) , and ˜ A ∈ R m × n is constructed conditional on A by zeroing out its last row.

Notice that the minimizer in Eq 4.1 is not entirely of the form we originally desired, which was h ( A, y t , t ) = E [ x 0 | A x t , A ] . In place of A , we now have ˜ A , which is a further degraded matrix.

Indeed, one trivial way to satisfy the condition in Theorem 4.1 is by forming ˜ A completely independently of A , e.g. by always setting ˜ A = 0 . However, in this case, the function we learn is not very useful. For this reason, we would like to add as little further noise as possible and ensure that ˜ A is close to A . In natural noise models such as the inpainting noise model, by letting the additional corruption probability δ approach 0 , we can indeed ensure that ˜ A follows a distribution very close to that of A .

5 Experimental Evaluation

5.1 Training from scratch on corrupted data

Our first experiment is to train diffusion models from scratch using corrupted training data at different levels of corruption. The corruption model we use for these experiments is random inpainting: we form our dataset by deleting each pixel with probability p . To create the matrix ˜ A , we further delete each row of A with probability δ - this removes an additional δ -fraction of the surviving pixels. Unless mentioned otherwise, we use δ = 0 . 1 . We train models on CIFAR-10, AFHQ, and CelebAHQ. All our models are trained with corruption level p ∈ { 0 . 0 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 0 . 9 } . We use the EDM [30] codebase to train our models. We replace convolutions with Gated Convolutions [64] which are known to perform better for inpainting-type problems. To use the mask ˜ A as an additional input to the model, we simply concatenate it with the image x . The full training details can be found in the Appendix, Section C.

We first evaluate the restoration performance of our model for the task it was trained on (random inpainting and noise). We compare with state-of-the-art diffusion models that were trained on clean data. Specifically, for AFHQ we compare with the state-of-the-art EDM model [30] and for CelebA we compare with DDIM [56]. These models were not trained to denoise, but we can use the prior learned in the denoiser as in [60, 29] to solve any inverse problem. We experiment with the state-of-the-art reconstruction algorithms: DDRM [34] and DPS [11].

We summarize the results in Table 1. Our model performs similarly to other diffusion models, even though it has never been trained on clean data. Further, it does so by requiring only one step, while all the baseline diffusion models require hundreds of steps to solve the same task with inferior or comparable performance. The performance of DDRM improves with more function evaluations at the cost of more computation. For DPS, we did not observe significant improvement by increasing the number of steps to more than 100 . We include results with noisy inpainted measurements and comparisons with a supervised method in the Appendix, Section E, Tables 3, 5. We want to emphasize that all the baselines we compare against have an advantage: they are trained on uncorrupted data. Instead, our models were only trained on corrupted data. This experiment indicates that: i) our training algorithm for learning the conditional expectation worked and ii) that the choice of corruption that diffusion models are trained to reverse matters for solving inverse problems.

Next, we evaluate the performance of our diffusion models as generative models. To the best of our knowledge, the only generative baseline with quantitative results for training on corrupted data is AmbientGAN [7] which is trained on CIFAR-10. We further compare with a diffusion model trained without our further corruption algorithm. We plot the results in Figure 3. The diffusion model trained without our further corruption algorithm performs well for low corruption levels but collapses entirely for high corruption. Instead, our model trained with further corruption maintains reasonable corruption scores even for high corruption levels, outperforming the previous state-of-theart AmbientGAN for all ranges of corruption levels.

For CelebA-HQ and AFHQ we could not find any generative baselines trained on corrupted data to compare against. Nevertheless, we report FID and Inception Scores and summarize our results in Table 4 to encourage further research in this area. As shown in the Table, 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.

5.2 Finetuning foundation models on corrupted data

We can apply our technique to finetune a foundational diffusion model. For all our experiments, we use Deepfloyd's IF model [2], which is one of the most powerful open-source diffusion generative

models available. We choose this model over Stable Diffusion [51] because it works in the pixel space (and hence our algorithm directly applies).

Memorization. We show that we can finetune a foundational model on a limited dataset without memorizing the training examples. This experiment is motivated by the recent works of Carlini et al. [8], Somepalli et al. [55], and Jagielski et al. [27] that show that diffusion generative models memorize training samples and they do it significantly more than previous generative models, such as GANs, especially when the training dataset is small. Specifically, Somepalli et al. [55] train diffusion models on subsets of size { 300 , 3000 , 30000 } of CelebA and they show that models trained on 300 or 3000 memorize and blatantly copy images from their training set.

We replicate this training experiment by finetuning the IF model on a subset of CelebA with 3000 training examples. Results are shown in Figure 1. Standard finetuning of Deepfloyd's IF on 3000 images memorizes samples and produces almost exact copies of the training set. Instead, if we corrupt the images by deleting 80% of the pixels prior to training and finetune, the memorization decreases sharply and there are distinct differences between the generated images and their nearest neighbors from the dataset. This is in spite of finetuning until convergence.

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Corruption Probability

4

5

6

7

8

9

10

Inception Score

CIFAR-10 Inception Scores

AmbientGAN

Diffusion No Further Corruption

Ambient Diffusion

To quantify the memorization, we follow the methodology of Somepalli et al. [55]. Specifically, we generate 10000 images from each model and we use DINO [9]-v2 [46] to compute top1 similarity to the training images. Results are shown in Figure 6. Similarity values above 0 . 95 roughly correspond to the same person while similarities below 0 . 75 typically correspond to random faces. The standard finetuning (Red) often generates images that are near-identical with the training set. Instead, finetuning with corrupted samples (blue) shows a clear shift to the left. Visually we never observed a near-copy generated from our process - see also Figure 1.

We repeat this experiment for models trained on the full CelebA dataset and at different levels of corruption. We include the results in Figure 8 of the Appendix. As shown, the more we increase the corruption level the more the distribution of similarities shifts to the left, indicating less memorization. However, this comes at the cost of decreased performance, as reported in Table 4.

New domains and different corruption. We show that we can also finetune a pre-trained foundation model on a new domain given a limited-sized dataset in a few hours in a single GPU. Figure 5 shows generated samples from a finetuned model on a dataset containing 155 examples of brain tumor MRI images [25]. As shown, the model learns the statistics of full brain tumor MRI images while only trained on brain-tumor images that have a random box obfuscating 25% of the image. The training set was resized to 64 × 64 but the generated images are at 256 × 256 by simply leveraging the power of the cascaded Deepfloyd IF.

0.60

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Similarity Value

0

100

200

300

400

500

600

Number of generated samples

Distribution of Similarity Values to the nearest training sample

Corruption probability = 0.8

Corruption probability = 0.0

Limitations. Our work has several limitations. First, there is a tradeoff between generator quality and corruption levels. For higher corruption, it is less likely that our generator memorizes parts of training examples, but at a cost of degrading quality. Precisely characterizing this trade-off is an open research problem. Further, in this work, we only experimented with very simple approximation algorithms to estimate E [ x 0 | x t ] using our trained models. Additionally, we cannot make any strict privacy claim about the protection of any training sample without making assumptions about the data distribution. We show in the Appendix that it is possible to recover E [ x 0 | x t ] exactly using our restoration oracle, but we do not have an algorithm to do so. Finally, our method cannot handle measurements that also have noise. Future work could potentially address this limitation by exploiting SURE regularization as in [1].

Acknowledgments and Disclosure of Funding

The authors would like to thank Tom Goldstein for insightful discussions that benefited this work. This research has been supported by NSF Grants CCF 1763702, AF 1901292, CNS 2148141, Tripods CCF 1934932, NSF AI Institute for Foundations of Machine Learning (IFML) 2019844, the Texas Advanced Computing Center (TACC) and research gifts by Western Digital, WNCG IAP, UT Austin Machine Learning Lab (MLL), Cisco and the Archie Straiton Endowed Faculty Fellowship. Giannis Daras has been supported by the Onassis Fellowship (Scholarship ID: F ZS 012-1/2022-2023), the Bodossaki Fellowship and the Leventis Fellowship. Aravind Gollakota was at UT Austin while this work was done.

References

Asad Aali, Marius Arvinte, Sidharth Kumar, and Jonathan I Tamir. 'Solving Inverse Problems with Score-Based Generative Priors learned from Noisy Data'. In: arXiv preprint arXiv:2305.01166 (2023) (page 10).

Stability AI. Deepfloyd IF . https://github.com/deep-floyd/IF . 2013 (pages 7, 18).

Brian D.O. Anderson. 'Reverse-time diffusion equation models'. In: Stochastic Processes and their Applications 12.3 (1982), pp. 313-326 (page 3).

Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S Li, Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. 'Cold diffusion: Inverting arbitrary image transforms without noise'. In: arXiv preprint arXiv:2208.09392 (2022) (page 4).

Joshua Batson and Loic Royer. 'Noise2self: Blind denoising by self-supervision'. In: International Conference on Machine Learning . PMLR. 2019, pp. 524-533 (page 2).

Felix JB Bäuerlein and Wolfgang Baumeister. 'Towards visual proteomics at high resolution'. In: Journal of Molecular Biology 433.20 (2021), p. 167187 (page 2).

Ashish Bora, Eric Price, and Alexandros G Dimakis. 'AmbientGAN: Generative models from lossy measurements'. In: International conference on learning representations . 2018 (pages 2, 3, 7, 9).

Nicholas Carlini, Florian Tramèr, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Úlfar Erlingsson, Alina Oprea, and Colin Raffel. 'Extracting Training Data from Large Language Models'. In: 30th USENIX Security Symposium (USENIX Security 21) . 2021, pp. 2633-2650. ISBN: 978-1-939133-24-3. URL: https://www.usenix.org/conference/usenixsecurity21/presentation/ carlini-extracting (visited on 11/10/2022) (pages 1, 2, 8).

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. 'Emerging properties in self-supervised vision transformers'. In: Proceedings of the IEEE/CVF international conference on computer vision . 2021, pp. 9650-9660 (page 9).

Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. 'Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions'. In: arXiv preprint arXiv:2209.11215 (2022) (pages 16, 17).

Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. 'Diffusion Posterior Sampling for General Noisy Inverse Problems'. In: The Eleventh International Conference on Learning Representations . 2023. URL: https://openreview. net/forum?id=OnD9zGAGT0k (pages 1, 7, 8, 22).

The Event Horizon Telescope Collaboration et al. 'First M87 Event Horizon Telescope Results. IV. Imaging the Central Supermassive Black Hole'. In: The Astrophysical Journal Letters 875.1 (Apr. 2019), p. L4. DOI: 10.3847/2041-8213/ab0e85 . URL: https://dx.doi. org/10.3847/2041-8213/ab0e85 (page 1).

Zhuo-Xu Cui, Chentao Cao, Shaonan Liu, Qingyong Zhu, Jing Cheng, Haifeng Wang, Yanjie Zhu, and Dong Liang. 'Self-score: Self-supervised learning on score-based models for mri reconstruction'. In: arXiv preprint arXiv:2209.00835 (2022).

Giannis Daras, Mauricio Delbracio, Hossein Talebi, Alex Dimakis, and Peyman Milanfar. 'Soft Diffusion: Score Matching with General Corruptions'. In: Transactions on Machine Learning Research (2023). ISSN: 2835-8856. URL: https://openreview.net/forum?id= W98rebBxlQ (pages 3, 4).

Mauricio Delbracio and Peyman Milanfar. 'Inversion by direct iteration: An alternative to denoising diffusion for image restoration'. In: arXiv preprint arXiv:2303.11435 (2023) (page 2).

Jeffrey M Ede. 'Deep learning in electron microscopy'. In: Machine Learning: Science and Technology 2.1 (2021), p. 011004 (page 2).

Bradley Efron. 'Tweedie's formula and selection bias'. In: Journal of the American Statistical Association 106.496 (2011), pp. 1602-1614 (page 4).

Yonina C. Eldar. 'Generalized SURE for Exponential Families: Applications to Regularization'. In: IEEE Transactions on Signal Processing 57.2 (2009), pp. 471-481. DOI: 10.1109/ TSP.2008.2008212 (page 2).

Angela F Gao, Oscar Leong, He Sun, and Katherine L Bouman. 'Image Reconstruction without Explicit Priors'. In: ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) . IEEE. 2023, pp. 1-5 (page 1).

Shi Guo, Zifei Yan, Kai Zhang, Wangmeng Zuo, and Lei Zhang. 'Toward convolutional blind denoising of real photographs'. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition . 2019, pp. 1712-1722 (page 2).

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. 'Masked autoencoders are scalable vision learners'. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2022, pp. 16000-16009 (pages 19, 20).

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. 'Gans trained by a two time-scale update rule converge to a local nash equilibrium'. In: Advances in neural information processing systems 30 (2017) (page 18).

Jonathan Ho, Ajay Jain, and Pieter Abbeel. 'Denoising diffusion probabilistic models'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 6840-6851 (pages 1, 18).

Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. 'Cascaded Diffusion Models for High Fidelity Image Generation.' In: J. Mach. Learn. Res. 23.47 (2022), pp. 1-33 (page 18).

'Huggingface Brain Tumor MRI Dataset'. In: (2020). URL: https://huggingface.co/ datasets/miladfa7/Brain-MRI-Images-for-Brain-Tumor-Detection/ (page 9).

Niels Bruun Ipsen, Pierre-Alexandre Mattei, and Jes Frellsen. 'not-MIWAE: Deep generative modelling with missing not at random data'. In: arXiv preprint arXiv:2006.12871 (2020) (page 2).

Matthew Jagielski, Om Thakkar, Florian Tramèr, Daphne Ippolito, Katherine Lee, Nicholas Carlini, Eric Wallace, Shuang Song, Abhradeep Thakurta, Nicolas Papernot, and Chiyuan Zhang. 'Measuring Forgetting of Memorized Training Examples'. In: arxiv:2207.00099[cs] (June 2022). DOI: 10.48550/arXiv.2207.00099 . arXiv: 2207.00099 [cs] . URL: http: //arxiv.org/abs/2207.00099 (visited on 11/10/2022) (pages 1, 2, 8).

Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir. 'Robust compressed sensing mri with deep generative priors'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 14938-14954 (page 1).

Zahra Kadkhodaie and Eero P Simoncelli. 'Solving linear inverse problems using the prior implicit in a denoiser'. In: arXiv preprint arXiv:2007.13640 (2020) (page 7).

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. 'Elucidating the design space of diffusion-based generative models'. In: arXiv preprint arXiv:2206.00364 (2022) (pages 1, 7, 8, 17-19).

Tero Karras, Miika Aittala, Samuli Laine, Erik Härkönen, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. 'Alias-free generative adversarial networks'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 852-863 (page 17).

Tero Karras, Samuli Laine, and Timo Aila. 'A style-based generator architecture for generative adversarial networks'. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition . 2019, pp. 4401-4410 (page 17).

Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. 'Analyzing and improving the image quality of stylegan'. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition . 2020, pp. 8110-8119 (page 17).

Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. 'Denoising Diffusion Restoration Models'. In: Advances in Neural Information Processing Systems (pages 1, 7, 8, 22).

Bahjat Kawar, Noam Elata, Tomer Michaeli, and Michael Elad. 'GSURE-Based Diffusion Model Training with Corrupted Data'. In: arXiv preprint arXiv:2305.13128 (2023).

Bahjat Kawar, Gregory Vaksman, and Michael Elad. 'SNIPS: Solving noisy inverse problems stochastically'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 2175721769 (page 1).

Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and Il-Chul Moon. 'Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation'. In: International Conference on Machine Learning . PMLR. 2022, pp. 1120111228 (page 17).

Alexander Krull, Tim-Oliver Buchholz, and Florian Jug. 'Noise2void-learning denoising from single noisy images'. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition . 2019, pp. 2129-2137 (page 2).

Jaakko Lehtinen, Jacob Munkberg, Jon Hasselgren, Samuli Laine, Tero Karras, Miika Aittala, and Timo Aila. 'Noise2Noise: Learning image restoration without clean data'. In: arXiv preprint arXiv:1803.04189 (2018) (page 2).

Steven Cheng-Xian Li, Bo Jiang, and Benjamin Marlin. 'Misgan: Learning from incomplete data with generative adversarial networks'. In: arXiv preprint arXiv:1902.09599 (2019) (page 19).

Hongyu Liu, Bin Jiang, Yi Xiao, and Chao Yang. 'Coherent Semantic Attention for Image Inpainting'. In: 2019 IEEE/CVF International Conference on Computer Vision (ICCV) (Oct. 2019). DOI: 10.1109/iccv.2019.00427 . URL: http://dx.doi.org/10.1109/ICCV. 2019.00427 (page 2).

Pierre-Alexandre Mattei and Jes Frellsen. 'MIWAE: Deep generative modelling and imputation of incomplete data sets'. In: International conference on machine learning . PMLR. 2019, pp. 4413-4423 (page 2).

Sachit Menon, Alexandru Damian, Shijia Hu, Nikhil Ravi, and Cynthia Rudin. 'Pulse: Selfsupervised photo upsampling via latent space exploration of generative models'. In: Proceedings of the ieee/cvf conference on computer vision and pattern recognition . 2020, pp. 24372445 (page 17).

Charles Millard and Mark Chiew. 'A Theoretical Framework for Self-Supervised MR Image Reconstruction Using Sub-Sampling via Variable Density Noisier2Noise'. In: IEEE Transactions on Computational Imaging 9 (2023), pp. 707-720. DOI: 10.1109/TCI.2023.3299212 .

Alex Nichol, Aditya Ramesh, Pamela Mishkin, Prafulla Dariwal, Joanne Jang, and Mark Chen. DALL · E 2 Pre-Training Mitigations . June 2022. URL: https://openai.com/blog/dalle-2-pre-training-mitigations/ (visited on 11/10/2022) (page 1).

Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. 'DINOv2: Learning Robust Visual Features without Supervision'. In: arXiv preprint arXiv:2304.07193 (2023) (page 9).

Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros. 'Context encoders: Feature learning by inpainting'. In: Proceedings of the IEEE conference on computer vision and pattern recognition . 2016, pp. 2536-2544 (page 2).

Elad Richardson, Yuval Alaluf, Or Patashnik, Yotam Nitzan, Yaniv Azar, Stav Shapiro, and Daniel Cohen-Or. 'Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation'. In: arXiv preprint arXiv:2008.00951 (2020) (page 2).

Trevor W Richardson, Wencheng Wu, Lei Lin, Beilei Xu, and Edgar A Bernal. 'Mcflow: Monte carlo flow models for data imputation'. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2020, pp. 14205-14214 (page 2).

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. 'High-resolution image synthesis with latent diffusion models'. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2022, pp. 10684-10695 (page 1).

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. 'High-resolution image synthesis with latent diffusion models'. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2022, pp. 10684-10695 (page 8).

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. 'Photorealistic text-to-image diffusion models with deep language understanding'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 36479-36494 (page 1).

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. 'Photorealistic text-to-image diffusion models with deep language understanding'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 36479-36494 (page 21).

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. 'Deep unsupervised learning using nonequilibrium thermodynamics'. In: International Conference on Machine Learning . PMLR. 2015, pp. 2256-2265 (page 1).

Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein. 'Diffusion Art or Digital Forgery? Investigating Data Replication in Diffusion Models'. In: arXiv preprint arXiv:2212.03860 (2022) (pages 1, 2, 8, 9).

Jiaming Song, Chenlin Meng, and Stefano Ermon. 'Denoising diffusion implicit models'. In: arXiv preprint arXiv:2010.02502 (2020) (pages 7, 8, 19).

Yang Song and Stefano Ermon. 'Generative modeling by estimating gradients of the data distribution'. In: Advances in Neural Information Processing Systems 32 (2019) (page 1).

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 'Score-based generative modeling through stochastic differential equations'. In: arXiv preprint arXiv:2011.13456 (2020) (pages 1, 3, 18).

Julián Tachella, Dongdong Chen, and Mike Davies. 'Unsupervised Learning From Incomplete Measurements for Inverse Problems'. In: arXiv preprint arXiv:2201.12151 (2022) (page 2).

Singanallur V Venkatakrishnan, Charles A Bouman, and Brendt Wohlberg. 'Plug-and-play priors for model based reconstruction'. In: 2013 IEEE Global Conference on Signal and Information Processing . IEEE. 2013, pp. 945-948 (page 7).

Pascal Vincent. 'A connection between score matching and denoising autoencoders'. In: Neural computation 23.7 (2011), pp. 1661-1674 (page 4).

Ge Wang, Jong Chul Ye, and Bruno De Man. 'Deep learning for tomographic image reconstruction'. In: Nature Machine Intelligence 2.12 (2020), pp. 737-748 (page 2).

Jiahui Yu, Zhe Lin, Jimei Yang, Xiaohui Shen, Xin Lu, and Thomas Huang. 'Free-Form Image Inpainting With Gated Convolution'. In: 2019 IEEE/CVF International Conference on Computer Vision (ICCV) (Oct. 2019). DOI: 10.1109/iccv.2019.00457 . URL: http: //dx.doi.org/10.1109/ICCV.2019.00457 (page 2).

Jiahui Yu, Zhe Lin, Jimei Yang, Xiaohui Shen, Xin Lu, and Thomas S Huang. 'Free-form image inpainting with gated convolution'. In: Proceedings of the IEEE/CVF international conference on computer vision . 2019, pp. 4471-4480 (pages 7, 17).

Yide Zhang, Yinhao Zhu, Evan Nichols, Qingfei Wang, Siyuan Zhang, Cody Smith, and Scott Howard. 'A poisson-gaussian denoising dataset with real fluorescence microscopy images'. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition . 2019, pp. 11710-11718 (page 2).

A Proofs

Proof of Theorem 4.1. Let h θ ∗ be a minimizer of equation 3.2, and for brevity let

be the difference between h θ ∗ and the claimed optimal solution. We will now argue that f must be identically zero.

First, by adding and subtracting A E [ x 0 | ˜ A x t , ˜ A ] , we can expand the objective value achieved at θ = θ ∗ in equation 3.2 as follows:

Here the first term is the irreducible error, while the third term vanishes by the tower law of expectations:

Thus the only part of J corr ( θ ∗ ) that actually depends on the parameter value θ ∗ is the second term. We now show that the second term can be made to vanish, and that this occurs precisely when f is identically 0:

For every ˜ A and x 0 , x t , by assumption we have that E A | ˜ A [ A T A ] is full-rank, and so the inner quadratic form is minimized when f ( ˜ A x t , ˜ A ) = 0 . Further, the term as a whole vanishes exactly when this holds for every ˜ A and x 0 , x t in the support, which means f must be identically zero.

Corollary A.1 (Inpainting noise model) . Consider the following inpainting noise model: A ∈ R n × n is a diagonal matrix where each entry A ii ∼ Ber(1 -p ) for some p > 0 (independently for each i ), and the additional noise is generated by drawing ˜ A | A such that ˜ A ii = A ii Ber(1 -δ ) for some small δ > 0 (again independently for each i ). Then the unique minimizer of the objective in equation 3.2 is

Proof. By Theorem 4.1, what we must show is that for any ˜ A in the support, E A | ˜ A [ A T A ] is full-rank. Fix any particular realization of ˜ A , which will be a diagonal matrix with only 0 s and 1 s. For indices i where ˜ A ii = 1 , we know that for any A drawn conditional on ˜ A , A ii = 1 as well, i.e.

And for indices i where ˜ A ii = 0 , by Bayes' rule we have

Thus we see that E A | ˜ A [ A T A ] = E A | ˜ A [ A ] is a diagonal matrix whose entries are 1 wherever ˜ A ii = 1 and q wherever ˜ A ii = 0 . This is clearly of full rank since q > 0 .

Corollary A.2 (Gaussian measurements noise model) . Consider the following noise model where we only observe m independent Gaussian measurements of the ground truth, and then one of the measurements is further omitted: A ∈ R m × n consists of m rows drawn independently from N (0 , I n ) , and ˜ A ∈ R m × n is constructed condtional on A by zeroing out its last row. Then the unique minimizer of the objective in equation 3.2 is

Proof. Again by Theorem 4.1, we must show that for any ˜ A in the support, E A | ˜ A [ A T A ] is full-rank. Fix any realization of ˜ A , which will have the following form:

Then it is clear that conditional on ˜ A , A has the following distribution:

Here b m is drawn entirely independently from N (0 , I n ) . Elementary manipulations now reveal that E A | ˜ A [ A T A ] = E b m ∼N (0 ,I n ) [ ˜ A T ˜ A + b m b T m ] = ˜ A T ˜ A + I n , which is clearly full rank (indeed, it is PSD with strictly positive eigenvalues).

A.1 Reduction

In this section we argue that if there is an algorithm that recovers the target distribution p 0 ( x 0 ) from i.i.d. samples ( A x 0 , A ) where A ∼ p ( A ) and x 0 ∼ p 0 ( x 0 ) , then, there is an algorithm that recovers p 0 ( x 0 ) without sample access, but instead, using access to an oracle that given t, x and A in the support of p ( A ) , returns E [ x 0 | A x t ] .

Indeed, for any A , Chen et al. [10] show that it is possible to recover the distribution of A x 0 given access to E [ A x 0 | A x t ] for any t and x under some minimal assumptions on the data distribution p 0 ( x 0 ) , see [10, Assumptions 1-3]. Using our oracle and using this theorem, we can recover the distribution of A x 0 for all A in the support. By sampling from these distributions, one can as well obtain samples of A x 0 for A ∼ p ( A ) and x 0 ∼ p 0 ( x 0 ) . Hence, if these samples are sufficient to recover p 0 ( x 0 ) , then, having an oracle to these conditional expectations is sufficient as well.

This intuition can be formalized as follows. Fix a distribution p A ( A ) over corruption matrices. For a distribution p 0 ( x 0 ) , denote by corrupt( p A , p 0 ) the distribution over pairs ( A,A x 0 ) where A ∼ p A ( A ) and x 0 ∼ p 0 ( x 0 ) . We say that it is possible to reconstruct p 0 ( x 0 ) from random corruptions A ∼ p A ( A ) if the following holds: for any two distributions, p 0 ( x 0 ) and p ′ 0 ( x ′ 0 ) that satisfy Assumptions 1-3 of Chen et al. [10], if corrupt( p A , p 0 ) = corrupt( p A , p ′ 0 ) , then p 0 = p ′ 0 . Similarly, we say that it is possible to reconstruct p 0 ( x 0 ) from conditional expectations given A ∼ p ( A ) if the following holds: for any distribution p 0 ( x 0 ) and p ′ 0 ( x ′ 0 ) that satisfy Assumptions 1-3 of Chen et al. [10], if for all x , t and A in the support of p A ,

then p 0 = p ′ 0 . Here, p 0 ,t ( x 0 , x t ) is obtained by sampling x 0 ∼ p 0 and x t = x 0 + σ t η where η ∼ N (0 , I ) . Similarly, p ′ 0 ,t ( x ′ 0 , x ′ t ) is obtained by the same process where x ′ 0 is instead sampled from p ′ 0 . We state the following lemma:

LemmaA.3. Fix a distribution p A ( A ) . If it is possible to reconstruct p 0 ( x 0 ) from random corruptions y 0 = A x 0 , A ∼ p A ( A ) , then it is possible to reconstruct p 0 ( x 0 ) given access to an oracle that computes the conditional expectations E [ x 0 | A x t , A ] , for A ∼ p A ( A ) and x t = x 0 + σ t η .

Proof. Assume that it is possible to reconstruct p 0 ( x 0 ) from random corruptions A ∼ p A ( A ) and we will prove that it is possible to reconstruct p 0 ( x 0 ) from conditional expectations given A ∼ p A ( A ) . To do so, let p 0 ( x 0 ) and p ′ 0 ( x ′ 0 ) be two distributions that satisfy Assumptions 1-3 of Chen et al. [10]. Assume that equation A.1 holds. We will prove that p 0 = p ′ 0 . Fix some A in the support of p A . By Chen et al. [10], there is an algorithm that samples from the distribution of A x 0 , x 0 ∼ p 0 ( x 0 ) , that only has access to E [ A x 0 | A x t , A ] and similarly, there is an algorithm that samples from the distribution of A x ′ 0 that only has access to E [ A x ′ 0 | A x ′ t , A ] . Since these two conditional expectations are assumed to be the same, then the distribution of A x 0 equals the distribution of A x ′ 0 . Consequently, corrupt( p A , p 0 ) = corrupt( p A , p ′ 0 ) . By the assumption that it is possible to reconstruct p 0 ( x 0 ) from random corruptions A ∼ p A ( A ) , this implies that p 0 = p ′ 0 . This completes the proof.

B Broader Impact and Risks

Generative models in general hold the potential to have far-reaching impacts on society in a variety of forms, coupled with several associated risks [43, 32, 33, 31]. Among other potential applications, they can be utilized to create deceptive images and perpetuate societal biases. To the best of our knowledge, our paper does not amplify any of these existing risks. Regarding the included MRI results, we want to clarify that we make no claim that such results are diagnostically useful. This experiment serves only as a toy demonstration that it can be potentially feasible to learn the distribution of MRI scans with corrupted samples. Significant further research must be done in collaboration with radiologists before our algorithm gets tested in clinical trials. Finally, we want to underline again that even though our approach seems to mitigate the memorization issue in generative models, we cannot guarantee the privacy of any training sample unless we make assumptions about the data distribution. Hence, we strongly discourage using this algorithm in applications where privacy is important before this research topic is investigated further.

C Training Details

We open-source our code and models to facilitate further research in this area: https://github.com/giannisdaras/ambient-diffusion.

For all the experiments in this paper, we concatenated the image with the mask along the channel dimension to condition the network on the measurement operator, i.e. the inpainting mask is an additional input channel.

Models trained from scratch. We trained models from scratch at different corruption levels on CelebA-HQ, AFHQ and CIFAR-10. The resolution of the first two datasets was set to 64 × 64 and for CIFAR-10 we trained on 32 × 32 .

We started from EDM's [30] official implementation and made some necessary changes. Architecturally, the only change we made was to replace the convolutional layers with Gated Convolutions [64] that are known to perform well for inpainting problems. We observed that this change stabilized the training significantly, especially in the high-corruptions regime. As in EDM, we use the architecture from the DDPM++ [37] paper.

To avoid additional design complexity, we tried to keep our hyperparameters as close as possible to the EDM paper. We observed that for high corruption levels, it was useful to add gradient clipping, otherwise, the training would often diverge. For all our experiments, we use gradient clipping with max-norm set to 1 . 0 . We underline that unfortunately, even with gradient clipping, the training at high corruption levels ( p ≥ 0 . 8 ), still diverges sometimes. Whenever this happened, we restarted the training from an earlier checkpoint. We list the rest of the hyperparameters we used in Table 2.

Training diffusion models from scratch is quite computationally intensive. We trained all our models for 200000 iterations. Our CIFAR-10 models required ≈ 2 days of training each on 6 A100 GPUs. Our AFHQ and CelebA-HQ models required ≈ 6 days of training each on 6 A100 GPUs. These numbers roughly match the performance reported in the EDM paper, indicating that the extra corruption we need to do on matrix A does not increase training time.

Due to the increased computational complexity of training these models, we could not extensively optimize the hyperparameters, e.g. the δ probability in our extra corruption. For higher corruption, e.g. for p = 0 . 9 , we noticed that we had to increase δ in order for the model to learn to perform well on the unobserved pixels. For a small-scale ablation study on CIFAR-10, see Table 7.

Finetuning Deepfloyd's IF. We access Deepfloyd's IF [2] model through the diffusers library. The model is a Cascaded Diffusion Model [24]. The first part of the pipeline is a text-conditional diffusion model that outputs images at resolution 64 × 64 . Next in the pipeline, there are two diffusion models that are conditioned both in the input text and the low-resolution output of the previous stage. The first upscaling module increases the resolution from 64 × 64 to 256 × 256 and the final one from 256 × 256 to 1024 × 1024 .

To reduce the computational requirements of the finetuning, we only finetune the first text-conditional diffusion model that works with 64 × 64 resolution images. Once the finetuning is completed, we use again the whole model to generate high-resolution images.

For the finetuning, we set the training batch size to 32 and the learning rate to 3e -6 . We train for a maximum of 15000 steps and we keep the checkpoint that gives the lowest error on the pixels that we further corrupted. To further reduce the computational requirements, we use an 8-bit Adam Optimizer and we train with half-precision.

For our CelebA finetuning experiments, we set δ = 0 . 1 and p = 0 . 8 . We experiment with the full training set, a subset of size 3000 (see Figure 1) and a subset of 300 . For the model trained with only 300 heavily corrupted samples, we did not observe memorization but the samples were of very low quality. Intuitively, our algorithm provides a way to control the trade-off between memorization and fidelity. Fully exploring this trade-off is a very promising direction for future work. For our MRI experiments, we use two non-overlapping blocks that each obfuscates 25% of the image and we evaluate the model in one of them.

All of our fine-tuning experiments can be completed in a few hours. Training for 15000 iterations takes ≈ 10 hours on an A100 GPU, but we usually get the best checkpoints earlier in the training.

D Evaluation Details

FID evaluation. Our FID [22] score is computed with respect to the training set, as is standard practice, e.g. see [23, 30, 58]. For each of our models trained from scratch, we generate 50000 images using the seeds 0 -49999 . Once we generate our images, we use the code provided in the official implementation of the EDM [30] paper for the FID computation.

E Additional Experiments

E.1 Restoration performance with noisy measurements

In Table 1 of the main paper, we compare the restoration performance of our models and vanilla diffusion models (trained with uncorrupted images). We compare the restoration performance in the task of random inpainting because it is straightforward to use our models to solve this inverse problem.

It is potentially feasible to use our trained generative models to solve any (linear or non-linear) inverse problem but we leave this direction for future work.

In this section, we further compare our models (trained with randomly inpainted images) in the task of inpainting with noisy measurements. Concretely, we want to predict x 0 , given measurements y t = A ( x 0 + σ y t η ) , η ∼ N (0 , I ) . As in the rest of the paper, we assume that the mask matrix A is known. We can solve this problem with one model prediction using our trained models since according to Eq. 4.1, we are learning: h θ ( y t , A, t ) = E [ x 0 | A x t = y t , A ] .

We use the EDM [30] state-of-the-art model trained on AFHQ as our baseline (as we did in the main paper). To use this pre-trained model to solve the noisy random inpainting inverse problem, we need a reconstruction algorithm. We experiment again with DPS and DDRM which can both handle inverse problems with noise in the measurements. We present our results in Table 3. As shown, our models significantly outperform the EDM models that use the DDRM reconstruction algorithm and perform on-par with the EDM models that use the DPS reconstruction algorithm.

E.2 Additional Comparisons

We incude additional comparisons on CIFAR-10 and CelebA with MisGAN [40]. The MisGAN work generalizes AmbientGAN in the case where the measurement operator is unknown. Specifically, the authors propose an additional generator that learns to model the corruption mechanism with adversarial training. The results are shown in Table 4.

E.3 Comparison with Supervised Methods

For completeness, we include a comparison with Masked AutoEncoders [21], a state-of-the-art supervised method for solving the random inpainting problem. The official repository of this paper does not include models trained on AFHQ. We compare with the available models that are trained on the iNaturalist dataset which is the most semantically close dataset we could find. We emphasize that this model was trained with access to uncorrupted images. Results are shown in 5. As shown, our method and DPS outperform this supervised baseline. We underline that this experiment is included for completeness and does not exclude the possibility that there are more performant supervised alternatives for random inpainting.

E.4 Sampler Ablation

By Lemma A.3, if it is possible to learn p 0 ( x 0 ) using corrupted samples then it is also possible to use our learned model to sample from p 0 ( x 0 ) . Even though such an algorithm exists, we do not know which one it is.

In the paper, we proposed a simple idea for sampling, the Fixed Mask Sampler. The Fixed Mask Sampler fixes a mask throughout the sampling process. Hence, sampling with this algorithm is equivalent to first sampling some pixels from the marginals of p 0 ( x 0 ) and then completing the rest of the pixels with the best reconstruction (the conditional expectation) in the last step. This simple sampler performs remarkably well and we use it throughout the main paper.

Sampling with Reconstruction Guidance. In the Fixed Mask Sampler, at any point the prediction is a convex combination of the current value and the predicted denoised image. As t → 0 , γ t → 0 . Hence, for the masked pixels, the fixed mask sampler outputs the conditional expectation of their value given the observed pixels. This leads to averaging effects as the corruption gets higher. To correct for this problem, we add one more term in the update: the Reconstruction Guidance term. The issue with the previous sampler is that the model never sees certain pixels. We would like to evaluate the model using different masks. However, the model outputs for the denoised image might be very different when evaluated with different masks. To account for this problem, we add an additional term that enforces updates that lead to consistency on the reconstructed image. The update of the sampler with Reconstruction Guidance becomes:

This sampler is inspired by the Reconstruction Guidance term used in Imagen [53]. In Table 6 we ablate the performance of this alternative sampler. For the reconstruction guidance sampler, we select each time four masks at random and we add an extra update term to the Fixed Mask Sampler that ensures that the predictions of the Fixed Mask Sampler are not very different compared to the predictions with the other four masks (that have different context regarding the current iterate x t ). We set the guidance parameter w t to the value 5 e -4 . As shown, this sampler improves the performance, especially for the low corruption probabilities where the extra masks give significant information about the current state to the predictions given only one fixed mask. However, the benefits of this sampler are vanishing for higher corruption. Given that the two samples perform on par and that the Reconstruction Guidance Sampler is much more computationally intensive (we need one extra prediction per step for each extra mask), we choose to use the Fixed Mask Sampler for all the experiments in the paper.

Table 6: Comparison between the Fixed Mask sampler and the Reconstruction Guidance Sampler.

E.5 δ Ablation

A hyperparameter of our method is δ , which controls the strength of the additional corruption. For the random inpainting case, δ needs to be positive in order to learn to predict the clean image for the missing pixels. On the other hand, as δ increases, so does the difficulty of the prediction task (which finally controls generation performance). In Table 7 we show how performance changes on CIFAR-10 for different values of δ . The recommended setting is δ = 0 . 1 but hyperparameter search for the value of δ . might lead to performance improvements.

E.6 Additional Figures

Figure 7 shows reconstructions of AFHQ corrupted images with the EDM AFHQ model trained on clean data (columns 3, 4) and our model trained on corrupted data (column 5). The restoration task is random inpainting at probability p = 0 . 8 . The last two rows also have measurement noise with σ y 0 = 0 . 05 .

In Figure 8, we repeat the experiment of Figure 6 of the main paper but for models trained on the full CelebA dataset and at different levels of corruption. As shown, increasing the corruption level leads to a clear shift of the distribution to the left, indicating less memorization. This comes at the cost of decreased performance, as reported in Table 4.

Input

Ground Truth

DDRM [34]

DPS [11]

Ours

Distribution of Similarity Values to the nearest training sample

Corruption probability = 0.0

Corruption probability = 0.1

Corruptionprobability=0.4

Number of generated samples

400

Corruptionprobability=0.6

Corruptionprobability=0.8

300

200

100

0

0.60

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

SimilarityValue

In the remaining pages, we include uncurated unconditional generations of our models trained at different corruption levels p . Results are shown in Figures 9,10,11.

3

AFHQ dataset, p = 0 . 9 , δ = 0 . 1

Uncurated samples from our model

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The authors propose the first diffusion-based method that can restore data from corruption without training from clean data. \n\n- The method can reuse regular diffusion sampler without much modification. \n\n- It is guaranteed theoretically for the model to recover clean data with some rank assumptions.\n\n- The method alleviates the memorization issues of diffusion models by considering corruption schemes.\n\nWEAKNESSES:\n- 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?\n\n- 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? \n\n- Line 248 the authors forgot line break.\n\n- 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.?\n\n- 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?\n\n- Why is there no comparison with AmbientGAN for Table 1, as it is an important baseline?\n\nQUESTIONS:\nMy questions and concerns are listed in the section above.\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: 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?

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThere are a number of fundamental technical and conceptual flaws which need to be considered and fixed: \n\n1) 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.\nIf this is the training setup, please clarify in the text. If not, please describe what was the training setup. \n\n2) 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. \n\n2) 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 .\n\n3) 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). \n\n4) 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?). \n\n5) 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). \n\n6) 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.\n\nQUESTIONS:\nIn addition to theoretical questions and comments I made under weaknesses: \n\nWhy do you think you achieve optimal sampling performance in one iteration? \nHow does the training time compares to uncorrupted models? Do you need to train on more samples? \n________________________________________________________\nNote: I am willing to raise the score if the questions asked under Weaknesses and here are addressed.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* The problem of handling corrupted datasets is significant even in generative model learning. This paper could be seen as one that explores this directions.\n* The paper is well-structured and easy to follow.\n* 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.\n\nWEAKNESSES:\n* Sampling Process\n  * It would be helpful to provide an algorithm or a detailed explanation of the sampling process from scratch.\n  * 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.\n\n* Related work\n  * 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].\n\n[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.\n\n[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.\n\n[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.\n\n[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).\n\n* Experiments\n  * It would be beneficial to provide the FID results for Figure 5 to demonstrate the differences.\n  * In addition to Figure 5, it would be valuable to compare the proposed model with existing generative models such as AmbientGAN in various experiments.\n\n* Presentation\n  * It would be helpful to have better spacing between subfigures in Figure 2 to improve the clarity of the captions.\n  * Proper citations are needed in the background section.\n  * Figure 3 would be better placed within a paragraph rather than between paragraphs.\n\nQUESTIONS:\nPlease see Weaknesses part.\n\nLIMITATIONS:\nThey provided in the last paragraph in the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}