A Stabilized Path-Space Approach to Diffusion-Based Posterior Sampling
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Posterior Sampling | X-ray tomography | Mean Error0.143 | 7 | |
| Posterior Sampling | Random linear sensing | Mean Error0.239 | 5 | |
| Posterior Sampling | Inpainting | Mean Error0.513 | 5 | |
| Bayesian Inference | Random linear sensing | Control Error0.0044 | 4 | |
| Bayesian Inference | Inpainting | Control Error0.0052 | 4 | |
| Posterior Sampling | Phase retrieval | Mean Error0.504 | 4 | |
| Posterior Sampling | Linear-Gaussian problems | Diffusion Drift (Train)1.5 | 4 | |
| Inverse Problem Posterior Sampling | Phase retrieval | NESS0.581 | 3 |