Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration

Introduction

Computing distances between probability measures on metric spaces, or more generally between point clouds, plays an increasingly preponderant role in machine learning [SL11, MJ15, LG15, JSCG16, ACB17], statistics [FCCR16, PZ16, SR04, BGKL17] and computer vision [RTG00, BvdPPH11, SdGP+15]. A prominent example of such distances is the earth mover’s distance introduced in [WPR85] (see also [RTG00]), which is a special case of Wasserstein distance, or optimal transport (OT) distance [Vil09].

While OT distances exhibit a unique ability to capture geometric features of the objects at hand, they suffer from a heavy computational cost that had been prohibitive in large scale applications until the recent introduction to the machine learning community of Sinkhorn Distances by Cuturi [Cut13]. Combined with other numerical tricks, these recent advances have enabled the treatment of large point clouds in computer graphics such as triangle meshes [SdGP+15] and high-resolution neuroimaging data [GPC15]. Sinkhorn Distances rely on the idea of entropic penalization, which has been implemented in similar problems at least since Schrödinger [Sch31, Leo14]. This powerful idea has been successfully applied to a variety of contexts not only as a statistical tool for model selection [JRT08, RT11, RT12] and online learning [CBL06], but also as an optimization gadget in first-order optimization methods such as mirror descent and proximal methods [Bub15].

Related work. Computing an OT distance amounts to solving the following linear system:

where $\mathbf{1}$ is the all-ones vector in ${\rm I}\kern-1.79993pt{\rm R}^{n}$, $C\in{\rm I}\kern-1.79993pt{\rm R}_{+}^{n\times n}$ is a given cost matrix, and $r\in{\rm I}\kern-1.79993pt{\rm R}^{n},c\in{\rm I}\kern-1.79993pt{\rm R}^{n}$ are given vectors with positive entries that sum to one. Typically $C$ is a matrix containing pairwise distances (and is thus dense), but in this paper we allow $C$ to be an arbitrary non-negative dense matrix with bounded entries since our results are more general. For brevity, this paper focuses on square matrices $C$ and $P$, since extensions to the rectangular case are straightforward.

This paper is at the intersection of two lines of research: a theoretical one that aims at finding (near) linear time approximation algorithms for simple problems that are already known to run in polynomial time and a practical one that pursues fast algorithms for solving optimal transport approximately for large datasets.

Practical algorithms for computing OT distances include Orlin’s algorithm for the Uncapacitated Minimum Cost Flow problem via a standard reduction. Like interior point methods, it has a provable complexity of $O(n^{3}\log n)$. This dependence on the dimension is also observed in practice, thereby preventing large-scale applications. To overcome the limitations of such general solvers, various ideas ranging from graph sparsification [PW09] to metric embedding [IT03, GD04, SJ08] have been proposed over the years to deal with particular cases of OT distance.

Our work complements both lines of work, theoretical and practical, by providing the first near-linear time guarantee to approximate (1) for general non-negative cost matrices. Moreover we show that this performance is achieved by algorithms that are also very efficient in practice. Central to our contribution are recent developments of scalable methods for general OT that leverage the idea of entropic regularization [Cut13, BCC+15, GCPB16]. However, the apparent practical efficacy of these approaches came without theoretical guarantees. In particular, showing that this regularization yields an algorithm to compute or approximate general OT distances in time nearly linear in the input size $n^{2}$ was an open question before this work.

Our contribution. The contribution of this paper is twofold. First we demonstrate that, with an appropriate choice of parameters, the algorithm for Sinkhorn Distances introduced in [Cut13] is in fact a near-linear time approximation algorithm for computing OT distances between discrete measures. This is the first proof that such near-linear time results are achievable for optimal transport. We also provide previously unavailable guidance for parameter tuning in this algorithm. Core to our work is a new and arguably more natural analysis of the Sinkhorn iteration algorithm, which we show converges in a number of iterations independent of the dimension $n$ of the matrix to balance. In particular, this analysis directly suggests a greedy variant of Sinkhorn iteration that also provably runs in near-linear time and significantly outperforms the classical algorithm in practice. Finally, while most approximation algorithms output an approximation of the optimum value of the linear program (1), we also describe a simple, parallelizable rounding algorithm that provably outputs a feasible solution to (1). Specifically, for any $\varepsilon>0$ and bounded, non-negative cost matrix $C$, we describe an algorithm that runs in time $\widetilde{O}(n^{2}/\varepsilon^{3})$ and outputs $\hat{P}\in\mathcal{U}_{r,c}$ such that

