Provable Diffusion Posterior Sampling for Bayesian Inversion (2512.08022v1)

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.

• The paper introduces a provably convergent diffusion framework for sampling Bayesian posteriors in inverse problems using score-based models and Monte Carlo estimation.
• It employs a warm-start Langevin initialization and time-reversal diffusion process to balance error sources while ensuring non-asymptotic convergence guarantees.
• Empirical evaluations on deblurring tasks show improved PSNR/SSIM and robust uncertainty quantification, even with multimodal posteriors and prior mismatches.

Provable Diffusion Posterior Sampling for Bayesian Inversion

Overview and Motivation

The paper "Provable Diffusion Posterior Sampling for Bayesian Inversion" (2512.08022) introduces a diffusion-based framework for sampling from complex Bayesian posterior distributions in statistical inverse problems. The approach leverages recent advances in score-based diffusion models and Monte Carlo estimation, addressing deficiencies of traditional MCMC schemes when confronted with multi-modal or non-log-concave posteriors typical in modern applications such as medical imaging, physics, and machine learning.

The contribution centers on a plug-and-play methodology: a learned score function models the unconditional prior, independent of the likelihood, while posterior sampling employs a Monte Carlo estimator informed by Langevin dynamics. Critically, the proposed method is supported by non-asymptotic convergence guarantees, even in multimodal settings — a limitation for previous heuristics.

Formulation and Methodology

Sampling from the posterior $p_{X_0|Y}(x_0|y)$ is achieved by constructing a probability flow from a tractable terminal distribution back to the target posterior (Fig. 1). This is realized via:

• Forward and reverse diffusion processes: The Ornstein-Uhlenbeck SDE generates a smooth interpolant from prior to a Gaussian, while the time-reversed SDE enables transformation from a terminal distribution to the posterior.
• Posterior score estimation: Applying a conditional Tweedie’s formula, the posterior score $\nabla \log q_t(x|y)$ is recast as a function of a conditional mean (posterior denoiser), leveraging knowledge of the forward kernel and prior score.
• Warm-start initialization: Rather than approximating the terminal distribution by a standard Gaussian (which is only valid for large $T$), a Langevin dynamics is used to sample directly from the precise terminal density $q_T(\cdot|y)$, improving accuracy for moderate $T$.

The algorithm integrates these components: warm-start via Langevin, time-reversal sampling with the Monte Carlo score estimator, and an early-stopping procedure to balance estimation error and sampling bias.

(Figure 1)

Figure 1: Schematic of the dual log-concavity regimes: for small $T$ the denoising density is log-concave (enabling posterior score estimation); for large $T$ the terminal density is log-Sobolev (enabling warm-start convergence). The method selects $T$ to lie in the intersection for guaranteed performance.

Theoretical Guarantees

A hallmark of this work is a complete non-asymptotic convergence analysis in 2-Wasserstein distance, covering the case where the posterior is multi-modal and only semi-log-concave. The error analysis decomposes the total approximation error into three distinct sources:

1. Posterior score estimation error: Quantifies the discrepancy between the true posterior score and its Monte Carlo approximation; depends critically on the prior score matching error and the problem condition number $\kappa_y$.
2. Warm-start error: Arises from approximating the true terminal posterior with samples generated by Langevin dynamics using the estimated score.
3. Early-stopping error: Due to termination of the reverse process short of $t = 0$ to avoid divergence as the variance shrinks.

The analysis demonstrates existence of a regime for terminal time $T$ where, simultaneously, the posterior denoising density is log-concave (ensuring convergence of the score estimator) and the terminal density satisfies a log-Sobolev inequality (needed for the outer Langevin sampler in warm-start). Hyperparameters such as simulation horizon, number of samples, and $T$ are optimally selected by balancing these competing error terms.

Empirical Evaluation

Experimental results span several challenging inverse problems: Gaussian deblurring, motion deblurring (nonlinear forward models), and cross-domain prior mismatch. Models are trained on FFHQ human faces with evaluation including AFHQ animal faces to test robustness.

Figure 2: Gaussian deblurring case studies comparing naive input (blurred), TV, DPS, and PDPS. Bottom: mean absolute error and standard deviation maps (across 24 samples) display the method’s uncertainty quantification capability.

Compared to TV and DPS, the PDPS algorithm substantially improves both average PSNR/SSIM (Table below) and qualitative artifacts, particularly restoring fine details under nonlinear measurement noise.

Method	Gaussian Deblur (PSNR/SSIM)	Motion Deblur (PSNR/SSIM)	Nonlinear Deblur (PSNR/SSIM)
TV	23.95 / 0.81	24.65 / 0.80	19.70 / 0.53
DPS	24.15 / 0.81	26.66 / 0.88	20.93 / 0.68
PDPS	26.42 / 0.87	28.86 / 0.92	28.44 / 0.91

Uncertainty quantification is addressed via pixel-wise standard deviation maps, showing that regions with higher reconstruction error correspond to elevated sampling variance — a desirable property missing from classical TV estimators.

Figure 3: Nonlinear deblurring case: PDPS recovers detailed structure with less artifact than DPS and TV; uncertainty maps accurately reflect reconstruction ambiguity.

Figure 4: Out-of-domain robustness: Deblurring AFHQ animal faces using a human face prior. PDPS manages to recover plausible details (e.g., whiskers) even when the prior is severely mismatched, outperforming heuristics.

Ablation studies on the terminal time $T$ reveal a sharp performance peak at intermediate $T$: small $T$ impedes warm-start convergence, large $T$ deteriorates posterior score estimation due to loss of log-concavity.

Figure 5: Ablation paper on terminal time $T$: Optimal PSNR/SSIM is achieved for $T$ in the intermediate regime ($\approx 0.05$ to $1.0$), confirming the theoretical requirement for dual regime convergence.

Implications and Future Directions

Practically, the PDPS method enables reliable posterior sampling in ill-posed inverse problems where multimodality and prior-likelihood mismatch are the norm. The plug-and-play architecture, requiring only prior score learning independent of the forward operator, offers wide applicability without retraining. The non-asymptotic theoretical analysis bestows confidence in deploying these samplers to new scientific and engineering domains.

The framework's rigor opens several research avenues, including:

• Extensions to infinite-dimensional inverse problems (e.g., PDE-constrained inference) via operator learning;
• Derivative-free Bayesian inference where gradients of the likelihood are unavailable;
• Temporal or sequential Bayesian inverse problems (data assimilation);
• Convergence guarantees for alternative multi-modal posterior constructions.

Conclusion

This research establishes a principled, provably convergent approach to Bayesian posterior sampling via diffusion models, supporting multi-modal and non-log-concave posteriors. By replacing heuristic approximations with a Monte Carlo score estimator and constructing a theoretically sound warm-start mechanism, the method achieves both robustness and uncertainty quantification. Theoretical and empirical evidence supports its superiority, with immediate impact for a wide range of applications in inverse problems and scientific machine learning.