# Unbalanced Low-rank Optimal Transport Solvers

Meyer Scetbon∗
Microsoft Research
<EMAIL_ADDRESS>
&Michael Klein∗
Apple
<EMAIL_ADDRESS>
&Giovanni Palla
Helmholtz Center Munich
<EMAIL_ADDRESS>
&Marco Cuturi
Apple
<EMAIL_ADDRESS>
###### Abstract

The relevance of optimal transport methods to machine learning has long been
hindered by two salient limitations. First, the $O(n^{3})$ computational cost
of standard sample-based solvers (when used on batches of $n$ samples) is
prohibitive. Second, the mass conservation constraint makes OT solvers too
rigid in practice: because they must match all points from both measures,
their output can be heavily influenced by outliers. A flurry of recent works
in OT has addressed these computational and modelling limitations, but has
resulted in two separate strains of methods: While the computational outlook
was much improved by entropic regularization, more recent $O(n)$ linear-time
low-rank solvers hold the promise to scale up OT further. On the other hand,
modelling rigidities have been eased owing to unbalanced variants of OT, that
rely on penalization terms to promote, rather than impose, mass conservation.
The goal of this paper is to merge these two strains, to achieve the promise
of both versatile/scalable unbalanced/low-rank OT solvers. We propose custom
algorithms to implement these extensions for the linear OT problem and its
Fused-Gromov-Wasserstein generalization, and demonstrate their practical
relevance to challenging spatial transcriptomics matching problems.

##  1 Introduction

Recent machine learning (ML) works have witnessed a flurry of activity around
optimal transport (OT) methods. The OT toolbox provides convenient, intuitive
and versatile ways to quantify the difference between two probability
measures, either to quantify a distance (the Wasserstein and Gromov-
Wasserstein distances), or, in more elaborate scenarios, by computing a push-
forward map that can transform one measure into the other (Peyré and Cuturi,
2019). Recent examples include, e.g., single-cell omics (Bunne et al., 2021,
2022; Demetci et al., 2020; Nitzan et al., 2019; Cang et al., 2023; Klein et
al., 2023), attention mechanisms (Tay et al., 2020; Sander et al., 2022),
self-supervised learning(Caron et al., 2020; Oquab et al., 2023), and learning
on graphs (Vincent-Cuaz et al., 2023).

On the challenges of using OT. Despite their long presence in ML (Rubner et
al., 2000), OT methods have long suffered from various limitations, that arise
from their statistical, computational, and modelling aspects. The statistical
argument is commonly referred to as the curse-of-dimensionality of OT
estimators: the Wasserstein distance between two probability densities, and
its associated optimal Monge map, is poorly approximated using samples as the
dimension $d$ of observation grows (Dudley et al., 1966; Boissard and Le
Gouic, 2014). On the computational side, computing OT between a pair of $n$
samples involves solving a (generalized) matching problem, with a price of
$O(n^{3})$ and above (Kuhn, 1955; Ahuja et al., 1993). Finally, the original
model for OT rests on a mass conservation constraint: all observations from
either samples must be accounted for, including outliers that are prevalent in
machine learning datasets. Combined, these weaknesses have long hindered the
use of OT, until a more recent generation of solvers addressed these three
crucial issues.

The Entropic Success Story. The winning approach, so far, to carry out that
agenda has been entropic regularization methods (Cuturi, 2013). The
computational virtues of the Sinkhorn algorithm when solving OT (Altschuler et
al., 2017; Peyré et al., 2016; Solomon et al., 2016) come with statistical
efficiency (Genevay et al., 2019; Mena and Niles-Weed, 2019; Chizat et al.,
2020), and can also be seamlessly combined with unbalanced formulations by
penalizing – rather than constraint – mass conversation, both for the linear
(Frogner et al., 2015; Chizat et al., 2018; Séjourné et al., 2022; Fatras et
al., 2021) and quadratic (Séjourné et al., 2021) problems. These developments
have all been implemented in popular OT packages (Feydy et al., 2019; Flamary
et al., 2021; Cuturi et al., 2022).

The Low-Rank Alternative. A recent strain of solvers relies instead on low-
rank (LR) properties of cost and coupling matrices (Forrow et al., 2018;
Scetbon and Cuturi, 2020; Scetbon et al., 2021). Much like entropic solvers,
these LR solvers have a better statistical outlook (Scetbon and Cuturi, 2022)
and extend to GW problems (Scetbon et al., 2022). In stark contrast to
entropic solvers, however, LR solvers benefit from linear complexity $O(nrd)$
w.r.t sample size $n$ (using rank $r$ and cost dimension $d$) that can scale
to ambitious tasks where entropic solvers fail (Klein et al., 2023).

The Need for Unbalanced Low-Rank Solvers. LR solvers do suffer, however, from
a major practical limitation: their inability to handle unbalanced problems.
Yet, unbalancedness is a crucial ingredient for OT to be practically relevant.
This is exemplified by the fact that unbalancedness played a crucial role in
the seminal reference (Schiebinger et al., 2019), where it is used to model
cell birth and death.

Our Contributions We propose in this work to lift this last limitation for LR
solvers to:

  * •

Incorporate unbalanced regularizers to define a LR linear solver (§ 3.1);

  * •

Provide accelerated algorithms, inspired by some of the recent corrections
proposed by (Séjourné et al., 2022), to isolate translation terms that appear
in dual subroutines (§ 3.2);

  * •

Carry over and adapt these approaches to the GW (§ 3.3) and Fused-GW problems
(§ 3.4);

  * •

Carry out an exhaustive hyperparameter selection procedure within large scale
OT tasks (spatial transcriptomics, brain imaging), and demonstrate the
benefits of our approach (§ 4).

##  2 Reminders on Low-Rank Transport and Unbalanced Transport

We consider two metric spaces $(\mathcal{X},d_{\mathcal{X}})$ and
$(\mathcal{Y},d_{\mathcal{Y}})$, as well as a cost function
$c:\mathcal{X}\times\mathcal{Y}\to[0,+\infty[$. The simplex $\Delta_{n}^{+}$
holds all positive $n$-vectors summing to $1$. For $n,m\geq
1,a\in\Delta_{n}^{+}$, and $b\in\Delta_{m}^{+}$, given points
$x_{1},\dots,x_{n}\in\mathcal{X}$ and $y_{1},\dots,y_{m}\in\mathcal{Y}$, we
define two discrete probability measures $\mu$ and $\nu$ as
$\mu:=\sum_{i=1}^{n}a_{i}\delta_{x_{i}}$,
$\nu:=\sum_{j=1}^{m}b_{j}\delta_{y_{j}}$ where $\delta_{z}$ is the Dirac mass
at $z$.

Cost matrices. For $q\geq 1$, consider first two square pairwise cost
matrices, each encoding the geometries of points within $\mu$ and $\nu$, and a
rectangular matrix that studies that across their support:

$A:=[d_{\mathcal{X}}^{q}(x_{i},x_{i^{\prime}})]_{1\leq i,i^{\prime}\leq
n},\;B:=[d_{\mathcal{Y}}^{q}(y_{j},y_{j^{\prime}})]_{1\leq j,j^{\prime}\leq
m}\,,\;C:=[c(x_{i},y_{j})]_{1\leq i,j\leq n,m}\,.$

The Kantorovich Formulation of OT is defined as the following linear program,
defined by $C$:

$\displaystyle\text{OT}(\mu,\nu):=\min_{P\in\Pi_{a,b}}\langle
C,P\rangle\,,\quad\text{where}\quad\Pi_{a,b}:=\left\\{P\in\mathbb{R}_{+}^{n\times
m},~{}\text{s.t.}~{}~{}P\mathbf{1}_{m}=a,~{}P^{T}\mathbf{1}_{n}=b\right\\}\,.$
(1)

The Low-Rank Formulation of OT is best understood as a variant of (1) that
rests on a low-rank property for cost matrix $C$, and low-rank constraints for
couplings $P$. More precisely, Scetbon et al. (2021) propose to constraint the
set of admissible couplings to those, within $\Pi_{a,b}$, that have a non-
negative rank of $r\geq 1$. That set can be equivalently reparamaterized as

$\displaystyle\Pi_{a,b}(r)=\\{P\in\mathbb{R}^{n\times
m}_{+}|P=Q\operatorname{diag}(1/g)R^{T},~{}~{}Q\in\Pi_{a,g},~{}~{}R\in\Pi_{b,g},~{}~{}\text{and}~{}~{}g\in\Delta_{r}^{+}\\}.$

The low-rank optimal transport (LOT) problem simply uses that restriction in
(1) to define :

$\displaystyle\text{LOT}_{r}(\mu,\nu):=\min_{P\in\Pi_{a,b}(r)}\langle
C,P\rangle=\min_{Q\in\Pi_{a,g},R\in\Pi_{a,g},g\in\Delta^{+}_{r}}\langle
C,Q\operatorname{diag}(g)R\rangle\,.$ (2)

Scetbon et al. (2021) propose and prove the convergence of a mirror-descent
scheme to solve (2), and obtain linear time and memory complexities with
respect to the number of samples, where each iteration in that descent scales
as $(n+m)rd$, where $d$ is the rank of $C$.

The Unbalanced Formulation of OT starts from (1) as well, but proposes to do
without $\Pi_{a,b}$ and its marginal constraints (Frogner et al., 2015; Chizat
et al., 2018), and rely instead on two regularizers:

$\text{UOT}(\mu,\nu):=\min_{P\in\mathbb{R}^{n\times m}_{+}}\langle
C,P\rangle+\tau_{1}\textrm{KL}(P\mathbf{1}_{m}|a)+\tau_{2}\textrm{KL}(P^{T}\mathbf{1}_{n}|b).$
(3)

This formulation is solved using entropic regularization, with modified
Sinkhorn updates (Frogner et al., 2015). Proposing an efficient algorithm able
to merge (2) with (3) is the first goal of this paper.

Gromov-Wasserstein (GW) Considerations. The GW problem (Mémoli, 2011) is a
generalization of (1) where the energy $\mathcal{Q}_{A,B}$ is a quadratic
function of $P$ defined through inner cost matrices $A$, $B$:

$\displaystyle\\!\\!\mathcal{Q}_{A,B}(P)\\!:=\\!\\!\\!\\!\\!\sum_{i,j,i^{\prime},j^{\prime}}\\!\\!(A_{ii^{\prime}}-B_{jj^{\prime}})^{2}P_{ij}P_{i^{\prime}j^{\prime}}\\!=\\!\mathbf{1}_{m}^{T}P^{T}A^{\odot
2}P\mathbf{1}_{m}+\mathbf{1}_{n}^{T}PB^{\odot 2}P^{T}\mathbf{1}_{n}-2\langle
APB,P\rangle\;.$ (4)

To minimize (4), the default approach rests on entropic regularization
(Solomon et al., 2016; Peyré et al., 2016) and variants (Sato et al., 2020;
Blumberg et al., 2020; Xu et al., 2019; Li et al., 2023). Scetbon et al.
(2022) adapted the low-rank framework to minimize $\mathcal{Q}_{A,B}$ over
low-rank matrices $P$, achieving a linear-time complexity when $A$ and $B$ are
themselves low-rank. Independently, (Séjourné et al., 2021) proposed an
unbalanced generalization that also applies to GW and which can be implemented
practically using entropic regularization. Finally, the minimization of a
composite objective involving both $\mathcal{Q}_{A,B}$ and $\langle
C,\cdot\rangle$ is known as the fused GW problem (Vayer et al., 2018).

##  3 Unbalanced Low-Rank Transport

###  3.1 Unbalanced Low-rank Linear Optimal Transport

We incorporate unbalancedness to low-rank solvers (Scetbon et al., 2021,
2022), moving gradually from the linear problem to the more involved GW and
FGW problem. Using the framework of (Frogner et al., 2015; Chizat et al.,
2018), we can first extend the definition of LOT, introduced in (2), to the
unbalanced case by considering the following optimization problem:

$\displaystyle\text{ULOT}_{r}(\mu,\nu):=\min_{P\text{:}~{}\text{rk}_{+}(P)\leq
r}\langle C,P\rangle+$
$\displaystyle\tau_{1}\text{KL}(P\mathbf{1}_{m}|a)+\tau_{2}\text{KL}(P^{T}\mathbf{1}_{n}|b),$
(5)

where $\text{rk}_{+}(P)$ denotes the nonnegative rank of $P$. Therefore by
denoting $\Pi_{r}:=\\{(Q,R,g)\in\mathbb{R}_{+}^{n\times
r}\times\mathbb{R}_{+}^{m\times
r}\times\mathbb{R}_{+}^{r}\text{:}~{}Q^{T}\mathbf{1}_{n}=R^{T}\mathbf{1}_{m}=g\\}$,
and using the repamatrization of low-rank couplings, we obtain the following
equivalent formulation of ULOT:

$\displaystyle\text{ULOT}_{r}(\mu,\nu)=\min_{(Q,R,g)\in\Pi_{r}}\underbrace{\langle
C,Q\operatorname{diag}(1/g)R^{T}\rangle}_{\mathcal{L}_{C}(Q,R,g)}+\underbrace{\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)}_{\mathcal{G}_{a,b}(Q,R,g)}\;.$
(6)

We introduce slightly more compact notations for
$\mathcal{G}_{a,b}(Q,R,g)=F_{\tau_{1},a}(Q\mathbf{1}_{r})+F_{\tau_{2},b}(R\mathbf{1}_{r}),$
where $F_{\tau,z}(s)=\tau\text{KL}(s|z)$ for $\tau>0$ and $z\geq 0$
coordinate-wise. To solve (6), and using this split, we move away from mirror-
descent and apply instead proximal gradient-descent with respect to the KL
divergence. At each iteration, we consider a linear approximation of
$\mathcal{L}_{C}$ where a KL penalization is added to the objective, as in the
classical mirror descent, however, we leave $\mathcal{G}_{a,b}$ intact at each
iteration. Borrowing notations from (Scetbon et al., 2021), we must solve at
each iteration the convex optimization problem:

$\displaystyle(Q_{k+1},R_{k+1},g_{k+1}):=\operatorname*{argmin}_{\bm{\zeta}\in\Pi_{r}}\frac{1}{\gamma_{k}}\text{KL}(\bm{\zeta},\bm{\xi}_{k})+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$
(7)

where $(Q_{0},R_{0},g_{0})\in\Pi_{r}$ is an initial point,
$\bm{\xi}_{k}:=(\xi_{k}^{(1)},\xi_{k}^{(2)},\xi_{k}^{(3)})$ holds running
costs matrices defined as

$\xi_{k}^{(1)}:=Q_{k}\odot
e^{-\gamma_{k}CR_{k}\operatorname{diag}(1/g_{k})},\xi_{k}^{(2)}:=R_{k}\odot
e^{-\gamma_{k}C^{T}Q_{k}\operatorname{diag}(1/g_{k}))},\xi_{k}^{(3)}:=g_{k}\odot
e^{\gamma_{k}\omega_{k}/g_{k}^{2}},$

