DeepAI AI Chat
Log In Sign Up

Variational Wasserstein gradient flow

by   Jiaojiao Fan, et al.
Georgia Institute of Technology
University of Washington

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.


The back-and-forth method for Wasserstein gradient flows

We present a method to efficiently compute Wasserstein gradient flows. O...

A Proximal-Gradient Algorithm for Crystal Surface Evolution

As a counterpoint to recent numerical methods for crystal surface evolut...

Fisher information regularization schemes for Wasserstein gradient flows

We propose a variational scheme for computing Wasserstein gradient flows...

Data driven gradient flows

We present a framework enabling variational data assimilation for gradie...

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...

A Wasserstein Minimum Velocity Approach to Learning Unnormalized Models

Score matching provides an effective approach to learning flexible unnor...

Primal Dual Interpretation of the Proximal Stochastic Gradient Langevin Algorithm

We consider the task of sampling with respect to a log concave probabili...