Scaling Algorithms for Unbalanced Transport Problems

Introduction

Optimal transport (OT) is a standard way to lift a metric defined on some “ground” space $X$ to a metric on probability distributions (positive Radon measures with unit mass) $\mathcal{M}_{+}(X)$. Initially formulated by Monge as a non-convex optimization problem over transport maps, its modern formulation as a linear program is due to Kantorovitch , and it has been revitalized thanks to the groundbreaking work of Brenier . We refer to the monographs for a more detailed background on the theory of OT.

While initially developed by theoreticians, OT is now becoming popular in applied fields, and we refer for instance to application for color manipulation in image processing , reflectance interpolation in computer graphics , image retrieval in computer vision and statistical inference in machine learning . A key limitation of classical OT is that it requires the input measures to be normalized to unit mass, which is a problematic assumption for many applications that require either to handle arbitrary positive measures and/or to allow for only partial displacement of mass. All these applications might indeed benefit from OT algorithms that can handle mass variation (creation or destruction) as well as mass transportation.

While many proposals have been made to account for these “unbalanced” optimal transport problems, they used to be tailored for specific applications. There have been several recent works (reviewed below) that try to put all these proposals (and some new ones) into a common generic framework. These emerging theoretical advances call for algorithms that extend existing fast OT methods to these new unbalanced problems. It is precisely the goal of this paper to show how a popular numerical approach to OT – namely entropic regularization – does extend in a very natural and efficient way to solve a variety of unbalanced problems, including unbalanced OT, barycenters and gradient flows.

The Kantorovitch formulation of OT as a linear program, when restricted to sums of Diracs, can be directly tackled using simplex or interior point methods. For the special case of optimal linear assignment (i.e. for sums of the same number of uniform-mass Diracs) one can also use combinatorial algorithms such as the Hungarian method or the auction algorithm . The time complexity of these algorithms is roughly cubic in the number of Diracs, and hence they do not scale to very large problems.

In the specific case of the squared Euclidean cost, it is possible to make use of the geodesic structure of the OT distance, and reparametrize this as a convex optimization problem, as proposed by Benamou and Brenier , see also for a discussion on the use of first order non-smooth optimization schemes. In a multi-scale algorithm is developed that consistently leverages the structure of geometric transport problems to accelerate linear program solvers.

A last class of approaches deals with semi-discrete problems, when one measure has a density, and the second one is a weighted sum of Dirac masses. This problem, introduced by Alexandrov and Pogorelov as a theoretical tool, can be solved with geometric tools when using the squared Euclidean cost (see also for the development of efficient algorithms using methods from computational geometry and for 3-D computations).

These computational methods thus cannot cope with large scale problems with arbitrary transportation costs. A recent class of approaches, initiated and revitalized by the paper of Marco Cuturi , proposes to compute an approximate transport coupling using entropic regularization. This idea has its origins in many different fields, most notably it is connected with Schrödinger’s problem in statistical physics and with the iterative scaling algorithm by Sinkhorn (also known as IPFP ) which, given a square matrix with positive entries, aims at finding two vectors of positive numbers—so-called scalings—that makes it a bistochastic matrix after multiplying rows and columns by these vectors. This entropic smoothing can be interpreted as a strictly convex barrier for positivity, but its main computational advantage is that it leads to very simple closed form expressions for all steps of the algorithm, which would not be possible when using different regularization functionals. Several follow-up articles to have shown that the same strategy can also be used to tackle the computation of barycenters for the Wasserstein distance (as initially formulated by ), and for solving OT problems on geometric domains (such as regular grids or triangulated meshes) using convolution and the heat diffusion kernel . Some theoretical properties of this regularization are studied in , including the $\Gamma$-convergence of the regularized problem toward classical transport when regularization vanishes.

The success of this entropic regularization scheme is tightly linked with use of the Kullback-Leibler (KL) divergence as a natural Bregman divergence for the optimization on the space of positive Radon measures. Not only is this divergence quite natural, but it also leads to simple formulas for the computation of projectors and so-called proximal operators (see below for a definition) for many functions typically involved in OT. The most simple algorithm, which is actually at the heart of Sinkhorn’s iterations, is the iterative projection on affine subspaces for the KL divergence. A refined version of these iterations, which works for arbitrary convex sets (not just affine spaces) is the so-called Dykstra’s algorithm , which can be interpreted (just like iterative projections) as an iterative block-coordinates minimization on a dual problem. Dykstra’s algorithm is known to converge when used in conjunction with Bregman divergences (see for details on the underlying idea for sums of two arbitrary functions). Many other first order proximal methods for Bregman divergences exists. The most simple one is the proximal point algorithm , but most proximal splitting schemes have been extended to this setting, such as for instance ADMM , primal-dual schemes and forward-backward .

The algorithm we propose in this article can be seen as special instance of Dykstra’s iterations, but with an extremely simple (both conceptually and algorithmically) structure, which we refer to as a “scaling” algorithm. It extends Sinkhorn’s iterations to more complex problems. This structure is due to the fact that the functions involved in OT problems make use of the marginals of the couplings that are being optimized.

There has been a large number of proposals to extend OT methods to arbitrary “unbalanced” positive measures. Let us for instance quote the Kantorovitch-Rubinstein dual-Lipschitz norms , optimal partial transport and geodesic computations with source terms . Most of these approaches, and much more, can be seen as special instances of a generic class of OT-like problems, that have been proposed independently in and . These unbalanced problems can be formulated in several ways, that are equivalent (under some restrictive conditions on the cost): a geodesic (dynamic) formulation with a source term , a static formulation with two semi-couplings and a static formulation with approximate marginal constraints . From a numerical perspective, the last formulation (approximate marginals constraints) is the most simple to handle, since, as we show in the following section, it only involves a minor modification of the initial linear program, that has to be turned into a convex problem involving $\varphi$-divergences (see Definition 2.2). This is the one that we consider in this article. Note that the use of such a relaxed formulation, in conjunction with entropic smoothing has been introduced, without proof of convergence, in for application in machine learning.

Beyond OT problems and barycenter problems, a popular use of Wasserstein distances is to study gradient flows, following the formalism of “minimizing movements” detailed in . This corresponds to discrete implicit stepping (i.e. proximal maps) for the Wasserstein distance (instead of more common Euclidean or Hilbertian metrics). In some cases, these time-discrete flows can be shown to converge to continuous flows that solve a suitable PDE, as the temporal stepsize tends to 0. The most famous example is the gradient flow of the entropy for the Wasserstein metric, which solves the diffusion equation . Non-linear PDEs are considered for example in . An application to imaging can be found in . Another use of these implicit steps is to construct minimizing flows of non-smooth functionals, for instance to model crowd motions .

A large variety of dedicated numerical schemes has been proposed for spatial discretization and solving of these time-discrete flows, such as for instance finite differences , finite volumes and Lagrangian schemes . A bottleneck of these schemes is the high computational complexity due to the resolution of an OT-like problem at each time step. The use of entropic regularization has been proposed recently in and studied theoretically in . A chief advantage of this approach is that each step can be solved efficiently on a regular grid with fast Sinkhorn iterations involving only Gaussian convolutions, but this comes at the price of additional diffusivity introduced by the approximation, which makes it unsuitable to capture sharp features of solutions.

It is possible to extend these gradient flows by replacing the Wasserstein distance by more general unbalanced distances. This allows to define flows over arbitrary positive measures, hence involving creation and destruction of mass, which is crucial to model growth phenomena, such as for instance the Hele-Shaw model of tumor evolution . An analysis of such flows based on a splitting scheme has been recently provided in . It is one of the goals of this article to propose a versatile algorithm, based on iterative scalings, to approximate numerically these unbalanced flows.

