SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 3 ULR-TI-Dykstra  $(a,b,\xi ,\gamma ,\tau_1,\tau_2,\delta)$
Inputs:  $a, b, \xi = (\xi^{(1)}, \xi^{(2)}, \xi^{(3)})$ ,  $\gamma, \tau_1, \tau_2, \delta$ $v_1 = v_2 = 1_r, u_1 = 1_n, u_2 = 1_m$   
repeat  
\[
\begin{aligned}
& \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. 2)} \\
& 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. 2)} \\
& g = \exp(\gamma(\lambda_1 + \lambda_2))^{1/3} (\xi^{(3)} \odot (\xi^{(1)})^T u_1 \odot (\xi^{(2)})^T u_2)^{1/3}, v_1 = \frac{g}{(\xi^{(1)})^T u_1}, v_2 = \frac{g}{(\xi^{(2)})^T u_2} \\
& \text{until } \frac{1}{\gamma} \max(\| \log(u_i / \tilde{u}_i) \|_{\infty}, \| \log(v_i / \tilde{v}_i) \|_{\infty}) < \delta;
\end{aligned}
\]  
Result:  $d i a g(u_{1})\xi_{k}^{(1)} d i a g(v_{1}), d i a g(u_{2})\xi_{k}^{(2)} d i a g(v_{2}), g$
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
$$
\operatorname {L G W} _ {r} ((a, A), (b, B)) = \min  _ {P \in \Pi_ {a, b} (r)} \mathcal {Q} _ {A, B} (P), \tag {12}
$$
Building upon § 3.1 and leveraging the TI variant presented in § 3.2, we introduce the unbalanced low-rank Gromov-Wasserstein (ULGW) problem. There is, however, a significant challenge that appears when introducing unbalanced regularizers in (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 discarded when optimizing. In an unbalanced setting, these terms vary and must be accounted for:
$$
\begin{array}{l} \operatorname {U L G W} _ {r} ((a, A), (b, B)) := \min  _ {(Q, R, g) \in \Pi_ {r}} \left\langle A ^ {\odot 2} Q \mathbf {1} _ {r}, Q \mathbf {1} _ {r} \right\rangle + \left\langle B ^ {\odot 2} R \mathbf {1} _ {r}, R \mathbf {1} _ {r} \right\rangle \tag {13} \\ - 2 \left\langle A Q \operatorname {d i a g} (1 / g) R ^ {T} B, Q \operatorname {d i a g} (1 / g) R ^ {T} \right\rangle + \tau_ {1} \mathrm {K L} \left(Q \mathbf {1} _ {r} | a\right) + \tau_ {2} \mathrm {K L} \left(R \mathbf {1} _ {r} | b\right) \\ \end{array}
$$
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 same as that solving ULOT, however, the kernels  $\xi_{k}$  now take into account the quadratic terms of the original LGW problem. More formally, at each iteration  $k$  of the outer loop, we propose to solve
$$
\left(Q _ {k + 1}, R _ {k + 1}, g _ {k + 1}\right) := \underset {\zeta \in \Pi_ {r}} {\operatorname {a r g m i n}} \frac {1}{\gamma_ {k}} \mathrm {K L} (\boldsymbol {\zeta} | \boldsymbol {\xi} _ {k}) + \tau_ {1} \mathrm {K L} (Q \mathbf {1} _ {r} | a) + \tau_ {2} \mathrm {K L} (R \mathbf {1} _ {r} | b), \tag {14}
$$
where  $(Q_0,R_0,g_0)\in \Pi_r$  is the initialization,  $(\gamma_{k})_{k\geq 0}$  a sequence of positive step sizes. Using notation  $P_{k} = Q_{k}\mathrm{diag}(1 / g_{k})R_{k}^{T}$ , the synthetic cost matrices  $\xi_{k}\coloneqq (\xi_{k}^{(1)},\xi_{k}^{(2)},\xi_{k}^{(3)})$  are updated as:
$$
\xi_ {k} ^ {(1)} := Q _ {k} \odot \exp (- 2 \gamma_ {k} A ^ {\odot 2} Q _ {k} \mathbf {1} _ {r} \mathbf {1} _ {r} ^ {T}) \odot \exp (- 4 \gamma_ {k} A P _ {k} B R _ {k} \operatorname {d i a g} (1 / g _ {k}))),
$$
$$
\xi_ {k} ^ {(2)} := R _ {k} \odot \exp (- 2 \gamma_ {k} B ^ {\odot 2} R _ {k} \mathbf {1} _ {r} \mathbf {1} _ {r} ^ {T}) \odot \exp (- 4 \gamma_ {k} B P _ {k} ^ {T} A Q _ {k} \operatorname {d i a g} (1 / g _ {k}))),
$$
$$
\xi_ {k} ^ {(3)} := g _ {k} \odot \exp (4 \gamma_ {k} \omega_ {k} / g _ {k} ^ {2}) \quad \text {w i t h} \quad [ \omega_ {k} ] _ {i} := [ Q _ {k} ^ {T} A P _ {k} B R _ {k} ] _ {i, i} \in \mathbb {R} ^ {r}.
$$
Note that (14) is the exact same optimization problem as (7), where only  $\pmb{\xi}_k$  has changed and therefore can be solved using Algorithm 3. Algorithm 4 summarizes our strategy to solve (13).
Algorithm 4 ULGW(A, B, a, b, r,  $\gamma_0$ ,  $\tau_1$ ,  $\tau_2$ ,  $\delta$ )
Inputs:  $A,B,a,b,r,\gamma_0,\tau_1,\tau_2,\delta$ $Q,R,g\gets$  Initialization as proposed in [Scetbon and Cuturi, 2022]   
repeat   
 $\tilde{Q} = Q,\tilde{R} = R,\tilde{g} = g,$ $\nabla_{Q} = 4AQ\mathrm{diag}(1 / g)R^{T}BR\mathrm{diag}(1 / g) + 2A^{\odot 2}Q\mathbf{1}_{r}\mathbf{1}_{r}^{T},$ $\nabla_R = 4BR\mathrm{diag}(1 / g)Q^T AQ\mathrm{diag}(1 / g) + 2B^{\odot 2}R\mathbf{1}_r\mathbf{1}_r^T,$ $\omega \leftarrow \mathcal{D}(Q^T AQ\mathrm{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)}\gets Q\odot \exp (-\gamma \nabla_Q),\xi^{(2)}\gets R\odot \exp (-\gamma \nabla_R),\xi^{(3)}\gets g$ $Q,R,g\gets$  ULR-TI-Dykstra(a,b,  $\pmb {\xi},\gamma ,\tau_1,\tau_2,\delta)$  (Alg. 3)   