with $[\omega_{k}]_{i}:=[Q_{k}^{T}CR_{k}]_{i,i}$ for all
$i\in\\{1,\dots,r\\}$, and $(\gamma_{k})_{k\geq 0}$ is a sequence of positive
step sizes.

Reformulation using Duality. To solve (7), we follow (Scetbon et al., 2021)
and apply Dykstra’s algorithm (Dykstra, 1983). The iterations of this
algorithm takes a very simple form that can be obtained as an alternating
maximization on the dual formulation of (7), provided as follows.

######  Proposition 1.

The convex optimization problem defined in (7) admits the following dual:

$\displaystyle\sup_{f_{1},h_{1},f_{2},h_{2}}$
$\displaystyle\mathcal{D}_{k}(f_{1},h_{1},f_{2},h_{2}):=-F_{\tau_{1},a}^{*}(-f_{1})-\frac{1}{\gamma_{k}}\langle
e^{\gamma_{k}(f_{1}\oplus h_{1})}-1,\xi^{(1)}_{k}\rangle$ (8) $\displaystyle-
F_{\tau_{2},b}^{*}(-f_{2})-\frac{1}{\gamma_{k}}\langle
e^{\gamma_{k}(f_{2}\oplus
h_{2})}-1,\xi^{(2)}_{k}\rangle-\frac{1}{\gamma_{k}}\langle
e^{-\gamma_{k}(h_{1}+h_{2})}-1,\xi^{(3)}_{k}\rangle$

where $h_{1},h_{2}\in\mathbb{R}^{r}$, $f_{1}\in\mathbb{R}^{n}$,
$f_{2}\in\mathbb{R}^{m}$, $F_{\tau,z}^{*}(\cdot):=\sup_{y}\\{\langle
y,\cdot\rangle-F_{\tau,z}(y)\\}$ is the convex conjugate of $F_{\tau,z}$. In
addition strong duality holds and the primal problem admits a unique
minimizer.

######  Remark 1.

While we stick to KL regularizers in this work for simplicity, it is worth
noting that this can be extended to more generic regularizers $F_{\tau_{1},a}$
and $F_{\tau_{2},b}$, as considered by Chizat et al. (2018).

We use an alternating maximization scheme to solve (8). Starting from
$h_{1}^{(0)}=h_{2}^{(0)}=\mathbf{0}_{r}$, we apply for $\ell\geq 0$ the
following updates (dropping iteration number $k$ in (7) for simplicity):

$\displaystyle
f_{1}^{(\ell+1)}:=\arg\sup_{z}\mathcal{D}(z,h_{1}^{(\ell)},f_{2}^{(\ell)},h_{2}^{(\ell)}),\,f_{2}^{(\ell+1)}:=\arg\sup_{z}\mathcal{D}(f_{1}^{(\ell+1)},h_{1}^{(\ell)},z,h_{2}^{(\ell)}),$
$\displaystyle(h_{1}^{(\ell+1)},h_{2}^{(\ell+1)}):=\arg\sup_{z_{1},z_{2}}\mathcal{D}(f_{1}^{(\ell+1)},z_{1},f_{2}^{(\ell+1)},z_{2}).$

These maximizations can all be obtained in closed form, to result in the
closed-form updates:

$\displaystyle\exp(\gamma
f_{1}^{(\ell+1)})=\left(\frac{a}{\xi^{(1)}\exp(\gamma
h_{1}^{(\ell)})}\right)^{\frac{\tau_{1}}{\tau_{1}+1/\gamma}},\quad\exp(\gamma
f_{2}^{(\ell+1)})=\left(\frac{b}{\xi^{(2)}\exp(\gamma
h_{2}^{(\ell)})}\right)^{\frac{\tau_{2}}{\tau_{2}+1/\gamma}}$ $\displaystyle
g_{\ell+1}:=\left(\xi^{(3)}\odot(\xi^{(1)})^{T}\exp(\gamma
f_{1}^{(\ell+1)})\odot(\xi^{(2)})^{T}\exp(\gamma
f_{2}^{(\ell+1)})\right)^{1/3}$ $\displaystyle\exp(\gamma
h_{1}^{(\ell+1)})=\frac{g_{\ell+1}}{(\xi^{(1)})^{T}\exp(\gamma
f_{1}^{(\ell+1)})},\quad\exp(\gamma
h_{2}^{(\ell+1)})=\frac{g_{\ell+1}}{(\xi^{(2)})^{T}\exp(\gamma
f_{2}^{(\ell+1)})}$

When using “scaling” representations for these dual variables, $\ell\geq 0$,
$u^{(\ell)}_{i}:=\exp(\gamma f_{i}^{(\ell)})$ and $v^{(\ell)}_{i}:=\exp(\gamma
h_{i}^{(\ell)})$ for $i\in\\{1,2\\}$, we obtain a simple update, provided in
the appendix (Alg. 6).

Initialization and Termination. We use the stopping criterion proposed in
(Scetbon et al., 2021) to terminate the algorithm,
$\Delta(\bm{\zeta},\bm{\tilde{\zeta}},\gamma):=\frac{1}{\gamma^{2}}(\mathrm{KL}(\bm{\zeta},\bm{\tilde{\zeta}})+\mathrm{KL}(\bm{\tilde{\zeta}},\bm{\zeta}))$.
Combined with practical improvements proposed in (Scetbon and Cuturi, 2022) to
initialize the algorithm, and adapt the choice of $\gamma_{k}$ at each
iteration $k$ of the outer loop, we can now summarize our proposal in
Algorithm 1, which can be seen as an extension of (Scetbon et al., 2021,
Alg.2).

Convergence and Complexity. The proof of convergence of the Dykstra algorithm
(Alg. 6) can be found for example in (Bauschke and Combettes, 2008)). In
addition, (Scetbon et al., 2021) show the convergence of their scheme towards
a stationary points w.r.t to the criterion $\Delta(\cdot,\cdot,\gamma)$ for
$\gamma$ fixed along the iterations of the outer loop. In terms of complexity,
given $\bm{\xi}$, solving Eq. (7) requires a time and memory complexity of
$\mathcal{O}((n+m)r)$. However computing $\bm{\xi}$ requires in general
$\mathcal{O}((n^{2}+m^{2})r)$ time and $\mathcal{O}(n^{2}+m^{2})$ memory. In
(Scetbon et al., 2021), the authors propose to consider low-rank approximation
of the cost matrix $C$ of the form $C\simeq C_{1}C_{2}^{T}$ where
$C_{1}\in\mathbb{R}^{n\times d}$ and $C_{2}\in\mathbb{R}^{m\times d}$: in that
case computing $\bm{\xi}$ can be done in $\mathcal{O}((n+m)rd)$ time and
$\mathcal{O}((n+m)(r+d))$ memory. Such approximations can be obtained using
the algorithm in (Indyk et al., 2019) which guarantees that for any distance
matrix $C\in\mathbb{R}^{n\times m}$ and $\alpha>0$ it can outputs matrices
$C_{1}\in\mathbb{R}^{n\times d}$, $C_{2}\in\mathbb{R}^{m\times d}$ in
$\mathcal{O}((m+n)\text{poly}(\frac{d}{\alpha}))$ algebraic operations such
that with probability at least $0.99$ that
$\|C-C_{1}C_{2}^{T}\|_{F}^{2}\leq\|C-C_{d}\|_{F}^{2}+\alpha\|C\|_{F}^{2}$,
where $C_{d}$ denotes the best rank-$d$ approximation to $C$.

Inputs: $C,a,b,r,\gamma_{0},\tau_{1},\tau_{2},\delta$
$Q,R,g\leftarrow\text{Initialization as proposed
in~{}\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{scetbon2022lot}{\@@citephrase{,
}}{})}}$
repeat

$\tilde{Q}=Q,~{}~{}\tilde{R}=R,~{}~{}\tilde{g}=g,\\\
\nabla_{Q}=CR\operatorname{diag}(1/g),~{}~{}\nabla_{R}=C^{\top}Q\operatorname{diag}(1/g),\\\
\omega\leftarrow\mathcal{D}(Q^{T}CR),~{}~{}\nabla_{g}=-\omega/g^{2},\\\
\gamma\leftarrow\gamma_{0}/\max(\|\nabla_{Q}\|_{\infty}^{2},\|\nabla_{R}\|_{\infty}^{2},\|\nabla_{g}\|_{\infty}^{2}),\\\
\xi^{(1)}\leftarrow Q\odot\exp(-\gamma\nabla_{Q}),~{}\xi^{(2)}\leftarrow
R\odot\exp(-\gamma\nabla_{R}),~{}\xi^{(3)}\leftarrow
g\odot\exp(-\gamma\nabla_{g}),\\\ Q,R,g\leftarrow\text{ULR-
Dykstra}(a,b,\bm{\xi},\gamma,\tau_{1},\tau_{2},\delta)~{}(\text{Alg.~{}\ref{alg-
dykstra-ulot})}$

until  _$\Delta((Q,R,g),(\tilde{Q},\tilde{R},\tilde{g}),\gamma) <\delta$_;

Result: $Q,R,g$

Algorithm 1 $\text{ULOT}(C,a,b,r,\gamma_{0},\tau_{1},\tau_{2},\delta)$

###  3.2 Improvements on the Unbalanced Dykstra Algorithm

A well documented source of instability of unbalanced formulations of OT lies
in capturing efficiently what optimal mass is targeted by such formulations.
Séjourné et al. (2022) have proposed a technique to address this issue and
lower significantly computational costs. They propose first a dual objective
that is _translation_ invariant. We take inspiration from this strategy and
adapt to our problem, to propose the following variant of (8):

$\sup_{\tilde{f}_{1},\tilde{h}_{1},\tilde{f}_{2},\tilde{h}_{2}}\left(\mathcal{D}_{\textrm{TI}}(\tilde{f}_{1},\tilde{h}_{1},\tilde{f}_{2},\tilde{h}_{2}):=\sup_{\lambda_{1},\lambda_{2}\in\mathbb{R}}\mathcal{D}(\tilde{f}_{1}+\lambda_{1},\tilde{h}_{1}-\lambda_{1},\tilde{f}_{2}+\lambda_{2},\tilde{h}_{2}-\lambda_{2})\right)$
(9)

