Massively scalable Sinkhorn distances via the Nyström method

Introduction

Optimal transport is a fundamental notion in probability theory and geometry (Villani, 2008), which has recently attracted a great deal of interest in the machine learning community as a tool for image recognition (Li et al., 2013; Rubner et al., 2000), domain adaptation (Courty et al., 2017, 2014), and generative modeling (Bousquet et al., 2017; Arjovsky et al., 2017; Genevay et al., 2016), among many other applications (see, e.g., Peyré and Cuturi, 2017; Kolouri et al., 2017).

The growth of this field has been fueled in part by computational advances, many of them stemming from an influential proposal of Cuturi (2013) to modify the definition of optimal transport to include an entropic penalty. The resulting quantity, which Cuturi (2013) called the Sinkhorn “distance”We use quotations since it is not technically a distance; see (Cuturi, 2013, Section 3.2) for details. The quotes are dropped henceforth. after Sinkhorn (1967), is significantly faster to compute than its unregularized counterpart. Though originally attractive purely for computational reasons, the Sinkhorn distance has since become an object of study in its own right because it appears to possess better statistical properties than the unregularized distance both in theory and in practice (Genevay et al., 2018; Montavon et al., 2016; Peyré and Cuturi, 2017; Schiebinger et al., 2019; Rigollet and Weed, 2018). Computing this distance as quickly as possible has therefore become an area of active study.

for a parameter $\eta>0$. We stress that we use the squared Euclidean cost in our formulation of the Sinkhorn distance. This choice of cost—which in the unregularized case corresponds to what is called the $2$-Wasserstein distance (Villani, 2008)—is essential to our results, and we do not consider other costs here. The squared Euclidean cost is among the most common in applications (Schiebinger et al., 2019; Genevay et al., 2018; Courty et al., 2017; Forrow et al., 2018; Bousquet et al., 2017).

Many algorithms to compute $W_{\eta}(\mathbf{p},\mathbf{q})$ are known. Cuturi (2013) showed that a simple iterative procedure known as Sinkhorn’s algorithm had very fast performance in practice, and later experimental work has shown that greedy and stochastic versions of Sinkhorn’s algorithm perform even better in certain settings (Genevay et al., 2016; Altschuler et al., 2017). These algorithms are notable for their versatility: they provably succeed for any bounded, nonnegative cost. On the other hand, these algorithms are based on matrix manipulations involving the $n\times n$ cost matrix $C$, so their running times and memory requirements inevitably scale with $n^{2}$. In experiments, Cuturi (2013) and Genevay et al. (2016) showed that these algorithms could reliably be run on problems of size $n\approx 10^{4}$.

We show that a simple algorithm can be used to approximate $W_{\eta}(\mathbf{p},\mathbf{q})$ quickly on massive data sets. Our algorithm uses only known tools, but we give novel theoretical guarantees that allow us to show that the Nyström method combined with Sinkhorn scaling provably yields a valid approximation algorithm for the Sinkhorn distance at a fraction of the running time of other approaches.

We establish two theoretical results of independent interest: (i) New Nyström approximation results showing that instance-adaptive low-rank approximations to Gaussian kernel matrices can be found for data lying on a low-dimensional manifold (Section 3). (ii) New stability results about Sinkhorn projections, establishing that a sufficiently good approximation to the cost matrix can be used (Section 4).

2 Prior work

Computing the Sinkhorn distance efficiently is a well studied problem in a number of communities. The Sinkhorn distance is so named because, as was pointed out by Cuturi (2013), there is an extremely simple iterative algorithm due to Sinkhorn (1967) which converges quickly to a solution to (1). This algorithm, which we call Sinkhorn scaling, works very well in practice and can be implemented using only matrix-vector products, which makes it easily parallelizable. Sinkhorn scaling has been analyzed many times (Franklin and Lorenz, 1989; Linial et al., 1998; Kalantari et al., 2008; Altschuler et al., 2017; Dvurechensky et al., 2018), and forms the basis for the first algorithms for the unregularized optimal transport problem that run in time nearly linear in the size of the cost matrix (Altschuler et al., 2017; Dvurechensky et al., 2018). Greedy and stochastic algorithms related to Sinkhorn scaling with better empirical performance have also been explored (Genevay et al., 2016; Altschuler et al., 2017). Another influential technique, due to Solomon et al. (2015), exploits the fact that, when the distributions are supported on a grid, Sinkhorn scaling performs extremely quickly by decomposing the cost matrix along lower-dimensional slices.

Other algorithms have sought to solve (1) by bypassing Sinkhorn scaling entirely. Blanchet et al. (2018) proposed to solve (1) directly using second-order methods based on fast Laplacian solvers (Cohen et al., 2017; Allen-Zhu et al., 2017). Blanchet et al. (2018) and Quanrud (2019) have noted a connection to packing linear programs, which can also be exploited to yield near-linear time algorithms for unregularized transport distances.

