A fast approach to optimal transport: the back-and-forth method

Abstract

We present a method to efficiently solve the optimal transportation problem for a general class of strictly convex costs. Given two probability measures supported on a discrete grid with n points we compute the optimal transport map between them in O(n log(n)) operations and O(n) storage space. Our approach allows us to solve optimal transportation problems on spatial grids as large as 4096 × 4096 and 384 × 384 × 384 in a matter of minutes.

The back-and-forth method for Wasserstein gradient flows

“…We begin in Section 3.1 with the case where U is convex with respect to ρ. In this case, the JKO scheme has an equivalent dual problem that we solve using an adaptation of the back-and-forth method from [21]. In Section 3.2, we show that the algorithm is gradient stable in a properly weighted H 1 space for convex energies of the form…”

“…The reformulations I and J genuinely simplify the task of finding maximizers. On a regular discrete grid, the c-transform can be computed very efficiently [21,23]. As a result, it is much more tractable to maximize I and J, rather than trying work with (1.4) directly.…”

“…We will find the maximizers φ * and ψ * by building upon the BFM algorithm introduced in [21]. The original BFM gives a very efficient scheme for finding the maximizers in the special case where U * is a linear functional.…”

“…In between gradient steps, information in one space (φ-space or ψ-space) is propagated back to the other by taking a forward/backward c-transform. As noted in [21], the advantage of the back-and-forth approach is that certain features of the optimal solution pair (φ * , ψ * ) may be easier to build in one space compared to the other. As a result, the back-and-forth method converges far more rapidly than vanilla gradient ascent methods that operate only on φ-space or only on ψ-space.…”

“…In this paper, we solve problem (1.2) by adapting the back-and-forth method (BFM) introduced in [21]. BFM is a state-of-the-art algorithm for computing optimal transport maps between two fixed densities.…”