Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sampling in High-Dimensions using Stochastic Interpolants and Forward-Backward Stochastic Differential Equations

Anand Jerry George, Nicolas Macris

TL;DR

This work tackles the problem of drawing samples from high-dimensional distributions given unnormalized densities by constructing finite-time diffusion-based samplers that transport a Gaussian prior to a target distribution via stochastic interpolants. The authors connect time-indexed densities to a diffusion process whose marginals follow the interpolant through Hamilton–Jacobi–Bellman equations, solved using forward–backward SDEs and neural network-based solvers. They introduce half- and full-interpolant schemes, solving the associated PDEs with FBSDEs to estimate gradients and moments needed for sampling, while detaching gradients to improve efficiency. Numerical experiments on Gaussian mixtures, Neal’s Funnel, double-well landscapes, and spin-glass models demonstrate state-of-the-art capability to sample from challenging targets in finite time and to estimate normalization constants without importance sampling.

Abstract

We present a class of diffusion-based algorithms to draw samples from high-dimensional probability distributions given their unnormalized densities. Ideally, our methods can transport samples from a Gaussian distribution to a specified target distribution in finite time. Our approach relies on the stochastic interpolants framework to define a time-indexed collection of probability densities that bridge a Gaussian distribution to the target distribution. Subsequently, we derive a diffusion process that obeys the aforementioned probability density at each time instant. Obtaining such a diffusion process involves solving certain Hamilton-Jacobi-Bellman PDEs. We solve these PDEs using the theory of forward-backward stochastic differential equations (FBSDE) together with machine learning-based methods. Through numerical experiments, we demonstrate that our algorithm can effectively draw samples from distributions that conventional methods struggle to handle.

Sampling in High-Dimensions using Stochastic Interpolants and Forward-Backward Stochastic Differential Equations

TL;DR

This work tackles the problem of drawing samples from high-dimensional distributions given unnormalized densities by constructing finite-time diffusion-based samplers that transport a Gaussian prior to a target distribution via stochastic interpolants. The authors connect time-indexed densities to a diffusion process whose marginals follow the interpolant through Hamilton–Jacobi–Bellman equations, solved using forward–backward SDEs and neural network-based solvers. They introduce half- and full-interpolant schemes, solving the associated PDEs with FBSDEs to estimate gradients and moments needed for sampling, while detaching gradients to improve efficiency. Numerical experiments on Gaussian mixtures, Neal’s Funnel, double-well landscapes, and spin-glass models demonstrate state-of-the-art capability to sample from challenging targets in finite time and to estimate normalization constants without importance sampling.

Abstract

We present a class of diffusion-based algorithms to draw samples from high-dimensional probability distributions given their unnormalized densities. Ideally, our methods can transport samples from a Gaussian distribution to a specified target distribution in finite time. Our approach relies on the stochastic interpolants framework to define a time-indexed collection of probability densities that bridge a Gaussian distribution to the target distribution. Subsequently, we derive a diffusion process that obeys the aforementioned probability density at each time instant. Obtaining such a diffusion process involves solving certain Hamilton-Jacobi-Bellman PDEs. We solve these PDEs using the theory of forward-backward stochastic differential equations (FBSDE) together with machine learning-based methods. Through numerical experiments, we demonstrate that our algorithm can effectively draw samples from distributions that conventional methods struggle to handle.

Paper Structure

This paper contains 39 sections, 11 theorems, 68 equations, 14 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Our Contributions
  3. Our Techniques
  4. Related Works
  5. MCMC methods:
  6. Diffusion-based methods:
  7. Stochastic interpolants:
  8. Schrödinger bridge:
  9. Preliminaries: Relation between diffusions and PDEs
  10. Main Ideas
  11. Sampling using half interpolants
  12. Solving the FBSDE
  13. Sampling
  14. Sampling using full interpolants
  15. Numerical Results
  16. ...and 24 more sections

Key Result

Lemma 1

The probability density function of $x_t$ defined in (eqn:half_interpolant) satisfies a PDE given by where $b(t,x) = \dot g(t)\mathbb{E}\mleft[\, x^* \;\middle\vert\; x_t=x\, \mright]-r(t)\dot r(t)s(t,x)$. Equivalently, $\rho$ satisfies the following Fokker-Planck equation: with initial condition $\rho(0,\cdot) = \mathcal{N}\left(0,r^2(0)I_d\right)$.

Figures (14)

  • Figure 1: Sample trajectories of diffusion process for sampling from Gaussian mixture.
  • Figure 2: Estimates of $\log Z$ as a function of training steps along with $95\%$ confidence intervals.
  • Figure 3: LossFBSDE($t,X,m,\delta,f,h$)
  • Figure 4: LossHalfInterpolant()
  • Figure 5: SampleHalfInterpolant()
  • ...and 9 more figures

Theorems & Definitions (18)

  • Definition 1: Linear interpolants
  • Lemma 1
  • Lemma 2
  • Lemma 3
  • Lemma 4
  • Lemma 5
  • Lemma 6
  • proof
  • Lemma 7
  • proof
  • ...and 8 more