A Stabilized Path-Space Approach to Diffusion-Based Posterior Sampling

2026-06-10

AI Summary

Existing diffusion-based posterior sampling methods struggle to ensure sampling accuracy and reliable uncertainty quantification under nonlinear operators or multimodal posteriors due to heuristic guidance approximations. This work proposes a path-space stochastic optimal control framework that formulates posterior sampling as learning a controlled stochastic process whose trajectory distribution matches the likelihood-weighted target measure. By employing a time reparameterization, the method eliminates bias arising from the unknown initial value function without requiring auxiliary training. It unifies guidance-based sampling and control learning, and incorporates importance sampling correction to yield asymptotically unbiased estimates of posterior expectations. Evaluated on multiple inverse problem benchmarks, the proposed approach significantly outperforms existing methods in sampling accuracy, robustness, and uncertainty quantification.

Abstract

Diffusion models provide expressive data-driven priors for Bayesian inverse problems, but many diffusion posterior samplers rely on heuristic guidance approximations that can fail for nonlinear operators and multimodal posteriors. In this work, we develop a stabilized path-space framework for diffusion-based posterior sampling. Starting from a base diffusion process whose terminal marginal represents the prior, we define a likelihood-weighted target measure on trajectories and cast posterior sampling as learning a controlled stochastic process whose path measure matches this target. This formulation connects diffusion posterior sampling to stochastic optimal control while preserving the Bayesian structure needed for uncertainty quantification. We introduce a time reparameterization that makes the path-space control problem well posed by removing the bias induced by the unknown initial value function, without auxiliary training. We then learn the control via a trust-region path-space optimization method with log-variance objectives. The path-space perspective also unifies our learned control approach with existing guidance-based samplers, quantifies the sampling error induced by approximate controls, and yields importance sampling corrections for asymptotically exact posterior expectations. We evaluate the proposed framework on a suite of benchmark inverse problems with analytically characterized or high-quality reference posteriors, enabling principled assessment of sampling accuracy and uncertainty quantification. These experiments provide insight into the behavior of diffusion-based posterior samplers and demonstrate improved accuracy and robustness over leading approaches.

Problem

• diffusion models
• posterior sampling
• Bayesian inverse problems
• nonlinear operators
• multimodal posteriors

Innovation

• path-space control
• diffusion posterior sampling
• stochastic optimal control
• time reparameterization
• uncertainty quantification