PaperlayerView exact revision on arXiv

Source-linked AI summary

Ambient Diffusion: Learning Clean Distributions from Corrupted Data

Giannis Daras, Kulin Shah, Yuval Dagan, Aravind Gollakota, Alexandros G. Dimakis, Adam Klivans

arXiv:2305.19256v1cs.LGcs.AIcs.CVcs.ITmath.IT

TL;DR

How can clean-data distributions be learned when high-quality samples are costly or impossible to obtain, while addressing diffusion-model memorization concerns? The method further corrupts observed measurements during diffusion and trains a model to predict the originally corrupted image, targeting a conditional expectation under further corruption. The method learns the distribution of uncorrupted images reasonably well with up to 90% missing pixels, outperforming AmbientGAN and natural baselines.

  • Problem

    How can clean-data distributions be learned when high-quality samples are costly or impossible to obtain, while addressing diffusion-model memorization concerns?

    Source: §1 · p. 1

  • Method

    The method further corrupts observed measurements during diffusion and trains a model to predict the originally corrupted image, targeting a conditional expectation under further corruption.

    Source: §1 · pp. 2–3, §3.1 · p. 5

  • Results

    The method learns the distribution of uncorrupted images reasonably well with up to 90% missing pixels, outperforming AmbientGAN and natural baselines.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The reported results show that clean distributions can be learned from corrupted samples without observing clean training images, including through finetuning on small corrupted datasets.

    Source: §1 · p. 3, §5.2 · p. 10

    Although theory says E[x0|Axt,A] can suffice whenever corrupted samples identify p0, the authors report no practical algorithm for that reconstruction route.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

We present the first diffusion-based framework that can learn an unknown distribution using only highly-corrupted samples. This problem arises in scientific applications where access to uncorrupted samples is impossible or expensive to acquire. Another benefit of our approach is the ability to train generative models that are less likely to memorize individual training samples since they never observe clean training data. Our main idea is to introduce additional measurement distortion during the diffusion process and require the model to predict the original corrupted image from the further corrupted image. We prove that our method leads to models that learn the conditional expectation of the full uncorrupted image given this additional measurement corruption. This holds for any corruption process that satisfies some technical conditions (and in particular includes inpainting and compressed sensing). We train models on standard benchmarks (CelebA, CIFAR-10 and AFHQ) and show that we can learn the distribution even when all the training samples have $90\%$ of their pixels missing. We also show that we can finetune foundation models on small corrupted datasets (e.g. MRI scans with block corruptions) and learn the clean distribution without memorizing the training set.

1 Introduction

Ambient Diffusion targets learning clean-data distributions when high-quality training images are unavailable or costly, while also addressing concerns that diffusion models can memorize examples. It further corrupts observed measurements during diffusion and reports distribution learning from highly corrupted samples.

1 Introduction

  • High-quality training data may be expensive or impossible to obtain, while diffusion-model memorization has raised privacy, security, and copyright concerns.The introduction gives black-hole imaging and high-quality MRI as examples of settings with data-access constraints.Source: §1 · p. 1
  • The paper distinguishes its goal—sampling fresh, unseen examples from the clean distribution—from restoration methods that reconstruct individual corrupted samples.It notes that restoration methods are effective only at relatively low corruption levels.Source: §1 · p. 2
  • The method adds measurement distortion during diffusion and trains the model to predict the originally corrupted image.Its stated theoretical target is a conditional expectation of the uncorrupted image under further corruption.Source: §1 · pp. 2–3, Figure 2 · p. 3
  • With up to 90% of pixels missing, the method learns image distributions reasonably well and outperforms AmbientGAN and natural baselines.The paper also reports performance on CelebA, CIFAR-10, and AFHQ.Source: §1 · p. 3
  • On certain inverse problems, the models perform on par with or outperform state-of-the-art diffusion models using one prediction step rather than hundreds.The compared baselines were trained on uncorrupted data.Source: §1 · p. 3
  • The approach can finetune pretrained foundation models on small corrupted datasets, and sufficiently corrupted training is reported to prevent memorization.The paper reports finetuning in a few hours on a single GPU.Source: §1 · p. 3

2 Background

