Ambient Diffusion: Learning Clean Distributions from Corrupted Data (2305.19256v1)

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 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.

An Essay on "Ambient Diffusion: Learning Clean Distributions from Corrupted Data"

The paper presents an innovative diffusion-based framework for learning unknown distributions utilizing highly corrupted samples, a scenario prevalent in scientific applications where uncorrupted data acquisition is either infeasible or costly. This method offers a strategic advantage by training generative models that are less prone to memorizing individual training samples, as these models never observe clean data during training.

The central concept of the research involves introducing additional measurement distortion during the diffusion process. Specifically, the model attempts to predict the original corrupted image from a further corrupted version of the image. This methodology ensures that the model learns the conditional expectation of the uncorrupted image given the additional corruption. The method is theoretically supported for any corruption process satisfying specific conditions, including inpainting and compressed sensing. Underlying this approach is the premise that by adding controlled amounts of further corruption, the model remains uninformed about the true nature of the missing data, thereby necessitating accurate predictions on the full data according to the learned distribution.

Numerical evaluations reveal substantial promise, with the models trained using this framework on standard benchmarks (CelebA, CIFAR-10, and AFHQ). Remarkably, the algorithm succeeds in learning distributions even with 90% of pixels missing from training samples. In practice, this approach allows for the fine-tuning of foundation models using small, corrupted datasets such as MRI scans exhibiting block corruptions, enabling the learning of clean distributions without risking memorization of specific training instances.

When it comes to implications, the training of models with such significant levels of corruption indicates robust adaptability and generalization capacities, potentially paving the way for applications in privacy-preserving generative modeling. This is particularly relevant in medical imaging fields such as MRI, where the ability to generalize from compromised datasets can reduce data acquisition demands. The model's success also presents an avenue for exploring diffusion in more complex or heavily obscured datasets, potentially leading to innovation across diverse domains beyond imaging.

The work contrasts against various prior studies. Unlike ambient generative adversarial networks (AmbientGANs), which require an infinitely powerful discriminator for hypothesis validation, the proposed diffusion framework employs the restored conditional expectation using corrupted data. Moreover, while significant literature in supervised learning focuses on model-to-data restoration, the current work diverges by centering on achieving the sample distribution of clean data, thus alleviating the potential need for direct image restoration.

Future development could expand the scope of this method to address more complex corruptions or explore efficiency improvements in both the training process and the restoration mechanisms, aiming at broader practical applications and possibly real-time deployment. Discovering optimal strategies to balance model efficiency with minimal corruption during training could enhance the adaptability and utility in even more domains.

Overall, the paper introduces a compelling technique to learn clean data distributions with corrupted samples, representing a new pathway for generative modeling in constrained data scenarios and potentially transforming applications where clean data remains inaccessible.