Stochastic Optimization for Large-scale Optimal Transport

Introduction

Many problems in computational sciences require to compare probability measures or histograms. As a set of representative examples, let us quote: bag-of-visual-words comparison in computer vision , color and shape processing in computer graphics , bag-of-words for natural language processing and multi-label classification . In all of these problems, a geometry between the features (words, visual words, labels) is usually known, and can be leveraged to compare probability distributions in a geometrically faithful way. This underlying geometry might be for instance the planar Euclidean domain for 2-D shapes, a perceptual 3D color metric space for image processing or a high-dimensional semantic embedding for words. Optimal transport (OT) is the canonical way to automatically lift this geometry to define a metric for probability distributions. That metric is known as the Wasserstein or earth mover’s distance. As an illustrative example, OT can use a metric between words to build a metric between documents that are represented as frequency histograms of words (see for details). All the above-cited lines of work advocate, among others, that OT is the natural choice to solve these problems, and that it leads to performance improvement when compared to geometrically-oblivious distances such as the Euclidean or $\chi^{2}$ distances or the Kullback-Leibler divergence. However, these advantages come at the price of an enormous computational overhead. This is especially true because current OT solvers require to sample beforehand these distributions on a pre-defined set of points, or on a grid. This is both inefficient (in term of storage and speed) and counter-intuitive. Indeed, most high-dimensional computational scenarios naturally represent distributions as objects from which one can sample, not as density functions to be discretized. Our goal is to alleviate these shortcomings. We propose a class of provably convergent stochastic optimization schemes that can handle both discrete and continuous distributions through sampling.

Previous works. The prevalent way to compute OT distances is by solving the so-called Kantorovitch problem (see Section 2 for a short primer on the basics of OT formulations), which boils down to a large-scale linear program when dealing with discrete distributions (i.e., finite weighted sums of Dirac masses). This linear program can be solved using network flow solvers, which can be further refined to assignment problems when comparing measures of the same size with uniform weights . Recently, regularized approaches that solve the OT with an entropic penalization have been shown to be extremely efficient to approximate OT solutions at a very low computational cost. These regularized approaches have supported recent applications of OT to computer graphics and machine learning . These methods apply the celebrated Sinkhorn algorithm’s , and can be extended to solve more exotic transportation-related problems such as the computation of barycenters . Their chief computational advantage over competing solvers is that each iteration boils down to matrix-vector multiplications, which can be easily parallelized, streams extremely well on GPU, and enjoys linear-time implementation on regular grids or triangulated domains .

These methods are however purely discrete and cannot cope with continuous densities. The only known class of methods that can overcome this limitation are so-called semi-discrete solvers , that can be implemented efficiently using computational geometry primitives . They can compute distance between a discrete distributions and a continuous density. Nonetheless, they are restricted to the Euclidean squared cost, and can only be implemented in low dimensions (2-D and 3-D). Solving these semi-discrete problems efficiently could have a significant impact for applications to density fitting with an OT loss for machine learning applications, see . Lastly, let us point out that there is currently no method that can compute OT distances between two continuous densities, which is thus an open problem we tackle in this article.

Contributions. This paper introduces stochastic optimization methods to compute large-scale optimal transport in all three possible settings: discrete OT, to compare a discrete vs. another discrete measure; semi-discrete OT, to compare a discrete vs. a continuous measure; and continous OT, to compare a continuous vs. another continuous measure. These methods can be used to solve both classical OT problems and their entropic-regularized versions (which enjoy faster convergence properties). We show that the discrete OT problem can be tackled using incremental algorithms, and we consider in particular the stochastic averaged gradient (SAG) method . Each iteration of that algorithm requires $N$ operations ($N$ being the size of the supports of the input distributions), which makes it scale better in large-scale problems than the state-of-the-art Sinkhorn algorithm, while still enjoying a convergence rate of $O(1/k)$, $k$ being the number of iterations. We show that the semi-discrete OT problem can be solved using averaged stochastic gradient descent (SGD), whose convergence rate is $O(1/\sqrt{k})$. For large-scale problems, this approach is numerically advantageous over the brute force approach consisting in sampling first the continuous density to solve next a discrete OT problem. Lastly, for continuous optimal transport, we propose a novel method which makes use of an expansion of the dual variables in a reproducing kernel Hilbert space (RKHS). This allows us for the first time to compute with a converging algorithm OT distances between two arbitrary densities, under the assumption that the two potentials belong to such an RKHS.