Diffusion models gradually corrupt data toward a distribution that can be sampled, then learn a score function for reversing that process. The score can be related to a conditional estimate of the original image, including in generalized corruption models.

2 Background

  • A diffusion model gradually transforms the data distribution into one that can be sampled, commonly using a stochastic differential equation, with a corresponding reverse process.The paper describes these forward and reverse processes as Ito SDEs.Source: §2 · p. 3
  • Training aims to learn the score function, and Tweedie’s Formula relates the optimal score estimate to the conditional expectation of the original image given the corrupted image.The stated optimal denoiser is E[x0|xt].Source: §2 · pp. 3–4
  • Generalized diffusion models use non-Markovian corruption xt = C_t x0 + σ_tη and a generalized Tweedie formula to recover the score.Soft Diffusion is cited as proposing a corresponding score-matching objective.Source: §2 · p. 4
  • In these generalized models, C_t is a design choice and becomes the identity at t = 0, when noise is zero and samples come from the true distribution.The paper compares choosing C_t with choosing the SDE functions f and g.Source: §2 · p. 4

3 Method

The method learns from non-invertible corrupted measurements by adding further corruption during diffusion training, so the model must recover pixels that were originally observed. At sampling time, it uses the resulting restoration predictions to approximate clean-data estimates; experiments produce samples reasonably close to the true distribution even with 90% of pixels missing.

3.1 Training

  • Further corruption makes the model predict the original corrupted sample; because it cannot distinguish naturally missing pixels from deliberately erased ones, it must also learn the latter.The naive objective leaves missing-pixel predictions unconstrained, whereas the further-corruption probability can be arbitrarily small as long as it is positive.Source: §3.1 · p. 5
  • The approach extends beyond random inpainting to Gaussian measurements and block inpainting, while aiming to keep the additional corruption minimal.For Gaussian measurements, one row can be deleted; for block inpainting, one additional non-overlapping block can be removed.Source: §3.1 · p. 5

3.2 Sampling

  • Fixed-mask sampling approximates E[x0|xt] with E[x0|˜Axt, ˜A], but estimates may be unreliable for corruption matrices in the tails of p(A).For random inpainting, the identity mask is technically in the support but has exponentially small probability when the corruption probability is at least a constant.Source: §3.2 · p. 5
  • With 90% of pixels missing, the fixed-mask sampler generates samples reasonably close to the true distribution by metrics such as FID and Inception.The paper uses this sampler for its main experiments unless otherwise noted.Source: §3.2 · p. 5
  • Reconstruction Guidance adds a cross-mask consistency term to the sampler, with modest improvements over fixed-mask sampling for certain corruption ranges.The fixed-mask sampler can average predictions for masked pixels; the added term addresses inconsistencies between denoised predictions from different masks.Source: §3.2 · p. 6
  • The paper proves that distribution recovery is possible in theory whenever recovery from corrupted samples is possible, but reports no practical algorithm for that general task.The result concerns recovery using access to E[x0|Axt, A].Source: §3.2 · p. 6

4 Theory

The theory identifies a condition on the corruption matrices under which the training objective has a population minimizer that predicts the clean image from the further-corrupted measurement. Inpainting and Gaussian measurements fit the framework, though the learned predictor uses the further-degraded matrix rather than the original one.

4 Theory

  • When E[A^T A|˜A] is full-rank for each ˜A in its support, the objective’s population minimizer is essentially the clean-image conditional expectation.The assumption formalizes variability in A conditional on ˜A, so predicting the measurement optimally entails predicting x0 itself.Source: §4 · p. 6
  • The framework includes random inpainting with positive pixel-erasure probability and Gaussian measurements formed by independently sampled rows, with one row removed for ˜A.For inpainting, ˜A further erases each surviving pixel with probability δ > 0.Source: §4 · p. 6
  • The minimizer conditions on ˜A rather than the originally desired A, so the method seeks a further corruption close to A.For inpainting, letting δ approach zero makes ˜A’s distribution very close to A’s while retaining positive additional corruption.Source: §4 · p. 6

5 Experimental Evaluation

