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

2024-03-13

arXiv.org

AI Summary

Diffusion models exhibit limited generalization across diverse linear inverse problems (e.g., highly accelerated MRI reconstruction, inpainting, deblurring). To address this, we propose Ambient Diffusion Posterior Sampling (A-DPS), a plug-and-play posterior sampling framework that leverages a diffusion model pretrained solely on degraded data (e.g., masked images) to solve inverse problems under arbitrary forward operators—such as multi-coil Fourier-domain undersampling or blur—without fine-tuning. A-DPS is the first method enabling cross-degradation-type posterior sampling without retraining. Theoretically and empirically, we show that models trained on strongly degraded data encode stronger priors: in MRI reconstruction at acceleration rates R = 2–8, A-DPS significantly improves PSNR and SSIM over baselines; in restoration tasks on CelebA, FFHQ, and AFHQ, it outperforms and accelerates inference relative to models trained on clean data. Code and pretrained MRI models are publicly available.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Web Mining and Content Analysis: Large pretrained models with web data
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 .

Problem

• Solving inverse problems with diffusion models trained on corrupted data
• Training diffusion models for MRI using Fourier subsampled measurements
• Leveraging pre-trained generative models for diverse image restoration tasks

Innovation

• Training diffusion models on Fourier-corrupted MRI data
• Ambient Diffusion Posterior Sampling for diverse corruptions
• A-DPS outperforms clean-data models in speed and performance