Structured Transforms Across Spaces
with Cost-Regularized Optimal Transport

Othmane Sebbouh Marco Cuturi Gabriel Peyré

ENS, CNRS, ENSAE Apple ENS and CNRS

###### Abstract

Matching a source to a target probability measure is often solved by
instantiating a linear optimal transport (OT) problem, parameterized by a
ground cost function that quantifies discrepancy between points. When these
measures live in the same metric space, the ground cost often defaults to its
distance. When instantiated across two different spaces, however, choosing
that cost in the absence of aligned data is a conundrum. As a result,
practitioners often resort to solving instead a quadratic Gromow-Wasserstein
(GW) problem. We exploit in this work a parallel between GW and cost-
regularized OT, the regularized minimization of a linear OT objective
parameterized by a ground cost. We use this cost-regularized formulation to
match measures across two different Euclidean spaces, where the cost is
evaluated between transformed source points and target points. We show that
several quadratic OT problems fall in this category, and consider enforcing
structure in linear transform (e.g. sparsity), by introducing structure-
inducing regularizers. We provide a proximal algorithm to extract such
transforms from unaligned data, and demonstrate its applicability to single-
cell spatial transcriptomics/multiomics matching tasks.

##  1 Introduction

Optimal Transport (OT) is by now an established tool in the machine learning
playbook, one that plays a key role when matching probability distributions.
OT has played a prominent role in generative modeling (Montavon et al., ,
2016; Genevay et al., , 2018; Tong et al., , 2023; Neklyudov et al., 2023a, ;
Lipman et al., , 2022; Neklyudov et al., 2023b, ), adversarial training (Sinha
et al., , 2018; Wong et al., , 2019), domain adaptation (Courty et al., ,
2017), neuroscience (Janati et al., , 2020), or single-cell modeling
(Schiebinger et al., , 2019; Tong et al., , 2020; Bunne et al., , 2023).
However, OT is not without drawbacks, and we target in this work an important
pain point that hinders the use of OT in practical tasks: the challenges of
choosing a ground cost function to compare two measures, particularly when
they are supported on heterogeneous and high-dimensional Euclidean spaces.

Choosing a cost or going quadratic? OT is most often used to match probability
distributions supported on the same space, in which case the ground cost is
often picked to be an (exponentiated) distance on that space. Yet, an
increasing number of applications require handling multimodality, i.e.
comparing distributions supported on heterogeneous spaces, as in multi-omics
singe-cell data matching (Demetci et al., , 2022; Klein et al., 2023a, ). For
such problems, defining a ground cost across two spaces (e.g. between vectors
of different dimensions), in the absence of other knowledge, is understandably
difficult. As a result, practitioners often bypass this issue by resorting to
the Gromov-Wasserstein (GW) framework (Mémoli, , 2011). GW instead
instantiates a quadratic problem in the space of couplings that involves two
within ground cost functions, one for each space. However, solving the GW
problem is fraught with computational challenges: it requires minimizing a
concave quadratic function on the space of couplings, an NP hard problem.
Moreover, some of the classic linear OT machinery, e.g. the existence of
Monge,  maps, does not directly translate to quadratic settings, as shown by
Vayer,  (2020); Dumont et al.,  (2022) who discussed their existence for a
variety of costs. This is a very active field of research, as shown by recent
contributions on the properties of entropic GW (Zhang et al., , 2022; Rioux et
al., , 2023); closed-form formulas for Gaussians (Salmona et al., , 2022; Le
et al., , 2022), or scalable computational schemes (Scetbon et al., , 2022;
Nekrashevich et al., , 2023).

High-dimensional issues. Another issue with OT comes from the significant
statistical challenges that arise when dealing with high-dimensional problems
(Dudley, , 1966; Weed and Bach, , 2019; Genevay et al., , 2019). One way
around this is performing offline dimensionality reduction, typically via PCA
or VAEs, or more elaborate schemes that carry out a joint estimation of a
projection and transport (Bonneel et al., , 2015; Niles-Weed and Rigollet, ,
2022; Paty and Cuturi, , 2019; Deshpande et al., , 2019; Le et al., , 2019). A
recent line of work proposes to leverage sparsity in displacements (Cuturi et
al., , 2023; Klein et al., 2023b, ). These challenges only add up when dealing
with heterogeneous spaces and taking the quadratic route advocated in GW. For
example, the standard iterative algorithm to solve the GW problem requires
solving a linear OT problem at each iteration (Peyré et al., , 2016), and is
only able to handle dimensionality as a way to speed up these routines
(Scetbon et al., , 2022)

Contributions. In this work, we target both issues in the case where
practitioners are faced with the task of transporting points across two
modalities, where each lies in a high-dimensional space.

  * •

We use a result by Paty and Cuturi,  (2020), who showed that minimizing a
concave function over couplings can be expressed as the minimization of a
cost-regularized OT problem, to show that most common quadratic OT
formulations can be unified under that “cost-regularized OT” lens.

  * •

We propose to leverage this fact to impose explicitly structure on costs using
a suitable convex regularization. Focusing on dot-product costs parameterized
by a linear transform across spaces, we propose a simple alternated
minimization algorithm, Prox-ROT, to solve cost-regularized OT when the
regularizer is evaluated on such transforms. Prox-ROT is particularly
efficient when the regularization function has a closed-form proximity
operator.

  * •

We use this framework to introduce sparsifying norms on such transforms and
show this is equivalent to variants of GW which only depend on an adaptively
selected subset of dimensions: Using Prox-ROT with such norms consists in
iteratively performing a feature selection step followed by an OT plan
computation step until convergence. This makes cost-regularized OT
particularly well-suited to high-dimensional problems, where linear OT and GW
usually struggle.

  * •

Building on Dumont et al.,  (2022), we show that there exist Monge maps for
the linear cost-regularized OT problem. Extending the framework of entropic
Monge maps (Pooladian and Niles-Weed, , 2021), we derive entropic Monge maps
across spaces, for linear cost-regularized OT, and demonstrate that they
converge to the ground truth Monge maps under suitable assumptions.

  * •

We apply our methods to toy and real-world data. We use sparsifying and low-
rank regularizations to solve high-dimensional mutliomics single-cell data
integration tasks while improving over the best-known OT baselines, and we
take advantage of the cost-regularized OT objective structure to use SGD to
solve a large-scale spatial transcriptomics problem.

The closest setting to ours appeared in Alvarez-Melis et al.,  (2019), who
proposed to use a constraint on the spectrum of the linear transforms and
showed the equivalence between the Gromov-Wasserstein problem and solving a
linear OT problem with a dot product cost using a Frobenius norm constraint on
the linear operator. The Prox-ROT algorithm in Section 3 using the nuclear or
the rank regularization results (up to reparametrization) in a similar
algorithm to the one proposed in Alvarez-Melis et al.,  (2019). We take the
alternative _regularization_ route and consider general regularizations.
Notably, we pay particular attention to sparsifying norms acting directly on
the entries of the linear operator, and draw links with GW-like problems on a
subset of the dimensions.

Notations. In what follows, we consider
$\left(\mathcal{X},d_{\mathcal{X}}\right)$ and
$\left(\mathcal{Y},d_{\mathcal{Y}}\right)$, two metric spaces, and
$\left(\alpha,\beta\right)\in\mathcal{P}(\mathcal{X})\times\mathcal{P}(\mathcal{Y})$,
two probability measures with compact supports. The space of couplings
$\Pi(\alpha,\beta)$ is defined as the probability distributions on
$\left(\mathcal{X}\times\mathcal{Y}\right)$ whose marginals are $\alpha$ and
$\beta$ (Santambrogio, , 2015):

$\displaystyle\Pi\left(\alpha,\beta\right)=\left\\{\gamma\in\mathcal{P}\left(\mathcal{X}\times\mathcal{Y}\right):\left(\pi_{x}\right)_{\sharp}\gamma=\alpha,\left(\pi_{y}\right)_{\sharp}\gamma=\beta\right\\},$
(1)

where $\pi_{x}$ and $\pi_{y}$ are the projections of
$\mathcal{X}\times\mathcal{Y}$ onto $\mathcal{X}$ and $\mathcal{Y}$
respectively. We denote the linear OT cost between $\alpha$ and $\beta$ with
cost $c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})$ as

$\displaystyle\mathcal{W}_{c}\left(\alpha,\beta\right)\triangleq\min_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}\times\mathcal{Y}}c(x,y)\mathop{}\\!\mathrm{d}\pi(x,y),$
(2)

which is a linear problem in $\pi$. When $c=d_{\mathcal{X}}^{p}$ and
$\mathcal{X}=\mathcal{Y}$, $\mathcal{W}_{c}^{1/p}$ defines a distance between
probability measures for all $p\geq 1$ (Villani, , 2003).

##  2 OT Across Spaces as a Concave Minimization Problem

While the classical OT objective is a linear function of the coupling $\pi$,
existing approaches to comparing probability distributions across spaces all
require minimizing a concave function over $\Pi(\alpha,\beta)$.

######  Problem 2.1.

Let $\mathcal{Q}$ be a concave function over $\Pi(\alpha,\beta)$. We define
the concave minimization over couplings problem as

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\mathcal{Q}(\pi).$
($\mathcal{Q}\textsc{OT}$)

Remarkably, concave minimization problems over couplings are closely linked to
linear OT via convex duality.

######  Proposition 2.2 (Remark 2 in Paty and Cuturi,  (2020)).

Let $\mathcal{Q}$ be a proper usc concave function over $\Pi(\alpha,\beta)$.
Writing $\mathcal{Q}^{*}$ for its convex conjugate, one has:

$\displaystyle\inf_{\pi\in\Pi(\alpha,\beta)}\mathcal{Q}(\pi)$
$\displaystyle=\inf_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})\end{subarray}}\int
c(x,y)\mathop{}\\!\mathrm{d}\pi+(-\mathcal{Q})^{*}(-c)$ (3)
$\displaystyle=\inf_{c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})}\mathcal{W}_{c}\left(\alpha,\beta\right)+(-\mathcal{Q})^{*}(-c).$
(4)

In other words, concave minimization problems over couplings can be
reformulated as the minimization over cost functions of a linear OT problem
plus a cost regularization. Consequently, we can define the following cost-
regularized reformulation of Problem 2.1.

######  Problem 2.3.

Let $\mathcal{R}$ be a convex function over
$\mathcal{C}\left(\mathcal{X}\times\mathcal{Y}\right)$. We define the cost-
regularized optimal transport problem as

$\displaystyle\mathcal{R}\textsc{OT}(\alpha,\beta)\triangleq\inf_{c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})}\mathcal{W}_{c}\left(\alpha,\beta\right)+\mathcal{R}(c).$
($\mathcal{R}\textsc{OT}$)

Note that, trivially, $\mathcal{R}\textsc{OT}$ $=$ $\mathcal{Q}\textsc{OT}$
when $\mathcal{R}(c)=\left(-\mathcal{Q}\right)^{*}(-c)$. We show next how
well-known OT formulations used to compare probability distributions across
spaces can be reduced to either $\mathcal{Q}\textsc{OT}$ or
$\mathcal{R}\textsc{OT}$. We say that two minimization problems are
equivalent, or that one is an instance of the other, when they have the same
minimizers in $\Pi(\alpha,\beta)$.

###  2.1 Examples

Sturm’s distance. Sturm,  (2006) defined a distance between metric measured
spaces (MMS) that quantifies how “isometric“ they are. Given two MMS
$(\mathcal{X},d_{X},\alpha)$ and $(\mathcal{Y},d_{Y},\beta)$, Sturm,  defined
the squared distance:

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\min_{\begin{subarray}{c}(Z,d_{Z})\in\mathcal{C}\\\
\phi:\mathcal{X}\to\mathcal{Z}\\\
\psi:\mathcal{Y}\to\mathcal{Z}\end{subarray}}\int_{\mathcal{X}\times\mathcal{Y}}d_{Z}^{2}(\phi(x),\psi(y))\,\mathrm{d}\pi(x,y),$
(5)

where $\phi:\mathcal{X}\to\mathcal{Z}$ and $\psi:\mathcal{Y}\to\mathcal{Z}$
are constrained to be isometries and
$\left(\mathcal{Z},d_{\mathcal{Z}}\right)$ is a metric space. As such, Sturm’s
distance is an instance of $\mathcal{R}\textsc{OT}$ with

