Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

U-DAVI: Uncertainty-Aware Diffusion-Prior-Based Amortized Variational Inference for Image Reconstruction

Ayush Varshney, Katherine L. Bouman, Berthy T. Feng

TL;DR

This work extends the amortized variational inference framework by injecting spatially adaptive perturbations to measurements during training, guided by uncertainty estimates, to emphasize learning in the most uncertain regions.

Abstract

Ill-posed imaging inverse problems remain challenging due to the ambiguity in mapping degraded observations to clean images. Diffusion-based generative priors have recently shown promise, but typically rely on computationally intensive iterative sampling or per-instance optimization. Amortized variational inference frameworks address this inefficiency by learning a direct mapping from measurements to posteriors, enabling fast posterior sampling without requiring the optimization of a new posterior for every new set of measurements. However, they still struggle to reconstruct fine details and complex textures. To address this, we extend the amortized framework by injecting spatially adaptive perturbations to measurements during training, guided by uncertainty estimates, to emphasize learning in the most uncertain regions. Experiments on deblurring and super-resolution demonstrate that our method achieves superior or competitive performance to previous diffusion-based approaches, delivering more realistic reconstructions without the computational cost of iterative refinement.

U-DAVI: Uncertainty-Aware Diffusion-Prior-Based Amortized Variational Inference for Image Reconstruction

TL;DR

This work extends the amortized variational inference framework by injecting spatially adaptive perturbations to measurements during training, guided by uncertainty estimates, to emphasize learning in the most uncertain regions.

Abstract

Ill-posed imaging inverse problems remain challenging due to the ambiguity in mapping degraded observations to clean images. Diffusion-based generative priors have recently shown promise, but typically rely on computationally intensive iterative sampling or per-instance optimization. Amortized variational inference frameworks address this inefficiency by learning a direct mapping from measurements to posteriors, enabling fast posterior sampling without requiring the optimization of a new posterior for every new set of measurements. However, they still struggle to reconstruct fine details and complex textures. To address this, we extend the amortized framework by injecting spatially adaptive perturbations to measurements during training, guided by uncertainty estimates, to emphasize learning in the most uncertain regions. Experiments on deblurring and super-resolution demonstrate that our method achieves superior or competitive performance to previous diffusion-based approaches, delivering more realistic reconstructions without the computational cost of iterative refinement.
Paper Structure (11 sections, 13 equations, 3 figures, 2 tables)

This paper contains 11 sections, 13 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Background: The DAVI Framework
  4. Uncertainty Estimation via Temporal Inconsistency
  5. Uncertainty-Aware Perturbed Posterior Bridge
  6. Single-Step Inference
  7. Experiments
  8. Experimental Setup
  9. Discussion
  10. Conclusion
  11. Acknowledgments

Figures (3)

  • Figure 1: Uncertainty-Aware Diffusion-Prior-Based Amortized Variational Inference (U-DAVI) during (a) training and (b) inference. Notation: $x_0$ ground truth; $y$ noisy measurement; $z\!\sim\!\mathcal{N}(0,I)$; $u$ uncertainty estimate; $y_a$ uncertainty-aware bridge sample; $\hat{x}_0$ reconstruction; $\bar{x}$ persistent reconstruction memory. During training, the uncertainty map $u$ emphasizes ambiguous regions (e.g., around glasses). At inference, U-DAVI reconstructs $\hat{x}_0$ in a single pass and typically recovers fine details better than the original DAVI approach lee2024davi.
  • Figure 2: Qualitative results on FFHQ for (a) deblurring and (b) super-resolution. In both tasks, U-DAVI restores sharper details such as backgrounds, wrinkles, fingers, shadows, reflections, fabrics, and earrings compared to the iterative RED-Diff and amortized DAVI baselines.
  • Figure 3: Histograms of U-DAVI's improvements on deblurring and super-resolution for FFHQ (1K images) and CelebA-HQ (30K images) across 100 inference seeds. We plot $\Delta\text{PSNR}=\text{PSNR}_{\text{U-DAVI}}-\text{PSNR}_{\text{DAVI}}$ and $\Delta\text{FID}=\text{FID}_{\text{DAVI}}-\text{FID}_{\text{U-DAVI}}$ such that positive values indicate U-DAVI outperforms DAVI (higher PSNR / lower FID).