| Contents | Wasserstein gradient flows offer a unifying view of many diffusion equations: in particular, the Fokker–Planck equation describes the steepest descent of an entropy-type free-energy functional on the space of probability measures equipped with the Wasserstein-2 metric. The JKO scheme, an implicit time discretization of this flow, therefore gives a principled way to simulate such diffusions, but each step requires solving an optimization problem over probability measures that is computationally demanding, especially in high dimension. This paper proposes a scalable parametric approach aimed at machine learning applications. Exploiting Brenier's theorem, each JKO step is represented by the gradient of an input-convex neural network that pushes the previous measure forward to the next one, so the Wasserstein-2 proximity term reduces to an expectation over samples and no separate optimal transport problem needs to be solved. The resulting sequence of problems is trained with stochastic gradient descent. Unlike earlier methods, the approach requires neither domain discretization nor particle simulation, and since every step is an explicit map, one can both draw samples from the measure at each time step and evaluate its probability density. The method is validated on Fokker–Planck diffusions and applied to sampling from unnormalized densities and to nonlinear filtering.
|