$\displaystyle\mathcal{R}(c)=\left\\{\begin{array}[]{ll}0&\mbox{if
}\exists(\mathcal{Z},\phi,\psi):c(x,y)=d_{Z}^{2}(\phi(x),\psi(y))\\\
\infty&\mbox{otherwise.}\end{array}\right.$ (8)

Wasserstein Procrustes. A more tractable variant of Sturm’s distance when
$\mathcal{X},\mathcal{Y}\subset\mathbb{R}^{d}$ is the so-called Wasserstein
Procrustes problem (Zhang et al., , 2017; Grave et al., , 2019). It is defined
as

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\min_{\mathbf{C}\in\mathcal{O}_{d}}\int_{\mathcal{X}\times\mathcal{Y}}\left\lVert\mathbf{C}x-y\right\rVert^{2}\mathop{}\\!\mathrm{d}\pi(x,y).$
(9)

where $\mathcal{O}_{d}$ is the orthogonal manifold. Wasserstein Procrustes
defines a distance between $\mathcal{X}$ and $\mathcal{Y}$ up to a rotation.
It is an instance of $\mathcal{R}\textsc{OT}$ with

$\displaystyle\mathcal{R}(c)=\left\\{\begin{array}[]{ll}0&\mbox{if
}\exists\mathbf{C}\in\mathcal{O}_{d}:c(x,y)=\left\lVert\mathbf{C}x-y\right\rVert^{2}.\\\
\infty&\mbox{otherwise.}\end{array}\right.$ (12)

Remarkably, (9) is a weighted orthogonal Procrustes problem, and can be cast
as an instance of $\mathcal{Q}\textsc{OT}$.

######  Proposition 2.4.

Let $\mathcal{X},\mathcal{Y}\subset\mathcal{R}^{d}$. Then the Wasserstein
Procrustes distance is an instance of the concave minimization problem
($\mathcal{Q}\textsc{OT}$) with

$\displaystyle\mathcal{Q}(\pi)=\scaleobj{1.2}{\int}\left\lVert\mathbf{U}(\pi)\mathbf{V}(\pi)^{\top}x-y\right\rVert^{2}\mathop{}\\!\mathrm{d}\pi$

where $\mathbf{U}(\pi)\bm{\Sigma}(\pi)\mathbf{V}(\pi)^{\top}$ is an SVD of
${\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}\mathop{}\\!\mathrm{d}\pi}$.
Arbitrary Transformations. While Sturm considers isometries and the
Wasserstein Procrustes approach uses rotations, any type of a transformation
of $\alpha$ and/or $\beta$ could be used. For example, Cohen and Guibas,
(1999) studied a large range of transformations, including linear ones and
translations. Finally, learning generative models with OT (Genevay et al., ,
2018; Salimans et al., , 2018) with learned features to define costs is
another instance of $\mathcal{R}\textsc{OT}$.

The Gromov-Wasserstein distance. Given two costs $c_{\mathcal{X}}$ and
$c_{\mathcal{Y}}$, the Gromov-Wasserstein (GW) distance (Mémoli, , 2011) is
defined as

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{(\mathcal{X}\times\mathcal{Y})^{2}}\left(c_{\mathcal{X}}(x,x^{\prime})-c_{\mathcal{Y}}(y,y^{\prime})\right)^{2}\mathop{}\\!\mathrm{d}\pi\mathop{}\\!\mathrm{d}\pi.$
(13)

Mémoli,  (2011) showed that when $c_{\mathcal{X}}$ and $c_{\mathcal{Y}}$ are
distances, (13) defines a distance between metric measured spaces up to
isometry. This problem is not concave in general, but when $c_{\mathcal{X}}$
and $c_{\mathcal{Y}}$ define a conditionally negarive kernel, the function
$\pi\mapsto\langle Q(\pi),\pi\rangle$, where

$\displaystyle
Q\left((x,y),(x^{\prime},y^{\prime})\right):=\left(c_{\mathcal{X}}(x,x^{\prime})-c_{\mathcal{Y}}(y,y^{\prime})\right)^{2},$
(14)

is a concave quadratic function of $\pi$ (Dumont et al., , 2022, Proposition
1). In this case, (13) is an instance of ($\mathcal{Q}\textsc{OT}$) with
$\mathcal{Q}(\pi)=\langle Q(\pi),\pi\rangle$.

Two common cases that fit this setting are the inner product and the squared
Euclidean distance. In this case, it has been recently shown (Vayer, , 2020)
that the concave minimization problem has a simple cost-regularized
equivalent.

######  Proposition 2.5.

Let $\mathcal{X}\subset\mathbb{R}^{d_{x}}$ and
$\mathcal{Y}\subset\mathcal{R}^{d_{y}}$. For all
$\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}$, define
$c_{\mathbf{M}}(x,y)=-\langle\mathbf{M}x,y\rangle$.

• Let $c_{\mathcal{X}}=\langle\cdot,\cdot\rangle_{\mathcal{X}}$ and
$c_{\mathcal{Y}}=\langle\cdot,\cdot\rangle_{\mathcal{Y}}$. The problems

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int\left(\langle
x,x^{\prime}\rangle-\langle
y,y^{\prime}\rangle\right)^{2}\mathop{}\\!\mathrm{d}\pi\mathop{}\\!\mathrm{d}\pi$
(GW-IP) and $\displaystyle\min_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}\end{subarray}}\int-\langle
Mx,y\rangle\mathop{}\\!\mathrm{d}\pi+\frac{1}{2}\left\lVert
M\right\rVert_{F}^{2}$ (GW-IP-$\mathcal{R}\textsc{OT}$)

are equivalent: GW-IP is an instance of $\mathcal{R}\textsc{OT}$ with

$\displaystyle\mathcal{R}(c)=\left\\{\begin{array}[]{ll}\frac{1}{2}\left\lVert\mathbf{M}\right\rVert^{2}_{F}&\mbox{if
}\exists\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}:c=c_{\mathbf{M}}\\\
\infty&\mbox{otherwise.}\end{array}\right.$ (17)

• Let $c_{\mathcal{X}}={\left\lVert\cdot-\cdot\right\rVert}^{2}_{\mathcal{X}}$
and $c_{\mathcal{Y}}={\left\lVert\cdot-\cdot\right\rVert}^{2}_{\mathcal{Y}}$.
The problems

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int({\left\lVert
x-x^{\prime}\right\rVert}^{2}-{\left\lVert
y-y^{\prime}\right\rVert}^{2})^{2}\mathop{}\\!\mathrm{d}\pi\mathop{}\\!\mathrm{d}\pi$
(18) and $\displaystyle\min_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}\end{subarray}}\int-\langle
Mx,y\rangle-{\left\lVert x\right\rVert}^{2}{\left\lVert
y\right\rVert}^{2}\mathop{}\\!\mathrm{d}\pi+\frac{1}{2}\left\lVert
M\right\rVert_{F}^{2}$ (19)

are equivalent. Hence, (18) is an instance of $\mathcal{R}\textsc{OT}$ with

$\displaystyle\mathcal{R}(c)=\left\\{\begin{array}[]{ll}\frac{1}{2}\left\lVert\mathbf{M}\right\rVert^{2}_{F}&\mbox{if
}\exists\mathbf{M}:c=c_{\mathbf{M}}-{\left\lVert
x\right\rVert}^{2}{\left\lVert y\right\rVert}^{2}\\\
\infty&\mbox{otherwise.}\end{array}\right.$ (22)

###  2.2 Leveraging cost-regularized OT

In this work, we aim to incorporate additional structure to OT across spaces
by leveraging the $\mathcal{R}\textsc{OT}$ formulation. Inspired by GW, we
will constrain our costs to be linear and consider new regularization
functions. We will see that depending on the regularization function, we can
draw interesting links to GW-like problems by writing the
$\mathcal{R}\textsc{OT}$ problem in its equivalent $\mathcal{Q}\textsc{OT}$
form.

Entropic regularization. A preferred way to compute OT and GW distances in
practice is adding an entropic regularization to the linear OT term (Cuturi, ,
2013; Peyré et al., , 2016). Given an $\varepsilon>0$ we can follow the same
procedure and add a regularization to $\mathcal{R}\textsc{OT}$. The entropy-
regularized $\mathcal{R}\textsc{OT}$ problem writes

$\displaystyle\min_{c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})}\mathcal{W}_{c}^{\varepsilon}(\alpha,\beta)+\mathcal{R}(c)$
(23)
$\displaystyle\text{where}\;\,\mathcal{W}_{c}^{\varepsilon}(\alpha,\beta)\triangleq\min_{\pi\in\Pi(\alpha,\beta)}\int
c\mathop{}\\!\mathrm{d}\pi+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta),$
(24)

with
$\textsc{KL}(\pi||\alpha\otimes\beta)=\int\log\left(\frac{\mathop{}\\!\mathrm{d}\pi}{\mathop{}\\!\mathrm{d}\alpha\otimes\beta}\right)\mathop{}\\!\mathrm{d}\pi$.
In this case, the entropic version of $\mathcal{Q}\textsc{OT}$ is simply

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\mathcal{Q}(\pi)+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta).$
(25)

Fused $\mathcal{R}\textrm{OT}$. In some cases, we may even have an inter-space
cost that can be used alongside the learned cost. This is the case in the
context of the fused Gromov-Wasserstein cost (Vayer et al., , 2020). Such a
cost $\tilde{c}:\mathcal{C}(\mathcal{X}\times\mathcal{Y})\to\mathbb{R}$, which
isn’t learned, can be naturally integrated as (with $\eta\geq 0$)

$\displaystyle\mathcal{R}OT_{\varepsilon}^{\tilde{c}}(\alpha,\beta)=\min_{c\in\mathcal{C}(\mathcal{X}\times\mathcal{Y})}\mathcal{W}_{c+\eta\tilde{c}}^{\varepsilon}(\alpha,\beta)+\mathcal{R}(c).$
(26)

##  3 Structured OT across spaces

###  3.1 Linear cost-regularized OT

Let $\mathcal{X}\subset\mathbb{R}^{d_{x}}$ and
$\mathcal{Y}\subset\mathbb{R}^{d_{y}}$. From this point on, we consider linear
costs parameterized by a matrix $\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}$:
${c_{\mathbf{M}}(x,y)=-\langle\mathbf{M}x,y\rangle}$. In this case we can
write (23) as

$\displaystyle\min_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}\end{subarray}}\int-\langle\mathbf{M}x,y\rangle\mathop{}\\!\mathrm{d}\pi+\mathcal{R}(\mathbf{M})+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta),$
(L-$\mathcal{R}\textsc{OT}_{\varepsilon}$)

where we directly define $\mathcal{R}$ on $\mathbb{R}^{d_{y}\times d_{x}}$ for
brevity.

######  Proposition 3.1.

The linear $\mathcal{R}\textsc{OT}$ problem
(L-$\mathcal{R}\textsc{OT}_{\varepsilon}$) is equivalent to the
$\mathcal{Q}\textsc{OT}$ problem (25) with

$\displaystyle\mathcal{Q}(\pi)=-\mathcal{R}^{*}\left(\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}\mathop{}\\!\mathrm{d}\pi(x,y)\right),$
(27)

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

Since for low values of $\varepsilon$,
(L-$\mathcal{R}\textsc{OT}_{\varepsilon}$) is a concave minimization problem,
solving it exactly is generally out of reach. However, we can approximate it
using a simple alternated minimization procedure on the objective.
Initializing at a $\pi_{0}\in\Pi(\alpha,\beta)$, let

$\displaystyle\begin{split}\mathbf{M}_{k+1}&=\operatorname*{\arg\\!\min}_{\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}}\int-\langle\mathbf{M}x,y\rangle\,\mathrm{d}\pi_{k}+\mathcal{R}(\mathbf{M})\\\
\pi_{k+1}&=\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}-\langle\mathbf{M}_{k+1}x,y\rangle\,\mathrm{d}\pi+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta)\end{split}$
(28)

######  Proposition 3.2.

Let $\mathcal{X}=\left\\{x_{i}\right\\}_{i=1}^{m}\subset\mathbb{R}^{d_{x}}$,
$\mathcal{Y}=\left\\{y_{i}\right\\}_{i=1}^{n}\subset\mathbb{R}^{d_{y}}$,
$\alpha=\sum_{i=1}^{m}\bm{\alpha}_{i}\delta_{x_{i}}$ and
$\beta=\sum_{j=1}^{n}\bm{\beta}_{j}\delta_{y_{j}}$. Let $\varepsilon>0$ and
$\mathcal{R}$ be strongly convex. Then, any limit point of
$\left(M_{k},\pi_{k}\right)$ defined in (28) is a stationary point of the
objective in (L-$\mathcal{R}\textsc{OT}_{\varepsilon}$).

In practice, the $\pi$-step can be carried out efficiently using Sinkhorn’s
algorithm (Cuturi, , 2013) with the cost $c_{\mathbf{M}_{k}}$. The
$\mathbf{M}$-step is in general trickier since it might involve a difficult
optimization problem depending on the regularization function $\mathcal{R}$.

###  3.2 Proximal $\mathcal{R}\textsc{OT}$

We consider regularizations of the form
${\mathcal{R}(\mathbf{M})=\frac{1}{2}\left\lVert\mathbf{M}\right\rVert^{2}_{F}+\lambda
g(\mathbf{M})}$. We aim to solve

$\displaystyle\min_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}\end{subarray}}-$
$\displaystyle\int\langle\mathbf{M}x,y\rangle\,\mathrm{d}\pi+\frac{\left\lVert\mathbf{M}\right\rVert_{F}^{2}}{2}+\lambda
g(\mathbf{M})$ (29)
$\displaystyle+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta),$ (30)

where $g:\mathbb{R}^{d_{y}\times d_{x}}\to\mathbb{R}$. Note that when
$\lambda=0$, per Prop. 2.5, (29) is equivalent to (GW-IP). The minimization
with respect to $\mathbf{M}$ in (29) is a well-studied problem in the
optimization and machine learning literature (Bauschke and Combettes, , 2011).

######  Lemma 3.3.

For any $\pi$, the solution of (29) in $\mathbf{M}$ is

$\displaystyle\mathbf{M}(\pi)\triangleq\mathrm{prox}_{\lambda
g}\left(\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}\,\mathrm{d}\pi\right),$
(31)

where
$\mathrm{prox}_{h}(x)\triangleq\operatorname*{\arg\\!\min}_{z\in\mathbb{R}^{d}}\frac{1}{2}{\left\lVert
x-z\right\rVert}^{2}+h(z)$.

The alternated minimization scheme (28) becomes:
$\displaystyle\mathbf{M}_{k+1}=\mathrm{prox}_{\lambda g}$
$\displaystyle\left(\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}\mathop{}\\!\mathrm{d}\pi_{k}\right)$
(Prox-$\mathcal{R}\textsc{OT}$)
$\displaystyle\pi_{k+1}=\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}$
$\displaystyle-\int_{\mathcal{X}\times\mathcal{Y}}\langle\mathbf{M}_{k+1}x,y\rangle\mathop{}\\!\mathrm{d}\pi_{k}$
(32) $\displaystyle+\varepsilon\textsc{KL}(\pi_{k}||\alpha\otimes\beta)$ (33)
As a result, the $\mathbf{M}$-step in Prox-$\mathcal{R}\textsc{OT}$ can be
efficiently implemented when the proximal operator of $g$ can be computed in
closed form.

