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

TL;DR

This paper introduces a fast iterative method for solving optimal transportation problems with convex costs, achieving efficient computation on large spatial grids with exponential convergence.

Contribution

The paper proposes a novel back-and-forth iterative approach that significantly reduces computational complexity for large-scale optimal transport problems with convex costs.

Findings

Able to solve large grid problems in minutes

Uses linear storage space and near-linear time per iteration

Achieves exponential convergence rate

Abstract

We present an iterative method to efficiently solve the optimal transportation problem for a class of strictly convex costs which includes quadratic and p-power costs. Given two probability measures supported on a discrete grid with n points, we compute the optimal map using O(n) storage space and O(n log(n)) operations per iteration, with an approximately exponential convergence rate. Our approach allows us to solve optimal transportation problems on spatial grids as large as 4096x4096 and 384x384x384 in a matter of minutes.

Tables

Table 1. Table 1. Two discs in 2 2 2 dimensions

	Error ${10}^{-4}$		Error ${10}^{-8}$
Grid size	Iterations	Time (s)	Iterations	Time (s)
$512\times 512$	3	1.00	5	1.25
$1024\times 1024$	3	4.23	5	5.21
$2048\times 2048$	3	16.7	5	21.75
$4096\times 4096$	3	56.46	5	75.70

Table 2. Table 2. Two balls in 3 3 3 dimensions

	Error ${10}^{-4}$		Error ${10}^{-8}$
Grid size	Iterations	Time (s)	Iterations	Time (s)
$128\times 128\times 128$	6	17.89	10	27.00
$256\times 256\times 256$	6	156.94	9	205.5
$384\times 384\times 384$	6	472.14	9	699.21

Table 3. Table 3. One square to four squares

	Error ${10}^{-4}$		Error ${10}^{-5}$		Error ${10}^{-6}$
Grid size	Iterations	Time (s)	Iterations	Time (s)	Iterations	Time (s)
$512\times 512$	3	1.08	5	1.45	13	2.96
$1024\times 1024$	3	4.98	5	5.82	14	12.2
$2048\times 2048$	3	18.9	5	25.2	14	53.9
$4096\times 4096$	3	68.3	5	93.6	13	217

Table 4. Table 4. One cube to eight cubes

	Error ${10}^{-3}$		Error ${10}^{-5}$
Grid size	Iterations	Time (s)	Iterations	Time (s)
$128\times 128\times 128$	3	13.36	6	21.63
$256\times 256\times 256$	3	117.44	8	255.17
$384\times 384\times 384$	3	377.05	13	1577.23

Table 5. Table 5. BFM v.s. gradient ascent (two discs)

	Error ${10}^{-4}$		Error ${10}^{-5}$		Error ${10}^{-6}$
Method	Iterations	Time (s)	Iterations	Time (s)	Iterations	Time (s)
BFM	3	4.23	5	5.21	5	5.21
gradient ascent	50	19.0	149	47.9	321	94.5

Table 6. Table 6. BFM v.s. gradient ascent (1 square to 4 squares)

	Error ${10}^{-4}$		Error ${10}^{-5}$		Error ${10}^{-6}$
Method	Iterations	Time (s)	Iterations	Time (s)	Iterations	Time (s)
BFM	3	4.98	5	5.82	14	12.2
gradient ascent	34	13.6	106	34.0	744	207

Equations

$$\underset{\varphi ,\psi }{\sup }\int \varphi d\nu +\int \psi d\mu ,$$

$$\varphi (y)+\psi (x)\le c(x,y).$$

$${\varphi }_{\ast }(y)={\psi }_{\ast }^c(y):=\underset{x}{\inf }c(x,y)-{\psi }_{\ast }(x),$$

$${\psi }_{\ast }(x)={\varphi }_{\ast }^c(x):=\underset{y}{\inf }c(x,y)-{\varphi }_{\ast }(y).$$

$$J(\varphi )=\int \varphi d\nu +\int {\varphi }^{c}d\mu ,$$

$$I(\psi )=\int {\psi }^{c}d\nu +\int \psi d\mu ,$$

$$c(x,y)=h(y-x),$$

$$C(\mu,\nu)=\inf_{T}\int_{\Omega}c\big{(}x,T(x)\big{)}\,d\mu(x),$$