Experiments evaluate random-inpainting diffusion models trained from scratch and foundation-model finetuning on corrupted data. They report inverse-problem and generation results, reduced memorization, and MRI-image adaptation, alongside quality and measurement-noise limitations.

5.1 Training from scratch on corrupted data

  • On CelebA and AFHQ random inpainting, the corrupted-data model performs on par with clean-data diffusion solvers using one score evaluation, versus typically hundreds of reconstruction steps.Source: Table 1 · p. 7, §5.1 · p. 8
  • Across all CIFAR-10 corruption levels, Ambient Diffusion outperforms both a diffusion model without further corruption and AmbientGAN.Source: Figure 3 · p. 8, §5.1 · p. 8
  • CelebA-HQ and AFHQ retain decent FID with 90% of pixels deleted, while CIFAR-10 performance degrades faster, potentially due to lower resolution.Source: §5.1 · p. 8

5.2 Finetuning foundation models on corrupted data

  • On 3,000-image CelebA finetuning, deleting 80% of pixels sharply reduces memorization versus standard finetuning, even through convergence.Source: §5.2 · p. 9
  • Finetuning on 155 MRI scans with 25% block obfuscation learns full brain-tumor image statistics, generating 256 × 256 images from 64 × 64 training images.Source: Figure 5 · p. 9, §5.2 · p. 10
  • Higher corruption makes memorization less likely but degrades generator quality, and the method cannot handle noisy measurements.Source: §5.2 · p. 10

A Proofs

The proofs establish when the training objective identifies the conditional-mean predictor and show that conditional-expectation access can support distribution recovery under stated assumptions. The broader-impact discussion limits the MRI demonstration and cautions that memorization mitigation does not guarantee privacy.

A Proofs

  • The objective’s excess error vanishes precisely when the predictor’s difference from the claimed solution is zero, with full rank ruling out other minimizers.The proof isolates a quadratic term and shows it vanishes exactly when f is identically zero.Source: A Proofs · p. 14
  • For random inpainting with additional independent pixel omissions, uniqueness follows because the conditional expected Gram matrix is full-rank.Its diagonal entries are 1 at retained pixels and q > 0 at omitted pixels.Source: A Proofs · pp. 14–15
  • For Gaussian measurements with one measurement further omitted, the conditional expected Gram matrix is ˜A^T˜A + I_n, hence full-rank.The omitted measurement contributes an independent Gaussian row whose expected outer product is I_n.Source: A Proofs · p. 15

A.1 Reduction

  • Under Chen et al. [10]’s assumptions, an oracle for conditional expectations suffices to recover p0 when recovery from random corruptions is possible.The reduction uses the oracle to recover distributions of Ax0, then obtain samples of the corrupted pairs used by the assumed recovery procedure.Source: §A.1 · p. 15

B Broader Impact and Risks

  • The MRI experiment is only a toy feasibility demonstration, and the paper makes no privacy guarantee without assumptions on the data distribution.The authors call for radiologist collaboration before clinical trials and discourage privacy-sensitive use pending further research.Source: B Broader Impact and Risks · p. 16

C Training Details

Training details describe the from-scratch model setup, finetuning configuration, and practical constraints. High-corruption training can remain unstable, while very small finetuning datasets can yield low-quality samples.

C Training Details

  • From-scratch models used Gated Convolutions in the EDM implementation, which the authors report stabilized training, especially at high corruption.They trained for 200,000 iterations; CIFAR-10 took about two days and AFHQ and CelebA-HQ about six days each on six A100 GPUs.Source: C Training Details · p. 16
  • At corruption levels p ≥ 0.8, training sometimes diverged despite gradient clipping, and higher corruption sometimes required increasing δ.When divergence occurred, training was restarted from an earlier checkpoint.Source: C Training Details · pp. 16–17
  • Finetuning Deepfloyd IF modified only its 64 × 64 text-conditional stage before using the full cascade for high-resolution generation.A 15,000-step run took about 10 hours on an A100, though the best checkpoints usually came earlier.Source: C Training Details · p. 17
  • Finetuning on 300 heavily corrupted CelebA samples showed no observed memorization, but produced very low-quality samples.The authors frame the method as offering a memorization–fidelity trade-off; MRI finetuning used two non-overlapping blocks, each obscuring 25% of the image.Source: C Training Details · p. 17