Statistical benefits of sparsity. Since OT famously suffers from the curse of
dimensionality – the sample complexity for estimating OT maps is
$\mathcal{O}(n^{-\nicefrac{{1}}{{d}}})$ (Niles-Weed and Rigollet, , 2022;
Pooladian and Niles-Weed, , 2021) – the main factor we consider in choosing
the regularization, beyond computational efficiency, is using inductive biases
to reduce effective dimension of the OT problem.

$\ell_{1}$ and $\ell_{1,2}$ regularizations. Among the choices of $g$ where
the $\mathbf{M}$-step can be implemented efficiently are the $\ell_{1}$ and
and $\ell_{1,2}$ regularizations, where
$\left\lVert\mathbf{M}\right\rVert_{1}=\sum_{i=1}^{d_{x}}\sum_{j=1}^{d_{y}}\left\lvert\mathbf{M}_{ij}\right\rvert\;\text{and}\;\left\lVert\mathbf{M}\right\rVert_{1,2}=\sum_{i=1}^{d_{x}}\left\lVert\mathbf{M}_{:i}\right\rVert_{2}$.
Their proximal operators are given by (Chierchia et al., , 2023)

$\displaystyle\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{1}}\left(\mathbf{M}\right)_{ij}$
$\displaystyle=\mathrm{\
sign}(\mathbf{M}_{ij})\left(\left\lvert\mathbf{M}_{ij}\right\rvert-\lambda\right)_{+}$
(34)
$\displaystyle\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{1,2}}\left(\mathbf{M}\right)_{:i}$
$\displaystyle=\left(1-\frac{\lambda}{\max\left\\{\left\lVert\mathbf{M}_{:i}\right\rVert\,,\lambda\right\\}}\right)\mathbf{M}_{:i},$
(35)

where $(a)_{+}=\max(a,0)$ and $\mathbf{A}_{:i}$ denotes the $i$th column of
$\mathbf{A}$. Intuitively, the reason behind using $\ell_{1}$ regularization
is to select the features in the spaces (or point clouds) $\mathcal{X}$ and
$\mathcal{Y}$ that are the most helpful in aligning those spaces. Similarly,
using $\ell_{1,2}$ regularization, we can discard features in one of the
spaces alone.

This is best seen by examining the iterations of Prox-$\mathcal{R}\textsc{OT}$
in each of these cases: For $\ell_{1}$, the prox operator in the $\mathbf{M}$
iteration results in a matrix $\mathbf{M}_{k+1}$ where some entries
$\mathbf{M}_{k+1}(i,j)$ can be set to $0$, and as a result the update of the
OT plan $\pi_{k+1}$ is oblivious to the corresponding features $x_{i}$ and
$y_{j}$; similarly, for $\ell_{1,2}$ iterations, the prox sets some columns of
$\mathbf{M}_{:i}$ to $0$, and as a result the update of the OT plan doesn’t
use the corresponding dimensions $x_{i}$. In both cases,
Prox-$\mathcal{R}\textsc{OT}$ alternates between a feature selection step and
a linear OT step.

Rewriting (29) in its $\mathcal{Q}\textsc{OT}$ form, we show that these
regularizations result in problems that are closely related to the inner-
product GW problem (GW-IP).

######  Proposition 3.4.

Consider Problem (29) with $\varepsilon=0$.

• $\ell_{1,2}$-regularization: with $g(M)=\left\lVert M\right\rVert_{1,2}$,
(29) with ${\varepsilon=0}$ is equivalent to $\mathcal{Q}\textsc{OT}$ with

$\displaystyle\mathcal{Q}(\pi)=\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\left(\langle
x_{I_{\pi}},x_{I_{\pi}}^{\prime}\rangle-\langle
y,y^{\prime}\rangle\right)^{2}\mathop{}\\!\mathrm{d}\pi\mathop{}\\!\mathrm{d}\pi-\frac{\lambda^{2}}{2}|I_{\pi}|$
(36) $\displaystyle-\int_{\mathcal{X}}\langle
x_{I_{\pi}},x_{I_{\pi}}^{\prime}\rangle^{2}\mathop{}\\!\mathrm{d}\alpha+\lambda\sum_{i\in
I_{\pi}}\left\lVert\int_{\mathcal{X}\times\mathcal{Y}}x_{i}y\,\,\mathrm{d}\pi(x,y)\right\rVert,$
(37)

where
$I_{\pi}=\left\\{i\,:\;\left\lVert\int_{\mathcal{X}\times\mathcal{Y}}x_{i}y\mathop{}\\!\mathrm{d}\pi\right\rVert>\lambda\right\\}$.
Moreover, with $\left(\pi^{\star},M^{\star}\right)$ a solution to (29), (29)
is equivalent to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\left(\langle
x_{I_{\pi^{\star}}},x_{I_{\pi^{\star}}}^{\prime}\rangle-\langle
y,y^{\prime}\rangle\right)^{2}\mathop{}\\!\mathrm{d}\pi^{\star}\mathop{}\\!\mathrm{d}\pi.$
(38)

• $\ell_{1}$-regularization: With $g(M)=\left\lVert M\right\rVert_{1}$, (29)
with ${\varepsilon=0}$ is equivalent to $\mathcal{Q}\textsc{OT}$ with

$\displaystyle\mathcal{Q}(\pi)=\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\sum_{(i,j)\in
I_{\pi}}\left(x_{i}x_{i}^{\prime}-y_{j}y_{j}^{\prime}\right)^{2}\mathop{}\\!\mathrm{d}\pi\mathop{}\\!\mathrm{d}\pi-\frac{\lambda^{2}}{2}|I_{\pi}|$
(39) $\displaystyle-\sum_{(i,j)\in
I_{\pi}}\left(\int_{\mathcal{X}}x_{i}^{2}\mathop{}\\!\mathrm{d}\alpha+\int_{\mathcal{Y}}y_{j}^{2}\mathop{}\\!\mathrm{d}\beta-\lambda\left\lvert\int_{\mathcal{X}\times\mathcal{Y}}y_{j}x_{i}\,\mathrm{d}\pi(x,y)\right\rvert\right),$
(40)

where
$I_{\pi}=\left\\{(i,j)\,:\;\int_{\mathcal{X}\times\mathcal{Y}}x_{i}y_{j}\,\mathrm{d}\pi>\lambda\right\\}$.
Moreover, with $\left(\pi^{\star},M^{\star}\right)$ a solution to (29), (29)
is equivalent to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\sum_{(i,j)\in
I_{\pi^{\star}}}\left(x_{i}x_{i}^{\prime}-y_{j}y_{j}^{\prime}\right)^{2}\mathop{}\\!\mathrm{d}\pi^{\star}\mathop{}\\!\mathrm{d}\pi.$
(41)

Nuclear and Rank regularizations. Two other possible choices of
regularizations are the nuclear norm and rank regularizations. Let
$\mathbf{M}=\mathbf{U}^{\top}\bm{\Sigma}\mathbf{V}$ be an SVD of $\mathbf{M}$,
where $\bm{\Sigma}$ contains the vector of singular vectors in decreasing
order $\sigma=\left(\sigma_{i}\right)_{i\in[d_{x}]}$ of $\mathbf{M}$. The
nuclear norm of $\mathbf{M}$ is defined as
$\left\lVert\mathbf{M}\right\rVert_{*}=\left\lVert\sigma\right\rVert_{1}$. The
proximal operator of $\left\lVert\cdot\right\rVert_{*}$ and $\mathrm{rk\
}(\cdot)$ are respectively

$\displaystyle\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{*}}(\mathbf{M})$
$\displaystyle=\mathbf{U}^{\top}\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{1}}\left(\sigma\right)\mathbf{V}$
(42) $\displaystyle\mathrm{prox}_{\lambda\mathrm{rk\ }}(\mathbf{M})$
$\displaystyle=\mathbf{U}^{\top}\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{0}}\left(\sigma\right)\mathbf{V},$
(43)

where
$\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{0}}(\sigma)_{i}=\sigma_{i}\mathbbm{1}_{\left\\{\sigma_{i}^{2}>2\lambda\right\\}}$.
Consider a solution $(\mathbf{M}^{\star},\pi^{\star})$ of (29) with the
nuclear norm. Then
$\mathbf{M}^{\star}=\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{*}}\left(\int
yx^{\top}\mathop{}\\!\mathrm{d}\pi^{\star}\right)$. Let $\lambda$ be large
enough so that $\mathrm{rk\ }(\mathbf{M}^{\star})=r>0$. Denote by
${\tilde{\mathbf{U}}=\mathbf{U}_{:,:r}}$ (resp.
${\tilde{\mathbf{V}}=\tilde{\sigma}\odot\mathbf{V}_{:,:r}}$) the restriction
of $\mathbf{U}$ (resp. $\mathbf{V}$) to its first $r$ lines (resp. columns
multiplied elementwise with $\tilde{\sigma}$, the restriction of $\sigma$ to
its first $r$ elements). Then, we can write the cost of the linearized problem
as
$c_{\mathbf{M}^{\star}}(x,y)=-\langle\tilde{\mathbf{V}}x,\tilde{\mathbf{U}}y\rangle$.
$\tilde{\mathbf{V}}x,\tilde{\mathbf{U}}y\in\mathbb{R}^{r}$, thus the nuclear
norm allows to similarly reduce the effective dimension of the problem from
potential large $d_{x},d_{y}$ to a small $r$ determined by the magnitude of
the regularization $\lambda$.

Note that since computing the SVD at each iteration of
Prox-$\mathcal{R}\textsc{OT}$ can be costly when $d_{x}$ and $d_{y}$ are
large, we suggest in extremely high-dimensional problems to directly
parametrize $\mathbf{M}$ as a low-rank matrix
$\mathbf{M}=\mathbf{M}_{2}^{\top}\mathbf{M}_{1}$, where
$r\ll\min(d_{x},d_{y})$, and use the alternated minimization scheme (28). We
write the corresponding iterations in the Appendix. Despite making the problem
non-convex in $\mathbf{M}$, we see in our applications that the resulting
method works well.

##  4 Entropic Monge maps for $\mathcal{R}\textsc{OT}$

Optimal transport (Monge, , 1781) seeks a map ${T:\mathcal{X}\to\mathcal{Y}}$
that minimizes the average displacement cost $c(x,T(x))$ between two measures
$\alpha$ and $\beta$, i.e. finding an OT coupling of the form
$\pi^{\star}=\left(\textrm{id},T\right)_{\sharp}\alpha$. Although such a (so-
called) Monge map does not exist for all costs, Brenier,  (1991) showed that
in the case where $c(x,y)=\tfrac{1}{2}{\left\lVert x-y\right\rVert}^{2}$, if
$\alpha$ has a density, the optimal map exists, is unique, and can be written
as the gradient of a convex function. A convenient approach to approximate
that Monge map can be found in entropic regularization. Pooladian and Niles-
Weed,  (2021) showed that one can build a map $T_{\varepsilon}$ using entropy
regularized transport (with regularization strength $\varepsilon>0$) such that
$T_{\varepsilon}\xrightarrow[]{\varepsilon\to 0}T$. We show that there always
exists a Monge map for linear $\mathcal{R}\textsc{OT}$, and propose a
formulation for an entropic map for this problem.

###  4.1 Monge maps for linear $\mathcal{R}\textsc{OT}$

Dumont et al.,  (2022) recently showed that there exists a Monge map for the
GW-IP problem, which, as seen in Prop. 2.5 is a special case of linear
$\mathcal{R}\textsc{OT}$. Here, we extend their reasoning to show the
existence of Monge maps for the general linear $\mathcal{R}\textsc{OT}$
problem.

######  Proposition 4.1.

Let $\alpha\in\mathcal{P}(\mathbb{R}^{d_{x}})$ and
$\beta\in\mathcal{P}(\mathbb{R}^{d_{y}})$ with compact supports and $d_{x}\geq
d_{y}$. Assume that $\alpha\ll\mathcal{L}^{d_{x}}$, the Lebesgue measure. Then
there exists a map $T:\mathcal{X}\to\mathcal{Y}$ such that
$\pi^{\star}=\left(\textrm{id},T\right)_{\sharp}\alpha$ and
$\left(\pi^{\star},\mathbf{M}(\pi^{\star})\right)$ is optimal for
(L-$\mathcal{R}\textsc{OT}_{\varepsilon}$) with $\varepsilon=0$.

###### Proof.

Let $(\mathbf{M}^{\star},\pi^{\star})$ be optimal for
(L-$\mathcal{R}\textsc{OT}_{\varepsilon}$) with $\varepsilon=0$. Then
$\pi^{\star}$ is a solution to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}\times\mathcal{Y}}-\langle\mathbf{M}^{\star}x,y\rangle\mathop{}\\!\mathrm{d}\pi(x,y),$
(44)

and (Dumont et al., , 2022, Theorem 4) showed that for all
$\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}$, under the assumptions of Prop.
4.1, there exists a Monge map for the cost
${c_{\mathbf{M}}(x,y)=-\langle\mathbf{M}x,y\rangle}$ between $\alpha$ and
$\beta$. ∎

Now that the existence of a Monge map has been established, a natural question
is whether, and how, we can approximate it using entropy-regularization. In
the following sections, we define an entropic map for
L-$\mathcal{R}\textsc{OT}_{\varepsilon}$ and show its convergence to a Monge
map for this problem under suitable assumptions. Then, we discuss its
convergence for the sparse and low-rank regularizations we considered in
Section 3.2.

###  4.2 Entropic Monge maps for linear $\mathcal{R}\textsc{OT}$

######  Definition 4.2 (Entropic map for
L-$\mathcal{R}\textsc{OT}_{\varepsilon}$).

Let $\left(\pi_{\varepsilon}^{\star},M_{\varepsilon}^{\star}\right)$ be a
solution of L-$\mathcal{R}\textsc{OT}_{\varepsilon}$. Let
$\varepsilon^{\prime}>0$. Define

