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

2025-02-01

AI Summary

This work addresses efficient sampling from high-dimensional, unnormalized target distributions. We propose a deterministic transport framework operating in finite time: an intermediate density flow is constructed via stochastic interpolation and modeled as a coupled forward-backward stochastic differential equation (FBSDE) system—the first application of FBSDEs to solve Hamilton–Jacobi–Bellman-type partial differential equations (PDEs)—enabling exact, finite-step density matching from a Gaussian initial distribution to the target. Unlike conventional diffusion models reliant on asymptotic iterative procedures, our approach bypasses such limitations by parameterizing the PDE solution via neural networks, yielding a differentiable and trainable transport map. Empirically, on high-dimensional multimodal and strongly correlated distributions, our method achieves 3–5× higher sampling efficiency than SGLD and DDPM, while significantly mitigating mode collapse and improving mode coverage.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision
• Machine Learning: Learning with Manifolds

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• High-Dimensional Sampling
• Complex Distributions
• Gaussian Transformation

Innovation

• Diffusion Processes
• Machine Learning
• Stochastic Interpolation