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

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.

• The paper introduces a unified path-space framework that formulates diffusion-based sampling as a measure-matching problem.
• It employs a time reparameterization trick and trust-region optimization to stabilize control learning and correct posterior bias.
• Empirical results show significant error reduction and robust performance across linear and nonlinear Bayesian inverse problems.

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

Introduction and Motivation

The challenge of high-dimensional Bayesian inverse problems (BIPs) is central in scientific and engineering applications, necessitating accurate posterior sampling for robust uncertainty quantification (UQ). Classical priors, such as Gaussians, impose significant constraints, failing to capture the complexity of natural signals. With the advent of diffusion models, expressive data-driven priors have become viable, but integrating these rich priors into efficient and theoretically sound posterior samplers remains a nontrivial task, particularly in regimes with nonlinear operators and multimodal posteriors.

Recent diffusion-based posterior sampling methods either modify reverse diffusion dynamics via heuristic likelihood guidance or decouple prior/data-consistency steps. However, these typically rely on strong conditional distributional assumptions (e.g., Gaussianity), leading to failure modes in nonlinear/multimodal contexts. Addressing these shortcomings, this paper develops a unified, theoretically grounded framework for diffusion-based posterior sampling via path-space control and stabilization.

Path-Space Formulation of Diffusion Posterior Sampling

The core insight of the work lies in formulating diffusion-based posterior sampling as a measure-matching problem in the path space of the diffusion process. The process begins with a base diffusion whose terminal marginal is the prior, with the posterior induced by reweighting trajectories according to observation likelihood at the terminal state. Posterior sampling then becomes equivalent to constructing a controlled diffusion—by adding a drift control—whose path law matches the likelihood-weighted target.

Crucially, this path-space framework connects diffusion posterior sampling with stochastic optimal control and Schrödinger bridge problems, naturally accommodating importance sampling corrections for posterior expectations and quantifying the sampling errors induced by approximate controls. Unlike traditional, myopic score-guided methods, this global variational formulation allows manipulation of the entire path space distribution, avoiding local approximations and providing rigorous error control.

Time Reparameterization and Well-Posedness

A fundamental challenge arises from the “initial value function bias”: the induced path measure under likelihood weighting generally yields an initial marginal incompatible with the reference process, preventing exact measure matching. Prior work addressed this using memoryless noise schedules or auxiliary value-function learning.

This paper introduces a simple, general time reparameterization trick (TRT): the diffusion process is extended backward to $t = -1$ with a deterministic initial state, followed by evolution with “empty” drift/diffusion, before resuming the standard SDE at $t=0$. This adjustment secures exact independence between the initial and terminal marginals, fully resolving the bias problem and establishing well-posedness without auxiliary neural networks or asymptotic approximations.

Trust-Region Path-Space Optimization

Having established a well-posed measure matching problem, posterior sampling is formulated as a stochastic optimal control task: learn a drift control such that the controlled diffusion matches the likelihood-weighted target measure in path space. Practically, this is achieved by minimizing divergences (primarily KL) between the controlled path measure and the target, using a neural network parameterization for the control.

Optimization in path space is challenging due to instability in high dimensions and for multimodal posteriors. The method employs the trust-region path-space stochastic optimization of [5], which constrains KL divergence between successive iterates, stabilizing updates. This is complemented by an off-policy log-variance divergence objective, which offers lower gradient variance and robustness. The result is a numerically stable, sample-efficient, and principled method to learn approximate posterior controls from path samples.

In this formalism, even existing heuristic score-guidance methods—such as DPS and IIGDM—are interpretable as employing implicit, approximate path-space controls. The framework offers path-space importance weights (Radon-Nikodym derivatives) to correct for biases and enables computation of asymptotically exact posterior expectations under any such control.

Theoretical Error Bounds and Importance Sampling Corrections

The paper rigorously quantifies the error introduced by control approximation. It establishes a path-space KL divergence bound between the time-$T$ distributions of the controlled process and the true posterior, proportionate to the time-integrated squared control discrepancy. More importantly, regardless of the approximation, importance sampling with properly derived path-space weights leads to asymptotically unbiased posterior expectations.

Practically, this enables algorithmic differentiation between sampling and inference: any controlled diffusion process, either heuristic or learned, can be corrected to provide unbiased posterior statistics via importance weighting, provided the pathwise Radon-Nikodym derivative is available and its variance is manageable.

Empirical Evaluation and Results

The method is evaluated on a comprehensive suite of benchmark BIPs: linear random sensing, inpainting, nonlinear X-ray tomography, and phase retrieval. All benchmarks use Gaussian mixture priors to induce multimodal posteriors. The benchmarking rigorously compares the proposed trust-region path-space sampler (TR) to leading alternatives: DPS, IIGDM, and DAPS.

On both linear-Gaussian and nonlinear problems, the TR approach demonstrates significant improvements in mean error, covariance error, maximum mean discrepancy (MMD), and central moment discrepancy (CMD) relative to all baselines. In particular, the method achieves:

• Orders of magnitude reductions in mean and covariance errors compared to heuristic methods.
• Low control error and high normalized effective sample size (NESS), indicating both accurate control learning and stable importance reweighting.
• Consistent superiority across both easy (linear) and hard (nonlinear, underdetermined) inverse problems.

The ablation analyses further confirm the stabilizing effect of the trust-region constraint. Empirical importance reweighting is shown to effectively debias approximate samplers (including DPS and IIGDM) when path-space weights are well-behaved.

While the TR method incurs higher up-front training costs (notably, additional neural network optimization), its amortized per-sample cost is competitive, especially in regimes dominated by expensive likelihood or drift evaluations or when large sample sets are required.

Implications and Future Directions

The path-space control framework fundamentally unifies and extends diffusion-based posterior sampling, offering:

• A universal perspective: both heuristic and learned controls correspond to explicit path-space measures, which can be systematically corrected and compared.
• Practical error bounds: algorithmic control error translates directly into quantifiable and correctable posterior bias.
• Theoretically principled construction: connections to stochastic optimal control, Schrödinger bridges, and measure transport theory yield robust, interpretable algorithms.
• Versatility: the method accommodates arbitrary (learned or closed-form) priors, general nonlinear likelihoods, and applies to a wide range of inverse problems.

For future developments, several extensions are naturally suggested:

• Amortized controls: By parameterizing the control as a function of not only $(x, t)$ but also observation $y$ (and forward operator $F$), amortized inference becomes feasible, significantly decreasing per-observation training costs for repeated or streaming problem instances.
• High-dimensional benchmarks: Extending empirical validation to high-dimensional, real-world BIPs (e.g., imaging) and developing statistically well-characterized benchmarks.
• Integration with advanced prior learning: Incorporating manifold-constrained or equality-constrained priors for scientific applications, leveraging recent advances in generative model training for structured distributions.
• Discretization error analysis: Developing finite-time and discretization error bounds in the context of diffusion SDE solvers.

Conclusion

This paper establishes a stabilized, theoretically grounded path-space approach for diffusion-based posterior sampling, resolving longstanding issues of bias and instability endemic to heuristic diffusion samplers. The time reparameterization trick, combined with trust-region path-space optimization, enables effective, robust, and correctable inference across a wide spectrum of inverse problems. The framework provides both a conceptual advance and a practical improvement, facilitating rigorous UQ in contemporary high-dimensional Bayesian inference scenarios.

Reference: "A Stabilized Path-Space Approach to Diffusion-Based Posterior Sampling" (2606.12710)