$\displaystyle T_{\varepsilon,\varepsilon^{\prime}}(x)=\frac{\int
y\exp{\left(\left(g_{\varepsilon,\varepsilon^{\prime}}(y)+\langle
M_{\varepsilon}^{\star}x,y\rangle\right)/\varepsilon^{\prime}\right)}d\beta(y)}{\int\exp{\left(\left(g_{\varepsilon,\varepsilon^{\star}}(y)+\langle
M_{\varepsilon}^{\star}x,y\rangle\right)/\varepsilon^{\prime}\right)}d\beta(y)}.$
(45)

Here,
$(f_{\varepsilon,\varepsilon^{\prime}},g_{\varepsilon,\varepsilon^{\prime}})$
are Sinkhorn potentials for the inner product cost between
${M_{\varepsilon}^{\star}}_{\sharp}\alpha$ and $\beta$ with an
$\varepsilon^{\prime}$ entropic regularization: They are a solution of the
problem

$\displaystyle\max_{(f,g)\in\mathcal{C}(X)\times\mathcal{C}(Y)}\int_{X}f(x)d{\mathbf{M}_{\varepsilon}^{\star}}_{\sharp}\alpha(x)+\int_{Y}g(y)d\beta(y)$
(46) $\displaystyle-\varepsilon^{\prime}\int_{X\times
Y}\exp\left(\frac{f(x)+g(y)+\langle
x,y\rangle}{\varepsilon^{\prime}}\right)\mathop{}\\!\mathrm{d}{\mathbf{M}_{\varepsilon}^{\star}}_{\sharp}\alpha\otimes\beta(x,y).$
(47)

####  4.2.1 Computing the entropic map in practice

An important point in the theoretical definition of the entropic map (45) is
to use two different regularization parameters
$\varepsilon,\varepsilon^{\prime}>0$. In practice (as we do in our
applications), we can simplify the definition and use a single $\varepsilon$.
Given samples $\left\\{x_{i}\right\\}_{i=1}^{n}\sim\alpha$ and
$\left\\{y_{i}\right\\}_{i=1}^{n}\sim\beta$:
• Find a primal solution $(\bm{\pi}^{\star},\mathbf{M}_{\varepsilon}^{\star})$
and a dual solution $\left(\mathbf{f},\mathbf{g}\right)$ of the discrete
L-$\mathcal{R}\textsc{OT}_{\varepsilon}$ problem with
$\hat{\alpha}=\sum_{i=1}^{n}\bm{\alpha}_{i}\delta_{x_{i}}$ and
$\hat{\beta}=\sum_{j=1}^{m}\bm{\beta}_{j}\delta_{y_{j}}$ using (28).
• Define the entropic map as

$\displaystyle
T_{\varepsilon}(x)=\frac{\sum_{j=1}^{n}y_{j}\exp{\left(\left(\mathbf{g}_{j}+\langle\mathbf{M}_{\varepsilon}^{\star}x,y_{j}\rangle\right)/\varepsilon\right)}}{\sum_{k=1}^{n}\exp{\left(\left(\mathbf{g}_{k}+\langle\mathbf{M}_{\varepsilon}^{\star}x,y_{k}\rangle\right)/\varepsilon\right)}}.$
(48)

Figure 1: Illustration of the entropic map (48) for GW-IP from a 3D to a 2D
point cloud with $\varepsilon\approx 0$.

Illustration. In Figure 1, we showcase the action of the entropic map (48)
when $\varepsilon\approx 0$. We consider two point clouds $\alpha$ and $\beta$
on $\mathbb{R}^{3}$ and $\mathbb{R}^{2}$. The entropic map implicitly acts in
two steps: it first transforms the 3D point cloud $\alpha$ into a 2D point
cloud $\mathbf{M}_{\sharp}\alpha$ which is aligned with $\beta$. It then maps
the points from $\mathbf{M}_{\sharp}\alpha$ to $\beta$ using the inner-product
entropic map.

####  4.2.2 Convergence theory

The case where $\mathbf{M}^{\star}$ is full rank. Assume that $d_{x}\geq
d_{y}$ and consider the case where $rk(\mathbf{M}^{\star})=d_{y}$. Dumont et
al.,  (2022) and Vayer,  (2020) showed that if $\mathbf{M}^{\star}$ is full
rank, then a Monge map for GW-IP is given by $T=-\nabla
f\circ\mathbf{M}^{\star}$, where $f$ is a Kantorovitch potential for
$\mathcal{W}^{\varepsilon}_{\langle\cdot,\cdot\rangle}\left(M_{\sharp}\alpha,\beta\right)$.
The entropic map for L-$\mathcal{R}\textsc{OT}_{\varepsilon}$ verifies an
analog result.

######  Lemma 4.3.

Let
$(f_{\varepsilon,\varepsilon^{\prime}},g_{\varepsilon,\varepsilon^{\prime}})$
be Sinkhorn potentials for
$\mathcal{W}^{\varepsilon^{\prime}}_{\langle\cdot,\cdot\rangle}({\mathbf{M}_{\varepsilon}^{\star}}_{\sharp}\alpha,\beta)$.
Then, we can rewrite (45) as

$\displaystyle T_{\varepsilon,\varepsilon^{\prime}}=-\nabla
f_{\varepsilon,\varepsilon^{\prime}}\circ\mathbf{M}_{\varepsilon}^{\star}$
(49)

Using the previous lemma and results from Pooladian and Niles-Weed,  (2021);
Zhang et al.,  (2022), we have the following convergence result.

######  Proposition 4.4.

Let $\alpha,\beta\in\mathcal{P}(\mathcal{X})\times\mathcal{P}(\mathcal{Y})$ be
two measures with compact supports. Assume that they are regular enough
(namely that they verify assumptions (A1-3) from Pooladian and Niles-Weed,
(2021)). Let $\varepsilon,\varepsilon^{\prime}>0$ and consider the map
$T_{\varepsilon,\varepsilon^{\prime}}$ defined in (45). Then, we have that
$T_{\varepsilon,\varepsilon^{\prime}}\xrightarrow{L^{2}(\alpha)}T_{\varepsilon,0}$
as $\varepsilon^{\prime}\to 0$, where $T_{\varepsilon,0}$ is a Monge map for
$\mathcal{W}^{0}_{c_{\mathbf{M}_{\varepsilon}^{\star}}}\left(\alpha,\beta\right)$,
and $c_{\mathbf{M}_{\varepsilon}^{\star}}(x,y)=-\langle
M_{\varepsilon}^{\star}x,y\rangle$. Moreover, we have along a subsequence that
$T_{\varepsilon,0}\xrightarrow{L^{2}(\alpha)}T_{0,0}$ as $\varepsilon\to 0$,
where $T_{0,0}$ is a Monge map for L-$\mathcal{R}\textsc{OT}_{\varepsilon}$.
Hence, along a subsequence,

$\displaystyle\lim_{\varepsilon\to 0}\lim_{\varepsilon^{\prime}\to
0}T_{\varepsilon,\varepsilon^{\prime}}=T_{0,0}\quad\text{in}\;L^{2}(\alpha).$
(50)

The case where $\mathbf{M}^{\star}$ is not full rank. The case where
$\mathbf{M}^{\star}$ is not full rank is more involved and with limited use in
practice, as the theory in Dumont et al.,  (2022) requires solving an OT
problem between conditional probabilities, which are not accessible in
practice. In the appendix, we simulate problems where we explicitly constrain
$\mathbf{M}^{\star}$ to being rank-deficient. We test on Gaussians (for which
ground-truth OT maps are known (Salmona et al., , 2022)) and show that rank-
defficiency has no effect on the convergence of the entropic map.

Sparsifying transforms and entropic Monge maps. As we saw in the previous
section, we can show the convergence of the entropic map for
L-$\mathcal{R}\textsc{OT}_{\varepsilon}$ when $\mathbf{M}^{\star}$ is full
rank. With sparsifying and low-rank norms, it is unlikely that such a matrix
is full rank. However, when using $\ell_{1,2}$ regularization, our goal is to
operate feature selection on $\mathcal{X}$, so that we don’t actually care
about the dimensions on which the matrix $\mathbf{M}_{\varepsilon}^{\star}$ is
$0$ (which correspond to features we want to discard). Thus we can restrict
our study to the existence of Monge maps between ${P_{I}}_{\sharp}\alpha$ and
$\beta$ for the cost
$c_{\mathbf{M}^{\star}_{I}}(x,y)=-\langle\mathbf{M}_{I}^{\star}x,y\rangle$,
where $I=\left\\{i\in[d_{x}]:\left\lVert\mathbf{M}_{;i}\right\rVert\neq
0\right\\}$, $P_{I}$ is the projection operator on the dimensions $I$, and
$\mathbf{M}_{I}$ the restriction of $\mathbf{M}$ to the columns indexed by
$I$. Note, though, that this matrix could still be rank-deficient.

Similarly, nuclear norm regularization inherently makes $\mathbf{M}^{\star}$
low-rank. As we saw earlier it results in a linear OT problem with an inner
product cost between $\tilde{\mathbf{U}}_{\sharp}\alpha$ and
$\tilde{\mathbf{V}}_{\sharp}\beta$ (i.e.
$\mathcal{W}_{\langle\cdot,\cdot\rangle}(\tilde{U}_{\sharp}\alpha,\tilde{V}_{\sharp}\beta)$).
If we are interested in low-dimensional representations of $\alpha$ and
$\beta$, there always exists a Monge map between
$\tilde{\mathbf{U}}_{\sharp}\alpha$ and $\tilde{\mathbf{V}}_{\sharp}\beta$,
and the entropic map converges to the Monge map since it is associated with
the inner-product cost on $\mathbb{R}^{r}$, where $r=\mathrm{rk\
}(\mathbf{M}^{\star})$ (Pooladian and Niles-Weed, , 2021).

##  5 Applications

In all of the applications below, we use our proximal approach to compute
transforms, jointly with the entropic map (48) to displace points from one
space to the other in and/or out of sample. That map’s $\varepsilon$
regularization is selected using cross-validation on the training set using a
grid in [5e-3, 1e-3, 5e-4, 1e-4, 5e-5, 1e-5].

###  5.1 Multi-omics data Integration: Sparse and Low-Rank Transforms

Sparse Transforms. We consider the scGM dataset (Cheow et al., , 2016)
containing the gene expression and DNA methylation modalities for human
somatic cells (Welch et al., , 2017). The goal is to match the $177$ samples
of the dataset across modalities using an entropic map from the gene
expression ($d_{x}=34$) to the DNA methylation domain ($d_{y}=27$).
Performance is measured in terms of Label Transfer Accuracy, a common way to
evaluate single-cell data integration tasks (Demetci et al., , 2022) (see
Appendix).

Figure 2: Left: Convergence of Prox-$\mathcal{R}\textsc{OT}$ with different
sparsifying regularizations $\mathcal{R}$ on the scGM dataset. X-axis: #
updates of the OT plan. Y-axis: LTA (higher is better). $\lambda=0$
corresponds to standard GW-IP. $\varepsilon=0.05$. Right:
$\mathbf{M}_{\varepsilon}^{\star}$. % zeros: Top left $\lambda=0:6\%$, right
$\ell_{1,2}(\lambda=7.5):67\%$; Bottom left $\ell_{1}(\lambda=1.5):71\%$,
right $\ell_{1}(\lambda=2):84\%$.

We compare solving GW-IP using nested Sinkhorn iterations (Peyré et al., ,
2016) to Prox-$\mathcal{R}\textsc{OT}$ with the $\ell_{1}$ and $\ell_{1,2}$
regularizations in (29). Adding a sparsity prior to the map $\mathbf{M}$
results in a better performance (Figure 2). Using an $\ell_{1}$ or
$\ell_{1,2}$ regularization also outperforms all the methods in (Demetci et
al., , 2022, Figure 2).

Low-Rank Transforms. We consider the Neurips 2021 multimodal single-cell
integration dataset (Lance et al., , 2022). We use the Site 1/Donor 1 stored
in the MOSCOT package (Klein et al., 2023a, ). Like scGM, the task consists in
mapping $6,224$ cells from the chromatin accessibility $(d_{y}=8,000)$ to the
gene expression domain $(d_{x}=2,000)$.

We consider random subsets of the data of sizes $25,50,100,250,500,1000,6224$.
On each subset, we run Prox-$\mathcal{R}\textsc{OT}$ with rank constraint
between 5 and 12. We display for each subset the best-performing rank on
average in Figure 3 (that rank is always 11 or less). Performance is measured
in terms of FOSCTTM (fraction of samples closer than the true match) depending
on the subset size. Compared to GW, low-rank Prox-$\mathcal{R}\textsc{OT}$
handles settings where $n\ll d_{y}$ much better than GW. Note that
Prox-$\mathcal{R}\textsc{OT}$-$\mathrm{rk\ }$ still has benefits in the high
$n$ regime since the dimension of the OT problem solved at each iteration is
at most $12$ vs. $2,000$ for GW.

###  5.2 Spatial Transcriptomics: Stochastic Fused GW-IP

An important benefit of the cost-regularized formulation of GW-IP is that it
lends itself ideally to stochastic optimization. Indeed, let
$\alpha=\sum_{i=1}^{m}\alpha_{i}\delta_{x_{i}}$ and
$\beta=\sum_{j=1}^{n}\beta_{j}\delta_{y_{j}}$ be two discrtete distributions.
Then, using entropic regularization, we can rewrite GW-
IP-$\mathcal{R}\textsc{OT}$ using the dual formulation of linear OT as

$\displaystyle\min_{\begin{subarray}{c}\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}\\\
(f,g)\in\mathbb{R}^{m}\times\mathbb{R}^{n}\end{subarray}}\sum_{i=1}^{m}f_{i}\alpha_{i}+\sum_{j=1}^{n}g_{j}\beta_{j}+\frac{1}{8}\left\lVert\mathbf{M}\right\rVert_{F}^{2}$
(51)
$\displaystyle-\varepsilon\sum_{i,j}\exp\left(\frac{f_{i}+g_{j}+\langle\mathbf{M}x_{i},y_{j}\rangle+\eta\tilde{c}\left(\tilde{x}_{i},\tilde{y}_{j}\right)}{\varepsilon}\right)\alpha_{i}\beta_{j}.$
(52)

