We propose Decoupled Annealing Posterior Sampling (DAPS) algorithm, which incorporates a pre-trained diffusion model as a prior to solve general inverse problems. By iteratively sampling from the time-marginal distribution, our method creates approximate samples from the posterior distribution.
We compare DAPS with DPS on a nonlinear inverse problem with a 2D Gaussian mixture prior. DAPS is able to sample more accurately from the posterior distribution.
We evaluate DAPS with a series of inverse problems in image domain, including five linear inverse problems: (1) super-resolution, (2) Gaussian deblurring, (3) motion deblurring, (4) inpainting with a box mask, (5) inpainting on random pixels, and three nonlinear inverse problems: (1) phase retrieval, (2) nonlinear deblurring, (3) high dynamic range reconstruction. DAPS outperforms existing baselines in terms of perceptual quality (LPIPS) and peak signal-to-noise ratio (PSNR).
DAPS can also be applied to large-scale text-conditioned latent diffusion models, e.g., Stable Diffusion.
We apply DAPS to solve the compressed sensing multi-coil magnetic resonance imaging problem (CS-MRI).
DAPS is able to generate diverse samples given less information. For example, we show several generated samples for (1) inpainting with large boxes and (2) 16x super-resolution.
We run DAPS with various configurations to test its performance under different computing budgets in terms of the number of function evaluation (NFEs). DAPS is able to generate high-quality samples with a modest time cost.
We are grateful to Pika for providing the computing resources essential for this research. We also extend our thanks to the Kortschak Scholars Fellowship for supporting B.Z. and W.C. at Caltech. J.B. acknowledges support from the Wally Baer and Jeri Weiss Postdoctoral Fellowship. A.A. is supported in part by Bren endowed chair and by the AI2050 senior fellow program at Schmidt Sciences.
@misc{zhang2024improvingdiffusioninverseproblem,
title={Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing},
author={Bingliang Zhang and Wenda Chu and Julius Berner and Chenlin Meng and Anima Anandkumar and Yang Song},
year={2024},
eprint={2407.01521},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2407.01521},
}