Notations. In the following we consider two metric spaces $\mathcal{X}$ and $\mathcal{Y}$. We denote by $\mathcal{M}_{+}^{1}(\mathcal{X})$ the set of positive Radon probability measures on $\mathcal{X}$, and $\mathcal{C}(\mathcal{X})$ the space of continuous functions on $\mathcal{X}$. Let $\mu\in\mathcal{M}_{+}^{1}(\mathcal{X})$, $\nu\in\mathcal{M}_{+}^{1}(\mathcal{Y})$, we define

the set of joint probability measures on $\mathcal{X}\times\mathcal{Y}$ with marginals $\mu$ and $\nu$. The Kullback-Leibler divergence between joint probabilities is defined as

Optimal Transport: Primal, Dual and Semi-dual Formulations

We consider the optimal transport problem between two measures $\mu\in\mathcal{M}_{+}^{1}(\mathcal{X})$ and $\nu\in\mathcal{M}_{+}^{1}(\mathcal{Y})$, defined on metric spaces $\mathcal{X}$ and $\mathcal{Y}$. No particular assumption is made on the form of $\mu$ and $\nu$, we simply assume that they both can be sampled from to be able to apply our algorithms.

Primal, Dual and Semi-dual Formulations. The Kantorovich formulation of OT and its entropic regularization can be conveniently written in a single convex optimization problem as follows

Here $c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})$ and $c(x,y)$ should be interpreted as the “ground cost” to move a unit of mass from $x$ to $y$. This $c$ is typically application-dependent, and reflects some prior knowledge on the data to process. We refer to the introduction for a list of previous work where various examples (in imaging, vision, graphics or machine learning) of such costs are given.

When $\mathcal{X}=\mathcal{Y}$, $\varepsilon=0$ and $c=d^{p}$ for $p\geq 1$, where $d$ is a distance on $\mathcal{X}$, then $W_{0}(\mu,\nu)^{\frac{1}{p}}$ is known as the $p$-Wasserstein distance on $\mathcal{M}_{+}^{1}(\mathcal{X})$. Note that this definition can be used for any type of measure, both discrete and continuous. When $\varepsilon>0$, problem ($\mathcal{P}_{\varepsilon}$) is strongly convex, so that the optimal $\pi$ is unique, and algebraic properties of the $\operatorname{KL}$ regularization result in computations that can be tackled using the Sinkhorn algorithm .

For any $c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})$, we define the following constraint set

and define its indicator function as well as its “smoothed” approximation

For any $v\in\mathcal{C}(\mathcal{Y})$, we define its $c$-transform and its “smoothed” approximation

The proposition below describes two dual problems. It is central to our analysis and paves the way for the application of stochastic optimization methods.

Problem ($\mathcal{D}_{\varepsilon}$) is the convex dual of ($\mathcal{P}_{\varepsilon}$), and is derived using Fenchel-Rockafellar’s theorem. The relation between $u$ and $v$ is obtained by writing the first order optimality condition for $v$ in ($\mathcal{D}_{\varepsilon}$). Plugging this expression back in ($\mathcal{D}_{\varepsilon}$) yields ($\mathcal{S}_{\varepsilon}$). ∎

A key advantage of ($\mathcal{S}_{\varepsilon}$) over ($\mathcal{D}_{\varepsilon}$) is that, when $\nu$ is a discrete density (but not necessarily $\mu$), then ($\mathcal{S}_{\varepsilon}$) is a finite-dimensional concave maximization problem, which can thus be solved using stochastic programming techniques, as highlighted in Section 4. By contrast, when both $\mu$ and $\nu$ are continuous densities, these dual problems are intrinsically infinite dimensional, and we propose in Section 5 more advanced techniques based on RKHSs.

Stochastic Optimization Formulations. The fundamental property needed to apply stochastic programming is that both dual problems ($\mathcal{D}_{\varepsilon}$) and ($\mathcal{S}_{\varepsilon}$) must be rephrased as minimizing expectations:

where the random variables $X$ and $Y$ are independent and distributed according to $\mu$ and $\nu$ respectively, and where, for $(x,y)\in\mathcal{X}\times\mathcal{Y}$ and $(u,v)\in\mathcal{C}(\mathcal{X})\times\mathcal{C}(\mathcal{Y})$,

This reformulation is at the heart of the methods detailed in the remainder of this article. Note that the dual problem ($\mathcal{D}_{\varepsilon}$) cannot be cast as an unconstrained expectation maximization problem when $\varepsilon=0$, because of the constraint on the potentials which arises in that case.

Discrete Optimal Transport