where $\eta>0$, and as in (26), we can use an additional known inter-space
cost $\tilde{c}$, which corresponds to using Fused GW-IP. We can solve this
problem using (epoch) stochastic gradient descent ascent.

Figure 3: Average performance (FOSCTTM - lower is better) across sampled
subsets of GW vs. Prox-$\mathcal{R}\textsc{OT}$ depending on the size of
training data. Left: Training performance on each subset size. Right:
Evaluation on the full dataset of the map computed using a subset.

Alg. |  val $\rho$ |  test $\rho$ | F1 macro | F1 micro | F1 weighted
---|---|---|---|---|---
LFGW | 0.365 | 0.443 | 0.576 | 0.720 | 0.714
ULFGW | 0.379 | 0.463 | 0.582 | 0.733 | 0.724
SFGW-IP | 0.335 | 0.436 | 0.609 | 0.768 | 0.765

Table 1: Performance of different GW variants on a spatial transciptomics
task. Details for LFGW and ULFGW can be found in Scetbon et al.,  (2023).
SFGW-IP refers to solving (51) via epoch SGDA.

We reproduce the experimental setting of Scetbon et al.,  (2023). The goal is
to align cells from two coronal sections of a mouse brain (Shi et al., ,
2023). The linear OT term in the fused formulation comes from a 30d PCA
computed in gene expression space. We use validation Pearson correlation for
hyperparameter selection. We show in Table 1 that the stochastic approach,
while having a lower per-iteration cost, is competitive with all GW variants,
and significantly exceeds them in F1 scores (see appendix for more details).

Conclusion. We have leveraged a cost-regularized OT perspective on GW problems
to propose new methods for OT across spaces that can induce structure (e.g.
sparsity) on the transformation from one space to another. Our injection of
sparsity into the transform $\mathbf{M}$ is unrelated to the sparsity observed
naturally for couplings when solving the original OT problem, from LP duality
or from other regularizations, as in e.g.(Blondel et al., , 2018; Liu et al.,
, 2022), nor is it related to that obtained for same-space Monge displacements
by the MBO estimator (Cuturi et al., , 2023). Our use of low-rank constraints
for transform $\mathbf{M}$ is also unrelated to the low rank constraint on
couplings introduced by (Scetbon et al., , 2021, 2022). The ability to add
structural assumptions on across-space transforms opens up new perspectives to
use transport across high-dimensional modalities. We also showed the existence
of Monge maps for our formulation, and demonstrated their applicability with
entropic maps. We used a proximal alternated minimization algorithm together
with structure-inducing regularizations, demonstrating applicability to
single-cell multiomics data matching and spatial transcriptomics tasks.

## References

  * Alvarez-Melis et al., (2019) Alvarez-Melis, D., Jegelka, S., and Jaakkola, T. S. (2019).  Towards optimal transport with global invariances.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1870–1879. PMLR.
  * Bauschke and Combettes, (2011) Bauschke, H. H. and Combettes, P. L. (2011).  Convex analysis and monotone operator theory in hilbert spaces.  CMS Books in Mathematics.
  * Bertsekas, (1997) Bertsekas, D. P. (1997).  Nonlinear programming.  Journal of the Operational Research Society, 48(3):334–334.
  * Blondel et al., (2018) Blondel, M., Seguy, V., and Rolet, A. (2018).  Smooth and sparse optimal transport.  In Storkey, A. and Perez-Cruz, F., editors, Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pages 880–889. PMLR.
  * Bonneel et al., (2015) Bonneel, N., Rabin, J., Peyré, G., and Pfister, H. (2015).  Sliced and radon wasserstein barycenters of measures.  Journal of Mathematical Imaging and Vision, 51:22–45.
  * Brenier, (1991) Brenier, Y. (1991).  Polar factorization and monotone rearrangement of vector-valued functions.  Communications on pure and applied mathematics, 44(4):375–417.
  * Bunne et al., (2023) Bunne, C., Stark, S. G., Gut, G., Del Castillo, J. S., Levesque, M., Lehmann, K.-V., Pelkmans, L., Krause, A., and Rätsch, G. (2023).  Learning single-cell perturbation responses using neural optimal transport.  Nature Methods, pages 1–10.
  * Cheow et al., (2016) Cheow, L. F., Courtois, E. T., Tan, Y., Viswanathan, R., Xing, Q., Tan, R. Z., Tan, D. S., Robson, P., Loh, Y.-H., Quake, S. R., et al. (2016).  Single-cell multimodal profiling reveals cellular epigenetic heterogeneity.  Nature methods, 13(10):833–836.
  * Chierchia et al., (2023) Chierchia, G., Chouzenoux, E., Combettes, P. L., and Pesquet, J.-C. (2023).  The proximity operator repository. user’s guide.
  * Chizat et al., (2020) Chizat, L., Roussillon, P., Léger, F., Vialard, F.-X., and Peyré, G. (2020).  Faster wasserstein distance estimation with the sinkhorn divergence.  Advances in Neural Information Processing Systems, 33:2257–2269.
  * Cohen and Guibas, (1999) Cohen, S. and Guibas, L. (1999).  The earth mover’s distance under transformation sets.  In Proceedings of the Seventh IEEE International Conference on Computer Vision, volume 2, pages 1076–1083. IEEE.
  * Courty et al., (2017) Courty, N., Flamary, R., Habrard, A., and Rakotomamonjy, A. (2017).  Joint distribution optimal transportation for domain adaptation.  Advances in neural information processing systems, 30.
  * Cuturi, (2013) Cuturi, M. (2013).  Sinkhorn distances: Lightspeed computation of optimal transport.  Advances in neural information processing systems, 26:2292–2300.
  * Cuturi et al., (2023) Cuturi, M., Klein, M., and Ablin, P. (2023).  Monge, Bregman and occam: Interpretable optimal transport in high-dimensions with feature-sparse maps.  In Proceedings of the 40th International Conference on Machine Learning, volume 202, pages 6671–6682. PMLR.
  * De Philippis, (2013) De Philippis, G. (2013).  Regularity of optimal transport maps and applications, volume 17.  Springer Science & Business Media.
  * Demetci et al., (2022) Demetci, P., Santorella, R., Sandstede, B., Noble, W. S., and Singh, R. (2022).  Scot: single-cell multi-omics alignment with optimal transport.  Journal of Computational Biology, 29(1):3–18.
  * Deshpande et al., (2019) Deshpande, I., Hu, Y.-T., Sun, R., Pyrros, A., Siddiqui, N., Koyejo, S., Zhao, Z., Forsyth, D., and Schwing, A. G. (2019).  Max-sliced wasserstein distance and its use for gans.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10648–10656.
  * Dudley, (1966) Dudley, R. M. (1966).  Weak convergence of probabilities on nonseparable metric spaces and empirical measures on euclidean spaces.  Illinois Journal of Mathematics, 10(1):109–126.
  * Dumont et al., (2022) Dumont, T., Lacombe, T., and Vialard, F.-X. (2022).  On the existence of monge maps for the gromov-wasserstein distance.  arXiv preprint arXiv:2210.11945.
  * Genevay et al., (2019) Genevay, A., Chizat, L., Bach, F., Cuturi, M., and Peyré, G. (2019).  Sample complexity of sinkhorn divergences.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1574–1583.
  * Genevay et al., (2018) Genevay, A., Peyré, G., and Cuturi, M. (2018).  Learning generative models with sinkhorn divergences.  In International Conference on Artificial Intelligence and Statistics, AISTATS 2018, pages 1608–1617.
  * Grave et al., (2019) Grave, E., Joulin, A., and Berthet, Q. (2019).  Unsupervised alignment of embeddings with wasserstein procrustes.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1880–1890. PMLR.
  * Janati et al., (2020) Janati, H., Bazeille, T., Thirion, B., Cuturi, M., and Gramfort, A. (2020).  Multi-subject meg/eeg source imaging with sparse multi-task regression.  NeuroImage, 220:116847.
  * (24) Klein, D., Palla, G., Lange, M., Klein, M., Piran, Z., Gander, M., Meng-Papaxanthos, L., Sterr, M., Bastidas-Ponce, A., Tarquis-Medina, M., et al. (2023a).  Mapping cells through time and space with moscot.  bioRxiv, pages 2023–05.
  * (25) Klein, M., Pooladian, A.-A., Ablin, P., Ndiaye, E., Niles-Weed, J., and Cuturi, M. (2023b).  Learning costs for structured monge displacements.  arXiv preprint arXiv:2306.11895.
  * Lance et al., (2022) Lance, C., Luecken, M. D., Burkhardt, D. B., Cannoodt, R., Rautenstrauch, P., Laddach, A., Ubingazhibov, A., Cao, Z.-J., Deng, K., Khan, S., et al. (2022).  Multimodal single cell data integration challenge: Results and lessons learned.  In NeurIPS 2021 Competitions and Demonstrations Track, pages 162–176. PMLR.
  * Le et al., (2022) Le, K., Le, D. Q., Nguyen, H., Do, D., Pham, T., and Ho, N. (2022).  Entropic gromov-wasserstein between Gaussian distributions.  In International Conference on Machine Learning, pages 12164–12203. PMLR.
  * Le et al., (2019) Le, T., Yamada, M., Fukumizu, K., and Cuturi, M. (2019).  Tree-sliced variants of wasserstein distances.  Advances in neural information processing systems, 32.
  * Lipman et al., (2022) Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. (2022).  Flow matching for generative modeling.  In The Eleventh International Conference on Learning Representations.
  * Liu et al., (2022) Liu, T., Puigcerver, J., and Blondel, M. (2022).  Sparsity-constrained optimal transport.  In The Eleventh International Conference on Learning Representations.
  * Mémoli, (2011) Mémoli, F. (2011).  Gromov–wasserstein distances and the metric approach to object matching.  Foundations of computational mathematics, 11(4):417–487.
  * Monge, (1781) Monge, G. (1781).  Mémoire sur la théorie des déblais et des remblais.  Mem. Math. Phys. Acad. Royale Sci., pages 666–704.
  * Montavon et al., (2016) Montavon, G., Müller, K.-R., and Cuturi, M. (2016).  Wasserstein training of restricted boltzmann machines.  Advances in Neural Information Processing Systems, 29.
  * (34) Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. (2023a).  Action matching: Learning stochastic dynamics from samples.  In International conference on machine learning.
  * (35) Neklyudov, K., Brekelmans, R., Severo, D., and Makhzani, A. (2023b).  Action matching: Learning stochastic dynamics from samples.  In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 25858–25889. PMLR.
  * Nekrashevich et al., (2023) Nekrashevich, M., Korotin, A., and Burnaev, E. (2023).  Neural gromov-wasserstein optimal transport.  arXiv preprint arXiv:2303.05978.
  * Niles-Weed and Rigollet, (2022) Niles-Weed, J. and Rigollet, P. (2022).  Estimation of wasserstein distances in the spiked transport model.  Bernoulli, 28(4):2663–2688.
  * Paty and Cuturi, (2019) Paty, F.-P. and Cuturi, M. (2019).  Subspace robust wasserstein distances.  In International conference on machine learning, pages 5072–5081. PMLR.
  * Paty and Cuturi, (2020) Paty, F.-P. and Cuturi, M. (2020).  Regularized optimal transport is ground cost adversarial.  In International Conference on Machine Learning, pages 7532–7542. PMLR.
  * Peyré et al., (2016) Peyré, G., Cuturi, M., and Solomon, J. (2016).  Gromov-wasserstein averaging of kernel and distance matrices.  In International conference on machine learning, pages 2664–2672. PMLR.
  * Pooladian and Niles-Weed, (2021) Pooladian, A.-A. and Niles-Weed, J. (2021).  Entropic estimation of optimal transport maps.  arXiv preprint arXiv:2109.12004.
  * Rioux et al., (2023) Rioux, G., Goldfeld, Z., and Kato, K. (2023).  Entropic gromov-wasserstein distances: Stability, algorithms, and distributional limits.  arXiv preprint arXiv:2306.00182.
  * Salimans et al., (2018) Salimans, T., Zhang, H., Radford, A., and Metaxas, D. (2018).  Improving gans using optimal transport.  In International Conference on Learning Representations.
  * Salmona et al., (2022) Salmona, A., Delon, J., and Desolneux, A. (2022).  Gromov-wasserstein distances between Gaussian distributions.  Journal of Applied Probability, 59(4).
  * Santambrogio, (2015) Santambrogio, F. (2015).  Optimal transport for applied mathematicians.  Birkäuser, NY, 55(58-63):94.
  * Scetbon et al., (2021) Scetbon, M., Cuturi, M., and Peyré, G. (2021).  Low-rank sinkhorn factorization.  In International Conference on Machine Learning, pages 9344–9354. PMLR.
  * Scetbon et al., (2023) Scetbon, M., Klein, M., Palla, G., and Cuturi, M. (2023).  Unbalanced low-rank optimal transport solvers.  arXiv preprint arXiv:2305.19727.
  * Scetbon et al., (2022) Scetbon, M., Peyré, G., and Cuturi, M. (2022).  Linear-time gromov Wasserstein distances using low rank couplings and costs.  In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S., editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 19347–19365. PMLR.
  * 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. (2019).  Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming.  Cell, 176(4):928–943.
  * Shi et al., (2023) Shi, H., He, Y., Zhou, Y., Huang, J., Maher, K., Wang, B., Tang, Z., Luo, S., Tan, P., Wu, M., et al. (2023).  Spatial atlas of the mouse central nervous system at molecular resolution.  Nature, pages 1–10.
  * Sinha et al., (2018) Sinha, A., Namkoong, H., and Duchi, J. (2018).  Certifying some distributional robustness with principled adversarial training.  In International Conference on Learning Representations.
  * Sturm, (2006) Sturm, K.-T. (2006).  On the geometry of metric measure spaces.  Acta Mathematica, 196(1):65–131.
  * Tong et al., (2020) Tong, A., Huang, J., Wolf, G., Van Dijk, D., and Krishnaswamy, S. (2020).  Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics.  In International conference on machine learning, pages 9526–9536. PMLR.
  * Tong et al., (2023) Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., FATRAS, K., Wolf, G., and Bengio, Y. (2023).  Improving and generalizing flow-based generative models with minibatch optimal transport.  In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems.
  * Vaswani et al., (2019) Vaswani, S., Mishkin, A., Laradji, I., Schmidt, M., Gidel, G., and Lacoste-Julien, S. (2019).  Painless stochastic gradient: Interpolation, line-search, and convergence rates.  Advances in neural information processing systems, 32:3732–3745.
  * Vayer, (2020) Vayer, T. (2020).  A contribution to optimal transport on incomparable spaces.  arXiv preprint arXiv:2011.04447.
  * Vayer et al., (2020) Vayer, T., Chapel, L., Flamary, R., Tavenard, R., and Courty, N. (2020).  Fused gromov-wasserstein distance for structured objects.  Algorithms, 13(9):212.
  * Villani, (2003) Villani, C. (2003).  Topics in optimal transportation.  Graduate Studies in Mathematics.
  * Weed and Bach, (2019) Weed, J. and Bach, F. (2019).  Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance.  Bernoulli, 25(4A).
  * Welch et al., (2017) Welch, J. D., Hartemink, A. J., and Prins, J. F. (2017).  Matcher: manifold alignment reveals correspondence between single cell transcriptome and epigenome dynamics.  Genome biology, 18(1):1–19.
  * Wong et al., (2019) Wong, E., Schmidt, F., and Kolter, Z. (2019).  Wasserstein adversarial examples via projected sinkhorn iterations.  In International Conference on Machine Learning, pages 6808–6817. PMLR.
  * Zhang et al., (2017) Zhang, M., Liu, Y., Luan, H., and Sun, M. (2017).  Earth mover’s distance minimization for unsupervised bilingual lexicon induction.  In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pages 1934–1945.
  * Zhang et al., (2022) Zhang, Z., Goldfeld, Z., Mroueh, Y., and Sriperumbudur, B. K. (2022).  Gromov-wasserstein distances: Entropic regularization, duality, and sample complexity.  arXiv preprint arXiv:2212.12848.

