# Low-Rank Sinkhorn Factorization

Meyer Scetbon  Marco Cuturi  Gabriel Peyré

###### Abstract

Several recent applications of optimal transport (OT) theory to machine
learning have relied on regularization, notably entropy and the Sinkhorn
algorithm. Because matrix-vector products are pervasive in the Sinkhorn
algorithm, several works have proposed to approximate kernel matrices
appearing in its iterations using low-rank factors. Another route lies instead
in imposing low-rank constraints on the feasible set of couplings considered
in OT problems, with no approximations on cost nor kernel matrices. This route
was first explored by Forrow et al. (2018), who proposed an algorithm tailored
for the squared Euclidean ground cost, using a proxy objective that can be
solved through the machinery of regularized 2-Wasserstein barycenters.
Building on this, we introduce in this work a generic approach that aims at
solving, in full generality, the OT problem under low-rank constraints with
arbitrary costs. Our algorithm relies on an explicit factorization of low rank
couplings as a product of sub-coupling factors linked by a common marginal;
similar to an NMF approach, we alternatively updates these factors. We prove
the non-asymptotic stationary convergence of this algorithm and illustrate its
efficiency on benchmark experiments.

Machine Learning, ICML

Figure 1: Two Gaussian mixture densities evaluated on $n=200$ and $m=220$
sized grids in 1D, displayed as blue/red curves. Between them, $n\times m$
optimal coupling matrices obtained by our proposed low-rank OT method for
varying rank constraint values $r$ (in increasing order, top row) and the
Sinkhorn algorithm, for various $\varepsilon$ (in decreasing order, bottom
row). The ground cost is the 1.5-norm.

##  1 Introduction

By providing a simple and comprehensive framework to compare probability
distributions, optimal transport (OT) theory has inspired many developments in
machine learning (Peyré & Cuturi, 2019). A flurry of works have recently
connected it to other trending topics, such as normalizing flows or convex
neural networks (Makkuva et al., 2020; Korotin et al., 2021; Tong et al.,
2020), while the scope of its applications has now reached several fields of
science such as single-cell biology (Schiebinger et al., 2019; Yang et al.,
2020), imaging (Schmitz et al., 2018; Heitz et al., 2020) or neuroscience
(Janati et al., 2020; Koundal et al., 2020).

Challenges when computing OT. Solving optimal transport problems at scale
poses, however, formidable challenges. The most obvious among them is
computational: Instantiating the Kantorovich (1942) problem on discrete
measures of size $n$ can be solved with a linear program (LP) of complexity
$O(n^{3}\log n)$. A second and equally important challenge lies in the
statistical performance of using that LP to estimate OT between densities: the
LP solution between i.i.d samples converges very slowly to that between
densities (Fournier & Guillin, 2015). It is now increasingly clear that
regularizing OT in some way or another is the only way to mitigate these two
issues (Genevay et al., 2018; Chizat et al., 2020; Clason et al., 2021). A
popular approach consists in penalizing the OT problem with a strongly convex
function of the coupling (Cuturi, 2013; Dessein et al., 2018). We explore in
this work an alternative, and more direct approach to add regularity: we
restrict, instead, the set of feasible couplings to have a small nonnegative
rank.

Low-Rank Kernel Factorization. Low-rank factorizations are not new to
regularized OT. They have been used to speed-up the resolution of entropy
regularized OT with the Sinkhorn algorithm, pending some approximations: Given
a data-dependent $n\times m$ cost matrix $C$, the Sinkhorn iterations consist
in matrix-vector products of the form $Kv$ or $K^{T}u$ where
$K\triangleq\exp(-C/\varepsilon)$ and $u,v$ are $n,m$\- vectors. Altschuler et
al. (2018) and Altschuler & Boix-Adsera (2020) have proposed to approximate
the kernel $K$ with a product of thin rank $r$ matrices,
$\widetilde{K}=AB^{T}$. Naturally, the ability to approximate $K$ with a low-
rank $\widetilde{K}$ degrades as $\varepsilon$ decreases, making this approach
valid only for sufficiently large $\varepsilon$. Thanks to this approximation,
however, each Sinkhorn iteration is linear in $n$ or $m$, and the coupling
outputted by the Sinkorn algorithm is of the form $\widetilde{P}=CD^{T}$ where
$C=\operatorname{diag}(u)A$, $D=\operatorname{diag}(v)B$. This approximation
results therefore in a _low-rank_ solution that is not, however, rigorously
optimal for the original problem as defined by $K$ but rather that defined by
$\widetilde{K}$. Similarly, Scetbon & Cuturi (2020) consider instead
nonnegative low-rank approximations for $K$ of the form $\widetilde{K}=QR^{T}$
where $Q,R>0$.x

Low-Rank Couplings. To our knowledge, only Forrow et al. (2018) have used low
rank considerations for couplings, rather than costs or kernels. Their work
studies the case where the ground cost is the squared Euclidean distance. They
obtain for that cost a proxy for rank-constrained OT problems using
2-Wasserstein barycenters (Agueh & Carlier, 2011). Their algorithm blends
those in (Cuturi & Doucet, 2014; Benamou et al., 2015) and results in an
intuitive mass transfer plan that goes through a small number of $r$ points,
where $r$ is the coupling’s nonnegative rank.

Our Contributions. In this work, we tackle directly the low-rank problem
formulated by (Forrow et al., 2018) but make no assumption on the cost matrix;
we address instead the low-rank OT problem in its full generality. We consider
couplings $P=Q\operatorname{diag}(1/g)R^{T}$ decomposed as the product of two
sub-couplings $Q,R,$ with common right marginal $g$, and left-marginal given
by those of $P$ on each side. Each of these sub-couplings minimizes a
transport cost that involves the original cost matrix $C$ and the other sub-
coupling. We handle this problem by optimizing jointly on $Q$, $R$ and $g$
using a mirror-descent approach. We prove the non-asymptotic stationary
convergence of this approach. In addition, we show that the time complexity of
our algorithm can become linear when exploiting low rank assumptions on the
cost (not the kernel) involved in the OT problem.

Differences with previous work. Our approach borrows ideas from (Forrow et
al., 2018) but is generic as it applies to all ground costs. Our approach
constrains the non-negative rank of the coupling solution $P$ by construction,
rather than relying on a low rank approximation $\widetilde{K}$ for kernel
$K=e^{-C/\varepsilon}$. This is a crucial point, because the ability to
approximate $K$ with a low rank $\widetilde{K}$ significantly degrades as
$\varepsilon$ decreases. By contrast, our approach applies to all ranks, small
and large. Interestingly, we also show that a low-rank assumption on the cost
matrix (not on the kernel) can also be leveraged, providing therefore a “best
of both worlds” scenario in which both the coupling’s and the cost’s (not the
kernel) low rank properties can be enforced and exploited. Finally, a useful
parallel can be drawn between our approach and that of the vanilla Sinkhorn
algorithm, in the sense that they propose different regularization schemes.
Indeed, the (discrete) path of solutions obtained by our algorithm when
varying $r$ between 1 and $\min(n,m)$ can be seen as an alternative to the
entropic regularization path. Both paths contain at their extremes the
original OT solution (maximal rank and minimal entropy) and the product of
marginals (minimal rank and maximal entropy), as illustrated in Fig. 1.

##  2 Discrete Optimal Transport

OT as a linear program. Let $a$ and $b$ be two histograms in
$\Delta_{n},\Delta_{m}$, the probability simplices of respective size $n,m$.
Assuming $a>0$ and $b>0$, set $X\triangleq(x_{1},\dots,x_{n})$ and
$Y\triangleq(y_{1},\dots,y_{m})$ two families of points taken each within
arbitrary sets, and define discrete distributions
$\mu\triangleq\sum_{i=1}^{n}a_{i}\delta_{x_{i}}$ and
$\nu\triangleq\sum_{j=1}^{m}b_{j}\delta_{y_{j}}$. The set of couplings with
marginals $a,b$ is:

$\displaystyle\Pi_{a,b}\triangleq\\{P\in\mathbb{R}_{+}^{n\times m}\text{ s.t.
}P\mathbf{1}_{m}=a,P^{T}\mathbf{1}_{n}=b\\}\,.$

Given a cost function $c$ defined on pairs of points in $X,Y$ and writing
$C\triangleq[c(x_{i},y_{j})]_{i,j}$ its associated matrix, the optimal
transport (OT) problem can be written as follows:

$\displaystyle\text{OT}(\mu,\nu)\triangleq\min_{P\in\Pi_{a,b}}\langle
C,P\rangle\,.$ (1)

Entropic regularization. Several works have shown recently (Genevay et al.,
2018; Chizat et al., 2020) that when $X$ and $Y$ are sampled from a continuous
space, it is preferable to regularize (1) using, for instance, an entropic
regularizer (Cuturi, 2013) to achieve both better computational and
statistical efficiency,

$\displaystyle\text{OT}_{\varepsilon}(\mu,\nu)\triangleq\min_{P\in\Pi_{a,b}}\langle
C,P\rangle-\varepsilon H(P)\,.$ (2)

where $\varepsilon\geq 0$ and $H$ is the Shannon entropy defined as
$H(P)\triangleq-\sum_{ij}P_{ij}(\log P_{ij}-1)$. If $\varepsilon$ goes to 0,
one recovers the classical OT problem and for any $\varepsilon>0$, Eq. (2)
becomes $\varepsilon$-strongly convex on $\Pi_{a,b}$ and admits a unique
solution $P_{\varepsilon}$, of the form

$\exists
u_{\varepsilon}\in\mathbb{R}^{n}_{+},v_{\varepsilon}\in\mathbb{R}^{m}_{+}\text{
s.t.
}P_{\varepsilon}=\text{diag}(u_{\varepsilon})K\text{diag}(v_{\varepsilon})$
(3)

where $K\triangleq\exp(-C/\varepsilon)$. Cuturi (2013) shows that the scaling
vectors $u_{\varepsilon}$ and $v_{\varepsilon}$ can be obtained efficiently
thanks to the Sinkhorn algorithm (see Alg. 1, where $\odot$ and $/$ denote
entry-wise operation). Each iteration can be performed in $\mathcal{O}(nm)$
algebraic operations as it involves only matrix-vector products. The number of
Sinkhorn iterations needed to converge to a precision $\delta$ (monitored by
the difference between the column-sum of
$\text{diag}(u)\mathbf{K}\text{diag}(v)$ and $b$) is controlled by the scale
of elements in $C$ relative to $\varepsilon$ (Franklin & Lorenz, 1989). That
convergence deteriorates with smaller $\varepsilon$, as studied in more detail
by (Altschuler et al., 2017; Dvurechensky et al., 2018).

Inputs: $K,a,b,\delta,u$
repeat

$v\leftarrow b/K^{T}u,\;u\leftarrow a/Kv$

until  _$\|u\odot Kv-a\|_{1}+\|v\odot K^{T}u-b\|_{1} <\delta$_;

Result: $u,v$

Algorithm 1 $\text{Sinkhorn}(K,a,b,\delta)$

Mirror descent and $\varepsilon$ schedule. A possible interpretation of the
entropic regularization in the OT problem is that it can be seen as the
$k_{\varepsilon}$-th update of a Mirror Descent (MD) algorithm applied to the
objective (1) where $k_{\varepsilon}\geq 1$ depends on $\varepsilon$ and the
gradient steps used in the MD. Several works have proposed such links between
a gradual decrease in $\varepsilon$ to obtain a better approximation of the
unregularized OT problem (Schmitzer, 2019; Lin et al., 2019; Xie et al.,
2020). More precisely, the MD algorithm associated to the Kullback–Leibler
divergence (KL) applied to the objective (1) makes for all $k\geq 0$ the
following update:

$\displaystyle Q^{k+1}\triangleq\operatorname*{argmin}_{Q\in\Pi_{a,b}}\langle
C,Q\rangle+\frac{1}{\gamma_{k}}\text{KL}(Q,Q_{k})$ (4)

where $(\gamma_{k})_{k\geq 0}$ is a sequence of positive real numbers,
$Q_{0}\in\Pi_{a,b}$ is an initial point and KL is the Kullback–Leibler
divergence defined as
$\text{KL}(P,Q)\triangleq\sum_{i,j}P_{i,j}(\log(P_{i,j}/Q_{i,j})-1)$. If
$Q_{0}\triangleq ab^{T}$, then one obtains that for all $k\geq 0$, updating
the coupling according to Eq. (4) is the same as solving

$Q^{k+1}\triangleq\operatorname*{argmin}_{Q\in\Pi_{a,b}}\langle
C,Q\rangle-\varepsilon_{k}H(Q)$

where $\varepsilon_{k}\triangleq(\sum_{j=0}^{k}\gamma_{j})^{-1}$. Therefore
the MD algorithm applied to (1) produces the sequence
$(P_{\varepsilon_{k}})_{k\geq 0}$ of optimal couplings according to the
objective (2). We show next that this viewpoint can be applied when one adds
also some structures to the couplings considered in the OT problem (1),
leading to a new regularized approach.

##  3 Nonnegative Factorization of the Optimal Coupling

Here we aim at regularizing the OT problem by decomposing the couplings
involved into a product of two low-rank couplings. We introduce the associated
non-convex problem and develop a mirror-descent algorithm which operates by
solving a succession of convex programs.

###  3.1 Low Rank and Factored Couplings

We introduce low rank couplings and explain how they can be parameterized as
factored couplings.

######  Definition 1.

Given $M\in\mathbb{R}^{n\times m}$, the nonnegative rank of $M$ is the
smallest number of nonnegative rank-one matrices into which the matrix can be
decomposed additively:

$\textup{rk}_{+}(M)\triangleq\min\\!\left\\{q|M=\sum_{i=1}^{q}R_{i},\forall
i,\textup{rk}(R_{i})=1,R_{i}\geq 0\right\\}.$

Let $r\geq 1$, and let us denote

$\Pi_{a,b}(r)\triangleq\\{P\in\Pi_{a,b},\textup{rk}_{+}(P)\leq r\\}.$

From Definition 1, one has