It is clear from the reparametrization that both problems (8) and (9) have the
same value and also that
$(\tilde{f}_{1},\tilde{h}_{1},\tilde{f}_{2},\tilde{h}_{2})$ is solution of (9)
if and only if
$(\tilde{f}_{1}+\lambda_{1}^{*},\tilde{h}_{1}-\lambda_{1}^{*},\tilde{f}_{2}+\lambda_{2}^{*},\tilde{h}_{2}-\lambda_{2}^{*})$
is solution of (8) where $(\lambda_{1}^{*},\lambda_{2}^{*})$ solves
$\mathcal{D}_{\textrm{TI}}(\tilde{f}_{1},\tilde{h}_{1},\tilde{f}_{2},\tilde{h}_{2})$.
To solve (9), we show that the variational formulation of the translation
invariant dual objective targeted inside (9) can be obtained in closed form.

######  Proposition 2.

Let $\tilde{f}_{1}\in\mathbb{R}^{n}$, $\tilde{f}_{2}\in\mathbb{R}^{m}$ and
$\tilde{h}_{1},\tilde{h}_{2}\in\mathbb{R}^{r}$, then the inner problem defined
in (9) by
$\mathcal{D}_{\textrm{TI}}(\tilde{f}_{1},\tilde{h}_{1},\tilde{f}_{2},\tilde{h}_{2})$
admits a unique solution $(\lambda_{1}^{\star},\lambda_{2}^{\star})$ and we
have that

$\displaystyle\lambda_{1}^{\star}$
$\displaystyle:=\left(1-\frac{\tau_{1}\tau_{2}}{(1/\gamma+\tau_{1})(1/\gamma+\tau_{2})}\right)^{-1}\left(\frac{\tau_{1}/\gamma}{1/\gamma+\tau_{1}}c_{1}-\frac{\tau_{1}/\gamma}{1/\gamma+\tau_{1}}\frac{\tau_{2}}{1/\gamma+\tau_{2}}c_{2}\right)$
(10) $\displaystyle\lambda_{2}^{\star}$
$\displaystyle:=\left(1-\frac{\tau_{1}\tau_{2}}{(1/\gamma+\tau_{1})(1/\gamma+\tau_{2})}\right)^{-1}\left(\frac{\tau_{2}/\gamma}{1/\gamma+\tau_{2}}c_{2}-\frac{\tau_{1}/\gamma}{1/\gamma+\tau_{1}}\frac{\tau_{2}}{1/\gamma+\tau_{2}}c_{1}\right)$
(11)

where

$\displaystyle
c_{1}:=\log\left(\frac{\langle\exp(-\tilde{f}_{1}/\tau_{1}),a\rangle}{\langle\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2})),\xi^{(3)}\rangle}\right),\quad\text{and}\quad
c_{2}:=\log\left(\frac{\langle\exp(-\tilde{f}_{2}/\tau_{2}),a\rangle}{\langle\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2})),\xi^{(3)}\rangle}\right).$

We are now ready to perform an alternate maximization scheme on the
translation invariant formulation of the dual $\mathcal{D}_{\textrm{TI}}$.
Indeed using Danskin’s theorem (under the assumption that
$\lambda_{1}^{*},\lambda_{2}^{*}$ do not diverge), one obtains a variant of
Algorithm 6, summarized in Algorithm 3.

Inputs: $a,b,\xi^{(3)},u_{1},v_{1},u_{2},v_{2},\gamma,\tau_{1},\tau_{2}\\\
\tilde{u}_{1}\leftarrow u_{1}^{-1/\gamma/\tau_{1}},~{}\tilde{u}_{2}\leftarrow
u_{2}^{-1/\gamma/\tau_{2}}\\\
c_{1}\leftarrow\log(\langle\tilde{u}_{1},a\rangle)-\log(\langle\xi^{(3)},v_{1}^{-1}\odot
v_{2}^{-1}\rangle),~{}~{}c_{2}\leftarrow\log(\langle\tilde{u}_{2},b\rangle)-\log(\langle\xi^{(3)},v_{1}^{-1}\odot
v_{2}^{-1}\rangle)$
Result: $\lambda_{1}^{\star},~{}~{}\lambda_{2}^{\star}$ as in (10), (11)

Algorithm 2 $\text{compute-
lambdas}(a,b,\xi^{(3)},u_{1},v_{1},u_{2},v_{2},\gamma,\tau_{1},\tau_{2})$

Inputs:
$a,b,\bm{\xi}=(\xi^{(1)},\xi^{(2)},\xi^{(3)}),\gamma,\tau_{1},\tau_{2},\delta$
$v_{1}=v_{2}=\mathbf{1}_{r}$, $u_{1}=\mathbf{1}_{n}$, $u_{2}=\mathbf{1}_{m}$
repeat

$\tilde{v}_{1}=v_{1},~{}\tilde{v}_{2}=v_{2},\tilde{u}_{1}=u_{1},\tilde{u}_{2}=u_{2}\\\
\lambda_{1},\lambda_{2}\leftarrow\text{compute-
lambdas}(a,b,\xi^{(3)},u_{1},v_{1},u_{2},v_{2},\gamma,\tau_{1},\tau_{2})~{}(\text{Alg.~{}\ref{alg-
compute-lamb})}\\\
u_{1}=\left(\frac{a}{\xi^{(1)}v_{1}}\right)^{\frac{\tau_{1}}{\tau_{1}+1/\gamma}}\exp(-\lambda_{1}/\tau_{1})^{\frac{\tau_{1}}{1/\gamma+\tau_{1}}},\quad
u_{2}=\left(\frac{b}{\xi^{(2)}v_{2}}\right)^{\frac{\tau_{2}}{\tau_{2}+1/\gamma}}\exp(-\lambda_{2}/\tau_{2})^{\frac{\tau_{2}}{1/\gamma+\tau_{2}}},\\\
\lambda_{1},\lambda_{2}\leftarrow\text{compute-
lambdas}(a,b,\xi^{(3)},u_{1},v_{1},u_{2},v_{2},\gamma,\tau_{1},\tau_{2})~{}(\text{Alg.~{}\ref{alg-
compute-lamb})}\\\
g=\exp(\gamma(\lambda_{1}+\lambda_{2}))^{1/3}\left(\xi^{(3)}\odot(\xi^{(1)})^{T}u_{1}\odot(\xi^{(2)})^{T}u_{2}\right)^{1/3},~{}v_{1}=\frac{g}{(\xi^{(1)})^{T}u_{1}},~{}v_{2}=\frac{g}{(\xi^{(2)})^{T}u_{2}}$

until
_$\frac{1}{\gamma}\max(\|\log(u_{i}/\tilde{u}_{i})\|_{\infty},\|\log(v_{i}/\tilde{v}_{i})\|_{\infty})
<\delta$_;

Result:
$\operatorname{diag}(u_{1})\xi^{(1)}_{k}\operatorname{diag}(v_{1}),~{}~{}\operatorname{diag}(u_{2})\xi^{(2)}_{k}\operatorname{diag}(v_{2}),~{}~{}g$

Algorithm 3 $\text{ULR-TI-
Dykstra}(a,b,\bm{\xi},\gamma,\tau_{1},\tau_{2},\delta)$

###  3.3 Unbalanced Low-rank Gromov-Wasserstein

The low-rank Gromov-Wasssertein (LGW) problem (Scetbon et al., 2022) between
the two discrete metric measure spaces $(\mu,d_{\mathcal{X}})$ and
$(\nu,d_{\mathcal{Y}})$, written for compactness using $(a,A)$ and $(b,B)$,
reads

$\displaystyle\text{LGW}_{r}((a,A),(b,B))=\\!\\!\min_{P\in\Pi_{a,b}(r)}\\!\\!\mathcal{Q}_{A,B}(P),$
(12)

Following § 3.1, we introduce the unbalanced low-rank Gromov-Wasserstein
problem (ULGW). There is, however, an important difference with (12): When $P$
is constrained to be in $\Pi_{a,b}$, the first two terms of the RHS in (12)
simplify to $a^{T}A^{\odot 2}a+b^{T}B^{\odot 2}b$. Hence, they are constant
and can be discarded when optimizing. In an unbalanced setting, these terms
vary and must be accounted for.

$\displaystyle\text{ULGW}_{r}((a,A),(b,B))=\\!\\!\min_{(Q,R,g)\in\Pi_{r}}\\!\\!\langle
A^{\odot 2}Q\mathbf{1}_{r},Q\mathbf{1}_{r}\rangle+\langle B^{\odot
2}R\mathbf{1}_{r},R\mathbf{1}_{r}\rangle$ (13) $\displaystyle-2\langle
AQ\operatorname{diag}(1/g)R^{T}B,Q\operatorname{diag}(1/g)R^{T}\rangle+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$

To solve the problem, we apply the same scheme as proposed for ULOT, that is a
proximal gradient descent where we linearize $\mathcal{Q}_{A,B}$ and add a KL
penalization while leaving the soft marginal constraints unchanged. Therefore
the algorithm to solve ULGW is the exact same as the one solving ULOT,
however, the kernels $\bm{\xi}_{k}$ now take into account the quadratic scores
of the original LGW problem. More formally at each iteration $k$ of the outer
loop, we propose to solve

$\displaystyle(Q_{k+1},R_{k+1},g_{k+1}):=\operatorname*{argmin}_{\bm{\zeta}\in\Pi_{r}}\frac{1}{\gamma_{k}}\text{KL}(\bm{\zeta},\bm{\xi}_{k})+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$
(14)

where $(Q_{0},R_{0},g_{0})\in\Pi_{r}$ is an initial point,
$(\gamma_{k})_{k\geq 0}$ is a sequence of positive step sizes,
$P_{k}=Q_{k}\operatorname{diag}(1/g_{k})R_{k}^{T}$,
$\bm{\xi}_{k}:=(\xi_{k}^{(1)},\xi_{k}^{(2)},\xi_{k}^{(3)})$ and

$\displaystyle\xi_{k}^{(1)}$
$\displaystyle:=Q_{k}\odot\exp(-2\gamma_{k}Q_{k}\mathbf{1}_{r}\mathbf{1}_{r}^{T})\odot\exp(-4\gamma_{k}AP_{k}BR_{k}\operatorname{diag}(1/g_{k})))$
$\displaystyle\xi_{k}^{(2)}$
$\displaystyle:=R_{k}\odot\exp(-2\gamma_{k}R_{k}\mathbf{1}_{r}\mathbf{1}_{r}^{T})\odot\exp(-4\gamma_{k}BP_{k}^{T}AQ_{k}\operatorname{diag}(1/g_{k})))$
$\displaystyle\xi_{k}^{(3)}$
$\displaystyle:=g_{k}\odot\exp(4\gamma_{k}\omega_{k}/g_{k}^{2})\quad\text{with}\quad[\omega_{k}]_{i}:=[Q_{k}^{T}AP_{k}BR_{k}]_{i,i}~{}~{}\forall~{}i\in\\{1,\dots,r\\}.$

Note that (14) is the exact same optimization problem as (7), where only
$\bm{\xi}_{k}$ has changed and therefore can be solved using Algorithm 3.
Algorithm 4 summarizes our strategy to solve (13).

Inputs: $A,B,a,b,r,\gamma_{0},\tau_{1},\tau_{2},\delta$
$Q,R,g\leftarrow\text{Initialization as proposed
in~{}\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{scetbon2022lot}{\@@citephrase{,
}}{})}}$
repeat

$\tilde{Q}=Q,~{}~{}\tilde{R}=R,~{}~{}\tilde{g}=g,\\\
\nabla_{Q}=4AQ\operatorname{diag}(1/g)R^{T}BR\operatorname{diag}(1/g)+2Q\mathbf{1}_{r}\mathbf{1}_{r}^{T},\\\
\nabla_{R}=4BR\operatorname{diag}(1/g)Q^{T}AQ\operatorname{diag}(1/g)+2R\mathbf{1}_{r}\mathbf{1}_{r}^{T},\\\
\omega\leftarrow\mathcal{D}(Q^{T}AQ\operatorname{diag}(1/g)R^{T}BR),~{}~{}\nabla_{g}=-\omega/g^{2},\\\
\gamma\leftarrow\gamma_{0}/\max(\|\nabla_{Q}\|_{\infty}^{2},\|\nabla_{R}\|_{\infty}^{2},\|\nabla_{g}\|_{\infty}^{2}),\\\
\xi^{(1)}\leftarrow Q\odot\exp(-\gamma\nabla_{Q}),~{}\xi^{(2)}\leftarrow
R\odot\exp(-\gamma\nabla_{R}),~{}\xi^{(3)}\leftarrow
g\odot\exp(-\gamma_{k}\nabla_{g}),\\\ Q,R,g\leftarrow\text{ULR-TI-
Dykstra}(a,b,\bm{\xi},\gamma,\tau_{1},\tau_{2},\delta)~{}(\text{Alg.~{}\ref{alg-
dykstra-ti-ulot})}$