SUPPLEMENTARY MATERIAL
Structured Transforms Across Spaces
with Cost-Regularized Optimal Transport

##  Appendix A Appendix for Section 2

###  A.1 Proof of Proposition 2.4

Recall that the Wasserstein Procrustes problem (9) is

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\min_{\mathbf{C}\in\mathcal{O}_{d}}\int_{\mathcal{X}\times\mathcal{Y}}\left\lVert\mathbf{C}x-y\right\rVert^{2}\mathop{}\\!\mathrm{d}\pi(x,y).$
(53)

As noticed in e.g. Eq. 11 in Zhang et al.,  (2017), this formulation is an
extension of the orthogonal Procrustes problem. Given two matrices
$\mathbf{A},\mathbf{B}\in\mathbb{R}^{d\times d}$ The orthogonal Procrustes
problem aims to solve

$\displaystyle\min_{\mathbf{C}\in\mathcal{O}_{d}}\left\lVert\mathbf{C}A-B\right\rVert_{F}.$
(54)

The solution of this problem is
$\mathbf{C}^{\star}=\mathbf{U}\mathbf{V}^{\top}$, where
$\mathbf{B}\mathbf{A}^{\top}=\mathbf{U}\bm{\Sigma}\mathbf{V}^{\top}$. One can
follow a similar procedure to show that a solution to the inner problem in
(53) is given by $\mathbf{C}^{\top}=\mathbf{U}(\pi)\mathbf{V}(\pi)^{\top}$,
where $\int
yx^{\top}\mathop{}\\!\mathrm{d}\pi(x,y)=\mathbf{U}(\pi)\bm{\Sigma}(\pi)\mathbf{V}(\pi)^{\top}$
is an SVD. Replacing with this solution in the inner problem gives the desired
$\mathcal{Q}\textsc{OT}$ formulation:

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}\times\mathcal{Y}}\left\lVert\mathbf{U}(\pi)\mathbf{V}(\pi)^{\top}x-y\right\rVert^{2}\mathop{}\\!\mathrm{d}\pi(x,y).$
(55)

##  Appendix B Appendix for Section 3

###  B.1 Proof of Proposition 3.1

Starting from L-$\mathcal{R}\textsc{OT}_{\varepsilon}$, we have

$\displaystyle\min_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}\end{subarray}}\int-\langle\mathbf{M}x,y\rangle
d\pi+\mathcal{R}(\mathbf{M})+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta),$
(56)
$\displaystyle=\min_{\pi\in\Pi(\alpha,\beta)}\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta)+\min_{\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}}\int-\langle\mathbf{M}x,y\rangle d\pi+\mathcal{R}(\mathbf{M})$ (57)
$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta)+\min_{\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}}-\langle\mathbf{M},\int
yx^{\top}\mathop{}\\!\mathrm{d}\pi\rangle_{F}+\mathcal{R}(\mathbf{M})$ (58)
$\displaystyle=\min_{\pi\in\Pi(\alpha,\beta)}\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta)-\max_{\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}}-\langle\mathbf{M},\int
yx^{\top}\mathop{}\\!\mathrm{d}\pi\rangle_{F}+\mathcal{R}(\mathbf{M})$ (59)
$\displaystyle=\min_{\pi\in\Pi(\alpha,\beta)}\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta)-\mathcal{R}^{*}\left(\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}\mathop{}\\!\mathrm{d}\pi(x,y)\right),$
(60)

which is the desired result.

###  B.2 Proof of Proposition 3.2

Consider L-$\mathcal{R}\textsc{OT}_{\varepsilon}$:

$\displaystyle\min_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\mathbf{M}\in\mathbb{R}^{d_{y}\times
d_{x}}\end{subarray}}\int-\langle\mathbf{M}x,y\rangle\mathop{}\\!\mathrm{d}\pi+\mathcal{R}(\mathbf{M})+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta).$
(61)

Under the assumption of Proposition 3.2, the minimizer over $\pi$ is unique
for any fixed $\mathbf{M}$, and the minimizer over $\mathbf{M}$ is unique for
any fixed $\pi$. Thus, we can directly apply (Bertsekas, , 1997, Prop. 2.7.1
and Ex. 2.7.1) to conclude.

###  B.3 Proof of Lemma 3.3

For any $\pi\in\Pi(\alpha,\beta)$,

$\displaystyle\mathbf{M}(\pi)\triangleq\operatorname*{\arg\\!\min}_{\mathbf{M}}-\int_{\mathcal{X}\times\mathcal{Y}}\langle\mathbf{M}x,y\rangle
d\pi+\frac{\left\lVert\mathbf{M}\right\rVert_{F}^{2}}{2}+\lambda
g(\mathbf{M})$ (62)
$\displaystyle=\operatorname*{\arg\\!\min}_{\mathbf{M}}-\langle\mathbf{M},\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}\mathop{}\\!\mathrm{d}\pi\rangle_{F}+\frac{\left\lVert\mathbf{M}\right\rVert_{F}^{2}}{2}+\lambda
g(\mathbf{M})$ (63)
$\displaystyle=\operatorname*{\arg\\!\min}_{\mathbf{M}}\frac{1}{2}\left\lVert\mathbf{M}-\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}\mathop{}\\!\mathrm{d}\pi\right\rVert^{2}_{F}+\lambda
g(\mathbf{M})$ (64) $\displaystyle=\mathrm{prox}_{\lambda
g}\left(\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}d\pi\right),$ (65)

where in the third line we introduced an additional term
$\frac{1}{2}\left\lVert\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}d\pi\right\rVert_{F}^{2}$
which doesn’t depend on $M$.

###  B.4 Proof of Proposition 3.4

Recall the cost-regularized OT problem with $\varepsilon=0$

$\displaystyle\operatorname*{\arg\\!\min}_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}\end{subarray}}-$
$\displaystyle\int\langle\mathbf{M}x,y\rangle
d\pi+\frac{\left\lVert\mathbf{M}\right\rVert_{F}^{2}}{2}+\lambda
g(\mathbf{M}).$ (66)

Denote by $h:\Pi(\alpha,\beta)\times\mathbb{R}^{d_{y}\times
d_{x}}\to\mathbb{R}$ and $H:\Pi(\alpha,\beta)\to\mathbb{R}$:

$\displaystyle
h(\pi,\mathbf{M})=-\int_{\mathcal{X}\times\mathcal{Y}}\langle\mathbf{M}x,y\rangle
d\pi+\frac{\left\lVert\mathbf{M}\right\rVert_{F}^{2}}{2}+\lambda
g(\mathbf{M})\quad\text{and}\quad
H(\pi)=\min_{\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}}h(\pi,\mathbf{M}).$
(67)

As shown in Lemma 3.3, given $\pi$, the solution of 66 is

$\displaystyle\mathbf{M}(\pi)\triangleq\mathrm{prox}_{\lambda
g}\left(\int_{\mathcal{X}\times\mathcal{Y}}yx^{\top}d\pi\right),$ (68)

Proposition 3.4 follows by evaluating the objective (66) in the solution
$\mathbf{M}(\pi)$ using the corresponding proximal operator for each
regularization.

$\ell_{1}$-regularization. Consider the case where
$g(\mathbf{M})=\left\lVert\mathbf{M}\right\rVert_{1}$. Recall that for
$\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}$ and all
$i,j\in[d_{x}]\times[d_{y}]$,

$\displaystyle\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{1}}\left(\mathbf{M}\right)_{ij}$
$\displaystyle=\mathrm{\
sign}(\mathbf{M}_{ij})\left(\left\lvert\mathbf{M}_{ij}\right\rvert-\lambda\right)_{+}.$
(69)

Thus, for $\mathbf{M}(\pi)$,

$\displaystyle\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{1}}\left(\mathbf{M}(\pi)\right)_{ij}$
$\displaystyle=\mathrm{\
sign}\left(\int_{\mathcal{X}\times\mathcal{Y}}x_{i}y_{j}d\pi\right)\left(\int_{\mathcal{X}\times\mathcal{Y}}x_{i}y_{j}d\pi-\lambda\right)_{+}.$
(70)

Define
$I_{\pi}\triangleq\left\\{(i,j)\in[d_{x}]\times[d_{y}]\,:\;\int_{\mathcal{X}\times\mathcal{Y}}x_{i}y_{j}d\pi(x,y)>\lambda\right\\}$.
Plugging (76) in (66) gives, after reducing,

$\displaystyle\operatorname*{\arg\\!\min}_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\end{subarray}}\sum_{(i,j)\in
I_{\pi}}\int_{\mathcal{X}\times\mathcal{Y}}\int_{\mathcal{X}\times\mathcal{Y}}-x_{i}y_{j}x_{i}^{\prime}y_{j}^{\prime}\mathop{}\\!\mathrm{d}\pi(x,y)\mathop{}\\!\mathrm{d}\pi(x^{\prime},y^{\prime})+\lambda\sum_{(i,j)\in
I_{\pi}}\left\lvert\int_{\mathcal{X}\times\mathcal{Y}}x_{i}y_{j}\mathop{}\\!\mathrm{d}\pi\right\rvert-\frac{\lambda^{2}}{2}|I_{\pi}|$
(71)
$\displaystyle=\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\sum_{(i,j)\in
I_{\pi}}\left(x_{i}x_{i}^{\prime}-y_{j}y_{j}^{\prime}\right)^{2}\mathop{}\\!\mathrm{d}\pi\mathop{}\\!\mathrm{d}\pi-\sum_{(i,j)\in
I_{\pi}}\left(\int_{\mathcal{X}}x_{i}^{2}\mathop{}\\!\mathrm{d}\alpha+\int_{\mathcal{Y}}y_{j}^{2}\mathop{}\\!\mathrm{d}\beta-\lambda\left\lvert\int_{\mathcal{X}\times\mathcal{Y}}y_{j}x_{i}d\pi(x,y)\right\rvert\right)-\frac{\lambda^{2}}{2}|I_{\pi}|.$
(72)

Let $\left(\pi^{\star},\mathbf{M}(\pi^{\star})\right)$ be a solution to (29).
Then

$\displaystyle\pi\in\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}H(\pi)\iff\pi\in\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}h(\pi,\mathbf{M}(\pi^{\star})).$
(73)

Hence, using the same calculations as in (71), (66) is equivalent to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\sum_{(i,j)\in
I_{\pi^{\star}}}\left(x_{i}x_{i}^{\prime}-y_{j}y_{j}^{\prime}\right)^{2}\mathop{}\\!\mathrm{d}\pi^{\star}\mathop{}\\!\mathrm{d}\pi-\sum_{(i,j)\in
I_{\pi}^{\star}}\left(\int_{\mathcal{X}}x_{i}^{2}\mathop{}\\!\mathrm{d}\alpha+\int_{\mathcal{Y}}y_{j}^{2}\mathop{}\\!\mathrm{d}\beta-\lambda\left\lvert\int_{\mathcal{X}\times\mathcal{Y}}y_{j}x_{i}d\pi^{\star}(x,y)\right\rvert\right).$
(74)

Thus, ignoring the terms that don’t depend on $\pi$, (66) is equivalent to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\sum_{(i,j)\in
I_{\pi^{\star}}}\left(x_{i}x_{i}^{\prime}-y_{j}y_{j}^{\prime}\right)^{2}\mathop{}\\!\mathrm{d}\pi^{\star}\mathop{}\\!\mathrm{d}\pi,$
(75)

which is the desired result.

$\ell_{1,2}$-regularization. We proceed similarly. Recall that for
$\mathbf{M}\in\mathbb{R}^{d_{y}\times d_{x}}$ and all $i\in[d_{x}]$,

$\displaystyle\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{1,2}}\left(\mathbf{M}\right)_{:i}$
$\displaystyle=\left(1-\frac{\lambda}{\max\left\\{\left\lVert\mathbf{M}_{:i}\right\rVert\,,\lambda\right\\}}\right)\mathbf{M}_{:i},$
(76)