2. Contributions and Outline

The main contribution of this article is to define a class of iterative scaling algorithms to solve the entropic approximation of a variety of unbalanced optimal transport problems. First, in Section 2, we exhibit a common structure to most optimization problems related to OT: a nonnegativity constraint, a linear transport cost and convex functions acting on the marginals of the optimized couplings. For solving the entropic regularization of these problems, we then introduce in Section 3 a generic “scaling” algorithm which is a direct generalization of Sinkhorn’s algorithm. In a continuous setting, we show under some assumptions that the iterates are well defined and correspond to alternating optimization on the dual and we prove linear convergence in a particular key case. In a discrete setting, we show in Section 4 that this algorithm converges as soon as the dual problem is well-posed. We also give a simple description of the algorithm, introduce a stabilization scheme to reach very small values of the regularization parameter and sketch a generalization of this algorithm. Finally, in Section 5, we showcase the application of these methods to 1-D and 2-D unbalanced optimal transport, generalizations of Wasserstein barycenters and gradient flows with growth. The main advantages of our approach are that it is simple (it only involves matrix multiplication and pointwise elementary operations), quite generic (it applies to most known OT-like problems), enjoys fast convergence (linear convergence is observed, and show in particular cases) and is highly parallelizable. The code to reproduce the results of this article is available onlinehttps://github.com/lchizat/optimal-transport. Note also that an efficient numerical implementation of the scaling algorithm developed in this paper is studied further in , see Section 4.4 for a discussion.

3. Notation

The space of nonnegative finite Radon measuresBorel measures which are inner regular. If $T$ is Polish, all Borel measures are inner regular. on a (Hausdorff) topological space $T$ is denoted by $\mathcal{M}_{+}(T)$, and the vector space it generates by $\mathcal{M}(T)$. For a measure on a product space $\mu\in\mathcal{M}(X\times Y)$, $P^{X}_{\#}\mu$ (or sometimes $P^{1}_{\#}\mu$ if $X=Y$) denotes its first marginal and $P^{Y}_{\#}\mu$ (or $P^{2}_{\#}\mu$) its second marginal.

Divergence functionals (or $\varphi$-divergences) are denoted by a calligraphic letter $\operatorname{\mathcal{D}}$ when they act on measures (as defined in Definition 2.2), by straight letters $\operatorname{D}$ when they act on functions (as defined in (5.2)) and with an overline $\overline{\operatorname{D}}$ when they act on two real numbers (as in (5.2)). More generally, functionals on measures are denoted by calligraphic letters ($\mathcal{F},\mathcal{G},\dots$) and functionals on functions by straight capital letters ($F,G,\dots$).

(with $0\log(0/0)=0$) if $r,s\geqslant 0$ a.e. and $(s_{i}(x,y)=0)\Rightarrow(r_{i}(x,y)=0)$ a.e., and $\infty$ otherwise.

The generalization of some notations to families of functions is often implicit. For instance, if $(r_{i})_{i=1}^{n}$ and $(s_{i})_{i=1}^{n}$ are two families of functions, we write

and the projection operators are defined componentwise, for $i\in\{1,\dots,n\}$, as

If $X$ and $Y$ are finite spaces (i.e. contain only a finite number of points), we represent functions on $X$ by vectors denoted by bold letters $\mathbf{a}$ and functions on $X\times Y$ by matrices denoted by capital bold letters $\mathbf{A}$. In this context the notations $\mathbf{a}\odot\mathbf{b}$ and $\mathbf{a}\oslash\mathbf{b}$ denote, respectively entrywise multiplication and entrywise division with convention $0/0=0$ between vectors.

The conjugate of a convex function $f$ is denoted by $f^{*}$ and its subdifferential is denoted by $\partial f$. Some reminders on convex analysis are given in Appendix A.1. The indicator of some convex set $C$ is

Unifying Formulation of Transport-like Problems

We review below a variety of variational problems related to optimal transport that can all be recast as a generic variational problem, involving transport and functions on the marginals. The notion of “divergence” functionals, introduced in optimal transport by , makes this unification possible and is defined in the following.

Divergences are functionals which, by looking at the pointwise “ratio” between two measures, give a sense of how close they are. They have nice analytical and computational properties and are built from entropy functions.

If $\varphi^{\prime}_{\infty}=\infty$, then $\varphi$ grows faster than any linear function and $\varphi$ is said superlinear. Any entropy function $\varphi$ induces a $\varphi$-divergence (also known as Csiszár divergence) as follows.

if $\mu,\nu$ are nonnegative and $\infty$ otherwise.

The proof of the following Proposition can be found in [47, Thm 2.7].

If $\varphi$ is an entropy function, then $\operatorname{\mathcal{D}}_{\varphi}$ is jointly $1$-homogeneous, convex and weakly* lower semicontinuous in $(\mu,\nu)$.

The Kullback-Leibler divergence, also known as the relative entropy, plays a central role in this article.

The Kullback-Leibler divergence $\operatorname{\mathcal{KL}}\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}\operatorname{\mathcal{D}}_{\varphi_{\operatorname{KL}}}$ is the divergence associated to the entropy function $\varphi_{\operatorname{KL}}$, given by

When restricted to densities on a measured space, it is also the Bregman divergence associated to (minus) the entropy.

The following divergences will also be considered as examples.

The total variation distance $\operatorname{TV}(\mu|\nu)\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}|\mu-\nu|_{\operatorname{TV}}$ is the divergence associated to:

The equality constraint $\iota_{\{=\}}(\mu|\nu)$ which is if $\mu=\nu$ and $\infty$ otherwise, is the divergence associated to $\varphi=\iota_{\{1\}}$.

One can generalize the latter and define a “range constraint”, denoted $\operatorname{RG}_{[\alpha,\beta]}(\mu|\nu)$, as the divergence which is zero if $\alpha\,\nu\leqslant\mu\leqslant\beta\,\nu$ with $0\leqslant\alpha\leqslant\beta\leqslant\infty$, and $\infty$ else. This is the divergence associated to $\varphi=\iota_{[\alpha,\beta]}$.

2. Balanced Optimal Transport

Using divergences, the classical “balanced” optimal transport problem can be defined as follows.

If $c$ is lower bounded and (2.3) is feasible, then the infimum is attained.

An important special case is when $X=Y$ and $c$ is the power of a distance on $X$. Then, the optimal cost (i.e. the value of (2.3)) is itself the power of a distance on the space of probability measures on $X$. This distance is often referred to as “Wasserstein” distance and denoted by $\operatorname{W}$, although this denomination is disputed. We refer to [77, Chap. 6, Bibliographical Notes] for a discussion of the historical context.

3. Unbalanced Optimal Transport

Standard optimal transport only allows meaningful comparison of measures with the same total mass: whenever $\mu(X)\neq\nu(Y)$, there is no feasible $\gamma$ in (2.3). Several propositions have been made to circumvent this limitation by defining “unbalanced” transport problems in various dynamic and static formulations (see Section 1.1). The following formulation with relaxed marginal constraints is best suited for the numerical schemes presented in this article.

