Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing

Bingliang Zhang*,1, Wenda Chu*,1, Julius Berner1, Chenlin Meng2, Anima Anandkumar1, Yang Song3,
1California Institute of Technology    2Stanford University    3OpenAI
Paper arXiv Code
        

Abstract

Diffusion models have recently achieved success in solving Bayesian inverse problems with learned data priors. Current methods build on top of the diffusion sampling process, where each denoising step makes small modifications to samples from the previous step. However, this process struggles to correct errors from earlier sampling steps, leading to worse performance in complicated nonlinear inverse problems, such as phase retrieval. To address this challenge, we propose a new method called Decoupled Annealing Posterior Sampling (DAPS) that relies on a novel noise annealing process. Specifically, we decouple consecutive steps in a diffusion sampling trajectory, allowing them to vary considerably from one another while ensuring their time-marginals anneal to the true posterior as we reduce noise levels. This approach enables the exploration of a larger solution space, improving the success rate for accurate reconstructions. We demonstrate that DAPS significantly improves sample quality and stability across multiple image restoration tasks, particularly in complicated nonlinear inverse problems. For example, we achieve a PSNR of 30.72dB on the FFHQ 256 dataset for phase retrieval, which is an improvement of 9.12dB compared to existing methods.

Method

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.


Experimental Results

2D Synthetic Data

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.


Linear Inverse Problems

We evaluate DAPS with a series of linear inverse problems in image domain, including (1) super-resolution, (2) Gaussian deblurring, (3) motion deblurring, (4) inpainting with a box mask, and (5) inpainting on random pixels. DAPS outperforms existing baselines in terms of perceptual quality (LPIPS) and peak signal-to-noise ratio (PSNR).

Nonlinear Inverse Problems

We evaluate DAPS on three nonlinear tasks in image domain, including (1) phase retrieval, (2) nonlinear deblurring, and (3) high dynamic range (HDR) reconstruction. DAPS is remarkably stable when handling nonlinear inverse problems, especially in the highly ill-posed phase retrieval task.

Sample Diversity

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.

Acknowledgements

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.

BibTeX

@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}, 
}