until  $\Delta ((Q,R,g),(\tilde{Q},\tilde{R},\tilde{g}),\gamma) <   \delta ;$    
Result:  $Q,R,g$
Convergence and Complexity. Similarly to linear ULOT, the unbalanced Dykstra algorithm is guaranteed to converge [Bauschke and Lewis, 2000]. Because we use Algorithm 5, we retain exactly the same complexity, both in terms of time of memory, to solve these inner problems. The slight variation in kernel  $\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 both costs matrices  $A$  and  $B$  to reach linear complexity. Indeed, assuming  $A \simeq A_1A_2^T$  and  $B \simeq B_1B_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 impactful [Klein et al., 2023] fused-Gromov-Wasserstein problem, which merges linear and quadratic objectives [Vayer et al., 2018]:
$$
\operatorname {F G W} (\mu , \nu) := \min  _ {P \in \Pi_ {a, b}} \alpha \langle C, P \rangle + \bar {\alpha} Q _ {A, B} (P) \tag {15}
$$
where  $\alpha \in [0,1]$  and  $\bar{\alpha} \coloneqq 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',j,j'] \coloneqq (A_{i,i'} - B_{j,j'})^2$  appearing in (4) by a composite interpolated cost between the OT and GW geometries, redefined as  $M[i,i',j,j'] = \alpha C_{i,j} + \bar{\alpha}(A_{i,i'} - B_{j,j'})^2$ . Our proposed unbalanced and low-rank version of the FGW problem includes  $|P| \coloneqq \|P\|_1$  the mass of  $P$ , to homogenize linear and quadratic terms,
$$
\operatorname {U L F G W} _ {r} (\mu , \nu) := \min  _ {P: \mathrm {r k} + (P) \leq r} \alpha | P | \langle C, P \rangle + \bar {\alpha} \mathcal {Q} _ {A, B} (P) + \tau_ {1} \mathrm {K L} \left(P \mathbf {1} _ {m} | a\right) + \tau_ {2} \mathrm {K L} \left(P ^ {T} \mathbf {1} _ {n} | b\right), \tag {16}
$$
which is expanded through the explicit factorization of  $P$ , noticing that  $|P| = |g| \coloneqq \| g\|_1$ :
$$
\operatorname {U L F G W} _ {r} (\mu , \nu) := \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) \tag {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 kernels  $\pmb{\xi}_k$  are now defined as
$$
\boldsymbol {\xi} _ {k} := \left(\xi_ {k} ^ {(1)}, \xi_ {k} ^ {(2)}, \xi_ {k} ^ {(3)}\right),
$$
$$
\xi_ {k} ^ {(1)} := 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})
$$
$$
\nabla_ {Q} \mathcal {H} _ {k} := \alpha | g _ {k} | C R _ {k} \operatorname {d i a g} (1 / g _ {k}) + \bar {\alpha} \left(2 A ^ {\odot 2} Q _ {k} \mathbf {1} _ {r} \mathbf {1} _ {r} ^ {T} + 4 A P _ {k} B R _ {k} \operatorname {d i a g} (1 / g _ {k})\right)
$$
$$
\nabla_ {R} \mathcal {H} _ {k} := \alpha | g _ {k} | C ^ {T} Q _ {k} \operatorname {d i a g} \left(1 / g _ {k}\right) + \bar {\alpha} \left(2 B ^ {\odot 2} R _ {k} \mathbf {1} _ {r} \mathbf {1} _ {r} ^ {T} + 4 B P _ {k} ^ {T} A Q _ {k} \operatorname {d i a g} \left(1 / g _ {k}\right)\right)
$$
$$
\nabla_ {g} \mathcal {H} _ {k} := \alpha \left(\langle C, P _ {k} \rangle \mathbf {1} _ {r} - | g _ {k} | \omega_ {k} ^ {\text {l i n}} / g _ {k} ^ {2}\right) - 4 \bar {\alpha} \omega_ {k} ^ {\text {q u a d}} / g _ {k} ^ {2}
$$
$$
\left[ \omega_ {k} ^ {\text {l i n}} \right] _ {i} := \left[ Q _ {k} ^ {T} C R _ {k} \right] _ {i, i}, \quad \left[ \omega_ {k} ^ {\text {q u a d}} \right] _ {i} := \left[ Q _ {k} ^ {T} A P _ {k} B R _ {k} \right] _ {i, i} \forall i \in \{1, \dots , r \}.
$$
These steps are summarized in Alg. 6 (see appendix). These steps result in a quadratic complexity, both in time and memory, with respect to the number of points  $n$  and  $m$ . However, these complexities become linear when square matrices  $A$ ,  $B$  and rectangular  $C$  all admit a low-rank factorization.
(a) Gene Nrgn: measured vs. predicted expression
(b) Gene  $Slc24a2$ : measured vs. predicted expression
Figure 1: Exp. 2: Spatial visualization of two mouse brain sections, contrasting observed vs. predicted (using ULFGW) spatial distributions of expression levels, for two different genes.
We focus first in Exp. 1 on demonstrating the empirical benefits of the TI variant of our algorithm to solve linear ULOT, as implemented in Alg. 3 vs. Alg. 5; that algorithm is subsequently used as an inner routine to solve all quadratic ULR problems. We compare in Exp. 2 unbalanced low-rank (ULR) solvers to balanced low-rank (LR) counterparts on a spatial transcriptomics task, and follow in Exp. 3 by comparing ULR solvers to entropic (E) counterparts on a smaller task, to accommodate entropic solvers' quadratic complexity. We conclude in Exp. 4 by comparing ULR solvers to [Thual et al., 2022], which can learn a sparse transport coupling, in the unbalanced FGW setting.
Datasets. We consider two real-world datasets, described in B.1, and two synthetic datasets, that are large enough to showcase our solvers. The real-world datasets 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 that are specific to each source  $s$  and target  $t$  measures, and which are used to compute the costs matrices for the quadratic term in the GW and FGW settings. In Exp. 1, we sim
Figure 2: Visualization of measured and predicted tissue regions in the mouse brain in Exp. 2
ply consider two isotropic Gaussians in  $d = 30$  to evaluate the performance of the TI variant on a liner problem. We use the mouse brain STARmap spatial transcriptomics data from [Shi et al., 2022] for Exp. 2 and Exp. 3. We use data from the Individual Brain Charting dataset [Pinho et al., 2018], to replicate the settings of [Thual et al., 2022], in Exp. 4.
Metrics. Following Klein et al. [2023], we evaluate maps by focusing on the two following metrics: (i) Pearson correlation  $\rho$  computed between the (ground truth) source  $s$  feature matrix  $F^s \in \mathbb{R}^{n \times d}$ , and the barycentric projection of the target  $t$  to the source scaled by the target marginals  $b^t$ . Writing  $P$  as the transport matrix from source to target, this can be computed as  $P\mathrm{diag}\left(\frac{1}{b^t}\right)F^t$ ; (ii) F1 score when assessing class transfer (among 11 possible classes), computed between the original source vector of labels  $l^s$ , taken in  $\{1,\dots,11\}^n$ , and the inferred labels for the same points, predicted for each  $i$  by taking the  $\arg\max_{j} B_{i,j}$ , where  $B$  is a matrix of  $n \times 11$  row probabilities, each the barycentric projection of the target  $t$  one-hot encoded labels  $L^t \in \{0,1\}^{m \times 11}$ ,  $B := P\mathrm{diag}\left(\frac{1}{b^t}\right)L^t$ .
Experiment 1: Benchmarking The Translation Invariant Variant. We evaluate the effect of the proposed TI procedure on the computational cost of ULR solvers: We compare the time taken when solving unbalanced LR problems, with or without using the TI objective. In Figure 3, we compare the execution time (using our ott-jax implementation, and a single NVIDIA GeForce RTX 2080
Ti card) of unbalanced LR Sinkhorn on large and high dimensional Gaussian distributions. The results presented are averaged over 10 random seeds with error bars. We use a  $\delta = 10^{-9}$  convergence threshold and 1000 maximal number of iterations for Dykstra, in 64-bit precision. We observe that the use of our proposed TI objective is consistently beneficial when solving ULR problems. See also Appendix B.3 for additional experiments.
Figure 3: Execution time of unbalanced LR Sinkhorn, with (Alg. 3) or without (Alg. 5) the TI variant. We fix the rank to  $r = 10$ ;  $n$  points (displayed in thousands) are sampled from two Gaussian distributions in  $d = 30$  of means respectively -1.2 and 1.3, and standard deviations 1 and 0.2. (left) displays large  $\tau$  (close to balanced), (right) is smaller  $\tau$  (more unbalanced). We use the same convergence threshold for the outer loop, for all sample sizes. As  $n$  gets bigger, this results in a relatively looser threshold, explaining why timings can slightly decrease w.r.t.  $n$ . What matters is, therefore, the comparative performance of TI vs non-TI for a fixed  $n$ , not the behaviour w.r.t.  $n$ .
Experiment 2: ULOT vs. LOT on Gene Expression / Cell Type Annotation. 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. 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.
solver	mass %	val ρ	test ρ	F1 mac.	F1 mic.	F1 weig.
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: Exp.2, Results for spatial transcriptomics dataset (brain coronal section from Shi et al. [2022]).
Experiment 3: ULOT vs. UEOT. We compare the performance of ULOT solvers to their unbalanced entropic alternatives (UEOT). We use the same datasets as in Exp. 2, but must pick a smaller subset (Olfactory bulb), to avoid OOM errors for entropic UGW solvers, since they cannot handle the  $40k$  sizes considered in Exp. 2 (see B.1). This results in  $n \approx 20$ , 000 source and  $\approx 15$ , 000 target
Figure 4: Visualization of measured and predicted right auditory click contrast map in Exp.4.
cells, and 1000 genes. Similar to Exp. 2, the fused term  $C$  is a squared-Euclidean distance matrix on 30-D PCA space, computed on gene expressions. As done in Exp. 2, 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 in Exp. 2 and B.2. Table 2 shows that ULFGW outperforms entropic solvers w.r.t. Pearson correlation  $\rho$ , but is worse when considering F1 scores. On the other hand, ULFGW confirms its superiority compared to the balanced alternative LFGW. Taken together, these results suggest that while unbalanced
solver	mass %	val ρ	test ρ	F1 mac.	F1 mic.	F1 weig.
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
LR solvers are on par with unbalanced entropic solvers in terms of performance, in small data regimes, they remain much faster and can unlock the applications of unbalanced OT to larger scales.
Experiment 4: ULOT to align brain meshes. In this experiment, we compare the performance of our ULFGW solver to FUGW-sparse [Thual et al., 2022], an alternative approach to solve unbalanced FGW problems, using a two-scale (corse/fine grid) approach to handle large sample sizes. This method was demonstrated to be effective in aligning brain anatomies, encompassing both mesh structures and functional signals associated with each vertex. For their empirical analysis, they use the individual brain charting dataset [Pinho et al., 2018].
In the absence of other information in the original paper, we draw inspiration from Pinho et al.'s smaller scale notebook implementations: We embed the  $n \approx 160,000$  vertices of the fsaverage7 mesh, into a 30-d space, using an approximation of the geodesic distances with landmark multi-dimensional scaling [De Silva and Tenenbaum, 2004] where 2048 points were used as anchors. Each vertex has an associated functional signal that entails 22 features. For both the quadratic and linear terms, we compute the costs based on the squared Euclidean distance. The coarse grid for FUGW-sparse is built using one-
tenth of  $n$ , i.e.  $\approx 16\bar{k}$  points. We evaluate all methods by comparing the performance of the best hyperparameter combination, based on the average correlation between the barycentric projection and ground-truth value of 5 features, across a test set of 5 contrast maps. In Table 3, we observe that ULFGW and LFGW outperform FUGW-sparse. In this setting, there is no clear evidence that the unbalanced version performs better than its balanced counterpart for low-rank methods. See also Appendix B.2 for additional experimental details and results.
Table 2: Exp. 3: Results for spatial transcriptomics dataset (Olfactory bulb section from Shi et al. [2022]).
solver	mass	val ρ	test ρ
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.4.
Conclusion. The practical success of OT methods to natural sciences demonstrates the relevance of OT to their analysis pipelines. Practitioners must, however, often deal with the poor scalability of OT algorithms, as well as their rigid assumptions w.r.t. mass conservation. While Low-rank OT approaches hold the promise of scaling OT methods to large sizes, unbalanced formulations have proved useful to relax mass conservation for entropic OT solvers. We have proposed in this paper to merge these two strains, and demonstrated the practical relevance of these unbalanced low-rank solvers on various challenging alignment tasks.
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.
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.
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.
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.
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.
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.
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.
F Alexander Wolf, Philipp Angerer, and Fabian J Theis. SCANPY: large-scale single-cell gene expression data analysis. Genome Biology, 19(1), 2018.
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.