Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing

Bingliang Zhang, Wenda Chu, Julius Berner, Chenlin Meng, Anima Anandkumar, Yang Song

TL;DR

This work introduces Decoupled Annealing Posterior Sampling (DAPS), a diffusion-model-based framework to solve Bayesian inverse problems, especially those with nonlinear forward operators. By decoupling consecutive samples along the diffusion trajectory and sampling time-marginals $p({\mathbf{x}}_t|{\mathbf{y}})$ via a decoupled noise-annealing process, DAPS enables large, non-local corrections early in the sampling path while preserving convergence as noise vanishes. The method supports both pixel-space and latent-space diffusion priors (LatentDAPS), extends to large-scale text-conditioned latent diffusion models, and even accommodates discrete diffusion, demonstrating superior reconstruction quality and stability on image restoration, HDR, phase retrieval, and CS-MRI tasks. Empirically, DAPS achieves higher PSNR/LPIPS and lower measurement-error across nonlinear inverse problems and shows favorable efficiency, making it a practical, versatile tool for inverse problem solving with diffusion priors.

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.

Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing

TL;DR

This work introduces Decoupled Annealing Posterior Sampling (DAPS), a diffusion-model-based framework to solve Bayesian inverse problems, especially those with nonlinear forward operators. By decoupling consecutive samples along the diffusion trajectory and sampling time-marginals via a decoupled noise-annealing process, DAPS enables large, non-local corrections early in the sampling path while preserving convergence as noise vanishes. The method supports both pixel-space and latent-space diffusion priors (LatentDAPS), extends to large-scale text-conditioned latent diffusion models, and even accommodates discrete diffusion, demonstrating superior reconstruction quality and stability on image restoration, HDR, phase retrieval, and CS-MRI tasks. Empirically, DAPS achieves higher PSNR/LPIPS and lower measurement-error across nonlinear inverse problems and shows favorable efficiency, making it a practical, versatile tool for inverse problem solving with diffusion priors.

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.
Paper Structure (58 sections, 4 theorems, 34 equations, 28 figures, 10 tables, 3 algorithms)

This paper contains 58 sections, 4 theorems, 34 equations, 28 figures, 10 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Diffusion Models
  4. Bayesian Inverse Problems with Diffusion
  5. Method
  6. Decoupled Annealing Posterior Sampling
  7. Practical Design Choice of DAPS
  8. DAPS with Latent Diffusion Model
  9. Discussion with Existing Methods
  10. Experiments
  11. Experimental Setup
  12. Main Results
  13. Further results on MRI
  14. Ablation Studies
  15. Conclusion
  16. ...and 43 more sections

Key Result

Proposition 1

Suppose ${\mathbf{x}}_{t_1}$ is sampled from the time-marginal $p({\mathbf{x}}_{t_1}\mid{\mathbf{y}})$, then satisfies the time-marginal $p({\mathbf{x}}_{t_2}\mid {\mathbf{y}})$.

Figures (28)

  • Figure 1: Overview of Decoupled Annealing Posterior Sampling (DAPS). Our method provides a flexible and effective framework for solving inverse problems through a decoupled posterior sampling process. In (a)(b), we present DAPS visual results on FFHQ and ImageNet at a resolution of 256, and in (c), on natural images at a resolution of 768. In (d), we display DAPS results on compressed sensing multi-coil MRI (CS-MRI). DAPS effectively addresses nonlinear inverse problems as well as medical imaging MRI challenges. Additionally, DAPS can be enhanced using large-scale latent diffusion models (LDMs) rombach2022high, as shown in (c).
  • Figure 2: An illustration of our method on phase retrieval. Given a noisy sample ${\mathbf{x}}_t$, we first solve the reverse diffusion process to obtain $\hat{{\mathbf{x}}}_0( {\mathbf{x}}_t)$. Using this, we construct an approximation for $p({\mathbf{x}}_0\mid{\mathbf{x}}_t,{\mathbf{y}})$. Next, we sample ${\mathbf{x}}_{0\mid y} \sim p({\mathbf{x}}_0\mid{\mathbf{x}}_t,{\mathbf{y}})$ with multiple steps of Langevin dynamics, then perturb it with the forward diffusion process to obtain a sample with slightly less noise. We repeat this process until the noise reduces to zero.
  • Figure 3: Probabilistic graphical models of the decoupled noise annealing process in (a) pixel and (b) latent spaces respectively. Here ${\mathbf{x}}_{t}\sim p({\mathbf{x}}_{t};\sigma_t)$, and ${\mathbf{y}}$ is the observed measurement. We factorize $p({\mathbf{x}}_{t}\mid{\mathbf{y}})$ based on these probabilistic graphs.
  • Figure 4: DAPS vs. DPS on 2D synthetic data. Consecutive sampling steps are close to each other for DPS but not for DAPS. Here DAPS approximates the posterior better than DPS.
  • Figure 5: Sample diversity. We present several diverse samples generated by the DAPS under two sparse measurements whose posterior distributions contain multiple modes. DAPS produces a variety of samples with distinct features, including differences in expression, wearings, and hairstyles.
  • ...and 23 more figures

Theorems & Definitions (6)

  • Proposition 1
  • Proposition 2
  • Proposition 3: Restated
  • proof
  • Proposition 4: Restated
  • proof