Large-Scale Wasserstein Gradient Flows

• URL: http://arxiv.org/abs/2106.00736v1
• Date: Tue, 1 Jun 2021 19:21:48 GMT
• Title: Large-Scale Wasserstein Gradient Flows
• Authors: Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Evgeny Burnaev
• Abstract summary: We introduce a scalable scheme to approximate Wasserstein gradient flows. Our approach relies on input neural networks (ICNNs) to discretize the JKO steps. As a result, we can sample from the measure at each step of the gradient diffusion and compute its density.
• Score: 84.73670288608025
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated to each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Unlike previous work, our method does not require domain discretization or particle simulation. As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate our algorithm's performance by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering.

Related papers

• Diffusion Density Estimators [0.0]
  We introduce a new, highly parallelizable method that computes log densities without the need to solve a flow. Our approach is based on estimating a path integral by Monte Carlo, in a manner identical to the simulation-free training of diffusion models.
  arXiv  Detail & Related papers  (2024-10-09T15:21:53Z)
• Score-based Diffusion Models in Function Space [140.792362459734]
  Diffusion models have recently emerged as a powerful framework for generative modeling. We introduce a mathematically rigorous framework called Denoising Diffusion Operators (DDOs) for training diffusion models in function space. We show that the corresponding discretized algorithm generates accurate samples at a fixed cost independent of the data resolution.
  arXiv  Detail & Related papers  (2023-02-14T23:50:53Z)
• Monte Carlo Neural PDE Solver for Learning PDEs via Probabilistic Representation [59.45669299295436]
  We propose a Monte Carlo PDE solver for training unsupervised neural solvers. We use the PDEs' probabilistic representation, which regards macroscopic phenomena as ensembles of random particles. Our experiments on convection-diffusion, Allen-Cahn, and Navier-Stokes equations demonstrate significant improvements in accuracy and efficiency.
  arXiv  Detail & Related papers  (2023-02-10T08:05:19Z)
• Variational Wasserstein gradient flow [9.901677207027806]
  We propose a scalable proximal gradient type algorithm for Wasserstein gradient flow. Our framework covers all the classical Wasserstein gradient flows including the heat equation and the porous medium equation.
  arXiv  Detail & Related papers  (2021-12-04T20:27:31Z)
• Sliced-Wasserstein Gradient Flows [15.048733056992855]
  Minimizing functionals in the space of probability distributions can be done with Wasserstein gradient flows. This work proposes to use gradient flows in the space of probability measures endowed with the sliced-Wasserstein distance.
  arXiv  Detail & Related papers  (2021-10-21T08:34:26Z)
• Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization [94.18714844247766]
  Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport. We present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures.
  arXiv  Detail & Related papers  (2021-02-02T21:01:13Z)
• Variational Transport: A Convergent Particle-BasedAlgorithm for Distributional Optimization [106.70006655990176]
  A distributional optimization problem arises widely in machine learning and statistics. We propose a novel particle-based algorithm, dubbed as variational transport, which approximately performs Wasserstein gradient descent. We prove that when the objective function satisfies a functional version of the Polyak-Lojasiewicz (PL) (Polyak, 1963) and smoothness conditions, variational transport converges linearly.
  arXiv  Detail & Related papers  (2020-12-21T18:33:13Z)
• The Wasserstein Proximal Gradient Algorithm [23.143814848127295]
  Wasserstein gradient flows are continuous time dynamics that define curves of steepest descent to minimize an objective function over the space of probability measures. We propose a Forward Backward (FB) discretization scheme that can tackle the case where the objective function is the sum of a smooth and a nonsmooth geodesically convex terms.
  arXiv  Detail & Related papers  (2020-02-07T22:19:32Z)
• A Near-Optimal Gradient Flow for Learning Neural Energy-Based Models [93.24030378630175]
  We propose a novel numerical scheme to optimize the gradient flows for learning energy-based models (EBMs) We derive a second-order Wasserstein gradient flow of the global relative entropy from Fokker-Planck equation. Compared with existing schemes, Wasserstein gradient flow is a smoother and near-optimal numerical scheme to approximate real data densities.
  arXiv  Detail & Related papers  (2019-10-31T02:26:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.