Our main algorithm relies on constructing a low-rank approximation of a Gaussian kernel matrix from a small subset of its columns and rows. Computing such approximations is a problem with an extensive literature in machine learning, where it has been studied under many different names, e.g., Nyström method (Williams and Seeger, 2001), sparse greedy approximations (Smola and Schölkopf, 2000), incomplete Cholesky decomposition (Fine and Scheinberg, 2001), Gram-Schmidt orthonormalization (Shawe-Taylor and Cristianini, 2004) or CUR matrix decompositions (Mahoney and Drineas, 2009). The approximation properties of these algorithms are now well understood (Mahoney and Drineas, 2009; Gittens, 2011; Bach, 2013; Alaoui and Mahoney, 2015); however, in this work, we require significantly more accurate bounds than are available from existing results as well as adaptive bounds for low-dimensional data. To establish these guarantees, we follow an approach based on approximation theory (see, e.g., Rieger and Zwicknagl, 2010; Wendland, 2004; Belkin, 2018), which consists of analyzing interpolation operators for the reproducing kernel Hilbert space corresponding to the Gaussian kernel.

Finally, this paper adds to recent work proposing the use of low-rank approximation for Sinkhorn scaling (Altschuler et al., 2018; Tenetov et al., 2018). We improve upon those papers in several ways. First, although we also exploit the idea of a low-rank approximation to the kernel matrix, we do so in a more sophisticated way that allows for automatic adaptivity to data with low-dimensional structure. These new approximation results are the key to our adaptive algorithm, and this yields a significant improvement in practice. Second, the analyses of Altschuler et al. (2018) and Tenetov et al. (2018) only yield an approximation to $W_{\eta}(\mathbf{p},\mathbf{q})$ when $\eta\to\infty$. In the moderately regularized case when $\eta=O(1)$, which is typically used in practice, neither the work of Altschuler et al. (2018) nor of Tenetov et al. (2018) yields a rigorous error guarantee.

3 Outline of paper

Section 2 recalls preliminaries, and then formally states our main result and gives pseudocode for our proposed algorithm. The core of our theoretical analysis is in Sections 3 and 4. Section 3 presents our new results for Nyström approximation of Gaussian kernel matrices and Section 4 presents our new stability results for Sinkhorn scaling. Section 5 then puts these results together to conclude a proof for our main result (Theorem 1). Finally, Section 6 contains experimental results showing that our proposed algorithm outperforms state-of-the-art methods. The appendix contains proofs of several lemmas that are deferred for brevity of the main text.

Main result

Our goal is to approximate the Sinkhorn distance with parameter $\eta>0$:

to some additive accuracy $\varepsilon>0$. By strict convexity, this optimization problem has a unique minimizer, which we denote henceforth by $P^{\eta}$. For shorthand, in the sequel we write

Our approach is based on Sinkhorn scaling, an algorithm due to Sinkhorn (1967) and popularized for optimal transport by Cuturi (2013). We recall the following fundamental definition.

Since $\mathbf{p}$ and $\mathbf{q}$ remain fixed throughout, we abbreviate $\Pi^{\mathcal{S}}_{\mathcal{M}(\mathbf{p},\mathbf{q})}$ by $\Pi^{\mathcal{S}}$ except when we want to make the feasible set $\mathcal{M}(\mathbf{p},\mathbf{q})$ explicit.

Let $K$ have strictly positive entries, and let $\log K$ be the matrix defined by $(\log K)_{ij}:=\log(K_{ij})$. Then

Note that the strict convexity of $-H(P)$ and the compactness of $\mathcal{M}(\mathbf{p},\mathbf{q})$ implies that the minimizer exists and is unique.

This yields the following simple but key connection between Sinkhorn distances and Sinkhorn scaling.

where $K$ is defined by $K_{ij}=e^{-\eta C_{ij}}$.

Sinkhorn (1967) proposed to find $\Pi^{\mathcal{S}}(K)$ by alternately renormalizing the rows and columns of $K$. This well known algorithm has excellent performance in practice, is simple to implement, and is easily parallelizable since it can be written entirely in terms of matrix-vector products (Peyré and Cuturi, 2017, Section 4.2). Pseudocode for the version of the algorithm we use can be found in Appendix A.1.

2 Main result and proposed algorithm

Pseudocode for our proposed algorithm is given in Algorithm 1. Nys-Sink (pronounced “nice sink”) computes a low-rank Nyström approximation of the kernel matrix via a column sampling procedure. While explicit low-rank approximations of Gaussian kernel matrices can also be obtained via Taylor explansion (Cotter et al., 2011), our approach automatically adapts to the properties of the data set, leading to much better performance in practice.

As noted in Section 1, the Nyström method constructs a low-rank approximation to a Gaussian kernel matrix $K=e^{-\eta C}$ based on a small number of its columns. In order to design an efficient algorithm, we aim to construct such an approximation with the smallest possible rank. The key quantity for understanding the error of this algorithm is the so-called effective dimension (also sometimes called the “degrees of freedom”) of the kernel $K$ (Friedman et al., 2001; Zhang, 2005; Musco and Musco, 2017).