$${\int }_{\Omega }f(y)d(T_{\#}\mu )(y)={\int }_{\Omega }f(T(x))d\mu (x).$$

$$\int_{\Omega}\phi\big{(}T(x)\big{)}\,d\mu(x)=\int_{\Omega}\phi(y)\,d\nu(y)$$

$$\inf_{T_{\#}\mu=\nu}\int_{\Omega}c\big{(}x,T(x)\big{)}\,d\mu(x)=\inf_{T}\sup_{\phi}\int_{\Omega}c\big{(}x,T(x)\big{)}\,d\mu(x)-\phi\big{(}T(x)\big{)}\,d\mu(x)+\int_{\Omega}\phi(y)\,d\nu(y).$$

$$\sup_{\phi}\inf_{T}\int_{\Omega}\Big{(}c\big{(}x,T(x)\big{)}-\phi\big{(}T(x)\big{)}\Big{)}\,d\mu(x)+\int_{\Omega}\phi(y)\,d\nu(y).$$

$${\varphi }^c(x)=\underset{y\in \Omega }{\inf }c(x,y)-\varphi (y).$$

$$J(\varphi )={\int }_{\Omega }\varphi d\nu +{\int }_{\Omega }{\varphi }^{c}d\mu .$$

$$C(\mu ,\nu )=\underset{\varphi }{\sup }J(\varphi ),$$

$$T_{\varphi }(x)=x-(\nabla h{)}^{-1}(\nabla {\varphi }^c(x)),$$

$$\underset{ϵ\to 0}{\lim }\frac{(\varphi +ϵu{)}^c(x)-{\varphi }^c(x)}{ϵ}=-u(T_{\varphi }(x))$$

$$\underset{\varphi }{\sup }{\int }_{\Omega }\varphi d\nu +{\int }_{\Omega }{\varphi }^{c}d\mu =:J(\varphi )$$

$$F:\mathcal{H}\to \mathbb{R}.$$

$$\underset{\Vert h{\Vert }_{\mathcal{H}}\to 0}{\lim }\frac{\Vert F(\varphi +h)-F(\varphi )-\delta F_{\varphi }(h){\Vert }_{\mathcal{H}}}{\Vert h{\Vert }_{\mathcal{H}}}=0.$$

$$⟨{\nabla }_{\mathcal{H}}F(\varphi ),h{⟩}_{\mathcal{H}}=\delta F_{\varphi }(h)$$

$${\dot{H}}^1(\Omega )=\{\phi :\Omega \to \mathbb{R}:{\int }_{\Omega }\phi (x)dx=0\text{and}{\int }_{\Omega }|\nabla \phi (x){|}^{2}dx<\infty \},$$

$$⟨{\phi }_1,{\phi }_2{⟩}_{{\dot{H}}^1}={\int }_{\Omega }\nabla {\phi }_1(x)\cdot \nabla {\phi }_2(x)dx.$$

$$\Vert {\rho }_2-{\rho }_1{\Vert }_{{\dot{H}}^{-1}}^2={\int }_{\Omega }|\nabla \varphi (x){|}^{2}dx,$$

$$\delta F_{\varphi }(h)={\int }_{\Omega }h(x)f_{\varphi }(x)dx.$$

$${\nabla }_{{\dot{H}}^1}F(\varphi )=(-\Delta {)}^{-1}{\bar{f}}_{\varphi },$$

$${\varphi }_{n+1}={\varphi }_n+\sigma {\nabla }_{\mathcal{H}}F({\varphi }_n),$$

$$F({\varphi }_n)\to \sup F,$$

$$F(\varphi )\ge F(\hat{\varphi })+\delta F_{\hat{\varphi }}(\varphi -\hat{\varphi })-\frac{1}{2\sigma }\Vert \varphi -\hat{\varphi }{\Vert }_{\mathcal{H}}^2$$

$${\varphi }_{n+1}={\varphi }_n+\sigma {\nabla }_{\mathcal{H}}F({\varphi }_n)$$

Full text

M.J. and F.L. are supported by AFOSR MURI FA9550-18-1-0502 and ONR N00014-12-1-0838. M.J. is also supported by NSF DMS-1118971, DARPA FA8750-18-2-0066, and ONR N00014-18-1-2527

1. Introduction

The optimal transportation problem was first introduced by Monge in 1781, to find the most cost-efficient way to transport mass from a set of sources to a set of sinks. The theory was modernized and revolutionized by Kantorovich in 1942, who found a key link between optimal transport and linear programming. In recent years, there has been an explosion of interest in optimal transport thanks in part to the discovery of deep connections between the quadratic-cost optimal transport problem and a diverse class of partial differential equations (PDEs) arising in statistical mechanics and fluid mechanics ([Bre91, Bre89, BB00, Ott01, JKO98] to name just a few of the most prominent results). In addition, optimal transport has become popular in data science (particularly machine learning and image processing), where it provides a very natural way to compare and interpolate probability distributions [McC97, HZTA04, ACB17].

In this work, we are interested in computing the optimal transport problem for large-scale applications that arise in image processing and in numerical methods for solving PDEs (in both two and three dimensions). For these applications, one needs to accurately compute the optimal transport map on enormous computational domains (millions of grid points or pixels). Up to now, computing the optimal map has been a notoriously difficult task. To the best of our knowledge, all previously known methods for solving the optimal transport problem either do not scale linearly with respect to the problem size [BC89, BB00, PBtTB*+*15, LR17], cannot provide an accurate computation of the optimal map [Cut13], or are only applicable to a limited class of probability densities [BFO14]. Semi-discrete optimal transport methods [Lé15, Mé11] have been applied successfully to a wide range of optimal transport problems [GM18, MM16], but require expensive geometric processing overhead.

Our goal in this paper is to remedy the current situation and provide a simple, efficient, and accurate algorithm for computing optimal transport maps. To that end, we introduce a new method for solving optimal transport problems: the back-and-forth method. Given two probability densities $\mu$ and $\nu$ discretized on a grid with $n$ points and a strictly convex cost function, our approach computes the optimal map using $O(n)$ storage space and $O(n\log(n))$ computation time per iteration. The number of iterations required for $\epsilon$ accuracy grows like O\Big{(}\max(\lVert\mu\rVert_{\infty},\lVert\nu\rVert_{\infty})\log(\frac{1}{\epsilon})\Big{)} in all tested experiments. As a result, the method converges extremely rapidly. Notably, in contrast to many other approaches, the back-and-forth method does not require positive lower bounds on the probability densities. Thus, we can solve optimal transport problems with densities that vanish on large portions of the domain. In total, the efficiency and flexibility of the back-and-forth method allows us to solve optimal transport problems on grids of size $4096\times 4096$ and $384\times 384\times 384$ in a matter of minutes on a personal computer. Furthermore, the method could be accelerated even further by applying parallelization or other standard scientific computing techniques such as multigrid acceleration.

1.1. Overall approach

The back-and-forth method is based on solving the Kantorovich dual formulation of optimal transport. To introduce this formulation, suppose we are given two probability measures $\mu$ and $\nu$ and a cost function $c(x,y)$, which measures the cost to move a unit of mass from location $x$ to location $y$ (for now we shall be deliberately vague about underlying spaces). The optimal cost to transport $\mu$ to $\nu$ is given by the value of the Kantorovich dual problem:

[TABLE]

where the Kantorovich potentials $\phi$ and $\psi$ are two scalar functions constrained by the inequality

[TABLE]

In addition to computing the optimal cost, the dual problem also encodes information about the optimal map. If an optimal map from $\mu$ to $\nu$ exists (or vice-versa), it can be recovered from the maximizers $\phi_{*},\psi_{*}$ of the dual problem. Indeed, in this case, the pair $(x,y)$ is in the graph of the map if and only if $\phi_{*}(y)+\psi_{*}(x)=c(x,y).$ Thus, to compute optimal maps it is enough to solve the dual problem.

In what follows, we will consider two equivalent unconstrained formulations of the dual problem. These formulations are based upon the observation that if $\phi_{*}$ and $\psi_{*}$ are the maximizers of the dual problem, then one must have the relations

[TABLE]

and

[TABLE]

Here we make use of an operator $\phi\to\phi^{c}$ (resp. $\psi\to\psi^{c}$) called the $c$-transform, which plays an important role in optimal transport and in our proposed method. Note that if one views the dual problem as a linear programming problem, then the above two equations are precisely the complementary slackness condition.

Now we see that there are two equivalent ways to remove the constraint (1), we can either replace $\phi$ by $\psi^{c}$ or we can replace $\psi$ by $\phi^{c}$. This leads to the twin functionals

[TABLE]

and

[TABLE]

which encode exactly the same problem, just posed in different spaces. $J$ formulates the problem in “$\phi$-space” and $I$ formulates the problem in “$\psi$-space”. Here one can draw an analogy to the Fourier transform: indeed, the reader is probably already quite familiar with functionals that can be equivalently expressed in either physical space or Fourier space.

The back-and-forth method solves the Kantorovich dual problem by hopping back-and-forth between gradient ascent updates on $J$ in $\phi$-space and gradient ascent updates on $I$ in $\psi$-space (hence the name). Gradient are taken with respect to the $\dot{H}^{1}$ metric (see (2)). In between gradient steps, information in one space is propagated back to the other space by taking a $c$-transform (c.f. Algorithm 1).

The advantage of the back-and-forth approach is that certain features of the optimal solution pair $(\phi_{*},\psi_{*})$ may be easier to build in one space compared to the other. For example, the Hessian of $\psi=\phi^{c}$ is very closely related to the inverse of the Hessian of $\phi$. Thus, large Hessian eigenvalues in one space correspond to small Hessian eigenvalues in the other space, and smaller features can be built in fewer gradient steps. By hopping back-and-forth between the two spaces, we get the best of both worlds — there is always an opportunity to build features in the space where they are smaller. As a result, the back-and-forth method converges far more rapidly than vanilla gradient ascent methods that operate only on $\phi$-space or only on $\psi$-space. Indeed, on certain examples, our method produces a 10,000-fold reduction in error in just 2 to 4 additional iterations (c.f. Tables 1, 2).

In addition, let us highlight that the choice of gradient ascent steps in the $\dot{H}^{1}$ metric appears to be crucial for maintaining the stability of the algorithm. Indeed, it seems that it is not possible to guarantee an increase in the value of the dual problem if the gradient ascent steps are taken in any Hilbert space weaker than $\dot{H}^{1}$ (see Section 3.1 and Proposition 1 for more details).

Turning to the computational efficiency of the back-and-forth method, we see that the scheme amounts to performing $\dot{H}^{1}$ gradient ascent iterations alternatively on $J$ and $I$ (i.e. in $\phi$-space and $\psi$-space), and computing $c$-transforms. Crucially, the derivative of $J$ (and, by symmetry, the derivative of $I$) takes a simple form [Gan95] which can be efficiently computed. Furthermore, for a large class of costs, the $c$-transform can be computed extremely efficiently. On a discrete grid with $n$ points, an exact $c$-transform can be computed in $O(n\log(n))$ operations [Cor96, Luc97] (see Section 4.1 for more details).

Finally, let us note that a gradient ascent method on the Kantorovich dual problem was previously considered in [CWVB09]. In their paper, the authors took $L^{2}$ gradient ascent steps on only one form of the dual problem. Furthermore, the authors tried to avoid computing $c$-transforms, and thus were forced to use an inaccurate estimation of the true gradient. Unsurprisingly, the authors note that their scheme appears to be unstable. Indeed, we believe that the ideas behind the back-and-forth method are necessary to obtain a robust, efficient, and accurate gradient ascent type method for solving the Kantorovich dual problem.

1.2. Future work and paper outline

While the present work is focused on the optimal transport problem, we anticipate that the back-and-forth method will prove useful for many other problems. For instance, the back-and-forth method can be readily adapted to compute Wasserstein gradient flows. This opens the door to large-scale simulations of a wide class of important and interesting PDEs, [Ott01, JKO98, Bre89, JKM19] to name just a few. In addition, the method may prove useful to solve computational problems arising from the burgeoning area of mean field games [LL07, HMC06]. We look forward to exploring these applications in future work.

The rest of the paper is organized as follows. In Section 2, we recall the necessary background information for optimal transport and gradient-based optimization schemes. In Section 3, we introduce the back-and-forth method and provide arguments for its stability and efficiency. In Section 4, we discuss the numerical implementation of the algorithm, and conduct various experiments to demonstrate its performance.

2. Background

2.1. Optimal transport and the $c$-transform

Let $\Omega$ be a convex and compact subset of $\mathbb{R}^{d}$. A cost on $\Omega$ is a continuous function $c\colon\Omega\times\Omega\to\mathbb{R}$. In the theory of optimal transport fairly general costs can be considered, however in the present work we will focus on the case

[TABLE]

for a strictly convex and even function $h\colon\mathbb{R}^{d}\to\mathbb{R}$. Given two probability measures $\mu$ and $\nu$ supported in $\Omega$, the Monge formulation of the optimal transport problem is defined by

[TABLE]

where the infimum runs over maps $T\colon\Omega\to\Omega$ which transport $\mu$ to $\nu$, i.e. such that $T_{\#}\mu=\nu$. We recall that the pushforward measure $T_{\#}\mu$ is defined by $T_{\#}\mu(A)=\mu(T^{-1}(A))$ for any measurable subset $A\subset\Omega$. One can also characterize the pushforward by defining the integral of the pushforward measure against continuous test functions $f\colon\Omega\to\mathbb{R}$:

[TABLE]

This second formulation is extremely useful for studying pushforwards in the context of optimal transport.

Throughout the paper, we will concentrate on the special case where both probability measures $\mu$ and $\nu$ are absolutely continuous with respect to the Lebesgue measure. As a result, we will frequently conflate a measure and its density function. Under this assumption, there exists a unique optimal map $T_{*}$, which pushes $\mu$ to $\nu$, and its inverse $T^{-1}_{*}$ is the optimal map that pushes $\nu$ to $\mu$  [Bre91, Gan95]. Furthermore, one can find the optimal map by solving the Kantorovich dual formulation of the optimal transport problem.

The dual formulation can be derived by introducing Lagrange multipliers for the pushforward constraint. Indeed, the pushforward constraint $T_{\#}\mu=\nu$ holds if and only if

[TABLE]

for every continuous function $\phi$. Therefore,

[TABLE]

Under the assumption that $\mu$ is absolutely continuous, the interchange of infimum and supremum is valid [Vil03, San15]. If we then group the terms involving $T$ we arrive at

[TABLE]

Now it is clear that the optimal choice of $T$ at some point $x_{0}$ is completely decoupled from the choice at any other point $x$. Therefore, one can move the infimum inside of the integral, and it becomes clear that the operation \inf_{T(x)}c\big{(}x,T(x)\big{)}-\varphi\big{(}T(x)\big{)} plays an important role in the dual problem. This operation is known as the $c$-transform, and it is at the heart of optimal transport. The $c$-transform maps a function $\phi$ to another function $\phi^{c}$ and it can be seen as a generalization of the Legendre transform (convex conjugation) from convex analysis.

Definition 1**.**

Given a continuous function $\phi\colon\Omega\to\mathbb{R}$, we define its $c$-transform $\phi^{c}\colon\Omega\to\mathbb{R}$ by

[TABLE]

Additionally, we say that $\phi$ is $c$-concave if there exists a continuous function $\psi\colon\Omega\to\mathbb{R}$ such that $\phi=\psi^{c}$ and we say that a pair of functions $(\phi,\psi)$ is $c$-conjugate if $\phi=\psi^{c}$ and $\psi=\phi^{c}$.

We are now ready to introduce the Kantorovich dual functional

[TABLE]

Our above discussion then shows that

[TABLE]

where the supremum runs over continuous functions $\phi\colon\Omega\to\mathbb{R}$. Note that in this section we will just focus on one formulation of the dual problem. It is easy to see that any properties satisfied by $J$ must also be satisfied by the alternative formulation, which we previously denoted in the introduction as $I$.

The following lemma encapsulates the most important properties of the $c$-transform (from the perspective of optimal transport).

Lemma 1** ([GM96, Gan95]).**

Let $\phi:\Omega\to\mathbb{R}$ be a continuous function. The $c$-transform satisfies the following properties:

1. (i)

$\phi(x)\leq\phi^{cc}(x)$* for every $x\in\Omega$, and $\phi^{cc}=\phi$ if and only if $\phi$ is $c$-concave. As a result, $\phi^{ccc}=\phi^{c}$ for any continuous function $\phi$.* 2. (ii)

If $\phi$ is $c$-concave, the minimization problem $\inf_{y\in\Omega}c(x,y)-\phi(y)$ has a unique minimizer $T_{\phi}(x)$ for almost every $x$. Furthermore, we have the explicit formula:

[TABLE]

where we recall that $c(x,y)=h(y-x)$. 3. (iii)

If $\phi$ is $c$-concave and $u$ is a continuous function on $\Omega$, then

[TABLE]

for almost every $x\in\Omega$.

With these properties in hand, one can establish one of the most fundamental results in optimal transport:

Theorem 1** ([Bre91, GM96, Gan95]).**

The dual problem

[TABLE]

satisfies the following properties:

1. (i)

$J$* is concave with respect to $\phi$.* 2. (ii)

If $\mu$ is absolutely continuous with respect to the Lebesgue measure, then $J$ is maximized by a $c$-concave function $\phi_{*}$ and $T_{\phi_{*}}$ is the unique optimal map which pushes $\mu$ to $\nu$. 3. (iii)

If $\nu$ is also absolutely continuous, then $T_{\phi_{*}}$ is invertible almost everywhere, and $T_{\phi_{*}}^{-1}$ is the optimal map which pushes $\nu$ to $\mu$.

From Theorem 1, we see that the optimal map can be computed by solving a concave maximization problem. Thus standard techniques of optimization, such as gradient ascent, can be leveraged to solve the optimal transport problem. However, in order to use gradient ascent effectively, one must choose the correct notion of distance and step size. For this reason, we review gradient ascent in Hilbert spaces below.

2.2. Gradient Ascent

In this section we recall elementary results on constant step-size gradient ascent methods for concave functions. Let $(\mathcal{H},\lVert\cdot\rVert_{\mathcal{H}})$ be a separable Hilbert space and suppose that $F$ is a smooth convex functional

[TABLE]

We first recall the notions of differential map and gradients.

Definition 2**.**

Given a point $\phi\in\mathcal{H}$ we say that a bounded linear map $\delta F_{\phi}\colon\mathcal{H}\to\mathbb{R}$ is the first variation (Fréchet derivative) of $F$ at $\phi$ if

[TABLE]

Definition 3**.**

Let $\langle\cdot,\cdot\rangle_{\mathcal{H}}$ be the inner product associated with the Hilbert space $\mathcal{H}$. We say that a map $\nabla_{\mathcal{H}}F\colon\mathcal{H}\to\mathcal{H}$ is the $\mathcal{H}$-gradient of $F$ if

[TABLE]

for all $(\phi,h)\in\mathcal{H}\times\mathcal{H}$.

In the back-and-forth method, we will make use of the $\dot{H}^{1}$-gradient where $\dot{H}^{1}(\Omega)$ is the Hilbert space

[TABLE]

with inner product

[TABLE]

Let us also mention the dual space $\dot{H}^{-1}(\Omega)$: for two probability densities $\rho_{1}$ and $\rho_{2}$ we define

[TABLE]

where $\phi$ is the unique solution in $\dot{H}^{1}(\Omega)$ to the Laplace equation $-\operatorname{\Delta}\phi=\rho_{2}-\rho_{1}$ with zero Neumann boundary conditions (for more details we refer to [San15, Section 5.5.2]).

The following lemma shows that the $\dot{H}^{1}$-gradient has a particularly simple form. The proof is a straightforward application of integration by parts.

Lemma 2**.**

Suppose that $F\colon\dot{H}^{1}(\Omega)\to\mathbb{R}$ is a Fréchet-differentiable function such that for any $\phi\in\dot{H}^{1}(\Omega)$, the first variation $\delta F_{\phi}$ evaluated at any point $h\in\dot{H}^{1}(\Omega)$ can be written as integration against a function $f_{\phi}$, i.e.

[TABLE]

Then the $\dot{H}^{1}$-gradient of $F$ has the form

[TABLE]

where $(-\Delta)^{-1}$ is the inverse negative Laplacian operator taken with zero Neumann boundary conditions and $\bar{f}_{\phi}=f_{\phi}-\frac{1}{|\Omega|}\int_{\Omega}f_{\phi}$.

Now that we have seen an example of a $\mathcal{H}$-gradient, we are ready to discuss gradient ascent methods on general Hilbert spaces. Gradient ascent maximizes a Fréchet-differentiable functional $F\colon\mathcal{H}\to\mathbb{R}$ by iterating

[TABLE]

where the step-size $\sigma>0$ is a fixed constant. In order to obtain convergence to the supremum, i.e.

[TABLE]

one needs to control the continuity of the gradient mapping $\nabla_{\mathcal{H}}F$. The following convergence theorem for gradient ascent is one of the cornerstones of optimization.

Theorem 2** (See for instance [Nes04]).**

Let $F\colon\mathcal{H}\to\mathbb{R}$ be a Fréchet-differentiable concave functional and suppose that there exists $\sigma>0$ such that

[TABLE]

for all $\phi,\hat{\phi}\in\mathcal{H}$. Then the gradient ascent iterations

[TABLE]

satisfy the ascent property

[TABLE]

Furthermore, if $F$ has a unique maximizer $\phi_{*}$ and if $\sup_{n}\lVert\phi_{n}\rVert_{\mathcal{H}}<\infty$ then the sequence $\{\phi_{n}\}_{n=0}^{\infty}$ converges weakly to $\phi_{*}$.

For completeness, we give a short proof of Theorem 2 in the appendix. With these tools in hands we are now ready to introduce the back-and-forth method.

3. The back-and-forth method

Let $\Omega\subset\mathbb{R}^{d}$ be a convex and compact region and consider a cost $c\colon\Omega\times\Omega\to\mathbb{R}$ of the form

[TABLE]

for a strictly convex and even function $h\colon\mathbb{R}^{d}\to\mathbb{R}$. Given two probability densities $\mu$ and $\nu$ supported in $\Omega$, we are interested in the optimal transport problem in its dual form

[TABLE]

where the supremum runs over continuous functions $\phi\colon\Omega\to\mathbb{R}$, and the $c$-transform is defined by $\phi^{c}(x)=\inf_{y\in\Omega}c(x,y)-\phi(y)$. For more background on the optimal transport problem and on $c$-transforms we refer to Section 2.1.

Our goal in this section is to develop an efficient algorithm to compute the maximizer of the dual problem. In what follows, we will consider the dual functional in the two following equivalent forms:

[TABLE]

and

[TABLE]

Note that the functionals $J$ and $I$ are essentially identical, the only difference is that the roles of $\mu$ and $\nu$ are flipped. In order to proceed further, we need formulas for the variations of $J$ and $I$.

Lemma 3** ([Gan94, Gan95, GM96]).**

Consider the functionals $J$ and $I$ defined by (5) and (6) over the space of continuous functions $\phi\colon\Omega\to\mathbb{R}$ and $\psi\colon\Omega\to\mathbb{R}$ respectively. If $\phi$ is $c$-concave, the first variation of $J$ can be expressed as

[TABLE]

and if $\psi$ is $c$-concave, the first variation of $I$ can be expressed as

[TABLE]

Here we recall that for any $c$-concave function $\varphi:\Omega\to\mathbb{R}$,

[TABLE]

A short proof of Lemma 3 is provided in the appendix for the reader’s convenience.

Now we are ready to introduce our approach to solve the optimal transport problem: the back-and-forth method. The method is outlined in Algorithm 1. It is based on two main ideas:

Gradient ascent steps in the $\dot{H}^{1}$ metric, where

[TABLE] 2. 2.

A back-and-forth update scheme, alternating between gradient ascent steps on $J$ and $I$.

In the following two subsections, we will motivate the choice of $\dot{H}^{1}$-gradient steps and the back-and-forth updates. For information about the step size $\sigma$, see Section 4.2.

3.1. $\dot{H}^{1}$-gradient ascent

The main steps of the back-and-forth method are the $\dot{H}^{1}$-gradient ascent steps

[TABLE]

and

[TABLE]

In order to obtain the convergence of the back-and-forth method, we will need to know if these steps increase the values of the dual functionals $J$ and $I$ respectively. In what follows, we will focus exclusively on the $J$ update step, since any properties that hold for $J$ will also hold for $I$ by symmetry.

Recalling Theorem 2 from the previous section, one can show that for a general Hilbert space $\mathcal{H}$, the $\mathcal{H}$-gradient ascent steps

[TABLE]

satisfy the ascent property

[TABLE]

if the inequality

[TABLE]

holds for some $\sigma>0$ and any $\phi,\hat{\phi}\in\mathcal{H}$. Note that this inequality becomes easier to satisfy when the norm associated to $\mathcal{H}$ becomes stronger i.e. when right hand side becomes more negative.

To establish the inequality, we need to bound from below

[TABLE]

which can be recognized as the error in approximating $J$ with a first-order Taylor expansion about $\hat{\phi}$. The easiest way to make progress is to use the concave inequality $J(\hat{\phi})\leq J(\phi)+\delta J_{\phi}(\hat{\phi}-\phi)$ to get the bound

[TABLE]

We can use the explicit formula for the first variation in Lemma 3, to write

[TABLE]

Now inequality (8) will follow if we can show that

[TABLE]

Thus, our goal is to find a Hilbert space $\mathcal{H}$ and a parameter $\sigma$ that make this inequality true.

For any Hilbert space $\mathcal{H}$, we have the Cauchy–Schwartz inequality

[TABLE]

where $\mathcal{H}^{*}$ is the dual space to $\mathcal{H}$ with respect to the $L^{2}$ inner product (i.e. we dualize $\mathcal{H}$ with respect to the so-called pivot space $L^{2}(\Omega)$). If we can show that for an appropriate choice of Hilbert space $\mathcal{H}$,

[TABLE]

then combining this with our bound from Cauchy–Schwartz we will get (8) (note a more careful argument can eliminate the factor of $\frac{1}{2}$ on the right side of (9), but we won’t worry about this here). Once again, inequality (9) becomes easier to satisfy as the norm associated to $\mathcal{H}$ becomes stronger. Indeed, as $\mathcal{H}$ becomes stronger $\mathcal{H}^{*}$ becomes weaker, thus the left hand side gets smaller while the right hand side gets larger.

We can interpret inequality (9) as saying that the function $\phi\mapsto T_{\phi\,\#}\mu$ needs to be Lipschitz continuous as a map from $\mathcal{H}$ to $\mathcal{H}^{*}$. Thus, to get a gradient ascent scheme with the ascent property, we will need to choose a Hilbert space $\mathcal{H}$ so that if $\phi$ and $\hat{\phi}$ are close in $\mathcal{H}$ then $T_{\phi\,\#}\mu$ and $T_{\hat{\phi}\,\#}\mu$ are close in $\mathcal{H}^{*}$. This is not an easy task. For example, if $\phi$ is not $c$-concave, the map $\phi\mapsto T_{\phi}$ may not even be well-defined, let alone have certain continuity properties.

If one tries to deal with the above complication head-on, it is not clear how to proceed. To make progress, we show that an ascent property holds in the Hilbert space $\dot{H}^{1}(\Omega)=\{\phi\colon\Omega\to\mathbb{R}:\int_{\Omega}\phi=0\text{ and }\lVert\nabla\phi\rVert_{L^{2}(\Omega)}<\infty\}$ if we restrict our attention to the quadratic cost and certain well-behaved $c$-concave functions. Before we state our result, let us note that the choice of $\dot{H}^{1}$ probably cannot be weakened. Indeed, the formula for $T_{\phi}$ depends on $\nabla\phi^{c}$, thus one must have some control on the gradients of $\phi$ and $\hat{\phi}$ to have any hope of showing either (8) or (9).

Proposition 1**.**

Consider the Kantorovich dual problem $\sup_{\phi}J(\phi)$ with the quadratic cost $h(y-x)=\frac{1}{2}|x-y|^{2}$. If we know that there exists $\lambda>0$ such that the iterates

[TABLE]

for every $n$, then the gradient ascent step

[TABLE]

with step size $\sigma=\lVert\mu\rVert_{L^{\infty}}^{-1}\lambda^{d+1}$ satisfies the ascent property

[TABLE]

Moreover, the mass densities move closer to the target $\nu$,

[TABLE]

where we denote $\rho_{n}=T_{\phi_{n}\#}\mu$ and $\rho_{n+\frac{1}{2}}=T_{\phi_{n+\frac{1}{2}}\#}\mu$ and where the $\dot{H}^{-1}$ distance is defined by (3).

See appendix for proof.

Remark 1*.*

Proposition 1 shows that the descent property holds over $\dot{H}^{1}(\Omega)$ when all of the iterates $\phi_{n}$ are assumed to have upper and lower Hessian bounds. Unfortunately, we cannot expect this property to hold throughout the simple ascent scheme $\phi_{n+1}=\phi_{n}+\sigma\nabla_{\!\dot{H}^{1}}J(\phi_{n})$. Indeed, the gradient ascent updates may eventually push the iterates outside of the desired region. In fact, we cannot even guarantee that $c$-concavity is preserved by the gradient update steps.

Remark 2*.*

One way to avoid the aforementioned problem is to add a step where the iterates are projected back onto the set

[TABLE]

$\mathcal{S}_{\lambda}$ is convex, thus the projection step does not interfere with the ascent property. Unfortunately, it is expensive to compute projections onto $\mathcal{S}_{\lambda}$ (one must solve a semi-definite programming problem). Furthermore, when the densities do not have convex support, the maximizer of the dual problem may not lie in $\mathcal{S}_{\lambda}$ for any $\lambda>0$ (however it must be in $\mathcal{S}_{0}$ by $c$-concavity). In such a case, one would need to send $\lambda\to 0$ to obtain an arbitrarily accurate solution.

To circumvent some of the difficulties mentioned in the above remarks, we replace traditional gradient ascent steps with a back-and-forth scheme that alternates between maximizing $J$ and maximizing $I$.

3.2. Back-and-forth updates

Let us recall that the Kantorovich dual problem can be written in the form

[TABLE]

where the supremum runs over continuous functions $\phi$ and $\psi$ satisfying

[TABLE]

for all $x,y\in\Omega$. This formulation emphasizes the symmetric role played by the potentials $\phi$ and $\psi$. In what follows, it will be convenient to represent the dual functional in the form

[TABLE]

where the constraints are encoded by the convex indicator function $\iota_{\mathcal{C}}$ which takes value [math] on the convex set

[TABLE]

and $+\infty$ for all other pairs $(\phi,\psi)$. Then the Kantorovich problem (10) can be simply written $C(\mu,\nu)=\sup_{\phi,\psi}D(\phi,\psi)$. Moreover the functionals $J$ and $I$ can be obtained by either eliminating $\phi$ or $\psi$ from the symmetric representation (11). Indeed, one can check that

[TABLE]

and thus

[TABLE]

A vanilla gradient ascent scheme $\phi_{n+1}=\phi_{n}+\sigma\nabla_{\!\dot{H}^{1}}J(\phi_{n})$ focuses (arbitrarily) on $\phi$-space, iterating gradient steps on the functional $J$. Alternatively, one could work instead in $\psi$-space and write a gradient ascent scheme on the functional $I$. A perhaps better idea is to alternate between $\phi$-space and $\psi$-space. This is the core idea of the back-and-forth method given in Algorithm 1. Recall that a current iterate $(\phi_{n},\psi_{n})$ is updated as follows:

[TABLE]

The back-and-forth approach already corrects a difficulty we identified in the previous subsection. When one considers a pure gradient ascent scheme $\phi_{n+1}=\phi_{n}+\sigma\nabla_{\!\dot{H}^{1}}J(\phi_{n})$ on $\phi$ and $J$ only ( or analogously on $\psi$ and $I$ only), there is no reason that the iterates remain $c$-concave. If $\phi$ or $\psi$ is not $c$-concave, then the gradients $\nabla_{\!\dot{H}^{1}}J(\phi)$ or $\nabla_{\!\dot{H}^{1}}I(\psi)$ may not be well defined. In the back-and-forth method, whenever one takes a gradient, the function in question is always $c$-concave. Indeed, $\phi_{n}=\psi_{n-1}^{c}$ and $\psi_{n+\frac{1}{2}}=\phi_{n+\frac{1}{2}}^{c}$ thus, $\nabla_{\!\dot{H}^{1}}J(\phi_{n})$ and $\nabla_{\!\dot{H}^{1}}I(\psi_{n+\frac{1}{2}})$ are always well-defined.

In addition, the back-and-forth updates can never perform worse than a pure gradient ascent scheme. Indeed, the intermediate steps where we take a $c$-transform can only increase the value of the dual problem. It follows from property (i) in Lemma 1 that

[TABLE]

and

[TABLE]

Thus, the intermediate $c$-transform steps help increase the value of the dual functional. If we combine this with our previous result from Section 3.1, we can conclude the following result about the back-and-forth method.

Proposition 2**.**

Consider the back-and-forth method (Algorithm 1) for the quadratic-cost optimal transport problem. Assume that there exists $\lambda>0$ such that all the iterates $\phi_{n},\phi_{n+\frac{1}{2}}$ and $\psi_{n},\psi_{n+\frac{1}{2}}$ lie in the set of functions

[TABLE]

Then we have the chain of inequalities

[TABLE]

for all integers $n$. Equivalently, the values of $I$ and $J$ alternatively increase:

[TABLE]

The short proof of Proposition 2 is given in the appendix.

Note that even with the back-and-forth updates we cannot give a rigorous convergence proof without assuming convexity bounds on the iterates. Nonetheless, for reasons that we do not fully understand, the back-and-forth method appears to be extremely stable even for relatively large step sizes (a vanilla ascent method on either $J$ or $I$ behaves less favorably). Furthermore, the back-and-forth method converges extremely rapidly to the maximizer. Heuristically, we believe this is because it is easier to build certain features in either “$\phi$-space” or “$\psi$-space”. Indeed, if $(\phi,\psi)$ is a pair of $c$-conjugate functions for the quadratic cost, one has the relationship $D^{2}\phi(x-\nabla\phi(x))=(D^{2}\psi(x))^{-1}$ i.e. the quadratic $c$-transform inverts bounds on Hessians. By alternating between spaces, one can build features in the space where they are smaller, and thus converge to the solution more rapidly.

4. Numerical Implementation and Results

4.1. The fast $c$-transform

Computing the $c$-transform on a regular grid requires solving the optimization problem

[TABLE]

for each grid point $x$. In one dimension, the strict convexity of the cost implies that the $y$ derivative

[TABLE]

is a decreasing function of $x$. As a result, the minimizers $y(x)=\operatorname*{argmin}_{y}h(y-x)-\phi(y)$ are monotone increasing with respect to $x$. In other words, if we have $x_{1}\leq x_{2}$ then $y(x_{1})\leq y(x_{2})$. This observation can be exploited to design a divide-and-conquer algorithm that computes a 1-dimensional $c$-transform on $n$ points in $O(n\log(n))$ operations [Cor96].

The key idea is that if one knows the minimizers $y(x)$ for points the points $x=\frac{k}{n},\frac{k+2}{n}$ then

[TABLE]

i.e. the minimizer for the middle point is “trapped” between the other two. Thus, one can separately compute a $c$-transform for the even grid points and a $c$-transform for the odd grid points then use the above interlacing property to reconstruct the full solution. See [Cor96] for the explicit algorithm and more details.

If we restrict our attention to costs $h$ which decompose along each dimension, i.e.

[TABLE]

then we can compute a $d$-dimensional $c$-transform by repeated applications of the 1-dimensional $c$-transform. For example, in 2 dimensions, we have

[TABLE]

[TABLE]

If we write $\tilde{\phi}(x_{1},y_{2})=-\phi(\cdot,y_{2})^{c}(x_{1})$ then the last line is a $c$-transform with respect to $y_{2}$

[TABLE]

Thus, in 2 dimensions, one can compute the full $c$-transform by first computing one-dimensional $c$-transforms along all horizontal lines, and then along all vertical lines. This idea generalizes to arbitrary dimension, and is the exact same mechanism which is used to compute the multidimensional Fast Fourier Transform (FFT).

In the (extremely important) special case $h(y-x)=\frac{1}{2}|y-x|^{2}$, the $c$-transform can be computed even more efficiently. Splitting the quadratic part into 3 terms, we get

[TABLE]

If we then set $\varphi(y)=\frac{1}{2}|y|^{2}-\phi(y)$, and let $\varphi^{*}(x)$ denote the Legendre transform of $\varphi$:

[TABLE]

then

[TABLE]

Therefore, one can compute the quadratic $c$-transform by computing a Legendre transform instead.

A one-dimensional Legendre transform on a set of $n$ points can be computed in $O(n)$ operations [Luc97]. The idea of [Luc97] uses two very important properties of the Legendre transform, namely that $\varphi^{*}=\varphi^{***}=(\varphi^{**})^{*}$ and that $\varphi^{**}$ is the convex hull of $\varphi$. Therefore, computing (12) amounts to finding the values $y(x)$ where

[TABLE]

changes sign from positive to negative. $\varphi^{**}$ is convex, so the slopes $(\varphi^{**})^{\prime}(y)$ are increasing. Therefore, finding the sign change $y(x)$ that corresponds to each $x$ in (13) can be done in a single sweep through the values of $(\varphi^{**})^{\prime}(y)$. On a one-dimensional grid with $n$ regularly spaced points, the sweep step takes $O(n)$ operations and crucially, the convex hull $\varphi^{**}$ can also be computed in $O(n)$ operations [Luc97]. Finally, a multidimensional Legendre transform on a regular grid can be computed by decomposing the problem along each dimension in the same manner as described above for the $c$-transform. For the explicit algorithm and more details, we refer to [Luc97].

4.2. Step sizes

As we saw in Section 3, it is not easy to determine the Lipschitz constant of the gradient mapping at an arbitrary function $\phi$. As such, it is not clear how to optimally choose $\sigma$ for a fixed step size gradient ascent scheme. Thus, rather than choose a single fixed value $\sigma$ at the outset of the algorithm, we update $\sigma=\sigma_{n}$ throughout the algorithm using Armijo–Goldstein type update rules [Arm66].

After each gradient step, we compare the difference $J(\phi_{n+1})-J(\phi_{n})$ to the squared norm $\lVert\nabla_{\!\dot{H}^{1}}J(\phi_{n})\rVert_{\dot{H}^{1}}^{2}$ (and correspondingly for $I$ and $\psi_{n+1}$, $\psi_{n}$). Given parameters $0<\beta_{1}<\beta_{2}<1$, $\alpha_{1}>1$ and $\alpha_{2}<1$ we check whether

[TABLE]

If the upper inequality fails, we decrease $\sigma_{n}$ by taking $\sigma_{n+1}=\alpha_{2}\sigma_{n}$, and if the lower inequality fails then we increase $\sigma_{n}$ by taking $\sigma_{n+1}=\alpha_{1}\sigma_{n}$. In all of our experiments we take $\beta_{1}=\frac{1}{4},\beta_{2}=\frac{3}{4},\alpha_{1}=\frac{5}{4},$ and $\alpha_{2}=\frac{4}{5}$ and we choose $\sigma=8\min(\lVert\mu\rVert_{L^{\infty}}^{-1},\lVert\nu\rVert_{L^{\infty}}^{-1})$ as our starting value. To ensure that the step sizes stay bounded away from zero, we do not allow the update rules to drop the step size below a small constant $\sigma_{\textrm{min}}=0.01$.

Note that we do not perform any backtracking. The new value of $\sigma=\sigma_{n+1}$ is simply used in the next update step. Indeed, evaluating $J$ or $I$ at a function requires computing a $c$-transform, thus it is not worth trying to optimize the step size at a particular iterate.

4.3. Experiments

Throughout this section we will assume that $\Omega=[0,1]^{d}$ is the unit cube in $\mathbb{R}^{d}$. When implementing our algorithms, we discretize $\Omega$ using a regular finite $d$-dimensional grid. Recall that Laplace equations on this domain should be solved with zero Neumann boundary conditions. This ensures that there is no flux of mass outside of the computational domain.

All numerical algorithms were coded in C and executed on a single 1.6 GHz core with 8GB RAM. Inversion of the Laplace operator was performed using the Fast Fourier Transform (FFT). All FFTs were calculated using the free FFTW C library.

In what follows, we will consider three sets of experiments. In the first two sets of experiments, we restrict our attention to optimal transport with the quadratic cost $\frac{1}{2}|x-y|^{2}$. We start by computing optimal transport maps between geometric shapes where explicit answers are available. This allows us to verify the accuracy and convergence rate of the method. We also verify that the back-and-forth method substantially outperforms a gradient ascent method on the dual problem with no back-and-forth structure. Next, we apply our method to images, and use the resulting optimal map to compute the displacement interpolation [Fry93, McC97] between the source and target images. As we shall see, our computed maps are sufficiently accurate to capture the fine-scale details of the images. Finally, in the last set of experiments, we consider the more general class of costs outlined in Section 4.1, and study how the optimal map varies as the underlying cost changes.

Quadratic cost optimal transport

In this subsection, we focus on computing optimal transport with the quadratic cost

[TABLE]

which represents the most important special case of optimal transport.

We begin with a simple example where the two probability densities $\mu$ and $\nu$ differ only by a translation, i.e. $\mu(x)=\nu(x+a)$ for a fixed constant $a\in\mathbb{R}^{d}$. In this case, it is known that the optimal map $T_{*}$ pushing $\mu$ to $\nu$ is itself a translation: $T_{*}(x)=x+a$, and the transportation cost is $\frac{1}{2}|a|^{2}$. Thus, we can directly check the accuracy of our method on this class of elementary examples.

First, in two dimensions, suppose that $\mu$ is uniform and supported on a disc of radius $\frac{1}{8}$ centered at $(\frac{1}{4},\frac{1}{4})$ and $\nu$ is uniform and supported on a disc of the same radius centered at $(\frac{3}{4},\frac{3}{4})$ (both densities are normalized to have unit mass).

\mu$$\nu

It then follows that the quadratic transportation cost between these densities is

[TABLE]

Our results for this experiment are presented in Table 1. We run our algorithm until the difference between the computed transportation cost and the exact solution is less than the prescribed error tolerance. We see that to get an error of $10^{-4}$ we only need 3 iterations, and to get an error of $10^{-8}$ we only need 5 iterations. Let us also highlight that the performance of the algorithm is completely independent of the tested grid size.

Next we repeat the same experiment in three dimensions. In this case, $\mu$ is uniform and supported on a ball centered at $(\frac{1}{4},\frac{1}{4},\frac{1}{4})$ with radius $\frac{1}{8}$ and $\nu$ is uniform and supported on a ball centered at $(\frac{3}{4},\frac{3}{4},\frac{3}{4})$ with the same radius. The cost to transport the densities to one another is $\frac{3}{8}$. Our results for this experiment are presented in Table 2. Again, we run our algorithm until the difference between the computed transportation cost and the exact solution is less than the prescribed error tolerance. Once again we see that the algorithm converges extremely rapidly, and the convergence rate is independent of the grid size.

Now we turn to a more difficult example where the optimal transport map is discontinuous. In two dimensions, we take $\mu$ to be the renormalized characteristic function of a square with side lengths $\frac{1}{4}$ centered at $(\frac{1}{2},\frac{1}{2})$ and $\nu$ to be the renormalized characteristic function of the union of four squares each with side length $\frac{1}{8}$ and centers at $(\frac{3}{16}+\frac{5i}{8},\frac{3}{16}+\frac{5j}{8})$ for $i,j\in\{0,1\}$.

\mu$$\nu

In this case, there is an explicit solution for the optimal transport map. One cuts $\mu$ into 4 squares along the lines $y=\frac{1}{2}$ and $x=\frac{1}{2}$ and then each cut square moves to the closest square in $\nu$ by translation. Thus, we see that the map is discontinuous along the lines $y=\frac{1}{2}$ and $x=\frac{1}{2}$. Each cut square is translated by a vector of the form $(\pm\frac{1}{4},\pm\frac{1}{4})$, therefore, the transportation cost is \frac{1}{2}\big{(}(\frac{1}{4})^{2}+(\frac{1}{4})^{2}\big{)}=\frac{1}{16}. The results for this experiment are presented in Table 3. Note that due to the discontinuity in the optimal map, it is not possible to achieve the same level of accuracy as the first example. Nonetheless, we see that the algorithm still converges to a highly accurate solution in a small number of iterations.

Next, we consider an analogue of this discontinuous example in 3 dimensions. In this setting, we take $\mu$ to be the renormalized characteristic function of a cube with side length $\frac{1}{4}$ centered at $(\frac{1}{2},\frac{1}{2},\frac{1}{2})$ and $\nu$ to be the renormalized characteristic function of the union of eight cubes each with side length $\frac{1}{8}$ and centers at $(\frac{3}{16}+\frac{5i}{8},\frac{3}{16}+\frac{5j}{8},\frac{3}{16}+\frac{5k}{8})$ for $i,j,k\in\{0,1\}$. For this problem, the optimal transportation map is obtained by cutting $\mu$ into eight cubes along the lines $x=\frac{1}{2},y=\frac{1}{2}$, and $z=\frac{1}{2}$ and translating each of the eight cubes to the closest cube in $\nu$. Each cube is translated by a vector of the form $(\pm\frac{1}{4},\pm\frac{1}{4},\pm\frac{1}{4})$, thus, the total transportation cost is \frac{1}{2}\big{(}(\frac{1}{4})^{2}+(\frac{1}{4})^{2}+(\frac{1}{4})^{2}\big{)}=\frac{3}{32}. Once again, due to the discontinuity of the map, the accuracy depends on the grid resolution. As a result, for grids of size $384\times 384\times 384$ and smaller, the method becomes stationary once the accuracy drops slightly below $10^{-5}$. Nonetheless, we continue to see that the algorithm converges rapidly to a stationary state.

Finally, we compare the performance of the back-and-forth method (BFM) to a simple $\dot{H}^{1}$-gradient ascent method on the dual problem (4). The gradient ascent method alternates the following two steps

[TABLE]

hence there is no back-and-forth structure, and the method only considers one of the dual problems. We consider two test cases in 2-dimensions, the disc to disc and the square to four squares examples described above. In each experiment, we choose a grid of size $1024\times 1024$ and run each algorithm until a certain error tolerance is met. In both experiments, the back-and-forth method substantially outperforms the pure gradient ascent method (c.f. Tables 5 and 6). Note that each gradient ascent iteration only requires computing two Legendre transforms and a single FFT, while each back-and-forth iteration must compute 4 Legendre transforms and 2 FFTs. Even when accounting for this difference, the back-and-forth method is still vastly superior in both iteration count and computation time.

Optimal transport of images

In our next set of experiments, we consider optimal transport between two arbitrary black and white images. We shall assume that the image pixels take values in $[0,1]$, where [math] represents black and $1$ represents white. The images can then be converted into probability densities by renormalizing pixel values so that each image has total mass $1$.

In the context of image processing, one typically wishes to compute optimal transport maps to obtain realistic looking interpolations between images. We will focus on the quadratic cost, which is the most natural for this task. Given two images with associated probability densities $\mu,\nu$ and the optimal map $T_{*}$ between them, one can define for each time $t\in[0,1]$ the so called displacement interpolant [Fry93, McC97]

[TABLE]

where $S_{t}(x)=tT_{*}(x)+(1-t)x$. At time zero, $S_{0}$ is the identity map, so $\rho(0,x)=\mu(x)$, and at time 1, $S_{1}=T_{*}$, so $\rho(1,x)=\nu(x)$. The intermediate times $t\in(0,1)$ essentially give a “video” which shows how the optimal map deforms one probability density into the other.

We compute three different examples of image interpolation, Figures 2, 3 and 4, using the images in Figure 1 as our initial and final densities. In Figure 2, the starting image is the silhouette of a jack-o’-lantern and the final image is a silhouette of a winged horse. In Figure 3, the starting configuration is the union of 4 differently sized discs and the ending configuration is the silhouette of a man holding a sword. In Figure 4, the starting configuration is a portrait of Monge and the final configuration is a photo of Kantorovich. Note that in each of these examples, the intermediate time interpolations remain sharp and there is almost no diffusion. In addition, near times 0 and 1, the discontinuities of the optimal map are well localized and resolved. Figure 2 is particularly interesting, as one can see how the delicate features of the legs, wings, and head of the horse form over time.

Even though the examples are large ($1024\times 1024$ pixels) and computed to a high degree of accuracy, the total computation time to find the optimal map and construct the interpolation does not exceed 20 seconds in any case. To the best of our knowledge, the size of the examples and the accuracy that we achieved were out of reach for all previous computational optimal transport methods.

Optimal transport with general costs

In the final set of experiments, we consider optimal transport with more general cost functions. As we noted earlier, the $c$-transform corresponding to any cost of the form $c(x,y)=\sum_{i=1}^{d}h_{i}(y_{i}-x_{i})$ for $h_{i}$ strictly convex can be computed efficiently using the fast $c$-transform. The most interesting costs of this form are sums of $p^{th}$ powers, i.e.

[TABLE]

where $p_{i}>1$ for each coordinate index $i$. These costs are particularly interesting when the powers $p_{i}$ are different for each $i$. Indeed, in this case the optimal map can change considerably as the values of $p_{i}$ are varied.

We illustrate this phenomenon in two dimensions by choosing $\mu$ to be the union of two discs of radius $\frac{1}{8}$ centered at $(\frac{1}{4},\frac{1}{4})$ and $(\frac{3}{4},\frac{3}{4})$ respectively, and $\nu$ to be the union of two discs of radius $\frac{1}{8}$ centered at $(\frac{1}{4},\frac{3}{4})$ and $(\frac{3}{4},\frac{1}{4})$ respectively:

\mu$$\nu

We then consider the optimal transport trajectory between $\mu$ and $\nu$ for several different choices of $p_{1}$ and $p_{2}$ (c.f. Figure 5).

We start off with $p_{1}=p_{2}=2$. In this case, each of the discs in $\mu$ splits evenly along the line $y=x$. The portion of the discs that lie above the line are transported to the disc in $\nu$ centered at $(\frac{1}{4},\frac{3}{4})$ and the parts below the line are transported to the disc in $\nu$ centered at $(\frac{3}{4},\frac{1}{4})$. When $p_{1}<p_{2}$, it is cheaper to travel in the vertical direction as compared to the horizontal direction (for distances smaller than 1). As a result, when we decrease $p_{1}$ and increase $p_{2}$ the optimal map changes to transport more mass in the vertical direction and less mass in the horizontal direction. When the values become sufficiently extreme, (for example $p_{1}=1.1$ and $p_{2}=3$), horizontal transport becomes so expensive that the optimal cost map moves mass in the vertical direction only. See Figure 5 for a visualization of the optimal trajectories for various choices of $p_{1}$ and $p_{2}$.

Appendix A Proofs

Proof of Theorem 2.

The ascent property follows from plugging in the choices $\phi=\phi_{n+1}$ and $\hat{\phi}=\phi_{n}$ into the given inequality.

Now we suppose that $F$ has a unique maximizer $\phi_{*}$ and that the sequence $\{\phi_{n}\}_{n=0}^{\infty}$ lies in a bounded subset of $\mathcal{H}$. We can then extract a weak limit point $\tilde{\phi}$ of the sequence, and by weak upper semi-continuity $F(\tilde{\phi})\geq\lim_{n\to\infty}F(\phi_{n})$ (note weak upper semi-continuity is automatic since we assume that $F$ is concave and real-valued). The values $F(\phi_{n})$ are bounded above by $F(\phi_{*})$, therefore $\lim_{n\to\infty}\lVert\nabla_{\mathcal{H}}F(\phi_{n})\rVert_{\mathcal{H}}=0$. We can then establish the string of inequalities

[TABLE]

where the first inequality follows from the concavity of $F$ and the second follows from choosing $R=\max(\lVert\phi_{*}\rVert_{\mathcal{H}},\sup_{n}\lVert\phi_{n}\rVert_{\mathcal{H}})$. Taking a limit on the right hand side we can conclude that $F(\tilde{\phi})\geq F(\phi_{*})$, and thus, $\tilde{\phi}=\phi_{*}$ since the maximizer is unique.

∎

Proof of Fact 3.

Here we prove the result on $J$ only, the argument for $I$ being identical by symmetry. Given a continuous function $u\colon\Omega\to\mathbb{R}$ we consider the perturbation

[TABLE]

Using property iii of Lemma 1 we see that

[TABLE]

Thus,

[TABLE]

and the result then follows. ∎

We now turn our attention to Proposition 1. The notion of Bregman divergence will play an important role in its proof.

Definition 4**.**

Let $\mathcal{H}$ be a Hilbert space and consider a Fréchet-differentiable function $F\colon\mathcal{H}\to\mathbb{R}$. The Bregman divergence associated to $F$ is defined by

[TABLE]

for all $\phi_{1},\phi_{2}\in\mathcal{H}$. Similarly, if $\Omega$ is a closed and convex subset of $\mathbb{R}^{d}$ and $\xi\colon\Omega\to\mathbb{R}$ is a differentiable function, we define

[TABLE]

for all $x_{1},x_{2}\in\Omega$.

As hinted by the above definition, we will make use of Bregman divergences both for Kantorovich potentials (defined on $\Omega$) and dual functionals (defined on the Hilbert space $\dot{H}^{1}(\Omega)$). The next results gathers properties which are well-known in the theory of Bregman divergences.

Lemma 4**.**

Let $\xi$ be a convex and differentiable function defined on $\Omega$, and denote by $\xi^{*}$ its convex conjugate $\xi^{*}(p)=\sup_{x}p\cdot x-\xi(x)$. Then

1. (i)

For any $\lambda\geq 0$, $\xi$ is $\lambda$-strongly convex $\Leftrightarrow\xi(x_{2}|x_{1})\geq\frac{\lambda}{2}\lvert x_{2}-x_{1}\rvert^{2}$. 2. (ii)

$\xi(x_{2}|x_{1})=\xi^{*}(p_{1}|p_{2})$* for all $x_{1},x_{2}\in\Omega$, where we set $p_{i}=\nabla\xi(x_{i})$.* 3. (iii)

For any $\lambda>0$, if $\xi$ is $\lambda$-strongly convex then $\xi(x_{1}|x_{2})\leq\frac{\lambda^{-1}}{2}\lvert p_{2}-p_{1}\rvert^{2}$, for all $x_{1},x_{2}\in\Omega$, where $p_{i}=\nabla\xi(x_{i})$. 4. (iv)

Fix $x_{0}\in\Omega$. Let $f(u)=\xi(x_{0}+u|x_{0})$ and $g(v)=\xi^{*}(p_{0}+v|p_{0})$, where $p_{0}=\nabla\xi(x_{0})$. Then $f^{*}(v)=g(v)$.

Now, recall that that the dual functional $J$ is defined by $J(\phi)=\int\phi\,\nu+\int\phi^{c}\,\mu$. It is a concave functional, and note that the first term $\int\phi\,\nu$ is linear in $\phi$ and therefore plays no role in the convexity of $J$. We therefore set

[TABLE]

a convex functional which essentially has the same convexity properties as $J$ since they only differ by a linear term. For instance one can check directly that for any potentials $\phi_{1}$ and $\phi_{2}$ we have

[TABLE]

Finally, the convex conjugate of $F$ is defined here by

[TABLE]

We are at this point ready to state the next lemma which is at the core of obtaining stability estimates in the gradient schemes.

Lemma 5**.**

Let $\phi_{1}$ and $\phi_{2}$ be two twice-differentiable real-valued functions defined on $\Omega$ such that

[TABLE]

for $i=1,2$, for some $\lambda>0$. Moreover, let $\rho_{i}=T_{\phi_{i}\#}\mu$ the associated mass densities, for $i=1,2$. Then

[TABLE]

Additionally,

[TABLE]

Proof.

Step 1. We start the proof by establishing the equality

[TABLE]

To do so, first note that a simple variant of Fact 3 tells us that the derivative of $F$ is precisely

[TABLE]

Therefore the Bregman divergence associated to $F$ can be written

[TABLE]

The above expression can be simplified using the identity

[TABLE]

which follows immediately from the definitions of the $c$-transform and of $T_{\phi}$. Consequently,

[TABLE]

To proceed further we make use of another equality,

[TABLE]

which is again a direct consequence of the $c$-transform definition. Adding the term -\big{(}T_{\phi_{2}}(x)-x\big{)}\cdot\big{(}T_{\phi_{1}}(x)-T_{\phi_{2}}(x)\big{)}+\nabla\phi_{2}(T_{\phi_{2}}(x))\cdot\big{(}T_{\phi_{1}}(x)-T_{\phi_{2}}(x)\big{)} (which is thus always [math]) to the expression of $F(\phi_{2}|\phi_{1})$ is enough to obtain the desired equality (19).

Step 2. We now establish the inequality

[TABLE]

where $\rho_{1}=T_{\phi_{1}\#}\mu$. Set

[TABLE]

which is a $\lambda$-strongly convex function because of the assumption made on $\phi_{2}$. Then, expression (19) can be written

[TABLE]

By Lemma 4iii we derive the upper bound

[TABLE]

for any $t_{1},t_{2}\in\Omega$. Employing again identity (21) we see that on the one hand

[TABLE]

and on the other hand \nabla\xi_{2}\big{(}T_{\phi_{2}}(x)\big{)}=x. We deduce the upper bound

[TABLE]

which after a change of measure implies the desired inequality (22).

Step 3. The next part of the proof consists of estimating the $L^{\infty}$ norm of $\rho_{1}$,

[TABLE]

The strong convexity of $\phi_{1}$ implies that the map $T_{\phi_{1}}$ is injective. Since $\rho_{1}$ is the (density of) the pushforward measure $T_{\phi_{1}\#}\mu$, we have the change of variable formula

[TABLE]

which is valid for injective and differentiable maps. We now show that $DT_{\phi_{1}}$ is a symmetric positive-definite matrix whose eigenvalues are bounded below by $\lambda$. To do so, set

[TABLE]

By the assumption made on $\phi_{1}$ we have $D^{2}\xi_{1}(x)\leq\lambda^{-1}$. Making use of identity (21), we can write x=\nabla\xi_{1}\big{(}T_{\phi_{1}}(x)\big{)}, which can be inverted into $\nabla\xi_{1}^{*}(x)=T_{\phi_{1}}(x)$. Since the Hessian of $\xi_{1}$ has eigenvalues bounded above by $\lambda^{-1}$, its convex conjugate is $\lambda$-strongly convex. Therefore

[TABLE]

for any $x\in\Omega$. The bound (23) directly follows, and combining it with (22) concludes the proof of (17).

Step 4. Now that the upper bound (17) is proven, it directly implies the lower bound (18). Indeed, rewrite (17) as

[TABLE]

with $h=\phi_{2}-\phi_{1}$ and $C=\lVert\mu\rVert_{L^{\infty}}\lambda^{-(d+1)}$. Taking convex conjugation on both sides, which reverses the sign of inequalities, yields

[TABLE]

Indeed it is easy to see that $\frac{1}{2}\lVert\cdot\rVert_{\dot{H}^{1}}^{2}$ and $\frac{1}{2}\lVert\cdot\rVert_{\dot{H}^{-1}}^{2}$ are conjugate to each other. We then apply the Bregman divergence property iv from Lemma 4 to the functional $F$ and obtain

[TABLE]

where $\rho_{1}=\delta F_{\phi_{1}}=T_{\phi_{1}\#}\mu$. This is precisely the desired inequality when $u=\rho_{2}-\rho_{1}$. ∎

We have now at our disposition all the necessary tools to prove Proposition 1.

Proof of Proposition 1.

Ascent property. Consider two consecutive iterates $\phi_{n}$ and $\phi_{n+1}$. By Fact 3 the derivative of $J$ takes the form $\delta J_{\phi}=\nu-T_{\phi\#}\mu$, which implies that

[TABLE]

Here we introduced a Bregman divergence $J(\cdot|\cdot)$ (see Definition 4 in this appendix). Let $F$ be defined by (15), then as previously noted in (16) we have $J(\phi_{n+1}|\phi_{n})=-F(\phi_{n+1}|\phi_{n})$. Additionally the gradient step $\phi_{n+1}-\phi_{n}=\sigma(-\operatorname{\Delta})^{-1}(\nu-T_{\phi_{n}\#}\mu)$ implies that

[TABLE]

Combining these two expressions we arrive at

[TABLE]

and the choice $\sigma=\frac{\lambda^{d+1}}{\lVert\mu\rVert_{L^{\infty}}}$ together with Lemma 5 stated and proven right above is enough to obtain the ascent property

[TABLE]

Decrease in $H^{-1}$ norm. A simple computation, expanding and rearranging quadratic terms, shows that

[TABLE]

Since the iterate $\phi_{n+1}$ is defined by $\phi_{n+1}=\phi_{n}+\sigma(-\operatorname{\Delta})^{-1}(\nu-\rho_{n})$,we obtain

[TABLE]

Next, as was explained by (20) the derivative of the functional $F$ defined by (15) is $\delta F_{\phi}=\rho$, with $\rho=T_{\phi\#}\mu$, and since $F$ is a convex functional this relation can be inverted into

[TABLE]

where $F^{*}$ denotes the convex conjugate of $F$ defined by $F^{*}(\rho)=\sup_{\phi}\int\phi\,\rho-F(\phi)$. Thus we can write

[TABLE]

By Lemma 5 we have $F^{*}(\rho_{n+1}|\rho_{n})+F^{*}(\rho_{n}|\rho_{n+1})\geq\sigma\lVert\rho_{n+1}-\rho_{n}\rVert_{H^{-1}}^{2}$, which is enough to conclude the proof. ∎

Proof of Proposition 2.

Fix an iterate $\phi_{n}$. Then $J(\phi_{n})=D(\phi_{n},\phi_{n}^{c})$. By Proposition 1 we have

[TABLE]

since the step $\phi_{n}\to\phi_{n+\frac{1}{2}}$ is a gradient ascent step of $J$; also note that $J(\phi_{n+\frac{1}{2}})=D(\phi_{n+\frac{1}{2}},(\phi_{n+\frac{1}{2}})^{c})=D(\phi_{n+\frac{1}{2}},\psi_{n+\frac{1}{2}})$. Next, it is easy to see that the value of $D$ always increases (or plateaus) when taking $c$-transforms, thus

[TABLE]

Therefore, we have shown that $I(\psi_{n+\frac{1}{2}})\geq J(\phi_{n+\frac{1}{2}})$; the remaining inequalities can be treated in a similar fashion. ∎