We emphasize that our analysis does not require the cost matrix $C$ to come from an underlying metric; we only require $C$ to be non-negative. This implies that our results also give, for example, near-linear time approximation algorithms for Wasserstein $p$-distances between discrete measures.

Notation. We denote non-negative real numbers by ${\rm I}\kern-1.79993pt{\rm R}_{+}$, the set of integers $\{1,\dots,n\}$ by $[n]$, and the $n$-dimensional simplex by $\Delta_{n}:=\{x\in{\rm I}\kern-1.79993pt{\rm R}_{+}^{n}\;:\;\sum_{i=1}^{n}x_{i}=1\}$. For two probability distributions $p,q\in\Delta_{n}$ such that $p$ is absolutely continuous w.r.t. $q$, we define the entropy $H(p)$ of $p$ and the Kullback-Leibler divergence $\mathcal{K}(p\|q)$ between $p$ and $q$ respectively by

Similarly, for a matrix $P\in{\rm I}\kern-1.79993pt{\rm R}_{+}^{n\times n}$, we define the entropy $H(P)$ entrywise as $\sum_{ij}P_{ij}\log\frac{1}{P_{ij}}$. We use $\mathbf{1}$ and $\mathbf{0}$ to denote the all-ones and all-zeroes vectors in ${\rm I}\kern-1.79993pt{\rm R}^{n}$. For a matrix $A=(A_{ij})$, we denote by $\exp(A)$ the matrix with entries $(e^{A_{ij}})$. For $A\in{\rm I}\kern-1.79993pt{\rm R}^{n\times n}$, we denote its row and columns sums by $r(A):=A\mathbf{1}\in{\rm I}\kern-1.79993pt{\rm R}^{n}$ and $c(A):=A^{\top}\mathbf{1}\in{\rm I}\kern-1.79993pt{\rm R}^{n}$, respectively. The coordinates $r_{i}(A)$ and $c_{j}(A)$ denote the $i$th row sum and $j$th column sum of $A$, respectively. We write $\|A\|_{\infty}=\max_{ij}|A_{ij}|$ and $\|A\|_{1}=\sum_{ij}|A_{ij}|$. For two matrices of the same dimension, we denote the Frobenius inner product of $A$ and $B$ by $\langle A,B\rangle=\sum_{ij}A_{ij}B_{ij}$. For a vector $x\in{\rm I}\kern-1.79993pt{\rm R}^{n}$, we write $\mathbf{D}(x)\in{\rm I}\kern-1.79993pt{\rm R}^{n\times n}$ to denote the diagonal matrix with entries $(\mathbf{D}(x))_{ii}=x_{i}$. For any two nonnegative sequences $(u_{n})_{n},(v_{n})_{n}$, we write $u_{n}=\widetilde{O}(v_{n})$ if there exist positive constants $C,c$ such that $u_{n}\leq Cv_{n}(\log n)^{c}$. For any two real numbers, we write $a\wedge b=\min(a,b)$.

Optimal Transport in near-linear time

In this section, we describe the main algorithm studied in this paper. Pseudocode appears in Algorithm 1.

The core of our algorithm is the computation of an approximate Sinkhorn projection of the matrix $A=\exp(-\eta C)$ (Step 1), details for which will be given in Section 3. Since our approximate Sinkhorn projection is not guaranteed to lie in the feasible set, we round our approximation to ensure that it lies in $\mathcal{U}_{r,c}$ (Step 2). Pseudocode for a simple, parallelizable rounding procedure is given in Algorithm 2.

Algorithm 1 hinges on two subroutines: Proj and round. We give two algorithms for Proj: Sinkhorn and Greenkhorn. We devote Section 3 to their analysis, which is of independent interest. On the other hand, round is fairly simple. Its analysis is postponed to Section 4.