Let $\lambda_{j}(K)$ denote the $j$th largest eigenvalue of $K$ (with multiplicity). Then the effective dimension of $K$ at level $\tau>0$ is

where $d_{\textrm{eff}}(\cdot)$ is the effective rank for the kernel matrix $K:=e^{-\eta C}$.

We note also that although the present paper focuses specifically on the squared Euclidean cost $c(x_{i},x_{j})=\|x_{i}-x_{j}\|_{2}^{2}$ (corresponding to the $2$-Wasserstein case of optimal transport pervasively used in applications; see intro), our algorithm Nys-Sink readily extends to other cases of optimal transport. Indeed, since the Nyström method works not only for Gaussian kernel matrices $K_{ij}=e^{-\eta\|x_{i}-x_{j}\|_{2}^{2}}$, but in fact more generally for any PSD kernel matrix, our algorithm can be used on any optimal transport instance for which the corresponding kernel matrix $K_{ij}=e^{-\eta c(x_{i},x_{j})}$ is PSD.

We note that, while our algorithm is randomized, we obtain a deterministic guarantee that $\hat{P}$ is a good solution. We also note that runtime dependence on the radius $R$—which governs the scale of the problem—is inevitable since we seek an additive guarantee.

Crucially, we show in Section 3 that $r^{*}$—which controls the running time of the algorithm with high probability by (3d)—adapts to the intrinsic dimension of the data. This adaptivity is crucial in applications, where data can have much lower dimension than the ambient space. We informally summarize this behavior in the following theorem.

For any $k$-dimensional manifold $\Omega$ satisfying certain technical conditions and $\eta>0$, there exists a constant $c_{\Omega,\eta}$ such that for any $n$ points lying on $\Omega$,

The formal versions of these bounds appear in Section 3. The second bound is significantly better than the first when $k\ll d$, and clearly shows the benefits of an adaptive procedure.

Combining Theorems 1 and 2 yields the following time and space complexity for our algorithm.

Moreover, if $X$ lies on a $k$-dimensional manifold $\Omega$, then with high probability Algorithm 1 requires

Altschuler et al. (2017) noted that an approximation to the unregularized optimal transport cost is obtained by taking $\eta=\Theta\left(\varepsilon^{-1}\log n\right)$. Thus it follows that Algorithm 1 computes an additive $\varepsilon$ approximation to the unregularized transport distance in $O\left(n\left(\varepsilon^{-1}R^{2}\log n\right)^{O(d)}\right)$ time with high probability. However, a theoretically better running time for that problem can be obtained by a simple but impractical algorithm based on rounding the input distributions to an $\varepsilon$-net and then running Sinkhorn scaling on the resulting instance.We are indebted to Piotr Indyk for inspiring this remark.

Kernel approximation via the Nyström method

In this section, we describe the algorithm AdaptiveNyström used in line 4 of Algorithm 1 and bound its runtime complexity, space complexity, and error. We first establish basic properties of Nyström approximation and give pseudocode for AdaptiveNyström (Sections 3.1 and 3.2) before stating and proving formal versions of the bounds appearing in Theorem 2 (Sections 3.3 and 3.4).

We now turn to understanding the approximation error of this method. In this paper we will sample the set $X_{r}$ via approximate leverage-score sampling. In particular, we do this via Algorithm 2 of Musco and Musco (2017). The following lemma shows that taking the rank $r$ to be on the order of the effective dimension $d_{\textrm{eff}}(\tau)$ (see Definition 3) is sufficient to guarantee that $\widetilde{K}$ approximates $K$ to within error $\tau$ in operator norm.

Let $\tau,\delta>0$. Consider sampling $X_{r}$ from $X$ according to Algorithm 2 of Musco and Musco (2017), for some positive integer $r\geqslant 400d_{\textrm{eff}}(\tau)\log\tfrac{3n}{\delta}.$ Then:

The result follows directly from Theorem 7 of Musco and Musco (2017) and the fact that $d_{\textrm{eff}}(\tau)\leqslant\textrm{rank}(K)\leqslant n$ for any $\tau\geqslant 0$. ∎

2 Adaptive Nyström with doubling trick

Below, line 6 in Algorithm 2 denotes the approximate leverage-score sampling scheme of Musco and Musco (2017, Algorithm 2) when applied to the Gaussian kernel matrix $K_{ij}:=e^{-\eta\|x_{i}-x_{j}\|^{2}}$. We note that the BLESS algorithm of Rudi et al. (2018) allows for re-using previously sampled points when doubling the sampling rank. Although this does not affect the asymptotic runtime, it may lead to speedups in practice.

The algorithm used $O(nr)$ space and terminated in $O(nr^{2})$ time.

There exists a universal constant $c$ such that simultaneously for every $\delta>0$,

3 General results: data points lie in a ball

