Ambient Diffusion Posterior Sampling: Solving Inverse Problems with Diffusion Models trained on Corrupted Data

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

Introduction to Ambient Diffusion Posterior Sampling

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