We assume in this section that both $\mu$ and $\nu$ are discrete measures, i.e. finite sums of Diracs, of the form $\mu=\sum_{i=1}^{I}\bm{\mu}_{i}\delta_{x_{i}}\quad\text{and}\quad\nu=\sum_{j=1}^{J}\bm{\nu}_{j}\delta_{y_{j}},$ where $(x_{i})_{i}\subset\mathcal{X}$ and $(y_{j})_{j}\subset\mathcal{Y}$, and the histogram vector weights are $\bm{\mu}\in\Sigma_{I}$ and $\bm{\nu}\in\Sigma_{J}$. These discrete measures may come from the evaluation of continuous densities on a grid, counting features in a structured object, or be empirical measures based on samples. This setting is relevant for several applications, including all known applications of the earth mover’s distance. We show in this section that our stochastic formulation can prove extremely efficient to compare measures with a large number of points.

The state-of-the-art method to solve the discrete regularized OT (i.e. when $\varepsilon>0$) is Sinkhorn’s algorithm [6, Alg.1], which has linear convergence rate . It corresponds to a block coordinate maximization, successively optimizing ($\mathcal{\bar{D}}_{\varepsilon}$) with respect to either $\mathbf{u}$ or $\mathbf{v}$. Each iteration of this algorithm is however costly, because it requires a matrix-vector multiplication. Indeed, this corresponds to a “batch” method where all the samples $(x_{i})_{i}$ and $(y_{j})_{j}$ are used at each iteration, which has thus complexity $O(N^{2})$ where $N=\max(I,J)$. We now detail how to alleviate this issue using online stochastic optimization methods.

Incremental Discrete Optimization when $\varepsilon>0$. Stochastic gradient descent (SGD), in which an index $k$ is drawn from distribution $\bm{\mu}$ at each iteration can be used to minimize the finite sum that appears in in $\mathcal{\bar{S}}_{\varepsilon}$. The gradient of that term $\bar{h}_{\varepsilon}(x_{k},\cdot)$ is

When $\varepsilon>0$, the finite sum appearing in ($\mathcal{\bar{S}}_{\varepsilon}$) suggests to use incremental gradient methods—rather than purely stochastic ones—which are known to converge faster than SGD. We propose to use the stochastic averaged gradient (SAG) . As SGD, SAG operates at each iteration by sampling a point $x_{k}$ from $\mu$, to compute the gradient corresponding to that sample for the current estimate $\mathbf{v}$. Unlike SGD, SAG keeps in memory a copy of that gradient, until that particular point is sampled again, at which point the copy is updated. Unlike SGD, SAG applies a fixed length update, in the direction of the average of all gradients stored so far, which provides a better proxy of the gradient corresponding to the entire sum. This improves the convergence rate to $|\bar{H}_{\varepsilon}(\mathbf{v}_{\varepsilon}^{\star})-\bar{H}_{\varepsilon}(\mathbf{v}_{k})|=O(1/k)$, where $\mathbf{v}_{\varepsilon}^{\star}$ is a minimizer of $\bar{H}_{\varepsilon}$, at the expense of storing the gradient for each of the $I$ points. This expense can be mitigated by considering mini-batches instead of individual points. Note finally that the SAG algorithm is adaptive to strong-convexity and will be linearly convergent around the optimum. The pseudo-code for SAG is provided in Algorithm 1, and we defer more details on SGD for Section 4, in which it will be shown to play a crucial role. Note that the Lipschitz constant of all these terms is upperbounded by $L=\max_{i}\bm{\mu}_{i}/\varepsilon$.

Semi-Discrete Optimal Transport

Figure 2 (a) shows the evolution of $\left\lVert\mathbf{v}_{k}-\mathbf{v}_{0}^{\star}\right\rVert_{2}/\left\lVert\mathbf{v}_{0}^{\star}\right\rVert_{2}$ as a function of $k$. It highlights the influence of the regularization parameters $\varepsilon$ on the iterates of SGD. While the regularized iterates converge faster, they do not converge to the correct unregularized solution. This figure also illustrates the convergence theorem of solution of $(\mathcal{S}_{\varepsilon})$ toward those $(\mathcal{S}_{0})$ when $\varepsilon\rightarrow 0$, which can be found in Appendix A.