For each $\tau\in(0,1]$, $d_{\textrm{eff}}(\tau)\leqslant 3\left(6+\frac{41}{d}\eta R^{2}+\frac{3}{d}\log\frac{1}{\tau}\right)^{d}$.

On the other hand, by the Eckart-Young-Mirsky Theorem,

Therefore by combining the above two displays, we conclude that

Proofs of the two claims follow by bounding this quantity. Details are in Appendix B.5.

Theorem 3 characterizes the eigenvalue decay and effective dimension of Gaussian kernel matrices in terms of the dimensionality of the space, with explicit constants and explicit dependence on the width parameter $\eta$ and the radius $R$ of the ball (see Belkin, 2018, for asymptotic results). This yields the following bound on the optimal rank for approximating Gaussian kernel matrices of data lying in a Euclidean ball.

Directly from the explicit bound of Theorem 3 and the definition of $r^{*}(X,\eta,\varepsilon^{\prime})$. ∎

4 Adaptivity: data points lie on a low dimensional manifold

Let $h_{X,\Omega}=\sup_{x^{\prime}\in\Omega}\inf_{x\in X}\|x-x^{\prime}\|$. Let $\mathcal{H}$ be the RKHS associated to the Gaussian kernel of a given width. There exist $c,h>0$ not depending on $X,n$, such that, when $h_{X,\Omega}\leqslant h$ the following holds

Let $K$ be the Gaussian kernel matrix associated to $X$. Then there exists a constant $c$ not depending on $X$ or $n$, for which

Let $\tau\in(0,1]$. Let $K$ be the Gaussian kernel matrix associated to $X$ and $d_{\textrm{eff}}(\tau)$ the effective dimension computed on $K$. There exists $c_{1},c_{2}$ not depending on $X$, $n$, or $\tau$, for which

and the space of $W^{m}_{p}(B)$ as $W^{m}_{p}(B)=\overline{C^{\infty}(B)}^{\|\cdot\|_{W^{m}_{p}(B)}}$.

For any $j\in[T]$, $u\in\mathcal{H}$ we have the following. By O, we have that there exists a constant $C_{d,k,R,r_{j}}$ such that for any $q\geqslant k$,

Now note that by Theorem 7.5 of Rieger and Zwicknagl (2010) we have that there exists a constant $C_{\eta}$ such that

Then, since $q^{m}\leqslant m^{m}(1+m)^{q}$, for any $q\geqslant 1$, we have

for a suitable constant $C_{d,k,R,r_{j},Q,\eta}$ depending on $C_{d,k,R,r_{j},Q}$, $C_{\eta}$ and $(d+1)/2$.

In particular we want to study $\|u\|_{L^{\infty}(\Omega)}$, for $u=f-\widehat{f}_{X}$. We have

Now for $j\in[T]$, denote by $Z_{j}$ the set $Z_{j}=\{\Psi_{j}(x)|x\in X\cap U_{j}\}$. By construction of $u=f-\widehat{f}_{X}$, we have

Define $h_{Z_{j},B^{k}_{r_{j}}}=\sup_{z\in B^{k}_{r_{j}}}\inf_{z^{\prime}\in Z_{j}}\|z-z^{\prime}\|$. We have established that there exists $C>0$, such that $\|u\circ\Psi_{j}^{-1}\|_{W^{q}_{2}(B^{k}_{r_{j}})}\leqslant C^{q}q^{\frac{3}{2}q}\|u\|_{\mathcal{H}}$, and by construction $(u\circ\Psi_{j})|_{Z_{j}}=0$. We can therefore apply Theorem 3.5 of Rieger and Zwicknagl (2010) to obtain that there exists a $c_{j},h_{j}>0$, for which, when $h_{Z_{j},B^{k}_{r_{j}}}\leqslant h_{j}$, then

Now, denote by $\bar{h}_{S,U}=\sup_{x^{\prime}\in U}\inf_{x\in S}d(x,x^{\prime})$ with $d$ the geodesic distance over the manifold $\Omega$. By applying Theorem 8 of Fuselier and Wright (2012), we have that there exist $C$ and $h_{0}$ not depending on $X$ or $n$ such that, when $\bar{h}_{X,\Omega}\leqslant h_{0}$, the inequality $\bar{h}_{X_{j}\cap U_{j},U_{j}}\leqslant C\bar{h}_{X,\Omega}$ holds for any $j\in[T]$. Moreover, since by Theorem 6 of the same paper $\|x-x^{\prime}\|\leqslant d(x,x^{\prime})\leqslant C_{1}\|x-x^{\prime}\|$, for $C_{1}>1$ and $x,x^{\prime}\in\Omega$, then

Finally, defining $c_{1}=c(2\max_{j}C_{j})^{-2/5},h=C_{1}^{-1}\min(h_{0},C^{-1}\min_{j}h_{j})$, when $h_{X,\Omega}\leqslant h$,

The proof of Points 2 and 3 now proceeds as in Theorem 3. Details are deferred to Appendix B.5. ∎

