Unleashing the Denoising Capability of Diffusion Prior for Solving Inverse Problems

2024-06-11

arXiv.org

AI Summary

This work addresses ill-posed inverse problems—including image restoration, source separation, and partial generation—by proposing a bilevel constrained optimization framework that explicitly incorporates pretrained diffusion models as structured, differentiable priors. Methodologically, it is the first to jointly formulate the diffusion prior and data consistency as a constrained optimization problem and introduces a gradient-truncated projected gradient descent algorithm to ensure both convergence stability and computational efficiency. Key contributions are: (1) moving beyond the conventional “sample-and-fine-tune” paradigm by embedding diffusion models as differentiable, plug-and-play structural priors; and (2) mitigating noise amplification and optimization oscillation in ill-posed settings via bilevel variable decoupling and gradient truncation. Extensive experiments demonstrate state-of-the-art performance on both linear and nonlinear inverse tasks. The implementation is publicly available.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Non-convex Optimization
• Constraint Satisfaction and Optimization: Constraint Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems

Abstract

The recent emergence of diffusion models has significantly advanced the precision of learnable priors, presenting innovative avenues for addressing inverse problems. Since inverse problems inherently entail maximum a posteriori estimation, previous works have endeavored to integrate diffusion priors into the optimization frameworks. However, prevailing optimization-based inverse algorithms primarily exploit the prior information within the diffusion models while neglecting their denoising capability. To bridge this gap, this work leverages the diffusion process to reframe noisy inverse problems as a two-variable constrained optimization task by introducing an auxiliary optimization variable. By employing gradient truncation, the projection gradient descent method is efficiently utilized to solve the corresponding optimization problem. The proposed algorithm, termed ProjDiff, effectively harnesses the prior information and the denoising capability of a pre-trained diffusion model within the optimization framework. Extensive experiments on the image restoration tasks and source separation and partial generation tasks demonstrate that ProjDiff exhibits superior performance across various linear and nonlinear inverse problems, highlighting its potential for practical applications. Code is available at https://github.com/weigerzan/ProjDiff/.

Problem

• Diffusion Models
• Denoising
• Complex Problem Solving

Innovation

• ProjDiff
• diffusion models
• gradient projection techniques