Provable Diffusion Posterior Sampling for Bayesian Inversion

TL;DR

This work tackles Bayesian inverse problems by delivering a provably convergent diffusion-based posterior sampling method that is plug-and-play: a data-trained prior score drives Langevin dynamics to estimate the posterior score, while a warm-start enables efficient sampling from a few diffusion steps at a small terminal time. It introduces a Monte Carlo posterior-score estimator using a Restricted Gaussian Oracle to avoid heuristic likelihood approximations, and provides a complete non-asymptotic 2-Wasserstein analysis that holds even for multi-modal posteriors, with explicit error decompositions for score estimation, warm-start, and early stopping. The method is validated on linear and nonlinear imaging inversions, showing improved reconstruction quality and credible uncertainty quantification, and demonstrates robustness to prior mismatch. Together, these contributions establish a theoretically grounded, flexible framework for diffusion-based Bayesian inversion with practical guidance for hyperparameters and broad applicability across forward models.

Abstract

This paper proposes a novel diffusion-based posterior sampling method within a plug-and-play (PnP) framework. Our approach constructs a probability transport from an easy-to-sample terminal distribution to the target posterior, using a warm-start strategy to initialize the particles. To approximate the posterior score, we develop a Monte Carlo estimator in which particles are generated using Langevin dynamics, avoiding the heuristic approximations commonly used in prior work. The score governing the Langevin dynamics is learned from data, enabling the model to capture rich structural features of the underlying prior distribution. On the theoretical side, we provide non-asymptotic error bounds, showing that the method converges even for complex, multi-modal target posterior distributions. These bounds explicitly quantify the errors arising from posterior score estimation, the warm-start initialization, and the posterior sampling procedure. Our analysis further clarifies how the prior score-matching error and the condition number of the Bayesian inverse problem influence overall performance. Finally, we present numerical experiments demonstrating the effectiveness of the proposed method across a range of inverse problems.

Figures (6)

• Figure 1: An illustrative schematic of the duality of convergence. Here, $\gamma_{d}$ denotes the density of the $d$-dimensional standard Gaussian distribution. (i) The log-Sobolev inequality regime corresponds to sufficiently large time values $t\in(\underline{t},+\infty)$, where the posterior density $q_{t}(\cdot|\mathbf{y})$ satisfies log-Sobolev inequality as established in Lemma \ref{['lemma:log:sobolev']}. (ii) The log-concave RGO regime encompasses sufficiently small time values $t\in(0,\bar{t})$, where the posterior denoising density $p_{t}(\cdot|\mathbf{x},\mathbf{y})$ is log-concave, thereby ensuring convergence of the score estimator $\widehat{\mathbf{s}}_{m}^{S}(t,\mathbf{x},\mathbf{y})$ as shown in Lemma \ref{['lemma:appendix:RGO:Hessian']}. (iii) Theorem \ref{['theorem:duality']} establishes the existence of an overlap between these regimes under mild assumptions, demonstrating that there exists a terminal time $T$ such that $T$ lies within the log-Sobolev inequality regime while the interval $(0,T)$ remains entirely contained in the log-concave RGO regime.
• Figure 2: Result of motion deblurring task (face 1-2). Comparison of naive input, TV, DPS, PDPS (Ours), and original ground truth are shown. Bottom rows show mean absolute error (mean err.) and standard deviation (std. dev.) maps computed over 24 runs for all methods.
• Figure 3: Gaussian deblurring examples (face 3-4). Comparison of naive input, TV, DPS, PDPS (Ours), and original ground truth. Bottom rows show mean absolute error (mean err.) and standard deviation (std. dev.) maps computed over 24 runs for TV, DPS, and PDPS.
• Figure 4: Nonlinear deblurring examples (face 5-6). Comparison of naive input, TV, DPS, PDPS (Ours), and original ground truth. Bottom rows show mean absolute error (mean err.) and standard deviation (std. dev.) maps computed over 24 runs for TV, DPS, and PDPS.
• Figure 5: Cross-dataset deblurring results on AFHQ animal faces ('cat 1', 'lion', 'dog', 'cat 2') using an FFHQ human face prior. Top two rows: Motion Deblurring. Bottom two rows: Nonlinear Deblurring. Comparison of Naive input, TV, DPS, PDPS (Ours), and original ground truth.

Theorems & Definitions (76)

• Lemma 4.1: Conditional Tweedie's formula
• Example 4.2: Linear forward operator with Gaussian noise
• Example 4.3: General forward operator with Gaussian noise
• Example 4.4: Gaussian mixture
• Example 4.5: Gaussian convolution
• Lemma 4.6: Log-concavity of RGO
• Remark 4.7: Choice of the terminal time of the diffusion model, I
• Example 4.8: Gaussian mixture
• Example 4.9: Gaussian convolution
• Lemma 4.10: Sub-Gaussian tails of the posterior