Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Authors: Peter Halmos
Julian Gold
Xinhao Liu
Benjamin Raphael
Venue: ICML 2025 oral (2025)
Keywords: Optimal transport
low-rank
linear complexity
sparse
full-rank
Primary area: general_machine_learning
TL;DR: Linear-complexity optimal transport, using low-rank optimal transport to progressively refine the solution to a Monge map.

Abstract: Optimal transport (OT) has enjoyed great success in machine learning as a principled way to align datasets via a least-cost correspondence, driven in large part by the runtime efficiency of the Sinkhorn algorithm (Cuturi, 2013). However, Sinkhorn has quadratic space and time complexity in the number of points, limiting scalability to larger datasets. Low-rank OT achieves linear complexity, but by definition, cannot compute a one-to-one correspondence between points. When the optimal transport problem is an assignment problem between datasets then an optimal mapping, known as the _Monge map_, is guaranteed to be a bijection. In this setting, we show that the factors of an optimal low-rank coupling co-cluster each point with its image under the Monge map. We leverage this invariant to derive an algorithm, _Hierarchical Refinement_ (`HiRef`), that dynamically constructs a multiscale partition of each dataset using low-rank OT subproblems, culminating in the bijective Monge map. Hierarchical Refinement runs in log-linear time and linear space, retaining the advantages of low-rank OT while overcoming its limited resolution. We demonstrate the advantages of Hierarchical Refinement on several datasets, including ones containing over a million points, scaling full-rank OT to problems previously beyond Sinkhorn's reach.

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Peter Halmos' 1*
Julian Gold 12*
Xinhao Liu Benjamin J. Raphael
Abstract
lem, instead seeks least-cost coupling 7 between μ and v. In the Kantorovich formulation, mass splitting is allowed and thus a solution always exists; in contrast, Monge map between μ and y/ may not exist. When a Monge map T does exist, the solution to the Kantorovich problem is a coupling Ty = (id x T); I supported on its graph, and the Monge and Kantorovich problems coincide (Brenier, 1991).
Optimal transport (OT) has enjoyed great suc- cess in machine learning as a principled way align datasets via least-cost correspondence, driven in large part by the runtime efficiency of the Sinkhorn algorithm (Cuturi, 2013). However, Sinkhorn has quadratic space and time complex- ity in the number of points, limiting scalability to larger datasets. Low-rank OT achieves linear com- plexity, but by definition, cannot compute a one- to-one correspondence between points. When the optimal transport problem is an assignment prob- lem between datasets then an optimal mapping, known as the Monge map, is guaranteed to be a bijection. In this setting, we show that the factors of an optimal low-rank coupling co-cluster each point with its image under the Monge map. We leverage this invariant to derive an algorithm, Hi- erarchical Refinement (HiRef), that dynamically constructs a multiscale partition of each dataset us- ing low-rank OT subproblems, culminating in the bijective Monge map. Hierarchical Refinement runs in log-linear time and linear space, retaining the advantages of low-rank OT while overcom- ing its limited resolution. We demonstrate the advantages of Hierarchical Refinement several datasets, including ones containing over a million points, scaling full-rank OT to problems previ- ously beyond Sinkhorn's 's reach.
When l and v are discrete uniform measures on points the optimal transport problem reduces to an assignment problem. Classical algorithms such as the Hungarian al- gorithm and Network Simplex (Tarjan, 1997; Orlin, 1997), solve this in cubic time. The Sinkhorn algorithm (Cuturi, 2013) solves the entropy-reeullarzzed Kantorovich problem with quadratic runtime, greatly expanding the applicabil- ity of computational OT. However, the Sinkhorn algorihtm requires quadratic space to store the coupling 7
In recent years, OT has found numerous applications in machine learning and across science, including: domain adaptation (Courty et al., 2014; Solomon et al., 2015), self- attention (Tay et al., 2020; Sander et al., 2022; Geshkovski et al., 2023), computational biology (Schiebinger al., 2019; Yang al., 2020; Zeira et al., 2022: Bunne al., 2023; Halmos et al., 2025b; Klein et al., 2025), unpaired data translation (Korotin et al., 2021; De Bortoli et al., 2024; Tong et al., 2024; Klein et al., 2024), and alignment prob- lems in transformers and large language models (Melnyk et al., 2024; Li et al., 2024). The least-cost principle of optimal transport is crucial for training high-quality gen- erative models using Schrödinger bridges, flow-matching, diffusion models, or neural ordinary differential equations (Finlay et al., 2020; Tong et al., 2024; De Bortoli et al., 2024; Kornilov et al., 2024; Klein et al., 2024). These models typi- cally require millions to hundreds of millions of data-points to achieve eigh-performance at scale (Ramesh et al., 2021), limiting the scope of OT for generative modeling.
1. Introduction
Optimal transport (OT) is mathematical framework for comparing probability distributions l and Given cost function c, the Monge problem is to find a mapping T trans- forming a distribution into (i.e. Ty = v) with least-cost. A relaxation of this problem, called the Kantorovich prob-
As modern datasets grow have tens of thousands even millions of points, the quadratic space and time complexity of Sinkhorn becomes increasingly prohibitive. This limita- tion is widely recognized in the machine learning literature, with (De Bortoli et al., 2024) noting that the quadratic com- plexity of optimal transport renders its application to modern datasets on the order of millions of points impractical. A number of approaches have been proposed address scal-
"Equal contribution 'Deparmment of Computer Science, Prince- University *Center for Statistics and Machine Learning, Princcton University. Correspondence to: Benjamin Raphael <braphael@ craphacl@princetooncdu.. princeton.edu>.
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
ing OT to massive datasets which avoid instantiating a full coupling matrix. Mini-batch OT (Genevay et al., 2018) im- proves scalability, but incurs significant biases (Sommerfeld et al., 2019; Korotin et al., 2021; Fatras et al., 2021a) la) as each mini-batch alignment is a poor representation of the global coupling. Multiple works have investigated the theo- retical properties of mini-batch estimators of the coupling (Fatras al.. 2020; 2021b), while others have attempted to mitigate this bias using partial unbalanced OT that allows mass variation between mini-batches (Nguyen et al., 2022a; Fatras et al., 2021a). However, these approaches introduce additional hyperparameters to control the degree of unbalancedness,, and ultimately remain biased, local ap- proximations of the global coupling.
of (Mérigot, 2011). Unlike (Mérigot, 2011; Gerber & Mag- gioni, 2017), Hierarchical Refinement (i) does not assume multiscale partitions for each dataset are given, instead con- structing them on the fly; and (ii) operates intrinsically to the data, without a mesh or anchor points in the ambient space of the data, avoiding the curse of dimensionality.
We demonstrate that Hierarchical Refinement computes OT maps efficiently in nigh-dimensional spaces, often matching or even outperforming Sinkhorn in terms of primal cost. Moreover, HiRef has linear space complexity and time complexity scaling log-linearly in the dataset size. Unlike low-rank OT, Hierarchical Refinement places X and Y in bijective correspondence. Hierarchical Refinement scales to over a million points, enabling the use of OT on massive datasets without incurring the bias of mini-batching.
Neural optimal transport methods (Makkuva et al.. 2020; Bunne et al., 2023; Fan et al., 2023; Korotin et al., 2023; Buzun et al., 2024), parametrize the Monge map as neural network instead of materializing a quadratic coupling matrix. However, these methods have noted limitations recovering faithful maps (Korotin et al., 2021).
2. Background and Related Work
Suppose X = {x and Y = {y;) are datasets in the same metric space (x, dzx). Let c X x X + R, be a cost function. This cost e is often assumed to satisty strict convexity or be a metric. Datasets X and Y are represented discretely supported probability measures μ = Σ a,ox, and >E= bjåy, for probability vec- tors a E An and b E Am- Throughout, An denotes the kismlle {p E R$ : Σ,P: = 1}. the set of probability vectors of length k.
Another approach to improve space complexity of OT is to introduce a low-rank constraint on the coupling matrix in the Kantorovich problem. This has been done by parameterizing the coupling through a set of low-rank factors (Scetbon et al., 2021; 2022; Scetbon & Cuturi, 2022; Scetbon et al., 2023; Halmos et al., 2024) or by using a proxy objective for the low-rank problem, factoring the transport through a small number of anchor points (Forrow al., 2019; Lin et al., 2021). For a given rank r these approaches have (nr) space complexity, enabling linear time and space scaling. Low-rank OT has been used successfully the order of 10y samples with ranks on the order of 10' (Scetbon et al., 2023; Halmos et al., 2024; 2025a; Klein et al., 2025), but computing full-rank couplings between datasets of sizes on the order of 10 and greater has not yet been accomplished.
Optimal transport has its origin in the Monge problem (Monge, 1781), concerned with finding an optimal map T X Y pushing μ forward to v:
Monge Problem
datasets
Mc(u,e)
min T.THK=
Ea T(x))
(1)
Above, Tapi is the pushforward of μI under T, the measure on with TB == (BT))(()) for any (measurable) set B C Y. In general, Monge map may not exist (e.g. if m > n). However, when |X| = Y = n and a b are uniform then the Monge problem becomes the assignment problem and has bijective solution (Thorpe, 2018).
Contributions We introduce Hicrarchical Refinement (HiRef), an algorithm to scalably compute a full-rank align- ment between two equally-sized input datasets X and Y by solving a hierarchy of low-rank OT sub-problems. The success of this refinement is driven by a theoretical result, Proposition 3.1, stating that factors of an optimal low-rank coupling between X and Y co-cluster points X with their image under the Monge map. We use Proposition 3.1 recur- sively to obtain increasingly fine partitions of X and Y. At each scale, the solutions to low-rank OT sub-problems are bijections between the partitions of X and Y. Iterating to the finest scale gives bijection between X and Y.
Kantorovich Problem The Kantorovich problem (Kan- torovich, 1942) was introduced as a relaxation of the Monge problem. In contrast to the Monge problem, the Kantorovich problem allows mass-splitting and solution is always guar- anteed to exist. Define the transport polytope lIab as the following set of coupling matrices
I&.b = [P maXm P1, = a, m =b}.
(2)
Hierarchical Refinement constructs a multiscale partition of each dataset, and thus is related to (Gerber & Maggioni, 2017), which introduced a general framework for multiscale optimal transport using such partitions, and the earlier work
respectively with left (or "source") marginal a and with right (or "target") marginal b. For the cost c(., define the cost matrix C by C₁₃ = c(2, 81). In this discrete setting, the Kantorovich problem seeks a least cost coupling matrix
2

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
P E llabb between the probability vectors a, b associated to each measure v:
Hierarchical and Multiscale Approaches to OT Hier- archical optimal transport (Schmitzer & Schnörr, 2013) is variant of OT modeling data and transport at two scales, using Wasserstein distances as the coarse-scale ground costs. It has been applied to document representation (Yurochkin et al., 2019), domain adaptation (EI Hamri et al., 2022), sliced Wasserstein distances (Bonneel et al., 2015; Nguyen et al., 2022b) and to give a discrete formulation of transport between Gaussian mixture models (Chen et al., 2018; De- lon & Desolneux, 2020). These works build interpretable, coarse graiied structure into single coupling, rather than solving for a sequence of couplings at progressively finer scales as in the present work.
We(u, = min (C, P) PElla,b
(3)
The optimal value V) of (3) is called the C- Wasserstein distance between μ and v.
Sinkhorn Algorithm and the e-schedule The Sinkhorn algorithm (Cuturi, 2013) relaxes the classical linear- programming formulation of optimal transport by solving an entropy regularized version of (3),
We(p, v) = min (C, (P), PElla,b
-
where HP)) := == (sy (log.s Py- P.ylo entropy, and the parameter 6 > 0 is the regularization strength. The Sinkhorn algorithm improved the On log n) time complexity of classical techniques used for OT such as the Hungarian algorithm (Kuhn, 1955) and Network Simplex (Orlin, 1997; Tarjan, 1997) to (n log n) (Luo et al., 2023). As 4 O, the optimal coupling P** for (4) converges to a sparse optimal coupling for (3) at an ex- tremal point of the transport polytope (c.f. (Peyré & Cu- turi, 2019)). However, the number of iterations required scales as poly(1/e), diverging as decreases. technique used to improve this scaling is the e-schedule, an adaptive, menetone- and step-dependent set of entropy pa- Ctam. This anneals Problem 4 from high-entropy to low-entropy, gradually driving dense initial condition to a sparse solution with a log (1/e) rate
is the Shannon
Multiscale approaches to OT generalize hierarchical OT to progression of scales. Building on the semidiscrete approach of (Aurenhammer et al., 1998), (Mérigot, 201 1) uses Lloyd's algorithm to progressively coarse-grain the target measure. More recently, using a regular family of multiscale partitions (Definition C.3) on each dataset, (Ger- ber & Maggioni, 2017) formalize a general hierarchical approach to the Kantorovich problem (3). They propose: (i) solving a Kantorovich problem between the coarsest partitions of X and Y in their respective multiscale fami- lies; and (ii) propagation of the optimal coupling at scale E f1...., K 1} to initialize the optimization at scale I + 1. They take as input a chain of partitions and mea- sures across scales (X(), (1) + (Y(1), v1) + -.... (Y(*), v2)) where each dataset X, Y is identified with the trivial partitions X(*) {{x] x E X} and Y(x) {(y) y E YJ. At the finest scale K, (Gerber & Maggioni, 2017) recover the original datasets and near optimal coupling for (3).
+ (X(*), p.)) and
rameters E1 > <g
(Chen et al., 2023).
Low-rank Optimal Transport The nonnegative rank (M) of nonnegative matrix M > 0 is the smallest number of nonnegative rank-1 matrices summing to M: i.e. rk4 (M) is the smallest integer z such that there exist nonnegative vectors q1...., 9= >> 0 and Iy > 0 satisfying M = Σ -r. Let (r,b(r) == (P Na,b xk4(P) = r be the set of rank-r couplings. The low-rank Wasserstein problem for general cost matrix C is:
A naive implementation of the above idea requires quadratic memory complexity, but (Gerber & Maggioni, 2017) pro- pose several propagation strategies to mitigate this, follow- ing (Glimm & Henscheid, 2013; Oberman & Ruan, 2015; Schmitzer, 2016). These strategies use the optimal coupling at scale / to restrict the support of the coupling computed at the next scale using local optimality criteria. In the next section, we give our own such criterion, Proposition 3.1.
P arg min PC.
Pellab(())
From (Cohen& & Rothblum, 1993), each P E (r) may be decomposed P E1/B:) where g E A,, Q E Il..g and R E Hb This factorization was introduced optimal transport by (Scetbon al., 2021) in the context of the general low-rank problem (5). The factors Q and R constitute co-clusterings of datasets X and Y onto the same set of r components. Other factorizations have recently been proposed (Halmos et al., 2024), using Q, R and an intermediate latent coupling T to solve (5) where X and Y have r1 and r2 components, respectively.
3. Hierarchical Refinement
3.1. Low-rank optimal transport co-clusters source-target pairs under the Monge map
== Qdiag(1/g)R (6)
We first show that under few assumptions, the optimal low- rank factors (Q*, RR) for a variant of the low-rank Wasser- stein problem (5) have qualities suited to our refinement strategy. Specifically, we parameterize low-rank couplings P of rank-r using the factorization P = [ddiag(//g)R of (Scethon et al., 2021), fixing g E A, to be uniform. Define
3

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
clusters the datasets optimally, in that x and (x) assigned the same label by q' and r*. Using the same label set to partition X and Y automatically places the blocks of each partition in bijective correspondence. One then recurses on each pair of corresponding blocks (which we call a co-cluster) at the previous scale, until all blocks have size one. This guarantee holds despite that optimal (Q", R*) for (7) may not constitute an optimal triple (Q", R", ."") for the original low-rank problem (5) under the (Scetbon et al., 2021) factorization.
((1)
Y(1)
(()
Y(2)
X
Y(3)
(())
(())
A hierarchy-depth &. denotes the total number of times Al- gorithm refines the initial trivial partitions (X), (Y). The effective rank at scale t is Pt == I Tax given rank- annealing schedule (ri, "2: .7.,:...,PP Fix) for which Pv. divides n. The base rank is Taase _2
Figure 1. Hierarchical Refinement algorithm: low-rank optimal transport is used to progressively refine partitions at the previous scale, with the coarsest scale partitions denoted X y(I), and the finest scale partitions X(*), y(*)) corresponding to the individual points the datasets.
p= Note that n/ Pt is also the
size of each partition at scale t: n/p: X = |Y(+)), and that any sequence of any factorization of n corresponds to a rank-annealing schedule.
Proposition 3.2. For any n. there exists rank-schedule (f1...., (T)) factorizing such that all partitions of Algo- rithm 1 at level t & [0 K - 1] satisfy strict 1++1-Monge separability (Definition B.2). Let LROT denote an optimal rank-r for (7) over hard hard-partitions For any satisfying rank-schedule, the map returned by Algorithm I is optimal and supported on the graph of the Monge map T
the following variant of (5):
(Q' R*) arg min (C, Qdiag RT) (Q.R)
(7)
s.t. Q E Ia..g R llb.g g (1/r)1,
Proposition 3.1 below is the main structural result behind Hicrarchical Refinement. says that when optimal Q and R for (7) correspond to hard-clusterings (partitions) of each dataset, given by clustering functions q* X + [r].r* one has q* r T*, where T is Monge map.
Proof. Existence follows from the trivial (r1) = (n) rank- schedule. For any schedule (ri,-.. .Tx) satisfying Monge separability, applying the invariant of Proposition 3.1 induc- tively on f to level yields tuples ((x, T (x))) containing each x E X and its image (x) under the Monge map.
Y - [r].
Proposition 3.1 (Optimal low-rank factors co-cluster Monge pairs). Let X, Y C TRd with |X|= Y| = n, with cost matrix C that is strictly r-Monge separable (Definition B.2). Let a, b E An be uniform so that a Monge map T X - Y exists. If (Q*, R*) are minimizers of (7) and correspond to clustering functions q² X + [r], r* Y + [r], then for all
If the black-box subroutine LROT in Algorithm solves (7) optimally, then Hierarchical Refinement is guaranteed to recover a Monge map. In practice, we implement LROT using the low-rank solver (Halmos et al., 2024) and enforce that inner marginal g is uniform.
X E X one has q" (x) = 1* (x)).)).
The proof of Proposition 3.1 is in two steps. First, we use the existence of a Monge map and its coupling Pi to per- mute the cost C to cost Ct (Definition B. for which the identity matrix is a Monge map. Second, supposing that strict r-Monge separability (Definition B.2) holds, we show
Let Fi denote the q-th co-cluster at scale t generated by Hierarchical Refinement:
Te.g = { (xx,) : xx y Y(0) q
(8)
the solution to Problem 7 with cost C¹ is symmetric, so that minq 8 Ct, QR ) Returning to the coordinate frame of the original cost C, we find that Q- PRR, R, implying Proposition 3.1. We note that when r = 2, optimal Q R are hard-partitions (Lemma B.5) automatically satisfying one of the assumptions of Proposi-
minqell,, (Q'QQ )F-
where x(1) = (x pp
|==1 Y(O) = (Y(*)) ) and define the
co-clustering I'r at scale by:
r { (xx() v())) }r
tion 3.1.
At scale E [s]] Hierarchical Refinement refines I, to I P by running rank T(+1 low-rank optimal transport problem between uniform g = (1/re41)1re4) and measures sup- ported on each pair (t) Y!!)) in I2 for g E [pt], yielding
3.2. Hierarchical Refinement Algorithm
The Hierarchical Refinement algorithm (Algorithm 1) uses Proposition 3.1 to guarantee that each low-rank step co-
4

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Algorithm 1 Hierarchical Refinement
Definition 3.3 (Hierarchical block-coupling). For each
Require: Data x, Y: Low-rank OT solver LROT(-);
scale & [s]] given the Hierarchical Refinement co-cluster
Rank schedule (T1, (2,,,,,); Base rank Ybase (=1). Initialize:
partition Ir., the hierarchical block-coupling at scale t is defined by the matrix
1: t 0, To {X, Y)}
PM
Pt.
8 (x
(10)
2: while 7(x) (1) y E TT, such that
3:
min[[X]] (
(
Thase do
4: T + Ø
Without loss of generality, P(() may be block diagonalized into Pt square blocks, as discussed in Appendix B (see Equation (S13)). By Proposition 3.1, for any rank-schedule (ry)ā=1 satisfying Monge separability, the final p(s) corre- sponds to an optimal coupling supported on the graph of the Monge map T", P(x) == (id x T F/X.. While these intermediate couplings are never instantiated, one can still use them to define a transport cost (C, P(t)) at each scale. In Appendix B.8, we show the following bounds on the cost difference across scales.
5:
for (X1) v(U)
E T do
6: if [ix (1) 1
S Thase then (1) Y(1)
7:
T + T++1 U ((()
8:
else
9:
x 1 Exext" dx
10:
(y = I Eyev!! By.
11:
g++1 + (1/ret1)lre-
12:
(Q, R) + -LROT( (1xit)) Hy(t), Bt+1)
13:
for z 1 - r++1 do
Proposition 3.4. Let c(,,)) be a strictly-convex and Lips- chitz cost function, let (Ti,72, .Tr) be a rank-schedule, and let p(t) denote the coupling defined in (10), obtained
(+++) + Assign(X (1), Q,
14:
y((+1) +- Assign (()), R. 2)
15:
:
T41 + T U (x (t)
{)
from step of Algorithm 1. Define A (C, PIO) (C, (C,(((+)) p. Then,
16:
17:
end for
> Assign(S, M, = fa E S| arg max, M,.' =
0 0 Sr.+1 S Vclloo
1 Σ diam (Tr.))
(11)
19:
end if end for
Pt
20:
21:
( 1 ++1 1
where q indexes co-clusters Fe.g scale t, defined in (8).
22:
end while
Thus, the lower-bound implies that each step of refinement improves the coarse partition, and the upper-bound implies that the difference in solution value is bounded above by factor depending on the Lipschitz constant and the mean diameter of the coarse partitions at each level f. The proof of Proposition 3.4 roughly follows that of Proposition I of (Gerber & Maggioni, 2017). In Remark B.9, we discuss how Proposition 3.4 compares, noting that our result makes fewer geometric assumptions on multiscale partitions (x()) and (Y())) )r= and therefore does not quantify the rate of decay of diam (Tr.g).
23:
Output: Tx = ((x;, T(x;))
> Mapped pairs.
factors specific to this q E [pt]:
(Q, R) + LROT( Hygn, g++1)
(9)
For each q E [pe] we use the Q R from (9) to co-cluster x(O with y(!!) using rt+1 labels. Within this pair, each X; E x is assigned label E [re+]] by taking the argmax over the i-th row of Q, and likewise each y, E Y(!) is assigned the argmax over the j-th row of R. This corre- sponds to the Assign step in Algorithm 1, and coincides with the hard assignment of qt and r" for an optimal (Q*, R*) (Lemma B.5).
3.3. On the Rank-Annealing Schedule
As observed by (Forrow et al., 2019; Scetbon et al., 2021), rank behaves like temperature parameter, inverse to the strength e of entropy regularization. The correspondence between small and large rank implies that annealing in the parameter is, from the perspective of rank, analogous to initializing the optimization a low-rank coupling, and then gradually increasing the rank constraint from low to full, In Hierarchical Refinement, this gradual rank increase is accomplished implicitly. At each scale I. = 1,..., R. the im- plicit coupling p(() is made explicit in the hierarchical block coupling defined in equation (10). A rank-annealing sched- ule (ri...., rx)) describes the sequence of multiplicative fac- tors by which the rank of this explicit coupling will increase
The uniform constraint g = / in (7) enforces an even split of the dataset, which by Lemma B.5 ensures a partition at optimality (for Tı. = 2). Repeating for all q E [pt], one obtains co-clustering with r++1 components within each co-cluster at the previous scale, leading to total of Pr+1 T++1pt co-clusters at scale + 1 (Fig. 1). If the base-case rank Tbase is one, Algorithm returns bijection between X and Y as a collection of n tuples.
Note that Hierarchical Refinement defines an implicit hier- archy of block-couplings at each scale
5

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Method Costs over Sample Complexities
must store Te which is set of subsets of X and Y. To derive the time-complexity of Hierarchical Refinement, note that if n = r rank-r schedule at each layer re- quires instances of LROT over rapidly decaying dataset y
2.3
2.2
2.1
sizes. The complexity of low-rank OT (Scetbon et al., 2021; 2022; Halmos et al., 2024) is linear (Kn) for con- stant K = O(BLrd) dependent on B the number of in- ner Sinkhorn (Halmos et al., 2024) or Dykstra (Scetbon et al., 2021) iterations, L the number of mirror-descent steps, the rank of the coupling, and d the rank of the factorization of the cost matrix C. In this setting, for n power of т, the runtime of Algorithm I is given by the sum (') + + (()) + r rr p (,,"",))))))"") .. = ((nr log, n) for i = log, n, achieving linear space with log-linear time for constant ranks r, d.
2.0
O
1.9
Methods
1.8
Hierarchical Refinement
1.7
Sinkhom
1.6
ProgET
1.5
102
10a
10a
10a
Sample Size
Figure 2. Primal OT cost for varying sample size on the synthetic half-moon S-curve dataset of (Buzun al., 2024) for BiRef, Sinkhorn, and ProgOT
In cases where the cost matrix does not admit low-rank factorization C = UVT, i.e., when d = O(n), one requires (nn²) space to store the cost matrix and Hierarchical Re- finement exhibits time complexity Ö(n²), as in Sinkhorn. For kernel costs such as squared Euclidean cost, as noted in (Scetbon et al., 2021), one may efficiently compute a (d + 2) dimensional factorization where d is the ambient dimen- sion, to achieve log-linear scaling with exact distances. We also use the sample-linear algorithm of (Indyk al., 2019) to compute approximate factorizations for distances c(.,)) satisfying metric properties such as the triangle inequality (c.g. Euclidean distance, see Appendix E.I). At each level, pairing such sample-linear approximations with each low- rank step only requires On log, n) time. We observe this scaling empirically, reported in Fig. S2.
at successive scales. The partial products of these, denoted (Pl,...:Pe)) are the ranks of the couplings P(1) p(,),..., P(*), Note that small values of T; generate coarse partitions of the points at the next scale, while large values of ri generate finer partitions at the next scale.
We now turn to the question of how to efficiently choose such a schedule under given memory constraints. For an integer n, Algorithm has log-linear complexity for depth log, n (Section 3.4). However, the large constants required by low-rank OT in practice encourage minimizing the number of calls to LROT as subroutine, that if memory permits, it may be advantageous decrease the depth by storing couplings of higher rank. If desired, mem- ory constraints can be enforced by imposing a maximum rank Fmax 2 F₁ for all f E [k] to ensure Hierarchical Refine- ment only requires Onnmaax space at each step. Thus, we seek factorizations with minimal partial sums of ranks while remaining below desired memory-capace
4. Experiments
We benchmark Hierarchical Refinement (HiRef) against the full-rank OT methods Sinkhorn (Cuturi, 2013), ProgOT (Kassraie et al., 2024), and mini-batch OT (Genevay et al., 2018: Fatras al., 2020; 2021b). We additionally bench- mark against the low-rank OT methods LOT (Scetbon et al., 2021) and FRLC (Halmos et al., 2024). We use the de- fault implementations of Sinkhorn, ProgOT, and LOT in the high-performance ott - jaax library (Cuturi et al., 2022). In particular, Sinkhorn is run with the default entropy regular- ization parameter of e 0.05. We also benchmark against the multiscale method MOP (Gerber & Maggioni, 2017). which requires multiscale partitions of the input datasets akin to family of dyadic cubes across scales alignments. This leads to transport cost that depends on the choice of this partition. For simplicity, we choose the default partitions of MOP which are computed from the GMRA (Geometric Resii-eesolttion Analysis) R package.
min (ru)
Px = n, ri S Tmax s.t.
(12)
The above optimization assumes base-rank Tbase of 1; we describe how to handle the general case in Appendix E.1. Importantly, the recursive structure min Σ enables a dynamic programming approach to (12), storing table of factors up to Tmax to optimize (12) in (f time. Assum- ing K. Tmax are small constants chosen to ensure that all matrices can fit within memory, determining the optimal rank-schedule with respect to K, n, Tmax is a simple linear- time procedure.
min (r1 + 2
compute
3.4. Complexity and Scaling of Hierarchical Refinement
For two datasets X, Y of size n, the space complexity of Hierarchical Refinement is (nn), since at each level, one
6

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Hierarchical- Rentreemelll Map
Sinkhorn Berycentric Map
Optimal Map (Dual Revised Simplex)
a.
3
b.
Figure 3. Comparison of the Hicrarchical Refinement Mapping, the Sinkhorn Barycentric Map, and an optimal map computing using dual revised simplex for the a. Half-moon and S-curve dataset (Buzun et al., 2024) of 4096 points (512 points for dual revised simplex) and b. Checkerboard dataset (Makkuva et al., 2020).
4.1. Evaluation on Synthetic Datasets.
to 16384 points. HiRef yields solutions with comparable primal cost to ProgOT and Sinkhorn on the sample sizes where all methods run.
We first evaluate the performance of Hierarchical Refine- ment against optimal transport methods returning primal couplings, namely Sinkhorn (Cuturi, 2013) (as implemented in ott-tjeee jax (Cuturi et al., 2022)) and ProgOT (Kassraic et al., 2024). We evaluate the methods with respect to the Wasserstein-1 and Wasserstein-2 distance on an alignment of 1024 pairs of samples on the Checkerboard (Makkuva et al., 2020), MAFMoons and Rings (Buzun et al., 2024), and Half-Moon and S- S-Curve (Buzun al., 2024) synthetic datasets (Fig. 3, Table S6).
We also find that HiRef achieves an OT cost that is com- petitive with the dual revised simplex solver (Huangfu & Hall, 2018), a solver which only scales up to 512 points (Ta- ble S4). This solver computes an optimal coupling, unlike ProgOT and Sinkhorn which rely on entropic regularization. While we benchmark Sinkhorn in place of mini-bater OT the synthetic datasets due to their limited complexity, we also evaluate the multi-scale method MOP on the 512 point instance (Table S4). Although MOP outputs a fast approximation to optimal transport, its primal cost on the Checkerboard (Makkuva et al., 2020) dataset is twice as high as that of the other methods, and it performs signifi- cantly worse on the MAF Moons & Rings and Half Moon & S S-Curve datasets (Buzun et al., 2024).
All methods are similarly effective at minimizing the pri- mal OT cost (C,P)) with small absolute difference in cost between the final couplings. Hierarchical Refinement achieves slightly lower primal cost on 4 out of the 6 evalua- tions. Notably, there is a massive difference the number of non-zero entries (defined as entries P.y 10 -5 in the cou- plings output by HiRef, Sinkhorn, and ProgOT (Table S3). Specifically, across the experiments Hi f outputs bijec- tion with exactly 1024 non-zero elements in the coupling matrix, equal to the number of aligned samples. In con- strast, Sinkhorn and ProgOT output couplings with 624733 to 678720 and 271087 to 337258 non-zero entries.
Lastly, we observe that Hierarchical Refinement scales to over a million points, two orders of magnitude greater than ProgOT and Sinkhorn, two full-rank OT methods that com- pute global alignments. We find HiRef scales linearly with the size of the problem instance (Fig. S2a) in contrast the quadratic scaling in time complexity of Sinkhorn (Fig. S2b).
We evaluate the scalability of Hierarchical Refinement rela- tive other full-rank solvers varying numbers of sam- ples from the Half Moon & S-Curve (Buzun et al., 2024) synthetic dataset. We vary the rank from 25 = 32 (64 points aligned) up to 220 = 1048576 points (2097152 points aligned) R²? the latter dataset of a size that is beyond the capabilities of current optimal transport solvers. We observe that Sinkhorn (Cuturi, 2013) and ProgOT meth- ods which produce dense mappings- - require coupling matrix with (n)²) non-zero entries and thus run only up
4.2. Large-scale Matching Problems and Transcriptomies
Recently, optimal transport has been applied to single-cell and spatial transcriptomics datasets to compute couplings between cells taken from different timepoints from develop- mental processes or perturbations (Schiebinger et al., 2019; Lavenant et al., 2024; Bunne et al., 2022; Huizing et al., 2024; Halmos et al., 2025b; Klein et al., 2025). However, the size of current datasets (Chen al., 2022) (> 100k cells)
7

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
logical replicates. Our "source" dataset (X+, ,,)) is replicate 3 of slice 2, while our "target" dataset (X2, S²) is replicate
Table 1. Cost Values (C,P)r Across Later Embryonic Stages
Method
E12-13.5
E13-14.5
E14-15.5
E15-16.5
2 of slice 2, following the expression transfer task described (Clifton et al.. 2023) between these two slices. Each dataset has roughly 84k spots, where memory constraints prohibit instantiation a full-rank alignment as a matrix. Thus, solvers such as Sinkhorn (Cuturi, 2013) and ProgOT (Kassraie et al., 2024) are unable to run on the dataset.
HiRef
14.35
13.78
14.29
12.79
Sinkhom
MB 128
14.86
14.14
14.75
13.32
MB 1024
14.45
13.86
14,43
12.91
FRLC
15.47
14.64
15.51
14.00
We use only spatial information when building map between the two slices, using the spatial Euclidean cost Cj == ||s - -sill.. after registering spatial coordinates S {s]) and S = { with an affine trans- formation. We gauged the quality of the HiRef align- ment (Fig. 4a), using gene expression abundances of five aeaatiallyvaryying genes. Specifically, we observe that ex- pression vector v of gene Slc/7a7 in the source slice ( Fig. 4b) when transferred target slice through the bijective mapping output by HiRef, denoted as v (Fig. 4c), closely matches the observed expression vector v² of Slc17a7 in the target slice (Fig. 4d) with cosine similarity equal to 0.8098. For genes Slc Grm4, Olig1, Gadl, Peg1O, the corre- sponding cosine similarities between the transferred and observed expression vectors are 0.8098, 0.7959, 0.7526, 0.4932, 0.6015, respectively.
has exceeded the capacity of existing full-rank solvers, re- quiring low-rank approximations of the coupling (Scetbon et al., 2023; Klein et al., 2025; Halmos et al., 2025a) to produce alignments.
We evaluate whether the full-rank solver of Hierarchical Re- finement exhibits competitive alignments for such datasets. Specifically, we analyze the mouse organogenesis spatiotem- poral transcriptomic atlas (MOSTA) datasets, which include spatial tramscriptomics data from mouse embryos at succes- sive 1 1-day iime-intervals with increasing number n of cells at each stage: E9.5 (n = 5913), E10.5 (n 18408), E11.5 (n = 30124), E12.5 (n = 51365), E13.5 (n 77369), E14.5 (n 102519), E15.5 (n 113350), and E16.5 (n = 121767). For the cost we use the Euclidean distance C. = ||x: - Yyllz in 60-dimensional PCA space of expres- sion vectors, so Y, E R
For comparison, we also ran the low-rank methods FRLC (Halmos et al., 2024) and LOT (Scetbon et al., 2021) with and without subsampling, reporting their best scores, as discussed in Section D.3. For the gene Slc17a7, FRLC's cosine similarity was 0.2373, while L.TT's cosine similar- ity was 0.3390. For all five genes Slc17a7, Grm4, Oligl, Gadl, Peg 10, FRLC's s scores were (0.2373, 0.2124, 0.1929, 0.0963, 0.1550, respectively, while LOT's scores were 0.3390, 0.2712, 0.3186, 0.1666, 0.1080. Across all five genes HiRef scores were at least twice those of FRLC or LOT (Table S7) with gene abundances shown in Fig. S1. On the same task, we compared against MOP, the method of (Gerber & Maggioni, 2017), whose scores for the five genes were: (0.5211, 0.4714, 0.5972, 0.3571, 0.2719). Finally, we also benchmarked against mini-batch OT using batch sizes ranging from 128 to 2048 in powers of two, whose best scores (0. (0.4434, 08722 0. 0.4912, 0.5683) were more comparable to that of the performance of HiRef. Across all methods and genes compared in Table S7, HiRef had greatest cosine similarity scores in the expression transfer task, while also having lowest transport cost. Further exper- imental details are in Section D.3.
Sinkhorn and ProgOT are unable to produce alignments for the stages beyond E10.5 (n 18408 cells), whereas HiRef, the low-rank solvers, and mini-batch OT (batch- sizes B- = 128 to B- 2048) are able to continue scaling to > 10 (Table 1, Table S6). We observe that the Kantorovich cost of fiRef is consistently lower than all other methods for all timepoints (Table 1, Table S6).
HiRef achieves a substantially lower cost than the low- rank solvers FRLC and LOT for rank 40, even though HiRef relies on low-rank optimal transport (FRLC) as subroutine. This result underscores the empirical trend observed in Fig. S3, where the refinement step of f progressively decreases the primal cost of coarser low-rank couplings (Proposition 3.4). While the mini-batch solvers exhibit competitive scaling up to the last pair, the primal cost of mini-batch is higher for all tested batch-sizes (Table S6). Unlike Hef,, mini-batch OT does not compute global alignment and exhibits batch-size dependent error.
4.3. MERFISH Brain Atlas Alignment
We ran Hi Ref on two slices of MERFISH Mouse Brain Receptor Map data from Vizgen to test whether Hi Ref can produce biologically valid alignments using the only spa- tial densities of each tissue. These spatial transeriptomics data consist of spatial and gene expression measurements at individual spots in three full coronal slices across three bio-
4.4. ImageNet Alignment
We demonstrate the scalability of Hierarchical Refinement on a large-scale and high-dimensional dataset by align- ing 2048-dimensional embeddings of 1.281 million images from the ImageNet ILSVRC dataset (Deng et al., 2009;
8

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
a.
b.
C.
d.
Slc1 Sic17a7 counts
30
Figure a. Hierarchical Refinement alignment on MERFISH mouse brain data, using only spatial coordinates. Abundance v of gene SIcI7a7 in the source slice. c. Predicted Sic17a7 abundance v from the source slice to the target slice, through the HiRef coupling. d. Abundance v² of the same gene the target slice. Transferred abundances & have cosine similarity 0.8098 with true abundances v² the target.
Optimal transport has been successfully applied in deep learning frameworks, such as OT flow- now-matching (Tong et al., 2024), computer vision and point cloud registration, (Yu et al., 2021; Qin et al., 2022), among many others. The mini-batch procedure used to train many of these methods involves sampling two datasets XB ~ μ and YB ~ v with batch-size B and aligning them with Sinkhorn at every train- ing iteration. Hi Ref suggests an alternative approach: one can precompute millions of globally aligned pairs and then sample XB fl and the optimal mapping T(XB) ~ y by in- dexing into these precomputed pairs. This approach applies to any loss function dependent on an OT alignment.
Table Cost Values (C, P)r x for ImageNet (Deng et al., 2009:
Russakovsky al., 2015) Alignment Task.
Method
Hidef
MB 128
MB 256
MB 512
MB 1024
FRLC
OT Cost
18.97
21.89
21.11
20.34
19.58
24.12
Russakovsky et al., 2015). Each image is embedded us- ing using the ResNet50 architecture (He et al., 2016), and we construct two datasets, X and Y, by taking a random 50:50 split of the embedded images. We align X and Y using Hi Ref, FRLC, and mini-batch OT with batch-sizes ranging from B = 128 to B = 1024. ProgOT, Sinkhorn, and LOT could not be run on the datasets due to memory constraints. f yielded a primal OT cost of 18.974, while FRLC (Halmos et al., 2024) solution had primal OT cost of 24.119 for rank r 40 and mini-batch OT has costs of 21.89 (B- 128) to 19.58 (B- = 1024) (Table 2).
Hierarchical Refinement may also be useful in neural OT approaches which learn a continuous Monge map between the densities of two datasets. For example, (Seguy et al., 2018) minimize a loss ming 1' (x;) T(xx 112 be- tween a neural network Ta with parameters 0 and Monge map T over samples x; f (Remark 1.1.). Thus. the procedure outlined above may be used to directly regress a neural network Ta on the Monge map T without the bias of mini mini-batching or entropy.
5. Discussion
Hicrarchical Refinement computes the Monge map between large-scale datasets in linear space, but has several limita- tions. First, we currently assume that the datasets X and Y have the same number of samples. In many machine learning applications, this is not a limiting factor, as one generally seeks to pair an equal number of source points x to target points y. Second, while Hicrarchical Refine- ment scales linearly in space and log-linearly in time, it still involves a constant dependent on the low-rank OT sub- procedure used- this underscores the need to accelerate and stabilize low-rank OT solvers further (Scetbon & Cu- turi, 2022; Halmos et al., 2024). Finally, while Hierarchical Refinement guarantees optimal solution given an optimal black-box low-rank solver (Proposition 3.1), the low-rank solvers (Scetbon et al., 2022; Halmos et al., 2024) used in practice are not necessarily optimal, owing to the non- convexity of low-rank problems.
6. Conclusion
We introduce Hierarchical Refinement (e) an algo- rithm to solve optimal transport with linear complexity in the number of points, making sparse, full-rank optimal transport feasible for large-scale datasets. Our algorithm leverages that low-rank optimal transport co-clusters points with their image under the Monge map, refining bijections between partitions of each dataset across hierarchy of scales, down to a bijective Monge map between the datasets at the finest scale. Hierarchical Refinement couplings achieve compa- rable primal cost to couplings obtained through full-rank entropic solvers, and scales to datasets with over a million points, opening the door to applications previously infeasi- ble for optimal transport.
9

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Acknowledgements
using neural optimal transport. Nature Methods, 20(1) 1): 1759-1768, 2023.
We thank Henri Schmidt for many helpful conversa- tions. This research was supported by NIH/NCI grant U24CA248453 B.J.R. J.G. supported by the Schmidt DataX Fund at Princeton University made possible through major gift from the Schmidt Futures Foundation.
Buzun, N., Bobrin, M., and Dylov, D. V. Expectile reg- ularization for fast and accurate training of neural op- timal transport. In Advances in Neural Information Processing Systems, volume 37, pp. 119811-119837, 2024. URL https:: wtw.pps://ppeeeeeiieeeee /pppneeeeeeeeeeeeeeeeeeeeee net forum? Ad455aPaP....
Impact Statement
Chen, A., Liao, S., Cheng, M., Ma, K., Wu, L.,, Lai, Y.,
Optimal transport has emerged as powerful tool in gener- ative modeling, yet its widespread use has been limited by scalability constraints. HiRes Ref overcomes this limitation by enabling the application of OT to datasets with millions of points. This advancement paves the way for integrating OT into large-scale deep generative models and modern vision and language tasks.
Qiu, X., Yang, J., Xu, J., Hao, S., et al. Spatiotempo- ral transcriptomic atlas of mouse organogenesis using DNA nanoball-pattemeet arrays. Cell, 185710): 1792, 2022.
Chen, J., Chen, L.,, Liu, Y. P., Peng, R., and Ramaswami, A. Exponential convergence of Sinkhorn under regular- ization scheduling. In SIAM Conference on Applied and Computational Discrete Algorithms, pp. 180-188. SIAM, 2023.
As with any computational tool which may advance large- scale generative modeling, there are potential issues with bias in training datasets and possibility of misuse. Use of HiRef f in applications should be careful and transparent about these risks and utilize appropriate mitigation strate- gies.
Chen, X. and Price, E. Condition number-free query and active learning of linear families. CoRR, abs/1711.10051, 24, 2017.
Chen, Y., Georgiou, T. T., and Tannenbaum, A. Optimal transport for Gaussian mixture models. IEEE Access, 6269-6278, 2018
Code Availability
Our implementation of Hierarchical Refinement is available at http://ggihubbccommmpphlllgguuup
Clifton, K., Anant, M., Aihara, G., Atta, L.,, Aimiuwu, O. K., Kebschull, J. M., Miller, M. L., Tward, D., and Fan, J. STalign: Alignment of spatial transcriptomics data using diffeomorphic metric mapping. Nature Communications, 14(1):8123, 2022.
References
Aurenhammer, F. Hoffmann, F. and Aronov, B. Minkowski-type theorems and least-squares clustering. Algorithmiica, 20:61-76, 1998.
Cohen, J. E. and Rothblum, U. G. Nonnegative ranks, de- compositions, and factorizations of nonnegative matrices. Linear Algebra and its Applications, 190:149-168 1993.
Birkhoff, G. Tres observaciones sobre el algebra lineal. Univ. Nuc. Tucumun, Ser. A, 5:147-154, 1946.
Courty, N., Flamary, R., and Tuia, D. Domain adaptation with regularized optimal transport. In European Confer- ence on Machine Learning and Knowledge Discovery in Databases, pp. 274-289, Springer, 2014.
Bonneel, N., Rabin, J., Peyré, G., and Pfister, H. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51:22-45, 2015.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26:2292- -3000, 2013.
Brenier, Y. Polar factorization and monotone rearrangement of vector- valued functions. Communications pure and applied mathematics, 44(4):375-417, 1991.
Cuturi, M., Mang-Papaxanthos, L., Tian, Y., Bunne, C.. Davis, G., and Teboul, O. Optimal Transport Tools (OTT); A JAX Toolbox for all things Wasserstein. arXiv preprint arXiv:2201.12324 2022.
Bunne, C., Papaxanthos, I.., Krause, A., and Cuturi, M. Proximal optimal transport modeling of population dy- namics. In International Conference on Artificial Intelli- gence and Statistics, volume 25, pp. 6511-6528. PMLR, 2022.
De Bortoli, V., Korshunova, I., Mnih, A., and Doucet, A. Schrödinger bridge flow for unpaired data translation. Advances in Neural Information Processing Systems, 37: 103384-103441, 2024. URL https: htap::://pppereee //pp///eee.....e/////// net/forum 1F221
Bunne, C., Stark, S. G., Gut, G., del Castillo, J. S., Levesque, M., Lehmann, K.,V., Pelkmans, L., Krause, A., and Rätsch, G. Learning single-cell perturbation responses
10

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
De Loera, J. A. and Kim, E. D. Combinatorics and geometry of transportation polytopes: An update. Discrete Geome- try and Algebraic Combinatorics, 625:37-76, 2013.
Genevay, A., Peyré, G., and Cuturi, M. Learning gen- erative models with Sinkhorn divergences. In In- ternational Conference Artificial Intelligence and Statistics, volume 84, pp. 1608 1617. PMLR, 2018. URL ptattps://proceeeee genevay html.
Delon, J. and Desolneux, A. A Wasserstein-tyie distance in the space of Gaussian mixture models. SIAM Journal on Imaging Sciences, ((2):936-9970) 2020.
ml mll. press/v84/
Gerber, S. and Maggioni, M. Multiscale strategies for com- puting optimal transport. Journal of Machine Learning Research, 18(72):1-32, 2017.
Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. ImageNet: A large-scale hierarchical image database. In IEEE Conference Computer Vision and Pattern Recognition, pp. 248-255. IEEE, 2009.
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. A mathematical perspective on Transformers. arXiv preprint arXiv:2312. 10794, 2023.
El Hamri, M., Bennani, Y., and Falih, Hierarchical optimal transport for unsupervised domain adaptation. Machine Learning, (1:1))444994488 2022.
Glimm, T. and Henscheid, N. Iterative scheme for solving optimal transportation problems arising in reflector de- sign. International Scholarly Research Notices, 2013(1): 635263, 2013.
Fan, J., Liu, S., Ma, S., Zhou, H.-M., and Chen, Y. Neural Monge map estimation and its applications. Transactions Machine Learning Research, 2023. URL https: //openrevien net / forum? idddmmms0ssss......
Halmos, P., Liu, X., Gold, J., and Raphael, B. Low- Rank Optimal Transport through Factor Relaxation with Latent Coupling. In The Thirty- Annual Conference Neural Information Processing Systems, 2024. URL https: ooperevveeeeeeeeeeeeeeeeeeeeeeee net/ forum? id=hGgkdFF2hR.
Fatras, K., Zine, Y., Flamary, R., Gribonval, R., and Courty, N. Learning with minibatch Wasserstein: asymptotic and gradient properties. In International Conference on Arti- ficial Intelligence and Statistics, volume 108, pp. 2131- 2141. PMLR, 2020. URL http: http::////coceeeeeg...e. mlr. press / /v10//00////// fat html.
Halmos, P., Gold, J., Liu, X., and Raphael, B. J. Learn- ing latent trajectories in developmental time series with -aidden-Marovv optimal transport. In International Con- ference on Research in Computational Molecular Biology, pp. 367-370. Springer, 2025a.
Fatras, K., Séjouméé, T., Flamary, R., and Courty, N. Un- balanced minibatch optimal transport; applications to do- main adaptation. In International Conference on Machine Learning, volume 139, pp. 3186-3197. PMLR, 2021a. URL http: /ttett://prrceee http://pccceee....e/e///// www.:ppeessvvll fatras21a html.
Halmos, P., Liu, X., Gold, J., Chen, F., Ding, L., and Raphael, B. J. DeST-OT: Alignment of spatiotemporal transcriptomics data. Cell Systems, 16(2). 2025b.
Fatras, K., Zine, Y., Majewski, S., Flamary, R., Gribon- val, R., and Courty, N. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792. 2021b.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual leam- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7770777.1000 2016.
Finlay, C., Jacobsen, J.-H., Nurbekyan, L., and Oberman, A. How to train your neural ODE: the world of Jacobian and kinetic regularization. In International Conference on Machine Learning, pp. 3154-3164. PMLR, 2020.
Huangfu, Q. and Hall, J. A. Parallelizing the dual revised simplex method. Mathematical Programming Computa- tion, 10()):11-142,, 2018.
Forrow, A., Hütter, J.-C., Nitzan, M., Rigollet, P., Schiebinger, G., and Weed, Statistical optimal transport via factored couplings. In International Conference Artificial Intelligence and Statistics, volume 89, pp. 2454 2465. PMLR, 2019, URL https: /ppoceedingg. ml r. prwss / v 8 / forrow1 9a. html.
Huizing. G.-J., Peyré, G., and Cantini, L. Learn- ing cell fate landscapes from spatial transcriptomies using Fused bioRtii:2022.07.26666055. 60.2241,, 2024.
bioRxiv preprint
Indyk, P. Vakilian, A., Wagner, T., and Woodruff, D. P. Sample-optimal low-rank approximation of distance ma- trices. In Conference on Learning Theory, volume 99, pp. 1723-1751. PMLR, 2019.
Frieze, A., Kannan, R., and Vempala, S. Fast Monte-Carlo Algorithms for Finding Low-rank Approximations. J. ACM, 51(6): 1041, nov 2004. ISSN 0004-5411. doi: 10.1145/1039488 1039494. URL https: htmps://dd... 1rg/100.144/1034888 10399444
Kantorovich, L. On the translocation of masses. Doklady Akademii Nauk SSSR, 37(7-8):227-229) 1942.

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Kassraie, P., Pooladian, A.-A., Klein, M., Thornton, J., Niles-Weed, J., and Cuturi, M. Progressive entropic opti- mal transport solvers. Advances in Neural Information Processing Systems, 37:19561-19599 2024.
Makkuva, A., Taghvaei, A., Oh, S., and Lee, J. Optimal transport mapping via input convex neural networks. In- ternational Conference on Machine Learning, 119:6672- 6681, 2020.
Klein, D., Uscidda, T., Theis, F. J., and Cuturi, M. Gen- erative entropic neural optimal transport to map within http///ppeeee...e/////e// petettss://ppeeeeeee
Melnyk, I., Mrouch, Y., Belgodere, B., Rigotti, M., Nit- sure, A., Yurochkin, M., Greenewald, K., Navratil, J., and Ross, J. Distributional preference alignment of LLMs via optimal transport. Advances in Neural In- formation Processing Systems, 2024. URL https: // //openreview www.:/e..eeeee.eee.eee....ee.e.e www.mrrrrdioooo Seezummid=2LLttgf66TY
and across space, 2024. URL https: htt::tt net wow.nmiiiiggLLBeEE0000
Klein, D., Palla, G., Lange, M., Klein, M., Piran, Z., Gander, M.,
L., Sterr, M., Saber, L., Jing. C., Bastidas- Ponce,, A., Cota, P., Tarquis-Meding M., Parikh, S., Gold, I., Lickert, H., Bakhti, M., Nitzan, M., Cuturi, M., and Theis, F. J. Mapping cells through time and space with moscot. Nature, pp. 1-11, 2025.
Mérigot, Q. A multiscale approach to optimal transport.
Computer Graphics Forum, 3065):1583-1592, 2011.
Monge, G. Mémoire sur la théorie des déblais et des rem- blais. Mem. Math. Phys. Acad. Royale Sci., pp. 666-704, 1781.
Kornloov, N., Mokrov, P., Gasnikov, A., and Korotin, A. Optimal flow matching: Learning straight trajectories in just one step. Advances in Neural Information Processing Systems,
Nguyen, K., Nguyen, D., Pham, T., and Ho, N. Improving mini-batch optimal transport via partial transportation. In Proceedings of the 39th International Conference Machine Learning, 2022a.
2024. URL https:// forumm wiciummiiii-mmaaa KgmuuDDVVV..
openreview beennrevewwnet//eee net
Korotin, A., Li, L., Genevay, A., Solomon, J. M., Filippov, A., and Burnaev, E. Do neural optimal transport solvers work? continuous Wasserstein-2 benehmark. Advances in Neural Information Processing Systems, 34:14593- 14605, 2021.
Nguyen, K., Ren, T., Nguyen, H., Rout, L., Nguyen, T. M., and Ho, N. Hierarchical sliced Wasserstein distance. International Conference on Learning Representations, 2022b.
Korotin, A., Selikhanovych, D., and Burnaev, E. Neural optimal transport. International Conference on Learning Representationss, 2023. URL https: oppnreviee.. net/ www.orummiiddBee I
Oberman, A. M. and Ruan, Y. An efficient linear program- ming method for optimal transportation. arXiv preprint arXiv: 100.0660,,,,68, 2015.
Orlin, J. B. A polynomial time primal network simplex algorithm for minimum cost flows Mathematical Pro- gramming, 78(2)1009-19,, 1997.
Kuhn, H. W. The Hungarian method for the assignment problem. Naval Research Logistics Quarterly, 2(1-2): 83-97, 1955.
Peyré, G. and Cuturi, M. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 1 11(56))355-607, 2019.
Lavenant, H., Zhang, S., Kim, Y.-H., and Schiebinger, G. Toward mathematical theory of trajectory inference. The Annals of Applied Probability, 34(1A):428 500, 2024.
Qin, Z., Yu, H., Wang, C., Guo, Y., Peng, Y., and Xu, K. Geometric transformer for fast and robust point cloud reg- istration. In IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 11143-11152, 2022.
Li, X., Chen, J., Chai, Y., and Xiong, H. GiLOT: Inter- preting generative language models via optimal trans- port. International Conference Machine Learning, 2024. URL https:://// wwwpeeeeeiwwweeeeeeerr..eee formm id=cKL25sGjxx.
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Rad- ford, A., Chen, M., and Sutskever, I. Zero-shot text-to- image generation. International Conference on Machine Learning, 139:8821-8831, 2021.
Lin, C.-H., Azabou, M., and Dyer, E. L. Making transport more robust and interpretable by moving data through small number of anchor points. International Conference on Machine Learning, 139:6631, 2021.
Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma. S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L. ImageNet large scale visual recognition challenge. International Journal of Computer Vision 115(3):211225) 2015.
Luo, J., Yang, D., and Wei, K. Improved complexity analysis of the sinkhorn and greenkhorn algorithms for optimal transport. arXiv preprint arXiv:2305.1499 2023.
12

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Sander, M. E., Ablin, P., Blondel, M., and Peyré, G. Sink- formers: Transformers with doubly stochastic attention. In International Conference Artificial Intelligence and Statistics, pp.
of gene expression in tissue sections by spatial transcrip- tomics. Science, 353(6294):78 -2, 2016.
Tarjan, R. E. Dynamic trees as search trees via Euler tours,
PMLR, 2022.
applied the network simplex algorithm. Mathematical Programming, 78(2):169-177, 1997.
Scetbon, M. and Cuturi, M. Low-rank optimal transport: Approximation, statisties and debiasing. Advances in Neural Information Processing Systems, 35:6802-6814, 2022.
Tay, Y., Bahri, D., Yang, L., Metzler, D., and Juan, D.-C. Sparse Sinkhorn attention. International Conference on Machine Learning, 19:9438 04488,947, 2020.
Scetbon, M., Cuturi, M., and Peyré, G. Low-rank Sinkhorn factorization. International Conference Machine Learning, pp. 9344-9354, 2021.
Thorpe, M. Introduction to optimal transport. Notes of Course at University of Cambridge, 2018.
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with mini- batch optimal transport. Transactions on Machine Learn- ing Research, 2024. URL https::////... /ttp:///epereiieeeeeemee net / forummruummmmmm 9dnnniddddddnnnc3m
Scetbon, M., Peyré, G., and Cuturi, M. Linear-time Gro- mov Wasserstein distances using low rank couplings and costs. International Conference on Machine Learning, pp. 19347-19365, 2022.
Scetbon, M., Klein, M., Palla, G., and Cuturi, M. Un- balanced low-rank optimal transport solvers. Advances in Neural Information Processing Systems, 36:52312- 52325, 2023.
Wolf, F. A., Angerer, P., and Theis, F. J. SCANPY: Large- scale single-cell gene expression data analysis. Genome Biology, 19:155, 2018.
Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subrama- nian, V., Solomon, A., Gould, J., Liu, S., Lin, S., and Berube, Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in repro- gramming. Cell, 176(4):928-943 2019.
Yang. K. D., Damodaran, K., Yenkatachalapathy, S., Soyle- mezoglu, A. C., Shivashankar, G. V., and Uhler, C. Pre- dicting cell lineages using autoencoders and optimal trans- port. PLOS Computational Biology, 16(4))e1007828, 2020.
Schmitzer, B. A sparse multiscale algorithm dense optimal transport. Journal of Mathematical Imaging and Vision, 56:2366:225555 2016.
Yu, H., Li, F., Saleh, M., Busam, B., and Ilie, S. CoFiNet: Reliable coarse-to-fine correspondences for robust point- cloud registration. Advances in Neural Information Pro- cessing Systems, 34:23872-23888 2021.
Schmitzer, B. and Schnörr, C. A hierarchical approach to optimal transport. In International Conference on Scale Space and Variational Methods in Computer Vision, pp. 452-464. Springer, 2013.
Yurochkin, M., Claici, S., Chien, E., Mirzazadeh, F., and Solomon, J. M. Hierarchical optimal transport for doc- ument representation. Advances in Neural Information Processing Systems, 32, 2019.
Seguy, V., Damodaran, B. B., Flamary, R., Courty, N., Ro- 1et, A., and Blondel, M. Large-scale optimal transport and mapping estimation. International Conference on Learning Representations, 2018.
Zeira, R., Land, M., Strzalkowski, A., and Raphael, B. J. Alignment and integration of spatial transcriptomics data. Nature Methods, 19(5):567-575, 2022.
Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics, (4):1-11, 2015.
Sommerfeld, M., Schrieber, J., Zemel, Y., and Munk, A. Optimal transport: Fast probabilistic approximation with exact solvers. Journal of Machine Learning Research, (105):1-23, 2019.
Ståhl, P. L., Salmén, F. Vickovic, S., Lundmark, A., Navarro, J. F., Magnusson, J., Giacomello, S., Asp, M., Westholm, O., and Huss, M. Visualization and analysis
13

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
A. Hiierarchical-Refeemeent Algorithm
Algorithm 2 Hierarchical Refinement for Full-Rank OT
Require:
Datasets X
{x: Y
{y: Low-rank OT solver LROT(+); Rank schedule (ri, 12 Tx):
Base rank Tbasc IT--1 Ti
(e.g. 1).
Initialize:
Tn {X, Y)}
2: while (x(()) Y(t) E ET, such that
3:
min[[X(0] |
Thase do
4: T 2
for (X4), y(1) E do
5:
6: if [ (1) 1
} Tbase then
T411 T4t U ((x() v(O))
7:
8:
else
9:
O 1
E ExEx(s 8x
10:
My(!) 1 Exev!" dy.
11:
ge+1 t (1/re41)1
12:
(Q, R) + (ROT(UxX) Hy(t), gt+1)
13:
for 1 r++1 do
((++1) + Assign(X') (1), Q.2)
14:
(((ign((()),aanoao y((+1) Assign (()), R. 2)
15:
F++1 + U (X (t+ v(++)))
16:
17:
end for
Assign(S, M, = {s
arg max, M. = =]
end if end for
20:
21:
t t + 1
22:
end while
23:
Output: Px = {(x;, ((x;)))
> Set of refined pairs.
B. Proofs
Datasets X and Y are represented discretely supported probability measures = >i a,0x, and v ンドー bjdy, for probability vectors a, b E An, which we assume to be uniform: a = b = Uh, (1/n)],, E An. We form the cost matrix
C' defined by
Cjj == c(x,,
(S1)
In all cases below, we are concerned with the assignment problem (1) for this cost matrix.
Let (erm(n) = (P E mxn, P1,, = PT = (1/n) 1..) denote the set (scaled) n x n permutation matrices. By the Birkhoff-voo von Neumann theorem (Birkhoff, 1946), an optimal solution to the n x n assignment problem attained at permutation matrix in perm(n).
Definition B.1. Say that cost matrix E jgnxn Monge rotated if the identity matrix is solution to the assignment problem associated C, i.e.
IE aig (C,P). Peperm(a)
For arbitrary cost matrix C Raxx let pl arg min Pεperm(n) P) and note that the column-permuted cost matrix Ct C Monge rotated by construction. This is consequence the following identity, which holds for any
14

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
permutation P perm(n).
(C,P)s==t(CCP)) = tr((P PPTP) =tr(PC PPT) - (CP )..
(S2)
Let [u,,, u,) = II,, uy denote the transport polytope between two uniform measures. For Q II Uy), say that row of Q is soft if at least two of its entries are positive, and call the row hard otherwise. For rank * such that divides n, let (u,,, u,) be the subset of II( U,, u,)) consisting of transport plans Q with only hard rows. Below, we consider two low-rank OT problems associated to Ct. The first low-rank problem considered is
min (CI,QR T )
(S3)
Q.RelIa! (Rell. Uup)
while the second low-rank problem considered restricted symmetric couplings:
min (v....))
(C², (Q")r.
(S4)
In either case, have omitted the constant factor of r coming from diag (1 /u,,). We next introduce technical condition on C. Let C E Raxx be cost matrix and let P1 E arg millpeperm(n) (C, P) corresponding to permutation af [v] +
[n], at E Gn. Given partitions I (I...... I and J (JI,..., Jr} of in] and a, [r). define the cost between two sets la. Jn to be
Cl =
Σ
Ciat(i)
(S5)
iELaJEJ.
We call partition I balanced if each block La I has the same number of elements, 44 = (n/r).
Definition B.2. Cost matrix C E Raxx is r-Monge separable if there exists a balanced partition I' = (Tg such that for any two permutations TI, 72 E Gn, one has
Con S Σ Ca1Up)saa
(S6)
We say that C is strictly M-Monge separable if (S6) holds with strict inequality (<) for any (T:) * m2())).
One interesting feature of this definition is that while the sum is r S terms, where it may occur that r n, this inequality must hold over all permutations тi and acting on the individual data points, rather than partition blocks. This captures the notion of finding low-rank or low-resolution solutions which are nevertheless compatible with the optimal bijective Monge map.
Remark B.3. If C is r-Monge separable, the distinguished partition I may be represented Q* E IT, (u,, u.) such that Q* is optimal for (S4) and the pair (Q", Q*) is optimal for (S3). After proving the next lemma, will relate MoMoge separability to cyclic monotonicity.
Proposition B.4. Let C E 1R be strictly r-Monge separable. 1fQ"R RR arg ReIl (C, QR then, for all i [n)..
arg max Q = arg max R
(S7)
(i)ar
2E(r]
2E(r]
where at [m] + [n] is the permutation corresponding to pt E arg min minpeperm(n) (C.P)
Proof. Let at pt be as in the statement of the lemma, and define Ct CPt.T The same reasoning as in (S2) implies
that if (Q*, RT) E arg milg..i.ea (u.,, (C, QR ) then
(Q' P'R E arg min
(ct, (R.RR..
(S8)
Q (u,
15

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
The membership (S8) follows from the identities
(C*, ) ) =
(CPtT PP""T tr( Tp..)) = =trCTQ'R* ((C,T===""C""""))
Remark B.3 above follows from the requirement that the variables Q, R. have all hard rows, and are subject to uniform marginal constraints, so that all non-zero entries of QR have the same value. Thus, if C is r-Monge separable, there exists Q E II U2) corresponding to distinguished balanced partition I from Definition B.2 such that
QQ Q) E
arg min Q.ReIl (u.,, u2))
(C², (T,QR.).
(S9)
Moreover, this pair (Q, Q) is the unique optimum when C is strictly r-Monge separable. From (S8), (S9), must have
Q = ==Q,, Q-P'R
from which (S7) follows immediately.
Let us now discuss how the notion of r-Monge separability is related to c-cyelic monotonicity, Recall that for a cost matrix C E 1=x= derived from ground cost the support of an optimal plan is e-cyclically monotone if for all permutations π [n] + [n], π E Ga, one has
C# S
Σ Cir(i)
(S10)
1=1
As it amounts reindexing of the sum the right side of (S10), can equivalently define the support of the optimal plan to be c-cyclically monotone if for any pair of permutations मI 12 Gn.
Σ (r)())a2()))
from which we see that e-cyclical monotonicity equivalent to r-Monge separability with n.
We next show that the optimal factors Q", RR for the rank-2 Wasserstein problem given in (5) correspond to hard- sard-partitions of each dataset, so that for this problem the optimal Q R E I[(u,, U2) satisfy Q*, R* II.(un, u,). Below, let supp, ) C [n] be the indices on which column i of Q* is supported, where 1,2.
Lemma B.5. Let R*) be optimal for the rank-2 Wasserstein problem (5) subject the additional constraint that b Un and g u2 uniform and symmetrically, is (supp, (R*), supp2((*) so (Q", [l)()))))))
is even. Then, (supp, (Q*), suppp((()))) is partition of [ni), and
Proof. The cost is linear in (Q, R) respectively: the minimization in each variable given the other fixed can be expressed as
arg min 2 (Q, CR) F. QeIl(u,.us)
arg min (RR, C' Q)
(S11)
REIl(u..ws)
Thus for any optimal Q' R fixed the minimization in the other variable linear optimal transport problem, where by Corollary 2.11 in (De Loera & Kim, 2013) it holds that since the constraint matrix is totally unimodular with marginals integral (on rescaling). the optima R. and Q' must be vertices on the transport polytope Tu,, with integral entries (on rescaling, by 2n or 2m). There are < n + 1 positive entries in any optimal rank r = 2 solution (De Loera & Kim, 2013; Peyré & Cuturi, 2019), so that n (resp. m) being even and the rescaled rows and columns summing to 2 and n implies that there are exactly n positive entries and thus that the vertices define partitions of [n] and [m]] Thus, solutions to S11 satisfy
(Q' R*) )) II (1)
16

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Notably, in the case of an odd number of points n or m this likewise implies that one has single row which has 2 entries (1/2m 1/2n), with all other rows of the form (0 1/mz) (l/n 0) defining partition of the remaining even subset of size (n 1) or (m 1). In the general case of ranks r A 2 there are maximally + r + 1 (Peyré & Cuturi, 2019) non-zero edges (so that the graph is acyclic),, and for n > r the optimal solution remains close to a partition given mild assumptions on C.
Lemma B.5 states optimal low-rank couplings (Q", R** for Problem 7 over [(,)..u)) are in Il ua) Thus, by Proposi- tion B.4 these solutions co-cluster points x X with their image under Monge map T (xx, supposing the cost is strictly 2-Monge separable (Definition B.2). This co-clustering is in the sense of the clustering functions q", ,r" from Proposition 3.1 corresponding to each factor Q*, RR. We note that when μ and v are discretely supported measures with supports of equal cardinality, a Monge map. T X + Y, is guaranteed to exist by Theorem 2.7 of (Thorpe, 2018).
On the Rank Schedule. At each intermediate scale t E [ks] the rank-schedule (ri,..., Tx) determines the effective rank of the coupling computed so far. For each E [ro], define the effective rank at scale t as Pr := IT.=1 Y... This effective rank corresponds to the number of partitions, which are placed in bijective correspondence
Y(O q
E [pc]
(S12)
at the 1-th step of HiRef. The size of the partitions at scale l is given by n/p: = X |Y]] Given these preliminaries. we show that for an appropriate rank-schedule Hierarchical Refinement yields optimal transport maps.
Proposition B.6 (Optimality of Hierarchical Refinement). Suppose the Monge-man exists between two datasets X, Y of size n. Then there exists a rank-schedule (ri,".. .Tx) which factorizes n such that all size n// Pt partitions generated by Hierarchical Refinement at level t satisfy strict -Mngee separability (Definition  t E [0 x - 1]. For any such rank-schedule given an optimal black-box low-rank solver over (t-),, Hierarchical Refinement returns the Monge-map.
Proof. For existence, observe that taking 71 = n implies the statement Σ Cig.t S Ek=r (n( For partitions 4 of size one, this is equivalent to the statement of c-cyclical monotonicity Σ C# S >E Cir(i), so that for the trivial rank-schedule (ri)= == (n) the cost is always n-Monge separable.
Given the existence of such a schedule (ri, ,Tx) with rr+1-Monge separability, we proceed by induction on & E [0, k.). For the base case of t. 0, as we assume the Monge map exists, for the initial partition To {(X, Y)} one has that Y= T()). We want to show the variant that I'2 contains sets which are co-clusters of sets with their image under T. As the inductive hypothesis, at scale > 0 with Pt co-clusters r = [(x!0) ) each satisfies Y!!)) = T(X("). As strict T+41-Mnnge separability holds for each size n/p, bipartition (x()) (1) ((0 E Пe, using Proposition B.4 each such set is divided into re+1 co-clusters
Thus, taking the union of these T1+1 bi-partitions across the Pt elements of T2 we form a set I r++ of size Pi+1 = Ti+1/t which maintains the invariant that
{(x((+)) } ve(++)))) )) which satisfy Y ((++) = -(X
1t+ v((+t)) E T++11 = v((+1)) T(X (1+1), 1) At the final level Monge separability holds, so one may conclude on singleton sets of the form I', = T(
Remark B.7. Striet Monge separability applies unconditionally at the terminal level. Observe that all sets in I rx-1 have size equal to the rank (n/Px-1) 1) = Tw. and that we have maintained the invariant that Yf(=1) = T(x
). Let JA C [n] denote the size Tx set of indices for X ((x-1) in X. By c-cyclical monotonicity, one has for all permutations (erm(n)
m
12 Σ Cjm()) = Σ Cer(i)
Σ
C. = CC.+ + 1=1
Σ JE(M)
C₁f S Σ Cir(e) + KEJ
TEJx
JE(N][
1=1
Thus, for the subset of permutations on where म : |[n] = id, we have EieJ. Oi S. EiEJ. Cirte) implying that one may solve a constant time (r?)) solution to the assignment problem on each size rx bipartition to recover the final map.
We call the effective rank because (to avoid quadratic space complexity) we never instantiate the transport coupling corresponding to the bijective mapping (S12) as a matrix T(O, Were we to instantiate T(t) it would have rank Pt. and moreover we can evaluate its transport cost by using T(O) to induce a transport coupling p(() between the full datasets X, Y.
pe/n2 0
if ((/pp)) < i,j s (q 1)(n/p))
(S13)
otherwise
17

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
where y [pe]. and where the mass Pr/n² //2 is simplified form of (p//n) We note that this is rewriting of th Eq= DETT, to have the indices ordered into contiguous block-structure. Using coupling (S13), which again never instantiate, one can define:
cost cost(T!0) (C,P
The next proposition shows that the costs (C, P(t)) decrease as t increases from 1 to K, and also provides a bound on their consecutive differences. Below, recall that each I denotes the co-clustering (x(9, Y(O), where
( (X
Jq-1
Y = =Vv (V())
1g=1:
and where co-cluster Fi.g is defined as:
Ti = ((x,) x X y yE
Proposition B.8 (Proposition 3.4). Let cost function e: Rdd x Rdd - TR+, be of the form cxx, y) h((x (xx--))oo -Y) some stricily convex function h Rd R+ and suppose that h is Lipschitz. Let p(t) be as defined above in (S13). Then one has the following bound on the difference in cost between iterations of refinement:
5 (C, (C, < //Vcl
diam (Tt.))
(S14)
Pt
where
diam (Tz.g) diam (x!1 UT(((1))
(xx, T(x;))
(xt. T(xi)
max X,
Xx. Xx ex(")
Proof. By definition (S13) of P(t..
Pr cxx y,
p++1
(C,P()) (C, p((+)))
Pt
Σ
PI+1 n²
2iy Mj
ò ix JeT
n²
1=1
q=1
06+1
Pt
cxix Vi)
6(x,y))ETT, - T1+1
(xi Vi)
6 y ) E
n²
i=l j=l
PL1
Pt
c(xi, c(x;) 1,)0 (0 T1+1 Σ
0(xx (Er
n²
By Proposition B.4, one then has:
Pt+1 n
Σ
Σ Σ c(x;, T(x)))
Σ
Σ (EXi++1) 4PL 4p
Σ
c(x,, T(x;))
(S15)
T+11
TEX jEX
JEX(t+D
average "Monge distortion" in F, over next scale
"Monge distortion" at t++
Note that the inner summands of (S15) (indexed by g) are non-negative by definition of the refinement step, where within each cluster, one has a minimization over larger set of couplings. This shows (C, (C, p(++1)) > 0. Towards an upper bound, we will bound each summand of (S15):
1 L Σ c(xi, T(x;)) Tt-1 TEX!!) jex!
- Σ Σ Σ
Ti+1
c(xi, T(x)))
(S16)
(EX(+1) jex(e-))
18

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Define 8141 11+1 well barycenters
x (4)
X;
y(0)
Σ
T(x:)
X; x 901+2
S1+1
xEX(++1)
81+1
and note that by Jensen's 's inequality, for convex cost c(,, :) one has:
T++1
Ye+1
Σ
E
E
(x,
= sit1
Σ
Σ
cxx, T(x)))
X, ex(++1)
xxex(tt1)
St+1
jEX((+1)
8141
2 12 ,,((1),
(x(1)
so that we may continue upper- boundin the difference (S16):
1
Σ
Σ
c(x, T(x)))
( ((0)
(S17)
rt+1
x:EXi
x, Ex(
1
Σ
n² c(x;, T(x)))
(S18)
r++1
Pt
xxEXX xyEX,
Σ XIEX
Σ ((xi;T(xx)) -c(x()) -(x())))(0)) .0(11))
(S19)
T++1
x)EX!")
Now, define the diameter of co-cluster Tr.g follows:
diam (Tz.) diam (x(*) UT(X!)))
(xx, T(x))) (xx, T(xi))
max X Me. xXx
Using our Lipschitz assumption on ht made at the beginning the section, where y) 1((x - y) (we will write || for //V)))>)), one has the inequality:
lo(x;, T(x:)) c(xx, T(x;))I S Vell (T
Thus, returning to the bound each summand (S16), obtain the upper bound:
Σ Σ
|Vd|o (xx, T(x,)) (xe(), 5(()
(S20)
Tt41
xxEX(()
x,EX("
As partition ((++1) is refinement of (() and y(++1) is a refinement of Y(O), it holds that (S16) is upper bounded by:
1
Σ
Σ
dial
(S21)
r1+1
(EX(
jEX!")
1
X
| V Vill
(S22)
T241
1
n²
Vell diam (Tr.g)
(S23)
p
r++1
To conclude, we plug these bounds into each summand of (S15), obtaining the following bound on the full sum:
P141
1
n² |
Velloo Pe
diam (Tt.))
(S24)
n²
re+1
P
Pe
= =//Velloc
diam (Tr.))-
(S25)
completing the proof.
19

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Remark B.9. Proposition B.8 should be considered a conditional result. Our proof follows that of (Proposition 1, (Gerber & Maggioni, 2017)), but they are able to provide sharper bounds between elements of cluster and the centroid of the cluster using the properties assumed to hold in their definition of multiscale family of partitions (Definition C.3), which mimick the structure of dyadic cubes in Euclidean space. As we do not make any geometric assumptions of our partitions, the above result a priori weaker, through we leave the exploration of the geometry of partitions induced by low-rank OT to future work.
Remark B.10. Note, if c(x;; T(x;)) is constant (i.e., if all points are equidistant in block), one has that refinement offers no gain from level Pt + r+11:
Pt
1
Pt Pt+1 n²
(n/p:)² - (141(1// 0.
P++1 n
12
7
Te+1
T1+1
9=1
9=1
Remark 1. The work (Seguy al., 2018) suggests a loss dependent on an (entropic) coupling ?. y sparse and supported on the graph of the Monge map that y (id x T) (4, this loss becomes regression of a neural network Te on the Monge map T over the support of i:: ming, B. (x,)),), 7 (x:)).
Proof. By linearity of the push-forward map one immediately obtains
JaN [[T((>) vil T): Σ P.o drdy
XxXY S. (Te(2)) - v112 Σ lsdid T) dady
Mii
(Ta(x)-
(a
Σ - -
T
JXXY
By integrating against the 6. As p4 > 0, it holds that this loss is identically zero if and only if To = T on the dataset (I;)
min To
(T)
v1gdd v) = [To(x;) - T(::)) = (x())) T(x:)
XXXY
In other words, when one minimizes the objective of (Seguy et al., 2018) using the bijective Monge map = (id x T'aa as opposed to an entropic coupling, the objective of (Seguy et al., 2018) reduces to an unbiased regression. That is, the neural map Ta directly matches 7 over the dataset support as if trained on supervised (æ, y) pairs y = T(x).
C. Background: Multiscale Optimal Transport
C.1. Multiscale Partitions
(Gerber& & Maggioni, 2017) describe a general multiscale strategy for computing OT couplings between metric measure spaces (X, dx, μ)) and (Y, dy, v). They state this in the Kantorovich setting, using a general cost function e ; X x Y + Rt. Their framework consists of several elements:
1. A way of coarsening the set of source points X and the measure across multiple scales:
(X. =: (z) + (X--1,P5-1)
(X1,
(S26)
as well as an analogous coarsening for the set of target points Y:
(Y, v) =: (Y), (Y3-1, +... (Y,,
(S27)
where at each scale j. supp(uş) = X, and supp(u,) = Y,, and the cardinality of each X₁ and Y, decreases with j.
2. A way of propagating coupling πy solving the transport problem 14, - Vi scale coupling Tz+1 at scale j 1.
3. A way of refining the coupling from scale j to an optimal solution at scale j
20

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
To derive approximation bounds for the error incurred by the multiscale transport problem at each scale, (Gerber & Maggioni, 2017) use regular families of multiscale partitions (Definition C.3 below) to define approximations to μ., v and at all scales.
For E X, define (r(r) := {x' E X x²) r) as the metric ball of radius centered at z. Functions f.g X + R have the same order of magnitude if there is c1, C2 > 0 with f (z) < g(z) I cef for all æ E X. and in this case we write f x g- Write (XX for the space of unsigned measures on X, and write P(X) for the subspace of probability measures,
Definition C.I. A metric space (X, dx) has doubling dimension if every (r(r) admits a covering by at most 24 balls of radius r/2.
metric space is said to be doubling if it has doubling dimension d for some d > A related notion doubling metric space is a doubling measure.
Definition C.2. Measure μ E M(X) is a doubling measure with dimension d if there is constant C₁ >> 0 such that for all E X and all > 0, one has G rn & ((B_(r)) S qra i.e. (r) x rdd
Note that if (x, dx, 14) is doubling, then dx is doubling, and up to modification of dx to an equivalent metric, the dimension can be taken as the same in either case.
Definition C.3. Given metric measure space (X, dx, p), regular family of multiscale partitions with scaling parameter A > 1 is family of sets
K3 1' {(0gA)
j=1
with each Cj,s C X such that:
1. For each scale j, the sets k2 partition X.
- 0 or CJ+1K C Cj.k- In this latter case, we say that (j (5+1,Kg is 2. For each scale j E [J - 1], either CJ+1K n ,, child of k), or equivalently that (j, k) is a parent of (j + 1,K), writing 1,k) x kj,
3. There is a constant A > 0 such that for all j, ki we have diameter diam (Cj.x) < A0-
4. Each Cj. contains a "center point" Cj.k such that B (0 C CCA.A
We take f = 2 for simplicity. As the child- tree structure, like dyadic cubes in Rdd Though the measure is not explicitly used in the above definition, the third and fourth points imply p(C),s) x 2-jd jd and K, x 2jd
terminology suggests, these partitions (through the second point) have
Coarsening spaces and measures Now suppose that each of X and Y are each discrete metric measure spaces, each equipped with regular families r(()), IYY) of multiscale partitions:
π(X) = (7 r₂(X: = (X)
r(Y) = (T)M)) I,(Y) == (O) K,
and these yield the coarsening chains in (S26), (S27) in the most natural way possible each scale j. defining the coarse-grained spaces X5, Yj to be the clusters at scale j:
X, = rx)), Y, TMM),
while the measures at scale j are defined from the measures at scale j + via:
p,(C,a(X)) Σ (X),,, vy(C,,(Y))
Σ (C+L
(41
(+11))
The fourth item of Definition C.3 requires that we define cluster centers (X) for each (X), At the finest scale j = J, all clusters (X)(()) correspond to singletons {æj,}. }. so we define (X) := IJ in this case. At coarser scales, these centers can be defined recursively from the next finest scale, depending on the structure of X.
21

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
For example, if X has vector space structure (in addition being metric measure space), natural choice for cluster centers Jij.k scale o,....,-- 1 is the weighted average jj,k Cy where
[x = Σ
(X)
(+IR)+U0A)
On the other hand, in the absence of vector space structure, one can still define
(X) arg min acX
Σ
(x,
(+IN)~GGA)
with analogous constructions for Y yielding centers yj,k-
Coarsening the cost function (Gerber & Maggioni, 2017) suggest three ways to coarsen the cost function using the multiscale partition. To condense the notation slightly, let us write j in place of (x, and V3, in place of C),(() and (Y).
(c-i) The pointwise value
C 4z = c(Tj,
(S28)
using centers 2 and Vj defined in any of the ways above.
(c-ii) The local average
9 Fy,ss Vj )= = E0EC,A(0))EC,,,0)) c(,, (x) v)
C33x
(c-iii) The local weighted average:
Cj Vitee ) EnECsMO) c(z, ( Mj 1)
EuEC ( Ui 1,81)
where - is the optimal (or approximately optimal) OT coupling scale j 1. defined below. The indices by and Ky are defined using the tree structure of the partition: k the unique index among [K]-1XX]] such that kj,k) < (j 11), and likewise kj is unique among [X)-
C.2. Propagation of OT solutions across scales
For each scale j. consider the OT problem given as follows.
Ti. == arg min cost(T;),, where:
*Ell(p, vs))
(S29)
cost(TT))
E
c(E My, MiNN (1)
(Gerber & Maggioni, 2017) show bounds on (oost(z;) (m))) | of a constant times 2 Vc/loo but note that this only implies closeness of the couplings in terms of their cost, not necessarily in any other sense.
Given an optimal coupling nī at scale j. (Glimm & Henscheid, 2013) proposed direct propagation strategy to initialize the problem scale + 1, distributing the mass лt (Ty,ks V),rr) equally to all combinations of paths between children (x),.) and children(yj.e.) In this context, path is understood mean source-targe tauree-target pair at the next scale, e.g. a pair of the form Wi+,1, To formalize this, let
A₁ = {I UyAA.) E [K2()], V' e (K)()))}
denote all paths between points in X, and Y,. The drawback of this warm-start procedure is that if supp(u;) C. Aj, which always the case, the refinement procedure still requires quadratic space complexity at the finest scale.
22

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
To mitigate the ultimate quadratic space complexity of retaining all possible paths at all scales, (Gerber & Maggioni, 2017) allow for refinement procedure where the support of couplings scale j + is restricted to subset Rj+11 C As+1 of all possible paths (with Rj+1 defined by the optimal coupling at the previous iteration). Given R3 C Aj, let Tř R2 denote the optimal solution to the path-restricter or restricted problem at scale j:
TJ|R, == arg cost(TT).
(S30)
TEl(()))))) Cup
Simple propagation. The simplest way to restrict the number of paths considered at subsequent scales is use paths at scale whose endpoints are children of mass-bearing paths scale 1:
R,+1= = (1)+1) 3(mr) Vs.Ar) supp(af) s.t (j + 1, < (J, k) and (j + 1, - (j,K)).
The optimal Kantorovich plan at scale j has at most (K2 (X) + (Y + 1) non-zero entries. Using the above simple
propagation strategy constrains plan at scale j + to be supported on at most
0 + (Y))
entries, where 0j is the maximum number of children of any (j,k) k) across both datasets. When the ambient space has doubling dimension d, for any j one has 02 x 24, yielding a plan with linear space complexity at the finest scale.
Capacity constraint propagation. This propagation strategy solves a modified minimum flow problem at scale j in order to include additional paths at scale j + likely to be included in the optimal solution - Concretely, one first computes an unconstrained optimal plan T] (R, at scale j. Then, a new OT plan #t]R,, is solved for at scale j now subject to the capacity constraint
Fjl,, S U
for each (Izjk: Vi,k)) E supp(r}lR,)) where the random variables Uk, are i.i.d. Uniform([O.. 0.9)). This can also be iterated several times, in all cases leading to linear space complexity in the optimization at the finest scale.
C.3. Refinement of the propagated solution
Propagation of solution to the restricted transport problem (S30) at scale j. in general cannot guarantee reaching optimal solution to the restricted problem at scale + 1, and can lead to accumulation of errors across all scales. Several refinement strategies are proposed in (Gerber & Maggioni, 2017) to address this.
Potential Refinement. One refinement strategy leverages the problem dual to (3), here stated at the finest scale:
m
max FER" gER" f SC
(((x;)) + Σ ((,i))g)
(S31)
The refinement strategy uses optimal dual variables f",g" to select paths to include at the next scale. From the dual formulation, an optimal solution (f" (f') to (S31) must have all nonnegative entries in the reduced cost matrix, defined as the matrix C - f + g with entries CAu f4 - gr Note that the dual to the restricted problem (S30) is well-defined, and for this dual we denote the optimal dual potentials by [RR, and gR|,,- With slight abuse of notation, let (f* () 1RR,, be
(p" + g)) = (f* @ O
where Mú) is the indicator matrix of the restricted set of paths R, at scale j, and where O denotes the Hadamard (entrywise) product. While the restricted set of paths R, is inherited from the previous scale, one can define a new set of paths 1 based on where the restricted reduced cost C- (f* + "g")) gR) is nonpositive:
(T)1,,, = (TIAA Vj E Ay CRR [(f* (gg")R,,kkr 0}.
23

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Table $1. Hyperparameters for Synthetic Experiments
Parameter Name
Variable
Value
Rank-Annealing Schedule
(T)...., Ti)
[2, 512]
Hierarchy Depth
K
Maximal Base Rank
2
Maximal Intermediate Rank
With a new of paths 2 : (m)) R, one can compute new optimal plan ojløy at scale j restricted these paths, as well as new optimal dual potentials f*1 luy and g'lvy leading to a new reduced cost C (f @ 22g") 1vg. This strategy can be iterated via
2j : -
(S32)
yielding the sequence of transport plans oo all at scale j. which converge solution whose reduced cost is nonnegative, necessarily making optimal. The potential refinement strategy was used by (Glimm & Henscheid, 2013), with (Schmitzer, 2016) introducing shielding neighborhoods in similar spirit, using dual potentials to locally verify global optimality.
D. Experimental Details
D.1. Synthetic Experiments
For all of the synthetic experiments, we first generate n = 1024 points from three datasets: the checkerboard dataset (Makkuua et al., 2020)), the MAFMoons and Rings dataset ((Buzun et al., 2024)), and the Half-moon and S-curve dataset ((uuzun al., 2024)). Following (Buzun et al., 2024) the random seed was set to for data- setta-genrration with jax. random. key (0). We evaluate the OT cost (C, P) of Rif Sinkhorn (Cuturi, 2013), and ProgOT (Kassraie al., 2024) each of these three datasets, where 11-112 (Table S2). We additionally quantify the number of non-zero entries in the plan and its entropy (Table S3).
(1) the Euclidean cost 1|- and (2) the squared Euclidean cost
We also compare the cost of couplings computed by Hierarchical Refinement to low-rank couplings (Scetbon et al., 2021; Halmos et al., 2024) of varying rank. We observe that as the latent rank r + n, the OT cost (C, Pr). asymptotically approaches the cost achieved by Hierarchical Refinement (Figure S3). In the limit lim,-a (C, Pr) low-rank OT recovers Sinkhorn (Scetbon Cuturi, 2022) and approaches quadratic memory complexity, while Hi Re remains linear in space.
Checkerboard
The checkerboard dataset (Makkuva et al., 2020) is defined by random variables Y ~ Q sampled from the source distribution according to Y = X + Z where X and Z are sampled from Uniform distributions defined by
X Uniform (((0,)) 0), (1, 1), (1,-1), (-1,1), (-1,-1))), Z ~ Uniform (l-0.5, 0.5] x [-0.5, 0.5)).
the target distribution has random variable Y' where the random variable Y' is defined as Y' X' + Z with components
X' Uniform (f(0. 1), (0, -0), (1,0), 0))),
Z Uniform (l-0.5, 0.5) x [-0.5,0.5]
MAFMoons and Rings
The MAFMoon dataset (Buzun et al., 2024) defines source distribution Q by sampling X N(, 12) and applying the non-linear transformation defined by
0.5(X, + x2)-5) -5]
Y
X₂
this introduces a quadratic dependency on the Gaussian randomly variable to generate crescent shape.
24

---

Hicrarchical Refinement: Optimal Transport to Infinity and Beyond
The target distribution P representing concentric rings is generated by first sampling 0 Uniform(27). with fixed radii (0.25, 0.55, 0.9, 1.2) from which one transforms to Cartesian coordinates I; 3r; cos 0; and yi 3r; sin 0,. Gaussian noise is added to each of these, as NO Lo²) for o 0.08.
Half-moon and S-Curve
make.moons The Half-moon and S-curve dataset (Buzun et al., 2024) is generated from Y make moons and make. S.cuve from the scikit -iikkeatrrrn library. Both datasets are transformed further with a rotation R(0), a scaling X, and translation applied as Y (AY) + p.
Table S2. Comparison Table for Coupling-Basee OT Methods on Primal Cost (C. P)r for ||-1|2 and 11-1
Method
Checkerboard (Makkuva 2020)
MAFMoons Rings (Buzun 2024)
Half Moon & S-Curve (Buzun 2024)
1|-1|2
11:12
1|-1/2
||
||-1|a
11:11
Sinkhomn
0.3573
0.1319
0.4422
0.4440
0.5663
0.5663
ProgOT
N/A
0.1320
N/A
0.4443
N/A
0.S709
HiRef
0.3533
0.1248
0.4398
0.4414
0.5741
0.5737
Table S3. Entropy and Non-Zero Entries (> 10 ) of Coupling Matrices for Each Method and Dataset (Wasserstein-2 distance cost,
11:113)
Method
Checkerboard (Makkuva 2020)
MAFMoons & Rings (Buzun 2024)
Half Moon & S-Curve (Buzun 2024)
Entropy
Non-Zeros
Entropy
Non-Zeros
Entropy
Non-Zeros
Sinkhorn
12.8509 12,3830
624733 271087
12.6117
678720 327764
12.7776 12.1170
652993 337258
ProgOT
11.6158
Hi Ref
6.9314
1024
6.9314
1024
6.9314
1024
Table S4. Comparison of Coupling-Banl OT Methods on Primal Cost (C, P) (Wasserstein-2)) on 512 point small instance.
Method
Checkerboard
MAF Moons & Rings
Half Moon & S-Curve
MOP (Gerber Maggioni, 2017)
0.393
0.276
0.401
Sinkhorn (ott-jax)
0.136
0.221
0.338
ProgOT
0.136
0,216
0.334
HiRe
0.129
0.216
0.334
Dual Revised Simplex Solver
0.127
0.214
0.332
25

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Table S5. Hyperparameters for Mouse-Embryo Spatial Transcriptomics Experiment (F15.5-16.5)
Parameter Name
Variable
Value
Rank-Annealing Schedule
(ri,
(ri....,Te)
[2, 86, 659]
Hierarchy Depth
R
3
Maximal Base Rank
Q
2n
Maximal Intermediate Rank
128
D.2. Large-scale Transcriptomies Matching on Mouse-Embryo
In this problem, we use fiRef to find a full-rank alignment matrix between successive pairs of spatial transeriptomics (ST) (Stâhl et al., 2016) slices. These are from a dataset of whole mouse embryogenesis (Chen et al., 2022) on the Stereo-Seq platform. These datasets have been measured at successive 1 1-a time-intervale F9.5 (n = 5913), E10.5 (n = 18408), E11.5 (n- 30124), E12.5 (n = 51365), E13.5 (n 77369), E14.5 (n 102519), B15.5 (n 113350), and E16.5 (n 121767), where the embryonic mouse is growing across the stages that the complexxy n increases with the numeric stage. For each pair of datasets of size n and m, we sub-sample the datasets so that the size of the two datasets is given as n < min(n,m)
In the context of spatial transcriptomics, an experiment conducted on a two-dimensional tissue slice produces a data pair Z). Here, X E R=*P represents the gene expression matrix, where n denotes the number of cells (or spatial spots) analyzed on the slice, and p signifies the number of genes measured. Specifically, the entry X., E R.+ corresponds the expression level of gene j in cell i, with higher values indicating greater expression intensity. Concurrently, Z IR is the spatial coordinate matrix, where each row contains the (x, y) coordinates of cell on the tissue slice. Consequently, every cell on the slice is characterized by a gene expression vector of length p. capturing its molecular features, and coordinate vector of length two, detailing its spatial position within the slice.
We utilize the extensive, real-world dataset on mouse embryo development presented in (Chen et al., 2022), which encompasses eight temporal snapshots of spatial transcriptomics (ST) slices throughout the entire mouse embryo development process. And align all consecutive timepoints. The preprocessing of this dataset is conducted using the standard SCANPY (Wolf al., 2018) workflow. Initially, we ensure that both slices contain an identical set of genes by filtering, which results in a common gene set across all cells for each pair of timepoints. Subsequently, we apply nogmaaamaiization to the gene expression data of all cells from the two slices. To compress the data, we perform Principal Component Analysis (PCA). reducing the dimensionality of the gene expression profiles to d = 60 PCs. Finally, we compute the Euclidean distances between gene expression vectors in the -redffranfor space to construct the cost matrix C, on which we solve a Wasserstein problem to obtain the optimal coupling P of full-rank. We offer hyperparameters for the 115-16.5 experiment (the largest alignment) in Table S5. For the other experiments, the maximal intermediate rank r = 16 up to E10.5, r 32 to E11.5, 64 up to E13.5, and 128 for E14.5-16.5. The anneaann schedule is generated according to the dynamic program in each case by the rank lank-annealll ing. - optimam www. n, hierarchy.depthh max_rank function.
max-Q
In this experiment, we benchmark against the default implementation of Sinkhorn in ott-jax (Cuturi et al., 2022) with entropy parameter c = 0.05, and additionally benchmark against the default implementations of ProgOT (Kassraie et al., 2024) and LCT (Scethon et al., 2021) in ott-jax. For the low-rank methods LOT and FRIC (Halmos et al., 2024) we fix constant rank of 40 for these experiments, While LOT (Scetbon et al., 2021) provides a robust, scalable low-rank procedure for the Wasserstein-2 distance, the LOT solver with point cloud input on Wasserstein-1 cost only runs for the first pair (E9.5.., For subsequent pairs we input the cost C directly, resulting in the LOT solver running up to the third pair (E1 1 15:1112...). Mini-batch OT is run with batch-sizes ranging from 128 to 2048, and is performed without replacement. As noted in prior works (Fatras et al., 2020; 202 la;)), this is a standard choice for instantiating a full-rank coupling using mini-batch OT. Sinkhorn is used to solve each mini-batch coupling, as implemented in ott -jax ja with a default setting of the entropy parameter e = 0.05.
D.3. Brain Atlas Spatial Alignment
We took inspiration from MERFISH-MERFIEE alignment experiments of (Clifton et al., 2023), particularly gene abundance transfer tasks that STalign is exhibited on. The data are available on the Vizgen website for MERFISH Mouse Brain Receptor Map data release (https://info. wwwgeei.nonnneeeeeeeeee wwwpl/ouuuu-bbb--maaaaa The two spatial transcriptomies slices used for the experiment
26

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Table S6. Cost Values (P. C) for Different Methods Across Embryonic Stages
Method
E9.5-E10.5
E10.5-E11.5
E11.5-E12.5
E12.5-E13.5
E13.5-E14.5
E14.5-E15.5
E15.5-E16.5
HiRef
21.81
14.81
16.14
14.35
13.78
14.29
12.79
Sinkhorn
21.91
14.89
ProgOT
22.56
15.35
MB 128
22.44
15.35
16.69
14.86
14.14
14.75
13.32
MB 512
22.15
15.05
16.33
14.54
13.92
14.50
13.01
MB 1024
22.05
15.02
16.24
14.45
13.86
14.43
12.91
MB 2048
21.98
14.98
16.18
14.39
13.81
14.39
12.85
FRLC
23.14
16.09
17.74
15.47
14.64
15.51
14.00
LOT
26.06
19.06
21.64
are slice 2, replicate 3 ("source" dataset) and slice 2, replicate 2 ("target" dataset). The datasets will be denoted (X), St) for the source and (x?, for the target.
The source dataset consists of 85, 958 spots, while the target dataset consists of 84, 172 spots. To apply Re these data, we subsampled the source dataset to have 84, 172 spots also (uniformly at random), removing total of 1786 spots. We use this samplad n X dataset for all methods, but as discussed below, note that this sub-sampling incurs little error. We ran Hi Ref using the settings max_rank = 1 and h hierar wwwerarchy.dppphhat. for a total runtime of 10 minutes 6 seconds, on an A100 GPU. The random seed was set to 44. For the cost function used by HiRef, f, we only use the spatial modalities S S of the two datasets. We centered the spatial coordinates of both datasets, and applied rotation by 45 degrees to the first dataset. With these registered spatial data, here denoted S = {s} )2=1 and $2 {s Ji=1 we formed the cost matrix C given by:
i
C₁j = ||s
8 12
where || |l denotes the Euclidean distance between the spatial coordinates. This cost C was used as input to HiRef, f, which produced as output 1-1 mapping T between the two datasets (a permutation matrix is too large to instantiate).
We then evaluated the performance of HiRef through cosine similarity of predicted gene abundance with target gene abundance, across five ceaatiallyppateermn genes" (using the terminology of (Clifton et al., 2023)): Slc17a7, Grm4, Oligl, Gadl, Peg/0. /0. Writing g to stand in for any of these genes, we formed the abundance vectors v and v2.0 using the raw counts for gene g in each datasets' expression component X X². Using HiRef output T. also formed the predicted abundance vector vo, which maps the counts from v the spots in the second dataset through T.
Moreover, to compute cosine similarities between predicted and true expression abundances, (Clifton et al., 2023) employ a spatial binning on their output, using windows of 200 to tile each slice. The diameter of each slice is roughly 10, 000μm, and to make our output comparable, we used the spatial coordinates to bin and average the vectors 22.0 and v locally. We used a total of 5625 bins, corresponding to a 15-to-1 mapping from spots to bins. Averaging the abundance of gene g in each bin, we obtain spatially smoothed versions of v2.8 and v, in (Clifton et al., 2023). Denote these smoothed vectors by w2.9 and w. For each gene g among (Slc/7a7, Grm4, Oligi, Gadl, Peg }. we computed the cosine similarity results in Table $7. In the same table, we list scores obtained by the low-rank methods FRLC (Halmos et al., 2024) and LOT (Scetbon al., 2021) for comparison. While HiRef is restricted to running on datasets of the same size, LOT and FRLC have no such restriction, and can run on the pair of MERFISH slices without any subsampling. To address this, in each case of LOT and FRLC,, we give results from the methods run on the datasets with and without subsampling, reporting the highest scores for each method in main. In particular, we compared the cosine similarities for the original and sub-sampled dataset on downstream task, as the primal OT cost is no longer directly comparable. Without the sub-sampling, the cosine score is only slightly higher than with: (0.3390, 0,2712, 0.3186, 0. 0.1666. 0. 0.1080) vs (0.3241, 0.2279, 0.3029, O. 1653, 0.0719). These scores remain significantly lower than those of hierarchical refinement on the sub-sampled data: (0.8098, 0. 0.7959, 0.7526, 0.4932, 0.6015).
between w2.0 and w, listing
For the FRLC algorithm, we set a 0, 200, On = 500, rank r = 500, using 20 outer iterations and 300 inner iterations. The runtime of FRLC was 1 minute 26 seconds on an 1 GPU. For the LOT algorithm, we were unable to pass a low-rank factorization of the distance matrix, so we had to use a smaller rank r = 20 in order to avoid exceeding GPU memory (the choice 20 led memory usage of 30GB). We set 6 = 0.01 and otherwise used the default settings of the method. The
27

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Table S7. Cosine Similarity Scores for Expression Transfer & Spatial Transport Cost
Method
Slc17a7
Grm4
Oligl
Gadl
Pegl0
Transport Cost
HiRef (this work)
0.8098
0.7959
0.7526
0.4932
0.6015
330.3301
FRLC (Halmos et al., 2024)
0.2180
0.2124
0.1929
0.0963
0.0991
415.0683
FRLC, no subsampling
0.2373
0.1896
0.1579
0.0644
0.1550
634.4158
LOT (Scethon al., 2021)
0.3241
0.2279
0.3029
0.1653
0.0719
3722.3171
LOT, no subsampling
0.3390
0.2712
0.3186
0.1666
0.1080
3722.1360
MOP (Gerber & Maggioni, 2017)
0,5211
0,4714
0.5972
0.3571
0.2719
2479.6117
Mini-batch (128)
0.6693
0.6637
0.6442
0.4150
0.4932
653.0491
Mini-batch (512)
0.7089
0.7383
0.6771
0.4562
0,5383
438.1703
Mini-batch (1,024)
0.7256
0.7621
0.6918
0.4733
0,5557
384.2498
Mini-batchs (2,048)
0.7434
0.7822
0.7056
0.4912
0.5683
349.2964
Table S8. Cost Values (C, F) for ImageNet (Deng et al., 2009; Russakovsky et al., 2015) Alignment Task.
Method
HiRef
Sinkhorn
MB 128
MB 256
MB 512
MB 1024
FRLC
LOT
ProgOT
OT Cost
18.97
N/A
21.89
21.11
20.34
19.58
24.12
N/A
N/A
total runtime was 36 minutes 8 seconds on an A100 GPU. To form a spot-to-spot mapping from each transport plan output by FRLC and LOT, we mapped the spot with index i in the first slice to the index argmax of the i-th row of the transport plan. Note that we ran LOT using the squared Euclidean cost as default, as passing PointClour raised an error. The discrepancy in transport cost between the two low rank
Bucl idean () as
an argument to ott-jattee.. methods reported in Table S7 is explained by (i) needing to rank-20 plan of LOT versus the rank-500 plan of FRLC. We applied the exact same spatial averaging to the outputs of all methods. We plot the ground- and HiRef- MeRe--preet -reediccd abundances in Figure S1.
Euclidean cost in the case of LOT, and (ii) using a
D.4. Alignment of ImageNet Embeddings
To demonstrate the scalability of и Ref to massive and high-dimensional datasets, we perform an alignment unprece- dented for OT solvers: aligning 1.281 million images from the ImageNet ILSVRC dataset (Russakovsky et al., 2015; Deng et al., 2009). A negligible amount of sub-sampling, 167 points out of 1281167, was applied so that n divided into two integers n/2 = 640500 of which neither is prime. From this, rank-anneal ing. opt wottimmllankschheeee n, hierarchy.depta.
( max_Q max_rank was called to generate the depth rank-annealing schedule of (r), = (7,50, 1830) for HiRef. We used the ResNet50 architecture (He et al., 2016) available at https://downloud.pyooot haapecchorggmooeeeeeeeetttO
to generate embeddings of each image of dimension d 2048. We then took 50:50 split of the dataset the two image datasets X, Y to be aligned, where we used random permutation of the indices of the dataset using torch. randperm so that the splits approximately represent the same distribution over images. We then aligned these image datasets using Hi fi FRLC, and mini-batch OT. For 2,, Zj E R22048 we used the standard Euclidean cost defined by
Ciy
1z -2jl22
We use the sample-linear algorithm (Indyk et al., 2019) to factorize C into low-rank factors of dimensions (d1, d2, dz) = (ri, T2, r3) = (7, 50, 1830) paralleling the rank-schedule. The final cost values for each are shown in Table S8.
E. Additional Information
There are number of additional practical details regarding Algorithm 1 in its actual implementation. In particular, to achieve linear scaling, one must also have sample-linear approximation of the distance matrix C. We use the algorithm of (Indyk et al., 2019) to accomplish this, as discussed in Section E. In addition, one requires parallel sequence of ranks for the distance matrices used at each step, (d,. d,). As a default, we set (d,,... d.) = (ri, .rx) so that the ranks of the distance matrices parallel those of the coupling matrices. Moreover, fiRef has the capacity to be heavily parallelized:
since Algorithm 1 breaks each instance into independent partitions, one may also parallelize the low-rank sub-problems of
28

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
Table $9. Hyperparameters for ImageNet Experiment
Parameter Name
Variable
Value
Rank-Annealing Schedule
(i....,T))
[7.50, 50, 1830]
Hierarchy Depth
8
Maximal Base Rank
Q
2
Maximal Intermediate Rank C
Algorithm across compute nodes.
E.1. Optimizing the Rank-Annealing Schedule
As discussed in Section 3.3, the large constants required by low-rank OT (LROT) in practice encourage factorizations which have minimal partial sums. In particular, one seeks a factorization which minimizes the number of times LROT is run as a sub-procedure. Suppose one defines the maximal admissible rank of the low-rank solutions to be C' E Z+, the hierarchy-depth to K, the number of data-points to be n, and the maximal-rank permissible for the base-case alignment n/Q,
be Q. If Q * 1, then one may take 1 + 11 +... + II;=1 rs. where the ranks factor the sample-size as II r, Thus, optimize the number of LROT calls for given dierarchy-depte K, can optimize for the rank-anne schedule by minimizing the sum of partial products defined by mint()),,
1, to observe that the total number of runs required is
- ra subject to [[ ri n, E 2+, S C. Observing that this equals min(r)) T1 + == I? Ta implies a standard Jynamic-programmingg approach and store a table of factors up to C' to optimize this in (CCm)) time for C, k generally pre-fixed constants.
Low-rank distance matrix C. A key work (Indyk et al., 2019) showed that one may approximately factor a distance matrix C with linear complexity in the number of points n (Algorithm E. For certain costs, e.g. squared Euclidean, this factorization can be given for free (Scetbon et al., 2021). We rely on both of these for low-rank factorizations the distance matrix, so that both the space of the coupling and pairwise distance matrix scale linearly.
Algorithm 3
Low-Rank approximation for distance matrix C
Input point sets {x; )x {y in metric space A' and metric d Pick indices i* E [m] E [m]] uniformly at random
for i
do Update sample probability p: d(x;, d(xx-, y3. )2 ) 1 Σラ d(xe,,,)² end for
Sample O(r/s) rows C, ~ Calegorical L Pi
Compute U using (Frieze et al., 2004)
Compute V using (Chen Price, 2017)
return V. U
29

---

Hicrarchical Refinement: Optimal Transport to Infinity and Beyond
a.
Siice (5283) Sic17a7 abundance
Siee (S282) Sic17a7 abundance
slice 2 (5282) predictec Sic17a7
b.
Slice (52R3) Gmmd abundance
Slice 2 (52R2) GMmH aburdance
Sike 2 (S2R21 predicted Gmm
C.
Siice! ($2R3) ol abundance
Siice (52R2) Olig1 alou Sdancc
Sice 2 (52R2) predicted Oliy
d. Siice (S2R3) Gad1 abundarce
Slice (52R2) Codi aburdence
Slice (52R2) predicted Gec1
e. Slice (5203) Peg10 abundance
Siic 2 (5202) Peg10 abundance
slice 152R21 predicted Peg10
Figure S1. Abundance of 5 genes (a. Slc17a7, b. Crm4, e. Oligl, d. Gadl, e. Peg/0) in Allen Brain Atlas MERFISH dataset (Clifton al., 2023). From left to right are plotted (1) abundance in the first dataset, (2) abundance in the second dataset, and (3) predicted abundance via transfer of the abundances in the first dataset under the mapping of HiRef.
30

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
a.
b.
bustines Sarrple Compleates
Purtino una Sample Carpkotics
HR07.A
200
158
Gorgle
HiRef Runtimes over Sample Complexities (GPU)
C.
120
100
Runtime (s)
80
60
40
20
20000
40000
60000
80000
100000
120000
Samples (n)
Figure S2. Runtime scaling across sample-complexities n of a. Hicrarchical Refinement (Hi Re and b. Sinkhorn for Euclidean cost, Il (single CPU core). Hierarchical Refinement exhibits linear scaling increasing n, whereas Sinkhorn exhibits quadratic scaling is unable to run beyond 16k points. c. Runtime scaling of HiRef (GPU).
OT va Coupling
Oue Pla)
Hicnetie Behwaas Can (CC, Posis)
2133
OF Ooe (c.)
0.00
5.133
20
20
100
Tank of Cocoling Merne (tuk(P))
Figure S3. HiRef cost and the cost of the low-rank OT solution of FRLC (Halmos al., 2024) across the coupling rank E [5,1.
31

---

Hierarchical Refinement: Optimal Transport to Infinity and Beyond
a.
Hierarchical Refinement Map
Sinkhom Barycentric Map
b.
C.
Figure S4. Comparison of optimal transport maps under (1) the HiRef alignment, and (2) the Sinkhorn (Cuturi, 2013) barycentric projection. The checkerboard dataset of (Makkuva al., 2020), b. the Half-moon and S-curve dataset of (Buzun et al., 2024), and e. the MAF-Moone Rings dataset of (Buzun et al., 2024).
32

---

Hicrarchical Refinement: Optimal Transport to Infinity and Beyond
a.
b.
C.
Figure S5. Alignments of the synthetic datasets of (Makkuva et al., 2020; Buzun et al., 2024) using the optimal dual revised simplex
(Huangfu & Hall, 2018) algorithm for small instances (512 points). The checkerboard dataset of (Makkuva al., 2020). b. the Half molf- and S-curve dataset of (Buzun et al., 2024), and e. the Rings dataset (Buzun et al., 2024).
33


Reviews:
[{"review_id": "wnaHbTcMuw", "reviewer": "Reviewer_dsLw", "summary": "This work concerns the use of hierarchical refined version of low-rank optimal transport (HR-LOT). In previous works, low-rank transportation plans have been explore to reduce (high) computational costs for solving the optimal transport problem, and hierarchically refined version of this problem was also studied. This work combines the two,", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer _dsLw_ for their feedback and careful reading.\n\n> The theoretical result Proposition 3.1 seems to have strong assumptions about the transportation plan, although a more general result might be simply very challenging to formulate and prove.\n\nYes, two important generalizations of  Proposition 3.1 are: (1) extend to datasets of unequal sizes $n \\neq m$;  (2) extend to general schedules of factors (as is done in practice in the algorithm) rather than just powers of 2. We briefly discuss (1) in response to reviewer 9vCU below. Regarding (2), we note requiring that $n$ be a power of two is without loss of generality in the Proposition. We will add the following statement to Section 3 to highlight this: _\"We note that the assumption that the datasets are of size $2^k$ is without loss of generality. For a dataset of size $n$, let $q = \\min$_ { _$2^{t} \\mid t \\in \\mathbb{N}, 2^{t} > n $_ } _and add $q-n$ 'dummy' points at 'infinite' distance from $\\mathsf{X}, \\mathsf{Y}$, and mutual distance zero.\"_ \n\nBy construction, the optimal mapping for this augmented data is given as the product measure $\\gamma = \\gamma _ {1}^{ * } \\otimes \\gamma _ {2}^{ * }$, where $\\gamma _ {1}^{ * }$ is the optimal mapping between the original points and $\\gamma _ {2}^{ * }$ pairs the dummy points. Thus $\\gamma_{1}^{ * }$ agrees with the original optimal coupling when restricted to the datasets $\\mathsf{X}, \\mathsf{Y}$, and one may thus invoke Proposition 3.1 directly for any $n$. The only computational cost is, in the worst case, one would extend the size of the dataset by a factor of 2.\n\n> Hierarchical low-rankness of the transportation map seems to share common properties with so-called Hierarchical Matrices used to solve partial differential equations (most notably the Helmholtz equation). Are there some commonality in their approaches?\n\nYes! Thank you for highlighting this, as there appears to be an interesting parallel between the two. The coupling at each iteration $\\mathbf{P} ^ {(t)}$ can be viewed as having a hierarchical block structure according to each partition. Similarly to how hierarchical matrices in PDE reduce the complexity of dense linear operators from $O(n^{2})$ to $O(n \\log{n})$, hierarchical refinement uses a multiscale block structure to reduce the complexity of OT to $O(n)$ in space, and $O(n\\log{n})$ in time. For iterations $t \\in [0, \\log{n} ]$ earlier iterations capture coarser structure, while later iterations capture fine-grained structure.\n\nA noteworthy difference is that the linear operators approximated by hierarchical matrices in PDE are inherently dense, while the optimal solution to OT is guaranteed to be sparse. In OT, dense matrices are prevalent because of the computational value of entropic (or rank) regularization as a means to explore the solution space by annealing from a dense initial condition to an optimal sparse solution. \n\n> I found the discussion regarding the rank-scheduling and annealing scheduling $\\epsilon_n$ in Section 3.2 a bit hard to follow. Is there an alternative explanation for their relation?\n\nThank you for this feedback. First, to better explain the connection between entropy- and rank-regularization, we refer to Chapter 5 in the thesis of Scetbon, based on the seminal work (Scetbon et al. '21):\n\n* _\"A key observation when entropy is added to the coupling is that the more entropy is added, the lower the rank.\"_ (29);\n* _\"a useful parallel can be drawn between (LOT) and that of the vanilla Sinkhorn algorithm, in the sense that they propose different regularization schemes. Indeed, the (discrete) path of solutions obtained by (LOT) when varying $r$ between $1$ and $\\min(n,m)$ can be seen as an alternative to the entropic regularization path. Both paths contain at their extremes the original OT solution (maximal rank and minimal entropy) and the product of marginals (minimal rank and maximal entropy).\"_ (30)\n\nNow because of this correspondence of small epsilon with large rank, and large epsilon with low-rank, the analogue of annealing in the parameter $\\epsilon$ in entropy-regularized OT, i.e. gradually decreasing $\\epsilon \\to 0$ according to some schedule, is to initialize at a low-rank plan, and then to gradually increase the rank from low to full.\n\nIn our case, this gradual rank increase is accomplished implicitly. At each scale $t = 1, \\dots, \\kappa$ this plan $\\mathbf{P}^t$ is made explicit in our supplement, equation (S8). Our rank-annealing schedule $(r_1, \\dots, r_\\kappa)$ describes the sequence of multiplicative factors by which the rank of this explicit plan will increase at each successive scale. The partial products of these, denoted by \n$(\\rho_1, \\dots, \\rho_\\kappa)$, are the ranks of the plans $\\mathbf{P}^1, \\dots, \\mathbf{P}^\\kappa$. We will clarify this relationship in future versions, and distinguish the notion of a rank-annealing schedule from the secondary discussion of how to efficiently choose such a schedule under given memory constraints.", "claims_and_evidence": "The authors claim that for optima transportation problems that do not necessarily possess the low-rank structure as assumed by LOT, HR-LOT enriches the space of transport maps while maintaining the linear runtime when low-rankness is present.", "methods_and_evaluation_criteria": "The examples are illustrative and convincing.", "theoretical_claims": "The Propositions 3.1-3 appear to be correct, although the assumptions in the statements are strong.", "experimental_designs_or_analyses": "The experiments are performed mostly on existing datasets, ranging from synthetic to realistic.", "relation_to_broader_scientific_literature": "The paper is an advance in LOT by adding the hierarchical partitioning. This in itself is novel, using algorithms both in hierarchical OT and LOT, and putting it together with a rank-annealing scheduling in Algorithm 1.", "essential_references_not_discussed": "I'm not aware of paper the authors missed.", "other_strengths_and_weaknesses": "Strengths\n\nThe paper is very well motivated and clearly written, and arguments are convincing. The experiments are thorough and the discussion is sufficiently detailed.\n\nWeaknesses\n\nThe theoretical result Proposition 3.1 seems to have strong assumptions about the transportation plan, although a more general result might be simply very challenging to formulate and prove.", "other_comments_or_suggestions": "- Hierarchical low-rankness of the transportation map seems to share common properties with so-called Hierarchical Matrices used to solve partial differential equations (most notably the Helmholtz equation). Are there some commonality in their approaches?", "questions_for_authors": "I found the discussion regarding the rank-scheduling and annealing scheduling $\\epsilon_n$ in Section 3.2 a bit hard to follow. Is there an alternative explanation for their relation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uKoOmc6q3i", "reviewer": "Reviewer_nCzR", "summary": "This paper focuses on the solving of Optimal Transformer (OT) problems.  To this end, it derives an algorithm, HR-OT that leverages the invariant under Monge map and dynamically constructs a multi-scale partition of each dataset using low-rank OT subproblems. By doing that, it could use linear space and achieve runtime ranging from log-linear to quadratic. Experiments have been conducted on several datasets, even including large scale data with over a million points. The results demonstrate its advantages over the original Sinkhorn.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer _nCzR_ for their feedback and careful reading.\n\n> Yes. The experimental designs are acceptable. However, it is better to include large scale real data, e.g., 3DMatch (https://3dmatch.cs.princeton.edu/}{https://3dmatch.cs.princeton.edu/) or KITTI (https://www.cvlibs.net/datasets/kitti/eval_odometry.php). Moreover, as Sinkhorn has been widely used in point cloud registration algorithms, e.g., CoFiNet (NeurIPS 2021) and GeoTransformer (CVPR 2022), I would like to see the promotion of those methods by replacing the Sinkhorn part with the proposed method.\n\nThank you for your response. We will investigate using Hierarchical Refinement as a module in CoFiNet and GeoTransformer, but unfortunately have limited time in the review period to benchmark it. We will, however, make sure to cite both of these methods and highlight them as a major application area for hierarchical refinement in our discussion with the following sentence:\n\n* _\"Optimal transport has also been successfully integrated into deep-learning frameworks for computer vision and point cloud registration, notably in methods such as CoFiNet (Yu et al. '21) and GeoTransformer (Qin et al. '22), suggesting that hierarchical refinement could help scale existing deep-learning methods based on OT.\"_\n\nOn the scalability of hierarchical refinement on real datasets, we note that the size of the real datasets used was 85958 and 84172 for the two MERFISH datasets, 5913, 18408, 30124, 51365, 77369, 102519, 113350, and 121767 for the seven Stereo-Seq spatial transcriptomics datasets, and 1.281 million for the ImageNet ILSVRC dataset.\n\n> On the MERFISH experiments, the score of HR-OT results are much higher than that of ther methods. However the reason under this should be further analyzed.\n\nIn the MERFISH experiments, as classical full-rank methods were unable to scale, we benchmarked against low-rank OT methods for a fixed rank. Low-rank OT is unable to compute one-to-one alignments in the case that the Monge map exists, and in this case an alignment of spatial coordinates lacks any low-rank cluster structure, which likely explains why full-rank performs significantly better. To address questions raised by other reviewers, we added a comparison to mini-batch OT on this task. Mini-batch OT is also a full-rank method, like hierarchical refinement, and we observe scores which are much closer to it. For example, for the largest batch-size (2048) we found mini-batch cosine scores across 5 genes of (0.7434, 0.7822, 0.7056, 0.4912, 0.5683), compared to the hierarchical refinement scores (0.8098, 0.7959, 0.7526, 0.4932, 0.6015). This implies the gap may largely explained by the difference in expressivity of full-rank ($r=84,172$) versus low-rank couplings ($r=20$ for LOT, $r=500$ for FRLC). \n\n> It has been mentioned that the number of samples ... what is the results compared with other methods?\n\nThank you for your question. To your point, if the datasets are of slightly different sizes one could randomly subsample the number of points in the larger one to ensure an $n$ to $n$ alignment. In all comparisons, the datasets aligned are either $n \\times n$ or sub-sampled to be $n \\times n$ so that all methods (Sinkhorn, ProgOT, FRLC, LOT, HR-OT, mini-batch) are compared on the exact same point clouds. To gauge the effect of subsampling the MERFISH data, we ran LOT without sub-sampling to $n$ points and then LOT with sub-sampling, comparing the cosine similarities on a downstream task, as the primal OT cost is no longer directly comparable. Without the sub-sampling, the cosine score is only slightly higher than with: (0.3390, 0.2712, 0.3186, 0.1666, 0.1080) vs (0.3241, 0.2279, 0.3029, 0.1653, 0.0719). These scores remain significantly lower than those of hierarchical refinement on the sub-sampled data: (0.8098, 0.7959, 0.7526, 0.4932, 0.6015). While sub-sampling incurred little error on this comparison, generalizing hierarchical refinement to directly handle datasets of unequal sizes (i.e. without sub-sampling) is an important direction for future work.\n\n**Table: Cosine Similarity Scores for Expression Transfer**\n\n| Method  | *Slc17a7* | *Grm4*   | *Olig1*  | *Gad1*   | *Peg10*  |\n|-|-|-|-|-|-|\n| HR-OT   | **0.8098**| **0.7959**| **0.7526**| **0.4932**| **0.6015**|\n| FRLC (Halmos et al. '24)  | 0.2180    | 0.2124   | 0.1929   | 0.0963  | 0.0991   |\n| FRLC, no subsampling | 0.2373    | 0.1896   | 0.1579   | 0.0644 | 0.1550   |\n| LOT (Scetbon et al. '21)  | 0.3241    | 0.2279   | 0.3029   | 0.1653  | 0.0719   |\n| LOT, no subsampling  | 0.3390    | 0.2712   | 0.3186   | 0.1666 | 0.1080   |\n| MOP (Gerber et al. '17)  | 0.5211    | 0.4714 | 0.5972   | 0.3571 | 0.2719   |\n| Mini-batch (128)  | 0.6693 | 0.6637   | 0.6442 | 0.4150   | 0.4932 |\n| Mini-batch (512)  | 0.7089 | 0.7383   | 0.6771 | 0.4562   | 0.5383 |\n| Mini-batch (1,024)  | 0.7256 | 0.7621| 0.6918 | 0.4733   | 0.5557 |\n| Mini-batch (2,048)  | 0.7434 | 0.7822| 0.7056 | 0.4912   | 0.5683 |", "claims_and_evidence": "Seems yes.", "methods_and_evaluation_criteria": "Seems yes.", "theoretical_claims": "Yes. Should be correct.", "experimental_designs_or_analyses": "Yes. The experimental designs are acceptable. However, it is better to include large scale real data, e.g., 3DMatch (https://3dmatch.cs.princeton.edu/) or  KITTI (https://www.cvlibs.net/datasets/kitti/eval_odometry.php). Moreover, as Sinkhorn has been widely used in point cloud registration algorithms, e.g., CoFiNet (NeurIPS 2021) and GeoTransformer (CVPR 2022), I would like to see the promotion of those methods by  replacing the Sinkhorn part with the proposed method.", "relation_to_broader_scientific_literature": "As a promoted version of Sinkhorn, it could use linear space and achieve runtime ranging from log-linear to quadratic. Validation is also provided via theoretical proof as well as experiments.", "essential_references_not_discussed": "I have no suggestion about the reference list.", "other_strengths_and_weaknesses": "Strengths:\n1. The writing is fluent and the organization is also good. \n2. The proposed method could be considered novel. \n3. The proposed method could be applied on data with over a million points, which is highly valuable, if no one could achieve that before.", "other_comments_or_suggestions": "See above. I highly suggest the authors to add the metioned experiments.", "questions_for_authors": "1. On the MERFISH experiments, the score of HR-OT results are much higher than that of ther methods. However the reason under this should be further analyzed.\n\n2. It has been metioned that the numer of sample should be the same (|X| = |Y| = n) for the proposed method，and it does not hold, the data could be slightly modified. Could the authors give an example how this could be done? And under this case, what is the results compared with other methods?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HzDnNpNeU0", "reviewer": "Reviewer_j2d6", "summary": "This paper proposes a hierarchical framework to obtain Kantarovich plans in Optimal Transport. The authors conceptually build on many prior works of low-rank OT (notably Scetbon 2021, Halmos 2024) and propose a rank annealing schedule to obtain the full rank Kantarovich plan using many low-rank solutions of sequentially partitioned sets. The main premise is founded on a theoretical result that when a Monge map exists between the datasets, it co-clusters the optimal low-rank solution (albeit with minor differences like uniform assumption in Eq 6). Altogether the proposed framework is claimed to be linear in space complexity and log-linear to quadratic in time complexity. The authors contrast this with well-known OT solvers mainly Sinkhorn (which has quadratic space and quadratic time) and ProgOT.\n\nExperiments are reported for (1.) Synthetic Datasets (Checkerboard, MAFMoons and Rings, and Half-Moon and S-Curve etc) (2.) Large-scale Matching Problems and Transcriptomics (3.) MERFISHBrainAtlasAlignment and (4.) ImageNet Alignment. The authors compare with Sinkhorn (Cuturi, 2013), ProgOT (Kassraieetal.,2024), as well as low rank OT solvers like (Scetbon et al.,2021) and FRLC(Halmos et al.,2024). The main message from the experiments is to highlight that the proposed approach scales better (in space complexity) than existing full-rank solvers and produces a smaller kantarovich cost, establishing a good operating point between accuracy and complexity.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer _j2d6_ for their feedback and careful reading.\n\n> At a high level, the main message of the paper is to propose the use of low rank OT solvers to sequentially obtain a full rank solution. This submission does NOT propose a new low rank OT solver. If we credit the authors for focusing on this novelty - however, I find that the authors do make specific choices (1.) Use of Halmos et al in Algorithm 1 (2.) the specific Select subroutine in Algorithm 1. While the choices are not unreasonable, I do find the whole setup arguably lacking generality. For e.g., I could not place any experiment that validates replacing with Scetbon 2021 et al, etc. I feel it is most impactful when the novel idea is shown to work for different LROT solvers, which I am not yet convinced.\n\nWe agree on the benefit of demonstrating with other LROT solvers.  We did not have time to implement this in the short response window, but aim to offer this option in the code and demonstrate results with the Scetbon et al. 2021 solver. Theoretically, our framework will work for the Scetbon et al. solver because our Proposition 3.1 is tailored to the low-rank factorization introduced by (Scetbon et al. '21). Proposition 3.1 hinges on Lemma B1, which shows _optimal_ factors $(\\mathbf{Q}, \\mathbf{R})$ for the LOT problem (Scetbon et al. '21) are vertices on the transport polytope and thus have $\\leq n+r-1=n+1$ non-zero entries for $r=2$. Thus for $n = 2^t$ the solutions correspond to partitions in Proposition 3.1. The key point is these arguments are independent of the solver, and apply to any optimal solution. \n\nRegarding (2), our choice of the argmax as the Select sub-routine is motivated by our Proposition 3.1, and is correct if $(\\mathbf{Q}, \\mathbf{R})$ are optimal, but other choices may be appropriate if the solutions are entropic and do not define exact partitions.\n\n> I do feel the lack of any comparison to Gerber & Maggioni, 2017 to be underwhelming. Even though I agree it is not apples to apples comparison, this is a very strong and similar prior work, and perhaps an experiment highlighting the difference in performance for a reasonable choice of initial partition would be very enlightening.\n\nWe agree, Gerber & Maggioni is a seminal work in multi-scale OT, and we have added comparisons with their method MOP for our 2-dimensional datasets where MOP uses the GMRA (Geometric Multi-Resolution Analysis) R package to generate multiscale partitions. In particular, we benchmarked against MOP for our three synthetic datasets and our MERFISH expression transfer task. MOP's performance for MERFISH is given in the table above in response to reviewer _nCzR_. MOP's performance on synthetic data is given in the table below:\n\n**Table: Comparison of Coupling-Based OT Methods on Primal Cost $\\langle \\mathbf{C}, \\mathbf{P} \\rangle_F$ (Wasserstein-2) on 512 point small instance**\n\n| Method | Checkerboard | MAF Moons & Rings | Half Moon & S-Curve |\n|-|--|--|--|\n| MOP (Gerber et al. '17)  | 0.393 | 0.276 | 0.401   |\n| Sinkhorn (`ott-jax`)   | 0.136   | 0.221 | 0.338  |\n| ProgOT (Kassraie et al. '24)  | 0.136  | 0.216 | 0.334   |\n| HR-OT  | 0.129        | 0.216  | 0.334               |\n| Dual Revised Simplex Solver   | **0.127** | **0.214** | **0.332**   |\n\n> Eq 1 incomplete (missing b); Line 119, define $\\delta_r$\n\nThank you for catching these mistakes. We will fix them in the final version.\n\n> Are the pointclouds in Figure 3 for comparing the 3 methods the same?\n\nThe left two pointclouds for  hierarchical refinement and Sinkhorn are the same, since these methods are able to run on 4096 points. The rightmost figure is for an optimal LP-solver, which scaled only to 512 points and thus is plotted with an identically distributed but smaller dataset. We will add this distinction to the legend.\n\n> Is the Sinkhorn used in the experiments (ott-jax) with the $\\epsilon$ annealer?\n\nRegarding your question about the Sinkhorn implementation: we use `ott-jax` as the Sinkhorn solver, but do not use the $\\epsilon$-annealer. We rely on the default implementation of Sinkhorn, which uses a fixed value of $\\epsilon = 0.05$. This default value produces relatively sparse solutions. We will note this detail in the experimental section.\n\n> This citation is missing: Cuturi, M., Meng-Papaxanthos, L., Tian, Y., Bunne, C., Davis, G., \\& Teboul, O. (2022). Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324.\n\nThank you for catching this, we will add this citation after our references of `ott-jax`.", "claims_and_evidence": "The paper is written well. The main contributions are highlighted nicely and the experiments reflect the nuanced benefits of the proposed approach. The main method is convincing - and I found it nice that the algorithm builds on Proposition 3.1", "methods_and_evaluation_criteria": "Yes. The experiments are diverse and fairly comprehensive", "theoretical_claims": "Yes. Propositions 3.1 3.2 and 3.3.", "experimental_designs_or_analyses": "Yes. All of the 4 experiments reported in the summary.", "relation_to_broader_scientific_literature": "- At a high level, the main message of the paper is to propose the use of low rank OT solvers to sequentially obtain a full rank solution. This submission does NOT propose a new low rank OT solver. If we credit the authors for focusing on this novelty - however, I find that the authors do make specific choices (1.) Use of Halmos et al in Algorithm 1 (2.) the specific Select subroutine in Algorithm 1. While the choices are not unreasonable, I do find the whole setup arguably lacking generality. For e.g., I could not place any experiment that validates replacing with Scetbon 2021 et al, etc. I feel it is most impactful when the novel idea is shown to work for different LROT solvers, which I am not yet convinced.", "essential_references_not_discussed": "- I do feel the lack of *any* comparison to Gerber&Maggioni, 2017 to be underwhelming. Even though I agree it is not apples to apples comparison, this is a very strong and similar prior work, and perhaps an experiment highlighting the difference in performance for a reasonable choice of initial partition would be very enlightening", "other_strengths_and_weaknesses": "Overall, barring some important clarifications (e.g., using the correct sinkhorn implementation), I am inclined positively. The paper has been compiled well, has an interesting message, and has comprehensive experiments.", "other_comments_or_suggestions": "- Eq 1 incomplete (missing b)\n- Line 119, define  \\delta_r", "questions_for_authors": "- Are the pointclouds in Figure 3 for comparing the 3 methods the *same*? \n- Is the Sinkhorn used in the experiments (ott-jax) with the \\eps annealer? \n- This citation is missing: Cuturi, M., Meng-Papaxanthos, L., Tian, Y., Bunne, C., Davis, G., & Teboul, O. (2022). Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uNy5erEZ8M", "reviewer": "Reviewer_9vCU", "summary": "This paper proposes a new large-scale hierarchical optimal transport algorithm, between distributions with the same number of samples. The authors algorithm is based on multi-scale partitions of the source and target datasets, as well as recent development in low-rank optimal transport. Through a series of experiments on synthetic and real-world datasets, the authors show that their approach 1) has better alignment performance than previously proposed methods and 2) can scale to larger datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer _9vCU_ for their feedback and careful reading.\n\n> As I mentioned ... authors could have cited a few papers ...\n\nThank you for these suggestions, we will include the following sentences in our Background: \n_\"Hierarchical OT (Schmitzer and Schnorr '13) is a variant of OT modeling data and transport across multiple scales, using Wasserstein distances as coarse-scale ground costs. It has been applied to document representation (Yurochkin et al. '19), domain adaptation (El Hamri et al. '22), sliced Wasserstein distances (Bonneel et al. '15; Nguyen et al. '22b) and to give a discrete formulation of transport between Gaussian mixture models (Chen et al. '18; Delon and Desolneux '20) . These works build interpretable, coarse-grained structure into a single transport plan, rather than solving for a sequence of plans at progressively finer scales as in the present work.\"_\n\nIn addition, we have added experimental validation against mini-batch OT. We will add the following sentences to our Introduction:\n_\"Mini-batch OT (Genevay et al. '18) improves scalability, but incurs significant biases (Sommerfeld et al. '19; Korotin et al. '21; Fatras et al. '21a) as each mini-batch alignment is often a poor representation of the global one. Several works have investigated the theoretical properties of mini-batch estimators of the plan (Fatras et al. '20b; '21c), while others have attempted to mitigate bias using partial or unbalanced OT (Nguyen et al. '22a; Fatras et al. '21b) by allowing for mass variation between the mini-batches.\"_\n\n>  ..deeper discussion of minibatch OT..\n\nThank you for pointing this out. We re-ran the large-scale real data experiments with mini-batch OT without-replacement, which is more conventional as noted in (Fatras et al. '21c), and will include these values in the updated manuscript. See the tables for the mouse embryo spatial transcriptomics, ImageNet below and MERFISH experiments above in response to _nCzR_, including mini-batch couplings. We notate with (MB $B$) for $B$ the batch size in each case. One value in the HR-OT table (Cost 14.35 for 12.5-13.5) has been changed, and reflects a slightly higher setting of $r _ {max} = 128$ in the rank annealer. We also will add more discussion to our paper on practical considerations for mini-batch OT, as well as existing theoretical results on convergence and complexity (Fatras et al. '20a; '21a; '21c) relative to full-rank OT.\n\n**Table: Cost Values for Embryo**\n\n|Method|9.5-10.5|10.5-11.5|11.5-12.5|12.5-13.5|13.5-14.5|14.5-15.5|15.5-16.5|\n|-|-|-|-|-|-|-|-|\n|HR-OT|**21.81**|**14.81**|**16.14**|**14.35**|**13.78**|**14.29**|**12.79**|\n|MB 128|22.44|15.35|16.69|14.86|14.14|14.75|13.32|\n|MB 512|22.15|15.05|16.33|14.54|13.92|14.50|13.01|\n|MB 1024|22.05|15.02|16.24|14.45|13.86|14.43|12.91|\n|MB 2048|21.98|14.98|16.18|14.39|13.81|14.39|12.85|\n\n**Table: Cost Values for ImageNet**\n\n|Method|HR-OT|MB 128|MB 512|MB 1024|\n|-|-|-|-|-|\n| |**18.97**|21.89|20.34|19.58|\n\n\nWe will also log-transform Fig.2's $x$-axis for the next version!\n\n> Do the linear space complexity..?\n\nGreat question. While you are correct that bijections can be stored with linear space, our linear space complexity does not rely on $X$ and $Y$ having the same number of elements. For classical approaches tackling OT as an assignment problem, while the final coupling returned is linear in space, the space complexity of the algorithm is quadratic. In comparison, HR-OT returns a solution with linear space and has linear space complexity. In addition, the runtime complexity is log-linear for squared-Euclidean cost $\\lVert \\cdot \\lVert_{2}^{2}$ and remains log-linear for sample-linear approximations of the distance matrix (Indyk et al. '19). We refer to the bullet below on your question concerning whether hierarchical refinement could retain its properties with an unequal number of points.\n\n> What would be the challenges of adapting their method to handle distributions with unequal number of samples?\n\nThere are two challenges to extend hierarchical refinement to datasets with $n$ source points and $m$ target points.\n\n1. If $n > m$, the challenge is to extend Proposition 3.1 from 1-1 Monge maps (the assignment problem) to many-to-one Monge maps (i.e. with smaller target dataset). Specifically, one needs to account for the possibility that each target index $j \\in [m]$ is co-clustered with a _set_ $ S _ {j} \\subset [n]$ of indices indexing the preimage of the Monge map $T^{-1}( \\mathbf{x} _ {j} )$. \n\n2. If $m>n$, no Monge map exists from the source to target dataset. While Sinkhorn approaches the Kantorovich problem, hierarchical refinement approaches the Monge formulation of optimal transport. Thus, it inherits any limitations of the Monge framework, including its asymmetry and inability to account for mass-splitting. However, assuming the extension of the proposition discussed in point (1.) holds, one can reverse the role of the two datasets and infer the Monge map from target to source.", "claims_and_evidence": "There are 2 main claims in this paper.\n\n1) The proposed strategy successfully aligns source and target distributions,\n2) The algorithm can scale to larger datasets\n\nBoth of them are supported in the experiments. Furtheremore, there are some theoretical results supporting 1).", "methods_and_evaluation_criteria": "The experiments are well designed and use relevant benchmarks. The authors could have considered other problems where OT has contributed to, such as generative modeling or domain adaptation, but overall the experiments are comprehensive.", "theoretical_claims": "There are 3 main theoretical results (Propositions 3.1, 3.2 and 3.3). I could not assess the validity of these proofs. I have some questions down below my review.", "experimental_designs_or_analyses": "The experiments are valid and well designed. As I mentioned previously, they are comprehensive.", "relation_to_broader_scientific_literature": "This paper fits into the broader litreature of Hierarchical OT (Chen, Georgiu and Tannenbaum, 2017; Yurochkin et al, 2018; Delon and Desolneux, 2019; El Hamri, Bennani and Falih, 2021), which consider transportation plans at different levels of representations of probability measures through clustering. Arguably, the paper pushes this idea further, as it consider partitions of increasing depth of the source and target datasets.", "essential_references_not_discussed": "As I mentioned in the previous point, I think the paper could benefit from a broader discussion with hierarchical optimal transport, especially these references,\n\n(Chen, Georgiu and Tannenbaum, 2017) Chen, Yongxin, Tryphon T. Georgiou, and Allen Tannenbaum. \"Optimal transport for Gaussian mixture models.\" IEEE Access 7 (2018): 6269-6278.\n\n(Yurochkin et al, 2018) Yurochkin, Mikhail, et al. \"Hierarchical optimal transport for document representation.\" Advances in neural information processing systems 32 (2019).\n\n(Delon and Desolneux, 2019) Delon, Julie, and Agnes Desolneux. \"A Wasserstein-type distance in the space of Gaussian mixture models.\" SIAM Journal on Imaging Sciences 13.2 (2020): 936-970.\n\n(El Hamri, Bennani and Falih, 2021) El Hamri, Mourad, Younes Bennani, and Issam Falih. \"Hierarchical optimal transport for unsupervised domain adaptation.\" Machine Learning 111.11 (2022): 4159-4182.\n\nFurthermore, while the authors do mention mini-batch OT, the authors could have cited a few papers on the problem, such as,\n\n(Nguyen et al., 2022) Khai Nguyen, Dang Nguyen, The-Anh Vu-Le, Tung Pham, Nhat Ho Proceedings of the 39th International Conference on Machine Learning, PMLR 162:16656-16690, 2022.\n\n(Fatras et al., 2022) Kilian Fatras, Thibault Sejourne, Rémi Flamary, Nicolas Courty Proceedings of the 38th International Conference on Machine Learning, PMLR 139:3186-3197, 2021.\n\n(Fatras et al., 2020) Fatras, Kilian, et al. \"Learning with minibatch Wasserstein: asymptotic and gradient properties.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2020.\n\nincluding hierarchical approaches\n\n(Nguyen et al., 2022) Khai Nguyen, Dang Nguyen, Quoc Dinh Nguyen, Tung Pham, Hung Bui, Dinh Phung, Trung Le, Nhat Ho Proceedings of the 39th International Conference on Machine Learning, PMLR 162:16622-16655, 2022.", "other_strengths_and_weaknesses": "Here, I make a summary of my review.\n\n**Strenghts**\n\n1. The paper provides a practical and scalable algorithm for large-scale optimal transport\n2. The idea of using multi-level partitions is quite interesting\n3. The discussion at the end of the paper is quite comprehensive and deals with natural questions that arise from the reading of the paper.\n\n**Weaknesses**\n\n1. The authors could have considered minibatch-OT in their experimental comparisons. There could be a deeper discussion of minibatch OT in the main paper\n\n----\n\nAs a result, I am leaning towards acceptance, as I think the paper is good.\n\n---\n\n__Post Rebuttal:__ The authors addressed my concerns in their rebuttal. As a result, I raise my score from __3. Weak Accept__ to __4. Accept__", "other_comments_or_suggestions": "In Figure 2, the x-axis should be in log-scale", "questions_for_authors": "Here are a few questions for the authors,\n\n- Do the linear space complexity stems from the assumption that $X$ and $Y$ have the same number of elements? For instance, in principle, if $X$ and $Y$ have the same number of elements and have uniform weights, one can store $\\{x_{i}, T(x_{i})\\}_{i=1}^{n}$ with classic OT.\n- What would be the challenges of adapting their method to handle distributions with unequal number of samples?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a scalable and elegant framework for computing full-rank optimal transport using a hierarchical refinement of low-rank solutions. The method combines insights from low-rank OT and multiscale partitioning to construct bijective transport plans while maintaining linear space complexity and favorable runtime scaling. Reviewers consistently noted the novelty of this approach, the solid theoretical grounding—including formal propositions—and the wide range of convincing experiments across synthetic, biological, and large-scale visual datasets. The rebuttal successfully addressed concerns about missing comparisons and limitations in the discussion of related literature. In particular, the authors added benchmarks against mini-batch OT, provided further clarity on implementation details, and drew connections to hierarchical matrix methods. They also committed to broadening the experimental coverage and referencing additional hierarchical OT and domain adaptation works in the final version. These responses were well-received by the reviewers. Given the conceptual clarity, practical impact, and strong reviewer support, I recommend acceptance.
Decision: Accept (oral)