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

TL;DR

This work tackles inverse problems when training data is corrupted, introducing Ambient Diffusion Posterior Sampling (Ambient DPS) which leverages diffusion models trained on linearly corrupted data as priors. It extends Ambient Diffusion to MRI by training on Fourier-domain subsampled data and developing a posterior-sampling procedure (A-DPS) that uses ambient scores and a measurement likelihood under a potentially different forward operator. The authors demonstrate that diffusion models trained on highly corrupted data can outperform models trained on clean data in high-acceleration MRI settings and achieve competitive or superior results on natural-image tasks, often with speed advantages. These findings broaden the practical utility of diffusion priors for ill-posed inverse problems when fully observed training data is unavailable, with implications for fast, robust medical imaging and beyond.

Abstract

We provide a framework for solving inverse problems with diffusion models learned from linearly corrupted data. Firstly, we extend the Ambient Diffusion framework to enable training directly from measurements corrupted in the Fourier domain. Subsequently, we train diffusion models for MRI with access only to Fourier subsampled multi-coil measurements at acceleration factors R= 2,4,6,8. Secondly, we propose Ambient Diffusion Posterior Sampling (A-DPS), a reconstruction algorithm that leverages generative models pre-trained on one type of corruption (e.g. image inpainting) to perform posterior sampling on measurements from a different forward process (e.g. image blurring). For MRI reconstruction in high acceleration regimes, we observe that A-DPS models trained on subsampled data are better suited to solving inverse problems than models trained on fully sampled data. We also test the efficacy of A-DPS on natural image datasets (CelebA, FFHQ, and AFHQ) and show that A-DPS can sometimes outperform models trained on clean data for several image restoration tasks in both speed and performance.

Figures (26)

• Figure 1: Ambient Diffusion Posterior Sampling (Ambient DPS). During training, we only have access to linearly corrupted data from a forward operator $A_{\mathrm{train}}$. We use the Ambient Diffusion framework to learn a generative model, $G_{\mathrm{ambient}}$, for the uncorrupted distribution, $p({\bm{x}}_0)$. At inference time, we sample from the posterior distribution $p({\bm{x}}_0 | {\bm{y}}_{\mathrm{A_{\mathrm{inf}}}})$, for measurements ${\bm{y}}_{\mathrm{inf}}$ coming from a different forward operator, $A_{\mathrm{inf}}$.
• Figure 2: Prior samples from diffusion models trained on MRI scans. Columns $1-4$: Diffusion models trained on subsampled MRI scans at acceleration factors $R=2,4,6,8$, using the Ambient Diffusion framework extended for Fourier subsampled training. Columns $5-8$: EDM models trained with L1-wavelet reconstructions of subsampled scans at $R=2,4,6,8$. Column $9$: NCSNV2 trained with fully sampled scans. Column $10$: EDM trained with fully sampled scans. We observe that Ambient Diffusion models consistently produce high-quality and realistic MRI scans even in high acceleration regimes.
• Figure 3: Posterior sampling reconstructions for MRI scans using models trained on Fourier subsampled data at various acceleration factors (Ambient DPS, columns $2-5$) and a model trained on clean data (FS-DPS dps, column $6$). Rows $1$ and $3$ show reconstructions at $R=4$ and $R=8$, respectively, while rows $2$ and $4$ display the difference to the ground truth on a $10\times$ scale. At high inference acceleration (R=$8$), Ambient DPS, outperforms FS-DPS, despite that the underlying models were trained solely on corrupted data.
• Figure 4: Posterior sampling reconstructions for FFHQ, showing: (1) Corrupted Input: Examples from FFHQ dataset corrupted using Resolution Downscaling ($4\times$) and Gaussian Noise ($\sigma=0.05$), (2) FS-DPS: Reconstruction using Diffusion Posterior Sampling (DPS), with a diffusion model trained with clean fully-sampled data ($p=0.0$), (3) A-DPS: Reconstruction using Ambient-DPS, with a diffusion model trained on randomly inpainted data with erasure probability ($p=0.6$), (4) Ground Truth: Original uncorrupted examples from FFHQ dataset. We observe that Ambient DPS, provides better reconstructions (qualitatively) even though the underlying models were trained on corrupted data ($p=0.6$).
• Figure 5: Compressed Sensing results, AFHQ: performance metric and standard deviation. As shown, the model trained with clean data ($p=0.0$) only outperforms the models trained with corrupted data for more than $1000$ measurements, in both LPIPS and MSE.

Theorems & Definitions (3)

• Theorem 2.1: (Informal)
• Theorem A.1: Formal Statement of Theorem \ref{['th:main_theorem']}
• proof