until  _$\Delta((Q,R,g),(\tilde{Q},\tilde{R},\tilde{g}),\gamma) <\delta$_;

Result: $Q,R,g$

Algorithm 4 $\text{ULGW}(A,B,a,b,r,\gamma_{0},\tau_{1},\tau_{2},\delta)$

Convergence and Complexity. Similarly to linear ULOT, the unbalanced Dykstra
algorithm is guaranteed to converge (Bauschke and Lewis, 2000). in addition,
(Scetbon et al., 2022) prove the convergence of their scheme to a stationary
point of the problem. Concerning the complexity, as we use Algorithm 6 with
the same complexity, we obtain therefore the exact same complexity in terms of
time of memory to solve these inner problems. The slight variation in kernel
$\bm{\xi}$ compared to ULOT still retains the same
$\mathcal{O}((n^{2}+m^{2})r)$ time and $\mathcal{O}(n^{2}+m^{2})$ memory
complexities. However as in ULOT, we can take advantage of low-rank
approximations of the costs matrices $A$ and $B$ in order to reach linear
complexity. Indeed, assuming $A\simeq A_{1}A_{2}^{T}$ and $B\simeq B_{1}B_{2}$
where $A_{1},A_{2}\in\mathbb{R}^{n\times d_{X}}$ and
$B_{1},B_{2}\in\mathbb{R}^{m\times d_{Y}}$, then the total time and memory
complexities become respectively $\mathcal{O}(mr(r+d_{Y})+nr(r+d_{X}))$ and
$\mathcal{O}((n+m)(r+d_{X}+d_{Y}))$. Again, when $A$ and $B$ are distance
matrices, we use the algorithms from (Indyk et al., 2019).

###  3.4 Unbalanced Low-rank Fused-Gromov-Wasserstein

We finally focus on the increasingly popular (Klein et al., 2023) Fused-
Gromov-Wasserstein problem, which merges linear and quadratic objectives
(Vayer et al., 2018):

$\displaystyle\text{FGW}(\mu,\nu)$
$\displaystyle:=\min_{P\in\Pi_{a,b}}\alpha\langle
C,P\rangle+\bar{\alpha}\mathcal{Q}_{A,B}(P)$ (15)

where $\alpha\in[0,1]$ and $\bar{\alpha}:=1-\alpha$ allows interpolating
between the GW and linear OT geometries. This problem remains a GW problem,
where one replaces the 4-way cost
$M[i,i^{\prime},j,j^{\prime}]:=(A_{i,i^{\prime}}-B_{j,j^{\prime}})^{2}$
appearing in (4) by a composite interpolated cost between the OT and GW
geometries, redefined as $M[i,i^{\prime},j,j^{\prime}]=\alpha
C_{i,j}+\bar{\alpha}(A_{i,i^{\prime}}-B_{j,j^{\prime}})^{2}$. Our proposed
unbalanced and low-rank version of the FGW problem includes $|P|:=\|P\|_{1}$
the mass of $P$, to homogenize linear and quadratic terms,

$\displaystyle\text{ULFGW}_{r}(\mu,\nu)$
$\displaystyle:=\\!\\!\\!\\!\min_{P\text{:}~{}\text{rk}_{+}(P)\leq
r}\\!\\!\\!\\!\alpha|P|\langle
C,P\rangle+\bar{\alpha}\mathcal{Q}_{A,B}(P)+\tau_{1}\text{KL}(P\mathbf{1}_{m}|a)+\tau_{2}\text{KL}(P^{T}\mathbf{1}_{n}|b)\,,$
(16)

which is expanded through the explicit factorization of $P$, noticing that
$|P|=|g|:=\|g\|_{1}$:

$\displaystyle\text{ULFGW}_{r}(\mu,\nu)$
$\displaystyle:=\min_{(Q,R,g)\in\Pi_{r}}\alpha|g|\mathcal{L}_{C}(Q,R,g)+\bar{\alpha}\mathcal{Q}_{A,B}(Q,R,g)+\mathcal{G}_{a,b}(Q,R,g)$
(17)

Then by linearizing again
$\mathcal{H}:(Q,R,g)\to\alpha|g|\mathcal{L}_{C}(Q,R,g)+\bar{\alpha}\mathcal{Q}_{A,B}(Q,R,g)$
with an added KL penalty and leaving $\mathcal{G}_{a,b}$ unchanged, we obtain
at each iteration, the same optimization problem as in (14) where the kernel
$\bm{\xi}_{k}$ is now defined as

$\displaystyle\bm{\xi}_{k}$
$\displaystyle:=(\xi_{k}^{(1)},\xi_{k}^{(2)},\xi_{k}^{(3)}),$
$\displaystyle\xi_{k}^{(1)}$
$\displaystyle:=Q_{k}\odot\exp(-\gamma_{k}\nabla_{Q}\mathcal{H}_{k}),~{}\xi_{k}^{(2)}:=R_{k}\odot\exp(-\gamma_{k}\nabla_{Q}\mathcal{H}_{k}),~{}\xi_{k}^{(3)}:=g_{k}\odot\exp(-\gamma_{k}\nabla_{g}\mathcal{H}_{k})$
$\displaystyle\nabla_{Q}\mathcal{H}_{k}$
$\displaystyle:=\alpha|g_{k}|CR_{k}\operatorname{diag}(1/g_{k})+\bar{\alpha}\left(2Q_{k}\mathbf{1}_{r}\mathbf{1}_{r}^{T}+4AP_{k}BR_{k}\operatorname{diag}(1/g_{k})\right)$
$\displaystyle\nabla_{R}\mathcal{H}_{k}$
$\displaystyle:=\alpha|g_{k}|C^{T}Q_{k}\operatorname{diag}(1/g_{k})+\bar{\alpha}\left(2R_{k}\mathbf{1}_{r}\mathbf{1}_{r}^{T}+4BP_{k}^{T}AQ_{k}\operatorname{diag}(1/g_{k})\right)$
$\displaystyle\nabla_{g}\mathcal{H}_{k}$ $\displaystyle:=\alpha\left(\langle
C,P_{k}\rangle\mathbf{1}_{r}-|g_{k}|\omega_{k}^{\text{lin}}/g_{k}^{2}\right)-4\bar{\alpha}\omega_{k}^{\text{quad}}/g_{k}^{2}$
$\displaystyle[\omega_{k}^{\text{lin}}]_{i}$
$\displaystyle:=[Q_{k}^{T}CR_{k}]_{i,i},~{}~{}[\omega_{k}^{\text{quad}}]_{i}:=[Q_{k}^{T}AP_{k}BR_{k}]_{i,i}~{}~{}\forall~{}i\in\\{1,\dots,r\\}\;.$

See Algorithm 5 for our implementatoin for the unbalanced and low-rank fused
GW problem.

Inputs: $A,B,C,a,b,r,t,\gamma_{0},\tau_{1},\tau_{2},\delta$
$Q,R,g\leftarrow\text{Initialization as proposed
in~{}\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{scetbon2022lot}{\@@citephrase{,
}}{})}}$
repeat

$\tilde{Q}=Q,~{}~{}\tilde{R}=R,~{}~{}\tilde{g}=g,\\\
\nabla_{Q}=\alpha|g|CR\operatorname{diag}(1/g)+\bar{\alpha}\left(2Q\mathbf{1}_{r}\mathbf{1}_{r}^{T}+4AQ\operatorname{diag}(1/g)R^{T}BR\operatorname{diag}(1/g)\right),\\\
\nabla_{R}=\alpha|g|C^{T}Q\operatorname{diag}(1/g)+\bar{\alpha}\left(2R\mathbf{1}_{r}\mathbf{1}_{r}^{T}+4BR\operatorname{diag}(1/g)Q^{T}AQ\operatorname{diag}(1/g)\right),\\\
\omega^{\text{lin}}\leftarrow\mathcal{D}(Q^{T}CR),~{}~{}\omega^{\text{quad}}\leftarrow\mathcal{D}(Q^{T}AQ\operatorname{diag}(1/g)R^{T}BR)\\\
\nabla_{g}=\alpha\left(\langle
C,Q\operatorname{diag}(1/g)R^{T}\rangle\mathbf{1}_{r}-|g_{k}|\omega^{\text{lin}}/g^{2}\right)-4\bar{\alpha}\omega^{\text{quad}}/g^{2},\\\
\gamma\leftarrow\gamma_{0}/\max(\|\nabla_{Q}\|_{\infty}^{2},\|\nabla_{R}\|_{\infty}^{2},\|\nabla_{g}\|_{\infty}^{2}),\\\
\xi^{(1)}\leftarrow Q\odot\exp(-\gamma\nabla_{Q}),~{}\xi^{(2)}\leftarrow
R\odot\exp(-\gamma\nabla_{R}),~{}\xi^{(3)}\leftarrow
g\odot\exp(-\gamma_{k}\nabla_{g}),\\\ Q,R,g\leftarrow\text{ULR-TI-
Dykstra}(a,b,\bm{\xi},\gamma,\tau_{1},\tau_{2},\delta)~{}(\text{Alg.~{}\ref{alg-
dykstra-ti-ulot})}$

until  _$\Delta((Q,R,g),(\tilde{Q},\tilde{R},\tilde{g}),\gamma) <\delta$_;

Result: $Q,R,g$

Algorithm 5 $\text{ULFGW}(A,B,a,b,r,\gamma_{0},\tau_{1},\tau_{2},\delta)$

######  Remark 2.

Note again that here, we have in general a quadratic complexity both in time
and memory, and as soon as we are provided a low-rank approximation of the
matrices $C,A,B$, our proposed algorithm scales linearly with respect to the
number of points $n$ and $m$.

##  4 Experiments

Our goal in Exp. 1 is to compare unbalanced and low-rank (ULR) solvers to
balanced and low-rank (LR) counterparts, in Exp. 2, compare ULR solvers to
entropic (E) counterparts, and in Exp. 3, compare our ULR solvers to (Thual et
al., 2022), which can learn a sparse transport map, in the unbalanced FGW
setting.

Datasets. We run the experiments on two real world datasets, described in B.1,
that are large enough to showcase our solvers. In particular, they consist of
both a shared feature space, used to compute the costs matrices for the linear
term in the OT and FGW settings, as well as geometries specific to each source
$s$ and target $t$ data, that are used to compute the costs matrices for the
quadratic term in the GW and FGW settings. We leverage mouse brain STARmap
spatial transcriptomics data from (Shi et al., 2022) for Exp. 1 and Exp. 2.
For Exp. 3 we use data from the Individual Brain Charting dataset (Pinho et
al., 2018), to recapitulate the settings of (Thual et al., 2022).

Metrics. Following Klein et al. (2023), we evaluate maps by focusing on the
two following metrics: (i) pearson correlation $\rho$ computed between the
source $s$ feature matrix $F^{s}$ and the barycentric projection of the target
$t$ to the source scaled by the target marginals $b^{t}$: $T^{T}_{t\rightarrow
s}\left(F^{t}\right.\frac{1}{b^{t}})$; (ii) F1 score computed between the
original source $s$ labels $l^{s}$ and the inferred source labels, computed by
taking the $\operatorname*{argmax}_{j}B_{i,j}$ of the barycentric projection
of the target $t$ one hot encoded labels $L^{t}$, scaled by the target
marginal $b^{t}$, to the source $T^{T}_{t\rightarrow
s}\left(L^{t}\right.\frac{1}{b^{t}})$.

###  4.1 Experiment 1: ULOT vs. LOT on gene expression / cell type annotation
mapping

solver | mass pct | val $\rho$ | test $\rho$ | F1 macro | F1 micro | F1 weighted
---|---|---|---|---|---|---
LOT | 1.000 | 0.282 | 0.386 | 0.210 | 0.411 | 0.360
ULOT | 0.889 | 0.301 | 0.409 | 0.200 | 0.425 | 0.363
LGW | 1.000 | 0.227 | 0.288 | 0.487 | 0.716 | 0.692
ULGW | 1.001 | 0.222 | 0.287 | 0.463 | 0.701 | 0.665
LFGW | 1.000 | 0.365 | 0.443 | 0.576 | 0.720 | 0.714
ULFGW | 0.443 | 0.379 | 0.463 | 0.582 | 0.733 | 0.724

Table 1: Results for spatial transcriptomics dataset (brain coronal section
from Shi et al. (2022)).

(a) Visualization of measured and predicted gene expression of Nrgn.

(b) Visualization of measured and predicted gene expression of Slc24a2.

Figure 1: Spatial visualization of the two mouse brain sections used in Exp. 1