Point 1 of the result above is new, to our knowledge, and extends interpolation results on manifolds (Wendland, 2004; Fuselier and Wright, 2012; Hangelbroek et al., 2010), from polynomial to exponential decay, generalizing a technique of Rieger and Zwicknagl (2010) to a subset of real analytic manifolds. Points 2 and 3 are a generalization of Theorem 3 to the case of manifolds. In particular, the crucial point is that now the eigenvalue decay and the effective dimension depend on the dimension of the manifold $k$ and not the ambient dimension $d\gg k$. We think that the factor $5/2$ in the exponent of the eigenvalues and effective dimension is a result of the specific proof technique used and could be removed with a refined analysis, which is out of the scope of this paper.

We finally conclude the desired bound on the optimal rank in the manifold case.

By the definition of $r^{*}(X,\eta,\varepsilon^{\prime})$ and the bound of Theorem 4, we have

Since $\log\tfrac{2n}{\varepsilon^{\prime}}\geqslant 1$, we may set $c_{\Omega,\eta}=\max\left\{(8c_{1}\eta R^{2})^{5k/2},c_{2}\right\}$ to obtain the claim. ∎

Sinkhorn scaling an approximate kernel matrix

The running time bound in Theorem 5 for the time required to produce $D_{1}$ and $D_{2}$ follows directly from prior work which has shown that Sinkhorn scaling can produce an approximation to the Sinkhorn projection of a positive matrix in time nearly independent of the dimension $n$.

The remainder of the section is devoted to proving the error bounds in Theorem 5. Subsection 4.1 proves stability bounds for using an approximate kernel matrix, Subsection 4.2 proves stability bounds for using an approximate Sinkhorn projection, and then Subsection 4.3 combines these results to prove the error bounds in Theorem 5.

In words, Proposition 2 establishes that the Sinkhorn projection operator is Lipschitz on the “logarithmic scale.” By contrast, we show in Appendix C that the Sinkhorn projection does not satisfy a Lipschitz property in the standard sense for any choice of matrix norm.

2 Using an approximate Sinkhorn projection

Here we present the second ingredient for the proof of Theorem 5: that the objective function $V_{C}(\cdot)$ for Sinkhorn distances in (1) is stable with respect to the target row and column sums $\mathbf{p}$ and $\mathbf{q}$ of the outputted matrix.

3 Proof of Theorem 5

Proof of Theorem 1

In this section, we combine the results of the preceding three sections to prove Theorem 1.

Next, we prove (3b). By Proposition 1, $P^{\eta}=\operatorname*{argmin}_{P\in\mathcal{M}(\mathbf{p},\mathbf{q})}V_{C}(P)$. Thus

where above the first inequality is by (7), the equality is by Lemma F, and the final inequality is by first-order KKT conditions which give $\nabla V_{C}(P^{\eta})(\hat{P}-P^{\eta})\geqslant 0$. After rearranging, we conclude that $\mathsf{KL}(\hat{P}\|P^{\eta})\leqslant\eta\varepsilon$, proving (3b).

Experimental results

In this section we empirically validate our theoretical results. To run our experiments, we used a desktop with 32GB ram and 16 cores Xeon E5-2623 3GHz. The code is optimized in terms of matrix-matrix and matrix-vector products using BLAS-LAPACK primitives.

Fig. 1 plots the time-accuracy tradeoff for Nys-Sink, compared to the standard Sinkhorn algorithm. This experiment is run on random point clouds of size $n\approx 20000$, which corresponds to cost matrices of dimension approximately $20000\times 20000$. Fig. 1 shows that Nys-Sink is consistently orders of magnitude faster to obtain the same accuracy.

Next, we investigate Nys-Sink’s dependence on the intrinsic dimension and ambient dimension of the input. This is done by running Nys-Sink on distributions supported on $1$-dimensional curves embedded in higher dimensions, illustrated in Fig. 2, left. Fig. 2, right, indicates that an approximation rank of $r=300$ is sufficient to achieve an error smaller than $10^{-4}$ for any ambient dimension $5\leqslant d\leqslant 100$. This empirically validates the result in 4, namely that the approximation rank – and consequently the computational complexity of Nys-Sink – is independent of the ambient dimension.

Finally, we evaluate the performance of our algorithm on a benchmark dataset used in computer graphics: we measure Wasserstein distance between 3D cloud points from “The Stanford 3D Scanning Repository”http://graphics.stanford.edu/data/3Dscanrep/. In the first experiment, we measure the distance between armadillo ($n=1.7\times 10^{5}$ points) and dragon (at resolution 2, $n=1.0\times 10^{5}$ points), and in the second experiment we measure the distance between armadillo and xyz-dragon which has more points ($n=3.6\times 10^{6}$ points). The point clouds are centered and normalized in the unit cube. The regularization parameter is set to $\eta=15$, reflecting the moderate regularization regime typically used in practice.