Our main theorem about Algorithm 1 is the following accuracy and runtime guarantee. The proof is postponed to Section 4, since it relies on the analysis of Proj and round.

Algorithm 1 returns a point $\hat{P}\in\mathcal{U}_{r,c}$ satisfying

in time $O(n^{2}+S)$, where $S$ is the running time of the subroutine $\textsc{Proj}(A,\mathcal{U}_{r,c},\varepsilon^{\prime})$. In particular, if $\|C\|_{\infty}\leq L$, then $S$ can be $O(n^{2}L^{3}(\log n)\varepsilon^{-3})$, so that Algorithm 1 runs in $O(n^{2}L^{3}(\log n)\varepsilon^{-3})$ time.

The time complexity in the above theorem reflects only elementary arithmetic operations. In the interest of clarity, we ignore questions of bit complexity that may arise from taking exponentials. The effect of this simplification is marginal since it can be easily shown [KLRS08] that the maximum bit complexity throughout the iterations of our algorithm is $O(L(\log n)/\varepsilon)$. As a result, factoring in bit complexity leads to a runtime of $O(n^{2}L^{4}(\log n)^{2}\varepsilon^{-4})$, which is still truly near-linear.

Linear-time approximate Sinkhorn projection

The core of our OT algorithm is the entropic penalty proposed by Cuturi [Cut13]:

The solution to (2) can be characterized explicitly by analyzing its first-order conditions for optimality.

[Cut13] For any cost matrix $C$ and $r,c\in\Delta_{n}$, the minimization program (2) has a unique minimum at $P_{\eta}\in\mathcal{U}_{r,c}$ of the form $P_{\eta}=XAY$, where $A=\exp(-\eta C)$ and $X,Y\in{\rm I}\kern-1.79993pt{\rm R}_{+}^{n\times n}$ are both diagonal matrices. The matrices $(X,Y)$ are unique up to a constant factor.

We call the matrix $P_{\eta}$ appearing in Lemma 1 the Sinkhorn projection of $A$, denoted $\Pi_{\mathcal{S}}(A,\mathcal{U}_{r,c})$, after Sinkhorn, who proved uniqueness in [Sin67]. Computing $\Pi_{\mathcal{S}}(A,\mathcal{U}_{r,c})$ exactly is impractical, so we implement instead an approximate version $\textsc{Proj}(A,\mathcal{U}_{r,c},\varepsilon^{\prime})$, which outputs a matrix $B=XAY$ that may not lie in $\mathcal{U}_{r,c}$ but satisfies the condition $\|r(B)-r\|_{1}+\|c(B)-c\|_{1}\leq\varepsilon^{\prime}$. We stress that this condition is very natural from a statistical standpoint, since it requires that $r(B)$ and $c(B)$ are close to the target marginals $r$ and $c$ in total variation distance.

Given a matrix $A$, Sinkhorn proposed a simple iterative algorithm to approximate the Sinkhorn projection $\Pi_{\mathcal{S}}(A,\mathcal{U}_{r,c})$, which is now known as the Sinkhorn-Knopp algorithm or RAS method. Despite the simplicity of this algorithm and its good performance in practice, it has been difficult to analyze. As a result, recent work showing that $\Pi_{\mathcal{S}}(A,\mathcal{U}_{r,c})$ can be approximated in near-linear time [AZLOW17, CMTV17] has bypassed the Sinkhorn-Knopp algorithm entirely.Replacing the Proj step in Algorithm 1 with the matrix-scaling algorithm developed in [CMTV17] results in a runtime that is a single factor of $\varepsilon$ faster than what we present in Theorem 1. The benefit of our approach is that it is extremely easy to implement, whereas the matrix-scaling algorithm of [CMTV17] relies heavily on near-linear time Laplacian solver subroutines, which are not implementable in practice. In our work, we obtain a new analysis of the simple and practical Sinkhorn-Knopp algorithm, showing that it also approximates $\Pi_{\mathcal{S}}(A,\mathcal{U}_{r,c})$ in near-linear time.

