Variational Wasserstein gradient flow

12/04/2021
by   Jiaojiao Fan, et al.
24

The gradient flow of a function over the space of probability densities with respect to the Wasserstein metric often exhibits nice properties and has been utilized in several machine learning applications. The standard approach to compute the Wasserstein gradient flow is the finite difference which discretizes the underlying space over a grid, and is not scalable. In this work, we propose a scalable proximal gradient type algorithm for Wasserstein gradient flow. The key of our method is a variational formulation of the objective function, which makes it possible to realize the JKO proximal map through a primal-dual optimization. This primal-dual problem can be efficiently solved by alternatively updating the parameters in the inner and outer loops. Our framework covers all the classical Wasserstein gradient flows including the heat equation and the porous medium equation. We demonstrate the performance and scalability of our algorithm with several numerical examples.

READ FULL TEXT
research
11/16/2020

The back-and-forth method for Wasserstein gradient flows

We present a method to efficiently compute Wasserstein gradient flows. O...
research
06/22/2020

A Proximal-Gradient Algorithm for Crystal Surface Evolution

As a counterpoint to recent numerical methods for crystal surface evolut...
research
07/03/2019

Fisher information regularization schemes for Wasserstein gradient flows

We propose a variational scheme for computing Wasserstein gradient flows...
research
05/24/2022

Data driven gradient flows

We present a framework enabling variational data assimilation for gradie...
research
09/16/2022

Solving Fredholm Integral Equations of the First Kind via Wasserstein Gradient Flows

Solving Fredholm equations of the first kind is crucial in many areas of...
research
02/18/2020

A Wasserstein Minimum Velocity Approach to Learning Unnormalized Models

Score matching provides an effective approach to learning flexible unnor...
research
07/04/2023

Wasserstein medians: robustness, PDE characterization and numerics

We investigate the notion of Wasserstein median as an alternative to the...

Please sign up or login with your details

Forgot password? Click here to reset