Here, we evaluate the accuracy of ULOT solvers for a large-scale spatial
transcriptomics task, using gene expression mapping and cell type annotation.
We compare it to the balanced LR alternative using the Pearson correlation
$\rho$ as described in the metrics section. We leverage two coronal sections
of the mouse brain profiled by STARmap spatial transcriptomics by (Shi et al.,
2022). They consist of $n\approx 40,000$ cells in both the source and target
brain section. Each cell is described by 1000 gene features, in addition to 2D
spatial coordinates. As a result $A,B$ are $\approx 40k\times 40k$, and the
fused term $C$ is a squared-Euclidean distance matrix on 30D PCA space
computed on the gene expression space. We selected 10 marker genes for the
validation and test sets from the HPF_CA cluster.

Figure 2: Visualization of measured and predicted tissue regions in the mouse
brain in Exp. 1

We run an extensive grid search as reported in B.2, we pick the best
hyperparameters combination using performance on the $10$ validation genes as
a criterion, and we report that metric on the other genes in Table 1, as well
as qualitative results in Figure 1 and Figure 2. Clearly, ULFGW is the best
performing solver across all metrics. Interestingly, the ULOT does not
consistently outperforms its balanced version, and unbalancedness seems to
hurt performance for the LGW solvers. Nevertheless, both solvers display
inconsistent performance across metrics, whereas the ULFGW and LFGW are
consistently superior to the rest of the solvers. These results highlight how
the flexibility given by the FGW formulation to leverage common and disparate
geometries, paired with the unbalancedness relaxation, can provide state of
the art algorithms for matching problems in large-scale, real world biological
problems.

###  4.2 Experiment 2: ULOT vs. UEOT

In this experiment, we evaluate the performance of ULOT solvers to the
unbalanced entropic alternative (UEOT). We use the same datasets as 4.1, but
pick a smaller subset (Olfactory bulb), to avoid OOM errors for entropic UGW
solvers, that cannot handle the $40k$ sizes considered previously (see B.1).

solver | mass | val $\rho$ | test $\rho$ | F1-mac | F1-mic | F1-wei
---|---|---|---|---|---|---
UEOT | 1.012 | 0.368 | 0.479 | 0.511 | 0.763 | 0.751
LOT | 1.000 | 0.335 | 0.440 | 0.511 | 0.760 | 0.751
ULOT | 0.998 | 0.356 | 0.461 | 0.518 | 0.770 | 0.762
UEFGW | 1.015 | 0.343 | 0.475 | 0.564 | 0.839 | 0.831
LFGW | 1.000 | 0.348 | 0.453 | 0.512 | 0.762 | 0.753
ULFGW | 0.339 | 0.368 | 0.491 | 0.556 | 0.826 | 0.818

Table 2: Results for spatial transcriptomics dataset (Olfactory bulb section
from Shi et al. (2022)).

They consist of $n\approx 20,000$ cells in the source and $\approx 15,000$
cells in the target sections, and 1000 genes. Similar to Exp.1, the fused term
$C$ is a squared-Euclidean distance matrix on 30D PCA space computed on the
gene expression space. As done in the previous experiment, we select 10 marker
genes for the validation and 10 genes for the test set, from cluster OB_1. We
run an extensive grid search as described for Exp.2 in 4.1 and B.2. In Table
2, we see that ULFGW outperforms entropic solvers w.r.t. $\rho$, but is worse
when considering the F1 scores. On the other hand, ULFGW confirms its
superiority compared to the balanced alternative LFGW. Taken together, these
results suggest that while unbalanced LR solvers are on par with unbalanced
entropic solvers in terms of performance on small data regimes, they unlock
the applications of unbalanced OT to large scale datasets.

###  4.3 Experiment 3: ULOT to align brain meshes

Thual et al. (2022) proposed a novel formulation for the unbalanced FGW
problem. Their proposal is showcased to align brain anatomies, and their
functional signal (FUGW). We compare it to our ULFGW solver, using the same
experimental setting, see Table 3.

solver | mass | val $\rho$ | test $\rho$
---|---|---|---
FUGW-sparse | 0.999 | 0.492 | 0.472
LFGW | 1.000 | 0.513 | 0.663
ULFGW | 0.981 | 0.533 | 0.643

Table 3: Results on the brain anatomy with functional signal data from Pinho
et al. (2018) in Exp.3. Figure 3: Visualization of measured and predicted
right auditory click contrast map in Exp.3.

### Conclusion

Recent practical successes of OT methods to natural sciences have demonstrated
the relevance of OT to their analysis pipelines, but have also shown,
repeatedly, that a certain degree of freedom to depart from the rigid
assumption of mass conservation is needed in practice. On the other hand, and
across the same range of applications, low-rank approaches can hold the
promise of scaling OT methods to relevant sample sizes for natural sciences.
This paper merges these two strains and demonstrate the practical relevance of
these novel algorithms.

## References

  * Ahuja et al. (1993) Ravindra K Ahuja, Thomas L Magnanti, and James B Orlin.  _Network Flows: Theory, Algorithms, and Applications_.  Prentice hall, 1993.
  * Altschuler et al. (2017) Jason Altschuler, Jonathan Weed, and Philippe Rigollet.  Near-linear time approximation algorithms for optimal transport via sinkhorn iteration.  _arXiv preprint arXiv:1705.09634_ , 2017.
  * Bauschke and Combettes (2008) Heinz H Bauschke and Patrick L Combettes.  A Dykstra-like algorithm for two monotone operators.  _Pacific Journal of Optimization_ , 4(3):383–391, 2008.
  * Bauschke and Lewis (2000) Heinz H Bauschke and Adrian S Lewis.  Dykstras algorithm with bregman projections: A convergence proof.  _Optimization_ , 48(4):409–427, 2000.
  * Blumberg et al. (2020) Andrew J Blumberg, Mathieu Carriere, Michael A Mandell, Raul Rabadan, and Soledad Villar.  Mrec: a fast and versatile framework for aligning and matching point clouds with applications to single cell molecular data.  _arXiv preprint arXiv:2001.01666_ , 2020.
  * Boissard and Le Gouic (2014) Emmanuel Boissard and Thibaut Le Gouic.  On the mean speed of convergence of empirical and occupation measures in wasserstein distance.  In _Annales de l’IHP Probabilités et statistiques_ , volume 50, pages 539–563, 2014.
  * Bunne et al. (2021) Charlotte Bunne, Stefan G. Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong-Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch.  Learning single-cell perturbation responses using neural optimal transport.  _bioRxiv_ , 2021.  doi: 10.1101/2021.12.15.472775.  URL https://www.biorxiv.org/content/early/2021/12/15/2021.12.15.472775.
  * Bunne et al. (2022) Charlotte Bunne, Andreas Krause, and Marco Cuturi.  Supervised Training of Conditional Monge Maps.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Cang et al. (2023) Zixuan Cang, Yanxiang Zhao, Axel A Almet, Adam Stabell, Raul Ramos, Maksim V Plikus, Scott X Atwood, and Qing Nie.  Screening cell–cell communication in spatial transcriptomics via collective optimal transport.  _Nature Methods_ , 20(2):218–228, 2023.
  * Caron et al. (2020) Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin.  Unsupervised learning of visual features by contrasting cluster assignments.  _Advances in neural information processing systems_ , 33:9912–9924, 2020.
  * Chizat et al. (2018) Lénaïc Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard.  Unbalanced optimal transport: Dynamic and kantorovich formulations.  _Journal of Functional Analysis_ , 274(11):3090–3123, 2018.
  * Chizat et al. (2020) Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xavier Vialard, and Gabriel Peyré.  Faster wasserstein distance estimation with the sinkhorn divergence.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Cuturi (2013) Marco Cuturi.  Sinkhorn distances: Lightspeed computation of optimal transport.  In _Advances in neural information processing systems_ , pages 2292–2300, 2013.
  * Cuturi et al. (2022) Marco Cuturi, Laetitia Meng-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul.  Optimal transport tools (ott): A jax toolbox for all things wasserstein.  _arXiv preprint arXiv:2201.12324_ , 2022.
  * De Silva and Tenenbaum (2004) Vin De Silva and Joshua B Tenenbaum.  Sparse multidimensional scaling using landmark points.  Technical report, technical report, Stanford University, 2004.
  * Demetci et al. (2020) Pinar Demetci, Rebecca Santorella, Björn Sandstede, William Stafford Noble, and Ritambhara Singh.  Gromov-wasserstein optimal transport to align single-cell multi-omics data.  _bioRxiv_ , 2020.  doi: 10.1101/2020.04.28.066787.
  * Dudley et al. (1966) Richard Mansfield Dudley et al.  Weak convergence of probabilities on nonseparable metric spaces and empirical measures on euclidean spaces.  _Illinois Journal of Mathematics_ , 10(1):109–126, 1966.
  * Dykstra (1983) Richard L Dykstra.  An algorithm for restricted least squares regression.  _Journal of the American Statistical Association_ , 78(384):837–842, 1983.
  * Fatras et al. (2021) Kilian Fatras, Thibault Sejourne, Rémi Flamary, and Nicolas Courty.  Unbalanced minibatch optimal transport; applications to domain adaptation.  In Marina Meila and Tong Zhang, editors, _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 3186–3197. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/fatras21a.html.
  * Feydy et al. (2019) Jean Feydy, Pierre Roussillon, Alain Trouvé, and Pietro Gori.  Fast and scalable optimal transport for brain tractograms.  In _Medical Image Computing and Computer Assisted Intervention–MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13–17, 2019, Proceedings, Part III 22_ , pages 636–644. Springer, 2019\.
  * Flamary et al. (2021) Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer.  Pot: Python optimal transport.  _Journal of Machine Learning Research_ , 22(78):1–8, 2021.  URL http://jmlr.org/papers/v22/20-451.html.
  * Forrow et al. (2018) Aden Forrow, Jan-Christian Hütter, Mor Nitzan, Philippe Rigollet, Geoffrey Schiebinger, and Jonathan Weed.  Statistical optimal transport via factored couplings, 2018.
  * Frogner et al. (2015) Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio.  Learning with a Wasserstein loss.  In _Advances in Neural Information Processing Systems_ , pages 2053–2061, 2015.
  * Genevay et al. (2019) Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré.  Sample Complexity of Sinkhorn Divergences.  In _International Conference on Artificial Intelligence and Statistics (AISTATS)_ , volume 22, 2019.
  * Indyk et al. (2019) Piotr Indyk, Ali Vakilian, Tal Wagner, and David Woodruff.  Sample-optimal low-rank approximation of distance matrices, 2019.
  * Klein et al. (2023) Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng-Papaxanthos, Michael Sterr, Aimee Bastidas-Ponce, Marta Tarquis-Medina, Heiko Lickert, Mostafa Bakhti, Mor Nitzan, Marco Cuturi, and Fabian J. Theis.  Mapping cells through time and space with moscot.  _bioRxiv_ , 2023.  doi: 10.1101/2023.05.11.540374.  URL https://www.biorxiv.org/content/early/2023/05/11/2023.05.11.540374.
  * Kuhn (1955) Harold W Kuhn.  The hungarian method for the assignment problem.  _Naval research logistics quarterly_ , 2(1-2):83–97, 1955.
  * Li et al. (2023) Jiajin Li, Jianheng Tang, Lemin Kong, Huikang Liu, Jia Li, Anthony Man-Cho So, and Jose Blanchet.  A convergent single-loop algorithm for relaxation of gromov-wasserstein in graph data.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Mémoli (2011) Facundo Mémoli.  Gromov–Wasserstein distances and the metric approach to object matching.  _Foundations of Computational Mathematics_ , 11(4):417–487, 2011.
  * Mena and Niles-Weed (2019) Gonzalo Mena and Jonathan Niles-Weed.  Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Nitzan et al. (2019) Mor Nitzan, Nikos Karaiskos, Nir Friedman, and Nikolaus Rajewsky.  Gene expression cartography.  _Nature_ , 576(7785), 2019.
  * Oquab et al. (2023) Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al.  Dinov2: Learning robust visual features without supervision.  _arXiv preprint arXiv:2304.07193_ , 2023.
  * Peyré et al. (2016) Gabriel Peyré, Marco Cuturi, and Justin Solomon.  Gromov-wasserstein averaging of kernel and distance matrices.  In _International Conference on Machine Learning_ , pages 2664–2672, 2016.
  * Peyré and Cuturi (2019) Gabriel Peyré and Marco Cuturi.  Computational optimal transport.  _Foundations and Trends in Machine Learning_ , 11(5-6), 2019\.  ISSN 1935-8245.
  * Pinho et al. (2018) Ana Luísa Pinho, Alexis Amadon, Torsten Ruest, Murielle Fabre, Elvis Dohmatob, Isabelle Denghien, Chantal Ginisty, Séverine Becuwe-Desmidt, Séverine Roger, Laurence Laurier, Véronique Joly-Testault, Gaëlle Médiouni-Cloarec, Christine Doublé, Bernadette Martins, Philippe Pinel, Evelyn Eger, Gael Varoquaux, Christophe Pallier, Stanislas Dehaene, Lucie Hertz-Pannier, and Bertrand Thirion.  Individual Brain Charting, a high-resolution fMRI dataset for cognitive mapping.  _Scientific Data_ , 5:180105, June 2018.  doi: 10.1038/sdata.2018.105.  URL https://hal.science/hal-01817528.
  * Rockafellar (1970) R Tyrrell Rockafellar.  _Convex analysis_.  Princeton university press, 1970.
  * Rubner et al. (2000) Yossi Rubner, Carlo Tomasi, and Leonidas J. Guibas.  The earth mover’s distance as a metric for image retrieval.  _International Journal of Computer Vision_ , 40(2):99–121, November 2000.
  * Sander et al. (2022) Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré.  Sinkformers: Transformers with doubly stochastic attention.  In _International Conference on Artificial Intelligence and Statistics_ , pages 3515–3530. PMLR, 2022.
  * Sato et al. (2020) Ryoma Sato, Marco Cuturi, Makoto Yamada, and Hisashi Kashima.  Fast and robust comparison of probability measures in heterogeneous spaces.  _arXiv preprint arXiv:2002.01615_ , 2020.
  * Scetbon and Cuturi (2020) Meyer Scetbon and Marco Cuturi.  Linear time sinkhorn divergences using positive features.  _Advances in Neural Information Processing Systems_ , 33:13468–13480, 2020.
  * Scetbon and Cuturi (2022) Meyer Scetbon and Marco Cuturi.  Low-rank optimal transport: Approximation, statistics and debiasing.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, _Advances in Neural Information Processing Systems_ , volume 35, pages 6802–6814. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/2d69e771d9f274f7c624198ea74f5b98-Paper-Conference.pdf.
  * Scetbon et al. (2021) Meyer Scetbon, Marco Cuturi, and Gabriel Peyré.  Low-rank sinkhorn factorization.  In _International Conference on Machine Learning_ , pages 9344–9354. PMLR, 2021.
  * Scetbon et al. (2022) Meyer Scetbon, Gabriel Peyré, and Marco Cuturi.  Linear-time Gromov-Wasserstein distances using low rank couplings and costs.  _ICML_ , 2022.
  * Schiebinger et al. (2019) Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al.  Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming.  _Cell_ , 176(4):928–943, 2019.
  * Séjourné et al. (2021) Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré.  The unbalanced gromov wasserstein distance: Conic formulation and relaxation.  _Advances in Neural Information Processing Systems_ , 34:8766–8779, 2021.
  * Séjourné et al. (2022) Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré.  Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe.  In _International Conference on Artificial Intelligence and Statistics_ , pages 4995–5021. PMLR, 2022.
  * Shi et al. (2022) Hailing Shi, Yichun He, Yiming Zhou, Jiahao Huang, Brandon Wang, Zefang Tang, Peng Tan, Morgan Wu, Zuwan Lin, Jingyi Ren, Yaman Thapa, Xin Tang, Albert Liu, Jia Liu, and Xiao Wang.  Spatial atlas of the mouse central nervous system at molecular resolution.  _bioRxiv_ , 2022.  doi: 10.1101/2022.06.20.496914.  URL https://www.biorxiv.org/content/early/2022/06/22/2022.06.20.496914.
  * Solomon et al. (2016) Justin Solomon, Gabriel Peyré, Vladimir G Kim, and Suvrit Sra.  Entropic metric alignment for correspondence problems.  _ACM Transactions on Graphics (TOG)_ , 35(4):1–13, 2016.
  * Tay et al. (2020) Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan.  Sparse Sinkhorn attention.  In Hal Daumé III and Aarti Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 9438–9447. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/tay20a.html.
  * Thual et al. (2022) Alexis Thual, Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion.  Aligning individual brains with fused unbalanced Gromov-Wasserstein.  _arXiv_ , 2022.
  * Vayer et al. (2018) Titouan Vayer, Laetita Chapel, Rémi Flamary, Romain Tavenard, and Nicolas Courty.  Fused gromov-wasserstein distance for structured objects: theoretical foundations and mathematical properties.  _arXiv preprint arXiv:1811.02834_ , 2018.
  * Vincent-Cuaz et al. (2023) Cédric Vincent-Cuaz, Rémi Flamary, Marco Corneli, Titouan Vayer, and Nicolas Courty.  Semi-relaxed gromov-wasserstein divergence and applications on graphs.  In _International Conference on Learning Representations_ , 2023.
  * Wolf et al. (2018) F Alexander Wolf, Philipp Angerer, and Fabian J Theis.  SCANPY: large-scale single-cell gene expression data analysis.  _Genome Biology_ , 19(1), 2018.
  * Xu et al. (2019) Hongteng Xu, Dixin Luo, and Lawrence Carin.  Scalable gromov-wasserstein learning for graph partitioning and matching.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper_files/paper/2019/file/6e62a992c676f611616097dbea8ea030-Paper.pdf.