Pseudocode for the Sinkhorn-Knopp algorithm appears in Algorithm 3. In brief, it is an alternating projection procedure which renormalizes the rows and columns of $A$ in turn so that they match the desired row and column marginals $r$ and $c$. At each step, it prescribes to either modify all the rows by multiplying row $i$ by $r_{i}/r_{i}(A)$ for $i\in[n]$, or to do the analogous operation on the columns. (We interpret the quantity $0/0$ as $1$ in this algorithm if ever it occurs.) The algorithm terminates when the matrix $A^{(k)}$ is sufficiently close to the polytope $\mathcal{U}_{r,c}$.

2 Prior work

3 New analysis of the Sinkhorn algorithm

Our new analysis allows us to obtain a dimension-independent bound on the number of iterations beyond the uniform case.

We first define some notation. Given a matrix $A$ and desired row and column sums $r$ and $c$, we define the potential (Lyapunov) function $f:{\rm I}\kern-1.79993pt{\rm R}^{n}\times{\rm I}\kern-1.79993pt{\rm R}^{n}\to{\rm I}\kern-1.79993pt{\rm R}$ by

This auxiliary function has appeared in much of the literature on Sinkhorn projections [KLRS08, CMTV17, KK96, KK93]. We call the vectors $x$ and $y$ scaling vectors. It is easy to check that a minimizer $(x^{*},y^{*})$ of $f$ yields the Sinkhorn projection of $A$: writing $X=\mathbf{D}(\exp(x^{*}))$ and $Y=\mathbf{D}(\exp(y^{*}))$, first order optimality conditions imply that $XAY$ lies in $\mathcal{U}_{r,c}$, and therefore $XAY=\Pi_{\mathcal{S}}(A,\mathcal{U}_{r,c})$.

The following lemma exactly characterizes the improvement in the potential function $f$ from an iteration of Sinkhorn, in terms of our current divergence to the target marginals.

If $k\geq 2$, then $\displaystyle f(x^{k-1},y^{k-1})-f(x^{k},y^{k})=\mathcal{K}(r\|r(A^{(k-1)}))+\mathcal{K}(c\|c(A^{(k-1)}))\,.$

Assume without loss of generality that $k$ is odd, so that $c(A^{(k-1)})=c$ and $r(A^{(k)})=r$. (If $k$ is even, interchange the roles of $r$ and $c$.) By definition,

where we have used that: $\|A^{(k-1)}\|_{1}=\|A^{(k)}\|_{1}=1$ and $Y^{(k)}=Y^{(k-1)}$; for all $i$, $r_{i}(x^{k}_{i}-x^{k-1}_{i})=r_{i}\log\frac{r_{i}}{r_{i}(A^{(k-1)})}$; and $\mathcal{K}(c\|c(A^{(k-1)}))=0$ since $c=c(A^{(k-1)})$. ∎

The next lemma has already appeared in the literature and we defer its proof to the Appendix.

For any probability measures $p$ and $q$, $\|p-q\|_{1}\leq\sqrt{2\mathcal{K}(p\|q)}$.

Let $k^{*}$ be the first iteration such that $\|r(A^{(k^{*})})-r\|_{1}+\|c(A^{(k^{*})})-c\|_{1}\leq\varepsilon^{\prime}$. Pinsker’s inequality implies that for any $k<k^{*}$, we have

4 Greedy Sinkhorn

In addition to a new analysis of Sinkhorn, we propose a new algorithm Greenkhorn which enjoys the same convergence guarantee but performs better in practice. Instead of performing alternating updates of all rows and columns of $A$, the Greenkhorn algorithm updates only a single row or column at each step. Thus Greenkhorn updates only $O(n)$ entries of $A$ per iteration, rather than $O(n^{2})$.

In this respect, Greenkhorn is similar to the stochastic algorithm for Sinkhorn projection proposed by [GCPB16]. There is a natural interpretation of both algorithms as coordinate descent algorithms in the dual space corresponding to row/column violations. Nevertheless, our algorithm differs from theirs in several key ways. Instead of choosing a row or column to update randomly, Greenkhorn chooses the best row or column to update greedily. Additionally, Greenkhorn does an exact line search on the coordinate in question since there is a simple closed form for the optimum, whereas the algorithm proposed by [GCPB16] updates in the direction of the average gradient. Our experiments establish that Greenkhorn performs better in practice; more details appear in the Appendix.