$\displaystyle\Pi_{a,b}(r)=\begin{aligned}
\Big{\\{}&\sum_{i=1}^{r}g_{i}q_{i}r_{i}^{T}\text{ s.t.
}\forall~{}i~{}q_{i}\in\Delta_{n},~{}r_{i}\in\Delta_{m},\\\
&g\in\Delta_{r},~{}\sum_{i=1}^{r}g_{i}q_{i}=a\text{ and
}\sum_{i=1}^{r}g_{i}r_{i}=b~{}\Big{\\}}\end{aligned}$

from which we deduce directly that $\Pi_{a,b}(r)$ is compact. Moreover for
$g\in\Delta_{r}^{*}\triangleq\\{h\in\Delta_{r}\text{ s.t. }\forall
i~{}h_{i}>0\\}$, we write

$\displaystyle\Pi_{a,g,b}\triangleq\begin{aligned}
\Big{\\{}&P\in\mathbb{R}_{+}^{n\times m},P=Q\operatorname{diag}(1/g)R^{T},\\\
&Q\in\Pi_{a,g},\text{ and }R\in\Pi_{b,g}\Big{\\}}.\end{aligned}$

Note that $\Pi_{a,g,b}$ is compact and a subset of $\Pi_{a,b}(r)$ since for
all $P\in\Pi_{a,g,b}$, $P\in\Pi_{a,b}$ and one has
$\textup{rk}(P)\leq\textup{rk}_{+}(P)\leq r$. Moreover, for any
$P\in\Pi_{a,b}$ such that $\textup{rk}_{+}(P)\leq r$, there exists
$g\in\Delta_{r}^{*}$, $Q\in\Pi_{a,g}$ and $R\in\Pi_{b,g}$ such that
$P=Q\operatorname{diag}(1/g)R^{T}$ (Cohen & Rothblum, 1993). Therefore

$\displaystyle\bigcup\limits_{g\in\Delta_{r}^{*}}\Pi_{a,g,b}=\Pi_{a,b}(r).$
(5)

We exploit next this identity to build an efficient algorithm in order to
solve the optimal transport problem under low nonnegative rank constraints.

###  3.2 The Low-rank OT Problem (LOT)

The problem of interest in this work is:

$\displaystyle\text{LOT}_{r}(\mu,\nu)\triangleq\min_{P\in\Pi_{a,b}(r)}\langle
C,P\rangle.$ (6)

Here the minimum is always attained as $\Pi_{a,b}(r)$ is compact and the
objective is continuous. Thanks to (5), problem (6) is equivalent to

$\displaystyle\min_{(Q,R,g)\in\mathcal{C}(a,b,r)}\langle
C,Q\operatorname{diag}(1/g)R^{T}\rangle$ (7)

where
$\mathcal{C}(a,b,r)\triangleq\mathcal{C}_{1}(a,b,r)\cap\mathcal{C}_{2}(r)$,
with

$\displaystyle\mathcal{C}_{1}(a,b,r)\triangleq\begin{aligned}
\Big{\\{}&(Q,R,g)\in\mathbb{R}_{+}^{n\times r}\times\mathbb{R}_{+}^{m\times
r}\times(\mathbb{R}_{+}^{*})^{r}\\\ &\text{ s.t.
}Q\mathbf{1}_{r}=a,R\mathbf{1}_{r}=b\Big{\\}}\end{aligned}$

and

$\displaystyle\mathcal{C}_{2}(r)\triangleq\begin{aligned}
\Big{\\{}&(Q,R,g)\in\mathbb{R}_{+}^{n\times r}\times\mathbb{R}_{+}^{m\times
r}\times\mathbb{R}^{r}_{+}\\\ &\text{ s.t.
}Q^{T}\mathbf{1}_{n}=R^{T}\mathbf{1}_{m}=g\Big{\\}}.\end{aligned}$

In the following, we also consider regularized version of the problem (7) by
adding an entropic term to the objective which leads for all $\varepsilon\geq
0$ to the following problem

$\begin{aligned}
\text{LOT}_{r,\varepsilon}(\mu,\nu)\triangleq\inf_{(Q,R,g)\in\mathcal{C}(a,b,r)}\langle
C,Q\operatorname{diag}(1/g)R^{T}\rangle\\\ -\varepsilon
H((Q,R,g)).\end{aligned}$ (8)

Here the entropy of $(Q,R,g)$ is to be understood as that of the values of the
three respective entropies evaluated for each term. We will see that adding an
entropic term to the objective allows to stabilize the MD scheme employed to
solve (6). For all $\varepsilon\geq 0$, the objective function defined in (8)
is lower semi-continuous, and admits therefore a minimum in
$\overline{\mathcal{C}_{1}(a,b,r)}\cap\mathcal{C}_{2}(r)$ where
$\overline{\mathcal{C}_{1}(a,b,r)}$ is the closure of
$\mathcal{C}_{1}(a,b,r)$. However, the existence of a solution for problem (8)
requires more care, as shown in the following proposition.

######  Proposition 1.

If $\varepsilon=0$ then the infimum of (8) is always attained. If
$\varepsilon>0$, then if $r=1$, the infimum of (8) is attained and for $r\geq
2$, problem (8) admits a minimum if
$\text{LOT}_{r,\varepsilon}(\mu,\nu)<\text{LOT}_{r-1,\varepsilon}(\mu,\nu)$.

#### Stabilized Formulation using Lower Bounds

In order to ensure stability of the mirror descent method, and enable its
theoretical analysis, we introduce a lower bound $\alpha$ on the weight vector
$g$.

Let us assume in the following that we consider $(r,\varepsilon)$ satisfying
the conditions of Proposition 1. In particular if $\varepsilon=0$, $r$ can be
arbitrarily chosen and we recover the problem defined in (6). Under this
assumption, there exists
$(Q^{*}_{\varepsilon},R^{*}_{\varepsilon},g_{\varepsilon}^{*})\in\mathcal{C}_{1}(a,b,r)\cap\mathcal{C}_{2}(r)$
solution of Eq. (8) from which follows the existence of
$\frac{1}{r}\geq\alpha^{*}>0$, such that $g_{\varepsilon}^{*}\geq\alpha^{*}$
coordinate-wise. Let us now define for any $\frac{1}{r}\geq\alpha>0$, the
following set

$\displaystyle\mathcal{C}_{1}(a,b,r,\alpha)\triangleq\begin{aligned}
\Big{\\{}&(Q,R,g)\in\mathbb{R}_{+}^{n\times r}\times\mathbb{R}_{+}^{m\times
r}\times\mathbb{R}_{+}^{r}\\\ &\text{ s.t.
}Q\mathbf{1}_{r}=a,R\mathbf{1}_{r}=b,~{}g\geq\alpha\Big{\\}}.\end{aligned}$

Then if $\alpha$ is sufficiently small (i.e. $\alpha\leq\alpha^{*}$) we have
that the problem (8) is equivalent to

$\text{LOT}_{r,\varepsilon,\alpha}(\mu,\nu)=\min_{(Q,R,g)\in\mathcal{C}(a,b,r,\alpha)}\langle
C,Q\operatorname{diag}(1/g)R^{T}\rangle\\\ -\varepsilon H((Q,R,g)),$ (9)

where
$\mathcal{C}(a,b,r,\alpha)\triangleq\mathcal{C}_{1}(a,b,r,\alpha)\cap\mathcal{C}_{2}(r)$.
Note that for any $\frac{1}{r}\geq\alpha>0$, the set of constraints is not
empty, compact and the minimum always exists.

###  3.3 Mirror Descent Optimization Scheme

#### Mirror descent outer loop.

We propose to use a Mirror Descent scheme with a KL divergence to solve Eq.
(9). It leads, for all $k\geq 0$, to the following updates which necessitate
the solution of a convex problem at each step

$(Q_{k+1},R_{k+1},g_{k+1})\triangleq\\!\\!\operatorname*{argmin}_{\bm{\zeta}\in\mathcal{C}(a,b,r,\alpha)}\\!\\!\text{KL}(\bm{\zeta},\bm{\xi}_{k})$
(10)

where $(Q_{0},R_{0},g_{0})\in\mathcal{C}(a,b,r,\alpha)$ is an initial point
such that $Q_{0}>0$ and $R_{0}>0$,
$\bm{\xi}_{k}\triangleq(\xi_{k}^{(1)},\xi_{k}^{(2)},\xi_{k}^{(3)})$,
$\xi_{k}^{(1)}\triangleq\exp(-\gamma_{k}CR_{k}\operatorname{diag}(1/g_{k})-(\gamma_{k}\varepsilon-1)\log(Q_{k}))$,
$\xi_{k}^{(2)}\triangleq\exp(-\gamma_{k}C^{T}Q_{k}\operatorname{diag}(1/g_{k})-(\gamma_{k}\varepsilon-1)\log(R_{k}))$,
$\xi_{k}^{(3)}\triangleq\exp(\gamma_{k}\omega_{k}/g_{k}^{2}-(\gamma_{k}\varepsilon-1)\log(g_{k}))$
with $[\omega_{k}]_{i}\triangleq[Q_{k}^{T}CR_{k}]_{i,i}$ for all
$i\in\\{1,\dots,r\\}$ and $(\gamma_{k})_{k\geq 0}$ is a sequence of positive
step sizes. Note that for all $k\geq 0$, $(Q_{k},R_{k},g_{k})$ live in
$(\mathbb{R}_{+}^{*})^{n\times r}\times(\mathbb{R}_{+}^{*})^{m\times
r}\times(\mathbb{R}_{+}^{*})^{r}$, and therefore $\bm{\xi}_{k}$ is well
defined and lives also in $(\mathbb{R}_{+}^{*})^{n\times
r}\times(\mathbb{R}_{+}^{*})^{m\times r}\times(\mathbb{R}_{+}^{*})^{r}$.

#### Dykstra’s inner loop.

In order to solve Eq. (10), we use the Dykstra’s Algorithm (Dykstra, 1983).
Given a closed convex set $\mathcal{C}\subset\mathbb{R}_{+}^{n\times
r}\times\mathbb{R}_{+}^{m\times r}\times\mathbb{R}_{+}^{r}$, we denote for all
$\bm{\xi}\in(\mathbb{R}_{+}^{*})^{n\times
r}\times(\mathbb{R}_{+}^{*})^{m\times r}\times(\mathbb{R}_{+}^{*})^{r}$ the
projection according to the Kullback-Leibler divergence as

$\displaystyle\mathcal{P}_{\mathcal{C}}^{\text{KL}}(\bm{\xi})\triangleq\operatorname*{argmin}_{\bm{\zeta}\in\mathcal{C}}\text{KL}(\bm{\zeta},\bm{\xi}).$

Starting from $\bm{\zeta}_{0}\triangleq\bm{\xi}$ and
$\bm{q}_{0}=\bm{q}_{-1}=(\mathbf{1},\mathbf{1},\mathbf{1})\in\mathbb{R}_{+}^{n\times
r}\times\mathbb{R}_{+}^{m\times r}\times\mathbb{R}_{+}^{r}$, the Dykstra’s
Algorithm consists in computing for all $j\geq 0$,

$\displaystyle\bm{\zeta}_{2j+1}$
$\displaystyle=\mathcal{P}_{\mathcal{C}_{1}(a,b,r,\alpha)}^{\text{KL}}(\bm{\zeta}_{2j}\odot\bm{q}_{2j-1})$
$\displaystyle\bm{q}_{2j+1}$
$\displaystyle=\bm{q}_{2j-1}\odot\frac{\bm{\zeta}_{2j}}{\bm{\zeta}_{2j+1}}$
$\displaystyle\bm{\zeta}_{2j+2}$
$\displaystyle=\mathcal{P}_{\mathcal{C}_{2}(r)}^{\text{KL}}(\bm{\zeta}_{2j+1}\odot\bm{q}_{2j})$
$\displaystyle\bm{q}_{2j+2}$
$\displaystyle=\bm{q}_{2j}\odot\frac{\bm{\zeta}_{2j+1}}{\bm{\zeta}_{2j+2}}.$

As $\mathcal{C}_{1}(a,b,r,\alpha)$ and $\mathcal{C}_{2}(r)$ are closed convex
subspaces and $\bm{\xi}\in(\mathbb{R}_{+}^{*})^{n\times
r}\times(\mathbb{R}_{+}^{*})^{m\times r}\times(\mathbb{R}_{+}^{*})^{r}$, one
can show that $(\bm{\zeta}_{j})_{j\geq 0}$ converges towards the unique
solution of Eq. (10), (Bauschke & Lewis, 2000). The following propositions
detail how to compute the relevant projections involved in the Dykstra’s
Algorithm.

######  Proposition 2.

For
$\tilde{\bm{\xi}}\triangleq(\tilde{Q},\tilde{R},\tilde{g})\in(\mathbb{R}_{+}^{*})^{n\times
r}\times(\mathbb{R}_{+}^{*})^{n\times r}\times(\mathbb{R}_{+}^{*})^{r}$, one
has, denoting $\hat{g}\triangleq\max(\tilde{g},\alpha)$

$\displaystyle\mathcal{P}_{\mathcal{C}_{1}(a,b,r,\alpha)}^{\text{KL}}(\tilde{\bm{\xi}})=\left(\operatorname{diag}\left(\frac{a}{\tilde{Q}\mathbf{1}_{r}}\right)\tilde{Q},\operatorname{diag}\left(\frac{b}{\tilde{R}\mathbf{1}_{r}}\right)\tilde{R},\hat{g}\right).$

Let us now show the solution of the projection on $\mathcal{C}_{2}(r)$.

######  Proposition 3.

For
$\tilde{\bm{\xi}}\triangleq(\tilde{Q},\tilde{R},\tilde{g})\in(\mathbb{R}_{+}^{*})^{n\times
r}\times(\mathbb{R}_{+}^{*})^{n\times r}\times(\mathbb{R}_{+}^{*})^{r}$, the
projection
$(Q,R,g)=\mathcal{P}_{\mathcal{C}_{2}(r)}^{\text{KL}}(\tilde{\bm{\xi}})$
satisfies