Let $X$ and $Y$ be Hausdorff topological spaces, let $c:X\times Y\rightarrow[0,\infty]$ be a lower semi-continuous function and let $\operatorname{\mathcal{D}}_{\varphi_{1}}$, $\operatorname{\mathcal{D}}_{\varphi_{2}}$ be two divergences over $X$ and $Y$, as in Definition 2.2. For $\mu\in\mathcal{M}_{+}(X)$ and $\nu\in\mathcal{M}_{+}(Y)$, the unbalanced soft-marginal transport problem is

Assume that (2.4) is feasible. If $\varphi_{1}$ and $\varphi_{2}$ are superlinear, then the infimum is attained. This is also the case if $c$ has compact sublevel sets and $\varphi^{\prime}_{1\infty}+\varphi^{\prime}_{2\infty}+\inf c>0$.

The proof of this result as well as a thorough study of duality properties of this problem can be found in . Note that (2.3) is a particular case of (2.4) when the domains of the entropy functions $\varphi_{i}$ are the singleton $\{1\}$. More generally, if the functions $\varphi_{i}$ admit unique minima at $1$, (2.4) can be viewed as a relaxation of the initial problem (2.3) where the “hard” marginal constraints are replaced by “soft” constraints, penalizing the deviation of the marginals of $\gamma$ from $\mu$ and $\nu$. Now we discuss a specific case of particular interest.

Take $X=Y$, let $d$ be a distance on $X$ and $\lambda\geqslant 0$. For the cost

with $\cos_{+}:z\mapsto\cos(z\wedge\frac{\pi}{2})$ and $\operatorname{\mathcal{D}}_{\varphi_{1}}=\operatorname{\mathcal{D}}_{\varphi_{2}}=\lambda\operatorname{\mathcal{KL}}$, we define $\operatorname{WFR}_{\lambda}(\mu,\nu)$ as the square root of the minimum in (2.4) (as a function of the measures $(\mu,\nu)\in\mathcal{M}_{+}(X)^{2}$). We simply write $\operatorname{WFR}\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}\operatorname{WFR}_{1}$ for $\lambda=1$.

As shown in , $\operatorname{WFR}$ defines a distance on $\mathcal{M}_{+}(X)$, which is equal to the Wasserstein-Fisher-Rao geodesic distance introduced simultaneously and independently in (it is named Hellinger-Kantorovich in ). An alternative static formulation is given in . Other important cases include:

The Gaussian Hellinger-Kantorovich distance $\operatorname{G\!H\!K}$ is obtained by taking the cost $c=d^{2}$ ($d$ still a metric) and $\operatorname{D}_{\varphi_{i}}=\lambda\operatorname{KL}$ with $\lambda>0$. It has been introduced in where it is also shown that when $X$ is a geodesic space, $\operatorname{WFR}$ is the geodesic distance generated by $\operatorname{G\!H\!K}$.

The optimal partial transport problem, which is obtained by taking a cost function bounded from below by $-2\lambda$ and $\operatorname{D}_{\varphi_{i}}=\lambda\operatorname{TV}$, with $\lambda>0$. Originally, optimal partial transport refers to a problem where the marginal constraints are replaced by a constraint on the total mass of the marginals (and domination constraints); here $\lambda$ is the Lagrange multiplier associated to the mass constraint (see ).

With the divergence $\operatorname{RG}$ as in example 2.7, one imposes a range constraint on the marginal:

The following result is a minor extension of results from and is useful if one wishes to vary the parameter $\lambda$ in numerical applications.

$\operatorname{WFR}_{\lambda}$ defines a distance if $0\leqslant\lambda\leqslant 1$ (degenerate if $\lambda=0$). The upper bound on $\lambda$ is necessary when $X$ is a geodesic space of diameter greater than $\pi/2$.