We emphasize that our algorithm is an extremely natural modification of Sinkhorn, and greedy algorithms for the scaling problem have been proposed before, though these do not come with with explicit near-linear time guarantees [PL82]. However, whereas previous analyses of Sinkhorn cannot be modified to extract any meaningful rates of convergence for greedy algorithms, our new analysis of Sinkhorn from Section 3.3 applies to Greenkhorn with only trivial modifications.

The choice of $\rho$ is justified by its appearance in Lemma 5, below. While $\rho$ is not a metric, it is easy to see that $\rho$ is nonnegative and satisfies $\rho(a,b)=0$ iff $a=b$.

We note that after $r(A)$ and $c(A)$ are computed once at the beginning of the algorithm, Greenkhorn can easily be implemented such that each iteration runs in only $O(n)$ time.

The analysis requires the following lemma, which is an easy modification of Lemma 2.

Let $A^{\prime}$ and $A^{\prime\prime}$ be successive iterates of Greenkhorn, with corresponding scaling vectors $(x^{\prime},y^{\prime})$ and $(x^{\prime\prime},y^{\prime\prime})$. If $A^{\prime\prime}$ was obtained from $A^{\prime}$ by updating row $I$, then

and if it was obtained by updating column $J$, then

We also require the following extension of Pinsker’s inequality (proof in Appendix).

For any $\alpha\in\Delta_{n},\beta\in{\rm I}\kern-1.79993pt{\rm R}_{+}^{n}$, define $\rho(\alpha,\beta)=\sum_{i}\rho(\alpha_{i},\beta_{i})$. If $\rho(\alpha,\beta)\leq 1$, then

We follow the proof of Theorem 2. Since the row or column update is chosen greedily, at each step we make progress of at least $\frac{1}{2n}(\rho(r,r(A))+\rho(c,c(A)))$. If $\rho(r,r(A))$ and $\rho(c,c(A))$ are both at most $1$, then under the assumption that $\|r(A)-r\|_{1}+\|c(A)-c\|_{1}>\varepsilon^{\prime}$, our progress is at least

Proof of Theorem 1

If $r,c\in\Delta_{n}$ and $F\in{\rm I}\kern-1.79993pt{\rm R}_{+}^{n\times n}$, then Algorithm 2 takes $O(n^{2})$ time to output a matrix $G\in\mathcal{U}_{r,c}$ satisfying

The proof of Lemma 7 is simple and left to the Appendix. (We also describe in the Appendix a randomized variant of Algorithm 2 that achieves a slightly better bound than Lemma 7). We are now ready to prove Theorem 1.

Error analysis. Let $B$ be the output of $\textsc{Proj}(A,\mathcal{U}_{r,c},\varepsilon^{\prime})$, and let $P^{*}\in\operatorname*{argmin}_{P\in\mathcal{U}_{r,c}}\langle P,C\rangle$ be an optimal solution to the original OT program.

We first show that $\langle B,C\rangle$ is not much larger than $\langle P^{*},C\rangle$. To that end, write $r^{\prime}:=r(B)$ and $c^{\prime}:=c(B)$. Since $B=XAY$ for positive diagonal matrices $X$ and $Y$, Lemma 1 implies $B$ is the optimal solution to

By Lemma 7, there exists a matrix $P^{\prime}\in\mathcal{U}_{r^{\prime},c^{\prime}}$ such that

Moreover, since $B$ is an optimal solution of (3), we have

where we have used the fact that $0\leq H(B),H(P^{\prime})\leq 2\log n$.

Lemma 7 implies that the output $\hat{P}$ of $\textsc{round}(B,\mathcal{U}_{r,c})$ satisfies the inequality $\|B-\hat{P}\|_{1}\leq 2\left(\|r^{\prime}-r\|_{1}+\|c^{\prime}-c\|_{1}\right)$. This fact together with (4) and Hölder’s inequality yields

Applying the guarantee of $\textsc{Proj}(A,\mathcal{U}_{r,c},\varepsilon^{\prime})$, we obtain

Plugging in the values of $\eta$ and $\varepsilon^{\prime}$ prescribed in Algorithm 1 finishes the error analysis.

Empirical results

