Stochastic Transport Maps in Diffusion Models and Sampling

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.