We compare our algorithm (Nys-Sink)—run with approximation rank $r=2000$ for $T=20$ iterations on a GPU—against two algorithms implemented in the library GeomLosshttp://www.kernel-operations.io/geomloss/. These algorithms are both highly optimized and implemented for GPUs. They are: (a) an algorithm based on an annealing heuristic for $\eta$ (controlled by the parameter $\alpha$, such that at each iteration $\eta_{t}=\alpha\eta_{t-1}$, see Kosowsky and Yuille, 1994b) and (b) a multiresolution algorithm based on coarse-to-fine clustering of the dataset together with the annealing heuristic (Schmitzer, 2019). Table 1 reports the results, which demonstrate that our method is comparable in terms of precision, and has computational time that is orders of magnitude smaller than the competitors. We note the parameters $r$ and $T$ for Nys-Sink are chosen by hand to balance precision and time complexity.

We note that in these experiments, instead of using Algorithm 2 to choose the rank adaptively, we simply run experiments with a small fixed choice of $r$. As our experiments demonstrate, Nys-Sink achieves good empirical performance even when the rank $r$ is smaller than our theoretical analysis requires. Investigating this empirical success further is an interesting topic for future study.

Appendix A Pseudocode for subroutines

Moreover, the matrices $\log(D_{1})$ and $\log(D_{2})$ can each be formed in $O(n)$ time, so computing $\hat{W}$ takes time $O(\textsc{T}_{\textsc{mult}}+n)$, as claimed. ∎

A.2 Pseudocode for rounding algorithm

For completeness, here we briefly recall the rounding algorithm Round from (Altschuler et al., 2017) and prove a slight variant of their Lemma 7 that we need for our purposes.

Moreover, the algorithm only uses $O(1)$ matrix-vector products with $F$ and $O(n)$ additional processing time.

The runtime claim is clear. Next, let $\Delta:=\|F\|_{1}-\|F^{\prime\prime}\|_{1}=\sum_{i=1}^{n}(r_{i}(F)-\mathbf{p}_{i})_{+}+\sum_{j=1}^{n}(c_{j}(F^{\prime})-\mathbf{q}_{j})_{+}$ denote the amount of mass removed from $F$ to create $F^{\prime\prime}$. Observe that $\sum_{i=1}^{n}(r_{i}(F)-\mathbf{p}_{i})_{+}=\operatorname{\frac{1}{2}}\|r(F)-\mathbf{p}\|_{1}$. Since $F^{\prime}\leqslant F$ entrywise, we also have $\sum_{j=1}^{n}(c_{j}(F^{\prime})-\mathbf{q}_{j})_{+}\leqslant\sum_{j=1}^{n}(c_{j}(F)-\mathbf{q}_{j})_{+}=\operatorname{\frac{1}{2}}\|c(F)-\mathbf{q}\|_{1}$. Thus $\Delta\leqslant\operatorname{\frac{1}{2}}(\|r(F)-\mathbf{p}\|_{1}+\|c(F)-q\|_{1})$. The proof is complete since $\|F-G\|_{1}\leqslant\|F-F^{\prime\prime}\|_{1}+\|F^{\prime\prime}-G\|_{1}=2\Delta$. ∎

Appendix B Omitted proofs

Let $P,Q\in\Delta_{n\times n}$. If $\|P-Q\|_{1}\leqslant\delta\leqslant 1$, then

By Ho and Yeung (2010, Theorem 6), |H(P)-H(Q)|\leqslant\frac{\delta}{2}\log(n^{2}-1)+h\big{(}\frac{\delta}{2}\big{)}, where $h$ is the binary entropy function. If $\delta\leqslant 1$, then $h(\frac{\delta}{2})\leqslant\delta\log\frac{2}{\delta}$, which yields the claim. ∎

B.2 Bregman divergence of Sinkhorn distances

The remainder in the first-order Taylor expansion of $V_{C}(\cdot)$ between any two joint distributions is exactly the KL-divergence between them.

Observing that $\nabla V_{C}(P)$ has $ij$th entry $C_{ij}+\eta^{-1}(1+\log P_{ij})$, we expand the right hand side as $[\langle C,P\rangle+\eta^{-1}\sum_{ij}P_{ij}\log P_{ij}]+[\langle C,Q-P\rangle+\eta^{-1}\sum_{ij}(Q_{ij}-P_{ij})\log P_{ij}\rangle]+[\eta^{-1}\sum_{ij}Q_{ij}\log\frac{Q_{ij}}{P_{ij}}]=\langle C,Q\rangle+\eta^{-1}\sum_{ij}Q_{ij}\log Q_{ij}=V_{C}(Q)$. ∎

B.3 Hausdorff distance between transport polytopes

where $d_{H}(A,B)$ is the Hausdorff distance between $A$ and $B$ with respect to $\|\cdot\|$.

Interchanging the role of $A$ and $B$ yields the claim. ∎

B.4 Miscellaneous helpful lemmas

where $\|\cdot\|_{*}$ denotes the dual norm to $\|\cdot\|$.

which implies the claim via the definition of the dual norm. ∎