Cuturi [Cut13] already gave experimental evidence that using Sinkhorn to solve (2) outperforms state-of-the-art techniques for optimal transport. In this section, we provide strong empirical evidence that our proposed Greenkhorn algorithm significantly outperforms Sinkhorn.

We run experiments on two datasets: real images, from mnist, and synthetic images, as in Figure 1.

We first compare the behavior of Greenkhorn and Sinkhorn on real images. To that end, we choose $10$ random pairs of images from the MNIST dataset, and for each one analyze the performance of ApproxOT when using both Greenkhorn and Sinkhorn for the approximate projection step. We add negligible noise $0.01$ to each background pixel with intensity . Figure 2 paints a clear picture: Greenkhorn significantly outperforms Sinkhorn both in the short and long term.

2 Random images

To better understand the empirical behavior of both algorithms in a number of different regimes, we devised a synthetic and tunable framework whereby we generate images by choosing a randomly positioned “foreground” square in an otherwise black background. The size of this square is a tunable parameter varied between 20%, 50%, and 80% of the total image’s area. Intensities of background pixels are drawn uniformly from $$; foreground pixels are drawn uniformly from$$. Such an image is depicted in Figure 1, and results appear in Figure 2.

We perform two other experiments with random images in Figure 3. In the first, we vary the number of background pixels and show that Greenkhorn performs better when the number of background pixels is larger. We conjecture that this is related to the fact that Greenkhorn only updates salient rows and columns at each step, whereas Sinkhorn wastes time updating rows and columns corresponding to background pixels, which have negligible impact. This demonstrates that Greenkhorn is a better choice especially when data is sparse, which is often the case in practice.

In the second, we consider the role of the regularization parameter $\eta$. Our analysis requires taking $\eta$ of order $\log n/\varepsilon$, but Cuturi [Cut13] observed that in practice $\eta$ can be much smaller. Cuturi showed that Sinkhorn outperforms state-of-the art techniques for computing OT distance even when $\eta$ is a small constant, and Figure 3 shows that Greenkhorn runs faster than Sinkhorn in this regime with no loss in accuracy.

Appendix A Omitted proofs

The proof of the first inequality is similar to the proof of Lemma 2:

where $\mathcal{K}(A^{(1)}\|A^{(0)})$ denotes the divergence between $A^{(1)}$ and $A^{(0)}$ viewed as elements of $\Delta_{n^{2}}$.

for all $i,j\in[n]$. Thus because $r$ and $c$ are both probability vectors,

A.2 Proof of Lemma 5

We prove only the case where a row was updated, since the column case is exactly the same.

Observe that $A^{\prime}$ and $A^{\prime\prime}$ differ only in the $I$th row, and $x^{\prime\prime}$ and $x^{\prime}$ differ only in the $I$th entry, and $y^{\prime\prime}=y^{\prime}$. Hence

where we have used the fact that $r_{I}(A^{\prime\prime})=r_{I}$ and $x^{\prime\prime}_{I}-x^{\prime}_{I}=\log(r_{I}/r_{I}(A^{\prime}))$. ∎

A.3 Proof of Lemma 6

Let $s=\sum_{i}\beta_{i}$, and write $\bar{\beta}=\beta/s$. The definition of $\rho$ implies

Note that both $s-1-\log s$ and $\mathcal{K}(\alpha\|\bar{\beta})$ are nonnegative. If $\rho(\alpha,\beta)\leq 1$, then in particular $s-1-\log s\leq 1$, and it can be seen that $s-1-\log s\geq(s-1)^{2}/5$ in this range. Applying Lemma 4 (Pinsker’s inequality) yields

By the triangle inequality and convexity,

The claim follows from the above two displays. ∎

A.4 Proof of Lemma 7

Let $G$ be the output of $\textsc{round}(F,\mathcal{U}_{r,c})$. The entries of $F^{\prime\prime}$ are nonnegative, and at the end of the algorithm $\text{err}_{r}$ and $\text{err}_{c}$ are both nonnegative, with $\|\text{err}_{r}\|_{1}=\|\text{err}_{c}\|_{1}=1-\|F^{\prime\prime}\|_{1}$. Therefore the entries of $G$ are nonnegative and