## Appendix

##  Appendix A Algorithms

Inputs:
$a,b,\bm{\xi}=(\xi^{(1)},\xi^{(2)},\xi^{(3)}),\gamma,\tau_{1},\tau_{2},\delta$
$v_{1}=v_{2}=\mathbf{1}_{r}$, $u_{1}=\mathbf{1}_{n}$, $u_{2}=\mathbf{1}_{m}$
repeat

$\tilde{v}_{1}=v_{1},~{}\tilde{v}_{2}=v_{2},\tilde{u}_{1}=u_{1},\tilde{u}_{2}=u_{2}\\\
u_{1}=\left(\frac{a}{\xi^{(1)}v_{1}}\right)^{\frac{\tau_{1}}{\tau_{1}+1/\gamma}},\quad
u_{2}=\left(\frac{b}{\xi^{(2)}v_{2}}\right)^{\frac{\tau_{2}}{\tau_{2}+1/\gamma}},\\\
g=\left(\xi^{(3)}\odot(\xi^{(1)})^{T}u_{1}\odot(\xi^{(2)})^{T}u_{2}\right)^{1/3},~{}v_{1}=\frac{g}{(\xi^{(1)})^{T}u_{1}},~{}v_{2}=\frac{g}{(\xi^{(2)})^{T}u_{2}}$

until
_$\frac{1}{\gamma}\max(\|\log(u_{i}/\tilde{u}_{i})\|_{\infty},\|\log(v_{i}/\tilde{v}_{i})\|_{\infty})
<\delta$_;

Result:
$\operatorname{diag}(u_{1})\xi^{(1)}_{k}\operatorname{diag}(v_{1}),~{}~{}\operatorname{diag}(u_{2})\xi^{(2)}_{k}\operatorname{diag}(v_{2}),~{}~{}g$

Algorithm 6 $\text{ULR-Dykstra}(a,b,\bm{\xi},\gamma,\tau_{1},\tau_{2},\delta)$

##  Appendix B Experiments

###  B.1 Datasets and preprocessing

We downloaded the two publicly available datasets from the respective
publications:

  * •

STARmap mouse brain sections from [Shi et al., 2022]

  * •

Brain mesh anatomy and functional signal from [Pinho et al., 2018]

We reprocessed the datasets using standard tools from the SCANPY pipeline
[Wolf et al., 2018]. Specifically, we log-normalized gene expression of all
genes present in dataset. We selected two brain coronal sections for Exp.1 and
two Coronal Olfactory Bulb (OB) sections for Exp.2, from the STARmap dataset.
For Exp.3, we used the meshes together with their functional signal of the
brains to recapitulate Exp.1 in [Thual et al., 2022]. A visualization of the
STARmap dataset for the two subsets used in Exp.1 and Exp.2 can be seen in
Figure 4 and an overview of the cell type proportions present in each of the
section pairs can be see in Figure 5. These visualization highlight the
differences in terms of spatial organization and cell type proportions of the
brain sections used in the experiment.

(a) Visualization of the brain coronal sections used in Exp.1.

(b) Visualization of the OB sections used in Exp.2.

Figure 4: Spatial visualization of the two mouse brain sections used in Exp.1.

(a) Visualization of cell type frequencies for the brain coronal sections used
in Exp.1.

(b) Visualization of cell type frequencies for the OB sections used in Exp.2.

Figure 5: Cell type frequencies of the datasets used in Exp.1 and Exp.2.

###  B.2 Experimental settings

For the brain signal alignment benchmark in Exp.3, we use the fsaverage7 mesh
containing 160k vertices. We compute an approximation of the geodesic
distances with Landmark Multi-Dimensional Scaling De Silva and Tenenbaum
[2004] using 2048 points in 30-dimensional embedding space. For fugw-sparse,
we compute the coupling in 2 stages:

  1. 1.

Similarly as in Thual et al. [2022], we subsample the mesh to 10% of the
points using Ward’s algorithm and compute the coarse optimal transport
coupling.

  2. 2.

We then use this coarse coupling to define a sparsity mask on the full mesh by
selecting for each source (target) vertex the most coupled target (source)
vertex and its neighbors within $\frac{4}{max\\_distance}$ radius using the
approximation of the geodesic distances. This mask is then used to compute the
fine-grained sparse coupling.

For all experiments, we ran the grid search as defined by 4 and selected the
best set of hyperparameters based on the validation correlation. We report
results of top performing hyperparameters for the evaluated algorithms in
Table 5 for Exp.1, Table 6 for Exp.2 and Table 7 for Exp.3

| values
---|---
rank | 10, 50, 100
reg (ours) | 0.0, 0.001, 0.01
reg (fugw-sparse) | 0.0001, 0.001, 0.01
tau1 | 0.1, 1.0, 100.0
tau2 | 0.1, 1.0, 100.0
Table 4: Hyperparameters considered in our grid-search.

solver | rank | tau1 | tau2 | temp | reg | mass | val $\rho$ | test $\rho$ | F1-mac | F1-mic | F1-wei
---|---|---|---|---|---|---|---|---|---|---|---
lot | 10 | - | - | 0.200 | 0.010 | 1.000 | 0.282 | 0.386 | 0.210 | 0.411 | 0.360
ulot | 10 | 1.000 | 1.000 | 0.200 | 0.010 | 0.889 | 0.301 | 0.409 | 0.200 | 0.425 | 0.363
lgw | 100 | - | - | 0.200 | 0.001 | 1.000 | 0.227 | 0.288 | 0.487 | 0.716 | 0.692
ulgw | 100 | 100.000 | 100.000 | 0.200 | 0.010 | 1.001 | 0.222 | 0.287 | 0.463 | 0.701 | 0.665
lfgw | 50 | - | - | 0.400 | 0.010 | 1.000 | 0.365 | 0.443 | 0.576 | 0.720 | 0.714
ulfgw | 100 | 0.100 | 0.100 | 0.400 | 0.001 | 0.443 | 0.379 | 0.463 | 0.582 | 0.733 | 0.724

Table 5: Results on the large spatial transcriptomics dataset (brain coronal
section from [Shi et al., 2022]).

solver | rank | tau1 | tau2 | temp | reg | mass | val $\rho$ | test $\rho$ | F1-mac | F1-mic | F1-wei
---|---|---|---|---|---|---|---|---|---|---|---
uot | - | 0.909 | 0.999 | 0.400 | 0.100 | 1.012 | 0.368 | 0.479 | 0.511 | 0.763 | 0.751
lot | 10 | - | - | 0.200 | 0.010 | 1.000 | 0.335 | 0.440 | 0.511 | 0.760 | 0.751
ulot | 10 | 1.000 | 100.000 | 0.200 | 0.010 | 0.998 | 0.356 | 0.461 | 0.518 | 0.770 | 0.762
ufgw | - | 0.500 | 0.999 | 0.600 | 0.100 | 1.015 | 0.343 | 0.475 | 0.564 | 0.839 | 0.831
lfgw | 10 | - | - | 0.600 | 0.010 | 1.000 | 0.348 | 0.453 | 0.512 | 0.762 | 0.753
ulfgw | 10 | 0.100 | 0.100 | 0.600 | 0.001 | 0.339 | 0.368 | 0.491 | 0.556 | 0.826 | 0.818

