Diffusion Prior-Based Amortized Variational Inference for Noisy Inverse Problems (2407.16125v1)

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.

• The paper introduces a diffusion prior-based AVI method that learns a direct mapping from measurements to implicit posterior distributions, reducing computational costs.
• It leverages an integral KL divergence and a Perturbed Posterior Bridge to ensure robust convergence and improved generalization in noisy environments.
• The approach outperforms state-of-the-art techniques in image restoration tasks, even under unknown noise scales, highlighting both theoretical and practical benefits.

Overview of Diffusion Prior-Based Amortized Variational Inference for Noisy Inverse Problems

The paper, "Diffusion Prior-Based Amortized Variational Inference for Noisy Inverse Problems" (DAVI), introduces a novel method for addressing noisy inverse problems by leveraging diffusion models within the framework of amortized variational inference (AVI). This method aims to mitigate the high computational costs and generalization limitations commonly associated with existing approaches to inverse problems, which typically involve iterative posterior samplings and measurement-specific optimization.

Key Contributions and Methodology

1. Amortized Variational Inference (AVI): The authors propose using AVI to directly learn a mapping from measurements to implicit posterior distributions of clean data. This enables single-step sampling from the posterior, which is both computationally efficient and capable of generalized inference on unseen measurements.
2. Integral KL Divergence (IKL): To ensure stable convergence and improve the robustness of their approximation, the loss function incorporates an integral of the Kullback-Leibler (KL) divergence across time points. This approach reduces the risk associated with traditional KL divergence methods where divergences might diverge due to disjoint supports.
3. Perturbed Posterior Bridge (PPB): A novel component introduced in the paper, PPB enhances the generalization capability by interpolating between clean and noisy measurements, effectively creating a bridge supported by strategically induced perturbations.
4. Alternating Optimization Framework: By alternating the training of two neural networks—one for the implicit distribution and one for the score function—the method directly minimizes the described loss function. This allows for efficient training and inference with reduced computational costs.

Results

The DAVI method was extensively tested on several image restoration tasks, including Gaussian deblurring, super-resolution, and box inpainting. The notable performance in terms of PSNR, LPIPS, and FID scores over recent state-of-the-art methods like DDRM, DDNM$^+$, and RED-diff demonstrates its superior capability in achieving higher-quality image restorations with fewer neural network evaluations.

The authors also explored robustness to unknown noise scales, showcasing that DAVI performs well even when the measurement noise level is not provided, a significant advantage over methods requiring precise noise information.

Implications and Future Work

The practical implications of the DAVI framework are significant, considering its potential for real-time applications and deployment on resource-constrained devices due to its efficiency. Theoretical implications include advancing the broader adoption and enhancement of AVI in image processing and other domains necessitating inverse problem solutions.

This work opens avenues for further research in extending the amortized approach to other types of inverse problems beyond image restoration, exploring different architecture choices for the implicit distributions, and studying the interplay between diffusion models and alternative probabilistic inference methodologies.

In conclusion, the authors present a well-substantiated advancement in the field, grounded in the strategic utilization of diffusion models within an amortized framework, offering compelling advantages in both efficacy and computational efficiency for addressing noisy inverse problems.