and likewise $c(G)=c$. This establishes that $G\in\mathcal{U}_{r,c}$.

Let us analyze both of the sums in (5). First, a simple calculation shows

Next, upper bound the second sum in (5) using the fact that the vector $c(F)$ is entrywise larger than $c(F^{\prime})$

Finally, we prove the $O(n^{2})$ runtime bound follows by observing that each rescaling and computing the matrix $\text{err}_{r}\text{err}_{c}^{\top}/\|\text{err}_{r}\|_{1}$ both require at most $O(n^{2})$ time. ∎

A.5 Randomized variant of rounding algorithm (Algorithm 2)

This asymmetry between $\|r(F)-r\|_{1}$ and $\|c(F)-c\|_{1}$ arises because Algorithm 2 creates $F^{\prime\prime}$ by first removing mass from rows of $F$, and then from columns. Consider modifying Algorithm 2 to create $F^{\prime\prime}$ by first removing mass from columns of $F$, and then from rows. Then a symmetrical argument gives the bound

Together the above two displays suggest the following simple randomized variant of Algorithm 2: with probability $1/2$, perform Algorithm 2; otherwise, perform the above-described column-then-row version of Algorithm 2. Combining the above two displays then gives the following improved bound for this randomized algorithm

A.6 Comparison with [GCPB16]

In this Section, we present an empirical comparison of the performance of Greenkhorn with the stochastic algorithm proposed by [GCPB16]. Their algorithm—which we call Stochastic Sinkhorn for convenience—uses a Stochastic Averaged Gradient (SAG) algorithm to optimize a dual version of the entropic penalty program (2).

We have noted in the main text that Greenkhorn and Stochastic Sinkhorn both attempt to solve the scaling problem via coordinate descent in the dual problem. Stochastic Sinkhorn does so via the method proposed in [SLRB17], whereas Greenkhorn greedily chooses a good coordinate to update, and then leverages an explicit closed form to perform an exact line search on this coordinate. One difference between our algorithms is their starting point: Greenkhorn is initialized with $A/\|A\|_{1}$, whereas the starting primal solution corresponding to the initialization of Stochastic Sinkhorn is the matrix obtained by first multiplying each column of $A$ by the corresponding entry of $c$ and then scaling the rows of the resulting matrix so they agree with $r$. This is equivalent to performing a full update step of Sinkhorn on the matrix $A\mathbf{D}(c)$ at the beginning of this algorithm. In simulations, this starting point is of better quality than the matrix $A/\|A\|_{1}$ which Greenkhorn uses as its first iterate; however, this advantage quickly disappears. Since our goal is to compare Greenkhorn and Stochastic Sinkhorn in terms of the number of required row or column updates, we also initialize Greenkhorn at this point instead of at $A/\|A\|_{1}$ to facilitate an apples-to-apples comparison.

To compare the performance of Greenkhorn with Stochastic Sinkhorn, we use an experiment on random images with $20$% foreground pixels, as in Section 5.2. We initialize both algorithms with the same primal solution and used Algorithm 2 to round iterates of each algorithm to the feasible polytope $\mathcal{U}_{r,c}$. Implementing Stochastic Sinkhorn requires choosing a step size, denoted by $C$ in [GCPB16]. That paper suggests choosing $C=1/(Ln)$, $3/(Ln)$, or $5/(Ln)$, where $L$ is an upper bound on the Lipschitz constant of the semi-dual problem they consider.In fact, they propose the step sizes $C=1/L,3/L,5/L$ in the main text, but the extra factor of $n$ is present in the simulation code posted online, so we have opted to retain it in our experiments. Our experimental results indicate that without the factor of $n$, the resulting algorithm is quite unstable. We compare all three choices of step size with our implementation of the Greenkhorn algorithm in Figure 4 with two different values of the parameter $\eta$.

Acknowledgments

We thank Michael Cohen, Adrian Vladu, Jon Kelner, Justin Solomon, and Marco Cuturi for helpful discussions. We are grateful to Pablo Parrilo for drawing our attention to the fact that Greenkhorn is a coordinate descent algorithm, and to Alexandr Andoni and Inderjit Dhillon for references.

References