Unbiased Diffusion Variational Inversion via Principled Posterior Matching

2026-05-24

AI Summary

Existing score-based methods for inverse problems often suffer from mode collapse and inaccurate uncertainty quantification due to their reliance on heuristic approximations that minimize the Kullback–Leibler (KL) divergence. This work proposes the Principled Posterior Matching (PPM) framework, which, for the first time, rigorously optimizes the KL divergence via an integral formulation of the Fisher divergence, thereby eliminating ad hoc approximations and enabling unbiased variational inference. PPM unifies coverage-oriented variational inference with efficient single-step reconstruction and generalizes naturally to a broader family of divergences. Demonstrated across diverse applications—including image inpainting, super-resolution fluorescence microscopy, and black hole radio interferometric imaging—PPM achieves high-fidelity reconstructions, recovers multimodal posteriors, and provides well-calibrated uncertainty estimates.

Abstract

Existing score-based methods for inverse problems often resort to approximate minimization of the KL divergence between the inversion distribution and the Bayesian posterior. Such an approximation leads to severe mode collapse and unreliable uncertainty quantification. In this paper, we propose Principled Posterior Matching (PPM), a framework that returns to the fundamentals of variational inference, rather than using tricky approximations. Instead of relying on heuristic approximations, we rigorously formulate the exact optimization of the KL divergence via the integration of Fisher divergence. We derive a tractable, equivalent gradient form of this integral, enabling precise optimization without the biases introduced by prior approximations. Our analysis clearly reveals that the mode collapse in previous methods stems directly from this approximation gap. Supported by our theoretical solution, PPM unifies two complementary paradigms: (1) In variational inference, PPM adopts mass-covering divergences that significantly improve the inversion diversity and uncertainty quantification; (2) In amortized inference, it enables the training of an efficient reconstruction network for rapid, single-step reconstruction. Furthermore, our formulation naturally extends to a broader family of divergence measures by generalizing the integral of the Fisher divergence. We validate PPM across challenging computational imaging tasks, including inpainting, super-resolution fluorescent microscopy, and radio interferometric black-hole imaging. In all experiments, PPM achieves superior reconstruction fidelity, faithful multimodal posterior recovery, and well-calibrated uncertainty estimates, establishing a robust framework for scientific imaging.

Problem

• inverse problems
• mode collapse
• uncertainty quantification
• variational inference
• Bayesian posterior

Innovation

• Principled Posterior Matching
• Variational Inference
• Fisher Divergence
• Uncertainty Quantification
• Score-based Generative Models