Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing

2024-07-01

Computer Vision and Pattern Recognition

AI Summary

Diffusion models often suffer from error accumulation and high reconstruction failure rates when solving strongly nonlinear Bayesian inverse problems—such as phase retrieval—due to conventional stepwise denoising. To address this, we propose Decoupled Annealed Posterior Sampling (DAPS), the first noise-annealing mechanism for posterior sampling in diffusion frameworks. DAPS decouples sampling steps and introduces variable-step annealing, progressive marginal distribution constraints, and posterior consistency regularization to ensure that the time-marginal distributions converge to the true posterior. This design significantly enhances exploration of the solution space and improves sampling robustness. Experiments demonstrate that DAPS consistently improves sample quality and convergence stability across multiple image restoration tasks. Notably, on phase retrieval, it achieves a substantial increase in reconstruction success rate compared to baseline diffusion samplers.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Calibration & Uncertainty Quantification

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs

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.

Problem

• Solving Bayesian inverse problems with diffusion models
• Correcting errors in nonlinear inverse problem reconstruction
• Improving sample quality and stability in image restoration

Innovation

• Decoupled noise annealing process
• Diffusion sampling trajectory variation
• True posterior marginal annealing