$\displaystyle Q$
$\displaystyle=\tilde{Q}\operatorname{diag}(g/\tilde{Q}^{T}\mathbf{1}_{n}),\quad
R=\tilde{R}\operatorname{diag}(g/\tilde{R}^{T}\mathbf{1}_{m})$ $\displaystyle
g$
$\displaystyle=(\tilde{g}\odot\tilde{Q}^{T}\mathbf{1}_{n}\odot\tilde{R}^{T}\mathbf{1}_{m})^{1/3}.$

#### Efficient computation of the updates.

The projection obtained in Proposition 2, 3 lead to simple updates of the
couplings. Indeed, starting with
$\bm{\zeta}_{0}\triangleq\bm{\xi}=(\xi^{(1)},\xi^{(2)},\xi^{(3)})$ the
Dysktra’s Algorithm applied to our problem (10) needs only to compute scaling
vectors as presented in Alg. 2. We have denoted $p_{1}\triangleq a$ and
$p_{2}\triangleq b$ to simplify the notations. See Appendix B for more
details.

Inputs:
$\xi^{(1)},\xi^{(2)},\tilde{g}\triangleq\xi^{(3)},p_{1},p_{2},\alpha,\delta,q^{(3)}_{1}=q^{(3)}_{2}=\mathbf{1}_{r},\forall
i\in\\{1,2\\},~{}\tilde{v}^{(i)}=\mathbf{1}_{r},q^{(i)}=\mathbf{1}_{r}$
repeat

$u^{(i)}\leftarrow p_{i}/\xi^{(i)}\tilde{v}^{(i)}~{}\forall i\in\\{1,2\\},\\\
g\leftarrow\max(\alpha,\tilde{g}\odot
q^{(3)}_{1}),~{}q^{(3)}_{1}\leftarrow(\tilde{g}\odot
q^{(3)}_{1})/g,~{}\tilde{g}\leftarrow g,\\\ g\leftarrow(\tilde{g}\odot
q^{(3)}_{2})^{1/3}\prod_{i=1}^{2}(v^{(i)}\odot
q^{(i)}\odot(\xi^{(i)})^{T}u^{(i)})^{1/3},\\\ v^{(i)}\leftarrow
g/(\xi^{(i)})^{T}u^{(i)}~{}\forall i\in\\{1,2\\},\\\
q^{(i)}\leftarrow(\tilde{v}^{(i)}\odot q^{(i)})/v^{(i)}~{}\forall
i\in\\{1,2\\},~{}q^{(3)}_{2}\leftarrow(\tilde{g}\odot q^{(3)}_{2})/g,\\\
\tilde{v}^{(i)}\leftarrow v^{(i)}~{}\forall
i\in\\{1,2\\},~{}\tilde{g}\leftarrow g$

until  _$\sum_{i=1}^{2}\|u^{(i)}\odot\xi^{(i)}v^{(i)}-p_{i}\|_{1} <\delta$_;

$Q\leftarrow\operatorname{diag}(u^{(1)})\xi_{k}^{(1)}\operatorname{diag}(v^{(1)})$
$R\leftarrow\operatorname{diag}(u^{(2)})\xi_{k}^{(2)}\operatorname{diag}(v^{(2)})$
Result: $Q,R,g$

Algorithm 2 $\text{LR-Dykstra}((\xi^{(i)})_{1\leq i\leq
3},p_{1},p_{2},\alpha,\delta)$

Let us now introduce the proposed MD algorithm applied to (9). By denoting
$\mathcal{D}(\cdot)$ the operator extracting the diagonal of a square matrix
we obtain Alg. 3.

Inputs: $C,a,b,(\gamma_{k})_{k\geq 0},Q,R,g,\alpha,\delta$
for  _$k=1,\dots$_ do

$\xi^{(1)}\leftarrow\exp(-\gamma_{k}CR\operatorname{diag}(1/g)-(\gamma_{k}\varepsilon-1)\log(Q)),\\\
\xi^{(2)}\leftarrow\exp(-\gamma_{k}C^{T}Q\operatorname{diag}(1/g)-(\gamma_{k}\varepsilon-1)\log(R)),\\\
\omega\leftarrow\mathcal{D}(Q^{T}CR),\\\
\xi^{(3)}\leftarrow\exp(\gamma_{k}\omega/g^{2}-(\gamma_{k}\varepsilon-1)\log(g)),\\\
Q,R,g\leftarrow\text{LR-Dykstra}((\xi^{(i)})_{1\leq i\leq
3},a,b,\alpha,\delta)~{}(\text{Alg.~{}\ref{alg-Dykstra})}$

end for

Result: $\langle C,Q\operatorname{diag}(1/g)R^{T}\rangle$

Algorithm 3 $\text{LOT}(C,a,b,r,\alpha,\delta)$

#### Computational Cost.

Note that $(\xi^{(i)})_{1\leq i\leq 3}$ considered in Alg. 3 live in
$\mathbb{R}_{+}^{n\times r}\times\mathbb{R}_{+}^{m\times
r}\times\mathbb{R}_{+}^{r}$ and therefore given those matrices, each iteration
of Alg. 2 requires $\mathcal{O}((n+m)r)$ algebraic operations, since it
involves only matrix/vector multiplications of the form $\xi^{(i)}v_{i}$ and
$(\xi^{(i)})^{T}u_{i}$. However without any assumption on the cost matrix $C$,
computing $(\xi^{(i)})_{1\leq i\leq 3}$ requires $\mathcal{O}(nmr)$ algebraic
operations since $CR$ and $C^{T}Q$ must be evaluated. We show in §3.5 how to
reduce the quadratic cost of computing $(\xi^{(i)})_{1\leq i\leq 3}$ to a
linear cost with respect to the number of samples if one assumes that the
considered cost matrix can be factored, either exactly (ensured with a squared
Euclidean distance cost) or approximately if that cost is a distance. Writing
$N$ the number of iterations of the MD scheme and $T$ the number of iterations
considered in Algorithm 2 at each step of the MD, we end up with a total
computational cost of $\mathcal{O}(NT(n+m)r+Nnmr)$.

###  3.4 Convergence of the Mirror Descent

Even if the objective (9) is not convex in $(Q,R,g)$, we obtain the non-
asymptotic stationary convergence of the MD algorithm in this setting. For
that purpose we introduce a stronger convergence criterion than the one
presented in (Ghadimi et al., 2013) to obtain non-asymptotic stationary
convergence of the MD scheme. Indeed let $F_{\varepsilon}$ be the objective
function of the problem (9) defined on $\mathcal{C}(a,b,r,\alpha)$ and let us
denotes for any $\gamma>0$ and $\bm{\xi}\in\mathcal{C}(a,b,r,\alpha)$

$\displaystyle\mathcal{G}_{\varepsilon,\alpha}(\bm{\xi},\gamma)\triangleq\operatorname*{argmin}_{\bm{\zeta}\in\mathcal{C}(a,b,r,\alpha)}\\{\langle\nabla
F_{\varepsilon}(\bm{\xi}),\bm{\zeta}\rangle+\frac{1}{\gamma}KL(\bm{\zeta},\bm{\xi})\\}.$

Then the criteron used in (Ghadimi et al., 2013) to show the stationary
convergence of the MD scheme is defined as the square norm of the following
vector:

$\displaystyle
P_{\mathcal{C}(a,b,r,\alpha)}(\bm{\xi},\gamma)\triangleq\frac{1}{\gamma}(\bm{\xi}-\mathcal{G}_{\varepsilon,\alpha}(\bm{\xi},\gamma)).$

This vector can be seen as a generalized projected gradient of
$F_{\varepsilon}$ at $\bm{\xi}$. Indeed if $X=\mathbb{R}^{d}$ and by replacing
the _prox-function_ $\text{KL}(u,x)$ by $\frac{1}{2}\|u-x\|_{2}^{2}$, we would
have $P_{X}(x,\gamma)=\nabla F_{\varepsilon}(x)$. Here we consider instead the
following criterion to establish convergence:

$\displaystyle\Delta_{\varepsilon,\alpha}(\bm{\xi},\gamma)\triangleq\frac{1}{\gamma^{2}}(\mathrm{KL}(\bm{\xi},\mathcal{G}_{\varepsilon,\alpha}(\bm{\xi},\gamma))+\mathrm{KL}(\mathcal{G}_{\varepsilon,\alpha}(\bm{\xi},\gamma),\bm{\xi})).$

Such criterion is in fact stronger than the one used in (Ghadimi et al., 2013)
as we have

