The back-and-forth method for Wasserstein gradient flows

TL;DR

This paper introduces an efficient computational method for Wasserstein gradient flows by extending the back-and-forth method to solve dual problems, enabling large-scale simulations of complex energies.

Contribution

It generalizes the back-and-forth method to compute Wasserstein gradient flows via dual problem formulation, improving efficiency and scalability.

Findings

Enables large-scale simulations of Wasserstein gradient flows.

Handles singular and non-convex energies effectively.

Provides a more well-behaved dual problem for computation.

Abstract

We present a method to efficiently compute Wasserstein gradient flows. Our approach is based on a generalization of the back-and-forth method (BFM) introduced by Jacobs and L\'eger to solve optimal transport problems. We evolve the gradient flow by solving the dual problem to the JKO scheme. In general, the dual problem is much better behaved than the primal problem. This allows us to efficiently run large-scale simulations for a large class of internal energies including singular and non-convex energies.