SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Giannis Daras
giannisdaras@utexas.edu
Kulin Shah
kulinshah@utexas.edu
Aravind Gollakota
aravindg@cs.utexas.edu
Alexandros G. Dimakis
dimakis@austin.utexas.edu
Yuval Dagan
yuvald@berkeley.edu
Adam Klivans
klivans@utexas.edu
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.
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.
Figure 1: Left panel: Baseline method of vanilla finetuning Deepfloyd IF using 3000 images from CelebA-HQ. We show generated sample images and nearest neighbors from the finetuning set. As shown, the generated samples are often near-identical copies from training data. This verifies related work Carlini et al. [8], Somepalli et al. [55], and Jagielski et al. [27] that pointed out that diffusions often generate training samples. Right panel: We finetune the same foundation model (Deepfloyd IF) using our method and 3000 highly corrupted training images. The corruption adds noise and removes 80 percent random pixels. We show generated samples and nearest neighbors from the training set. Our method still learns the clean distribution of faces (with some quality deterioration, as shown) but does not memorize training data. We emphasize that our training is performed without ever accessing clean training data.
Corrupted Training Images
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  $\mathcal{D}$  when the training set only contains highly-corrupted examples drawn from  $\mathcal{D}$ . Specifically, we consider the problem of learning to sample from the target distribution  $p_0(x_0)$  given corrupted
Random painting
in-Further corruption
Figure 2: Illustration of our method: Given training data with deleted pixels, we corrupt further by erasing more (illustrated with green color). We feed the learner the further corrupted images and we evaluate it on the originally observed pixels. We can do this during training since the green pixel values are known to us. The score network learner has no way of knowing whether a pixel was missing from the beginning or whether it was corrupted by us. Hence, the score network learns to predict the clean image everywhere. Our method is analogous to grading a random subset of the questions in a test, but the students not knowing which questions will be graded.
Block inpaint-Further corrup-ting tion
samples  $A\mathbf{x}_0$  where  $A \sim p(A)$  is a random corruption matrix (with known realizations and prior distribution) and  $\mathbf{x}_0 \sim p_0(\mathbf{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  $\mathbb{E}[\pmb{x}_0|\tilde{A} (\pmb{x}_0 + \sigma_t\pmb{\eta}),\tilde{A} ]$ , for all noise levels  $t$  and for  $\tilde{A}\sim p(\tilde{A}\mid A)$  being a further corrupted version of  $A$ . The result holds for a general family of corruption processes  $A\sim p(A)$ . For various corruption processes, we show that the further degradation introduced by  $\tilde{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(\pmb{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.
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:  $\mathrm{d}\pmb{x} = \pmb{f}(\pmb{x},t)\mathrm{d}t + g(t)\mathrm{d}\pmb{w}$ , where  $\pmb{w}$  is the standard Wiener process. Such corruption processes are reversible and the reverse process is also described by an Ito SDE [3]:  $\mathrm{d}\pmb{x} = \left(\pmb{f}(\pmb{x},t) - g^{2}(t)\nabla_{\pmb{x}}\log p_{t}(\pmb{x})\right)\mathrm{d}t + g(t)\mathrm{d}\pmb{w}$ . The designer of the diffusion model is usually free to choose the drift function  $\pmb{f}(\cdot,\cdot)$  and the diffusion function  $g(\cdot)$ . Typical choices are setting  $\pmb{f}(\pmb{x},t) = \mathbf{0}$ ,  $g(t) = \sqrt{\frac{\mathrm{d}\sigma_t^2}{\mathrm{d}t}}$  (Variance Exploding SDE) or setting  $\pmb{f}(\pmb{x},t) = -\beta(t)\pmb{x}$ ,  $g(t) = \sqrt{\beta(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  $\nabla_{\pmb{x}}\log p_t(\pmb{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  $\boldsymbol{x}_t = \boldsymbol{x}_0 + \sigma_t\boldsymbol{\eta}$ .
Vincent [61] showed that we can learn the score function at level  $t$  by optimizing for the score-matching objective:
$$
J (\theta) = \frac {1}{2} \mathbb {E} _ {\left(\boldsymbol {x} _ {0}, \boldsymbol {x} _ {t}\right)} \left\| \boldsymbol {h} _ {\theta} \left(\boldsymbol {x} _ {t}, t\right) - \boldsymbol {x} _ {0} \right\| ^ {2}. \tag {2.1}
$$
Specifically, the score function can be written in terms of the minimizer of this objective as:
$$
\nabla_ {\boldsymbol {x} _ {t}} \log p _ {t} (\boldsymbol {x} _ {t}) = \frac {\boldsymbol {h} _ {\theta^ {*}} (\boldsymbol {x} _ {t} , t) - \boldsymbol {x} _ {t}}{\sigma_ {t}}. \tag {2.2}
$$
This result reveals a fundamental connection between the score-function and the best restoration model of  $\pmb{x}_0$  given  $\pmb{x}_t$ , known as Tweedie's Formula [17]. Specifically, the optimal  $h_{\theta^*}(\pmb{x}_t,t)$  is given by  $\mathbb{E}[\pmb{x}_0|\pmb{x}_t]$ , which means that
$$
\nabla_ {\boldsymbol {x} _ {t}} \log p _ {t} \left(\boldsymbol {x} _ {t}\right) = \frac {\overbrace {\mathbb {E} \left[ \boldsymbol {x} _ {0} \mid \boldsymbol {x} _ {t} \right]} ^ {\text {b e s t r e s t o r a t i o n}} - \boldsymbol {x} _ {t}}{\sigma_ {t}}. \tag {2.3}
$$
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:  $\pmb{x}_t = C_t\pmb{x}_0 + \sigma_t\pmb{\eta}$ . Specifically, Soft Diffusion proposes the Soft Score Matching objective:
$$
J _ {\text {s o f t}} (\theta) = \frac {1}{2} \mathbb {E} _ {\left(\boldsymbol {x} _ {0}, \boldsymbol {x} _ {t}\right)} \left| \left| C _ {t} \left(\boldsymbol {h} _ {\theta} \left(\boldsymbol {x} _ {t}, t\right) - \boldsymbol {x} _ {0}\right) \right| \right| ^ {2}, \tag {2.4}
$$
and shows that it is sufficient to recover the score function via a generalized Tweedie's Formula:
$$
\nabla_ {\boldsymbol {x} _ {t}} \log p _ {t} (\boldsymbol {x} _ {t}) = \frac {C _ {t} \mathbb {E} [ \boldsymbol {x} _ {0} | \boldsymbol {x} _ {t} ] - \boldsymbol {x} _ {t}}{\sigma_ {t}}. \tag {2.5}
$$
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  $\sigma_t$  becomes zero, i.e. we observe samples from the true distribution.
As explained in the introduction, in many cases we do not observe uncorrupted images  $\boldsymbol{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.  $\boldsymbol{y}_0 = A\boldsymbol{x}_0$ , and the corruption matrices  $A: \mathbb{R}^{m\times n}$ . We note that we are interested in non-invertible corruption matrices. We ask two questions:
1. Is it possible to learn  $\mathbb{E}[\pmb{x}_0|A(\pmb{x}_0 + \sigma_t\pmb{\eta}), A]$  for all noise levels  $t$ , given only access to corrupted samples ( $\pmb{y}_0 = A\pmb{x}_0, A$ )?
2. If so, is it possible to use this restoration model  $\mathbb{E}[\pmb{x}_0|A(\pmb{x}_0 + \sigma_t\pmb{\eta}), A]$  to recover  $\mathbb{E}[\pmb{x}_0|\pmb{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  $\mathbb{E}[\pmb{x}_0|\pmb{x}_t]$  directly using  $\mathbb{E}[\pmb{x}_0|A\pmb{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  $\boldsymbol{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_\theta$  which receives a corruption matrix  $A$  and a noisy version of a corrupted image,  $y_t = A \underbrace{\left( x_0 + \sigma_t \eta \right)}_{x_t}$  where  $\eta \sim \mathcal{N}(0, I)$ , and produces an estimate for the conditional expectation. The simplest idea would be to simply ignore the missing pixels and optimize for:
$$
J _ {\text {n a i v e}} ^ {\text {c o r r}} (\theta) = \frac {1}{2} \mathbb {E} _ {\left(\boldsymbol {x} _ {0}, \boldsymbol {x} _ {t}, A\right)} \left| \left| A \left(\boldsymbol {h} _ {\theta} (A, A \boldsymbol {x} _ {t}, t) - \boldsymbol {x} _ {0}\right) \right| \right| ^ {2}, \tag {3.1}
$$
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_{\theta'}$  satisfying  $A h_{\theta'}(A, y_t, t) = A \mathbb{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  $\tilde{A} = BA$  for some matrix  $B$  that is selected randomly given  $A$ . In our example of missing pixels,  $\tilde{A}$  is obtained from  $A$  by randomly erasing an additional fraction  $\delta$  of the pixels that survive after the corruption  $A$ . Here,  $B$  will be diagonal where each element is 1 with probability  $1 - \delta$  and 0 w.p.  $\delta$ . We will penalize the model on recovering all the pixels that are visible in the sample  $Ax_0$ : this includes both the pixels that survive in  $\tilde{A}x_0$  and those that are erased by  $\tilde{A}$ . The formal training objective is given by minimizing the following loss:
$$
J ^ {\operatorname {c o r r}} (\theta) = \frac {1}{2} \mathbb {E} _ {\left(\boldsymbol {x} _ {0}, \boldsymbol {x} _ {t}, A, \tilde {A}\right)} \left| \left| A \left(\boldsymbol {h} _ {\theta} (\tilde {A}, \tilde {A} \boldsymbol {x} _ {t}, t) - \boldsymbol {x} _ {0}\right) \right| \right| ^ {2}, \tag {3.2}
$$
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  $\delta$  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  $\tilde{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  $\tilde{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  $\mathbb{E}[\pmb{x}_0|\tilde{A}\pmb{x}_t,\tilde{A}]$  using our algorithm and samples  $(\pmb{y}_0 = A\pmb{x}_0,A)$ . Our goal is to satisfy this condition while adding minimal further corruption, i.e., while keeping  $\tilde{A}$  close to  $A$ .
3.2 Sampling
To sample from  $p_0(\pmb{x}_0)$  using the standard diffusion formulation, we need access to  $\nabla_{\pmb{x}_t} \log p_t(\pmb{x}_t)$ , which is equivalent to having access to  $\mathbb{E}[\pmb{x}_0|\pmb{x}_t]$  (see Eq. 2.3). Instead, our model is trained to predict  $\mathbb{E}[\pmb{x}_0|\tilde{\pmb{A}}\pmb{x}_t,\tilde{\pmb{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  $\boldsymbol{x}_t$ . Hence, we should not expect our model to give good estimates of  $\mathbb{E}[\boldsymbol{x}_0|\tilde{\boldsymbol{A}}\boldsymbol{x}_t,\tilde{\boldsymbol{A}}]$  for corruption matrices  $A$  that belong to the tails of the distribution  $p(A)$ .
The simplest idea is to sample a mask  $\tilde{A} \sim p(\tilde{A})$  and approximate  $\mathbb{E}[\pmb{x}_0|\pmb{x}_t]$  with  $\mathbb{E}[\pmb{x}_0|\tilde{A}\pmb{x}_t,\tilde{A}]$ . Under this approximation, the discretized sampling rule becomes:
$$
\boldsymbol {x} _ {t - \Delta t} = \underbrace {\frac {\sigma_ {t - \Delta t}}{\sigma_ {t}}} _ {\gamma_ {t}} \boldsymbol {x} _ {t} + \underbrace {\frac {\sigma_ {t} - \sigma_ {t - \Delta t}}{\sigma_ {t}}} _ {1 - \gamma_ {t}} \underbrace {\mathbb {E} [ \boldsymbol {x} _ {0} | \tilde {A} \boldsymbol {x} _ {t} , \tilde {A} ]} _ {\hat {\boldsymbol {x}} _ {0}}. \tag {3.3}
$$
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)
Input: noise schedule  $\{\sigma_t\}_0^T$  , model parameters  $\theta$  , and step size  $\Delta t$  for discretizing the Reverse ODE. Sample  $x_{T}\sim \mathcal{N}(0,\sigma_{T}^{2}I)$  ▷ Sample initial noise Sample  $\tilde{A}\sim p(\tilde{A})$  ▷ Sample a mask  $t\gets T$    
repeat   
 $y_{t}\gets \tilde{A} x_{t}$ $\hat{x}_0\gets h_\theta (\pmb {y}_t,\tilde{\pmb{A}},t)$ $x_{t - \Delta t}\gets \frac{\sigma_{t - \Delta t}}{\sigma_t}\pmb {x}_t + \frac{\sigma_t - \sigma_{t - \Delta t}}{\sigma_t}\hat{x}_0$ $\mathbb{E}[\pmb {x}_0|\pmb {x}_t]$  with  $\hat{x}_0$ $t\gets t - \Delta t$    
until  $t\leq \Delta t$    
return  $\hat{x}_0$
We underline that the Fixed Mask Sampler does not come with a theoretical guarantee of sampling from  $p_0(\pmb{x}_0)$ . In the Appendix, we prove that whenever it is possible to reconstruct  $p_0(\pmb{x}_0)$  from corrupted samples, it is also possible to reconstruct it using access to  $\mathbb{E}[\pmb{x}_0|A\pmb{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  $\pmb{h}(A, \pmb{y}_t, t) = \mathbb{E}[\pmb{x}_0 | \pmb{y}_t, A]$ . We now show that under a certain assumption on the distribution of  $A$  and  $\tilde{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  $\tilde{A}$ ,  $A$  has considerable variability, and the latter ensures that the best way to predict  $A\pmb{x}_0$  as a function of  $\tilde{A}\pmb{x}_t$  and  $\tilde{A}$  is to optimally predict  $\pmb{x}_0$  itself. All proofs are deferred to the Appendix.
Theorem 4.1. Assume a joint distribution of corruption matrices  $A$  and further corruption  $\tilde{A}$ . If for all  $\tilde{A}$  in the support it holds that  $\mathbb{E}_{A|\tilde{A}}[A^T A]$  is full-rank, then the unique minimizer of the objective in equation 3.2 is given by
$$
\boldsymbol {h} _ {\theta^ {*}} (\tilde {A}, \boldsymbol {y} _ {t}, t) = \mathbb {E} [ \boldsymbol {x} _ {0} \mid \tilde {A} \boldsymbol {x} _ {t}, \tilde {A} ] \tag {4.1}
$$
Two simple examples that fit into this framework (see Corollaries A.1 and A.2 in the Appendix) are:
- Inpainting:  $A \in \mathbb{R}^{n \times n}$  is a diagonal matrix where each entry  $A_{ii} \sim \operatorname{Ber}(1 - p)$  for some  $p > 0$  (independently for each  $i$ ), and the additional noise is generated by drawing  $\tilde{A}|A$  such that  $\tilde{A}_{ii} = A_{ii} \cdot \operatorname{Ber}(1 - \delta)$  for some small  $\delta > 0$  (again independently for each  $i$ ).<sup>1</sup>
- Gaussian measurements:  $A \in \mathbb{R}^{m \times n}$  consists of  $m$  rows drawn independently from  $\mathcal{N}(0, I_n)$ , and  $\tilde{A} \in \mathbb{R}^{m \times n}$  is constructed conditional on  $A$  by zeroing out its last row.