The case $\lambda=0$ is trivial. If $\lambda>0$, when dividing (2.4) by $\lambda$, one obtains the new cost $-2\log(\cos_{+}^{1/\lambda}(d(x,y)))$. But the function $f:d\mapsto\arccos(\cos^{\frac{1}{\lambda}}(d))$ defined on $[0,\pi/2]$ is increasing, positive, satisfies $\{0\}=f^{-1}(0)$ and for $x\in]0,\pi/2[$ it holds

From the convexity inequality $\text{sign}[(1-X)^{\frac{1}{\lambda}}-1+\frac{1}{\lambda}X]=\text{sign}(\frac{1}{\lambda}-1)$ it follows that $f$ is strictly concave on $[0,\pi/2]$ if $\lambda<1$ and strictly convex if $\lambda>1$. Thus if $\lambda\leqslant 1$, $f\circ d$ still defines a distance on $X$ and consequently $\operatorname{WFR}_{\lambda}$ too.

If $X$ is a geodesic space of diameter greater than $\pi/2$, take $(x_{1},x_{2})\in X^{2}$ such that $d(x_{1},x_{2})=\pi/2$. From [47, Corollary 8.3], $(\mathcal{M}(X),\operatorname{WFR}_{1})$ is itself a geodesic space. Consequently, there exists a midpoint $\mu\in\mathcal{M}_{+}(X)$, i.e. such that

From [47, Theorem 8.6] and the characterization of geodesics in $\text{Cone}(X)$ it holds $0<d(x,x_{i})<\pi/2$ for $\mu$ a.e. $x\in X$. This implies, for $\lambda>1$, and $\mu$ a.e. $x$, $-\log(\cos_{+}^{1/\lambda}(d(x_{i},x)))<-\log(\cos_{+}(d(x_{i},x)))$. Thus $(\operatorname{WFR}^{2}_{\lambda}/\lambda)(\delta_{x_{i}},\mu)<\operatorname{WFR}^{2}_{1}(\delta_{x_{i}},\mu)$. But, for $\lambda\geqslant 1$, $(\operatorname{WFR}^{2}_{\lambda}/\lambda)(\delta_{x_{1}},\delta_{x_{2}})=2\operatorname{KL}(0|1)=2$ this leads to

and the triangle inequality property is lost. ∎

4. Barycenter Problem and Extensions

The problem of finding an “average” measure $\sigma\in\mathcal{M}_{+}(Y)$ which minimizes the sum of the (possibly unbalanced) transport cost toward every measure of a family $(\mu_{k})_{k=1}^{n}\in\mathcal{M}_{+}(X)^{n}$ is of theoretical and practical interest (see Section 1.1). To formalize this problem, consider a family of costs functions $(c_{k})_{k=1}^{n}$ on $X\times Y$ and a family of divergences $(\operatorname{\mathcal{D}}_{\varphi_{k,1}},\operatorname{\mathcal{D}}_{\varphi_{k,2}})_{k=1}^{n}$. The problem is to solve

By exchanging the infima, this is equivalent to

Note that while the object of interest is the minimizer $\sigma$ and not the family of couplings, we will see in Section 5.2 that the computation of $\sigma$ is a byproduct of the “scaling” algorithm defined below.

The two following examples are specific instances of so-called Fréchet means, defined in any complete metric space $(E,d)$ as solutions to

where $(\alpha_{k})_{k=1}^{n}$ is a family of nonnegative weights and $(\mu_{k})_{k=1}^{n}\in E^{n}$.

Let $X=Y$ and let $c$ be the quadratic cost $(x,y)\mapsto d(x,y)^{2}$ for a metric $d$ and define $c_{k}=\alpha_{k}\,c$. Let all the divergences be the equality constraints. Then (2.6) is a formulation of the Fréchet means in the Wasserstein space, also known as Wasserstein barycenters (see ).

Let $c$ be as in (2.5), define $c_{k}=\alpha_{k}\,c$ and let $\operatorname{\mathcal{D}}_{\varphi_{k,1}}=\operatorname{\mathcal{D}}_{\varphi_{k,1}}=\alpha_{i}\operatorname{\mathcal{KL}}$ for all $k\in\{1,\dots,n\}$. Then (2.6) is a formulation of Fréchet means for the $\operatorname{WFR}$ distance.

5. Gradient Flows

Initiated by , the study of such flows when $T$ is the space of probability measures endowed with the Wasserstein metric $d=\operatorname{W}$ (see Section 2.2) has led to considerable advances in the theoretical study of PDEs and their numerical resolution (see Section 1.1). The recent introduction of “unbalanced transport” metrics paves the way for further applications, since it is now possible to consider the whole space of nonnegative measures, endowed for instance with the metric $\operatorname{WFR}$, as considered in .

For gradient flows based on an optimal transport metric, such as $\operatorname{W}$, $\operatorname{WFR}$ or $\operatorname{G\!H\!K}$, each step requires to solve, after swapping the two infima, a problem of the form

where $\varphi_{1}$, $\varphi_{2}$ are entropy functions and $\mathcal{G}$ is a lower semicontinuous functional which we also assume convex. This problem, as well as variants, involving so-called “splitting” techniques (see ) fit into the framework developed below. In particular, we show in Section 5.3 that the computation of $\mu_{k+1}$ is a byproduct of the minimization of (2.7) with the algorithm defined below.

The distance $\operatorname{WFR}$ was only introduced recently, so a sound theoretical analysis of the corresponding gradient flows and their limit PDEs is not yet available (and is not the subject of the present article). Meanwhile, we present here heuristic arguments in a smooth setting (see also [41, Section 3.2]) which lead to the evolution equation

When restricted to measures $\mu$ with positive density of Sobolev regularity, $\operatorname{WFR}$ is the weak Riemannian metric associated to the tensor

where, for a small variation $\delta\mu$, one searches over the decompositions $\delta\mu=-\operatorname{div}(v\mu)+\alpha\mu$ into displacement (given by the velocity field $v\in\text{L}^{2}(X,\mu)^{d}$) and growth (given by the rate of growth $\alpha\in\text{L}^{2}(X,\mu)$) (see ). A new step $\mu_{k+1}^{\tau}$ is given from $\mu^{\tau}_{k}$ through the resolution of

Searching the minimizer in the form $\mu=\mu^{\tau}_{k}+\tau(-\operatorname{div}(v\mu^{\tau}_{k})+\alpha\mu_{k}^{\tau})$ with unknown $(v,\alpha)$, this can be rewritten, in first order of $(\tau v,\tau\alpha)$, as

The first order optimality conditions yield $v=-\nabla\mathcal{G}^{\prime}(\mu)$ and $\alpha=-4\mathcal{G}^{\prime}(\mu)$. One thus obtains

6. Generic Formulation

Consider two convex and lower semicontinuous functions $\mathcal{F}_{1}$ and $\mathcal{F}_{2}$ defined on $\mathcal{M}_{+}(X)^{n}$ and $\mathcal{M}_{+}(Y)^{n}$ respectively. The variety of problems reviewed above can be seen as special cases of

Balanced OT (2.3): $\mathcal{F}(\sigma)=\iota_{\{=\}}(\sigma|\mu)$;

Unbalanced OT (2.4): $\mathcal{F}(\sigma)=\operatorname{\mathcal{D}}_{\varphi}(\sigma|\mu)$;

Barycenters (2.6): $\mathcal{F}(\sigma)=\inf_{\omega\in\mathcal{M}^{n}(X)}\left\{\sum_{k=1}^{n}\operatorname{\mathcal{D}}_{\varphi_{k}}(\sigma_{k}|\omega_{k})+\iota_{D}(\omega)\right\}$, where $D$ is the set of families of measures for which all components are equal,

Gradient flows (2.7): $\mathcal{F}(\sigma)=\inf_{\omega\in\mathcal{M}(X)}\left\{\operatorname{\mathcal{D}}_{\varphi}(\sigma|\omega)+2\tau\mathcal{G}(\omega)\right\}$.

It is remarkable that, even if $\mathcal{F}$ is sometimes defined through an auxiliary minimization problem, it can still be handled efficiently in some cases by the scaling algorithm detailed below. For instance, functions of the form

where $\mathcal{G}$ is convex and $(\varphi_{k})_{n}$ is a family of entropy functions, are still convex and can model a great variety of problems. A graphical interpretation of these problems is suggested in Figure 1 along with some examples.

Entropic Regularization and Iterative Scaling Algorithm

In this section, we introduce and describe an iterative scaling algorithm which solves a regularized version of the generic variational problem (2.9). This analysis is carried out in a continuous setting.

Recalling that $\mathcal{J}(\gamma)=\langle c,\gamma\rangle+\mathcal{F}_{1}(P^{X}_{\#}\gamma)+\mathcal{F}_{2}(P^{Y}_{\#}\gamma)\,$, this can be rewritten, up to a constant, as

Adding the entropy term can be interpreted in the following ways, detailed for $n=1$ for simplicity.

where $\mathcal{H}(\bar{\gamma})$ is finite since it is a finite sum of real numbers. As $\mathcal{H}$ is coercive and lower semicontinuous, this set is compact. This implies the existence of cluster points for $(\gamma_{k})$: let $\gamma^{*}$ be one of them. One has that $\gamma^{*}$ belongs to $A$ since for all $k$, $|\mathcal{J}(\gamma_{k})-\mathcal{J}(\bar{\gamma})|\leqslant 2\varepsilon_{k}\mathcal{H}(\bar{\gamma})\to 0$ and $\mathcal{J}(\gamma^{*})\leqslant\liminf_{k\to\infty}\mathcal{J}(\gamma_{k})$. Moreover, as $\mathcal{H}(\gamma^{*})\leqslant\mathcal{H}(\bar{\gamma})$ and $\bar{\gamma}$ is arbitrarily chosen in $A$, it holds $\gamma^{*}\in\operatorname*{argmin}_{\gamma\in A}\mathcal{H}(\gamma)$. By strict convexity, this cluster point is unique and $\gamma_{k}\to\gamma^{*}$. ∎

Beyond computational aspects, for the specific case of standard OT with quadratic cost, the entropic regularization admits several interpretations as a stochastic version of OT. For instance, it is obtained by considering an optimal matching problem where there is (a specific model of) unknown heterogeneity in the preferences within each class to be matched . What is more, it amounts to computing the law of motion of the (indistinguishable) particules of a gaz which follow a Brownian motion, conditionally to the observation of its density at times $t=0$ and $t=1$, the so-called Schrödinger bridge problem .

Also, for the barycenter (2.6) or the gradient flow problems (2.7), one does not have access in general to a reference measure according to which an optimizer is absolutely continuous. For numerics, the choice of the reference measures then corresponds to the choice of a discretization grid on which we find approximate solutions to the original problem.

2. Reformulation using Densities and Duality

for $k\in\{1,\dots,n\}$, and with the convention $\exp(-\infty)=0$, the projection operator acts on each component of $r$ and $\operatorname{KL}$ is the sum of the Kullback-Leibler divergences on each component (see the notations in Section 1.3). In this Section, we only make the following general assumptions on the objects involved in ($P_{\varepsilon}$):

$K\in\text{L}^{\infty}_{+}(X\times Y)^{n}$ and $\varepsilon>0$.

We begin with a general duality result, which is an application of Fenchel-Rockafellar Theorem (see Appendix A.1). It is similar to duality results that can be found in the literature on entropy minimization but the functional we consider is more general.

The dual problem of ($P_{\varepsilon}$) is

where $u\oplus v:(x,y)\mapsto u(x)+v(y)$. Strong duality holds, i.e. $\min\eqref{eq-general-regul-KL}=\sup\eqref{eq-dual-pbm}$ and the minimum of ($P_{\varepsilon}$) is attained for a unique $r=(r_{k})_{k=1,\dots,n}\in\text{L}^{1}(X\times Y)^{n}$. Moreover, $u$ and $v$ maximize ($D_{\varepsilon}$) if and only if

are equal, the latter being exactly ($P_{\varepsilon}$) since $G$ is infinite outside of $\text{L}^{1}(X\times Y)^{n}$. It states also that if $(u,v)$ maximizes ($D_{\varepsilon}$), then any minimizer of (3.1) satisfies $r\in\partial G^{*}(A(u,v)/\varepsilon)$ and the expression for the subdifferential of $G^{*}$ is an application of the result in Appendix A.2. Finally, uniqueness of the minimizer for ($P_{\varepsilon}$) comes from the strict convexity of $G$. ∎

3. Scaling Algorithm

The specific splitting of the problem ($P_{\varepsilon}$) makes it suitable for the well-known Dykstra’s algorithm (see Section 1.1 for more background on this algorithm). Since the functions $F_{1}$ and $F_{2}$ operate on the marginals only, Dykstra’s iterations take a very simple form which is related to the celebrated Sinkhorn algorithm. They are obtained as an alternating maximization on ($D_{\varepsilon}$).

where we used the fact that, by Fubini-Tonelli, one has

where the proximal operator for the $\operatorname{KL}$ divergence is defined for $F_{1}$ (and similarly for $F_{2}$) as

The following Proposition shows that, as long as they are well defined, these iterates are related to the alternate dual maximization iterates.

The proof of this proposition makes use of the following Lemma.

4. Existence of the iterates for integral functionals

Our next step is to give conditions on $F_{1}$ and $F_{2}$ that guarantee the existence of the scaling iterates (S) and an equivalence with alternate maximization on the dual (3.6). This is provided by Theorem 3.8, where it is required that $F_{1}$ and $F_{2}$ are integral functionals, as we define now.

where $f$ is a normal integrand and $f(x,\cdot)$ is convex for all $x\in X$. In this paper, $F$ is an admissible integral functional if moreover for all $x\in X$, $f(x,\cdot)$ takes nonnegative values, has a domain which is a subset of $[0,\infty[^{n}$ and if there exists $s\in\text{L}^{1}(X)^{n}$ such that $I_{f}(s)<\infty$.

The concept of normal integrands allows to deal conveniently with measurability issues. For finite dimensional problems (when $X$ and $Y$ have a finite number of points), integral functionals are simply sums of pointwise lower semicontinuous functions. The following proposition shows that for such functionals, conjugation and subdifferentiation can be performed pointwise.

If $F$ is an admissible integral functional associated to the convex normal integrand $f$, then $F$ is convex and weakly lower semicontinuous, $f^{*}$ is also a normal convex integrand, $F^{*}=I_{f^{*}}$ and

where conjugation and subdifferentiation on $f$ are w.r.t. the second variable.

This property can be found in under the assumption of existence of a feasible point $s^{\star}\in\text{L}^{1}(X)^{n}$ for $I_{f}$ and a feasible point $u^{\star}\in\text{L}^{\infty}(X)^{n}$ for $I_{f^{*}}$. Our admissibility criterion requires the existence of $s^{\star}$ and one has

since $\inf_{s}f(x,s)\in[0,f(x,s^{\star}(x))]$. ∎

The function $(x,z)\mapsto\operatorname{\overline{\operatorname{KL}}}(z|s(x))$ is a convex normal integrand by [65, Prop. 14.30 and 14.45c]. Thus $g$ is itself a normal convex integrand, as the sum of normal convex integrands [65, Prop. 14.44]. Then a minimization interchange result [65, Thm. 14.60]] states that minimizing $I_{g}$ is the same as minimizing $g$ pointwise. ∎

By Proposition 3.6, if $F_{1}$ and $F_{2}$ are admissible integral functionals then $F^{*}_{1}$ and $F^{*}_{2}$ are also integral functionals. So the alternating optimization on ($D_{\varepsilon}$) can be relaxed to the space of measurable functions and still have a meaning:

The following Theorem gives existence, uniqueness of this iterates and a precise relation with the scaling iterates (S).

5. Convergence Analysis

This Section gives a fixed point result and a convergence result for a particular case. The finite dimensional case is postponed to the next Section.

The following proposition sheds some light on the name “scaling” given to iterations (S). It comes from the fact that these iterations allow to recover a solution to ($P_{\varepsilon}$) by multiplying the kernel $K$ with positive functions, interpreted as scalings.

Under the assumptions of Theorem 3.8, if the scaling iterations (S) admit a fixed point $(a,b)$ such that $\log a\in\text{L}^{\infty}(X)^{n}$ and $\log b\in\text{L}^{\infty}(Y)^{n}$ then $(\varepsilon\log a,\varepsilon\log b)$ is the unique solution of ($D_{\varepsilon}$) and the function $r$ defined for each $k=1,\ldots,n$ by $r_{k}(x,y)=a_{k}(x)K_{k}(x,y)b_{k}(y)$ is the unique solution of ($P_{\varepsilon}$).

As a consequence of Proposition 3.7, on can write the optimality condition of a fixed point of (3.9) for almost every $x\in X$ as

for some $u(x)\in-\partial f_{1}(x,a(x)\cdot\mathcal{K}b(x))$. Thus $-\varepsilon\log(a)\in\partial I_{f_{1}}(P^{X}_{\#}r)$ because $P^{X}_{\#}r=a\cdot\mathcal{K}b$ (the dot denotes componentwise multiplication). Similar derivations for $b$ show that the couple $(\varepsilon\log a,\varepsilon\log b)$ and $r$ satisfies the primal dual optimality conditions (3.4). ∎

Sinkhorn’s algorithm (the special case of the scaling iterations (S) obtained when $F_{1}=\iota_{\{p\}}(s)$ and $F_{2}=\iota_{\{q\}}(s)$ are the convex indicators of equality constraints) is known to converge at a linear rate. This property is usually shown (see for instance ) by using the contraction property of the operator $\mathcal{K}$ for the Hilbert metric (a projective metric on the cone of nonnegative functions). Using a similar approach, based on the related Thompson metric, we now show the linear convergence of the scaling iterates (S) in another special case, which is central in unbalanced optimal transport: when $F_{1}$ and $F_{2}$ are $\operatorname{KL}$ divergences with respect to fixed densities.

An approach involving the Hilbert metric would still be possible, but the use of the Thompson metric allows for a very short proof and a convergence rate which does not depend on a bound on the cost. This metric is defined as follows.

\text{L}^{\infty}_{+} ). For $r,s\in\text{L}^{\infty}_{+}(X)$, let $M(r/s)\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}\inf\left\{\alpha\geqslant 0\;;\;r\leqslant\alpha s\right\}$ (or $\infty$ if that set is empty). The Thomson part metric is defined by

if $T:\text{L}^{\infty}(X)\to\text{L}^{\infty}(Y)$ is a $z$-homogeneous order-preserving (or order-reversing) operator, then $d(Tx,Ty)\leqslant|z|\cdot d(x,y)$;

the equivalence relation $r\sim s\Leftrightarrow d(r,s)<+\infty$ generates a partition of $(\text{L}^{\infty}(X),d)$ and each part is a complete metric space. In particular, the set $\left\{s\in\text{L}^{\infty}_{+}(X)\;;\;\log s\in\text{L}^{\infty}(X)\right\}$ endowed with the Thompson metric form a complete metric space.

Let $p$ and $q$ be such that $\log p\in\text{L}^{\infty}(X)$ and $\log q\in\text{L}^{\infty}(Y)$ and define

Let $z_{i}\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}\lambda_{i}/(\lambda_{i}+\varepsilon)\in\,]0,1[$ for $i\in\{1,2\}$ and let $b^{(0)}=1$. Following a simple application of Proposition 3.7 (or see Table 1) the iterates in this specific case read

Our assumptions are such that $d(b^{(1)},b^{(0)})$ and $d(a^{(1)},a^{(0)})$ are finite (this is direct since the logarithms of $b^{(0)}$, $K$, $p$ and $q$ are bounded). Using the properties of the Thompson metric, it holds

Algorithm for Discrete Measures

In this section, we take a closer look at the scaling iterations (S) in the specific case of discrete finite spaces. We give a convergence result, we adapt the notation to obtain an algorithm that can be implemented in a straightforward manner and also discuss how to numerically stabilize the algorithm for small values of the regularization parameter $\varepsilon$, to reach a higher precision.

In finite dimensionwhen $X$ and $Y$ are finite sets, or when the reference measures are finite sums of atoms., the scaling iterations (S) converge in general. Note that the given convergence rate is pessimistic compared to that observed in practice (which is linear, see Figure 5).

In order to check (i), take any feasible point $s$ for $F$. The Hölder inequality $F(s)+F^{*}(-(u,v)/\varepsilon)\geqslant\langle s,-(u,v)/\varepsilon\rangle$ gives

As a consequence, the components of $A(u,v)$ are upper bounded on the set of dual variables $(u,v)$ satisfying $I(u,v)\leqslant I(u^{(1)},v^{(0)})$ and so $I$ is uniformly Lipschitz on this set. In order to check (ii), remark that a primal minimizer $r^{*}$ exists and since $K$ takes positive values, one can build a minimizer by simply inverting the primal-dual relationship $(u^{*},v^{*})\in A^{-1}(\varepsilon\log r^{*}/K)$.

It remains to prove the first part of the inequality of the Theorem. We follow the strategy developed in [32, Thm. 1]. By direct computations (standard for Bregman divergences), one has for any pair $(z,z^{*})$ in the domain of $G^{*}$, by defining $r(z)(x,y)=K(x,y)\exp(z(x,y))$,

which is similar to a strict convexity estimate. Moreover, for any $m\in\partial F^{*}(-(u^{*},v^{*})/\varepsilon)$, one has by definition of the subdifferential

2. Numerical Algorithm for Discrete Measures

We now adopt a more “implementation oriented” point of view, which leads to algorithms which are straightforward to implement. For the sake of clarity, we focus on the case of a single unknown coupling (i.e. $n=1$, the extension to $n>1$ merely requires to add an index).

Given this function, which is easy to compute in many cases (see e.g. Table 1), the scaling algorithm is straightforward to implement.

By virtue of Theorem 4.1, Algorithm 1 is guaranteed to stop and to return an approximate minimizer of (4.1) if one uses a consistant stopping criterion.

3. Log-domain Stabilization

and remark that it holds, after direct computations,

We adopt the same notation as in (4.2) because it is just the special case when $\mathbf{u}=0$. The main numerical algorithm thus obtained is displayed in Algorithm 2.

The issue of extreme numerical values is not completely remedied by the absorption steps in Algorithm 2 since the $\operatorname{proxdiv}$ operation still involves the potentially extreme factor $e^{-\mathbf{u}/\varepsilon}$. In practice however, we find that for many problems $\operatorname{proxdiv}$ can be computed without evaluating the exponential $e^{-\mathbf{u}/\varepsilon}$ and the formula remains numerically stable in the limit of small $\varepsilon$. Several examples for this are given in Section 5.

4. Comments on Implementation

We wish to emphasize the simplicity of Algorithm 1 and Algorithm 2. When an optimization problem of the form (2.9) is given one just has to (i) choose the reference measures (which also determines a discretization grid in practice) (ii) determine the functions $F_{i}$ by going to the space of densities and (iii) find a way to efficiently compute $\operatorname{proxdiv}_{F_{i}}$ or $\operatorname{proxdiv}_{F_{i}}$ (in many cases, this operator has a closed form, or can be computed with a few parallelizable iterations). Let us briefly discuss some details on the practical implementation of Algorithm 2.

When solving for a very small $\varepsilon$, most of the entries of $\mathbf{K}$ are below machine precision, so one need to first “estimate” the dual variables $(\mathbf{u},\mathbf{v})$ by performing several iterations with higher values of $\varepsilon$. Reduction of $\varepsilon$ should be performed between lines 8 and 9 in Algorithm 2: after line 8, $(\mathbf{u},\mathbf{v})$ are “approximations” of the dual variable of the unregularized problem so one can change $\varepsilon$ and start solving for a different $\varepsilon$ with $(\mathbf{u},\mathbf{v})$ as a starting point. We use this heuristic in Section 5 (when mentioned) as follows: starting from $\varepsilon=1$, after every $100$ iteration we perform an absorption step and divide $\varepsilon$ by factor chosen so that the final value $\varepsilon$ is reached after $10$ divisions. Then we run the standard Algorithm 2 until the desired convergence criterion is met.

An efficient numerical implementation of Algorithm 2 is studied further in . In addition to the log-domain stabilization and gradually decreasing $\varepsilon$, it is proposed to approximate $\mathbf{K}$ by a sparse matrix, obtained by adaptive truncation. Thus, one can avoid storing of and multiplication by the dense kernel matrix, while keeping the inflicted truncation error negligible. This is more flexible than for instance the Gaussian convolution trick and can easily be extended to more general cost functions. In addition, this can be directly combined with the log-domain stabilization and therefore allows to solve larger problems with small regularization parameter (and hence, with little entropic blur). We choose however not to use these additional tricks Section 5 in order to display results which are easily reproducible.

5. Generalization: more spaces and pushforward operators

Scaling algorithms similar to Algorithm (S) can be formulated for solving problems of more general form than ($P_{\varepsilon}$). There can be more than $2$ functionals, more than $2$ spaces involved and the projection operators $P^{X}_{\#}$ and $P^{Y}_{\#}$ can be replaced by more general linear operators, such as pushforwards of functions $t$ which are not necessarily projections (i.e. not of the form $t(x,y)=x$). Several examples of such extensions can be found in for the special case of classical optimal transport. Let us sketch this extension in the discrete setting and for the case of $n=1$ (one “coupling”) so as to remain simple and to stick close to implementation concerns. For brevity, we limit ourselves to giving the “scaling” form of the alternate maximization on the dual and an example, without proof.

with the convention $0\log(0/0)=0$, and the dual reads

With those operators, the rightmost term of (4.5) can be computed in a “marginalized” way, using the relation

valid for $k\in\{1,\dots,N\}$. The key feature for obtaining this relation is the fact that $((t^{k})^{-1}(i))_{i=1}^{I_{k}}$ forms a partition of $Z$, and this explains why the scaling algorithm generalizes naturally to linear operators which are “pushforward”. It is now simple, at least formally, to define the generalization of the scaling algorithm dispayed in Algorithm 3, by writing the alternate optimization on the dual problem, and taking again the dual, in the spirit of Proposition 3.3.

As a simple illustration, consider, in the setting of equation (4.1), an extension where is added a function of the total mass

Applications

Throughout this Section, we detail how to use the scaling iterations (S) for solving the problems discussed in Section 2. We first analyze the properties of the functionals on marginals $F_{i}$, then we derive the iterations in a continuous setting, and finally show numerical experiments. We extend the definition of the operator $\operatorname{proxdiv}$ (defined in (4.4) in the discrete setting) to the continuous setting as follows: for $s\in\text{L}^{1}(X)$, $u\in\text{L}^{\infty}(X)$ and $\varepsilon>0$

with the convention $0\times\infty=0$. Some properties of these divergences between functions are studied in Appendix A.2.

All reported runtimes were obtained with an implementation in Julia, on a standard laptop with CPU clock rate $2.5$ GHz.

as in (5.2) above. As shown in Appendix A.2, if $\varphi_{1}$ and $\varphi_{2}$ are a nonnegative entropy function (Definition 2.1) then $F_{1}$ and $F_{2}$ are admissible integral functionals (Definition 3.5). In order to compute the associated $\operatorname{proxdiv}$ operator, let us apply Proposition 3.7 in this precise case.

Let $\varphi$ be a nonnegative entropy function and $(s,p)\in\text{L}^{1}_{+}(X)^{2}$ such that $0\in\operatorname{dom}\varphi$ or $s(x)=0\Rightarrow p(x)=0$ a.e. Let $F(s)=\operatorname{D}_{\varphi}(s|p)$. Then $\operatorname{prox}^{\operatorname{KL}}_{F/\varepsilon}(s)$ is not empty and is the singleton $s^{\star}$ satisfying for a.e. $x\in X$,

It is the pointwise optimality conditions associated to Proposition 3.7. ∎

This formula allows to compute explicitly the $\operatorname{proxdiv}$ operators of the examples introduced in Section 2.1, as listed in Table 1. These entropy functions as well as the associated $\operatorname{proxdiv}$ operators are displayed on Figure 2. Note that, in Table 1 the first line corresponds to standard Sinkhorn iterations and these iterations are recovered in the second and third line by letting $\lambda\to+\infty$ and by setting $\alpha=\beta=1$ in the fourth line. In the context of the log-domain stabilization (Section 4.3), all four $\operatorname{proxdiv}$ operators remain stable in the limit of small $\varepsilon$: either $\operatorname{proxdiv}$ is independent of $u$, only a regularized exponential $e^{-u/(\lambda+\varepsilon)}$ must be evaluated, or extreme values are cut off by thresholding.

In general it is difficult to display optimal transport maps for three-dimensional problems. An interesting application which allows intuitive visualization is color transfer: a classical task in image processing where the goal is to impose the color histogram of one image onto another image. Optimal transport between histograms has proven useful for problems of this sort such as contrast adjustment and color transfer via 1D transportation . Indeed, optimal transport produces a correspondence between histograms which minimizes the total amount of color “distortion” (where the notion of distortion is specified by the cost function) and thus maintains maximal visual consistency.

In our experiments we represent colors in the three-dimensional “CIE-Lab” space (one coordinate for luminance and two for chrominance), resized to fit into a cuboid $X=Y=^{3}$, discretized into $I=J=64\times 32\times 32$ uniform bins and we choose the quadratic cost $c(x,y)=|x-y|^{2}$. The anisotropic discretization of $X$ account for the fact that the eye is more sensitive to variations in luminance than variations in chrominance.

On Figure 8, we display the color transfer between very dissimilar images, computed with parameter $\varepsilon=0.002$. The algorithm was stopped after $2000$ iterations and the running time was approximately $160$ seconds. This intentionally challenging example is insightful as it exhibits a strong effect of the choice of the divergence. There are no quantitative measures for the quality of a transformed image, but the application of unbalanced optimal transport allows to select the “right amount of colors” in the target histogram so as to match the modes of the initial histogram and yields meaningful results.

2. Unbalanced Barycenters

where $\varphi$ is a nonnegative entropy function, $(\alpha_{k})_{k=1}^{n}\in]0,\infty[^{n}$ are weights and $\lambda>0$ is a (redundant) parameter. It is also convenient to slightly modify (for this Section only) the definition of the $\operatorname{KL}$ divergence given in (1.1) by introducing weights as

No theoretical aspect is affected by this small change, but the definition of the kernel has to be adapted, for each component, as

so that one still has $\langle c,r\rangle+\varepsilon\operatorname{KL}(r|1)=\varepsilon\operatorname{KL}(r|K)+\textnormal{const}$. By equation (2.6), the barycenter problem with entropic regularization corresponds to defining

(for all $r\in\text{L}^{1}(X)^{n}$ and $s\in\text{L}^{1}(Y)^{n}$) in ($P_{\varepsilon}$). It is direct to see that the proximal operator of $F_{1}$ for the $\operatorname{KL}$ divergence (according to our specific definition of $\operatorname{KL}$) can be computed componentwise as

If $\varphi^{\prime}_{\infty}>0$ then $F_{2}$ is an admissible integral functional in the sense of Definition 3.5 and for all $s\in\text{L}^{1}(Y)^{n}$, there exists a minimizer $h\in\text{L}^{1}(Y)$. Moreover, if $r^{\star}\in\text{L}^{1}(X\times Y)^{n}$ minimizes ($P_{\varepsilon}$), then the associated minimizer $h^{\star}$ is a pointwise minimizer of (5.5) at the point $s=P^{Y}_{\#}r^{\star}$.

Let $G(s,h)$ be the function on the right side, which is an admissible integral functional (see Proposition A.2 in Appendix A.2). Let us verify the assumptions of the “reduced minimization Theorem” [64, Corollary 3B]. Assumption (i) is satisfied because $\varphi^{\prime}_{\infty}>0$ guarantees the growth condition [64, 2R] and assumption (ii) is guaranteed by the fact that if $u\in\text{L}^{1}(Y)$ and $\operatorname{D}_{\varphi}(u|v)<\infty$ then $v\in\text{L}^{1}(Y)$ (this is proven by adapting slightly the proof of Lemma 3.4, using the—at least linear— growth of $\varphi$). Thus the cited Corollary applies. ∎

According to Proposition 3.7, computing the $\operatorname{prox}$ and $\operatorname{proxdiv}$ operators requires to solve, for each point $y\in Y$, a problem of the form

If $\varphi$ is smooth, first order optimality conditions for (5.6) are simple to obtain. The next Proposition deals with the general case where more care is needed.

In Figure 10, we display Fréchet means-like experiments for a family of $4$ given marginals where $X=Y$ is the segment $$(discretized as above). The (discrete) densities of the marginals$(\mathbf{p}_{k})_{k=1}^{4}$consist each of the sum of three bumps (centered near the points$x=0.1$,$x=0.5$and$x=0.9$). These computations where performed with Algorithm 2 which was stopped after$1500$iterations (running time of$30$seconds approximately) and with$\varepsilon=10^{-5}$. We observe that relaxing the marginal constraints (Figures 10(c)-10(f)) allows to conserve this structure in three bumps in contrast to classical optimal transport (Figure 10(b)).

Figure 11 and 12 display barycenters for the Wasserstein and the $\operatorname{G\!H\!K}$ metric (defined in Section 2.3) between three densities on $^{2}$ discretized into $200\times 200$ samples. Computations where performed using Algorithm 1 and the “separable kernel” method (see Section 4.4) which was stopped after $1500$ iterations (running time of $70$ seconds approximately) with $\varepsilon=9.10^{-4}$. The barycenter coefficients are the following:

The input densities have a similar global structure: each is made of three distant “shapes” of varying mass. The comparison between Figures 11 and 12 lead to a similar remark than for Figure 10: relaxing the strict marginal constraints allows to maintain the global “structure” of the input densities.

3. Gradient Flows and Evolution of Densities

The basic framework of gradient flows has been briefly laid out in Section 2.5. This Section details the application of Algorithm 2 for solving them. As the transition from the measures formulation to the density formulation and further to the algorithm with pointwise optimality conditions was carefully detailed in Sections 5.1 and 5.2, we skip some of these intermediate steps here.

Scaling algorithms for solving Wasserstein gradient flows are not new , but our framework allows to simplify the derivation of the algorithm and the stabilized Algorithm 2 allows to use much smaller regularization parameter $\varepsilon$ yielding sharper and more precise flows. Given a convex, lower semicontinuous function on measures $\mathcal{G}$ with compact sublevel sets, each step requires to find the minimizer of

3.2. WFRWFR\operatorname{WFR} Gradient Flows

For $\operatorname{WFR}$ gradient flows, each step requires to solve

for which the first order optimality conditions read

One of the simplest functional which generates non-trivial gradient flows is

where $\alpha\in]0,\infty[$. Since the distance $\operatorname{WFR}$ measures both the displacement and the rate of growth, one can interpret the gradient flow of $\mathcal{G}$ as describing the evolution of a density of cells (a tumor, say) which have a tendency to multiply—hence increase the total mass—but which density cannot exceed $1$. One can solve the time discretized gradient flow with Algorithm 2 by choosing $F_{2}(s)=\inf_{p}\operatorname{KL}(s|p)-2\alpha\tau\,p+\iota_{\leqslant 1}(p)$. With the optimality conditions above, one obtains

A numerical illustration is given in Figure 13 where we used Algorithm 2 (stopped after $500$ iterations) for solving each step , with the following parameters: $X$ is the segment $$discretized into$3000$uniformly spaced samples, the initial density$p_{0}$is the black line on Figure 13(b),$\tau=0.006$,$\alpha=1$and$\varepsilon=10^{-8}$. The running time was$315$seconds. Remark how, by using a very small value for$\varepsilon$, the “smoothing” effect of the entropy disappears: the contours of the free-boundary which evolve with time remain sharp.

3.3. WFRWFR\operatorname{WFR} Gradient Flows with Multiple Species

The generic form (2.9) also includes gradient flows with multiple species with a mutual interaction (with $n>1$, similar to the barycenter problem). Such systems have been theoretically studied in . Here we consider this class of problems in order to illustrate the versatility of the algorithm and it is not our purpose to make a link with the theory of PDEs. Let us consider a simple example which is a direct extension of Example 5.6. Consider the following functional:

In this model, one has two species which have a tendency to grow in mass (with the same incentive $\alpha>0$, for simplicity of the algorithm), and their sum cannot exceed the reference measure. The corresponding $F_{2}$ is given by

where $\beta(x,y)\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}\max\left\{(x+y)^{\frac{1}{1+\varepsilon}},\,(1-2\tau\alpha)^{\frac{1}{\varepsilon}}\right\}$. Morevoer, given an optimal pair of couplings $r^{a},r^{b}$, the next densities are given by $(p^{a}_{k+1},p^{b}_{k+1})=(P^{2}_{\#}r^{a},P^{2}_{\#}r^{b})/\beta(P^{2}_{\#}r^{a},P^{2}_{\#}r^{b})$. Remark than in this model, as in Example 5.6, if the domain $X$ is compact and the initial densities are not null, a steady state is reached in finite time, where the sum of the two densities is constant and equal to $1$.

Some initial densities and the associated final steady state are shown on Figure 14. For this illustration, we started with input densities $p^{a}$ and $p^{b}$ on the segment $$discretized into$3000$uniform samples, as displayed on Figure 14(a) (where the red density$p^{b}$is layered over$p^{a}$). We computed the evolution with Algorithm 2 for solving each step of the discretized gradient flow, with the parameters$\tau=0.004$,$\alpha=1$and$\varepsilon=10^{-7}$.

Note that although the incentive of growing mass $\alpha$ is the same for the two species, the resulting interaction is non trivial: for instance the small amount of blue mass is pushed to the right by the action of the expanding red mass. This behavior is explained by the fact that for the $\operatorname{WFR}$ metric, it requires less effort (i.e. the distance is smaller) to add a given amount of mass to a high density than to a small one.

Acknowledgements

The work of Bernhard Schmitzer has been supported by the French National Research Agency (ANR) as part of the ‘Investissements d’avenir’, program-reference ANR-10-LABX-0098 via the Fondation Sciences Mathématiques de Paris. The work of Bernhard Schmitzer and Gabriel Peyré has been supported by the European Research Council (ERC project SIGMA-Vision).

Appendix A Appendix

The subdifferential operator is defined at a point $x\in E$ as

and is empty if $f(x)=\infty$. Those definitions admit their natural counterparts for functions defined on $E^{*}$.

Let $(E,E^{*})$ and $(F,F^{*})$ be two couples of topologically paired spaces. Let $A:E\to F$ be a continuous linear operator and $A^{*}:F^{*}\to E^{*}$ its adjoint. Let $f$ and $g$ be lower semicontinuous and proper convex functions defined on $E$ and $F$ respectively. If there exists $x\in\operatorname{dom}f$ such that $g$ is continuous at $Ax$, then

and the $\min$ is attained. Moreover, if there exists a maximizer $x\in E$ then there exists $y^{*}\in F^{*}$ satisfying $Ax\in\partial g^{*}(y^{*})$ and $A^{*}y^{*}\in\partial f(-x)$.

A.2. Properties of Divergence Functionals

Here we collect a few results on divergences functionals when they are defined on functions as in (5.2) (as opposed to Section 2.1 where they are defined between measures).

where $\varphi^{*}$ is the convex conjugate of $\varphi$.

Moreover, the subdifferential $\partial\operatorname{D}_{\varphi}(\cdot|v)$ at a point $u\in\text{L}^{1}(X)$ is the set of functions $a\in\text{L}^{\infty}(X)$ such that $\varphi^{\prime}_{\infty}-a$ is nonnegative and such that, for a.e. $x$ where $v(x)>0$, $a(x)\in\partial\varphi(u(x)/v(x))$ .

Similarly, the subdifferential $\partial\operatorname{D}^{*}_{\varphi}(\cdot|v)$ at a point $a\in\text{L}^{\infty}(X)$ bounded above by $\varphi^{\prime}_{\infty}$ is the set of nonnegative functions $u\in\text{L}^{1}(X)$ such that, for a.e. $x$, $u(x)\in\partial\varphi^{*}(a(x))v(x)$ if $v(x)>0$ and $u(x)=0$ if $v(x)=0$ and $a(x)<\varphi^{\prime}_{\infty}$.

A.3. Proof of the Iterates for the Barycenter Problems

This case is simple because solving (5.6) boils down to solving the one dimensional problem $\min_{h}\sum\alpha_{k}\overline{\operatorname{KL}}(h|s_{k})$, which is direct with first order optimality conditions.

If $s_{k}=0$ for some $k\in\{1,\dots,n\}$ then $h=0$ (this is the only feasible point). Otherwise, $h>0$ and the condition $\sum\alpha_{k}b_{k}=0$ gives the implicit equation.

References