Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stochastic Transport Maps in Diffusion Models and Sampling

Xicheng Zhang

TL;DR

The work develops a rigorous framework for constructing (stochastic) transport maps between distributions using diffusion processes, showing that the time-marginal of sums of independent diffusions satisfies a Fokker-Planck equation and leveraging the Ambrosio–Figalli–Trevisan superposition principle to obtain SDE solutions. It introduces a unified diffusion-based generative model framework with deterministic (ODE) and stochastic (SDE) transports, providing explicit drift representations and three particle-based schemes with convergence guarantees. The approach enables direct, scalable sampling from complex targets (including high-dimensional funnel distributions and Gaussian mixtures) by solving forward ODE/SDE dynamics or their particle approximations, bridging theory and practical algorithms. The results offer new tools for high-dimensional generative modeling and sampling, with rigorous performance guarantees for particle approximations and clear pathways for implementation. Overall, the paper contributes a principled link between transport maps and diffusion-based sampling, yielding both theoretical insights and computationally effective methods.

Abstract

In this work, we present a theoretical and computational framework for constructing stochastic transport maps between probability distributions using diffusion processes. We begin by proving that the time-marginal distribution of the sum of two independent diffusion processes satisfies a Fokker-Planck equation. Building on this result and applying Ambrosio-Figalli-Trevisan's superposition principle, we establish the existence and uniqueness of solutions to the associated stochastic differential equation (SDE). Leveraging these theoretical foundations, we develop a method to construct (stochastic) transport maps between arbitrary probability distributions using dynamical ordinary differential equations (ODEs) and SDEs. Furthermore, we introduce a unified framework that generalizes and extends a broad class of diffusion-based generative models and sampling techniques. Finally, we analyze the convergence properties of particle approximations for the SDEs underlying our framework, providing theoretical guarantees for their practical implementation. This work bridges theoretical insights with practical applications, offering new tools for generative modeling and sampling in high-dimensional spaces.

Stochastic Transport Maps in Diffusion Models and Sampling

TL;DR

The work develops a rigorous framework for constructing (stochastic) transport maps between distributions using diffusion processes, showing that the time-marginal of sums of independent diffusions satisfies a Fokker-Planck equation and leveraging the Ambrosio–Figalli–Trevisan superposition principle to obtain SDE solutions. It introduces a unified diffusion-based generative model framework with deterministic (ODE) and stochastic (SDE) transports, providing explicit drift representations and three particle-based schemes with convergence guarantees. The approach enables direct, scalable sampling from complex targets (including high-dimensional funnel distributions and Gaussian mixtures) by solving forward ODE/SDE dynamics or their particle approximations, bridging theory and practical algorithms. The results offer new tools for high-dimensional generative modeling and sampling, with rigorous performance guarantees for particle approximations and clear pathways for implementation. Overall, the paper contributes a principled link between transport maps and diffusion-based sampling, yielding both theoretical insights and computationally effective methods.

Abstract

In this work, we present a theoretical and computational framework for constructing stochastic transport maps between probability distributions using diffusion processes. We begin by proving that the time-marginal distribution of the sum of two independent diffusion processes satisfies a Fokker-Planck equation. Building on this result and applying Ambrosio-Figalli-Trevisan's superposition principle, we establish the existence and uniqueness of solutions to the associated stochastic differential equation (SDE). Leveraging these theoretical foundations, we develop a method to construct (stochastic) transport maps between arbitrary probability distributions using dynamical ordinary differential equations (ODEs) and SDEs. Furthermore, we introduce a unified framework that generalizes and extends a broad class of diffusion-based generative models and sampling techniques. Finally, we analyze the convergence properties of particle approximations for the SDEs underlying our framework, providing theoretical guarantees for their practical implementation. This work bridges theoretical insights with practical applications, offering new tools for generative modeling and sampling in high-dimensional spaces.

Paper Structure

This paper contains 18 sections, 21 theorems, 263 equations, 6 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Transport Maps
  3. Diffusion Models
  4. Main Results
  5. Related Results and Our Contributions
  6. Organization of this Paper
  7. Notations
  8. The Sum of Two Independent Diffusion Processes
  9. Diffusion Generative Models: ODE and SDE Approaches
  10. Forward ODE Approach: Starting from Absolutely Continuous Distributions
  11. Forward SDE Approach: Starting from Any Distribution
  12. Particle Approximations of Stochastic Transport Maps
  13. SDE: Zero-Order Particle Approximation
  14. SDE: First-Order Particle Approximation
  15. Application to Sampling from Known Distributions
  16. ...and 3 more sections

Key Result

Theorem 1.1

Suppose that for some $p \geqslant 1$, $\mathbb{E}|\xi_0|^p + \mathbb{E}|\eta|^p < \infty$, and for some constants $K, \kappa > 0$, the density $\rho_0\in C^1$ satisfies Then, for each initial point $x_0 \in \mathbb{R}^d$, the ODE ODE-0 admits a unique solution $(X_t(x_0))_{t \in [0, 1)}$ with $X_0 = x_0$. Moreover, the following properties hold:

Figures (6)

  • Figure 1: Deterministic transport map
  • Figure 2: Stochastic transport map
  • Figure 3: Sampling of high dimensional funnel distribution
  • Figure 4: Sampling of 20D, 50D, 100D-funnel distributions
  • Figure 5: Sampling of Gaussian distribution with degenerate covariance matrix
  • ...and 1 more figures

Theorems & Definitions (56)

  • Theorem 1.1
  • Remark 1.2
  • Remark 1.3
  • Theorem 1.4
  • Remark 1.5
  • Remark 1.6
  • Theorem 2.1
  • proof
  • Theorem 2.2
  • proof
  • ...and 46 more