Diffusion Path Samplers via Sequential Monte Carlo

Technical Summary: Diffusion Path Samplers via Sequential Monte Carlo

Problem Statement
Sampling from target distributions $\pi \propto \exp(-V_\pi)$ known only up to a normalising constant is a fundamental challenge in Bayesian inference, machine learning, and statistical physics. Direct sampling from complex, high-dimensional, or multimodal targets is often intractable. While sequential methods bridging a simple base distribution $\nu$ and the target $\pi$ offer a solution, standard geometric paths have been shown to suffer from "mass teleportation" and degraded log-Sobolev constants, hindering efficiency. Conversely, diffusion paths, popularized by generative modeling, offer better mathematical regularity (uniformly bounded functional inequalities and finite action) but require efficient score estimation along the path, which is typically intractable.

Methodology
The authors propose Diffusion Path Sequential Monte Carlo (DPSMC), a framework that combines diffusion paths with Sequential Monte Carlo (SMC) to estimate scores and sample from target distributions.

1. Diffusion Path Construction:
   The method utilizes a continuous-time family of distributions $(\mu_t)_{t \in [0,1]}$ interpolating between a base Gaussian $\nu$ and the target $\pi$. This path is defined via convolution:
   $$\mu_t(x) = \frac{1}{\sqrt{1-\lambda_t}} \nu\left(\frac{x}{\sqrt{1-\lambda_t}}\right) * \frac{1}{\sqrt{\lambda_t}} \pi\left(\frac{x}{\sqrt{\lambda_t}}\right)$$
   where $\lambda_t$ is a monotonically increasing schedule. The authors analyze three dynamics inducing this path:

   • Exact Diffusion SDE: Matches path marginals exactly for Gaussian bases (related to time-reversed Ornstein-Uhlenbeck processes).
   • Diffusion-Annealed Langevin Dynamics (DALD): Approximates the path for arbitrary bases with a controllable error parameter $\epsilon$.
   • Stochastic Interpolants (SI): A dynamics formulation that allows exact induction of the path with finite time horizons, avoiding critical hyperparameters like the time horizon $T$ required by OU processes.

2. Score Estimation via SMC:
   Since the score function $\nabla \log \mu_t$ is intractable, DPSMC estimates it using an SMC sampler that evolves auxiliary variables $Y$ from conditional (posterior) distributions $\varrho_{t,x}$ along the path. The authors derive score identities (Denoising Score Identity and Target Score Identity) and express the score as an expectation over these auxiliary variables. An SMC sampler propagates particles through these conditional distributions, updating importance weights to estimate the score at each step of the main diffusion trajectory.

3. Control Variate Schedules:
   To minimize the variance of score estimates, the authors introduce Control Variate (CV) schedules. They define a class of test functions combining the denoising and target score identities via a matrix $A_t$.

   • They derive optimal scalar ($\alpha_t^*$) and matrix ($A_t^*$) schedules that minimize the expected variance of the estimator.
   • The optimal matrix schedule, termed MMCVSI (Marginalised Matrix Control-Variate Score Identity), accounts for pairwise covariances of score components, offering advantages for anisotropic targets.
   • These schedules are estimated online using the auxiliary variables, incurring minimal computational overhead.

4. Algorithmic Implementation:
   The algorithm (Algorithm 1) alternates between:

   • Main Trajectory: Propagating samples $X_k$ using the estimated score (via Euler-Maruyama or Exponential Integrator).
   • Auxiliary SMC: Running an SMC sampler to target the conditional distributions $\varrho_{t_k, X_k}$, generating weighted particles to update the score estimate.
   • Initialization: For non-Gaussian starts or specific dynamics, a "Langevin-within-Langevin" (LL) initialisation is used to warm up the process.

Key Contributions

• Novel Framework: Development of a highly parallelisable DPSMC sampler that uses SMC for score estimation along diffusion paths, avoiding the sequential bottlenecks of MCMC-based score estimation.
• Variance Reduction: Introduction of principled, practical control variate schedules (MCVSI and MMCVSI) that minimize score estimation variance without requiring additional energy evaluations.
• Theoretical Guarantees:
  • A new bound on score estimation error extending previous results (He et al., 2024), showing that even with imperfect auxiliary sampling in early steps, the score estimator remains reliable.
  • Convergence guarantees for the KL divergence between the target and the final distribution for OU, SI, and DALD dynamics.
• Path Design: Derivation of an optimal cosine schedule ($\lambda_t = \sin^2(\pi t/2)$) that minimizes the upper bound of the action functional, ensuring efficient path traversal.

Results
The authors evaluate DPSMC on synthetic benchmarks (GMM40, Rings, Funnel) and real-world datasets (Bayesian logistic regression on Ionosphere and Sonar).

• Performance: DPSMC consistently outperforms or matches state-of-the-art samplers (AIS, SMC, RDMC, SLIPS) across all benchmarks.
• Efficiency: The method achieves speed-ups of several orders of magnitude compared to other samplers due to its ability to batch energy evaluations and leverage parallel computing (GPUs).
• Anisotropy: The matrix control variate (MMCVSI) demonstrates significant gains on the "Funnel" distribution, effectively handling anisotropy where scalar schedules or other methods struggle.
• Normalising Constants: The framework is adapted to estimate normalising constants using a nested SMC approach (RDSMC), providing improved estimates compared to existing methods while maintaining sample quality.

Significance and Claims
The paper claims that DPSMC provides a robust, theoretically grounded, and computationally efficient alternative for sampling from complex distributions. By leveraging the mathematical regularity of diffusion paths and the parallelisability of SMC, the method overcomes the limitations of geometric paths and the sequential bottlenecks of MCMC-based score estimation. The authors emphasize that their approach is "training-free" (unlike neural samplers) and provides unbiased estimates for path marginal densities, allowing for flexible integration into broader sampling frameworks. The work highlights the importance of batched energy evaluations in the era of parallel computing, arguing that methods capable of exploiting this axis can achieve superior wall-clock times and sample quality.