Table 6: Results on the small subset STARmap dataset (OB section from [Shi et al., 2022]). solver | rank | tau1 | tau2 | reg | reg | mass | val $\rho$ | test $\rho$
---|---|---|---|---|---|---|---|---
fugw-sparse | - | 1.000 | 0.100 | 0.200 | 0.01 | 0.999 | 0.492 | 0.472
lfgw | 100 | - | - | 0.600 | 0.000 | 1.000 | 0.513 | 0.663
ulfgw | 100 | 1.000 | 0.100 | 0.600 | 0.001 | 0.981 | 0.533 | 0.643
Table 7: Results on the brain anatomy and functional signal from [Pinho et al., 2018]). experiment |  | val $\rho$ | tst $\rho$ | F1-mac | F1-mic | F1-wei
---|---|---|---|---|---|---
Exp. 1 | mean | 0.362 | 0.449 | 0.546 | 0.687 | 0.677
std | 0.027 | 0.022 | 0.054 | 0.062 | 0.061
Exp. 2 | mean | 0.356 | 0.463 | 0.538 | 0.800 | 0.791
std | 0.008 | 0.018 | 0.021 | 0.031 | 0.032
Table 8: Effect of k-means initialization [Scetbon and Cuturi, 2022]. We
report mean and standard deviation of test criterion for ULFGW, with the best
hyperparameter on validation data for each experiment. We use 5 initial seeds
for Exp. 1. We observe more variability in validation performance for Exp. 2,
and therefore start with 10 seeds, pruning the lowest performing 5 seeds.

(a) Visualization of measured and predicted gene expression of Gad1.

(b) Visualization of measured and predicted gene expression of Gnrg4.

Figure 6: Measured and predicted gene expression for the small subset STARmap
dataset (OB section from [Shi et al., 2022]) for ULRFGW.

##  Appendix C Proofs

###  C.1 Proof of Proposition 1

Let $n,m\geq r\geq 1$, $\gamma>0$,
$,\bm{\xi}:=(\xi^{(1)},\xi^{(2)},\xi^{(3)})$ where
$\xi^{(1)}\in\mathbb{R}^{n\times r}_{+}$, $\xi^{(2)}\in\mathbb{R}^{m\times
r}_{+}$ and $\xi^{(3)}\in\mathbb{R}^{r}_{+}$ and let us recall that
$\text{KL}(\cdot,\cdot)$ is the generalized Kullback-Leibler divergence
defined as $\text{KL}(p|q):=\sum_{i}p_{i}\log(p_{i}/q_{i})+q_{i}-p_{i}$. Then
observe that

$\displaystyle\min_{(Q,R,g)\in\Pi_{r}}\frac{1}{\gamma}\left[\text{KL}(Q,\xi^{(1)})+\text{KL}(R,\xi^{(2)})+\text{KL}(g,\xi^{(3)})\right]+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$

is a convex problem satisfying the Slater’s condition and therefore strong
duality holds. Therefore we have:

$\displaystyle\min_{(Q,R,g)\in\Pi_{r}}\frac{1}{\gamma}\left[\text{KL}(Q,\xi^{(1)})+\text{KL}(R,\xi^{(2)})+\text{KL}(g,\xi^{(3)})\right]+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$
$\displaystyle=\sup_{\lambda_{1},\lambda_{2}}\min_{Q,R,g}\langle\lambda_{1},g-Q^{\top}\mathbf{1}_{n}\rangle+\langle\lambda_{2},g-R^{\top}\mathbf{1}_{m}\rangle+\frac{1}{\gamma}\left[\text{KL}(Q,\xi^{(1)})+\text{KL}(R,\xi^{(2)})+\text{KL}(g,\xi^{(3)})\right]$
$\displaystyle+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$
$\displaystyle=\sup_{\lambda_{1},\lambda_{2}}\min_{Q}\frac{1}{\gamma}\text{KL}(Q,\xi^{(1)})+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\langle\lambda_{1},-Q^{\top}\mathbf{1}_{n}\rangle$
$\displaystyle+\min_{R}\frac{1}{\gamma}\text{KL}(R,\xi^{(2)})+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)+\langle-\lambda_{2},R^{\top}\mathbf{1}_{m}\rangle+\min_{g}\frac{1}{\gamma}\text{KL}(g,\xi^{(3)})+\langle
g,\lambda_{1}+\lambda_{2}\rangle.$

Now consider

$\min_{g}\frac{1}{\gamma}\text{KL}(g,\xi^{(3)})+\langle
g,\lambda_{1}+\lambda_{2}\rangle$

and observe that this problem can be solved explicitly. The first-order
optimality condition gives us that
$g^{*}=\exp(-\gamma(\lambda_{1}+\lambda_{2}))\odot\xi^{(3)}$ solves the
problem and

$\displaystyle\min_{g}\frac{1}{\gamma}\text{KL}(g,\xi^{(3)})+\langle
g,\lambda_{1}+\lambda_{2}\rangle=-\frac{1}{\gamma}\langle\exp(-\gamma(\lambda_{1}+\lambda_{2})),\xi^{(3)}\rangle+\langle\xi^{(3)},\mathbf{1}\rangle.$

Let us now focus on the following convex optimization problem,

$\displaystyle\min_{Q}\frac{1}{\gamma}\text{KL}(Q,\xi^{(1)})+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\langle-\lambda_{1},Q^{\top}\mathbf{1}_{n}\rangle$
(18)

and note that it admits a unique solution due to the strict convexity of
$Q\to\text{KL}(Q,\xi^{(1)})$. Then by denoting
$F_{\tau,z}(s):=\tau\text{KL}(s|z)$ and $G_{\lambda}(s):=\langle
s,-\lambda\rangle$, and by applying the Fenchel-Rockafellar theorem
[Rockafellar, 1970], we obtain that strong duality holds, the dual problem of
(18) is

$\displaystyle\sup_{f_{1},h_{1}}-F_{\tau_{1},a}^{*}(-f_{1})-G_{\lambda_{1}}^{*}(-h_{1})-\frac{1}{\gamma}\langle\exp(\gamma(f_{1}+h_{1})),\xi^{(1)}\rangle$

and that $(f_{1},h_{1})$ solves the dual if and only if $-f_{1}\in\partial
F_{\tau_{1},a}(Q\mathbf{1}_{r})$,$-h_{1}\in\partial
G_{\lambda_{1}}(Q^{\top}\mathbf{1}_{n})$ and
$Q=\operatorname{diag}(\exp(\gamma
f_{1}))\xi^{(1)}\operatorname{diag}(\exp(\gamma h_{1}))$ where $Q$ is the
solution of (18). Recall that here we denote for any convex set
$X\in\mathbb{R}^{q}$ and function $f:X\to\mathbb{R}\cup\\{+\infty\\}$, $f^{*}$
its convex conjugate defined for any $y\in
X^{*}:=\\{x^{*}~{}\text{s.t.}~{}\sup_{x\in X}\langle
x,x^{*}\rangle-f(x)<+\infty\\}$ by $f^{*}(y):=\sup_{x\in X}\langle
x,y\rangle-f(x)$ and $\partial
f(x):=\\{y~{}\text{s.t.}~{}f(x^{\prime})-f(x)\geq\langle
y,x-x^{\prime}\rangle~{}\forall x^{\prime}\in X\\}$. Now remarks that

$\displaystyle G_{\lambda_{1}}^{*}(-h_{1})=\sup_{s}\langle
s,\lambda_{1}-h_{1}\rangle=\left\\{\begin{array}[]{ll}+\infty&\mbox{if
}\lambda_{1}\neq h_{1}\\\ 0&\mbox{otherwise .}\end{array}\right.$

therefore $G_{\lambda_{1}}^{*}$ ensures that $\lambda_{1}=h_{1}$. Similarly we
obtain that

$\displaystyle\min_{R}\frac{1}{\gamma}\text{KL}(R,\xi^{(2)})+\tau_{2}\text{KL}(r\mathbf{1}_{r}|b)+\langle-\lambda_{2},R^{\top}\mathbf{1}_{m}\rangle$
(19)

is equal to its dual defined as

$\displaystyle\sup_{f_{2},h_{2}}-F_{\tau_{2},b}^{*}(-f_{2})-G_{\lambda_{2}}^{*}(-h_{2})-\frac{1}{\gamma}\langle\exp(\gamma(f_{2}+h_{2})),\xi^{(2)}\rangle$

where again

$\displaystyle
G_{\lambda_{2}}^{*}(-h_{2})=\left\\{\begin{array}[]{ll}+\infty&\mbox{if
}\lambda_{2}\neq h_{2}\\\ 0&\mbox{otherwise}\end{array}\right.$

and with the primal-dual relationship $R=\operatorname{diag}(\exp(\gamma
f_{2}))\xi^{(2)}\operatorname{diag}(\exp(\gamma h_{2}))$ such that
$-f_{2}\in\partial F_{\tau_{2},b}(R\mathbf{1}_{r})$,$-h_{2}\in\partial
G_{\lambda_{2}}(R^{\top}\mathbf{1}_{m})$. Finally the dual can be written as

$\displaystyle\sup_{\lambda_{1},\lambda_{2}}\sup_{f_{1},h_{1}}-F_{\tau_{1},a}^{*}(-f_{1})-G_{\lambda_{1}}^{*}(-h_{1})-\frac{1}{\gamma}\langle\exp(\gamma(f_{1}+h_{1})),\xi^{(1)}\rangle$
$\displaystyle+\sup_{f_{2},h_{2}}-F_{\tau_{2},b}^{*}(-f_{2})-G_{\lambda_{2}}^{*}(-h_{2})-\frac{1}{\gamma}\langle\exp(\gamma(f_{2}+h_{2})),\xi^{(2)}\rangle$
$\displaystyle-\frac{1}{\gamma}\langle\exp(-\gamma(\lambda_{1}+\lambda_{2})),\xi^{(3)}\rangle+\langle\xi^{(3)},\mathbf{1}\rangle$

and using the definition of $G_{\lambda_{1}}^{*}(-h_{1})$ and
$G_{\lambda_{2}}^{*}(-h_{2})$, we obtain the desired dual up to an additive
constant $(\langle\xi^{(3)},\mathbf{1}\rangle)$ which does not affect the
solution of the problem and conclude the proof.

###  C.2 On the Iterations of the Dykstra’s Algorithm

Recall that we propose to consider an alternate maximization scheme to solve
(8). Starting from $h_{1}^{(0)}=h_{2}^{(0)}=\mathbf{0}_{r}$, we apply for
$\ell\geq 0$ the following updates (dropping iteration number $k$ in (7) for
simplicity):

$\displaystyle
f_{1}^{(\ell+1)}:=\arg\sup_{z}\mathcal{D}(z,h_{1}^{(\ell)},f_{2}^{(\ell)},h_{2}^{(\ell)}),\,f_{2}^{(\ell+1)}:=\arg\sup_{z}\mathcal{D}(f_{1}^{(\ell+1)},h_{1}^{(\ell)},z,h_{2}^{(\ell)}),$
$\displaystyle(h_{1}^{(\ell+1)},h_{2}^{(\ell+1)}):=\arg\sup_{z_{1},z_{2}}\mathcal{D}(f_{1}^{(\ell+1)},z_{1},f_{2}^{(\ell+1)},z_{2}).$

where

$\displaystyle\mathcal{D}(f_{1},h_{1},f_{2},h_{2})$
$\displaystyle=-F_{\tau_{1},a}^{*}(-f_{1})-\frac{1}{\gamma}\langle
e^{\gamma(f_{1}\oplus h_{1})}-1,\xi^{(1)}\rangle-
F_{\tau_{2},b}^{*}(-f_{2})-\frac{1}{\gamma}\langle e^{\gamma(f_{2}\oplus
h_{2})}-1,\xi^{(2)}\rangle$ $\displaystyle-\frac{1}{\gamma}\langle
e^{-\gamma(h_{1}+h_{2})}-1,\xi^{(3)}\rangle.$

Let us consider the first update of the scheme that consists in solving

$f_{1}^{(\ell+1)}:=\arg\sup_{z}\mathcal{D}(z,h_{1}^{(\ell)},f_{2}^{(\ell)},h_{2}^{(\ell)})$

To solve this problem, we again apply the Fenchel-Rockafellar theorem
[Rockafellar, 1970] and obtain that

$\displaystyle\sup_{f_{1}}-F_{\tau_{1},a}^{*}(-f_{1})-\frac{1}{\gamma}\langle\exp(\gamma(f_{1}+h_{1})),\xi^{(1)}\rangle=\min_{s}F_{\tau_{1},a}(s)+\frac{1}{\gamma}\text{KL}(s|\xi^{(1)}\exp(\gamma
h_{1}))$

and the optimality condition gives that $f_{1}^{*}$ is solution of the LHS if
and only if $s^{*}$ solves the RHS and belongs to the subdifferential of
$f_{1}\to\exp(\gamma(f_{1}+h_{1})),\xi^{(1)}\rangle$ at $f_{1}^{*}$, that is
$s^{*}=\exp(\gamma f_{1}^{*})\odot\xi^{(1)}\exp(\gamma h_{1})$. However the
RHS problem can can be solved exactly and one obtained that
$s^{*}=a^{(\tau_{1}/(1/\gamma+\tau_{1}))}\odot\xi^{(1)}\exp(\gamma
h_{1})^{(1/(1/1+\gamma\tau_{1}))}$, therefore when combined with the previous
equation on $s^{*}$ we obtain that

$\displaystyle\exp(\gamma f_{1}^{*})=\frac{s^{*}}{\xi^{(1)}\exp(\gamma
h_{1})}=\left(\frac{a}{\xi^{(1)}\exp(\gamma
h_{1})}\right)^{\frac{\tau_{1}}{1/\gamma+\tau_{1}}},$

Similarly, the solution of $\arg\sup_{z}\mathcal{D}(f_{1},h_{1},z,h_{2})$ is

$\displaystyle\exp(\gamma f_{2}^{*})=\left(\frac{b}{\xi^{(2)}\exp(\gamma
h_{2})}\right)^{\frac{\tau_{2}}{1/\gamma+\tau_{2}}}.$

Let us now consider the following optimization problem corresponding to the
last update if the alternate maximization scheme, that is

$(h_{1}^{(\ell+1)},h_{2}^{(\ell+1)}):=\arg\sup_{z_{1},z_{2}}\mathcal{D}(f_{1}^{(\ell+1)},z_{1},f_{2}^{(\ell+1)},z_{2}).$

In fact this problem can be solved directly using simply the first-order
condition of optimality that gives the two following equations:

$\displaystyle\exp(\gamma h_{1})\odot(\xi^{(1)})^{\top}\exp(\gamma
f_{1})-\exp(-\gamma h_{1})\odot(\xi^{(3)})\odot\exp(-\gamma
h_{2})=0\quad\text{and}$ $\displaystyle\exp(\gamma
h_{2})\odot(\xi^{(2)})^{\top}\exp(\gamma f_{2})-\exp(-\gamma
h_{2})\odot(\xi^{3)})\odot\exp(-\gamma h_{1})=0$

leading to

$\displaystyle g=(\xi^{3)}\odot(\xi^{(1)})^{\top}\exp(\gamma
f_{1})\odot(\xi^{(2)})^{\top}\exp(\gamma f_{2}))^{1/3}$