i.e.

$\displaystyle\mathbf{M}_{:i}=\left\\{\begin{array}[]{ll}\left(1-\frac{\lambda}{\mathbf{M}_{:i}}\right)\mathbf{M}_{:i}&\mbox{if
}\left\lVert M_{:i}\right\rVert>\lambda\\\
0&\mbox{otherwise.}\end{array}\right.$ (79)

Let $\mathbf{A}\in\mathbb{R}^{d_{y}\times d_{x}}$ and
$\tilde{\mathbf{A}}=\mathrm{prox}_{\lambda\left\lVert\cdot\right\rVert_{1,2}}(\mathbf{A})$.
Let
$I\triangleq\left\\{\left\lVert\mathbf{A}_{:i}\right\rVert>\lambda\right\\}$.
Then,

$\displaystyle-\langle\mathbf{A},\tilde{\mathbf{A}}\rangle_{F}+\frac{\left\lVert\tilde{\mathbf{A}}\right\rVert}{2}+\lambda\left\lVert\tilde{\mathbf{A}}\right\rVert_{1,2}$
(80) $\displaystyle=\lambda\sum_{i\in
I}\left\lVert\mathbf{A}_{:i}\right\rVert-\lambda^{2}+\sum_{i\in
I}\sum_{j}\left(-\mathbf{A}_{ij}^{2}\left(1-\frac{\lambda}{\left\lVert\mathbf{A}_{:i}\right\rVert}\right)+\frac{1}{2}\left(1-\frac{\lambda}{\left\lVert\mathbf{A}_{:i}\right\rVert}\right)^{2}\mathbf{A}_{ij}^{2}\right)$
(81) $\displaystyle=\lambda\sum_{i\in
I}\left\lVert\mathbf{A}_{:i}\right\rVert-\lambda^{2}|I|-\frac{1}{2}\sum_{i\in
I}\sum_{j}\left(\left\lVert\mathbf{A}_{:i}\right\rVert-\lambda^{2}\right)\frac{\mathbf{A}_{ij}^{2}}{\left\lVert\mathbf{A}_{:i}\right\rVert}$
(82) $\displaystyle=\lambda\sum_{i\in
I}\left\lVert\mathbf{A}_{:i}\right\rVert-\lambda^{2}|I|-\frac{1}{2}\sum_{i\in
I}\left(\left\lVert\mathbf{A}_{:i}\right\rVert^{2}-\lambda^{2}\right)$ (83)
$\displaystyle=-\frac{1}{2}\sum_{i\in
I}\left\lVert\mathbf{A}_{:i}\right\rVert^{2}+\lambda\sum_{i\in
I}\left\lVert\mathbf{A}_{:i}\right\rVert-\frac{\lambda^{2}}{2}|I|$ (84)

Thus, using $\mathbf{A}=\mathbf{M}(\pi)$ and
$I_{\pi}\triangleq\left\\{i\in[d_{x}]\,:\;\left\lVert\int_{\mathcal{Y}\times\mathcal{Y}}x_{i}y\mathop{}\\!\mathrm{d}\pi\right\rVert>\lambda\right\\}$,
(66) is equivalent to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\left(\langle
x_{I_{\pi}},x_{I_{\pi}}^{\prime}\rangle-\langle
y,y^{\prime}\rangle\right)^{2}\mathop{}\\!\mathrm{d}\pi\mathop{}\\!\mathrm{d}\pi-\int_{\mathcal{X}}\langle
x_{I_{\pi}},x_{I_{\pi}}^{\prime}\rangle^{2}\mathop{}\\!\mathrm{d}\alpha+\sum_{i\in
I_{\pi}}\left\lVert\int_{\mathcal{X}\times\mathcal{Y}}x_{i}yd\pi(x,y)\right\rVert-\frac{\lambda^{2}}{2}|I_{\pi}|.$
(85)

Proceding as we did for $\ell_{1}$ regularization, let
$\left(\pi^{\star},\mathbf{M}(\pi^{\star})\right)$ be a solution to (29). Then
(66) is equivalent to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\left(\mathcal{X}\times\mathcal{Y}\right)^{2}}\int_{\mathcal{X}\times\mathcal{Y}}\left(\langle
x_{I_{\pi^{\star}}},x_{I_{\pi^{\star}}}^{\prime}\rangle-\langle
y,y^{\prime}\rangle\right)^{2}\mathop{}\\!\mathrm{d}\pi^{\star}\mathop{}\\!\mathrm{d}\pi-\int_{\mathcal{X}}\langle
x_{I_{\pi^{\star}}},x_{I_{\pi^{\star}}}^{\prime}\rangle^{2}\mathop{}\\!\mathrm{d}\alpha+\lambda\sum_{i\in
I_{\pi^{\star}}}\left\lVert\int_{\mathcal{X}\times\mathcal{Y}}x_{i}yd\pi^{\star}(x,y)\right\rVert-\frac{\lambda^{2}}{2}|I_{\pi^{\star}}|$
(86)

Thus, ignoring the terms that don’t depend on $\pi$, (66) is equivalent to

$\displaystyle\min_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}\times\mathcal{Y}}\left(\langle
x_{I_{\pi^{\star}}},x_{I_{\pi^{\star}}}^{\prime}\rangle-\langle
y,y^{\prime}\rangle\right)^{2}\mathop{}\\!\mathrm{d}\pi^{\star}\mathop{}\\!\mathrm{d}\pi,$
(87)

which is the desired result.

##  Appendix C Appendix for Section 4

###  C.1 Proof for 4.1

Before moving to the proof of Proposition 4.1, we first clarify a point
concerning the existence and convergence of entropic maps for the inner-
product linear OT.

The results of Pooladian and Niles-Weed,  (2021), which define an entropic map
and prove its convergence to the Monge map, are dedicated to the linear OT
problem with the squared euclidean cost with
$\mathcal{X}=\mathcal{Y}\subset\mathbb{R}^{d}$:

$\displaystyle\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}\times\mathcal{Y}}\frac{1}{2}{\left\lVert
x-y\right\rVert}^{2}d\pi(x,y).$ (88)

In our case, an intermediary result that we will need is the convergence of
the entropic map for the inner-product cost $c(x,y)=-\langle x,y\rangle$.
However, by developing the square, one can see that the two problems are
equivalent:

$\displaystyle\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}\times\mathcal{Y}}\frac{1}{2}{\left\lVert
x-y\right\rVert}^{2}d\pi(x,y)$ (89)
$\displaystyle=\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}\times\mathcal{Y}}\frac{1}{2}{\left\lVert
x\right\rVert}^{2}\mathop{}\\!\mathrm{d}\pi(x,y)-\int_{\mathcal{X}\times\mathcal{Y}}\langle
x,y\rangle\mathop{}\\!\mathrm{d}\pi(x,y)+\frac{1}{2}{\left\lVert
y\right\rVert}^{2}\mathop{}\\!\mathrm{d}\pi(x,y)$ (90)
$\displaystyle=\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}\int_{\mathcal{X}}\frac{1}{2}{\left\lVert
x\right\rVert}^{2}\mathop{}\\!\mathrm{d}\alpha(x)-\int_{\mathcal{X}\times\mathcal{Y}}\langle
x,y\rangle\mathop{}\\!\mathrm{d}\pi(x,y)+\frac{1}{2}\int_{\mathcal{Y}}{\left\lVert
y\right\rVert}^{2}\mathop{}\\!\mathrm{d}\beta(y)$ (91)
$\displaystyle=\operatorname*{\arg\\!\min}_{\pi\in\Pi(\alpha,\beta)}-\int_{\mathcal{X}\times\mathcal{Y}}\langle
x,y\rangle\mathop{}\\!\mathrm{d}\pi(x,y).$ (92)

Thus, squared euclidean linear OT and inner-product OT have the same OT plan,
i.e. (since the OT plans for the squared euclidean cost are induced by Monge
maps) the same Monge maps. Hence, the results of Pooladian and Niles-Weed,
(2021) directly extend to the inner-product cost. An additional, easy to
verify fact, is the form of the Monge maps for the inner-product cost.

######  Lemma C.1.

Let $f$ be a Kantorovitch potential for the inner product cost between two
measures $\mu$ and $\nu$ with compact supports in $\mathbb{R}^{d}$. Then the
Monge map for the inner-product cost (or equivalently for the squared
euclidean cost) can be written as

$\displaystyle T(x)=-\nabla f(x).$ (93)

###### Proof.

One can show that if $f$ is a Kantorovitch potential for the inner product
cost, then $\tilde{f}:x\mapsto f(x)+\frac{{\left\lVert x\right\rVert}^{2}}{2}$
is a Kantorovitch potential for the squared Euclidean cost. And we know from
Brenier,  (1991) that given any Kantorovitch potential $\bar{f}$ for the
squared euclidean cost, the unique (since Kantorovitch potentials are equal up
to a constant) Monge map can be written as $T(x)=x-\nabla\bar{f}(x)$. Thus,
given a Kantorovitch potential $f$ for the IP cost, we can write the Monge map
for the inner product cost (or for squared euclidean costs since they have the
same Monge map) as $T(x)=x-\nabla(f(x)+\tfrac{{\left\lVert
x\right\rVert}^{2}}{2})=-\nabla f(x)$. ∎

For completeness, we rewrite the entropic map of Definition 4.2.

$\displaystyle T_{\varepsilon,\varepsilon^{\prime}}(x)=\frac{\int
y\exp{\left(\left(g_{\varepsilon,\varepsilon^{\prime}}(y)+\langle\mathbf{M}_{\varepsilon}^{\star}x,y\rangle\right)/\varepsilon\right)}d\beta(y)}{\int\exp{\left(\left(g_{\varepsilon,\varepsilon^{\star}}(y)+\langle\mathbf{M}_{\varepsilon}^{\star}x,y\rangle\right)/\varepsilon\right)}d\beta(y)}.$
(94)

######  Proposition C.2.

Let
$(f_{\varepsilon,\varepsilon^{\prime}},g_{\varepsilon,\varepsilon^{\prime}})$
be Sinkhorn potentials for the inner product cost between
${\mathbf{M}_{\varepsilon}^{\star}}_{\sharp}\alpha$ and $\beta$ with an
$\varepsilon^{\prime}$ entropic regularization. Then,

$\displaystyle T_{\varepsilon,\varepsilon^{\prime}}=-\nabla
f_{\varepsilon,\varepsilon^{\prime}}\circ\mathbf{M}_{\varepsilon}^{\star}$
(95)

###### Proof.

Using optimality conditions for the $\varepsilon^{\prime}$-entropy regularized
OT problem, we have that for all $\tilde{y}\in\mathcal{Y}$,

$\displaystyle
f_{\varepsilon,\varepsilon^{\prime}}(\tilde{y})=-\varepsilon\log\left(\int\exp{\left(\left(g_{\varepsilon,\varepsilon^{\prime}}(y)+\langle\tilde{y},y\rangle\right)/\varepsilon\right)}d\beta(y)\right)$
(96)

Differentiating w.r.t. $\tilde{y}$ gives

$\displaystyle\nabla
f_{\varepsilon,\varepsilon^{\prime}}(\tilde{y})=-\frac{\int
y\exp{\left(\left(g_{\varepsilon,\varepsilon^{\prime}}(y)+\langle\tilde{y},y\rangle\right)/\varepsilon\right)}d\beta(y)}{\int\exp{\left(\left(g_{\varepsilon,\varepsilon^{\star}}(y)+\langle\tilde{y},y\rangle\right)/\varepsilon\right)}d\beta(y)}.$
(97)

Hence for all $x\in\mathcal{X}$, evaluating at
$\tilde{y}=\mathbf{M}_{\varepsilon}^{\star}x$

$\displaystyle\nabla
f_{\varepsilon,\varepsilon^{\prime}}(\mathbf{M}_{\varepsilon}^{\star}x)=-\frac{\int
y\exp{\left(\left(g_{\varepsilon,\varepsilon^{\prime}}(y)+\langle\mathbf{M}_{\varepsilon}^{\star}x,y\rangle\right)/\varepsilon\right)}d\beta(y)}{\int\exp{\left(\left(g_{\varepsilon,\varepsilon^{\star}}(y)+\langle\mathbf{M}_{\varepsilon}^{\star}x,y\rangle\right)/\varepsilon\right)}d\beta(y)}=-T_{\varepsilon,\varepsilon^{\prime}}(x).$
(98)

That is, $T_{\varepsilon,\varepsilon^{\prime}}=-\nabla
f_{\varepsilon,\varepsilon^{\prime}}\circ\mathbf{M}_{\varepsilon}^{\star}$. ∎

Finally, before moving on to the proof, recall the following result from
Vayer,  (2020), also appearing in (Dumont et al., , 2022, Proposition 6).

######  Proposition C.3 (Theorem 4.2.3 in Vayer,  (2020), Informal.).

Finding a Monge map for the inner-product GW problem reduces to finding a
Monge map between $\alpha$ and $\beta$ for the cost
$c_{\mathbf{M}}(x,y)=-\langle\mathbf{M}^{\star}x,y\rangle$, where
$\mathbf{M}^{\star}=\int yx^{\top}\mathop{}\\!\mathrm{d}\pi^{\star}$, and
$\pi^{\star}$ is a solution to the inner-product GW problem. Such a map is
given by

$\displaystyle T=-\nabla f\circ\mathbf{M}^{\star},$ (99)

where $f$ is a Kantorovitch potential for
$\mathcal{W}_{0,\langle\cdot,\cdot\rangle}\left(\mathbf{M}_{\sharp}\alpha,\beta\right)$.

###### Proof of Proposition 4.4.

Assume that $d_{x}\geq d_{y}$ and consider the case where
$rk(\mathbf{M}^{*})=d_{y}$.

