SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Indeed, for any  $A$ , Chen et al. [10] show that it is possible to recover the distribution of  $Ax_0$  given access to  $\mathbb{E}[Ax_0 \mid Ax_t]$  for any  $t$  and  $\pmb{x}$  under some minimal assumptions on the data distribution  $p_0(\pmb{x}_0)$ , see [10, Assumptions 1-3]. Using our oracle and using this theorem, we can recover the distribution of  $Ax_0$  for all  $A$  in the support. By sampling from these distributions, one can as well obtain samples of  $Ax_0$  for  $A \sim p(A)$  and  $\pmb{x}_0 \sim p_0(\pmb{x}_0)$ . Hence, if these samples are sufficient to recover  $p_0(\pmb{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(\pmb{x}_0)$ , denote by  $\mathrm{corrupt}(p_A, p_0)$  the distribution over pairs  $(A, A\pmb{x}_0)$  where  $A \sim p_A(A)$  and  $\pmb{x}_0 \sim p_0(\pmb{x}_0)$ . We say that it is possible to reconstruct  $p_0(\pmb{x}_0)$  from random corruptions  $A \sim p_A(A)$  if the following holds: for any two distributions,  $p_0(\pmb{x}_0)$  and  $p_0'(\pmb{x}_0')$  that satisfy Assumptions 1-3 of Chen et al. [10], if  $\mathrm{corrupt}(p_A, p_0) = \mathrm{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 \sim p(A)$  if the following holds: for any distribution  $p_0(\pmb{x}_0)$  and  $p_0'(\pmb{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$ ,
$$
\mathbb {E} _ {\left(\boldsymbol {x} _ {0}, \boldsymbol {x} _ {t}\right) \sim p _ {0, t} \left(\boldsymbol {x} _ {0}, \boldsymbol {x} _ {t}\right)} \left[ \boldsymbol {x} _ {0} \mid A \boldsymbol {x} _ {t} = \boldsymbol {x} \right] = \mathbb {E} _ {\left(\boldsymbol {x} _ {0} ^ {\prime}, \boldsymbol {x} _ {t} ^ {\prime}\right) \sim p _ {0, t} ^ {\prime} \left(\boldsymbol {x} _ {0} ^ {\prime}, \boldsymbol {x} _ {t} ^ {\prime}\right)} \left[ \boldsymbol {x} _ {0} ^ {\prime} \mid A \boldsymbol {x} _ {t} ^ {\prime} = x \right] \tag {A.1}
$$
then  $p_0 = p_0'$ . Here,  $p_{0,t}(\boldsymbol{x}_0,\boldsymbol{x}_t)$  is obtained by sampling  $\boldsymbol{x}_0\sim p_0$  and  $\boldsymbol{x}_t = \boldsymbol{x}_0 + \sigma_t\boldsymbol{\eta}$  where  $\boldsymbol{\eta}\sim \mathcal{N}(0,I)$ . Similarly,  $p_{0,t}^{\prime}(\boldsymbol{x}_{0}^{\prime},\boldsymbol{x}_{t}^{\prime})$  is obtained by the same process where  $\boldsymbol{x}_0^{\prime}$  is instead sampled from  $p_0'$ . We state the following lemma:
Lemma A.3. Fix a distribution  $p_A(A)$ . If it is possible to reconstruct  $p_0(\pmb{x}_0)$  from random corruptions  $\pmb{y}_0 = A\pmb{x}_0$ ,  $A \sim p_A(A)$ , then it is possible to reconstruct  $p_0(\pmb{x}_0)$  given access to an oracle that computes the conditional expectations  $\mathbb{E}[\pmb{x}_0|A\pmb{x}_t, A]$ , for  $A \sim p_A(A)$  and  $\pmb{x}_t = \pmb{x}_0 + \sigma_t\pmb{\eta}$ .
Proof. Assume that it is possible to reconstruct  $p_0(\pmb{x}_0)$  from random corruptions  $A \sim p_A(A)$  and we will prove that it is possible to reconstruct  $p_0(\pmb{x}_0)$  from conditional expectations given  $A \sim p_A(A)$ . To do so, let  $p_0(\pmb{x}_0)$  and  $p_0'(\pmb{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\pmb{x}_0$ ,  $\pmb{x}_0 \sim p_0(\pmb{x}_0)$ , that only has access to  $\mathbb{E}[A\pmb{x}_0 \mid A\pmb{x}_t, A]$  and similarly, there is an algorithm that samples from the distribution of  $A\pmb{x}_0'$  that only has access to  $\mathbb{E}[A\pmb{x}_0' \mid A\pmb{x}_t', A]$ . Since these two conditional expectations are assumed to be the same, then the distribution of  $A\pmb{x}_0$  equals the distribution of  $A\pmb{x}_0'$ . Consequently,  $\mathrm{corrupt}(p_A, p_0) = \mathrm{corrupt}(p_A, p_0')$ . By the assumption that it is possible to reconstruct  $p_0(\pmb{x}_0)$  from random corruptions  $A \sim 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 \times 64$  and for CIFAR-10 we trained on  $32 \times 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 \geq 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.
Table 2: Training Hyperparameters
Dataset	Iter	Batch	LR	SDE	p	δ	Aug. Prob.	Ch. Multipliers	Dropout
CIFAR-10	200000	512	1e-3	VP	{0.2, 0.4, 0.6, 0.8}	0.12	(1, 1, 1)	0.13
AFHQ	256	2e-4	{0.2, 0.4, 0.6, 0.8, 0.9}	0.1	0.15	(1, 2, 2, 2)	0.25
CelebA-HQ	{0.2, 0.6, 0.8}	0.3	0.1
0.9
Training diffusion models from scratch is quite computationally intensive. We trained all our models for 200000 iterations. Our CIFAR-10 models required  $\approx 2$  days of training each on 6 A100 GPUs. Our AFHQ and CelebA-HQ models required  $\approx 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  $\delta$  probability in our extra corruption. For higher corruption, e.g. for  $p = 0.9$ , we noticed that we had to increase  $\delta$  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 \times 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 \times 64$  to  $256 \times 256$  and the final one from  $256 \times 256$  to  $1024 \times 1024$ .
To reduce the computational requirements of the finetuning, we only finetune the first text-conditional diffusion model that works with  $64 \times 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  $3\mathrm{e} - 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  $\delta = 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  $\approx 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.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  $\boldsymbol{x}_0$ , given measurements  $\boldsymbol{y}_t = A(\boldsymbol{x}_0 + \sigma_{\boldsymbol{y}_t}\boldsymbol{\eta})$ ,  $\boldsymbol{\eta} \sim \mathcal{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_\theta(\boldsymbol{y}_t, A, t) = \mathbb{E}[x_0 | A\boldsymbol{x}_t = \boldsymbol{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.
Dataset	Corruption Probability	Measurement Noise (σy0)	Method	LPIPS	PSNR	NFE
AFHQ	0.4	0.05	Ours	0.0861	29.46	1
DPS	0.0846	29.83	100
0.2061	24.47	35
DDRM	0.1739	25.38	99
0.1677	25.58	199
0.6	0.6	0.05	Ours	0.1031	27.40	1
DPS	0.0949	27.92	100
0.4066	18.73	35
DDRM	0.3626	19.49	99
0.3506	19.70	199
0.8	0.8	0.05	Ours	0.1792	23.21	1
DPS	0.1778	23.01	100
0.5879	13.65	35
DDRM	0.5802	13.99	99
0.5753	14.09	199
Table 3: Comparison of our model (trained on corrupted data) with state-of-the-art diffusion models on CelebA (DDIM [56] model) and AFHQ (EDM [30] model) for solving the random inpainting inverse problem with noise.
E.2 Additional Comparisons
We include 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(\pmb{x}_0)$  using corrupted samples then it is also possible to use our learned model to sample from  $p_0(\pmb{x}_0)$ . Even though such an algorithm exists, we do not know which one it is.
CelebA
Corruption Probability	Method	FID
0.6	MisGAN	37.42
0.6	Ambient Diffusion	6.08
0.8	MisGAN	100.0
0.8	Ambient Diffusion	11.19
0.9	MisGAN	141.11
0.9	Ambient Diffusion	25.53
CIFAR-10
Corruption Probability	Method	FID
0.4	MisGAN	18.95
0.4	Ambient Diffusion	18.85
0.6	MisGAN	49.30
0.6	Ambient Diffusion	28.88
0.8	MisGAN	111.50
0.8	Ambient Diffusion	46.27
Table 4: Comparison with MisGAN.
Corruption Probability	Method	LPIPS	PSNR	NFE
0.4	Ours	0.0304	33.27	1
DPS	0.0203	34.06	100
MAE	0.0752	28.88	1
0.6	Ours	0.0628	29.46	1
DPS	0.0518	30.03	100
MAE	0.0995	25.89	1
0.8	Ours	0.1245	25.37	1
DPS	0.1078	25.30	100
MAE	0.1794	22.01	1
Table 5: Comparison with the MAE [21], a state-of-the-art supervised method for solving the restoration task of random inpainting.
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 \to 0$ ,  $\gamma_t \to 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:
$$
\boldsymbol {x} _ {t - \Delta t} = \underbrace {\frac {\sigma_ {t - \Delta t}}{\sigma_ {t}}} _ {\gamma_ {t}} x _ {t} + \underbrace {\frac {\sigma_ {t} - \sigma_ {t - \Delta t}}{\sigma_ {t}}} _ {1 - \gamma_ {t}} \mathbb {E} [ \boldsymbol {x} _ {0} | \tilde {A} \boldsymbol {x} _ {t}, \tilde {A} ] - \frac {w _ {t}}{2} \nabla_ {\boldsymbol {x} _ {t}} \mathbb {E} _ {A ^ {\prime} \sim p (\tilde {A} ^ {\prime})} \left[ | | \mathbb {E} [ \boldsymbol {x} _ {0} | \tilde {A} \boldsymbol {x} _ {t}, \tilde {A} ] - \mathbb {E} [ \boldsymbol {x} _ {0} | \tilde {A} ^ {\prime} \boldsymbol {x} _ {t}, \tilde {A} ^ {\prime} ] | | ^ {2} \right]. \tag {E.1}
$$
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  $\boldsymbol{x}_t$ ). We set the guidance parameter  $w_t$  to the value  $5e - 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.
Sampler Type	Corruption Probability	FID	Inception Score
Fixed Mask	0.2	11.70	7.97
0.4	18.85	7.45
0.6	28.88	6.88
0.8	46.27	6.14
Reconstruction Guidance	0.2	11.59	8.01
0.4	18.52	7.51
0.6	28.90	6.91
0.8	46.31	6.13
E.5  $\delta$  Ablation
A hyperparameter of our method is  $\delta$ , which controls the strength of the additional corruption. For the random inpainting case,  $\delta$  needs to be positive in order to learn to predict the clean image for the missing pixels. On the other hand, as  $\delta$  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  $\delta$ . The recommended setting is  $\delta = 0.1$  but hyperparameter search for the value of  $\delta$ , might lead to performance improvements.
Table 6: Comparison between the Fixed Mask sampler and the Reconstruction Guidance Sampler.
p	δ	Inception Score
0.4	0.0	6.70 (from Figure 5)
0.4	0.1	7.45 (from Figure 6)
0.4	0.4	6.95 (experiment added for the rebuttal)
Table 7: Ablation Study for the value of  $\delta$  .
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  $\sigma_{\mathbf{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.
DDRM [34]
DPS [11]
Figure 7: 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.6$ . The last two rows also have measurement noise with  $\sigma_{y_0} = 0.05$ .
Figure 8: Distribution of similarity values to the nearest neighbor in the dataset for finetuned IF models on the full CelebA dataset at different levels of corruption. Please note that similarity values above 0.95 roughly correspond to the same person, while similarities below 0.75 correspond to almost random faces. As shown, increasing the corruption level leads to a clear shift of the distribution to the left, indicating less memorization.
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.
CelebA dataset,  $p = 0.6$ ,  $\delta = 0.1$
Uncurated samples from our model
CelebA dataset,  $p = 0.8$ ,  $\delta = 0.1$
CelebA dataset,  $p = 0.9$ ,  $\delta = 0.3$
Figure 9: Left column: CelebA-HQ training dataset with random inpainting at different levels of corruption  $p$  (the survival probability is  $(1 - p) \cdot (1 - \delta)$ ). Right column: Unconditional generations from our models trained with the corresponding parameters. As shown, the generations become slightly worse as we increase the level of corruption, but we can reasonably well learn the distribution even with  $93\%$  pixels missing (on average) from each training image.
AFHQ dataset,  $p = 0.4$ ,  $\delta = 0.1$
AFHQ dataset,  $p = 0.6$ ,  $\delta = 0.1$
AFHQ dataset,  $p = 0.8$ ,  $\delta = 0.1$
AFHQ dataset,  $p = 0.9$ ,  $\delta = 0.1$
Figure 10: Left column: AFHQ training dataset with random inpainting at different levels of corruption  $p$  (the survival probability is  $(1 - p) \cdot (1 - \delta)$ ). Right column: Unconditional generations from our models trained with the corresponding parameters. As shown, the generations become slightly worse as we increase the level of corruption, but we can reasonably well learn the distribution even with  $91\%$  pixels missing (on average) from each training image.
CIFAR-10 training dataset, Uncurated samples from our  $p = 0.2, \delta = 0.1$  model
CIFAR-10 training dataset, Uncurated samples from our  $p = 0.4, \delta = 0.1$  model
CIFAR-10 training dataset, Uncurated samples from our  $p = 0.6$ ,  $\delta = 0.1$  model
CIFAR-10 training dataset, Uncurated samples from our  $p = 0.8, \delta = 0.1$  model
Figure 11: Left column: CIFAR-10 training dataset with random inpainting at different levels of corruption  $p$  (the survival probability is  $(1 - p) \cdot (1 - \delta)$ ). Right column: Unconditional generations from our models trained with the corresponding parameters. As shown, the generations become slightly worse as we increase the level of corruption, but we can reasonably well learn the distribution even with high corruption.