Figure 2 (b) shows the evolution of $\left\lVert\mathbf{v}_{k}-\mathbf{v}_{\varepsilon}^{\star}\right\rVert_{2}/\left\lVert\mathbf{v}_{\varepsilon}^{\star}\right\rVert_{2}$ averaged over 40 runs as a function of $k$, for a fixed regularization parameter value $\varepsilon=10^{-2}$. It compares SGD to SAG using different numbers $N$ of samples for the empirical measures $\hat{\mu}_{N}$. While SGD converges to the true solution of the semi-discrete problem, the solution computed by SAG is biased because of the approximation error which comes from the discretization of $\mu$. This error decreases when the sample size $N$ is increased, as the approximation of $\mu$ by $\hat{\mu}_{N}$ becomes more accurate.

Continuous optimal transport using RKHS

In the case where neither $\mu$ nor $\nu$ are discrete, problem ($\mathcal{S}_{\varepsilon}$) is infinite-dimensional, so it cannot be solved directly using stochastic SGD. We propose in this section to solve the initial dual problem ($\mathcal{D}_{\varepsilon}$), using expansions of the dual variables in two reproducing kernel Hilbert spaces (RKHS). Recall that contrarily to the methods from previous sections, we can only solve the regularized problem here (i.e. $\varepsilon>0$), since ($\mathcal{D}_{\varepsilon}$) cannot be cast as an expectation maximization problem when $\varepsilon=0$.

The dual problem ($\mathcal{D}_{\varepsilon}$) is conveniently re-written in (3) as the maximization of the expectation of $f^{\varepsilon}(X,Y,u,v)$ with respect to the random variables $(X,Y)\sim\mu\otimes\nu$. The SGD algorithm applied to this problem reads, starting with $u_{0}=0$ and $v_{0}=0$,

where $(x_{k},y_{k})$ are i.i.d. samples from $\mu\otimes\nu$. The following proposition shows that these $(u_{k},v_{k})$ iterates can be expressed as finite sums of kernel functions, and that the coefficients of these expansions enjoy a particularly simple recursion formula.

The iterates $(u_{k},v_{k})$ defined in (5) satisfy

where $(x_{i},y_{i})_{i=1\dots k}$ are i.i.d samples from $\mu\otimes\nu$ and $\Pi_{B_{r}}$ is the projection on the centered ball of radius $r$. If the solutions of ($\mathcal{D}_{\varepsilon}$) are in the $\mathcal{H}\times\mathcal{G}$ and if $r$ is large enough, the iterates ($u_{k}$,$v_{k}$) converge to a solution of ($\mathcal{D}_{\varepsilon}$).

Rewriting $u(x)$ and $v(y)$ as scalar products in $f^{\varepsilon}(X,Y,u,v)$ yields

Numerical Illustrations. We consider optimal transport in 1D between a Gaussian $\mu$ and a Gaussian mixture $\nu$ whose densities are represented in Figure 3 (a). Since there is no existing benchmark for continuous transport, we use the solution of the semi-discrete problem $W_{\varepsilon}(\mu,\hat{\nu}_{N})$ with $N=10^{3}$ computed with SGD as a proxy for the solution and we denote it by $\hat{u}^{\star}$. We focus on the convergence of the potential $u$, as it is continuous in both problems contrarily to $v$. Figure 3 (b) represents the plot of ${\left\lVert\mathbf{u}_{k}-\hat{\mathbf{u}}^{\star}\right\rVert_{2}}/{\left\lVert\hat{\mathbf{u}}^{\star}\right\rVert_{2}}$ where $\mathbf{u}$ is the evaluation of $u$ on a sample $(x_{i})_{i=1\dots N^{\prime}}$ drawn from $\mu$. This gives more emphasis to the norm on points where $\mu$ has more mass. The convergence is rather slow but still noticeable. The iterates $u_{k}$ are plotted on a grid for different values of $k$ in Figure 3 (c), to emphasize the convergence to the proxy $\hat{u}^{\star}$. We can see that the iterates computed with the RKHS converge faster where $\mu$ has more mass, which is actually where the value of $u$ has the greatest impact in $F_{\varepsilon}$ ($u$ being integrated against $\mu$).

Conclusion

We have shown in this work that the computations behind (regularized) optimal transport can be considerably alleviated, or simply enabled, using a stochastic optimization approach. In the discrete case, we have shown that incremental gradient methods can surpass the Sinkhorn algorithm in terms of efficiency, taken for granted that the (constant) stepsize has been correctly selected, which should be possible in practical applications. We have also proposed the first known methods that can address the challenging semi-discrete and continuous cases. All of these three settings can open new perspectives for the application of OT to high-dimensional problems.

Acknowledgement