$\displaystyle\Delta_{\varepsilon,\alpha}(\bm{\xi},\gamma)$
$\displaystyle=\frac{1}{\gamma^{2}}(\langle\nabla
h(\mathcal{G}_{\varepsilon,\alpha}(\bm{\xi},\gamma))-\nabla
h(\bm{\xi}),\mathcal{G}_{\varepsilon,\alpha}(\bm{\xi},\gamma)-\bm{\xi}\rangle$
$\displaystyle\geq\frac{1}{2\gamma^{2}}\|\mathcal{G}_{\varepsilon,\alpha}(\bm{\xi},\gamma)-\bm{\xi}\|_{1}^{2}$
$\displaystyle=\frac{1}{2}\|P_{\mathcal{C}(a,b,r,\alpha)}(\bm{\xi},\gamma)\|_{1}^{2}$

where $h$ denotes the minus entropy function and the last inequality comes
from the strong convexity of $h$ on $\mathcal{C}(a,b,r,\alpha)$.

For any $\frac{1}{r}\geq\alpha>0$, we show in the following proposition the
non-asymptotic stationary convergence of the MD scheme applied to the problem
(9). To prove this result, we show that for any $\varepsilon\geq 0$, the
objective is smooth relatively to the negative entropy function (Bauschke et
al., 2017) and we extend the proof of (Ghadimi et al., 2013) to this case.

######  Proposition 4.

Let $\varepsilon\geq 0$, $\frac{1}{r}\geq\alpha>0$ and $N\geq 1$. By denoting

$L_{\varepsilon,\alpha}\triangleq\sqrt{3\left(2\frac{\|C\|_{2}^{2}}{\alpha^{4}}+\left(\frac{\varepsilon+2\|C\|_{2}}{\alpha^{3}}\right)^{2}\right)}$

and by considering a constant stepsize in the MD scheme (10) such that for all
$k=1,\dots,N$ $\gamma_{k}=\frac{1}{2L_{\varepsilon,\alpha}}$, we obtain that

$\displaystyle\min_{1\leq k\leq
N}\Delta_{\varepsilon,\alpha}((Q_{k},R_{k},g_{k}),\gamma_{k})\leq\frac{4L_{\varepsilon,\alpha}D_{0}}{N}.$

where $D_{0}\triangleq
F_{\varepsilon}(Q_{0},R_{0},g_{0})-\mathrm{LOT}_{r,\varepsilon,\alpha}$ is the
distance of the initial value to the optimal one.

Thanks to Proposition 4, for $\alpha$ sufficiently small (i.e.
$\alpha\leq\alpha^{*}$), we have
$\text{LOT}_{r,\varepsilon,\alpha}=\text{LOT}_{r,\varepsilon}$ and therefore
we obtain a stationary point of (8). In particular, if $\varepsilon=0$, the
proposed algorithm converges towards a stationary point of (6).

######  Remark 1.

We also propose an algorithm to directly solve (8). The main difference is
that the updates of the MD can be solved using the Iterative Bregman
Projections (IBP) Algorithm. See Appendix F for more details.

######  Remark 2.

For all $\varepsilon\geq 0$, the MD scheme implies that each iteration $k$ of
our proposed algorithm outputs
$(Q_{k},R_{k},g_{k})\in\mathcal{C}_{1}(a,b,r,\alpha)\cap\mathcal{C}_{2}(r)$,
and therefore the matrix obtained a each iteration
$P_{k}^{\textit{LOT}}=Q_{k}\operatorname{diag}(1/g_{k})R_{k}^{T}$ is a
coupling which sastifies the marginal constraints while in the Sinkhorn
algorithm, the matrix defined at each iteration by
$P_{k}^{\text{Sin}}=\operatorname{diag}(u_{k})K\operatorname{diag}(v_{k})$
becomes a coupling which satisfies the marginal constraints only at
convergence.

In the following section, we aim at accelerating our method in order to obtain
a linear time algorithm to solve (8).

Figure 2: In this experiment, we consider two Gaussian distributions evaluated
on $n=m=5000$ in 2D. The first one has a mean of $(1,1)^{T}$ and identity
covariance matrix $I_{2}$ while the other has 0 mean and covariance $0.1\times
I_{2}$. The ground cost is the squared Euclidean distance. Note that for this
cost, an exact low-rank factorization of the cost is available, and therefore
all low-rank methods, including ours, have a linear time complexity. _Left:_
we show that when $\varepsilon=0$ our method is able to quickly obtain the
exact OT by forcing the nonnegative rank of the coupling to be relatively
small compared to the number of samples. Note that in this setting, all the
other methods cannot be applied. _Middle left, middle right:_ In these plots,
we show that our method can obtain high accuracy for either estimate the true
OT or its regularized version with order of magnitude faster than the other
low-rank methods for any rank $r$. Moreover, our methods outperforms Sin in
these regimes of small regularizations. Note that Sin does not converge for
$\varepsilon=0.002$ as we do not consider its stabilized version using log-
sum-exp function but rather its classical version which is less costly to
compute. _Right:_ Here we change the scale of the $y$-axis of the plot. We see
that the regime of the entropic regularizations for the Sinkhorn algorithm and
our method differs. Indeed, the Sinkhorn algorithm has a larger range of
$\varepsilon$ such that it provides an efficient approximation of the OT,
whereas LOT is regularizing twice, namely with respect to both rank and
entropy.

###  3.5 Linear time approximation of the Low-Rank Optimal Transport

Here we aim at obtaining the optimal solution of Eq. (8) in linear time with
respect to the number of samples. For that purpose let us introduce our main
assumption on the cost matrix $C$.

######  Assumption 1.

Assume that $C$ admits a low-rank factorization, that is there exists
$A\in\mathbb{R}^{n\times d}$ and $B\in\mathbb{R}^{m\times d}$ such that
$C=AB^{T}.$

From the Assumption 1 one can in fact accelerate the computation in the
iterations of the proposed Alg. (3) and obtain a linear time algorithm with
respect to the number of samples. Indeed recall that given
$\bm{\xi}=(\xi^{(i)})_{1\leq i\leq 3}$, each iteration of the Dykstra’s Alg.
(2) can be performed in linear time. Moreover, thanks to Assumption 1, the
computation of $\bm{\xi}$, which requires to compute both $CR$ and $C^{T}Q$
can be performed in $\mathcal{O}((n+m)dr)$ algebraic operations and thus Alg.
(3) requires only a linear number of algebraic operations with respect to the
number of samples at each iteration.

Let us now justify why the Assumption 1 of a low-rank factorization for the
cost matrix is well suited in the problem of computing the Optimal Transport.

#### Squared Euclidean Metric.

In the specific case where $C$ is a Square Euclidean distance matrix, it
admits a low-rank decomposition. Indeed let
$X\triangleq[x_{1},\dots,x_{n}]\in\mathbb{R}^{d\times n}$, let
$Y\triangleq[y_{1},\dots,y_{m}]\in\mathbb{R}^{d\times m}$ and let
$D\triangleq(\|x_{i}-y_{j}\|_{2}^{2})_{i,j}$. Then by denoting
$p=[\|x_{1}\|_{2}^{2},\dots,\|x_{n}\|_{2}^{2}]^{T}\in\mathbb{R}^{n}$ and
$q=[\|y_{1}\|_{2}^{2},\dots,\|y_{m}\|_{2}^{2}]^{T}\in\mathbb{R}^{m}$ we can
rewrite $D$ as the following:

$\displaystyle D=p\mathbf{1}_{m}^{T}+\mathbf{1}_{n}q^{T}-2X^{T}Y.$

Therefore by denoting
$A=[p,\mathbf{1}_{n},-2X^{T}]\in\mathbb{R}^{n\times(d+2)}$ and
$B=[\mathbf{1}_{m},q,Y^{T}]\in\mathbb{R}^{n\times(d+2)}$ we obtain that

$\displaystyle D=AB^{T}.$

#### General Case: Distance Matrix.

In the following we denote a distance matrix $D\in\mathbb{R}^{n\times m}$, any
matrix such that there exists a metric space $(\mathcal{X},d)$,
$\\{x_{i}\\}_{i=1}^{n}\in\mathcal{X}^{n}$ and
$\\{y_{j}\\}_{j=1}^{m}\in\mathcal{X}^{m}$ which satisfy for all $i,j$,
$D_{i,j}=d(x_{i},y_{j})$. In fact it is always possible to obtain a low-rank
approximation of a distance matrix in linear time. In (Bakshi & Woodruff,
2018; Indyk et al., 2019), the authors proposed an algorithm such that for any
distance matrix $D\in\mathbb{R}^{n\times m}$ and $\gamma>0$ it outputs
matrices $M\in\mathbb{R}^{n\times d}$, $N\in\mathbb{R}^{m\times d}$ in
$\mathcal{O}((m+n)\text{poly}(\frac{d}{\gamma}))$ algebraic operations such
that with probability at least $0.99$ we have

$\displaystyle\|D-MN^{T}\|_{F}^{2}\leq\|D-D_{d}\|_{F}^{2}+\gamma\|D\|_{F}^{2}$

where $D_{d}$ denotes the best rank-$d$ approximation to $D$. Therefore one
can always obtain a low-rank factorization of a distance matrix in linear time
with respect to the number of samples. See Appendix D for more details.

##  4 Numerical Results

We consider three problems in which we study the time-accuracy trade-off as
well as the couplings obtained, by comparing our method with other low-rank
methods, as well as Sinkhorn’s algorithm. More precisely, we compare our
proposed method, LOT, with the factored Optimal Transport (Forrow et al.,
2018), FactoredOT, the Nystrom-based method (Altschuler et al., 2018), Nys,
the random features-based method (Scetbon & Cuturi, 2020), RF and the Sinkhorn
algorithm (Cuturi, 2013), Sin. For LOT, and in all experiments, we set the
lower bound on $g$ to $\alpha=10^{-5}$.

#### Time-accuracy Tradeoff

We consider two problems where the ground cost involved in the OT problem is
either the _squared Euclidean_ distance or the _Euclidean_ distance. In the
first one, we consider measures supported on $n=5000$ points in
$\mathbb{R}^{2}$, while the second we consider $n=10000$ samples in
$\mathbb{R}^{2}$. The method proposed by (Forrow et al., 2018) can only be
used with the squared Euclidean distance (2-Wasserstein) while ours works for
any cost. For all the low-ranks methods, we vary the ranks between 10 and 500.
For all the randomized methods, we consider the mean over 10 runs to estimate
the OT.

In Fig. 2, 3 we plot the ratio w.r.t. the (non-regularized) optimal transport
cost defined as $\text{R}:=\langle C,\widetilde{P}\rangle/\langle
C,P^{*}\rangle$ where $\widetilde{P}$ is the coupling obtained by the method
considered and $P^{*}$ is the ground truth (we ensure this optimal cost is
large enough to avoid spurious divisions by 0). We present the time-accuracy
tradeoffs of the methods for different regularizations $\varepsilon$ and ranks
$r$. We show that our method provides consistently a better approximation of
the OT while being much faster than the other low-rank methods for various
targeted rank values $r$. We also show that our method is able to approximate
arbitrarily well the OT and so faster than the Sinkhorn algorithm thanks to
the low-rank constraints. We compare the methods in the same setting but we
increase the dimensionality of the problems considered and we observe similar
results. See Appendix G for more details.

Figure 3: Here we consider two Gaussian mixture densities sampled with
$n=m=10000$ points in 2D (See Appendix G for more details). The ground cost is
the Euclidean distance. As this cost is a distance, we can apply our linear
version of the algorithm and we denote LOT Quad to refer to its quadratic
counterpart. We see that LOT and LOT Quad provide similar results while LOT is
faster. All kernel-based methods (Nys, RF) fail to converge in this setting.
As in Fig. 2, we see that our method is able to approximate faster than Sin
the true OT thanks to the low-rank constraint.

Figure 4: We illustrate in this plot the gaps between the OT objective (top)
and the couplings (bottom) reached by Sin and LOT for varying regularization
strengths. Measures were sampled on a complete graph obtained by sampling
$2n=2000$ points from a 2-D standard normal distribution, the edge weights set
to their squared Euclidean distances. The supports are obtained by randomly
splitting the nodes of the graphs into two subsets of same size. We vary the
entropic regularization $\varepsilon$ and the nonnegative rank $r$. We
consider $\varepsilon$ in log-scale ranging from $0.001$ to $1$ and $r$
ranging from 1 to 1000, represented as a fraction of $n$. The blue (resp. red)
curve stands for Sin (resp. LOT). We plot the absolute difference between the
OT obtained (top) and $\ell_{1}$ distance between their respective couplings
(bottom).

######  Remark 3.

Adding an entropic regularization in our objective allows to stabilize the MD
scheme and therefore obtain faster convergence. Indeed if $\varepsilon>0$,
then the number of iterations required to solve each iteration of the MD
scheme (10) by Algorithm (2) is monitored by $\varepsilon$ given a certain
precision $\delta$ while in the case where $\varepsilon=0$, the number of
iterations required for Algorithm 2 to reach the precision $\delta$ increases
as the number of iterations in the MD scheme increases.

#### Comparison of the Couplings

Seeking to take a deeper look at the phenomenon highlighted in Fig. 1, we
study differences in the regularization paths of LOT and Sin. We consider
distributions supported on graphs of $n=1000$ nodes, endowed with the shortest
path distance (Bondy et al., 1976). We consider LOT with no entropic
regularization (i.e. $\varepsilon=0$ in Eq. (9)) against Sin for various pairs
of regularizers. Results are displayed in Fig. 4, where the discrete path of
regularizations parameterized by the rank $r$ of LOT is compared with that
obtained by Sin when varying $\varepsilon$. The gaps in ratio R and couplings
(in $\ell_{1}$) between the two methods are displayed. Both methods are able
to approximate arbitrarily well the OT but offer two different paths to
interpolate from the independent coupling $ab^{T}$ of rank 1 to the optimal
one. More precisely, we see that the range of $\varepsilon$ for which the
entropic OT provides an efficient approximation of the true coupling is very
localized, while the rank $r$ needed for LOT to obtain such approximation is
wider. Moreover, we see that the decay of the ratio of LOT with respect to $r$
is faster than the decay of Sin w.r.t. $\varepsilon$.

######  Remark 4.

A comparative advantage of using the low-rank parameterization of OT over the
Sinkhorn approach lies in the simple bounds that $r$ admits, between $1$ and
$n$, and the fact that $r$ encodes directly, through an integer, a direct
property of the resulting coupling. In that sense, the same value $r$ can be
used across experiments that compare measures of various sizes and supports.
By contrast, selecting a suitable regularization strength $\varepsilon$ in the
Sinkhorn algorithm is usually challenging, as the parameter is continuous and
its magnitude depends directly on the cost matrix values, making a common
choice across experiments difficult.

#### Conclusion

We proposed a new approach to regularize the OT problem by restricting
solutions to have a small non-negative rank. Our algorithm leverages both low-
rank constraints and entropic smoothing. Our method can leverage the
factorization of the ground cost (and not that of the kernel usually
associated to Sinkhorn) to propose a linear time complexity alternative to
solve OT problems.

#### Acknowledgements

The work of G. Peyré was supported by the European Research Council (ERC
project NORIA) and by the French government under management of ANR as part of
the “Investissements d’avenir” program (ANR19-P3IA-0001, PRAIRIE 3IA
Institute).

## References

  * Agueh & Carlier (2011) Agueh, M. and Carlier, G.  Barycenters in the wasserstein space.  _SIAM Journal on Mathematical Analysis_ , 43(2):904–924, 2011.
  * Altschuler et al. (2017) Altschuler, J., Weed, J., and Rigollet, P.  Near-linear time approximation algorithms for optimal transport via sinkhorn iteration.  _arXiv preprint arXiv:1705.09634_ , 2017.
  * Altschuler et al. (2018) Altschuler, J., Bach, F., Rudi, A., and Niles-Weed, J.  Massively scalable sinkhorn distances via the nyström method, 2018.
  * Altschuler & Boix-Adsera (2020) Altschuler, J. M. and Boix-Adsera, E.  Polynomial-time algorithms for multimarginal optimal transport problems with structure, 2020.
  * Bakshi & Woodruff (2018) Bakshi, A. and Woodruff, D. P.  Sublinear time low-rank approximation of distance matrices, 2018.
  * Bauschke & Lewis (2000) Bauschke, H. H. and Lewis, A. S.  Dykstras algorithm with bregman projections: A convergence proof.  _Optimization_ , 48(4):409–427, 2000.
  * Bauschke et al. (2017) Bauschke, H. H., Bolte, J., and Teboulle, M.  A descent lemma beyond lipschitz gradient continuity: first-order methods revisited and applications.  _Mathematics of Operations Research_ , 42(2):330–348, 2017.
  * Benamou et al. (2015) Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, L., and Peyré, G.  Iterative Bregman projections for regularized transportation problems.  _SIAM Journal on Scientific Computing_ , 37(2):A1111–A1138, 2015.
  * Bondy et al. (1976) Bondy, J. A., Murty, U. S. R., et al.  _Graph theory with applications_ , volume 290.  Macmillan London, 1976.
  * Bregman (1967) Bregman, L. M.  The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming.  _USSR computational mathematics and mathematical physics_ , 7(3):200–217, 1967.
  * Chizat et al. (2020) Chizat, L., Roussillon, P., Léger, F., Vialard, F.-X., and Peyré, G.  Faster wasserstein distance estimation with the sinkhorn divergence.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Clason et al. (2021) Clason, C., Lorenz, D. A., Mahler, H., and Wirth, B.  Entropic regularization of continuous optimal transport problems.  _Journal of Mathematical Analysis and Applications_ , 494(1):124432, 2021.  ISSN 0022-247X.  doi: https://doi.org/10.1016/j.jmaa.2020.124432.
  * Cohen & Rothblum (1993) Cohen, J. E. and Rothblum, U. G.  Nonnegative ranks, decompositions, and factorizations of nonnegative matrices.  _Linear Algebra and its Applications_ , 190:149 – 168, 1993\.  ISSN 0024-3795.  doi: https://doi.org/10.1016/0024-3795(93)90224-C.  URL http://www.sciencedirect.com/science/article/pii/002437959390224C.
  * Cuturi (2013) Cuturi, M.  Sinkhorn distances: Lightspeed computation of optimal transport.  In _Advances in neural information processing systems_ , pp. 2292–2300, 2013.
  * Cuturi & Doucet (2014) Cuturi, M. and Doucet, A.  Fast computation of Wasserstein barycenters.  In _Proceedings of ICML_ , volume 32, pp. 685–693, 2014.
  * Dessein et al. (2018) Dessein, A., Papadakis, N., and Rouas, J.-L.  Regularized optimal transport and the rot mover’s distance.  _The Journal of Machine Learning Research_ , 19(1):590–642, 2018.
  * Dvurechensky et al. (2018) Dvurechensky, P., Gasnikov, A., and Kroshnin, A.  Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn’s algorithm.  In Dy, J. and Krause, A. (eds.), _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 1367–1376. PMLR, 10–15 Jul 2018.
  * Dykstra (1983) Dykstra, R. L.  An algorithm for restricted least squares regression.  _Journal of the American Statistical Association_ , 78(384):837–842, 1983.
  * Forrow et al. (2018) Forrow, A., Hütter, J.-C., Nitzan, M., Rigollet, P., Schiebinger, G., and Weed, J.  Statistical optimal transport via factored couplings, 2018.
  * Fournier & Guillin (2015) Fournier, N. and Guillin, A.  On the rate of convergence in Wasserstein distance of the empirical measure.  _Probability Theory and Related Fields_ , 162(3-4):707–738, 2015.
  * Franklin & Lorenz (1989) Franklin, J. and Lorenz, J.  On the scaling of multidimensional matrices.  _Linear Algebra and its Applications_ , 114:717–735, 1989\.
  * Genevay et al. (2018) Genevay, A., Chizat, L., Bach, F., Cuturi, M., and Peyré, G.  Sample complexity of sinkhorn divergences.  _arXiv preprint arXiv:1810.02733_ , 2018.
  * Ghadimi et al. (2013) Ghadimi, S., Lan, G., and Zhang, H.  Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization, 2013.
  * Heitz et al. (2020) Heitz, M., Bonneel, N., Coeurjolly, D., Cuturi, M., and Peyré, G.  Ground metric learning on graphs.  _Journal of Mathematical Imaging and Vision_ , pp. 1–19, 2020.
  * Indyk et al. (2019) Indyk, P., Vakilian, A., Wagner, T., and Woodruff, D.  Sample-optimal low-rank approximation of distance matrices, 2019.
  * Janati et al. (2020) Janati, H., Bazeille, T., Thirion, B., Cuturi, M., and Gramfort, A.  Multi-subject meg/eeg source imaging with sparse multi-task regression.  _NeuroImage_ , pp. 116847, 2020.
  * Kantorovich (1942) Kantorovich, L.  On the transfer of masses (in russian).  _Doklady Akademii Nauk_ , 37(2):227–229, 1942\.
  * Korotin et al. (2021) Korotin, A., Li, L., Solomon, J., and Burnaev, E.  Continuous wasserstein-2 barycenter estimation without minimax optimization.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=3tFAs5E-Pe.
  * Koundal et al. (2020) Koundal, S., Elkin, R., Nadeem, S., Xue, Y., Constantinou, S., Sanggaard, S., Liu, X., Monte, B., Xu, F., Van Nostrand, W., et al.  Optimal mass transport with lagrangian workflow reveals advective and diffusion driven solute transport in the glymphatic system.  _Scientific reports_ , 10(1):1–18, 2020.
  * Lin et al. (2019) Lin, T., Ho, N., and Jordan, M.  On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms.  In Chaudhuri, K. and Salakhutdinov, R. (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 3982–3991. PMLR, 09–15 Jun 2019.
  * Lu et al. (2017) Lu, H., Freund, R. M., and Nesterov, Y.  Relatively-smooth convex optimization by first-order methods, and applications, 2017.
  * Makkuva et al. (2020) Makkuva, A., Taghvaei, A., Oh, S., and Lee, J.  Optimal transport mapping via input convex neural networks.  In III, H. D. and Singh, A. (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 6672–6681. PMLR, 13–18 Jul 2020.
  * Peyré & Cuturi (2019) Peyré, G. and Cuturi, M.  Computational optimal transport.  _Foundations and Trends in Machine Learning_ , 11(5-6), 2019\.  ISSN 1935-8245.
  * Scetbon & Cuturi (2020) Scetbon, M. and Cuturi, M.  Linear time sinkhorn divergences using positive features, 2020.
  * Schiebinger et al. (2019) Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., et al.  Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming.  _Cell_ , 176(4):928–943, 2019.
  * Schmitz et al. (2018) Schmitz, M. A., Heitz, M., Bonneel, N., Ngole, F., Coeurjolly, D., Cuturi, M., Peyré, G., and Starck, J.-L.  Wasserstein dictionary learning: Optimal transport-based unsupervised nonlinear dictionary learning.  _SIAM Journal on Imaging Sciences_ , 11(1):643–678, 2018.
  * Schmitzer (2019) Schmitzer, B.  Stabilized sparse scaling algorithms for entropy regularized transport problems.  _SIAM Journal on Scientific Computing_ , 41(3):A1443–A1481, 2019.
  * Tong et al. (2020) Tong, A., Huang, J., Wolf, G., Van Dijk, D., and Krishnaswamy, S.  TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics.  In III, H. D. and Singh, A. (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 9526–9536. PMLR, 13–18 Jul 2020.
  * Xie et al. (2020) Xie, Y., Wang, X., Wang, R., and Zha, H.  A fast proximal point method for computing exact wasserstein distance.  In _Uncertainty in Artificial Intelligence_ , pp. 433–453. PMLR, 2020.
  * Yang et al. (2020) Yang, K. D., Damodaran, K., Venkatachalapathy, S., Soylemezoglu, A. C., Shivashankar, G., and Uhler, C.  Predicting cell lineages using autoencoders and optimal transport.  _PLoS computational biology_ , 16(4):e1007828, 2020.
  * Zhang et al. (2020) Zhang, K. S., Peyré, G., Fadili, J., and Pereyra, M.  Wasserstein control of mirror langevin monte carlo, 2020.

## Supplementary material

In Sec. A, we introduce some important notions linked to the mirror-descent
scheme. We also prove in this section a general result which states the non-
asymptotic stationary convergence of the mirror-descent according to a
specific criterion introcuded in this work. In Sec. B, we detail the
computation of the Dykstra’s algorithm 2 for which we have obtained a simple
expression of the updates of the couplings. In Sec. C, we provides all the
proofs of the Propositions introduced in this work in the main text. In Sec D,
we detail the algorithm presented in (Indyk et al., 2019). In Sec. E, F, we
give two variants of our algorithm when either the marginal $g$ is fixed or
when no lower bound is provided on the coordinates of $g$. In Sec. G, we
provides more experiment to illustrate our method.

##  Appendix A Mirror Descent Algorithm

Let $\mathcal{X}$ a closed convex subset in a Euclidean space
$\mathbb{R}^{q}$, $f:\mathcal{X}\rightarrow{\mathbb{R}}$ continuously
differentiable and let us consider the following problem

$\displaystyle\min_{x\in\mathcal{X}}f(x).$ (11)

Given a convex function $h:\mathcal{X}\rightarrow\mathbb{R}$ continuously
differentiable, one can define the _prox-function_ associated to $h$ as

$\displaystyle D_{h}(x,z):=h(x)-h(z)-\langle\nabla h(z),x-z\rangle.$

To solve Eq. (11), one can employ the mirror-descent (MD) algorithm. Given an
initial point $x_{0}\in\mathcal{X}$ and a sequence of positive step-size
$(\gamma_{k})_{k\geq 0}$, the mirror-descent scheme associated to the _prox-
function_ $D_{h}$ computes

$\displaystyle x_{k+1}=\operatorname*{argmin}_{x\in\mathcal{X}}\langle\nabla
f(x_{k}),x\rangle+\frac{1}{\gamma_{k}}D_{h}(x,x_{k}).$

In the following, we need to introduce two notions of relative strong
convexity and relative smoothness in order to prove non-asymptotic stationary
convergence of the MD scheme.

######  Definition (Relative smoothness.).

Let $L>0$ and $f$ continuously differentiable on $\mathcal{X}$. $f$ is said to
be $L$-smooth relatively to $h$ if

$\displaystyle f(y)\leq f(x)+\langle\nabla f(x),y-x\rangle+LD_{h}(y,x)$

######  Definition (Relative strong convexity).

Let $\alpha>0$ and $f$ continuously differentiable on $\mathcal{X}$. $f$ is
said to be $\alpha$-strongly convex relatively to $h$ if

$\displaystyle f(y)\geq f(x)+\langle\nabla f(x),y-x\rangle+\alpha
D_{h}(y,x)~{}\forall~{}x,y\in\mathcal{X}$

Note that $h$ is always 1-strongly convex relatively to $h$. Let us now prove
a general result to show non-asymptotic stationary convergence of the MD
scheme. For that purpose, we introduce for all $k\geq 0$ the following
criterion to establish convergence:

$\displaystyle\Delta_{k}\triangleq\frac{1}{\gamma_{k}^{2}}(D_{h}(x_{k},x_{k+1})+D_{h}(x_{k+1},x_{k})).$

######  Proposition 5.

Let $N\geq 1$, $f$ continuously differentiable on $\mathcal{X}$ which is
$L$-smooth relatively to $h$. By considering for all $k=1,\dots,N$,
$\gamma_{k}=1/2L$, and by denoting
$D_{0}=f(x_{0})-\min_{x\in\mathcal{X}}f(x)$, we have

$\displaystyle\min_{0\leq k\leq N-1}\Delta_{k}\leq\frac{4LD_{0}}{N}.$

###### Proof.

Let $k\geq 0$, then by $L$-smoothness of $f$, we have

$\displaystyle f(x_{k+1})\leq f(x_{k})+\langle\nabla
f(x_{k}),x_{k+1}-x_{k}\rangle+LD_{h}(x_{k+1},x_{k}),$

and by optimality of $x_{k+1}$, we have for all $x\in\mathcal{X}$,

$\displaystyle\langle\nabla f(x_{k})+\frac{1}{\gamma_{k}}[\nabla
h(x_{k+1})-\nabla h(x_{k})],x-x_{k+1}\rangle\geq 0,$

which implies, by taking $x=x_{k}$, that

$\displaystyle\langle\nabla f(x_{k}),x_{k}-x_{k+1}\rangle$
$\displaystyle\geq\frac{1}{\gamma_{k}}[-\langle\nabla
h(x_{k+1}),x_{k}-x_{k+1}\rangle-\langle\nabla h(x_{k}),x_{k+1}-x_{k}\rangle]$
$\displaystyle\geq\frac{1}{\gamma_{k}}[D_{h}(x_{k},x_{k+1})+D_{h}(x_{k+1},x_{k})].$

Then we have

$\displaystyle f(x_{k+1})\leq
f(x_{k})-\frac{1}{\gamma_{k}}[D_{h}(x_{k},x_{k+1})+D_{h}(x_{k+1},x_{k})]+LD_{h}(x_{k+1},x_{k})+LD_{h}(x_{k},x_{k+1})$

where the last term is added by positivity of $D_{h}(\cdot,\cdot)$ (as $h$ is
supposed to be convex on $\mathcal{X}$). Finally we obtain that

$\displaystyle\left(\sum_{k=0}^{N-1}\gamma_{k}(1-\gamma_{k}L)\Delta_{k}\right)\leq
f(x_{0})-f(x_{N})\leq D_{0},$

and as soon as $\gamma_{k}<\frac{1}{L}$, we have

$\displaystyle\min_{0\leq k\leq
N-1}\Delta_{k}\leq\frac{D_{0}}{\left(\sum_{k=0}^{N-1}\gamma_{k}(1-\gamma_{k}L)\right)}.$

Then by taking $\gamma_{k}=\frac{1}{2L}$, the result follows. ∎

In this paper, we consider $h$ to be the negative entropy function defined on
$\Delta_{q}^{*}$ as

$\displaystyle h(x)=\sum_{i=1}^{q}x_{i}\log(x_{i}).$ (12)

Therefore the _prox-function_ associated is just the Kullback–Leibler
divergence (KL) defined as,

$\displaystyle\text{KL}(x,z)=\sum_{i=1}^{q}x_{i}\log(x_{i}/z_{i}).$

Moreover if $\mathcal{X}\subset\prod_{i=1}^{p}\Delta_{q_{i}}^{*}$ for $p\geq
1$, we consider instead

$\displaystyle
h((x^{(1)},\dots,x^{(p)})):=\sum_{i=1}^{p}\sum_{j=1}^{q_{i}}x_{j}^{(i)}\log(x_{j}^{(i)})$

where the associated _prox-function_ is

$\displaystyle
D_{h}((x^{(1)},\dots,x^{(p)}),(z^{(1)},\dots,z^{(p)}))=\sum_{i=1}^{p}\text{KL}(x^{(i)},z^{(i)}).$

##  Appendix B The Dykstra’s Algorithm

In order to solve Eq. (10), we use the Dykstra’s Algorithm (Dykstra, 1983).
Given a closed convex set $\mathcal{C}\subset\mathbb{R}_{+}^{n\times
r}\times\mathbb{R}_{+}^{m\times r}\times\mathbb{R}_{+}^{r}$, we denote for all
$\bm{\xi}\in(\mathbb{R}_{+}^{*})^{n\times
r}\times(\mathbb{R}_{+}^{*})^{m\times r}\times(\mathbb{R}_{+}^{*})^{r}$ the
projection according to the Kullback-Leibler divergence as

$\displaystyle\mathcal{P}_{\mathcal{C}}^{\text{KL}}(\bm{\xi})\triangleq\operatorname*{argmin}_{\bm{\zeta}\in\mathcal{C}}\text{KL}(\bm{\zeta},\bm{\xi}).$

Starting from $\bm{\zeta}_{0}\triangleq\bm{\xi}$ and
$\bm{q}_{0}=\bm{q}_{-1}=(\mathbf{1},\mathbf{1},\mathbf{1})\in\mathbb{R}_{+}^{n\times
r}\times\mathbb{R}_{+}^{m\times r}\times\mathbb{R}_{+}^{r}$, the Dykstra’s
Algorithm 2 applied to our problem consists in computing for all $j\geq 0$,

$\displaystyle\bm{\zeta}_{2j+1}$
$\displaystyle=\mathcal{P}_{\mathcal{C}_{1}(a,b,r,\alpha)}^{\text{KL}}(\bm{\zeta}_{2j}\odot\bm{q}_{2j-1})$
$\displaystyle\bm{q}_{2j+1}$
$\displaystyle=\bm{q}_{2j-1}\odot\frac{\bm{\zeta}_{2j}}{\bm{\zeta}_{2j+1}}$
$\displaystyle\bm{\zeta}_{2j+2}$
$\displaystyle=\mathcal{P}_{\mathcal{C}_{2}(r)}^{\text{KL}}(\bm{\zeta}_{2j+1}\odot\bm{q}_{2j})$
$\displaystyle\bm{q}_{2j+2}$
$\displaystyle=\bm{q}_{2j}\odot\frac{\bm{\zeta}_{2j+1}}{\bm{\zeta}_{2j+2}}.$

In fact these operations can be simplified to simple matrix/vector
multiplications. More precisely, the Dykstra’s Algorithm produces the iterates
$(\bm{\zeta}_{j})_{j\geq 0}$ which satisfy for all $j\geq 0$
$\bm{\zeta}_{j}=(Q_{j},R_{j},g_{j})$ where

$\displaystyle Q_{j}$
$\displaystyle=\operatorname{diag}(u_{j}^{1})\xi^{(1)}\operatorname{diag}(v_{j}^{1})$
$\displaystyle R_{j}$
$\displaystyle=\operatorname{diag}(u_{j}^{2})\xi^{(2)}\operatorname{diag}(v_{j}^{2})$

for the sequences $(u_{j}^{i},v_{j}^{i})_{j\geq 0}$ initialized as,
$u_{0}^{i}\triangleq\mathbf{1}_{n}$, $v_{0}^{i}\triangleq\mathbf{1}_{m}$ for
all $i\in\\{1,2\\}$,
$q_{0,1}^{(3)}=q_{0,2}^{(3)}=q_{0}^{(1)}=q_{0}^{(2)}=\mathbf{1}_{r}$ and
computed with the iterations

$\displaystyle u_{n+1}^{k,i}$
$\displaystyle=\frac{p_{i}}{\xi_{k}^{i}v_{n}^{k,i}}$
$\displaystyle\tilde{g}_{n+1}$ $\displaystyle=\max(\alpha,g_{n}\odot
q_{n,1}^{(3)}),~{}q_{n+1,1}^{(3)}=(g_{n}\odot q_{n,1}^{(3)})/\tilde{g}_{n+1}$
$\displaystyle g_{n+1}$ $\displaystyle=(\tilde{g}_{n+1}\odot
q_{n,2}^{(3)})^{1/3}\prod_{i=1}^{2}(v_{n}^{k,i}\odot
q_{n}^{(i)}\odot(\xi_{k}^{i})^{T}u_{n}^{k,i})^{1/3}$ $\displaystyle
v_{n+1}^{k,i}$ $\displaystyle=\frac{g_{n+1}}{(\xi_{k}^{i})^{T}u_{n}^{k,i}}$
$\displaystyle q_{n+1}^{(i)}$ $\displaystyle=(v_{n}^{k,i}\odot
q_{n}^{(i)})/v_{n+1}^{k,i},~{}q_{n+1,2}^{(3)}=(\tilde{g}_{n+1}\odot
q_{n,2}^{(3)})/g_{n+1}$

##  Appendix C Proofs

###  C.1 Proof of Proposition 1

###### Proof.

The case when $\varepsilon=0$ is clear. Assume now that $\varepsilon>0$. When
$r=1$, note that $\mathcal{C}_{1}(a,b,r)\cap\mathcal{C}_{2}(r)$ is closed as
$g=1$ and bounded, therefore and by continuity of the objective the mininum
exists. Let $r\geq 2$. First remarks that we always have
$\text{LOT}_{r,\varepsilon}(\mu,\nu)\leq\text{LOT}_{r-1,\varepsilon}(\mu,\nu)$.
Let us assume that (8) does not admits a minimum. Because the objective
$F_{\varepsilon}$ is a lower semi-continuous function on
$\overline{\mathcal{C}_{1}(a,b,r)}\cap\mathcal{C}_{2}(r)$, and by compacity of
$\overline{\mathcal{C}_{1}(a,b,r)}\cap\mathcal{C}_{2}(r)$, the objective
function admits a minimum
$(Q,R,g)\in\overline{\mathcal{C}_{1}(a,b,r)}\cap\mathcal{C}_{2}(r)$ and we
have $\text{LOT}_{r,\varepsilon}(\mu,\nu)=F_{\varepsilon}(Q,R,g)$. But as the
minimum is not attained on $\mathcal{C}_{1}(a,b,r)\cap\mathcal{C}_{2}(r)$, it
means that there exists at least one coordinate $i\in\\{1,\dots,r\\}$ such
that $g_{i}=0$. Then because the constraints, $Q$ and $R$ both admit a column
which is the null vector. By deleting these coordinates in $Q,R,g$, we obtain
that
$\text{LOT}_{r,\varepsilon}(\mu,\nu)=\text{LOT}_{r-1,\varepsilon}(\mu,\nu)$. ∎

###  C.2 Proof of Proposition 2

###### Proof.

The first oder conditions of the projection gives that there exists
$(\lambda_{1},\lambda_{2},\lambda_{3})\in\mathbb{R}^{n}\times\mathbb{R}^{m}\times\mathbb{R}^{r}_{+}$
such that

$\displaystyle\log(Q/\tilde{Q})+\lambda_{1}\mathbf{1}^{T}$ $\displaystyle=0$
$\displaystyle\log(R/\tilde{R})+\lambda_{2}\mathbf{1}^{T}$ $\displaystyle=0$
$\displaystyle\log(g/\tilde{g})+\lambda_{3}$ $\displaystyle=0$

Moreover the conditions $Q\mathbf{1}=a$, $R\mathbf{1}=b$ and $g\geq\alpha$
imply that

$\displaystyle Q$ $\displaystyle=\text{Diag}(a/\tilde{Q}\mathbf{1})\tilde{Q}$
$\displaystyle R$ $\displaystyle=\text{Diag}(b/\tilde{R}\mathbf{1})\tilde{R}$
$\displaystyle g$ $\displaystyle=\max(\alpha,\tilde{g}).$

∎

###  C.3 Proof of Proposition 3

###### Proof.

The first order conditions of the projection states that there exists
$(\lambda_{1},\lambda_{2})\in\mathbb{R}^{r}\times\mathbb{R}^{r}$ such that

$\displaystyle\log(Q/\tilde{Q})+\mathbf{1}_{n}\lambda_{1}^{T}$
$\displaystyle=0$
$\displaystyle\log(R/\tilde{R})+\mathbf{1}_{m}\lambda_{2}^{T}$
$\displaystyle=0$ $\displaystyle\log(g/\tilde{g})-(\lambda_{1}+\lambda_{2})$
$\displaystyle=0$

Moreover the conditions $Q^{T}\mathbf{1}_{n}=R^{T}\mathbf{1}_{m}=g$ imply that

$\displaystyle Q$
$\displaystyle=\tilde{Q}\text{Diag}(g/\tilde{Q}^{T}\mathbf{1}_{n})$
$\displaystyle R$
$\displaystyle=\tilde{R}\text{Diag}(g/\tilde{R}^{T}\mathbf{1}_{m})$
$\displaystyle g^{3}$
$\displaystyle=\tilde{g}\odot\tilde{Q}^{T}\mathbf{1}_{n}\odot\tilde{R}^{T}\mathbf{1}_{m}$

from which the result follows. ∎

###  C.4 Proof of Proposition 4

###### Proof.

To show the result, we just need to show that

$F_{\varepsilon}:(Q,R,g)\in\mathcal{C}(a,b,r,\alpha)\rightarrow\langle
C,Q\operatorname{diag}(1/g)R^{T}\rangle-\varepsilon H(Q,R,g)$

is smooth relatively to

$H(Q,R,g):=\sum_{i,j}Q_{i,j}\log(Q_{i,j})+\sum_{i,j}R_{i,j}\log(R_{i,j})+\sum_{j}g_{j}\log(g_{j}),$

then by applying Proposition 5, the result will follow. Let us now show that
$F_{\varepsilon}$ is $L_{\varepsilon,\alpha}$-smooth. To do so, it is enough
to show that (Lu et al., 2017; Zhang et al., 2020)

$\displaystyle\|\nabla F_{\varepsilon}(Q_{1},R_{1},g_{1})-\nabla
F_{\varepsilon}(Q_{2},R_{2},g_{2})\|_{2}\leq
L_{\varepsilon,\alpha}\|H(Q_{1},R_{1},g_{1})-H(Q_{2},R_{2},g_{2})\|_{2}.$

We first have that

$\displaystyle\nabla
F_{\varepsilon}(Q,R,g)=\left(CR\operatorname{diag}(1/g)+\varepsilon(\log
Q+\mathbf{1}),C^{T}Q\operatorname{diag}(1/g)+\varepsilon(\log
R+\mathbf{1}),-\mathcal{D}(Q^{T}RC)/g^{2}+\varepsilon(\log g+1)\right)$

Now we have,

$\displaystyle\|\nabla F_{\varepsilon}(Q_{1})-\nabla
F_{\varepsilon}(Q_{2})\|_{2}^{2}$
$\displaystyle\leq\|CR_{1}\operatorname{diag}(1/g_{1})-CR_{2}\operatorname{diag}(1/g_{2})\|_{2}^{2}+\varepsilon^{2}\|\log
Q_{1}-\log Q_{2}\|_{2}^{2}$ $\displaystyle+2\varepsilon\|\log Q_{1}-\log
Q_{2}\|_{2}\|CR_{1}\operatorname{diag}(1/g_{1})-CR_{2}\operatorname{diag}(1/g_{2})\|_{2}$
$\displaystyle\leq\|C\|_{2}^{2}\|(R_{1}-R_{2})\operatorname{diag}(1/g_{1})+(\operatorname{diag}(1/g_{1})-\operatorname{diag}(1/g_{2}))R_{2}\|_{2}^{2}+\varepsilon^{2}\|\log
Q_{1}-\log Q_{2}\|_{2}^{2}$ $\displaystyle+2\varepsilon\|\log Q_{1}-\log
Q_{2}\|_{2}\|CR_{1}\operatorname{diag}(1/g_{1})-CR_{2}\operatorname{diag}(1/g_{2})\|_{2}$
$\displaystyle\leq\|C\|_{2}^{2}\left[\frac{\|R_{1}-R_{2}\|_{2}^{2}}{\alpha^{2}}+\|1/g_{1}-1/g_{2}\|_{2}^{2}+\frac{\|R_{1}-R_{2}\|\|1/g_{1}-1/g_{2}\|_{2}}{\alpha}\right]+\varepsilon^{2}\|\log
Q_{1}-\log Q_{2}\|_{2}^{2}$ $\displaystyle+2\varepsilon\|\log Q_{1}-\log
Q_{2}\|_{2}\|CR_{1}\operatorname{diag}(1/g_{1})-CR_{2}\operatorname{diag}(1/g_{2})\|_{2}.$

As $Q\rightarrow H(Q)$ is 1-strongly convex w.r.t to the $\ell_{2}$-norm on
$\Delta_{n\times r}$, we have

$\displaystyle\|Q_{1}-Q_{2}\|_{2}^{2}$ $\displaystyle\leq\langle\log
Q_{1}-\log Q_{2},Q_{1}-Q_{2}\rangle$ $\displaystyle\leq\|\log Q_{1}-\log
Q_{2}\|_{2}\|Q_{1}-Q_{2}\|_{2}$

from which follows that

$\displaystyle\|Q_{1}-Q_{2}\|_{2}\leq\log Q_{1}-\log Q_{2}\|_{2}.$

Moreover we have

$\displaystyle\|1/g_{1}-1/g_{2}\|_{2}\leq\frac{\|g_{1}-g_{2}\|_{2}}{\alpha^{2}}\leq\|\frac{\|\log
g_{1}-\log g_{2}\|_{2}}{\alpha^{2}}$

Therefore we obtain that

$\displaystyle\|\nabla F_{\varepsilon}(Q_{1})-\nabla
F_{\varepsilon}(Q_{2})\|_{2}^{2}$
$\displaystyle\leq\left(\frac{\|C\|_{2}}{\alpha}\|\log R_{1}-\log
R_{2}\|_{2}+\frac{\|C\|_{2}}{\alpha^{2}}\|\log g_{1}-\log
g_{2}\|_{2}+\varepsilon\|\log Q_{1}-\log Q_{2}\|_{2}\right)^{2}.$

An analogue proof leads to

$\displaystyle\|\nabla F_{\varepsilon}(R_{1})-\nabla
F_{\varepsilon}(R_{2})\|_{2}^{2}$
$\displaystyle\leq\left(\frac{\|C\|_{2}}{\alpha}\|\log Q_{1}-\log
Q_{2}\|_{2}+\frac{\|C\|_{2}}{\alpha^{2}}\|\log g_{1}-\log
g_{2}\|_{2}+\varepsilon\|\log R_{1}-\log R_{2}\|_{2}\right)^{2}.$

Let us now consider smoothness of $F_{\varepsilon}$ w.r.t $g$,

$\displaystyle\|\nabla F_{\varepsilon}(g_{1})-\nabla
F_{\varepsilon}(g_{2})\|_{2}^{2}$
$\displaystyle\leq\left\|\frac{\mathcal{D}(Q_{1}^{T}CR_{1})}{g_{1}^{2}}-\frac{\mathcal{D}(Q_{2}^{T}CR_{2})}{g_{2}^{2}}\right\|_{2}^{2}+\varepsilon^{2}\|\log
g_{1}-\log g_{2}\|_{2}^{2}$
$\displaystyle+2\varepsilon\left\|\frac{\mathcal{D}(Q_{1}^{T}CR_{1})}{g_{1}^{2}}-\frac{\mathcal{D}(Q_{2}^{T}CR_{2})}{g_{2}^{2}}\right\|_{2}\|\log
g_{1}-\log g_{2}\|_{2}.$

but we have that

$\displaystyle\left\|\frac{\mathcal{D}(Q_{1}^{T}CR_{1})}{g_{1}^{2}}-\frac{\mathcal{D}(Q_{2}^{T}CR_{2})}{g_{2}^{2}}\right\|_{2}^{2}$
$\displaystyle\leq\|(1/g_{1}^{2}-1/g_{2}^{2})\operatorname{diag}(Q_{1}^{T}CR_{1})\|_{2}^{2}+\|\mathcal{D}(Q_{1}^{T}CR_{1})-\mathcal{D}(Q_{2}^{T}CR_{2})/g_{2}^{2}\|_{2}^{2}$
$\displaystyle+2\|(1/g_{1}^{2}-1/g_{2}^{2})\operatorname{diag}(Q_{1}^{T}CR_{1})\|_{2}\|\mathcal{D}(Q_{1}^{T}CR_{1})-\mathcal{D}(Q_{2}^{T}CR_{2})/g_{2}^{2}\|_{2}$
$\displaystyle\leq\left(\frac{1\|C\|_{2}}{\alpha^{2}}\|\log g_{1}-\log
g_{2}\|_{2}+\frac{\|C\|_{2}}{\alpha^{2}}\left[\|Q_{1}-Q_{2}\|_{2}^{2}+\|R_{1}-R_{2}\|_{2}\right]\right)^{2}.$

Therefore we obtain that

$\displaystyle\|\nabla F_{\varepsilon}(g_{1})-\nabla
F_{\varepsilon}(g_{2})\|_{2}^{2}\leq\left(\left(\frac{\varepsilon+2\|C\|_{2}}{\alpha^{3}}\right)\|\log
g_{1}-\log
g_{2}\|_{2}+\frac{\|C\|}{\alpha^{2}}\|Q_{1}-Q_{2}\|_{2}++\frac{\|C\|}{\alpha^{2}}\|R_{1}-R_{2}\|_{2}\right)^{2}$

Finally we obtain that

$\displaystyle\|\nabla F_{\varepsilon}(Q_{1},R_{1},g_{1})-\nabla
F_{\varepsilon}(Q_{2},R_{2},g_{2})\|_{2}^{2}$ $\displaystyle\leq
3\left(\frac{\|C\|_{2}^{2}}{\alpha^{2}}+\frac{\|C\|_{2}^{2}}{\alpha^{4}}+\varepsilon^{2}\right)[\|\log
Q_{1}-\log Q_{2}\|_{2}^{2}+\|\log R_{1}-\log R_{2}\|_{2}^{2}]$
$\displaystyle+3\left(\frac{2\|C\|_{2}^{2}}{\alpha^{4}}+\left(\frac{\varepsilon+2\|C\|_{2}}{\alpha^{3}}\right)^{2}\right)\|\log
g_{1}-\log g_{2}\|_{2}^{2}$

Thus we obtain that

$\displaystyle\|\nabla F_{\varepsilon}(Q_{1},R_{1},g_{1})-\nabla
F_{\varepsilon}(Q_{2},R_{2},g_{2})\|_{2}\leq L_{\varepsilon,\alpha}\|\nabla
H(Q_{1},R_{1},g_{1})-\nabla H(Q_{2},R_{2},g_{2})\|_{2}$

and the result follows. ∎

##  Appendix D Low-Rank Factorization of Distance Matrix

In this section we present the algorithm used to perform a low-rank
approximation of a distance matrix (Bakshi & Woodruff, 2018; Indyk et al.,
2019). Given a metric space $(\mathcal{X},d)$,
$X=\\{x_{i}\\}_{i=1}^{n}\in\mathcal{X}^{n}$ and
$Y=\\{y_{j}\\}_{j=1}^{m}\in\mathcal{X}^{m}$ we aim at obtaining a low-rank
approximation of the distance matrix $D=(d(x_{i},y_{j}))_{i,j}$ with a
precision $\gamma>0$. Let us now present the algorithm considered where we
have denoted $t=\lfloor r/\gamma\rfloor$.

Inputs: $X,Y,r,\gamma$
Choose $i^{*}\in\\{1,\dots,n\\}$, and $j^{*}\\{1,\dots,m\\}$ uniformly at
random.
For $i=1,\dots,n$, $p_{i}\leftarrow
d(x_{i},y_{j}^{*})^{2}+d(x_{i}^{*},y_{j}^{*})^{2}+\frac{1}{m}\sum_{j=1}^{m}d(x_{i}^{*},y_{j})^{2}$.
Independently choose $i^{(1)},\dots,i^{(t)}$ according $(p_{1},\dots,p_{n})$.
$X^{(t)}\leftarrow[x_{i^{(1)}},\dots,x_{i^{(t)}}],~{}P^{(t)}\leftarrow[\sqrt{tp_{i^{(1)}}},\dots,\sqrt{tp_{i^{(t)}}}],~{}S\leftarrow
d(X^{(t)},Y)/P^{(t)}$
Denote $S=[S^{(1)},\dots,S^{(m)}]$,
For $j=1,\dots,m$, $q_{j}\leftarrow\|S^{(j)}\|_{2}^{2}/\|S\|_{F}^{2}$
Independently choose $j^{(1)},\dots,j^{(t)}$ according $(q_{1},\dots,q_{m})$.
$S^{(t)}\leftarrow[S^{j^{(1)}},\dots,S^{j^{(t)}}],~{}Q^{(t)}\leftarrow[\sqrt{tq_{j^{(1)}}},\dots,\sqrt{tq_{j^{(t)}}}],~{}W\leftarrow
S^{(t)}/Q^{(t)}$
$U_{1},D_{1},V_{1}\leftarrow\text{SVD}(W)$ (decreasing order of singular
values).
$N\leftarrow[U_{1}{(1)},\dots,U_{1}^{(r)}],~{}N\leftarrow
S^{T}N/\|W^{T}N\|_{F}$
Choose $j^{(1)},\dots,j^{(t)}$ uniformly at random in $\\{1,\dots,m\\}$.
$Y^{(t)}\leftarrow[y_{j^{(1)}},\dots,y_{j^{(t)}}],D^{(t)}\leftarrow
d(X,Y^{(t)})/\sqrt{t}$.
$U_{2},D_{2},V_{2}=\text{SVD}(N^{T}N),~{}U_{2}\leftarrow
U_{2}/D_{2},~{}N^{(t)}\leftarrow[(N^{T})^{(j^{(1)})},\dots,(N^{T})^{(j^{(t)})}],~{}B\leftarrow
U_{2}^{T}N^{(t)}/\sqrt{t},~{}A\leftarrow(BB^{T})^{-1}$.
$Z\leftarrow AB(D^{(t)})^{T},~{}M\leftarrow Z^{T}U_{2}^{T}$
Result: $M,N$

Algorithm 4 $\text{LR-Distance}(X,Y,r,\gamma)$

##  Appendix E Positive low-rank factorization with fixed marginal

Let $g\in\Delta_{r}^{*}$, and let us for now consider the following problem

$\displaystyle\text{LOT}_{r,g}(\mu,\nu):=\min_{P\in\Pi_{a,g,b}}\langle
C,P\rangle.$ (13)

By definition of $\Pi_{a,g,b}$, this problem can be formulated as follows:

$\displaystyle\text{LOT}_{r,g}(\mu,\nu)=\min_{\begin{subarray}{c}Q\in\Pi_{a,g}\\\
R\in\Pi_{b,g}\end{subarray}}\langle C,Q\text{Diag}(1/g)R^{T}\rangle.$ (14)

As in the classical OT problem, one can extend the above objective and
consider for any $\varepsilon\geq 0$ an entropic version of the problem
defined as

$\displaystyle\text{LOT}_{r,g,\varepsilon}(\mu,\nu):=\min_{\begin{subarray}{c}Q\in\Pi_{a,g}\\\
R\in\Pi_{b,g}\end{subarray}}\langle
C,Q\text{Diag}(1/g)R^{T}\rangle-\varepsilon H((Q,R))$ (15)

Note that for any $\varepsilon\geq 0$, the minimum always exists as the
objective is continuous and $\Pi_{a,g,b}$ is compact. Moreover we clearly have
that $\text{LOT}_{r,g,0}(\mu,\nu)=\text{LOT}_{r,g}(\mu,\nu)$. Applying a MD
method to the objective (14) leads for all $k\geq 0$ to the following updates

$\displaystyle Q_{k+1}$
$\displaystyle:=\operatorname*{argmin}_{Q\in\Pi_{a,g}}\langle
C_{k}^{(1)},Q\rangle-\frac{1}{\gamma_{k}}H(Q)$ $\displaystyle R_{k+1}$
$\displaystyle:=\operatorname*{argmin}_{R\in\Pi_{a,g}}\langle
C_{k}^{(2)}R\rangle-\frac{1}{\gamma_{k}}H(R)$

where, $(Q_{0},R_{0})\in\Pi_{a,g}\times\Pi_{b,g}$ is an initial point,
$C_{k}^{(1)}:=CR_{k}\text{Diag}(1/g)+(\varepsilon-\frac{1}{\gamma_{k}})\log(Q_{k})$,
$C_{k}^{(2)}:=C^{T}Q_{k}\text{Diag}(1/g)+(\varepsilon-\frac{1}{\gamma_{k}})\log(R_{k})$
and $\gamma_{k}$ is a sequence of positive real numbers. Therefore a MD method
bowls down to solve at each iteration two regularized OT problems which can be
done efficiently using the Sinkhorn algorithm (1).

#### Convergence of the Mirror Descent.

Even if the objective (14) is not convex in $(Q,R)$, one can obtain the non-
asymptotic stationary convergence of the MD algorithm in this setting.

Let $f_{\varepsilon}$ be the objective function of the problem (15) defined on
$X:=\Pi_{a,g}\times\Pi_{b,g}$ and let us denotes for any $\gamma>0$ and $x\in
X$

$\displaystyle\mathcal{G}_{\varepsilon}(x,\gamma):=\operatorname*{argmin}_{u\in
X}\\{\langle\nabla f_{\varepsilon}(x),u\rangle+\frac{1}{\gamma}KL(u,x)\\}.$

Let us now define the following criterion to establish convergence:

$\displaystyle\Delta_{\varepsilon}(x,\gamma):=\frac{1}{\gamma^{2}}(KL(x,\mathcal{G}_{\varepsilon}(x,\gamma))+KL(\mathcal{G}_{\varepsilon}(x,\gamma),x)).$

To show the non-asymptotic stationary convergence, we show that for any
$\varepsilon\geq 0$, the objective is smooth relative to the entropy function
(Bauschke et al., 2017) and we extend the proof of (Ghadimi et al., 2013) to
this case.

######  Proposition.

Let $\varepsilon\geq 0$ and $N\geq 1$. By denoting
$L_{\varepsilon}:=\sqrt{2(\|C\|_{2}^{2}\|\text{Diag}(1/g)\|_{2}^{2}+\varepsilon^{2})}$
and by considering a constant stepsize in the MD scheme such that for all
$k=1,\dots,N$ $\gamma_{k}=\frac{1}{L_{\varepsilon}}$, we obtain that

$\displaystyle\min_{1\leq k\leq
N}\Delta_{\varepsilon}((Q_{k},R_{k}),\gamma_{k})\leq\frac{2L_{\varepsilon}D_{0}}{N}.$

where $D_{0}:=f_{\varepsilon}(Q_{0},R_{0})-\text{LOT}_{r,g,\varepsilon}$ is
the distance of the initial value to the optimal one.

###### Proof.

A similar proof of the one given for Proposition 4 gives that
$f_{\varepsilon}$ is $L_{\varepsilon}$-smooth relatively to $H$. ∎

Let us now introduce our first algorithm (5) to compute a positive low-rank
factorization of the optimal coupling. Here we consider the case where
$g:=\mathbf{1}_{r}/r$. Before introducing our algorithm it is worth noting
that a trivial initialization may lead to a trivial fixed point in the MD
updates. Indeed if one initialize $Q:=ag^{T}$ and $R:=bg^{T}$, then
$CR\text{Diag}(1/g)=Ca\mathbf{1}^{T}$ and
$C^{T}Q\text{Diag}(1/g)=C^{T}b\mathbf{1}^{T}$ and therefore $(Q,R)$ is a fixed
point of the MD. To avoid this, we initialize our algorithm in the following
way: let $\lambda:=\min_{i,j,k}(a_{i},b_{j},g_{k})/2$,
$a_{1}\in\Delta_{n}^{*}\backslash\\{a\\}$, $a_{2}:=(a-\lambda
a_{1})/(1-\lambda)$, $b_{1}\in\Delta_{n}^{*}\backslash\\{b\\}$,
$b_{2}:=(b-\lambda b_{1})/(1-\lambda)$,
$g_{1}\in\Delta_{r}^{*}\backslash\\{g\\}$ and $g_{2}:=(g-\lambda
g_{1})/(1-\lambda)$. We can now define our initialization as $Q:=\lambda
a_{1}g_{1}^{T}+(1-\lambda)a_{2}g_{2}^{T}$, $R:=\lambda
b_{1}g_{1}^{T}+(1-\lambda)b_{2}g_{2}^{T}$.

Inputs: $C,a,b,\delta,Q,R,g,\gamma,\delta_{\text{S}}$
repeat

$Q_{\text{old}}\leftarrow Q,~{}R_{\text{old}}\leftarrow R\\\ C^{(1)}\leftarrow
CR\text{Diag}(1/g)-\frac{1}{\gamma}\log(Q),\\\ C^{(2)}\leftarrow
C^{T}Q\text{Diag}(1/g)-\frac{1}{\gamma}\log(R),\\\
K^{(1)}\leftarrow\exp(-\gamma C^{(1)}),\\\ K^{(2)}\leftarrow\exp(-\gamma
C^{(2)}),\\\
u,v\leftarrow\text{Sinkhorn}(K^{(1)},a,g,\delta_{\text{S}})~{}(\text{Algorithm}~{}(\ref{alg-
sink})),\\\ Q\leftarrow\text{Diag}(u)K^{(1)}\text{Diag}(v),\\\
u,v\leftarrow\text{Sinkhorn}(K^{(2)},a,g,\delta_{\text{S}})~{}(\text{Algorithm}~{}(\ref{alg-
sink})),\\\ R\leftarrow\text{Diag}(u)K^{(2)}\text{Diag}(v)$

until  _$\Delta((Q,R),\gamma) <\delta$_;

Result: $Q,R$

Algorithm 5 $\text{LOT-F}(C,a,b,\delta)$

#### Computational Cost.

Note that the kernels $(K^{(i)})_{1\leq i\leq 2}$ considered in algorithm (5)
live in $\mathbb{R}_{+}^{n\times r}\times\mathbb{R}_{+}^{m\times r}$ and
therefore each iteration of both Sinkhorn algorithms can be computed either in
$\mathcal{O}(nr)$ or in $\mathcal{O}(mr)$ algebraic operations as it involves
only matrix/vector multiplications of the form $K^{(i)}v$ and
$(K^{(i)})^{T}u$. However without any assumption on the cost matrix $C$,
computing $(K^{(i)})_{1\leq i\leq 2}$ costs $\mathcal{O}(nmr)$ algebraic
operations as it requires to compute both $CR$ and $C^{T}Q$ at each iteration.
Thanks to assumption 1, such multiplications can be performed in
$\mathcal{O}((n+m)dr)$ algebraic operations and thus algorithm (5) requires
only a linear number of algebraic operations with respect to the number of
samples at each iteration.

In the following, we will see that if we do not fix the marginal, the problem
can also be solved efficiently as each iteration of the MD algorithm can be
seen as a wasserstein barycenter problem.

##  Appendix F A Positive low-rank factorization with free marginal

Applying a MD method to the objective (8) leads, for all $k\geq 0$, to the
following updates

$(Q_{k+1},R_{k+1},g_{k+1}):=\operatorname*{argmin}_{\bm{\zeta}\in\mathcal{C}_{1}(a,b,r)\cap\mathcal{C}_{2}(r)}\text{KL}(\bm{\zeta},\bm{\xi}_{k})$
(16)

where $(Q_{0},R_{0},g_{0})\in\mathcal{C}_{1}(a,b,r)\cap\mathcal{C}_{2}(r)$ is
an initial point, $\bm{\xi}_{k}:=(\xi_{k}^{(1)},\xi_{k}^{(2)},\xi_{k}^{(3)})$,
$\xi_{k}^{(1)}:=\exp(-\gamma_{k}CR_{k}\text{Diag}(1/g_{k})_{k}-(\gamma_{k}\varepsilon-1)\log(Q_{k}))$,
$\xi_{k}^{(2)}:=\exp(-\gamma_{k}C^{T}Q_{k}\text{Diag}(1/g_{k})-(\gamma_{k}\varepsilon-1)\log(R_{k}))$,
$\xi_{k}^{(3)}:=\exp(\gamma_{k}\omega_{k}/g_{k}^{2}-(\gamma_{k}\varepsilon-1)\log(g_{k}))$
with $[\omega_{k}]_{i}:=[Q_{k}^{T}CR_{k}]_{i,i}$ for all $i\in\\{1,\dots,r\\}$
and $(\gamma_{k})_{k\geq 0}$ is a sequence of positive real numbers.

Eq. (16) is well defined. Indeed as the kernels $(\xi_{k}^{(i)})$ are matrices
with positive coefficients, the infimum is attained in
$\mathcal{C}_{1}(a,b,r)\cap\mathcal{C}_{2}(r)$ and the problem admits a unique
solution. Moreover solving Eq. (16) bowls down to solve

$(Q_{k+1},R_{k+1},g_{k+1}):=\\\
\operatorname*{argmin}_{\bm{\zeta}\in\overline{\mathcal{C}_{1}(a,b,r)}\cap\mathcal{C}_{2}(r)}\text{KL}(\bm{\zeta},\bm{\xi}_{k})$
(17)

In order to solve Eq. (17), we consider the Iterative Bregman Projections
(IBP) algorithm. Starting from $\bm{\zeta}^{(k)}_{0}:=\bm{\xi}_{k}$, the IBP
algorithm consists in computing for all $j\geq 0$,

$\displaystyle\bm{\zeta}_{2j+1}^{(k)}$
$\displaystyle=\mathcal{P}_{\overline{\mathcal{C}_{1}(a,b,r)}}^{\text{KL}}(\bm{\zeta}_{2j}^{(k)})$
$\displaystyle\bm{\zeta}_{2j+2}^{(k)}$
$\displaystyle=\mathcal{P}_{\mathcal{C}_{2}(r)}^{\text{KL}}(\bm{\zeta}_{2j+1}^{(k)}).$

As $\overline{\mathcal{C}_{1}(a,b,r)}$ and $\mathcal{C}_{2}(r)$ are affine
subspaces (note that nonnegativity constraints are already in the definition
of the objective) one can show that $\bm{\zeta}_{j}^{(k)}$ converges towards
the unique solution of Eq. (17), (Bregman, 1967). Remarks that the projection
on $\overline{\mathcal{C}_{1}(a,b,r)}$ can be computed very easily as one has
for any
$\tilde{\bm{\xi}}:=(\tilde{Q},\tilde{R},\tilde{g})\in\mathbb{R}_{+}^{n\times
r}\times\mathbb{R}_{+}^{n\times r}\times\mathbb{R}_{+}^{r}$,

$\displaystyle\mathcal{P}_{\overline{\mathcal{C}_{1}(a,b,r)}}^{\text{KL}}(\tilde{\bm{\xi}})=\left(\text{Diag}\left(\frac{a}{\tilde{Q}\mathbf{1}_{r}}\right)\tilde{Q},\text{Diag}\left(\frac{b}{\tilde{R}\mathbf{1}_{r}}\right)\tilde{R},\tilde{g}\right)$

and the solution of the projection on $\mathcal{C}_{2}(r)$ is already given in
Proposition 3.

#### Efficient computation of the updates.

For all $k\geq 0$, starting with $\bm{\zeta}_{0}^{(k)}:=\bm{\xi}_{k}$ the IBP
algorithm leads to a simple algorithm (6) which computes only scaling vectors.
More precisely, the IBP algorithm produces the iterates
$(\bm{\zeta}_{n}^{(k)})_{n\geq 0}$ which satisfy for all $n\geq 0$
$\bm{\zeta}_{n}^{(k)}=(Q_{n}^{(k)},R_{n}^{(k)},g_{n}^{(k)})$ where

$\displaystyle Q_{n}^{(k)}$
$\displaystyle=\text{Diag}(u_{n}^{k,1})\xi_{k}^{1}\text{Diag}(v_{n}^{k,1})$
$\displaystyle R_{n}^{(k)}$
$\displaystyle=\text{Diag}(u_{n}^{k,2})\xi_{k}^{2}\text{Diag}(v_{n}^{k,2})$

for the sequences $(u_{n}^{k,i},v_{n}^{k,i})$ initialized as
$v_{0}^{k,i}:=\mathbf{1}$ for all $i\in\\{1,2\\}$ and computed with the
iterations

$\displaystyle u_{n}^{k,i}$
$\displaystyle=\frac{p_{i}}{\xi_{k}^{i}v_{n}^{k,i}}$ $\displaystyle
g_{n+1}^{(k)}$
$\displaystyle=(g_{n}^{(k)})^{1/3}\prod_{i=1}^{2}(v_{n}^{k,i}\odot(\xi_{k}^{i})^{T}u_{n}^{k,i})^{1/3}$
$\displaystyle v_{n+1}^{k,i}$
$\displaystyle=\frac{g_{n+1}^{(k)}}{(\xi_{k}^{i})^{T}u_{n}^{k,i}}$

where we have denoted $p_{1}:=a$ and $p_{2}:=b$ to simplify the notations.

Inputs: $\xi^{(1)},\xi^{(2)},g:=\xi^{(3)},p_{1},p_{2},\delta,v^{(i)}$
repeat

$u^{(i)}\leftarrow p_{i}/\xi^{(i)}v^{(i)}~{}\forall i\in\\{1,2\\},\\\
g\leftarrow(g)^{1/3}\prod_{i=1}^{2}(v^{(i)}\odot(\xi^{(i)})^{T}u^{(i)})^{1/3},\\\
v^{(i)}\leftarrow g/(\xi^{(i)})^{T}u^{(i)}~{}\forall i\in\\{1,2\\}$

until  _$\sum_{i=1}^{2}\|u^{(i)}\odot\xi^{(i)}v^{(i)}-p_{i}\|_{1} <\delta$_;

$Q\leftarrow\text{Diag}(u^{(1)})\xi_{k}^{(1)}\text{Diag}(v^{(1)})$
$R\leftarrow\text{Diag}(u^{(2)})\xi_{k}^{(2)}\text{Diag}(v^{(2)})$
Result: $Q,R,g$

Algorithm 6 $\text{LR-IBP}((\xi^{(i)})_{1\leq i\leq 3},p_{1},p_{2},\delta)$

Let us now introduce the proposed MD algorithm applied to (7). By denoting
$\mathcal{D}(\cdot)$ the operator extracting the diagonal of a square matrix
we obtain the following algorithm (7) to solve Eq. (6). We initialize our
algorithm with the exact same procedure as in algorithm (5).

Inputs: $C,a,b,(\gamma_{k})_{k\geq 0},Q,R,g,\delta$
for  _$k=1,\dots$_ do

$\xi^{(1)}\leftarrow\exp(-\gamma_{k}CR\text{Diag}(1/g)-(\gamma_{k}\varepsilon-1)\log(Q)),\\\
\xi^{(2)}\leftarrow\exp(-\gamma_{k}C^{T}Q\text{Diag}(1/g)-(\gamma_{k}\varepsilon-1)\log(R)),\\\
\omega\leftarrow\mathcal{D}(Q^{T}CR),~{}\xi^{(3)}\leftarrow\exp(\gamma_{k}\omega/g^{2}-(\gamma_{k}\varepsilon-1)\log(g)),\\\
Q,R,g\leftarrow\text{LR-IBP}((\xi^{(i)})_{1\leq i\leq
3},a,b,\delta)~{}(\text{Algorithm~{}(\ref{alg-IBP})})$

end for

Result: $\langle C,Q\text{Diag}(1/g)R^{T}\rangle$

Algorithm 7 $\text{LOT}(C,a,b,r,\delta)$

#### Computational Cost.

Note that $(\xi^{(i)})_{1\leq i\leq 3}$ considered in algorithm (7) lives in
$\mathbb{R}_{+}^{n\times r}\times\mathbb{R}_{+}^{m\times
r}\times\times\mathbb{R}_{+}^{r}$ and therefore each iteration of algorithm
(6) can be computed in $\mathcal{O}((n+m)r)$ algebraic operations as it
involves only matrix/vector multiplications of the form $\xi^{(i)}v_{i}$ and
$(\xi^{(i)})^{T}u_{i}$. However without any assumption on the cost matrix $C$,
computing $(\xi^{(i)})_{1\leq i\leq 3}$ costs $\mathcal{O}(nmr)$ algebraic
operations as it requires to compute both $CR$ and $C^{T}Q$ at each iteration.
Thanks to assumption 1, such multiplications can be performed in
$\mathcal{O}((n+m)dr)$ algebraic operations and thus algorithm (7) requires
only a linear number of algebraic operations with respect to the number of
samples at each iterations.

##  Appendix G Addiational Experiments

In Fig. 3, we compare two Gaussian mixture densities sampled with $n=m=10000$
points in 2D. The two densities considered are

$\displaystyle f_{X}(x)$
$\displaystyle=\frac{1}{3}\frac{\exp\left((x-\mu_{1})^{T}\Sigma^{-1}(x-\mu_{1})\right)}{\sqrt{2\pi|\Sigma|}}+\frac{1}{3}\frac{\exp\left((x-\mu_{2})^{T}\Sigma^{-1}(x-\mu_{2})\right)}{\sqrt{2\pi|\Sigma|}}+\frac{1}{3}\frac{\exp\left((x-\mu_{3})^{T}\Sigma^{-1}(x-\mu_{3})\right)}{\sqrt{2\pi|\Sigma|}}$
$\displaystyle f_{Y}(x)$
$\displaystyle=\frac{1}{2}\frac{\exp\left((x-\nu_{1})^{T}\Sigma^{-1}(x-\nu_{1})\right)}{\sqrt{2\pi|\Sigma|}}+\frac{1}{2}\frac{\exp\left((x-\nu_{2})^{T}\Sigma^{-1}(x-\nu_{2})\right)}{\sqrt{2\pi|\Sigma|}}$

where

$\displaystyle\mu_{1}=[0,0],\quad~{}\mu_{2}=[0,1],\quad~{}\mu_{3}=[1,1],\quad~{}\nu_{1}=[0.5,0.5],\quad~{}\nu_{2}=[-0.5,0.5],~{}\quad\Sigma=0.05\times\text{Id}_{2}.$

We show in Fig. 5 a plot of the two distributions considered.

Figure 5: Plot of the Gaussian mixtures considered in Fig. 3.

In Fig. 6, we consider the exact same setting as the one presented in Fig. 3
but we increase the dimension of the problem. More precisely we consider two
Gaussian mixture densities samples with $n=m=10000$ points in 10D where

$\displaystyle\mu_{1}=[0,\dots,0],~{}\mu_{2}=[0,1,0,\dots,0],~{}\mu_{3}=[1,1,0,\dots,0],$
$\displaystyle~{}\nu_{1}=[0.5,0.5,0,\dots,0],~{}\nu_{2}=[-0.5,0.5,0,\dots,0],$
$\displaystyle\Sigma=0.05\times\text{Id}_{10}.$

Similarly as in Fig. 3, we observe that LOT and LOT Quad provide similar
results while LOT is faster. All kernel-based methods fail to converge in this
setting. Moreover we see that for small regularizations $\varepsilon$, our
method is able to approximate faster than Sin the true OT thanks to the low-
rank constraint. Note also that we observe again a difference between the two
entropic regularizations of the Sin objective and LOT objective. Indeed the
range of $\varepsilon$ where Sin provides an efficient approximation of the
true OT is larger than the one of LOT. Indeed recall that for LOT, we
regularize _twice_ as we constraint the nonnegative rank of the couplings and
we add an entropic term to regularize the objective.

Figure 6: Comparison of the time-accuracy tradeoff for different methods for
estimating the OT or its regularized version between two mixture of gaussians
in 10D.

In Fig. 2, we compare the time-accuracy tradeoff for different methods on a
synthetic problem where we aim at estimating either the OT or its regularized
version between two gaussians in 2D. Here we consider the exact same setting
but we increase the dimension of the problem: $d=10$. As in Fig. 2, our
proposed method obtains an efficient approximation of the OT or its
regularized version for all rank $r$ faster than other low-rank methods in the
regime of small $\varepsilon$. We also see that for all low-rank methods, a
rank of $r=500$ is not enough in this setting to obtain the exact OT, but as
the rank increases, the approximation gets better.

Figure 7: In this experiment, we consider two Gaussian distributions evaluated
on $n=m=5000$ in 10D. The first one has a mean of
$(1,\dots,1)^{T}\in\mathbb{R}^{10}$ and identity covariance matrix $I_{10}$
while the other has 0 mean and covariance $0.1\times I_{10}$. The ground cost
is the squared Euclidean distance.