Without loss of generality, assume $a\geqslant b$. Then $\log a-\log b=\log\tfrac{a}{b}\leqslant\tfrac{a}{b}-1=\tfrac{a-b}{\min\{a,b\}}\,,$ as claimed. ∎

Since $\|x_{i}\|_{2}\leqslant R$ for all $i\in[n]$, the matrix $K$ satisfies $K_{ij}=e^{-\eta\|x_{i}-x_{j}\|_{2}^{2}}\geqslant e^{-4\eta R^{2}}$ for all $i,j\in[n]$. Hence $\widetilde{K}_{ij}\geqslant\tfrac{\varepsilon^{\prime}}{2}e^{-4\eta R^{2}}$ for all $i,j\in[n]$ and thus by Lemma K,

and bound the three terms separately. First, the assumptions imply that $\eta\varepsilon\leqslant n$ and $2\|C\|_{\infty}+3\eta^{-1}\leqslant 5\|C\|_{\infty}$. We therefore have

Since $\|C\|_{\infty}\eta\geqslant 1$, we likewise obtain

Finally, the fact that $\tfrac{\eta^{-1}\delta}{\varepsilon}\leqslant\tfrac{1}{50}$ and $x\log\frac{1}{x}\leqslant\tfrac{1}{10}$ for $x\leqslant\tfrac{1}{50}$ yields

B.5 Supplemental results for Section 3

By the Eckart-Young-Mirsky Theorem, we have

Therefore by combining the above two displays, we conclude that

Let $T_{\tau}$ be such that $\varepsilon(T_{\tau})\leqslant\tau$. We can then bound $\varepsilon(T)/(\varepsilon(T)+\tau)$ above by $1$ for $T\leqslant T_{\tau}-1$ and by $\varepsilon(T)/\tau$ for $T\geqslant T_{\tau}$, obtaining

In particular, we can choose $T_{\tau}=d+2e^{2}\eta R^{2}+\log(1/\tau)$. Since $\log\frac{T_{\tau}}{2e\eta R^{2}}>1$, for any $T\geqslant T_{\tau}$, then $\varepsilon(T)\leqslant e^{-T\log\frac{T}{2e\eta R^{2}}}\leqslant e^{-T}$. Moreover since $M_{T+1}-M_{T}=dM_{T}/(T+1)$, and $M_{T}\leqslant e^{d}(1+T/d)^{d}$, we have

Finally, by changing variables, $x=u+T_{\tau}$ and $u=(d+T_{\tau})z$,

where for the last equality we used the characterization of the incomplete gamma function $\Gamma(a,z)=z^{-a}e^{-z}\int_{0}^{\infty}(1+t)^{a-1}e^{-zt}dt$ (see Eq. 8.6.5 of Olver et al., 2010). To complete the proof note that by P we have $\Gamma(a,z)\leqslant z/(z-a)z^{a-1}e^{-z}$, for any $z>a>0$. Since $\log(1/\tau)\geqslant 0$ for $\tau\in(0,1]$, we have $(d+T_{\tau})/(T_{\tau}-1)\leqslant 2$ and $(de^{-T_{\tau}})/(\tau T_{\tau})\leqslant 1$, so

B.5.2 Full proof of Theorem 4

The proof of Points 2 and 3 here is completely analogous to the proof of Points 1 and 2, respectively, in Theorem 3.

Since $B$ is of rank $p$, the the Eckart-Young-Mirsky Theorem again implies $\lambda_{p+1}(K)\leqslant ne^{-c\varepsilon^{-2/5}}$. We conclude by recalling that $\varepsilon\leqslant(p/C_{0})^{-1/k}$.

Let $M_{\tau}$ be such that $\lambda_{M_{\tau}+1}\leqslant n\tau$. By Point 2, this holds if we take $M_{\tau}=(c_{0}\log\tfrac{1}{\tau})^{5k/2}$ for a sufficiently large constant $c$. By definition of $d_{\textrm{eff}}(\tau)$ and the fact that $x/(x+\lambda)\leqslant\min(1,x/\lambda)$ for any $x\geqslant 0,\lambda>0$, we have

Denoting $\beta:=\tfrac{2}{5k}$ for shorthand, we can upper bound the sum as follows:

where above the second step was by the change of variables $u:=cx^{\beta}$, the third step was by Cauchy-Schwartz with respect to the inner product $\langle f,g\rangle:=\int_{0}^{\infty}f(u)g(u)e^{-u}du$, and the final line was for some constant $c_{k}$ only depending on $k$, whenever $c_{0}$ is taken to be at least $\tfrac{2}{c}$. This proves the claim. ∎

B.5.3 Additional bounds

Now note that by the properties of $l_{j},k_{j}$, we have that $|\nu|=|\sum_{j}|k_{j}|l_{j}|=\sum_{j}|k_{j}||l_{j}|,$ then

To conclude, denote by $S^{n}_{k}$ the Stirling numbers of the second kind. By Constantine and Savits (1996, Corollary 2.9) and Rennie and Dobson (1969) we have

