Ambient Diffusion Posterior Sampling: Solving Inverse Problems with Diffusion Models trained on Corrupted Data (2403.08728v1)

Abstract: We provide a framework for solving inverse problems with diffusion models learned from linearly corrupted data. Our method, Ambient Diffusion Posterior Sampling (A-DPS), leverages a generative model pre-trained on one type of corruption (e.g. image inpainting) to perform posterior sampling conditioned on measurements from a potentially different forward process (e.g. image blurring). We test the efficacy of our approach on standard natural image datasets (CelebA, FFHQ, and AFHQ) and we show that A-DPS can sometimes outperform models trained on clean data for several image restoration tasks in both speed and performance. We further extend the Ambient Diffusion framework to train MRI models with access only to Fourier subsampled multi-coil MRI measurements at various acceleration factors (R=2, 4, 6, 8). We again observe that models trained on highly subsampled data are better priors for solving inverse problems in the high acceleration regime than models trained on fully sampled data. We open-source our code and the trained Ambient Diffusion MRI models: https://github.com/utcsilab/ambient-diffusion-mri .

• The paper introduces a novel framework that trains diffusion models on corrupted data to solve inverse problems.
• It demonstrates that models trained on corrupted images excel in tasks like compressed sensing and super-resolution.
• It extends the approach to multi-coil MRI, outperforming traditional methods at high acceleration factors.

Exploring the Bounds of Diffusion Models in Solving Inverse Problems with Corrupted Training Data

Introduction to Ambient Diffusion Posterior Sampling

The field of generative models has seen significant advancement in recent years, notably with the adoption of Diffusion Models in solving complex inverse problems. Traditionally, these models have been trained on clean, fully-observed datasets. However, there exist situations where acquiring uncorrupted or fully observed data is challenging, urging the need for models that can learn from and solve inverse problems with corrupted data. Addressing this gap, the work on Ambient Diffusion Posterior Sampling (A-DPS) proposes a framework that not only trains Diffusion Models on corrupted data but also utilizes them to solve arbitrary inverse problems conditioned on linear measurements originating from different forward operators.

Training Diffusion Models on Corrupted Data

A significant contribution of this research is the methodological advancement in training generative models, specifically Diffusion Models, with linearly corrupted datasets. The methodology, dubbed Ambient Diffusion, explores the avenue of training these models with data subjected to various corruption processes (e.g., image inpainting). Through a nuanced training objective, it demonstrates that models trained on heavily corrupted samples can form better priors for inverse problem solving in highly corrupted regimes compared to their counterparts trained on clean data.

Evaluation on Standard Natural Image Datasets

The evaluation of A-DPS spans across several datasets of natural images, such as CelebA, FFHQ, and AFHQ. The experimentation targets two core tasks: compressed sensing and super-resolution, assessing both performance and speed. Interestingly, the models trained on corrupted data showcase an ability to outperform models fed with clean data when dealing with high levels of data corruption. This finding accentuates the potential of utilizing corrupted datasets for training, pushing the boundaries of generative models' applicability when clean data is scarce.

Extension to Multi-coil Fourier Subsampled MRI

One of the paper's pivotal extensions is applying the Ambient Diffusion framework to the domain of MRI imaging, a field where obtaining fully sampled data can be impractical. By innovatively training models directly on Fourier subsampled multi-coil MRI data, this work shed light on the efficacy of Ambient Diffusion Models in medical imaging. Remarkably, it was found that these models surpass ones trained on fully sampled data for inverse problems at high acceleration factors, thereby presenting an intriguing avenue for medical image reconstruction where data corruption is prevalent.

Theoretical Grounding and Practical Implications

The paper meticulously underpins its experimental findings with theoretical insights, proving that transforming corrupted into further corrupted representations and then training to predict the less corrupted original can mathematically lead to learning correct estimations. This theoretical foundation lends significant credibility to the ambient training methodology, providing a pathway for future research in generative model training with corrupted datasets.

Conclusion and Future Directions

The advent of Ambient Diffusion Posterior Sampling marks a crucial step towards leveraging corrupted or partially observed data to train powerful generative models for solving intricate inverse problems. Beyond natural image restoration tasks, its successful application to MRI reconstruction exemplifies the broader implications for medical imaging and potentially other domains where data corruption is a norm. As generative models continue to evolve, the exploration of training and application paradigms that embrace data imperfections will be paramount, and A-DPS stands as a seminal work guiding this endeavor.