and

$\displaystyle\exp(\gamma h_{1})=\frac{g}{(\xi^{(1)})^{\top}\exp(\gamma
f_{1})},\quad\exp(\gamma h_{2})=\frac{g}{(\xi^{(2)})^{\top}\exp(\gamma
f_{2})}.$

###  C.3 Proof of Proposition 2

Let us consider the following optimization problem

$\displaystyle\mathcal{D}_{\textrm{TI}}(\tilde{f}_{1},\tilde{h}_{1},\tilde{f}_{2},\tilde{h}_{2}):=\sup_{\lambda_{1},\lambda_{2}\in\mathbb{R}}\mathcal{D}(\tilde{f}_{1}+\lambda_{1},\tilde{h}_{1}-\lambda_{1},\tilde{f}_{2}+\lambda_{2},\tilde{h}_{2}-\lambda_{2})$

Therefore we have

$\displaystyle\sup_{\lambda_{1},\lambda_{2}\in\mathbb{R}}\mathcal{D}(\tilde{f}_{1}+\lambda_{1},\tilde{h}_{1}-\lambda_{1},\tilde{f}_{2}+\lambda_{2},\tilde{h}_{2}-\lambda_{2})$
$\displaystyle=-F_{\tau_{1},a}^{*}(-(\tilde{f}_{1}+\lambda_{1}))-F_{\tau_{2},b}^{*}(-(\tilde{f}_{2}+\lambda_{2}))-\frac{1}{\gamma}\langle
e^{-\gamma(\tilde{h}_{1}+\tilde{h}_{2})}\odot
e^{\gamma(\lambda_{1}+\lambda_{2})},\xi^{(3)}\rangle+C$

where $C$ does not depends on $\lambda_{1}$ and $\lambda_{2}$. Now observe
that

$\displaystyle F_{\tau_{1},a}^{*}(s)$ $\displaystyle=\sup_{x}\langle
x,s\rangle-\tau_{1}\text{KL}(s|a)$

and by applying the first-order optimality condition, we obtain that
$x^{*}=\exp(s/\tau_{1})\odot a$ solves the above optimization problem and

$\displaystyle
F_{\tau_{1},a}^{*}(s)=\tau_{1}\langle\exp(s/\tau_{1}),a\rangle.$

Similarly,

$\displaystyle
F_{\tau_{2},b}^{*}(s)=\tau_{2}\langle\exp(s/\tau_{2}),b\rangle,$

Then by appyling the first-order optimality condition we obtain the two
following equations

$\displaystyle\exp(-\lambda_{1}/\tau_{1})\langle\exp(-\tilde{f}_{1}/\tau_{1}),a\rangle-\exp(\gamma\lambda_{1})\langle\exp(\gamma\lambda_{2}),\xi^{(3)}\odot\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2}))\rangle=0\quad\text{and}$
$\displaystyle\exp(-\lambda_{2}/\tau_{2})\langle\exp(-\tilde{f}_{2}/\tau_{2}),b\rangle-\exp(\gamma\lambda_{2})\langle\exp(\gamma\lambda_{1}),\xi^{(3)}\odot\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2}))\rangle=0.$

which is equivalent to

$\displaystyle\exp\left(\lambda_{1}\frac{1/\gamma+\tau_{1}}{\tau_{1}/\gamma}\right)$
$\displaystyle=\frac{\langle\exp(-\tilde{f}_{1}/\tau_{1}),a\rangle}{\langle\xi^{(3)},\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2}))\rangle}\exp(-\gamma\lambda_{2})\quad\text{and}$
$\displaystyle\exp\left(\lambda_{2}\frac{1/\gamma+\tau_{2}}{\tau_{2}/\gamma}\right)$
$\displaystyle=\frac{\langle\exp(-\tilde{f}_{2}/\tau_{2}),b\rangle}{\langle\xi^{(3)},\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2}))\rangle}\exp(-\gamma\lambda_{1})$

Then applying $\log$ to the system, we obtain that

$\displaystyle\lambda_{1}\gamma\frac{1/\gamma+\tau_{1}}{\tau_{1}}$
$\displaystyle=c_{1}-\gamma\lambda_{2}\quad\text{and}$
$\displaystyle\lambda_{2}\gamma\frac{1/\gamma+\tau_{2}}{\tau_{2}}$
$\displaystyle=c_{2}-\gamma\lambda_{1}$

where

$\displaystyle
c_{1}:=\log\left(\frac{\langle\exp(-\tilde{f}_{1}/\tau_{1}),a\rangle}{\langle\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2})),\xi^{(3)}\rangle}\right),\quad\text{and}\quad
c_{2}:=\log\left(\frac{\langle\exp(-\tilde{f}_{2}/\tau_{2}),a\rangle}{\langle\exp(-\gamma(\tilde{h}_{1}+\tilde{h}_{2})),\xi^{(3)}\rangle}\right).$

Finally we obtain a simple linear system and the solution follows.

###  C.4 Double Regularizations: Low-rank Structure and Entropy

Our proposed procedure can be easily extended to the case where one wants to
add entropy in addition to the low-rank constraint to solve unbalanced low-
rank and entropic optimal transport problems. More precisely, let us consider
the general case where one aims at solving for any $\varepsilon>0$

$\displaystyle\text{ULOT}_{r,\varepsilon}(\mu,\nu):=\min_{(Q,R,g)\in\Pi_{r}}\underbrace{\langle
C,Q\operatorname{diag}(1/g)R^{T}\rangle}_{\mathcal{L}_{C}(Q,R,g)}+\underbrace{\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)-\varepsilon
H(Q,R,g)}_{\mathcal{G}_{a,b,\varepsilon}(Q,R,g)}$ (20)

where $H(Q,R,g)=H(Q)+H(R)+H(g)$ and $H(p):=-\sum_{i}p_{i}(\log(p_{i})-1)$.
Note that here, compared to (6), we have simply add en entropic term to the
objective to smooth the matrices $Q,R$ and the barycenter $g$. To solve this
problem, we propose to consider the exact same strategy as the one proposed to
solve (6) where we slightly modify $\mathcal{G}_{a,b,\varepsilon}$ and
explicitely show the dependency w.r.t. $\varepsilon$. Now by applying the
linearzation step of $\mathcal{L}_{C}(Q,R,g)$, we now aim to solve at
iteration $k$ the following optimization problem:

$\displaystyle(Q_{k+1},R_{k+1},g_{k+1}):=\operatorname*{argmin}_{\bm{\zeta}\in\Pi_{r}}\frac{1}{\gamma_{k}}\text{KL}(\bm{\zeta},\bm{\xi}_{k})+\varepsilon
H(\bm{\zeta})+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$
(21)

In fact, this problem can be reformulated as a problem of the form (14) where
we simply have to modify $\bm{\xi}_{k}$ and $\gamma$. Indeed observe that we
have

$\displaystyle\frac{1}{\gamma}\text{KL}(Q|\xi^{(1)})-\varepsilon H(Q)$
$\displaystyle=\frac{1}{\gamma_{\varepsilon}}\text{KL}(Q|\xi_{\varepsilon}^{(1)})$

where $\gamma_{\varepsilon}=\frac{1}{1/\gamma+\varepsilon}$ and
$\xi_{\varepsilon}^{(1)}:=(\xi^{(1)})^{\gamma_{\varepsilon}/\gamma}$.
Therefore we obtain that

$\displaystyle\operatorname*{argmin}_{\bm{\zeta}\in\Pi_{r}}\frac{1}{\gamma}\text{KL}(\bm{\zeta},\bm{\xi})+\varepsilon
H(\bm{\zeta})+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$
$\displaystyle=\operatorname*{argmin}_{\bm{\zeta}\in\Pi_{r}}\frac{1}{\gamma_{\varepsilon}}\text{KL}(\bm{\zeta},\bm{\xi}_{\varepsilon})+\tau_{1}\text{KL}(Q\mathbf{1}_{r}|a)+\tau_{2}\text{KL}(R\mathbf{1}_{r}|b)$

where
$\bm{\xi}_{\varepsilon}:=(\xi_{\varepsilon}^{(1)},\xi_{\varepsilon}^{(2)},\xi_{\varepsilon}^{(3)})$.
Therefore the entropic version of our problem can be solved using the exact
same solver as the one proposed in the main paper where only simple updates of
the gradient-step $\gamma$ and the kernels $\bm{\xi}$ are required at each
iteration. We summarize the proposed algorithm below.

Inputs: $C,a,b,\varepsilon,\gamma_{0},\tau_{1},\tau_{2},\delta$
$Q,R,g\leftarrow\text{Initialization as proposed
in~{}\cite[citep]{[\@@bibref{AuthorsPhrase1Year}{scetbon2022lot}{\@@citephrase{,
}}{}]}}$
repeat

$\tilde{Q}=Q,~{}~{}\tilde{R}=R,~{}~{}\tilde{g}=g,\\\
\nabla_{Q}=CR\operatorname{diag}(1/g),~{}~{}\nabla_{R}=C^{\top}Q\operatorname{diag}(1/g),\\\
\omega\leftarrow\mathcal{D}(Q^{T}CR),~{}~{}\nabla_{g}=-\omega/g^{2},\\\
\gamma\leftarrow\gamma_{0}/\max(\|\nabla_{Q}\|_{\infty}^{2},\|\nabla_{R}\|_{\infty}^{2},\|\nabla_{g}\|_{\infty}^{2}),\\\
\gamma\leftarrow\frac{1}{1/\gamma+\varepsilon}\\\ \xi^{(1)}\leftarrow
Q\odot\exp(-\gamma\nabla_{Q}),~{}\xi^{(2)}\leftarrow
R\odot\exp(-\gamma\nabla_{R}),~{}\xi^{(3)}\leftarrow
g\odot\exp(-\gamma\nabla_{g}),\\\
\xi^{(1)}\leftarrow(\xi^{(1)})^{\gamma_{\varepsilon}/\gamma},~{}\xi^{(2)}\leftarrow(\xi^{(2)})^{\gamma_{\varepsilon}/\gamma},~{}\xi^{(3)}\leftarrow(\xi^{(3)})^{\gamma_{\varepsilon}/\gamma},\\\
Q,R,g\leftarrow\text{ULR-
Dykstra}(a,b,\bm{\xi},\gamma,\tau_{1},\tau_{2},\delta)~{}(\text{Alg.~{}\ref{alg-
dykstra-ulot})}$

until  _$\Delta((Q,R,g),(\tilde{Q},\tilde{R},\tilde{g}),\gamma) <\delta$_;

Result: $Q,R,g$

Algorithm 7
$\text{ULOT}_{\varepsilon}(C,a,b,r,\gamma_{0},\tau_{1},\tau_{2},\delta)$