D Evaluation Details

The evaluation details specify how FID was computed for models trained from scratch.

D Evaluation Details

  • FID was computed against the training set using 50,000 generated images per from-scratch model and EDM’s official evaluation code.The images used seeds 0–49,999.Source: D Evaluation Details · p. 17

E.1 Restoration performance with noisy measurements

The section evaluates random inpainting from noisy measurements using models trained on randomly inpainted data. The authors report that their models outperform EDM with DDRM and perform on par with EDM using DPS.

E.1 Restoration performance with noisy measurements

  • The trained models predict x0 from noisy measurements in one model prediction, assuming the mask matrix A is known.The learned quantity is E[x0|Axt = yt, A].Source: §E.1 · p. 18
  • The evaluation uses noisy measurements yt = A(x0 + σytη) and compares restoration against an AFHQ-trained EDM baseline.The authors leave evaluation on other linear or nonlinear inverse problems for future work.Source: §E.1 · pp. 17–18
  • The models significantly outperform EDM with DDRM and perform on par with EDM with DPS on noisy random inpainting.The comparison uses the AFHQ-trained EDM baseline.Source: Table 3 · p. 18, §E.1 · p. 18

E.2 Comparison with Supervised Methods

The comparison reports that Ambient Diffusion and DPS outperform the available Masked AutoEncoders baseline, while noting that the comparison does not rule out stronger supervised alternatives.

E.2 Comparison with Supervised Methods

  • Ambient Diffusion and DPS outperform the available Masked AutoEncoders baseline on random inpainting.The authors caution that this comparison does not exclude more performant supervised alternatives.Source: §E.2 · p. 18, Table 4 · p. 18
  • The comparison uses available Masked AutoEncoders models trained on iNaturalist, the closest semantically related dataset the authors could find.The MAE model had access to uncorrupted images during training.Source: §E.2 · p. 18
  • The authors include this experiment for completeness and do not claim it establishes superiority over all supervised inpainting methods.They note the MAE repository lacked models trained on AFHQ.Source: §E.2 · p. 18

E.3 Sampler Ablation

The Fixed Mask Sampler is the paper’s primary sampler, while Reconstruction Guidance adds cross-mask consistency but brings limited benefit at higher corruption and greater computational cost.

E.3 Sampler Ablation

  • Lemma A.3 guarantees that if corrupted samples permit learning p0(x0), the learned model can sample from p0(x0), but the authors do not know which algorithm achieves this.Source: §E.3 · p. 19
  • The Fixed Mask Sampler keeps one mask throughout sampling, first sampling pixels from p0(x0) marginals and then completing the rest with a final conditional-expectation reconstruction.The paper uses this sampler throughout its main experiments.Source: §E.3 · p. 19
  • The Reconstruction Guidance Sampler improves sampling especially at low corruption, but its benefits vanish at higher corruption and it requires an extra prediction per step for each additional mask.The authors therefore use the Fixed Mask Sampler throughout the experiments.Source: §E.3 · p. 19, Table 5 · p. 19

E.4 Additional Figures

Additional figures compare reconstructions, examine nearest-neighbor similarity as corruption increases, and show unconditional generations across datasets at high corruption.

E.4 Additional Figures

  • AFHQ reconstruction figures compare an EDM model trained on clean data with the paper’s model trained on corrupted data; the last two rows include measurement noise.Source: §E.4 · p. 19, Figure 7 · p. 20
  • Increasing corruption shifts finetuned CelebA models’ nearest-neighbor similarity distribution left, indicating less memorization, alongside decreased performance.Similarities above 0.95 roughly correspond to the same person, while values below 0.75 correspond to almost-random faces.Source: §E.4 · p. 19, Figure 8 · p. 20
  • Unconditional generations become slightly worse with more corruption, yet remain reasonably good with 93% of CelebA-HQ pixels and 91% of AFHQ pixels missing on average.CIFAR-10 generations are also shown at high corruption.Source: Figure 9 · p. 21, Figure 10 · p. 23, Figure 11 · p. 24
Loading 2305.19256v1…