A Variational Perspective on Solving Inverse Problems with Diffusion Models

TL;DR

The paper tackles universal inverse-problem solving with diffusion priors by formulating a variational sampling framework (RED-diff). It derives a KL-based objective that couples a reconstruction term with a score-matching regularizer across diffusion timesteps, enabling stochastic-optimization-style sampling using off-the-shelf optimizers. A denoiser-weighting mechanism based on denoising SNR balances data fidelity and prior guidance, improving stability and performance. Empirical results on image inpainting and nonlinear inverse tasks show RED-diff achieves superior fidelity and perceptual quality with lightweight iterations compared to state-of-the-art methods, highlighting its practical impact for efficient diffusion-based restoration.

Abstract

Diffusion models have emerged as a key pillar of foundation models in visual domains. One of their critical applications is to universally solve different downstream inverse tasks via a single diffusion prior without re-training for each task. Most inverse tasks can be formulated as inferring a posterior distribution over data (e.g., a full image) given a measurement (e.g., a masked image). This is however challenging in diffusion models since the nonlinear and iterative nature of the diffusion process renders the posterior intractable. To cope with this challenge, we propose a variational approach that by design seeks to approximate the true posterior distribution. We show that our approach naturally leads to regularization by denoising diffusion process (RED-Diff) where denoisers at different timesteps concurrently impose different structural constraints over the image. To gauge the contribution of denoisers from different timesteps, we propose a weighting mechanism based on signal-to-noise-ratio (SNR). Our approach provides a new variational perspective for solving inverse problems with diffusion models, allowing us to formulate sampling as stochastic optimization, where one can simply apply off-the-shelf solvers with lightweight iterates. Our experiments for image restoration tasks such as inpainting and superresolution demonstrate the strengths of our method compared with state-of-the-art sampling-based diffusion models.

Figures (17)

• Figure 1: The schematic diagram of our proposed variational sampler (RED-diff). The forward denoising diffusion process gradually adds noise to the estimate $\mu$. The denoisers of the backward diffusion process apply score-matching regularization to the measurement matching loss. The refined estimate using optimization is then fed back to the forward process and the process repeats.
• Figure 2: Comparison of the proposed variational sampler with alternatives for inpainting representative ImageNet examples. Each sampler is tuned for the best performance.
• Figure 3: Ablation for denoiser weight tuning. Left: denoiser weight over timesteps (reversed); right: KID and PSNR vs. $\lambda$ for different monotonic functions of inverse SNR.
• Figure 4: Comparison of the proposed variational sampler with alternatives for superresolution of representative ImageNet examples. Each sampler is tuned for the best performance.
• Figure 5: LPIPS vs. PSNR (dB) trade-off for 4x superresolution.