Diffusion Prior-Based Amortized Variational Inference for Noisy Inverse Problems

TL;DR

This work tackles the challenge of noisy inverse problems by replacing per-sample optimization with a diffusion-prior, amortized variational inference framework. DAVI learns an implicit posterior mapping ${q_{\phi}}({\mathbf{x}_0}|{\mathbf{y}})$ from measurements to clean data, enabling single-step posterior sampling ${\hat{\mathbf{x}} = {\mathcal{I}}_{\phi}({\mathbf{y}} + h{\mathbf{z}})}$ and generalization to unseen measurements. It introduces an integral KL objective (IKL) and a Perturbed Posterior Bridge (PPB) to stabilize training and improve generalization, along with an alternating optimization scheme involving an implicit score ${s_\psi}$ and a pre-trained diffusion score ${s_\theta}$. Empirically, DAVI achieves state-of-the-art performance on Gaussian deblurring, 4× super-resolution, and box inpainting across FFHQ and ImageNet, delivering improved FID/LPIPS while maintaining competitive PSNR and delivering fast inference (≈0.04 s/image). The approach has practical implications for scalable, real-time inversion in imaging and related domains, with ethical considerations regarding potential privacy implications of restored imagery.

Abstract

Recent studies on inverse problems have proposed posterior samplers that leverage the pre-trained diffusion models as powerful priors. These attempts have paved the way for using diffusion models in a wide range of inverse problems. However, the existing methods entail computationally demanding iterative sampling procedures and optimize a separate solution for each measurement, which leads to limited scalability and lack of generalization capability across unseen samples. To address these limitations, we propose a novel approach, Diffusion prior-based Amortized Variational Inference (DAVI) that solves inverse problems with a diffusion prior from an amortized variational inference perspective. Specifically, instead of separate measurement-wise optimization, our amortized inference learns a function that directly maps measurements to the implicit posterior distributions of corresponding clean data, enabling a single-step posterior sampling even for unseen measurements. Extensive experiments on image restoration tasks, e.g., Gaussian deblur, 4$\times$ super-resolution, and box inpainting with two benchmark datasets, demonstrate our approach's superior performance over strong baselines. Code is available at https://github.com/mlvlab/DAVI.

Figures (19)

• Figure 1: Representative results of Diffusion prior-based Amortized Variational Inference (DAVI). The top row demonstrates the qualitative comparison between our method and baselines. The bottom two rows showcase that DAVI provides robust solutions with fine-grained details across various image restoration tasks, achieved with a single neural network evaluation.
• Figure 2: Illustration of Diffusion prior-based Amortized Variational Inference (DAVI). Our proposed method employs an alternative optimization procedure between the score function $\mathbf{s}_\psi$ and the neural network $\mathcal{I}_\phi$ to minimize the KL divergence between the implicit posterior distribution $q_\phi(\mathbf{x}_0|\mathbf{y})$ and the true posterior distribution $p(\mathbf{x}_0|\mathbf{y})$, where the true posterior distribution is approximated by the likelihood $p(\mathbf{y}|\mathbf{x}_0)$ and the diffusion prior $\mathbf{s}_\theta$.
• Figure 3: Visualization of integral KL divergence. IKL loss perturbs the implicit distribution into a smoother distribution using forward SDE to alleviate the disjoint support problem (see $t=t_1, t=t_2$, and $t_1 < t_2$). The gradient of IKL loss, $\nabla_\phi \mathcal{L}_{IKL}$, updates the parameters of the implicit distribution in a direction that minimizes the $\vartriangle \mathbf{s}_{\psi, \theta}$, which leads to minimizing the discrepancy between $q_\phi(\mathbf{x}_0|\mathbf{y})$ and $p(\mathbf{x}_0)$.
• Figure 4: Perturbed Posterior Bridge (PPB). (A) shows sampling distributions of $a$. (B) and (C) illustrate the PPB between two 1D samples, e.g., $\mathbf{x}=0$ and $\mathbf{y}=1$, with different perturbation schedule $\bar{\sigma}_a$. We plot the perturbation along the $y$ axis.
• Figure 5: Qualitative comparison. DAVI shows the most vivid and realistic solutions while maintaining intricate details of measurement, in contrast to baselines, which struggle to satisfy both aspects, as highlighted in red boxes.