First note that $\|\cdot\|_{L^{\infty}(B^{d}_{R})}\leqslant C_{d,R}\|\cdot\|_{W^{(d+1)/2}_{2}(B^{d}_{R})}$ (Adams and Fournier, 2003) for a constant $C_{d,R}$ depending only on $d$ and $R$. Therefore

Moreover note that $\|D^{\alpha}f\|_{W^{(d+1)/2}_{2}(B^{d}_{R})}\leqslant\|f\|_{W^{|\alpha|+(d+1)/2}_{2}(B^{d}_{R})}$. By N we have that

By definition of Sobolev space $W^{q}_{2}(B^{k}_{r_{j}})$, we have

where $C_{q}:=\sum_{|\alpha|\leqslant q}(2|\alpha|dQ)^{|\alpha|}$. Then,

The final result is obtained via the bound $C_{q}\leqslant(2qdQ)^{q}\binom{k+q}{k}\leqslant(2ek)^{k}(2qdQ)^{q}q^{k}$ for $q\geqslant k$. ∎

Denote by $\Gamma(a,x)$ the function defined as

In particular $\Gamma(a,x)\leqslant 2x^{a-1}e^{-x},$ for $x\geqslant 2(a-1)_{+}~{}\wedge~{}x>0$.

Assume $x>0$. When $a\leqslant 1$, the function $z^{a-1}e^{-z}$ is decreasing and in particular $z^{a-1}e^{-z}\leqslant x^{a-1}e^{-z}$ for $z\geqslant x$, so when $z\geqslant x$ we have

When $a>1$, for any $\tau\in(0,1)$, we have

Now note that the maximum of $z^{a-1}e^{-\tau z}$ is reached when $z=(a-1)/\tau$. When $x\geqslant a-1$, we can set $\tau=(a-1)/x$, so the maximum of $z^{a-1}e^{-\tau z}$ is exactly in $z=x$. In that case $\sup_{z\geqslant x}(z^{a-1}e^{-\tau z})=x^{a-1}e^{-\tau x}$ and

The final result is obtained by gathering the cases $a\leqslant 1$ and $a>0$ in the same expression. ∎

By the change of variable $x=(u/q_{2})^{1/b}$ we have

Appendix C Lipschitz properties of the Sinkhorn projection

We give a simple construction illustrating that the Sinkhorn projection operator is not Lipschitz in the standard sense. This stands in contrast with Proposition 2, which illustrates that this projection is Lipschitz on the logarithmic scale.

This non-Lipschitz result holds even for the following simple rescaling of the $2\times 2$ Birkhoff polyope:

By the equivalence of finite-dimensional norms, it suffices to prove this for $\|\cdot\|_{1}$, for which we will show

For $\varepsilon,\delta\in(0,1)$, define the matrix

and let $P_{\varepsilon,\delta}:=\Pi_{\mathcal{M}}^{\mathcal{S}}(K_{\varepsilon,\delta})$ denote the Sinkhorn projection of $K_{\varepsilon,\delta}$ onto $\mathcal{M}$. The polytope $\mathcal{M}$ is parameterizable by a single scalar as follows:

By definition, $P_{\varepsilon,\delta}$ is the unique matrix in $\mathcal{M}$ of the form $D_{1}K_{\varepsilon,\delta}D_{2}$ for positive diagonal matrices $D_{1}$ and $D_{2}$. Taking

for $\beta_{\varepsilon,\delta}=2(\sqrt{\delta(1-\varepsilon)}+\sqrt{\varepsilon(1-\delta)})$, we verify $D_{1}K_{\varepsilon,\delta}D_{2}=M_{a_{\varepsilon,\delta}}$, where $a_{\varepsilon,\delta}:=\frac{\sqrt{\delta(1-\varepsilon)}}{\beta_{\varepsilon,\delta}}$. Therefore $P_{\varepsilon,\delta}=M_{a_{\varepsilon,\delta}}$ for $\varepsilon,\delta\in(0,1)$.

Now parameterize $\varepsilon:=c\delta$ for some fixed constant $c\in(0,\infty)$ and consider taking $\delta\to 0^{+}$. Then $a_{c\delta,\delta}=\frac{\sqrt{1-c\delta}}{2\left[\sqrt{c(1-\delta)}+\sqrt{1-c\delta}\right]}$, which for fixed $c$ becomes arbitrarily close to

as $\delta$ approaches . Thus $\|P_{c\delta,\delta}-M_{g(c)}\|_{1}=o_{\delta}(1)$ and similarly $\|P_{\delta/c,\delta}-M_{g(1/c)}\|_{1}=o_{\delta}(1)$. We therefore conclude that for any constant $c\in(0,\infty)\setminus\{1\}$, although

vanishes as $\delta\to 0^{+}$, the quantity

does not vanish. Therefore combining the above two displays and taking, e.g., $c=2$ proves (9). ∎

References