Using Lemma C.1, we now that the Monge map for
$\mathcal{W}^{0}_{\langle\cdot,\cdot\rangle}\left({\mathbf{M}_{\varepsilon}^{\star}}_{\sharp}\alpha,\beta\right)$
is $T_{\varepsilon,0}(x)=-\nabla f_{\varepsilon,0}(x)$, where
$f_{\varepsilon,0}$ is a Kantorovitch potential for
$\mathcal{W}^{0}_{\langle\cdot,\cdot\rangle}\left({\mathbf{M}_{\varepsilon}^{\star}}_{\sharp}\alpha,\beta\right)$.
Since $f_{\varepsilon,\varepsilon^{\prime}}$ is a Sinkhorn potential for
$\mathcal{W}^{\varepsilon^{\prime}}_{\langle\cdot,\cdot\rangle}\left({\mathbf{M}_{\varepsilon}^{\star}}_{\sharp}\alpha,\beta\right)$,
we can apply (Pooladian and Niles-Weed, , 2021, Corollary 1) (here with the
inner product cost rather than the squared euclidean cost), with
$P={\mathbf{M}_{\varepsilon}}_{\sharp}\alpha$ and $Q=\beta$. We have that

$\displaystyle\int_{\mathcal{Y}}{\left\lVert\nabla
f_{\varepsilon,\varepsilon^{\prime}}(y)-\nabla
f_{\varepsilon,0}(y)\right\rVert}^{2}\mathop{}\\!\mathrm{d}{\mathbf{M}_{\varepsilon}}_{\sharp}\alpha(y)\leq{\varepsilon^{\prime}}^{2}I_{0}({\mathbf{M}_{\varepsilon}}_{\sharp}\alpha,\beta)+{\varepsilon^{\prime}}^{\left(\bar{\alpha}+1\right)/2},$
(100)

Hence,

$\displaystyle\int_{\mathcal{X}}{\left\lVert\nabla
f_{\varepsilon,\varepsilon^{\prime}}(\mathbf{M}_{\varepsilon}x)-\nabla
f_{\varepsilon,0}(\mathbf{M}_{\varepsilon}x)\right\rVert}^{2}d\alpha(x)\leq{\varepsilon^{\prime}}^{2}I_{0}({\mathbf{M}_{\varepsilon}}_{\sharp}\alpha,\beta)+{\varepsilon^{\prime}}^{\left(\bar{\alpha}+1\right)/2},$
(101)

that is, using Lemma C.1 and Proposition C.2,

$\displaystyle\int_{\mathcal{X}}{\left\lVert
T_{\varepsilon,\varepsilon^{\prime}}(x)-T_{\varepsilon,0}(x)\right\rVert}^{2}\mathop{}\\!\mathrm{d}\alpha(x)\leq{\varepsilon^{\prime}}^{2}I_{0}({\mathbf{M}_{\varepsilon}}_{\sharp}\alpha,\beta)+{\varepsilon^{\prime}}^{\left(\bar{\alpha}+1\right)/2}.$
(102)

And since
$I_{0}(\alpha_{\varepsilon}^{\prime},\beta_{\varepsilon}^{\prime})<\infty$ and
doesn’t depend on $\varepsilon^{\prime}$ Chizat et al.,  (2020), we have that
for any $\varepsilon>0$,
$T_{\varepsilon,\varepsilon^{\prime}}\xrightarrow{L^{2}(\alpha)}T_{\varepsilon,0}$
as $\;\varepsilon^{\prime}\to 0$.

Now for the second statement. Zhang et al.,  (2022) showed that
$\pi_{\varepsilon}^{*}\rightharpoonup_{\varepsilon}\pi_{0}^{*}$ along a
subsequence. Hence, $\mathbf{M}_{\varepsilon}:=\int
yx^{\top}\mathop{}\\!\mathrm{d}\pi_{\varepsilon}^{\star}\rightharpoonup_{\varepsilon}\mathbf{M}:=\int
yx^{\top}d\pi^{\star}(x,y)$, which implies that
${\mathbf{M}_{\varepsilon}}_{\sharp}\alpha\rightharpoonup\mathbf{M}_{\sharp}\alpha$
along a subsequence.

Since $\nabla f_{\varepsilon,0}$ is a Monge map between
${\mathbf{M}_{\varepsilon}}_{\sharp}\alpha$ and $\beta$, and $\nabla f_{0,0}$
is a Monge map between $\mathbf{M}_{\sharp}\alpha$ and $\beta$, it follows
from the proof of (De Philippis, , 2013, Theorem 4.2) that
$T_{\varepsilon},0\xrightarrow[\varepsilon]{L_{2}(\alpha)}T_{0,0}$. Indeed, it
is shown that this implies that $\nabla f_{\varepsilon,0}$ converges locally
uniformly to $\nabla f_{0,0}$, which in turn implies that $\nabla
f_{\varepsilon,0}$ converges uniformly to $\nabla f_{0,0}$ since $\mathcal{X}$
is compact. Thus, since $\mathbf{M}_{\varepsilon}\rightarrow\mathbf{M}$ and
$\mathcal{X}$ is compact, we have
$T_{\varepsilon,0}\xrightarrow[\varepsilon]{u}T_{0,0}$, and consequently
$T_{\varepsilon,0}\xrightarrow[\varepsilon]{L_{2}(\alpha)}T_{0,0}$ along a
subsequence, which gives the desired result. ∎

###  C.2 Using rank constraints in extremely high-dimensional problems

As we suggested at the end of Section 3, using the proximal operators of the
nuclear norm or of the rank can be infeasible in extremely high dimension
because they involve computing the SVD of a high-dimensional matrix in the
$\mathbf{M}$-step of Prox-$\mathcal{R}\textsc{OT}$. Thus, in such cases, we
suggest to choose a rank $r\ll d_{y}$ and use the following regulariztion:

$\displaystyle\mathcal{R}(c)=\left\\{\begin{array}[]{ll}\frac{1}{2}\left\lVert\mathbf{M}_{2}^{\top}\mathbf{M}_{1}\right\rVert_{F}^{2}&\mbox{if
}\exists\left(\mathbf{M}_{1},\mathbf{M}_{2}\right)\in\mathbb{R}^{d_{x}\times
r}\times\mathbb{R}^{d_{y}\times
r}:\forall(x,y),\;c(x,y)=\langle\mathbf{M}_{1}x,\mathbf{M}_{2}y\rangle\\\
\infty&\mbox{otherwise.}\end{array}\right.$ (105)

In other words, we directly parametrize the matrix $\mathbf{M}$ in
L-$\mathcal{R}\textsc{OT}_{\varepsilon}$ as a low-rank matrix:
$\mathbf{M}=\mathbf{M}_{2}^{\top}\mathbf{M}_{1}$. In this case,
L-$\mathcal{R}\textsc{OT}_{\varepsilon}$ rewrites:

$\displaystyle\min_{\begin{subarray}{c}\pi\in\Pi(\alpha,\beta)\\\
\left(\mathbf{M}_{1},\mathbf{M}_{2}\right)\in\mathbb{R}^{r\times
d_{x}}\times\mathbb{R}^{r\times
d_{y}}\end{subarray}}\int-\langle\mathbf{M}_{1}x,\mathbf{M}_{2}y\rangle\mathop{}\\!\mathrm{d}\pi+\frac{1}{2}\left\lVert\mathbf{M}_{2}^{\top}\mathbf{M}_{1}\right\rVert_{F}^{2}+\varepsilon\textsc{KL}(\pi||\alpha\otimes\beta).$
(L-$\mathcal{R}\textsc{OT}_{\varepsilon}$)

To solve this problem, we use alternated minimization (a.k.a. block coordinate
descent) on each of the variables. Notably, we never have to compute the
product $\mathbf{M}_{2}^{\top}\mathbf{M}_{1}$. Fixing $\pi$ and considering
the first-order conditions on $(\mathbf{M}_{1},\mathbf{M}_{2})$ gives

$\displaystyle\left(\mathbf{M}_{2}\mathbf{M}_{2}^{\top}\right)\mathbf{M}_{1}=\mathbf{M}_{2}\int
yx^{\top}\mathop{}\\!\mathrm{d}\pi(x,y)$ (106)
$\displaystyle\left(\mathbf{M}_{1}\mathbf{M}_{1}^{\top}\right)\mathbf{M}_{2}=\mathbf{M}_{1}\int
xy^{\top}\mathop{}\\!\mathrm{d}\pi(x,y).$ (107)

Thus, assuming that $\left(\mathbf{M}_{2}\mathbf{M}_{2}^{\top}\right)$ and
$\left(\mathbf{M}_{1}\mathbf{M}_{1}^{\top}\right)$ are invertible,

$\displaystyle\mathbf{M}_{1}=\left(\mathbf{M}_{2}\mathbf{M}_{2}^{\top}\right)^{-1}\int\left(\mathbf{M}_{2}y\right)x^{\top}\mathop{}\\!\mathrm{d}\pi(x,y)$
(108)
$\displaystyle\mathbf{M}_{2}=\left(\mathbf{M}_{1}\mathbf{M}_{1}^{\top}\right)^{-1}\int\left(\mathbf{M}_{1}x\right)y^{\top}\mathop{}\\!\mathrm{d}\pi(x,y).$
(109)

So in practice, given $\pi_{k}$, we successively update

$\displaystyle\mathbf{M}_{1}^{k+1}=\left(\mathbf{M}_{2}^{k}{\mathbf{M}_{2}^{k}}^{\top}\right)^{\dagger}\int\left(\mathbf{M}_{2}^{k}y\right)x^{\top}\mathop{}\\!\mathrm{d}\pi_{k}(x,y)$
(110)
$\displaystyle\mathbf{M}_{2}^{k+1}=\left(\mathbf{M}_{1}^{k+1}{\mathbf{M}_{1}^{k+1}}^{\top}\right)^{\dagger}\int\left(\mathbf{M}_{1}^{k+1}x\right)y^{\top}\mathop{}\\!\mathrm{d}\pi_{k}(x,y),$
(111)

where $\left(\mathbf{A}\right)^{\dagger}$ denotes the Moore pseudoinverse of
$\mathbf{A}$. We can see that here all the operations are linear in $d_{x}$
and $d_{y}$. Moreover, as we choose a can choose a small rank $r\ll d_{y}$,
the $\mathcal{O}(r^{3})$ dependency of computing the pseudoinverse can be much
smaller than $\mathcal{O}(d_{y})$.

Application of the updates (111). We test the updates of (111) in place of the
proximal rank operator we used in Section 5. These updates result in very
similar performance in practice (Figure 5). We display Figure 3 here again for
the convenience of the reader.

Figure 4: Copy of Figure 3. Figure 5: Reproduction of Figure 4 but with the
updates (111) instead of Prox-$\mathcal{R}\textsc{OT}$.

##  Appendix D Appendix for Section 5

###  D.1 Illustration of the action of the entropic map in the case where
$\mathbf{M}^{\star}$ is not full-rank

One condition that we need for the convergence of the entropic map in
Proposition 4.4 is that the matrix $\mathbf{M}^{\star}$ (where
$\left(\pi^{\star},\mathbf{M}^{\star}\right)$ is a solution of
L-$\mathcal{R}\textsc{OT}_{\varepsilon}$) is full-rank. In Section 5, we
showed through an experiment (see Figure 3) that even when the matrix
$\mathbf{M}^{\star}$ is explicitely constrained to being low-rank, the
entropic map still results in good performance.

Here, we show through a more controlled experiment that indeed the fact that
$\mathbf{M}^{\star}$ is not full rank doesn’t result in a divergent entropic
map (as measured by the performance of the map when $\varepsilon\approx 0$).

To do so, we simulate two point clouds coming from two Gaussian distributions,
one in 20D and the other in 10D. We then use rank regularizations and take
$\lambda=12.5$, which is equivalent to constraining the matrix $\mathbf{M}$ to
being of at most of rank $5$. The reason we use Gaussians is that we know
ground truth OT map between them (Salmona et al., , 2022).

Figure 6: Illustration of the action of the entropic map when
$\mathbf{M}^{\star}_{\varepsilon}$ is low-rank.

###  D.2 Details on the evaluation metrics for sparse and low-rank transforms
applications

For our sparse transforms experiment on the scGM dataset (Section 5.1), we
measured the performance of the algorithms using Label Transfer Accuracy, as
done in Demetci et al.,  (2022). Label Transfer Accuracy uses a k-NN
classifier fitted on the data in the anchor domain (the one where we map using
the entropic map, for us the data in the lower-dimensional data methylation
domain $\left\\{y_{i}\right\\}_{i=1}^{N}$). Then the score is given as the
test accuracy of the classifier on the mapped
$\left\\{T_{\varepsilon}(x_{i})\right\\}_{i=1}^{N}$.

For our low-rank transforms experiment on the Neurips 2021 multimodal single-
cell integration dataset (Lance et al., , 2022), we measured the performance
of the algorithms using the Fraction of Samples Closer than the True Match
(FOSCTTM) as done in the MOSCOT package (Klein et al., 2023a, ). The score
computes the distances between each mapped sample
$\left\\{T_{\varepsilon}(x_{i})\right\\}_{i=1}^{N}$ and the data
$\left\\{y_{i}\right\\}_{i=1}^{N}$ and, given ground truth correspondences
between data in $\mathcal{X}$ and data in $\mathcal{Y}$, computes the FOSCTTM
score as the average proportion of $y_{i}$’s closer to
$T_{\varepsilon}(x_{i})$ than its true corresponding datapoint in
$\left\\{y_{i}\right\\}_{i=1}^{N}$.

###  D.3 Details on the spatial transcriptomics experiment

As done in Scetbon et al.,  (2023), to determine the best hyperparameters for
the experiment, we ran a grid search and picked the best hyperparameters
combination using performance on 10 validation genes and use pearson
correlation as a validation metric. The hyperparameters we considered in our
grid search are the entropic regularization $\varepsilon$ and the fused cost
parameter $\eta$ (26). We used a single minibatch size of $1000$ and a single
number of ascent steps in $\left(f,g\right)$ of $1000$. For the SGD step in
$\mathbf{M}$, we use stochastic line search (Vaswani et al., , 2019).