The work of G. Peyré has been supported by the European Research Council (ERC project SIGMA-Vision). The work of A. Genevay has been supported by Région Ile-de-France. M. Cuturi gratefully acknowledges the support of JSPS young research A grant 26700002.

References

The convergence of the solution of ($\mathcal{P}_{\varepsilon}$) toward a solution of $(\mathcal{P}_{0})$ as $\varepsilon\rightarrow 0$ is proved in . The convergence of solutions of ($\mathcal{D}_{\varepsilon}$) toward solutions of $(\mathcal{D}_{0})$ as $\varepsilon\rightarrow 0$ is proved for the special case of discrete measures in . To the best of our knowledge, the behavior of ($\mathcal{S}_{\varepsilon}$) has not been studied in the literature, and we propose a convergence result in the case where $\nu$ is discrete, which is the setting in which this formulation is most advantageous.

We assume that $\forall y\in\mathcal{Y}$, $c(\cdot,y)\in L^{1}(\mu)$, that $\nu=\sum_{j=1}^{J}\bm{\nu}_{j}\delta_{y_{j}}$, and we fix $x_{0}\in\mathcal{X}$,. For all $\varepsilon>0$, let $v^{\varepsilon}$ be the unique solution of ($\mathcal{S}_{\varepsilon}$) such that $v^{\varepsilon}(x_{0})=0$. Then $(v^{\varepsilon})_{\varepsilon}$ is bounded and all its converging sub-sequences for $\varepsilon\rightarrow 0$ are solutions of $(\mathcal{S}_{0})$.

If $\forall y$, $x\mapsto c(x,y)\in L^{1}(\mu)$ then $H_{\varepsilon}$ converges pointwise to $H_{0}$.

Let $\alpha_{j}(x)\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}v_{j}-c(x,y_{j})$ and $j^{\star}\stackrel{{\scriptstyle\mbox{\tiny def.}}}{{=}}\operatorname*{arg\,max}_{j}\alpha_{j}(x)$. On the one hand, since $\forall j$, $\alpha_{j}(x)\leq\alpha_{j^{\star}}(x)$ we get

On the other hand, since $\log$ is increasing and all terms in the sum are non negative we have

Hence $\varepsilon\log(\sum_{j=1}^{J}e^{\frac{\alpha_{j}(x)}{\varepsilon}}\nu_{j})\overset{\varepsilon\to 0}{\longrightarrow}\alpha_{j^{\star}}(x)$ and $\varepsilon\log(\sum_{j=1}^{J}e^{\frac{\alpha_{j}(x)}{\varepsilon}}\nu_{j})\leq\alpha_{j^{\star}}(x)$. Since we assumed $x\mapsto c(x,y_{j})\in L^{1}(\mu)$, then $\alpha_{j^{\star}}\in L^{1}(\mu)$ and by dominated convergence we get that $H_{\varepsilon}(v)\overset{\varepsilon\to 0}{\longrightarrow}H_{0}(v)$. ∎

Besides, the regularized potentials are unique up to an additive constant. Hence we can set without loss of generality $v_{\varepsilon}(y_{0})=0$. So from the previous inequality yields :

Let $v^{\star}\in\operatorname*{arg\,max}_{v}H_{0}$. To prove that $\bar{v}$ is optimal, it suffices to prove that $H_{0}(v^{\star})\leq H_{0}(\bar{v})$. By optimality of $v_{\varepsilon}$,

The term on the left-hand side of the inequality converges to $H_{0}(v^{\star})$ since $H_{\varepsilon}$ converges pointwise to $H_{0}$. We still need to prove that the right-hand term converges to $H_{0}(\bar{v})$.

where $\pi_{i}(v)=\frac{\int_{\mathcal{X}}e^{\frac{v_{i}-c(x,y_{i})}{\varepsilon}}\nu_{i}d\mu(x)}{\int_{\mathcal{X}}\sum_{j=1}^{J}e^{\frac{v_{j}-c(x,y_{j})}{\varepsilon}}\nu_{j}d\mu(x)}$

It is the difference of two elements in the simplex thus it is bounded by a constant $C$ independently of $\varepsilon$.

By pointwise convergence of $H_{\varepsilon}$ we know that $H_{\varepsilon}(\bar{v})\rightarrow H_{0}(\bar{v})$, and since $\bar{v}$ is a limit point of $v_{\varepsilon}$ we can conclude that the left and right hand term of the inequality converge to $H_{0}(\bar{v})$. Thus we get $H_{\varepsilon}(v_{\varepsilon})\rightarrow H_{0}(\bar{v})$. ∎