Progressive Entropic Optimal Transport Solvers
Authors: Parnian Kassraie
Aram-Alexandre Pooladian
Michal Klein
James Thornton
Jonathan Niles-Weed
marco cuturi
Venue: NeurIPS 2024 poster (2024)
Keywords: Optimal Transport
Entropy Regularization
Primary area: other
TL;DR: We propose a progressive algorithm for estimation of OT maps and plans. We prove its statistically consistency and demonstrate its performance on synthetic and single-cell data.

Abstract: Optimal transport (OT) has profoundly impacted machine learning by providing theoretical and computational tools to realign datasets.
In this context, given two large point clouds of sizes $n$ and $m$ in $\mathbb{R}^d$, entropic OT (EOT) solvers have emerged as the most reliable tool to either solve the Kantorovich problem and output a $n\times m$ coupling matrix, or to solve the Monge problem and learn a vector-valued push-forward map. 
While the robustness of EOT couplings/maps makes them a go-to choice in practical applications, EOT solvers remain difficult to tune because of a small but influential set of hyperparameters, notably the omnipresent entropic regularization strength $\varepsilon$. Setting $\varepsilon$ can be difficult, as it simultaneously impacts various performance metrics, such as compute speed, statistical performance, generalization, and bias. In this work, we propose a new class of EOT solvers (ProgOT), that can estimate both plans and transport maps.
We take advantage of several opportunities to optimize the computation of EOT solutions by *dividing* mass displacement using a time discretization, borrowing inspiration from dynamic OT formulations, and *conquering* each of these steps using EOT with properly scheduled parameters. We provide experimental evidence demonstrating that ProgOT is a faster and more robust alternative to *standard solvers* when computing couplings at large scales, even outperforming neural network-based approaches. We also prove statistical consistency of our approach for estimating OT maps.

Progressive Entropic Optimal Transport Solvers
Parnian Kassraie ETH Zurich, Apple pkassraie@ethz.cm ch
Aram-Alexandre Pooladian New York University
Michal Klein
Apple mi michalkk@pplle.com com
aram-a -rammalexandre.eoo wttp:/aeeeeanreeppaaaaaaaaaaaaaa edu
James Thornton
Jonathan Niles-Weed
Marco Cuturi Apple cuturi@apple.com
Apple Jamestchoraton@apple.. com
New York University jnw@cims .ny...yue edu
Abstract
Optimal transport (OT) has profoundly impacted machine learning by providing the- oretical and computational tools realign datasets. In this context, given large point clouds of sizes n and m in R entropic OT (EOT) solvers have emerged as the most reliable tool to either solve the Kantorovitch problem and output pling matrix, or to solve the Monge problem and learn vector-valued push- purw-forward map. While the robustness of EOT couplings/maps makes them a go-to choice in practical applications, EOT solvers remain difficult to tune because of small but influential set hyperparameters, notably the omnipresent entropic regularization strength Setting can be difficult, it simultaneously impacts various perfor- mance metrics, such as compute speed, statistical performance, generalization, and bias. In this work, propose mate both plans and transport maps. We take advantage of several opportunities to optimize the computation of EOT solutions by dividing mass displacement using time discretization, borrowing inspiration from dynamic OT formulations [McCann, 1997], and conquering each of these steps using EOT with properly scheduled parameters. We provide experimental evidence demonstrating that PROGOT a faster and more robust altermative to EOT solvers when computing couplings and maps at large scales, even outperforming neural network- based approaches. We also prove the statistical consistency of PROGOT when estimating OT maps.
cou-
class of EOT solvers (PROGOT), that can esti-
Introduction
Many problems in generative machine learning and natural sciences-nottan biology [Schiebinger et al., 2019, Bunne et al., 2023], astronomy [Métivier et al., 2016] or quantum chemistry [Buttazzo et al., 2012] require aligning datasets learning map data points from source to target distribution. These problems stand at the core of optimal transport theory [Santambrogio, 2015] and have spurred the proposal of various solvers [Peyré al., 2019] to perform these tasks reliably. In these tasks, we are given n and m points respectively sampled from source and target probability distributions Rdd with the goal of either returning coupling matrix of size the so-called Kantorovitch problem), or vector-valued map estimator that extends to out-of-sample data (solving the Monge problem).
(which solves
In modern applications, where n. 2 10 a popular approach to estimating either coupling maps rely on regularization of the original Kantorovitch linear OT formulation using neg-entropy. This technique, referred to as entropic OT, can be traced back to Schrödinger and was popularized for MI. applications [Cuturi, 2013] (see Section 2). Crucially, HOT can be solved efficiently with Sinkhorn's 's algorithm (Algorithm 1), with favorable computational [Altschuler et al., 2017, Lin et al., 2022] and statistical properties [Genevay, 2019, Mena and Niles- Weed, 2019] compared to linear programs, Most couplings computed nowadays on large point clouds within ML applications
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

EOT
Debiased EOT
PROGOT
Xtrin
Xtol
Yieaia
Transported [
Niterpolate
Figure 1: (left) EOT solvers collapse when the value of is not properly chosen. This typically results in biased map estimators and in blurry couplings (see Fig. for the coupling matrix obtained between Xtmin and Ymii). (middle) Debiasing the output of EOT solvers can prevent collapse the mean seen in EOT estimators, but computes the same coupling. PROGOT (right) ameliorates these problems in various ways: by decomposing the resolution of the OT problem into multiple time steps, and using various forms of progressive scheduling, we recover both coupling whose entropy can be tuned automatically and a map estimator that is fast and reliable.
are obtained using EOT solvers that rely on variants of the Sinkhorn algorithm, whether explicitly, lower-level subroutine [Scetbon et al., 2021, 2022]. The widespread adoption of EOT has spurred many modifications of Sinkhorn's 's original algorithm (e.g.. through acceleration [Thibault al., 2021] or initialization [Thornton and Cuturi, 2023]), and encouraged its incorporation within neural-network OT approaches [Pooladian et al., 2023, Tong et al., 2023, Uscidda and Cuturi, 2023].
Though incredibly popular, Sinkhorn's 's algorithm is without its drawbacks. While a popular tool due its scalability and simplicity, its numerical behavior deeply impacted by the amount of neg-entropy regularization, driven by the hyperparameter e. Some practitioners suggest to have the parameter nearly vanish [Xie et al., 2020, Schmitzer, 2019], others consider the case where it diverges, highlighting links with the maximum mean discrepancy [Ramdas et al., 2017, Genevay et al., 2019].
EOT
PRocOT
0.0000
0.000
0.0002
0.0603
0.00034
Figure 2: Coupling matrices between train points in Fig. Comparison of EOT with fairly large e, and PROGOT which automati- cally tunes the entropy of its coupling accord- ing to the target point cloud's 's dispersion.
Several years after its introduction to the machine learning community [Cuturi, 2013], choosing suitable regularization term for EOT remains a thorny pain point. Common approaches are setting
> 0 to a default value (c.g., the max [Flamary al., 2021] ]] or mean [Cuturi et al., 2022b] normalization of the transport cost matrix), incorporating a form of cross-validation or an unsupervised criterion [Vacher and Vialard, 2022, Van Assel et al.,
2023], or scheduling [Lehmann et al., 2022, Feydy, 2020]. When e is too large, the algorithm converges quickly, but yields severely biased maps (Figure I, left), or blurry, uninformative couplings (Figure 2). Even theoretically and numerically debiasing the Sinkhorn solver (Figure 1, middle) does not seem to fully resolve the issue [Feydy et al., 2019, Pooladian et al., 2022]. To conclude, while strategies exist to alleviate this bias, there currently exists no one-size-fitte fits-all solution to this problem.
Our contribution: an EOT solver with a dynamic lens. Recent years have witnessed an explosion in neural-network approaches based on the so-called Benamou and Brenier dynamic formulation of OT [Lipman et al., 2022, Liu, 2022, Tong et al., 2023, Pooladian et al., 2023]. A benefit this perspective is the ability to split the OT problem into simpler too-problems that likely better conditioned than the initial transport problem. With this observation, we propose a novel family of progressive EOT solvers, called PROGOT, that meant to be sturdier and easier to parameterize than existing solvers. Our key idea is to exploit the dynamic nature of the problem, and vary parameters dynamically, such and convergence thresholds, along the transport. We show that PROGOT
can be used to recover both Kantorovitch couplings and Monge map estimators,
gives rise to novel, provably statistically consistent map estimator under standard assumptions.
strikes the right balance between computational and statistical tradeoffs,
can outperform other (including based) approaches on real datasets,

---

2 Background
Optimal transport. For domain S2 S RR let P2((2) denote the space of probability measures over Ω with finite second moment, and let (9t) be those with densities. Let μ, E P2(S), and let [(,, v) be the set of joint probability measures with mert-marginal μ and right-marginal v. We consider a translation invariant cost function c(x,y) == h(x y), with h a strictly convex function, and define the Wasserstein distance, parameterized by h, between and v
// h(2 y)d (2, y)
W (H, = inf
wEI(u.)
This formulation is due to Kantorovitch [1942], and we call the minimizers to (1) OT couplings or OT plans, and denote as To. A subclass of couplings are those induced by pushforward maps. We say that T: IRd -> Id pushes js forward to if T(X) for X J4, and write T# /1 v. Given a choice of cost, we can define the Monge [1781] formulation of OT
/ hix - (((=)T())(p)))
To
arg min T:Iap=oJ
where the minimizers are referred to as Monge maps, or OT maps from to u. Unlike OT couplings, OT maps are not always guaranteed to exist. Though, if y has density, we obtain significantly more structure on the OT map:
Theorem 1 (Brenier's 's Theorem [1991]). Suppose μ E P2 (S) and P2(()). Then there exists unique solution to (2) that is of the form To = Ld - vhh a V for where he is the convex-conjugate of h, i.e. h² (y) == max, (z, 3) h(z)) and
/ fdu++ f gdv,
(fo. E arg max (fagie
(3)
where F = (f19) E (1) x (v) f(x) + g(v) I Mxx - v). Va, y E S2.]. Moreover, the OT plan is given by To Lo(dz, dy) = (m) (d))
Importantly, (3) is the dual problem to (1) and the pair of functions (fo.go) are referred to as the optimal Kantorovich potentials. Lastly, we recall the notion of geodesies with respect to the Wasserstein distance. For pair of measures μ and with OT map To. the McCann interpolation
between y and v is defined as
Na (a-
where E [0, 1]. Equivalently, Mo is the law of X, = (1 a) aTT, where X ~ μ. In the case where = 1| | for p > the McCann interpolation is in fact a geodesic in the Wasserstein space [Ambrosio et al., 2005]. While this equivalence may not hold for general costs, the McCann interpolation still provides a natural path of measures between μ and v [Liu, 2022].
Entropic OT. Entropic regularization has be- come the de-facto approach to estimate all three variables To. fo and To using samples (xi.....X, xx) and (y1..... vm). both weighted by probability weight vectors a E R7, b RT summing to 1, to form approximations jnn Σ ax0x, and in Ene bj A com- mon formulation of the EOT problem the fol- lowing e-strongly concave program:
Algorithm 1 SINK(a, X, b, Y, finit Binit).
1: f,g t fin ginit (zero by default)
2: XI...., X, X,
,1,..., ym Y
3: C+ [h(x: yi)lij
4: while || |xp CC 1, - aal11 < do
5: ft e-egoga min (C f0 g) +
6: -- log ming (C g $ f) +
7: end while
f* g" FE3" RER
arg max (f) + (g, b)
8: return g. exp ((C f 1 g)/)
=/(e/e, KeR/es, (5)
where 6 > 0 and Ki, lexpeeexxxxx - /E)ls, E Ruxm We can verify that (5) is a regularized version of (3) when applied to empirical measures [Peyré et al., 2019, Proposition 4.4], Sinkhorn's algorithm presents iterative scheme for obtaining (f*,g"), and we recall it in Algorithm I, where for matrix S [S.
the notation ming (S) = |-s log (1 Ts..../))) and $ is the tensor
3

---

sum of two vectors, i.e. f@g:= Eyزر3. Note that solving (5) also outputs a valid coupling P = exp(-(ff + 8))/E), which upproximately solves the smite-sample counterpart of (1). K
Additionally, the optimal potential fo can be approximated by the entropic potential
fz(=) == min,(lg) h(x -y))b))
(6)
where an analogous expression can be written for 9e in terms of f* Using the entropic potential, we can also approximate the optimal transport map To by the entropic map
(x(x) = x - Vh o Vfe(x). This connection is shown in Pooladian and Niles-Weed [2021, Proposition 2] for h 2 -11-10 [Cuturi et al., 2023] for more general functions.
1|2 and
3 Progressive Estimation of Optimal Transport
We consider the problem of estimating the OT solutions To and To. given empirical measures jú and from n i.i.d. samples. Our goal is to design an algorithm which is numerically stable, computationally light, and yields a consistent estimator. The entropic map (7) is an attractive option to estimate OT maps compared to other consistent estimators [e.g.. Hütter and Rigollet, 2021, Manole et al., 2021]. In contrast to these methods, the entropic map is tractable since it is the output of Sinkhorn's 's algorithm. While Pooladian and Niles-Weed [2021] show that the entropie map is bona fide estimator of the optimal transport map, it hides the caveat that the estimator is always biased. For any pre-set ε > O, the estimator is never valid pushforward map i.e., (Te)#11 F v, and this holds true the number of samples tends to infinity. In practice, the presence of this bias implies that the performance of T. is sensitive to the choice of E, e.g. as in Figure 1. Instead of having Sinkhorn as the end-all solver, we propose to use it as a subroutine. Our approach is to iteratively move the source closer to the target, thereby creating a sequence of matching problems that are increasingly easier to solve. As a consequence, the algorithm is less sensitive to the choice of E for the earlier EOT problems, since it has time to correct itself at later steps. To move the source closer the target, we construct a McCann-type interpolator which uses the entropic map Te of the previous iterate, as outlined in the next section.
3.1 Method
As a warm-up, consider To the optimal transport map from μ to v. We let T(O) == To and define
g(0) = (1 ax) Id (o) This gives rise to the measure y(() S(0) μ., which traces out the
McCann interpolation between (n, v) as a varies in the interval (0, 1). Then, letting T')) be the optimal transport map for the pair (n(() (), 2/), it is straightforward to show that T( (()) = T())
In other words, in the idealized setting, composing the output of progressive sequence of Monge problems along the McCann interpolation path recovers the solution to the original Monge problem.
Building on this observation, we set up a progressive sequence of entropic optimal transport problems, along estimated in- terpolation path, between the empirical counterparts of (J4, v). We show that, long as we remain close to the true interpola- tion path (by not allowing a to be too large), the final output is close to v. Moreover, the algorithm progresses, choosing the parameters e, becomes a less arduous task, and computation of T. becomes more stable numerical problem.
cr(3)
c( g(0)
g(0)
((3
(11
At step zero, we set (0) = - and calculate the entropic map
((() = Te. from samples (āt") (0) ()) with regularization pa- rameter co > O. To set up the next EOT problem, we create an intermediate distribution via the McCann-type interpolation
Figure 3: Intuition of PROGOT: By iteratively fitting to the interpola- tion path, the final transport step is less likely to collapse, resulting in more stable solver.
A0 = 1 p! ((0) = (1 - a0) ld (a)
with ap E 1). In doing so, we are mimicking step along the interpolation path for the pair (μ, v)) In fact, we can show that fi! (1) is close to Ma, as defined in (4) (see Lemma 12). For the next iteration of the algorithm, we choose E and a, compute (!)) the entropic map for the pair (fL 2) with regularization E and (2) We repeat the same move along the estimated interpolation path by computing the distribution fic
process for K steps. The algorithm then outputs the progressive entropic map
J(K) = {(K) gg( (k) o... (())

---

where 8(8) = - (1-ag) Id +age is the McCann-type interpolator at step k. Figure 3 visualizes the one-step algorithm, and Definition 2 formalizes the construction of our progressive estimators.
Definition (PROGOT), For two empirical measures p. and given step and regularization schedules (ox k==o and (ex) the PROGOT map estimator T is defined as the composition
T Prog == g(K) o g(K 0 ....0. o(0)
where these maps are defined recursively, starting from Pe = ft and then at each iteration:
(()) is the entropic map Pew computed between samples (ite) 0) with regularization Es.
g(k) (1 (:k) Id +a+E(k, is McCann-type interpolating map time 0:k-
p( = (k)) pie)) the updated measure used in the next iteration.
Additionally, the PROGOT coupling matrix P between f and & is identified with the matrix solving the discrete EOT problem between ((K) and i.
The sequence of (ax) k=o characterizes the speed of movement along the path. By choosing Og a<k) we can recover a constant-speek curve, or an accelerating curve which initially takes large steps and as it gets closer to the target, the steps become finer, or a decelerating curve which does the opposite. This is discussed more detail Section and visualized in Figure (4-C). Though theoretical guarantee requires particular choice for the sequence (=x) and k experimental results reveal that the performance of our estimators not sensitive to this choice. We steering close
hypothesize that this behavior is due to the fact that PROGOT is serff-oorrecting" to the interpolation path, later steps in the trajectory can correct the biases introduced in earlier steps.
3.2 Theoretical Guarantees
By running PROGOT, we are solving a sequence of EOT problems, each building on the outcome of the previous one. Since еггor potentially accumulate across iterations, leads us to ask if the algorithm diverges from the interpolation path and whether the ultimate progressive map estimator Prog is consistent, focusing on the squared- Euclidenn cost of transport, i.e., h 1 1| To answer this question, we assume
(A1) p., v E P2..c((2) with Ω C Rdd convex and compact, with 0 < Mmmn S v(-) S Vmax and u(-) Mmax.
(A2) the inverse mapping + (To(z)) has at least three continuous derivatives,
(A3) there exists X, A > 0 such that AI < DTo(z) 3 AI, for all E E (D denotes Jacobian)
and prove that PROGOT yields a consistent map estimator. Our error bound depends on the number of iterations K, via a constant multiplicative factor. Implying that T(k) is consistent as long as K does not grow too quickly as function of n the number of samples. In experiments, we set K <
Theorem 3 (Consistency of Progressive Entropic Maps). Let h = 111 - 112. Suppose μ., and their optimal transport map To satisfy (A1)-(A3), further suppose we have i.i.d. samples from both j (K) be as defined in Definition 2, with parameters EN X n 2a ,rg x n for all k E [K]. and v. Let Prog Then, the progressive entropic map is consistent and converges to the optimal transport map as
(K)
Tol Tollz (m Slog(n),K -
Prog
where the notation implies that the inequality ignores terms of rate Poly(log(n), K).
The rate of convergence for PROGOT is slower than the convergence of entropic maps shown by Pooladian and Niles-Weed [2021] under the same assumptions, with the exception of convexity of Ω2. However, the rates that Theorem suggests for the parameters as and ε are set merely to demonstrate convergence and do not reflect how each parameter should be chosen function of when executing the algorithm. We will present practical schedules for (a and (sx) in Section 4. The proof is deferred to Appendix C; here present brief sketch.
5

---

Proof sketch. In Lemma 11, we show that
K
K
E
(k) Prog
- To
2 (m))) k=0
A =
Blle(k) - -1(()) (12()0)))
where (k) is a point the true interpolation path, and T(k), is the optimal transport map emanating from it. Here, g(K) is the entropic map estimator between the final target points and the data that has been pushed forward by earlier entropic maps. It suffices to control the term AR. Since {(k) and T(A) are calculated for different measures, we prove a novel stability property (Proposition 4) show that along the interpolation path, these intermediate maps remain close their unregularized population counterparts, if (g and e; are chosen prescribed. This result is based off the recent work by Divol et al. [2024] and allows us to recursively relate the estimation at the k-th iterate the estimation at the previous ones, down to Ao. Thus, Lemma 12 tells us that, under our assumptions and parameter choices og x n n-//d 1/d
and E n 1/2d, it holds that for all k 2 0
A Slog(n) 1/d
Since the stability bound allows us to relate As to Do. combined with the above, we have that
A Slog(n) Do Slog(x) n 1/d
where the penultimate inequality uses the existing estimation rates from Pooladian and Niles- Weed [2021], with our parameter choice for Eo.
Proposition 4 (Stability of entropic maps with variations in the source measure). Let h = 111 112 Let p. pL, p be probability measures over a compact domain with radius R. Suppose Te, T are, respectively, the entropic maps from μ p and y¹ p. both with the parameter 0. Then
I[. - TTll S 333
4 Computing Couplings and Map Estimators with PROGOT
Following the presentation and motivation of PROGOT in Section 3, here we outline a practical implementation. Recall that f., Ei= 2,0x, and 6m 1== bydy,, and we summarize the locations of these measures the matrices X= (x....., x.,) and size n x d and m x d. respectively. Our PROGOT solver, concretely summarized in Algorithm 2, takes as input two weighted point clouds, step-lengths (ax)a, regularization parameters (E)),, and threshold parameters (Tk)a J., to output two objects of interest: the final coupling matrix P of size X m. as illustrated in Figure 2, and the entities needed instantiate the TProg map estimator, where an implementation is detailed in Algorithm 3. We highlight that Algorithm 2 incorporates warm-starti surm-starting method when instantiating Sinkhorn solvers (Line 3). This step may be added improve the runtime.
(y1....: which are of
Algorithm 2 PROGOT(a, X, Y, (Ek, Ork, Th)s)
Algorithm Jprog [b, Y, (g( Sks ag)a
1: f 0 g (-1) 0m. 2: for k ===00 0. K do
input: Source point Rd 2: initialize: y = x, ag reset to 1.
3:
fiits gini < (1 - ax) f. ox (k- f(g()) ,P +- Sink(a, X, Y, Y, Thi finits gimi))
for k
0, ... K do
4:
4:
p [b, exp( x -My }]
5:
Q + [iag(1/P1,,) P
5:
p f p/1 E Im
6:
Z [h' (E, Qu Vh(xi - yz))le E Raxd
6:
A Vhly y,)1 Rmxd
7:
X X au
7:
Z + Vh"(p A) E Rdd
8: end for
y + y Ozz.
9: return: Coupling matrix P.
9: end for
10: Map estimator Jrro (b, Y, (g()) Eb, ax)
10: return:
Setting step lengths. We propose three scheduling schemes for (a+)x: decelerated, constant-speed and accelerated. Let l E [0, 1] denote the progress along the interpolation path at iterate k. At step zero, to cgp Then at the next step, we progress by fraction @1 of the remainder, and therefore I = to + (1 (1- - Okg). It is straightforward to show that 4 = I_ (1 - oz)). We call a schedule constant speed, if f t5 is constant function of k whereas accelerated (resp. decelerated)

---

schedule has fx++ t increasing (resp. decreasing) with k. Table 6 presents the specific choices of Og for each of these schedules. By convention, we set the last step to be a complete step, i.e. ak 1.
Setting regularization schedule. To set the regularization parameters (ex (Cx))=>> we propose Al- gorithm 4. To set Eo. we use the average of the values in the cost matrix [h(x; - yj) ls between source and target, multiplied by small factor, as implemented in [Cuturi et al., 2022b]. Then for EK, we make the following key observation. As the last iteration, the algorithm is computing e(,), an entropic map roughly between the target measure and itself. For this problem, we know trivially that the OT map should be the identity. Therefore, given set of values to choose from, we pick ex to be that which minimizes this error over hold-out evaluation set of Y test = (8,)
error(s; Yerstt) = Σ
|s, -000E000 (3))
Algorithm e-scheduler Y lest (8p) 30)
1: recover a, a, b, Y, (ax)k.
The intermediate values are then set by interpolating between ig and εk, according to the times
2: set En
1
Ei h(x; 8₁)-
am
3: set a + 20 m
1
Σ 12 h(y: - 8z).
4.- Figure 4-C visualizes the effect of applying Algorithm 4 for scheduling, as opposed to choosing default values for Ek.
4: for p = L...., P do
5:
+ 8p J.
6:
g + SINK(b, Y,b, E,,))
Setting threshold schedule. By setting the Sinkhorn stopping threshold Ty as a function of time k; one can modulate the amount of compute effort spent by the Sinkhorn subroutine at each step. This can be achieved by decreasing тę linearly w.r.t. the iteration number, from a loose initial value, e.g., 0.1, to final target value TK < 1. Doing so results naturally in sub-optimal couplings P and dual variables (k)) each step k. which might hurt performance. However, two comments are in order: (i) Because the last threshold TK can be set independently, the final coupling matrix P returned by PROGOT can be arbitrarily feasible, in the sense that its marginals can be made arbitrarily close to a, b by setting Tk to a small value. This makes possible to compare in a fair way to direct application of the Sinkhorn algorithm. (ii) Because the coupling is normalized by its own marginal in line (5) of Algorithm 2, we ensure that the barycentric projection computed at each step remains valid, i.e., the matrix Q is a transition kernel, with line vectors in the probability simplex.
7:
T = Prog bb Y. (g. 1)
8:
errorp +- Y T(Y test ||
9: end for
10: p" arg min, errorp.
11: 51 f Spt X 0.
12: (4&); (1 I1=_ (1 oz))
13: return: ((1 1xx 1,00 + 12E1)K.
Experiments
We run experiments to evaluate the performance of PROGOT across various datasets, on its ability to act as map estimator, and to produce couplings between the source and target points. The code for PROGOT, is included in the OTT-JAX package [Cuturi et al., 2022b].
5.1 PROGOT as a Map Estimator
In map experiments, unless mentioned otherwise, we run PROGOT for K = 16 steps, with a constant-speed schedule for ag, and the regularization schedule set via Algorithm 4 with 3n 5 and In these experiments, we fit the estimators on training data using the & transport cost, and report their performance on test data in Figure 4 and Table 1. To this end, we quantify the distance between two test point clouds (X, Y) with the Sinkhorn divergence [Genevay et al., 2018, Feydy et al., 2019], always using the l transport cost. Writing OT, (X, Y) for the objective value of Equation (5), the Sinkhorn divergence reads
sp E (2-33 ..... 23),
Den (X, Y) = OT (X, Y) (OT,,, X) + OT (Y. YY)
(8)
where dp is 5% of the mean (intra) cost seen within the target distribution (see Appendix B).
Exploratory Experiments on Synthetic Data, We consider a synthetic dataset where X d- dimensional point cloud sampled from a 3-component Gaussian mixture. The ground-truth To is the gradient of an input convex neural network (ICNN) previously fitted to push roughly X to a mixture of 10 Gaussians [Korotin et al., 2021]. From this map, we generate the target point cloud Y. Unless stated otherwise, we use Merain = 7000 samples to train a progressive map between the source and target point clouds and visualize some of its properties in Figure 4 using fhest 500 test points.
7

---

(A) Estimator Consistency
(B) Seiedhling oz
Scheduling ey
0.04
Entropic Prng(T
Cmettat Develrnate
Prgli Prggii
MSE
6.03
Stale Div
100-
Accdenute
Stak Dre
6.02
300
20-
0.01
10
11
16
mumber of samples (Reests)
items
(k)
1t0rs
Figure 4: (A) Convergence of "J Prog to the ground-truth map w.r.t. the empirical 1.2 norm, for d = 4.
(B) Effect of scheduling ogs, for d = 64. (C) Effect of scheduling Ek using Algorithm 4, for d = 64.
Figure 4-(A) demonstrates the convergence of JProy to the true map as the number of training points grows, in empirical L2 norm, that is, MSE Reees Eist [To(x,, - JProg (x,) Figure 4-(B) She
shows the progression of PROGOT from source to target as measured by DED (X()), Y) where X are the intermediate point clouds corresponding to ((*)) The curves reflect the speed of movement for three different schedules, e.g., the decelerated algorithm takes larger steps in the first iterations, resulting in DED (x()), Y) to initially drop rapidly. Across multiple evaluations, we observe that the ae schedule has little impact on performance and settle on the constant-sped sppeed schedule. Lastly, Figure 4-(C) plots DED (x(k) Y) for the last 6 steps of the progressive algorithm under two scenarios. PROGOT uses regularization parameters set according to Algorithm 4, and PROGOT without scheduling. sets every E as 5% of the mean of the cost matrix between the point clouds of (X Y). This experiment shows that Algorithm 4 can result in displacements ((k) that are "closer" to the target Y. potentially improving the overall performance.
Comparing Map Estimators on Single-Cell Data. We consider the sci-Plex single-cell RNA sequencing data from [Srivatsan et al., 2020] which contains the responses of cancer cell lines to 188 drug perturbations, as reflected in their gene expression. Visualized in Figure 8, we focus on 5 drugs (Belinostat, Dacinostat, Givinostat, Hesperadin, and Quisinostat) which have a significant impact on the cell population as reported by Srivatsan et al. [2020]. We remove genes which appear in less than 20 cells, and discard cells which have an incomplete gene expression of less than 20 genes, obtaining n z 10 source and m = 500 target cells, depending on the drug. We whiten the data, take it to log(1 + z) scale and apply PCA reduce the dimensionality d {16, 64, 256). This procedure repeats the pre-processing steps of Cuturi et al. [2023].
We consider four baselines: (1) training an input convex neural network (ICNN) [Amos et al., 2017] using the objective of Amos [2022] (2) training a feed-forward neural network regularized with the Monge Gap [Uscidda and Cuturi, 2023], (3) instantiating the entropic map estimator [Pooladian and Biles-Weed, 2021] and (4) its debiased variant [Feydy et al., 2019, Pooladian et al., 2022]. The first two algorithms use neural networks, and we follow pppre-parameter tuning in [Uscidda and Cuturi, 2023]. We choose the number of hidden layers for both as [128, 64, 64]. For the ICNN we use learning rate 17 = 10-3 3 batch size b = 256 and train using the Adam optimizer [Kingma and Ba, 2014] for 2000 iterations. For the Monge Gap we set the regularization constant AMG 10, Acous = 0.1 and the Sinkhorn regularization to E 0.01. We train the Monge Gap setting, except that we set 1 = 0.01. To choose & for entropic estimators, we split the training data to get an evaluation set and perform 5-fold cross-validation on the grid of {2 ........ 23) x Eg, where eo is computed as in line 2 of Algorithm 4.
similar
We compare the algorithms by their ability to align the population of control cells, cells treated with a drug. We randomly split the data into 80%- 20% train and test sets, and report the mean and standard error of performance over the test set, for an average of 5 runs. Detailed in Table 1, PROGOT outperforms the baselines consistently with respect to Drp (TT # X, Y). The table shows complete results for 3 drugs, and the overall ranking based on performance across all drugs. Table 5 presents the synthetic counterpart to Table 1, using Gaussian Mixture data for d = 128, 256.
5.2 PROGOT as a Coupling Solver
In this section, we benchmark the ability of PROGOT to return a coupling, and compare it to that of the Sinkhorn algorithm. Comparing coupling solvers is rife with challenges, as their time performance must be compared comprehensively by taking into account three crucial metrics: (i) the cost of transport according to the coupling P, that is, (P, C)2 (ū) the entropy F(P),, ) and (ii) satisfaction of
marginal constraints P1mm all1 + I
bll1. Due to our threshold schedule (Tx (Tr)., as detailed
8

---

Table 1: Performance of PROGOT compared baselines, w.r.t Den between source and target of the sci cii-Plex dataset. Reported numbers are the average of 5 runs, together with the standard error.
Drug
Belinostat
Givinostat
Hesperudin
5-drug
dPCA
16
64
256
16
64
256
16
64
256
rank
PROGOT
2.9=0.1
8.8+0.1
20.8+0.2
3.3=0.2
9.0=0.3
21.9+003
3.7+00.4
10.1+00.4
23.1=0.4
EOT
2.5.1.11
9.610.1
22,910.233,99000
10.011..1
24.7:0.91
4.110.4..1..44
10.4.10.5
26111.3
Debiased FOT
3.2+0.11
14.3+0.1
39.8+0.4
337--0.2
14.74-0.1
42.4+0088
4.0+0..5
15.2+0.6
41+1.1
4
Monge Gup
3.1=0.1
10.3=0.1
3..4.0.3
2.8-0.2
9.9-0.2
34,910.3
3.7.0..5
11.0:0.5
36111.1
3
ICNN
5.0
14.74 0.1
42+1
5.1 0.1 0.2
40.3+0.1
4.0+0.4
1..410.5
46+2.1
5
displatin
desalinib
everolinans
ProgOT
ProgOT
ProgUI
Sinkioom
Sinkhoen
9.5
Sinkhorn
Fotropy
9.6
8.6-
9.4
9.4
9.2-
9.0
2.50
2,555
2,560
2,570
2,575
9.0
12/92
12.93 12.94
9.0
4.55
4.05
Cost
Cost
Coet
Vitnleeine
staurosporize
decitalinn
9.6
10.4
9.6
10.2
W
Eatoop
9.44
10.0
M.4
9.8
9.3-
9.3
9.1-
9:1
9.0
4,001
4.006
4.00s
4.010
4.014
4.016
0.0
18.68
18:69
18
18.71
18.72 Creit
18.73
18.74
9.0
2.7.6
2.718
2.720
2,720 Cost
2.721
2.726
Cost
- 0.0
B 0.006
3 =0003
0.002
-2
Figure 5: Performance as a coupling solver on the 4i dataset. PROGOT returns better couplings, in terms of the OT cost and the entropy, for a fraction of Sinkhorn iterations, while still returning coupling that has the same deviation to the original marginals. The (top) row is computed using |1.12, the (bottom) row shows results for the cost h I1 where p 1.5.
in Section 4, both approaches are guaranteed to output couplings that satisfy the same threshold for criterion (iāi), leaving only three quantities to monitor: compute effort here quantified total number of Sinkhorn iterations, summed over all K steps for PROGOT), transport cost and entropy. While compute effort and transport cost should, ideally, be small as possible, certain applications requiring, e.g., differentiability [Cuturi et al., 2019] or better sample complexity [Genevay et al., 2019], may prefer higher entropies.
To monitor these three quantities, and cover interesting space of solutions that, we Sinkhorn's algorithm for a logarithmic grid of E = Beg values (here Eo is defined in Line of Algorithm 4), and compare it to constant-speed PROGOT with K f2, 4, 8}. Because ularizations, we explore many choices to schedule & within PROGOT. Following the default strategy used in OTT-JAX [Cuturi al., 2022a], mean of the cost matrix at that iteration, as detailed in Appendix B. We do not use Algorithm 4 since it returns regularization schedule that is tuned for map estimation, while the goal here is to recover couplings that are comparable to those outputted by Sinkhorn. We set the threshold for marginal constraint satisfaction for both algorithms as TK 7 0.001 and run all algorithms to convergence, with infinite iteration budget. For the coupling experiments, we use the single-cell multiplex data of Bunne al. [2023], reflecting morphological features and protein intensities of melanoma tumor cells. The data describes d = 47 features for n = 11, 000 control cells, and m z 2, 800 treated cells, for each of 34 drugs, of which we use only 6 random. To align the cell populations, we consider two ground costs: the squared- -uulid-n norm 1| as well as h ill lp. with y 1.5.
cannot directly compare reg-
at every iterate k. S = 0Cx, where O2 is 5% of the the
Results for drugs are displayed in Figure 5. The area of the marker reflects the total number of Sinkhorn iterations needed for either algorithm to converge to coupling with threshold - 10-3 The values for 3 and K displayed in the legend are encoded using colors. The global scaling parameter for PROGOT is set to 0 = 2 Figure 10 and 11 visualize other choices for 0. These results prove that PROGOT provides a competitive alternative Sinkhorn, to compute couplings that yield a small entropy and cost at low computational effort, while satisfying the same level of marginal constraints.

---

Figure 6: We consider the optimal assignment problem Table 2: Coupling recovery, quantified between all CIFAR images and their blurry CIFAR as trace, and KL divergence from counterparts using the lg loss. A small subset of 3 original images on the left can be compared with their blurred counterpart on the right, with = 4. The optimal coupling for this task is the identity, which we compare with couplings recovered by our methods at large scales.
identity matrix, for coupling matrices obtained with PROGOT and Sinkhorn, and blur strengths a 2, 4. PROGOT is run for K = 4 and with the constant- speed schedule.
2
4
p
Tr
0.9999
0.9954
Sinkhorn
KL
0.00008
0.02724
iters
10
2379
TT
1.000
0.9989
PROGOT
KL
0.00000
0.00219
iters
40
1590
5.3 Scalability of PROGOT
Real-world experiments on pre- processed single-cell data are often run with limited sample sizes (n 109) and medium data dimensionality (d S. 200). As a result they are not suitable benchmark OT solvers at larger scale. To address this limitation, we design a challenging large-scale (large n, large d) OT problem on real data, for which the ground-truth
known. We believe our approach can be replicated to create benchmarks for OT solvers. We consider the entire grayscale CIFAR10 dataset [Krizhevsky et al., 2009] for which n 60, 000 and d 32 x 32 1024. We consider the task of matching these n images to their blurred counterparts, using Gaussian blurs of varying width. To blur an image U E I X we use the isotropic Gaussian kernel K [exp (-(i - N²)) for j.j < N fc.f. Remark 4.17, Peyré et al., 2019], and define the Gaussian blur operator as G(U)= = KUK E NNX,, The crucial observation we make in Proposition is that, when using the squared Euclidean ground cost LZ, the optimal matching is necessarily equal to the identity (i.e. each image must be necessarily matched to its blurred counterpart), as pictured in (Figure 6).
Proposition 5. Lel î== h E. Sü. be the empirical distribution over n images and define i == Guj where G is the Gayssian blur operator with a < co. Then P the optimal coupling between (jt, 1) with the h = 1 11.112 cost is the normalized n-dimensional identity matrix Id /n.
In light of Proposition 5, we generate two blurred datasets using Gaussian kernel with o = 2 and = (see Figure 7). We then use PROGOT with and Sinkhorn's Algorithm to match the blurred dataset back to the original CIFARIO (de-blurring). The hyper-parameter configurations are the same = 2-4, We evaluate the performance of the OT solvers by checking how close the trace of the recovered coupling Tr(P) is to 1.0, or with the KL divergence from the ground-truth, that is, (TI(P)|P)) = log - n Eisn log( P.)).
as Section 5.2, with B
Table 2 compares the performance of PROGOT and Sinkhorn, along with the number of iterations needed to achieve this performance. Both algorithms scale well and show high accuracy, while requiring similar amount of computation. We highlight that this scale, simply storing the or coupling matrices would require about 30Gb. The experiment has to happen across multiple GPUs. Thanks to its integration in JAX and OTT-JAX [Cuturi et al., 2022b], PROGOT supports sharding by simply changing few lines of code. The algorithms scales seamlessly and each run takes about 15 minutes, on a single node of 8 A100 GPUs This experiment sets convincing example on how PROGOT scales to much larger (in n and d) problems than considered previously.
Conclusion
In this work, we proposed PROGOT, new family of EOT solvers that blend dynamic and static formu- lations of OT by using the Sinkhorn algorithm as subroutine within progressive scheme. PROGOT aims to provide practitioners with an alternative to the Sinkhorn algorithm that (i) does not fail when instantiated with uninformed or ill-informed regularization, thanks to its self-correcting behavior and our simple e-scheduling scheme that is informed by the dispersion of the larget distribution, (ii) performs at least as fast as Sinkhorn when used to compute couplings between point clouds, and (ii)) provides a reliable out-of-the-box OT map estimator that comes with a non-asymptotic convergence guarantee. We believe PROGOT can be used as strong baseline to estimate Monge maps.
10

---

References
J. Altschuler, J. Weed, and P. Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, 2017.
L. Ambrosio, N. Gigli, and G. Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
B. Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations, 2022.
B. Amos, L. Xu, and Z. Kolter. Input convex neural networks. Proceedings of the 34th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2017.
J.-D. Benamou and Y. Brenier. A computational fluid mechanics solution to the Mongc-Kantorovia mass transfer problem. Numerische Mathemarik, 2000.
Y. Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Comm. Pure Appl. Math., 1991.
C. Bunne, G. Stark, G. Gut, J. S. del Castillo, M. Levesque, K.-V Lehmann, Pelkmans, A. Krause, and G. Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature Methods, 2023.
G. Buttazzo, De Pascale, and P. Gori-Giorgi. Optimal-transport formulation of electronic density- functional theory. Phys. Rev. A, 2012.
G. Carlier, L. Chizat, and M. Laborde. Displacement smoothness of entropic optimal transport. ESAIM: COCV, 2024.
S. Chewi and A.-A. Pooladian. An entropic generalization of Caffarelli's contraction theorem via covariance inequalities. Comptes Rendus. Mathématique, 2023.
I. Csiszár. L-divergence geometry of probability distributions and minimization problems. Ann. Probability, 3:146-158, 1975.
M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26, 2013.
M. Cuturi, O. Teboul, and J.-P. Vert. Differentiable ranking and sorting using optimal transport. Advances in neural information processing systems, 32, 2019.
M. Cuturi, L. tools (ott): A JAX toolbox for all things Wasserstein. arXiv preprint, 2022a.
Y. Tian, C. Bunne, G. Davis, and O. Teboul. Optimal transport
M. Cuturi, L. tools (OTT): A JAX toolbox for all things Wasserstein. CoRR, 2022b.
Y. Tian, C. Bunne, G. Davis, and O. Teboul. Optimal transport
M. Cuturi, M. Klein, and P. Ablin. Monge, Bregman and occam: Interpretable optimal transport in high-dimensions with feature-sparse maps. In Proceedings of the 40th International Conference Machine Learning, Proceedings of Machine Learning Research. PMLR, 2023.
V. Divol, J. Niles Weed, and A.-A. Pooladian. Tight stability bounds for entropic Brenier maps. arXiv preprint, 2024.
S. Eckstein and M. Nutz. Quantitative stability of regularized optimal transport and convergence of Sinkhorn's 's algorithm. SIAM Journal Mathematical Analysis, 2022.
J. Feydy. Geometric data analysis, beyond convolutions. Applied Mathematics, 2020.
J. Feydy, T. Séjourné, F.-X. Vialard, S.-i. Amari, A. Trouvé, and G. Peyré. Interpolating between optimal transport and MMD using Sinkhorn divergences. The 22nd International Conference on Artificial Intelligence and Statistics. PMLR. 2019.

---

R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, L. Gautheron, N. T. Gayraud, H. Janati, A. Rakotomamoniy, Redko, A. Rolet, A. Schutz, V. Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 2021.
N. Fournier and A. Guillin, On the rate of convergence in Wasserstein distance of the empirical measure. Probability theory and related fields, 162(3):707-738, 2015.
A. Genevay, Entropy-regulaizeed optimal transport for machine learning. PhD thesis, Paris Sciences et Lettres (ComUE), 2019.
A. Genevay, G. Peyré, and M. Cuturi. Learning generative models with Sinkhorn divergences. In Proceedings of the 21st International Conference on Artificial Intelligence and Statistics, 2018.
A. Genevay, L. Chizat, F. Bach, M. Cuturi, and G. Peyré. Sample complexity of sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pages 1574-1583, PMLR, 2019.
Ghosal, M. Nutz, and E. Bernton. Stability of entropic optimal transport and Schrödinger bridges. Journal of Functional Analysis, 283(9):109622, 2022.
G. Gut, M. D. Herrmann, and L.. Pelkmans. Multiplexed protein maps link subcellular organization to cellular states. Science, 2018.
J.-C. Hütter and P. Rigollet. Minimax estimation of smooth optimal transport maps. The Annals of Statistics, 2021.
L. Kantorovitch. On the translocation of masses. C. R. R. (Doklady) Acad. Sei. URSS (N.S.), 1942.
D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412:6442 2014.
A. Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do neural optimal transport solvers work? A continuous Wasserstein-2 benchmark. Advances in neural information processing systems, 2021.
A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Lehmann, M.-K. von Renesse, A. Sambale, and A. Uschmajew. A note on overrelaxation in the Sinkhorn algorithm. Optimization Letters, 2022.
T. Lin, N. Ho, and M. I. Jordan. On the efficiency of entropic regularized algorithms for optimal transport. Journal of Machine Learning Research, 2022.
Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. arXiv preprint, 2022.
Q. Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint, 2022.
T. Manole, S. Balakrishnan, J. Niles Weed, and L. Wasserman. Plugin estimation of smooth optimal transport maps. arXiv preprint, 2021.
J. J. McCann. convexity principle for interacting gases. Advances in mathematics, 128(1):153-11) 1997.
G. Mena and Niles-Weed. Statistical bounds for entropic optimal transport:
sample complexity and the central limit theorem. Advances in neural information processing systems, 2019.
I. Métivier, R. Brossier, Q. Mérigot, F. Oudet, and J. Virieux. Measuring the misfit between seismo- grams using an optimal transport distance: Application to full waveform inversion. Geophysical Supplements to the Monthly Notices of the Royal Astronomical Society, 2016.
G. Monge. Mémoire sur théorie des déblais et des remblais. Histoire de l'Académie Royale des Sciences, 1781.
12

---

G. Peyré, M. Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 2019.
A.-A. Pooladian and J. Niles-Weed. Entropic estimation of optimal transport maps. arXiv preprint, 2021.
A.-A. Pooladian, M. Cuturi, and Niles-Weed Debiaser beware: Pitfalls of centering regularized transport maps. In International Conference on Machine Learning. PMLR, 2022.
A.-A Pooladian, H. Ham-H C. Domingo-Enrich, B. Amos, Y. Lipman, and R. T. Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In Proceedings of the 40th International Conference on Machine Learning. Proceedings of Machine Learning Research. PMLR, 2023.
A. Ramdas, N. García Trillos, and M. Cuturi. On Wasserstein two-sample testing and related families of nonparametric tests. Entropy, 2017.
F. Santambrogio. Optimal transport for applied mathematicianss Springer, 2015.
M. Scethon, M. Cuturi, and G. Peyré. Low-rank Sinkhorn factorization. In Proceedings of the 38th International Conference on Machine Learning. Proceedings of Machine Learning Research, 2021.
M. Scetbon, G. Peyré, and M. Cuturi. Linear-time Gromov Wasserstein distances using low rank couplings and costs. In Proceedings of the 39th International Conference on Machine Learning. PMLR, 2022.
G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, Berube, al. Optimal-transport analysis single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 2019.
B. Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing, 2019.
E. Schrödinger. Über die unkehrung der naturgesetze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u. ..... 1931.
R. Sinkhom. A relationship between arbitrary positive matrices and doubly stochastic matrices. The Annals of mathematical statistics, 1964.
S. Srivatsan, L. McFaline-Figuero. V. Ramani, Saunders, Cao, Packer, H. A. Pliner,
D. L. Jackson, R. M. Daza, L. Christiansen, et al. Massively multiplex chemical transcriptomics at single-cell resolution. Science, 2020.
A. Thibault, L. Chizat, C. Dossal, and N. Papadakis. Overrelaxed Sinkhom-Knopp algorithm for regularized optimal transport. Algorithms, 2021.
Thornton and M. Cuturi. Rethinking initialization of the sinkhorn algorithm. Proceedings of The 26th International Conference Artificial Intelligence and Statistics. PMLR, 2023.
A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, K. Fatras, G. Wolf, and Y. Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint, 2023.
Uscidda and M. Cuturi. The Monge gap: regularizer to learn all transport maps. International Conference Machine Leurning. PMLR, 2023.
A. Vacher and F.X. Vialard. Parameter tuning and model selection in optimal transport with semi-dual Brenier formulation. Advances in Neural Information Processing Systems, 2022.
H. Van Assel, T. Vayer, R. Flamary, and N. Courty. Optimal transport with adaptive regularisation. In NeurlPS 2023 Workshop Optimal Transport and Machine Learning, 2023.
C. Villani et al. Optimal transport: old and new. Springer, 2009.
Y. Xie, X. Wang, R. Wang, and H. Zha. A fast proximal point method for computing exact Wasserstein distance. In Proceedings of The 35th Uncertainty Artificial Intelligence Conference, Proceedings of Machine Learning Research. PMLR, 2020.
13

---

A Additional Experiments
In Section 5.1, we demonstrated the performance of PROGOT for map estimation on the sci-Plex dataset. Here, we present similar experiment the 4i data, and extend all map experiments the case of a general translation invariant cost function, h 1 1.5/|-1l -ll... In Table 3 and Table 4 we show DED ((TProg) # X, h) the sinkhorn divergence using the cost function h.
Table 3: Performance of algorithms on sci-Plex data, w.r.t Den (( # X, Y; h) with the 1.5-norm cost. Reported numbers are the average of runs, together with the standard emror.
Drug
Belinostat
Givinostat
Hesperadin
5-drug
decA
16
64
256
16
64
256
16
64
256
rank
PROGOT
2.03+
7.11 F0.03
08.6
2.04
0.07
7.1
19.5+0.1
2.4
7.7
20.2+0..4
BOT
2.00
0.22.1006
18.810.2
0.0
7.110.1
19.5100.
2.6.1..2
8.1.
0.66..0.6
2
Debiased BOT
3.90=00.04
13.8+0.1
37.6+0.2
4.2=0.1
14.4+0.1
38.7+0022
3.6+0.2
13.0=0.2
06.0++0.5
4
Monge Gap 2.4+0.1
8.6+0.11
34.2+00.3
2.3=0.1
8.5+0.2
33.8++0.3
3.7+0.55
00.4=0.5
36.0++0.8
3
Table 4: Performance algorithms on 4i data, w.r.t DED ((TPro) ) X, Y; where is reported in the table. Reported numbers are the average of 10 runs, logether with the standard error.
Drug/
cisplatin
dasatinib everolimus
vindesine
staurosporine
decitabine
overall
Cost
E2
82
8g
1.51|-1|
1.5||-| 15
1.5||-11
rank
PROGOT
1.6.10.004
2.710.1
1.66_0.06
2.2110.03
2.7410.08
1.6610.01
BOT
2.72:10.05
2.17:00033
4.16-0.11
2.91100.09
1.620.04
1.7310.04
3=
Debiased FOT
1.79++0.07
1.65+00.0
2.98
0.29
2.86:::0.25
1.81
1.64
Monge Gup
1.7+0.1
3.3
1.1=0.1
2.17::0.05
3.03=0008
1..05
3=
ICNN
1.74.10.08
3.810.77
1.88 0.05
GMM Benchmark. To provide further evidence on performance of PROGOT, we benchmark the map estimation methods on nigh-dimensional Gaussian Mixture data, using the dataset of Korotin al. [2021] for 128, 256. In this experiment the ground truth maps are known and allow us to compare the algorithms more rigorously using the empirical εg distance of the maps, as defined in section Section 5. Shown in Table 5, we consider Thest = 500 test points and 7ltrain = 8000 and 9000 training points, respectively for each dimension. We have also included a variant of the entropic estimator which uses the default value the OTT-JAX library for e, and unlike other EOT algorithms,, is not cross-validated. This is to demonstrate the significant effect that ε has Sinkhorn solvers.
Table 5: GMM benchmark. The Table shows the MSE, average of || Ytest 12 for Frtsst 500
points, where y = T (Xtest and the ground truth is Ytest-
d 128
= 256
PROGOT
0.099-000009
0.12:0000
EOT
0.12.00.01
0..62
Debiased BOT
0.11+:0.01
0.128-10.00
Untuned EOT
0.250=0.0023
0.276=0.006
Monge Gap
0.36+:0.02
0..273+0.005
ICNN
0.177+00023
0.117#0.005
B Details of Experiments
Generation of Figure 1 and Figure 2. We consider toy example where the target and source clouds are as shown in Figure 1. We visualize the entropic map [Pooladian and Niles-Weed, 2021] its debiased variant [Pooladian al., 2022] and PROGOT, where we consider decelerated schedule with 6 steps, and only visualize steps k = 3, 5 to avoid clutter. The hyperparameters of the algorithms are set as described in Section Figure shows the coupling matrix corresponding to the same data, resulting from the same solvers.
Sinkhorn Divergenee and its Regularization Parameter. In some of the experiments, we calculate the Sinkhorn divergence between two point clouds as measure of distance. In all experiments we set the value of Ep and according to the geometry of the target point cloud, In particular, we set E to be default value of the OTT-JAX Cuturi et al. [2022b] library for this point cloud via
14

---

Original: dog
Blurred: σ 2
Blurred: o
Figure Example of a CIFARIO image and blurred variant. We match blurry images to the originals.
Delinootat
Ducinostat
Civinostat
Hesperadin
Quisinnctat
2-
2-
-2
Phoas
Shor
Figure 8: Overview of the single cell dataset Srivatsan et al. [2020]. We show the first two PCA dimensions performed on the training data, and limit the figure 50 samples. The point cloud (Xuain Xtest) shows the control cells and (Ytrain Vrest) are the perturbed cells using a specific drug.
ott. Stt.geometry ceeneetry-poiintee www.tooood..d..o.. Pointlood (Y) epsilon, that is, 5% of the average cost matrix, within the target points.
Details of Scheduling (oo)). Table 6 specifies our choices of for the three schedules detailed in Section 4. Figure 9 compares the performance of PROGOT with constant-speed schedule in red, with the decelerated (Dec.) schedules in green. The figure shows results on the sci-Plex data (averaged across 5 drugs) and the Gaussian Mixture synthetic data. We observe that the algorithms perform roughly on par, implying that in practice PROGOT is robust to choice of a.
Table 6: Scheduling Schemes for ag- Or = a(k)
Schedule
t = (kk
1k+1 11+11
Decelerated
1/e
Constant-Speed
K
Accelerated
2k
(k)
(K+1)*-(0 -k)
Details of Scheduling (Ex) For map experiments on the sci-Plex data [Srivatsan et al., 2020], schedule the regularization parameters via Algorithm 4. We set f% 5 and consider the set 8p = (0.1,0.5, 1, 5, 10, 20}. For coupling experiments on the 4i data [Gut et al., 2018] we set the regularizations (according to Line 7, Algorithm 2) and recall that y .I....,y is the target point cloud. The scaled average cost this iterate č* Ei, li(x!!) -y;)/(20mn)) which is the default value of typically used for running Sinkhorn. Then for every l E [K], we set 6 = 0cg to make PROGOT compatible the 3 regularization levels of the bench-marked Sinkhorn algorithms. For Figure we have set 0 = 2 44 In Figure 10 and Figure 11, we visualize the results for 0 E {2-7.2 2,2 2--1- to give an overview of the results using small and larger scaling values.
follows. Let x; ,...;X, x,,) denote the interpolated point cloud at iterate
Compute Resources. Experiments were run on single Nvidia 100 GPU for a total of 24 hours. Smaller experiments and debugging was performed on single MacBook M2 Max.
15

---

GMM data
Sciple drage aggregpted
0.24
0.30
10-
0.28
000
0.20
04
206
16
216
PrufTT avedard
DrngOT
Or
PruOT w/o mazhee (Dou.)
PTgOT [Drc.]
Prg [Dec.]
Figure 9: Comparison of constant speed vs decelerated PROGOT.
cisplatin, 0 =0.008
dasatinib, 0 =0.008
everolimus, =0.008
Entropy
9.8
9.8
9.8
9.4
9.4
9.4
9.0
9.0
9.0
2.535
2.545
2.555
2.565
12.88
12.90
12.92
12.94
4.645
4.655
4.663
4.675
cisplatin, 0 =0.062
dasatinib, 0
=0.062
everolimus, =0.062
Entropy
9.8
Ite.
9.8
9.8
9.4
9.4
9.4
9.0
9.0
2.535
2.545
2.555
2.565
12.88
12.90
12:92
12.94
4.645
4.655
4.663
4.675
cisplatin, 0 =055
dasatinib, =0.5
everolimus, 0 =0.5
15.5
15.5
15.5
Entropp
14.5
ma
14.5
13.5
13.5
13.5
12.5
12.5
12.5
11.5
3.8
5.0
5.4
14.0
15.6
6.0 Cost
6.4 6.8
Cost
Cost
8 =0.008
=0.031
=0.125
B =0.5
=2.0
K
8 0.006
3 0.002
8 0.25
5 =10
K 2
K
Figure 10: Comparison of PROGOT and Sinkhorn as coupling solvers for h(.-) || 112 on the 4i dataset. Rows show different choices of regularization 0 for PROGOT detailed in Appendix B.
C Proofs
Preliminaries. Before proceeding with the proofs, we collect some basic definitions and facts.
First, we write the the y-W y-Wasserstein distance any p 2
wp(u,v) = (- inf (rET((u))))
// || -/ (llx-yl"dTTz,))
Moreover, it well-known that p- Wassersteiin distances ordered for p 2 for p S q. holds that W,(n, v) S W, (n(,, v) [ef. Remark 6.6, Villani et al., 2009].
For the special case of the 1- L-Wasserstein distance, have the following dual formulation
/ [a(u--) v).
W WH, sup fetip:
where Lip, is the space of 1-Lipschitz functions [cf. Theorem 5.10, Villani et al., 2009].
Returning to the 2- W-assersteiin distance, we will repeatedly use the following two properties of optimal transport maps. First, for any two measures μl, y and an L-Lipschitz map T, it holds that
(T Twv) S (W, v).
This follows from a coupling argument. similar vein, we will use the following upper bound on the Wasserstein distance between the pushforward of source measure y by two different optimal transport maps T, and T3:
((T) 4H, (Tb) I TTa Tbll2
(10)
16

---

vindesine, =0.008
staurosporine, =0.008
decitabine, =0.008
10.0
10.0
10.0
Entropy
Hess
9.0
9.0
9.0
4.006 4.010 4.014 4.018
18.675 18.685
18.695
18.705
2.718
2.722
2.726
2.730
vindesine,
=0.062
staurosporine,
=0.062
decitabine, =0.062
10.0
10.0
10.0
Eutropy
9.0
9.0
9.0
4.006
4.010 4.014 4.018
8 8.675 18.685 18.095 18.705
2.718
2.722
2.726
2.730
vindesine, 0 =0.5
slaurosporine, 0 =0.5
decitabine, 0 =0.5
15
15
15
Entropy
33
13
13
4.2
5.0
5.4
18.8
19.2 Cost
19.6
20.0
3.0
3.4 Cost.
Cost
A =0.008
3 =0.031
8
0.122
B =0.5
8 =2.0
K =4
-0.016
3 =0.062
8 =0.25
5 =1.0
K =2
K
Figure 11: Comparison of PROGOT and Sinkhorn coupling solvers for h(-) || |1|:/11.5 on the dataset. Rows show different choices regularization 0 for PROGOT as detailed in Appendix B.
Belino, =0.008
Givino, 8 =0.008
Hespera, =0.008
Entropy
9.6
9.6
9.6
10:00
Ipc
9.2
9.2
9.2
49.47
49.49
49.51
45.495
45.505 45.515 45.525
7.805
7.815
7.825
Belino, =0.062
Givino, 8 =0.062
Hospera, =0.062
Entrops
9.6
9.6
9.6
0.2
9.2
9.2
49.47
49.49
49.51
455
45.505 45.515 45.525
7.805
7.815
7.825
Belino, =0.5
Givino, =0.5
Hespera, 0 =055
Entropy
12.5
11.55
12.5
12.5
10.5
10.5
301
11.5-
100
0.5
9.5
8.5
8.5
10.5
49.6
50.4
45.5
46.5
8.5
9.5
10.5
Cost
Cost
Cost
K 2
=8
A =0.016
8 =0.002
8 =0.25
8 =100
4 B 0.008
-0.031
8 =0.5
-2.0
Figure 12: Comparison of PROGOT and Sinkhorn coupling solvers for h(-) |1 on the sci-Plex dataset. Rows show different choices regularization 0 for PROGOT as detailed in Appendix B.
17

---

Belino, 0 =0.008
Givino, 8 =0.008
Hespera, =0.000
10.0
Entropy
9.6
lias
+i
9.6
9.4
9.2
9.2
49.6
50.0
50.4
28
32
36
40
5.0
6.0
7.0
8.0
Belino, 0 =0.062
Givino. 8 =0.062
Hespera, =0.062
10.0
Entrog
9.6
9.6
9.4
9.2
9.2
50.4
36
6.0
Belino, 0 =0.5
Givino, 0 =0.5
Hespera, 0 =0.5
15
Entrope
13-
Iree
13
use
13
11
11
49.3
50.5
51.5
32.5 53.5
32
40
48
7.5
8.5
Cost
Cast
Cost.
K =2 K =8
=0.016
B =0.062
3 =0.25
=10
K
4
3 -00008
-0003
3 0..125
3 0.5
8 2.0
Figure 13: Comparison of PROGOT and Sinkhorn coupling solvers for h(-) = || |1 /1.5 on the sci-Plex dataset. Rows show different choices of regularization 0 for PROGOT as detailed in Appendix B.
Belino, 0 =0.008
Givino, 0 =0.008
Hespera, 0 =0.008
9.8
10.2
Entropy
9.6
9.8
9.6
9.4
9.2
9.2
9.0
40
44
28
32
48
5.145
5.155
5.165
5.175
Belino, 0 =0.022
Givino, 8 =0.062
Hespera, =0.062
9.8
10.2
Entropy
9.6
9.6
9.4
9.4
9.2
9.2
9.0
40
28
32
44 48 52
5.1
5.155
5.165
5.175
Belino, =0.5
Givino, =0.5
Hospera, 0 =0.5
12
Entropy
12
12
5XD
10
10
10
36
38
40
42
44
46 48
50
52
28
36
44
52
60
5.10
5.20 5.30 5.40 5.50 5.60
Cost
Cost
Cost.
K
200.166
0.006
3
-1100
K
0.0088
000031
00.125
3 0.5
=2.0
Figure 14: Comparison of PROGOT and Sinkhorn as coupling solvers for h(-) = |11:3/ /1.5 /1.5 on the sci-Plex dataset. Rows show different choices of regularization 0 for PROGOT detailed in Appendix B. The threshold for marginals is also scheduled here, starting from = 0.01 and reaching 0.001 to match Sinkhorn.
18

---

Notation conventions. For an integer K E N, = (O...,, We write S b to mean that there exists a constant C 0 such that a I Cb. A constant can depend on any of the quantities present in (A1) to (A3), as well as the support of the measures, and the number of iterations in Algorithm 2. The notation Slog(n) means that a I Ci (logg) n) for positive constants C1 and C2.
C.1 Properties of entropic maps
Before proving properties of the entropic map, first recall the generalized form of (5), which holds for arbitrary measures (cf. Genevay [2019]):
/ sdu+ gdv-=s sup ( ()))
/f ((())) -11== - -1 du( +e. (11)
When the entropic dual formulation admits maximixers, denote them by (fes 9e) and refer to them as optimal entropic Kantorovich potentials [e.g., Genevay, 2019, Theorem 7]. Such potentials always exist if μ and have compact support.
We can express an entropic approximation to the optimal transport coupling To as a function of the dual maximizers [Csiszár, 1975]:
(e 8) = 7/s (z, (I, == exp [+(x) + (v)
x u|l? du(r))d(v).
(12)
When necessary, will be explicit about the measures that give rise the entropic coupling. For example, in place of the above, we would write
(x, y) = "z," (2, v)d
(13)
The population counterpart to (7), the entropic map from is then expressed as
TH (x) == E,, YY| z].
and similarly the entropic map from μ
Ty) E.,, X =
We write the forward (resp. backward) entropic Brenier potentials as P2 = b11l 112 fe (resp. v/ + 1|2 ge). By dominated convergence, one can verify that
(x) Vile = Ty)"
We now collect some general properties of the entropic map, which we state over the ball but can be readily generalized.
Lemma 6. Let μ. be probability measures over B(O; R) in Rd. Then for a fixed >> 0, it holds that both T and T are Lipschitz with constant upper-boundee by R2 /E.
Proof of Lemma 6. We prove only the case for T as the proof for the other map is completely analogous. It is well-known that the Jacobian of the map is symmetric positive semi-definite matrix:
Vx² = e Coov, (Y|X- = x) (see e.g., Chewi and Pooladian [2023, Lemma 1J). Since the
probability measures are supported in a ball of radius R, it holds that sup, ||Cov=, (Y) =)]lop S R², which completes the claim.
We also require the following results from Divol et al. [2024], as well as the following object: for three measures P: μ., with finite second moments, write
i= [If log 26 v) i=
( - (I, (Z)))))))
"x ()((,,y)dTT...... ()))())))(
where æ is an optimal transport coupling for the Wasserstein distance between y and v, and 7: the density defined in (13).
19

---

Lemma 7. [Divol et al., 2024, Proposition 3.7 and Proposition 3.8] Suppose M4 have finite second moments, then
et ets ||(T0 (TE- -a) 2))
and
PM
S 27 sup ||Cov= (X|Y u)llop-
veRd
We are now ready to prove Proposition 4. We briefly note that stability of entropic maps and couplings has been investigated by many [c.g.. Ghosal et al., 2022, Eckstein and Nutz, 2022, Carlier al., 2024]. These works either present qualitative notions of stability, or give bounds that depend exponentially on l/e. In contrast, the recent work of Divol et al. [2024] proves that the entropic maps are Lipschitz with respect to variations of the target measure, where the underlying constant is linear in 1/s. We show that their result also encompasses variations in the source measure, which is of independent interest.
Let [(, n²) be the optimal transport coupling from μ to p². By disinte- grating and applying the triangle inequality, we have
Proof of Proposition 4.
(2)) // (px) -
(x)) - (
/// (TT) (x)- T.. (2 (2') (2')lldT(x,)))
+ || ( (z) T.. +0 (
(x)) Th. - (2 || | (I,
Re // ||: x' ||dm (2,
+ R² W2(u, u²),
1T (x) - (
where the penultimate inequality follows from Lemma 6, and the last step is due to Jensen's 's inequality. To bound the remaining term, recall that
sup HRR4
||Cov,g--<X|X = ollp S R²
and by the two inequalities in Lemma 7, we have (replacing with μ²)
[] (TT))(()) (a''""'') S 2TR
2R² (el)
2R²
||
- -T (2), - =)dT((,,2) 2) [TT - Ti
8
2R²
Tu' - ( x²) W2(,,)()
E
where we used Cauchy-Schware in the last line. An application of Jensen's inequality and rearranging results in the bound:
2R // [T)) -T (1d)ddd)) = E W2(u,,))
which completes the claim.
Finally, we require the following results from Pooladian and Weed [2021], which we restate for convenience but under our assumptions.
20

---

Lemma 8. [Two-sample bound: Pooladian and Niles-Weed, 2021. Theorem 3] Consider i..d. samples of size n from each distribution μ and v, resulting the empirical measures pa and v, with the corresponding. Let Te be the entropic map between it and D. Under (A1)-(A3), it holds that
E
T.
To (P())) Stor(n) E d n//2211/2 +
Moreover, i = =((n)= x n 11(444) then the overall rate of convergence is nt" -2/(6+44)
Lemma 9. /One-sample bound: Pooladian and Niles-Weed 2021, Theorem 4] Consider i...d. samples of size n from v, resulting in the empirical measure D, with full access a probability measure μ. Let Rg be the entropic map from μ to D. Under (AI)-(A3), it holds that
Stog(n) g¹ 1-1/22111222 522 Rs To (2(
Moreover, ife = E(n) = 1d42)) then the overall rate of convergence is 22/(2+)),
C.2 Remaining ingredients for the proof of Theorem 3
We start by analyzing our Progressive OT map estimator between the iterates. We will recurse these steps, and aggregate the total error at the end. We introduce some concepts and shorthand notations.
First, the ideal progressive Monge problem: Let 7 be the optimal transport map from ş to v, and write (0) p. Then write
((O) (1 ag) Id (0)
and consequently p()) (S(O) map from p1 to v, and consequently
# ( We iteratively define T be the optimal transport
(0)
(6) (1 a;) Id +a,T())
and thus p (S())) #(() The definition of McCann interpolation implies that these iterates all lie on the geodesic between μ and v. This ideal progressive Monge problem preciscly mimicks our progressive map estimator, though (1) these quantities are defined at the population level, and (2) the maps are defined solutions to the Monge problem, rather than its entropic analogue. Recall that we write js and as the empirical measures associated with and v, and recursively define g( to be the entropic map from fis)) to I where jag (0) Úú and
A ((+1))== ) (11aa) (a, (()
(14)
We also require f)( defined to be the the entropic map between (9 and e using regularization ex. This map can also be seen ""nne-sappl" estimator, which starts from iterates of the McCann interpolation, and maps to an empirical target distribution.
To control the performance of ((0) below, we want to use Lemma 9. To do so, we need to verify that (1) also satisfies the key assumptions (A1) to (A3). This is accomplished in the following lemma.
Lemma 10 (Error rates for (())) For any i 2 O. the measures μl() and continue to satisfy (A1) to (A3), and thus
E O - Tollz 55 1/ +
Proof. We verify that the conditions (A1) to (A3) hold for the pair (1(2), v); repeating the argument for the other iterates straightforward.
First, we recall that for two measures 140 M1 = // with support in convex subset Ω S RRd, the McCann interpolation (Ha))e(o,1) remains supported in Q; see Santambrogio [2015, Theorem 5.27].
Moreover, by Proposition 7.29 in Santambrogio [2015], it holds that
| 1. (0) S max [|lo||||| (()) I| (()} S mmxx(Mmxx 1mmax) )
for any a E [0, 1), recall that the quantities jmax: Vmax are from (A1). Thus, the density (1)
Moo = ((1 00) ld T' #/ uniformly upper bounded on Ω; altogether this covers (A1). For
21

---

(A2) and (A3), note that we are never leaving the geodesic. Rather than study the "forward" map, we can therefore instead consider the "reverse" map
T = (ag ld +(1 +(1-app -11-a))f-1) 1)
which satisfies T = =x(0) and hence T(1) = f-1. We now verify the requirements of (A2) and (A3). For (A2), since (T(D)-1 = T = og Id +(1 ag) and T is three-times continuously differentiable by assumption, the map (2-(1) is also three times continuously differentiable, with third derivative bounded by that of For (A3), we use the fact that T = Vioo for some function PP which is A-smooth and X-strongly convex. Basic properties of convex conjugation then imply ||=1|= + (1 pp) Since the conjugate of a A-strongly convex that T(1) = Vi1,, where ;1 = (on 2
function is A -smooth, and conversely. we obtain that the function ęP is (oo + (1 oo )
strongly convex and (as + (1 A 1)-- In particular, since (oo + (1-00)}1t)-- > min and (0kg + (1 . 1)-- s max (1, A), we obtain that DT(1) is uniformly bounded above and below.
We define the following quantities which we will recursively bound:
A == BlIg(0 (1) AR. = || -T(0)
W. = BW2 (i),, (15) [[ -T00)
as well as
A = 1 a, + -122
where recall R is the radius of the ball B(0; R) in Rd.
First, the following lemma:
Lemma 11. If the support of p. and v is contained in B(0; R) and a S fo i = O,... k. then
[
To (2()))
A
Proof. We prove this lemma by iterating over the quantity defined by
o EE = || 2(0) o g(k-1) o...o 8(6) - (1))) o g(k-1) o o... sti 2
when j < k 1 and Es = Ar. By adding and subtracting (ti and (( appropriately, we obtain for
j < k
Ej x (((0) g(k-1) o... c() - g()) ((- .o...................... .... g(i+1) sti
LA(WO))
+ +||e00 o ogtt)) o (ti) (k)) s(k-1) s(1) 2
(ipp(( Lip(8(-11)) ".. A, + Ej+1
II A
A + Eft1
S
where in the last inequality we have used the fact that oy S Ey. so that A₁ S for all k. Repeating this process yields
P
To
(2())) = Eo S Et Ag. which completes the proof.
To prove Theorem 3, it therefore suffices to bound Ag. We prove the following lemma by induction, which gives the proof.
22

---

Lemma 12. Assume d 2 4. Suppose (AI) to (A3) hold, and a, x z// and E x xn-1/2d for all i E [k]] Then it holds that for k 2 0,
Wx Slog(n) n 2/d and D Slog(n) n 1/d
Proof. We proceed by induction. For the base case = 0, the bounds of Fournier and Guillin [2015] imply that Wi B p)] & n n-2/d Similarly, by Lemma 8, we have Do Slog(n)
Eo --4/2 + eb Slog(n) n n-1//_
Now, assume that the claimed bounds hold for Wa and Ak. We have
W W++1 B[W ( (8(k)
)
5 <EW3 (8(*)) ) (( (8 ) (b) + E W (18(A)) ) (8(k)) #
((Lip(8(k)) 22 (p 4 (1)) E s S(A) 1|2
(
A/ + af
Wx Ax.Ax
S 2/d
where the last step follows by the induction hypothesis and the choice of ag. By Proposition 4 and the preceding bound, we have
A $ [1g(k+1)
AC 11 (2
A((+1)
+ A AR
S Exit Wk+1 ARREI
S Ek41 -2/d + A
Lemma 10 implies that A ARE Stor n -1/d, The choice of Ex41 therefore implies As41 Slog(n) n -1/d completing the proof.
C.3 Proofs for the CIFAR Benchmark
Proof of Proposition 5. By Proposition 13, the Gaussian blur map is a linear positive-definite operator. Considering the h 2 |1-112 cost, then G acts as Monge map between from distribution f over a finite set of images, onto their blurred counterparts & G This is a direct corollary of Brenier's Theorem, and follows the fact that G is the gradient of h(U) 1 (U, potential, and therefore Monge map [c.f. Section 1.3.1, Santambrogio, 2015]. Therefore, and again following Brenier, the optimal assignment between and their blurred counterparts i, necessarily that which maps an image to its blurred version regardless of the value of < oo and the optimal permutation is the identity.
(U,G(U) G(U)) the convex
Proposition 13. The gaussian blur operator G: - KUK is a linear positive-definite operator, where U, K NxN and the kernel K. is defined via
K = [exp (- j}>/(oN3)))
Vi,j N.
Proof. The linearity of the operator is implied by the linearity of matrix multiplication. As for the dositive-deiniteness. we show that the kernel matrix corresponding to the operator G positive- definite. Consider images UI...., U. in N and the corresponding kernel matrix A E Raxs defined as
Ag (U, G(U,)) (U,, KU,K) = (KU, KUJ) This is dot-product matrix (of all elements KU,), and is therefore always positive definite.
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have linked the numerical and theoretical evidence for the claims made in the abstract and introduction in the text. We have provided a theorem showing consistency of our method.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitation of theory are clearly stated in Assumptions Al-A3. We report the error and runtime of the algorithm, showing the numerical limitations, over numerous
runs and 3 different datasets.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-tex system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: This provided in Al-A3 and in the statement of the Theorem and lemmas in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
All experimental details are reported in the beginning of the experiment section. If pyper-parameter is chosen via cross-validation, the CV method is described. Otherwise, the exact values are reported. We provide detailed algorithm pseudo-code which allows the reader to re-implement the method without using our code.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
25

---

(a) If the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We do not share new dataset. We include the base implementation of our main algorithm Jax as supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm mublic/guides wttps//oooSSSmmaoooonoooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: ciaaps.cc/pubbicco http://ggiiiee//eeeubiiiiiieee
//oo//eoooemseeeeeeee wwwtimisssiooooooo onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/eetails
Question:
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: This is detailed in first paragraph of Section 5 and in later paragraph titled "Comparing Map Estimators on Single Cell Data".
Guidelines:
26

---

The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
We run every experiment for (or depending on the experiment) random seeds. All results are reported with error bars indicating the standard error. In our Table, we do the same, reporting only numbers that have statistical significance.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Our experiments are light and ran on a single GPU node. Details mentioned in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
27

---

The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: : pectps:/nneus /ttp//eurppp...// cttps.cc/ppllee
Answer: [Yes]
Justification:
We fully adhere to the NeurlPS code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper of theoretical and methodological nature. We do not foresee an application of our proposed algorithm which may have a direct social impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative socictal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
28

---

Justification:
We do not release model or dataset.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We use open sourced data and code. The reference packages and papers cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not introduce asset.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the dataset/code/moodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
29

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We do have involve crowdsourcing or human feedback.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do have involve crowdsourcing human feedback.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
30


Reviews:
[{"review_id": "xcMgqFSJKN", "reviewer": "Reviewer_Agip", "summary": "The authors proposed ProgOT, a method to solve a sequence of EOT problems so that practitioners do not have to tune the entropic regularizer parameter $\\varepsilon$ and strike a good balance between computational and statistical complexity.", "questions": "- How would $\\nu_{min}, \\nu_{max}$ affect the theoretical analysis?\n- I would recommend the authors to consider quadratic regularized OT/UOT formulations [1], [2] and entropic-UOT [3].\n- In addition, I would recommend the authors to cite and discuss the following works as well.\n\nReferences:\n[1] Smooth and Sparse Optimal Transport. Mathieu Blondel, Vivien Seguy, Antoine Rolet. In Proc. of AISTATS 2018. https://arxiv.org/abs/1710.06276\n\n[2] \"On Unbalanced Optimal Transport: Gradient Methods, Sparsity and Approximation Error\". Quang Minh Nguyen, Hoang Huy Nguyen, Lam Minh Nguyen, Yi Zhou. Journal of Machine Learning Research, (JMLR), 2023\n\n[3] Pham, K., Le, K., Ho, N., Pham, T., & Bui, H. (13--18 Jul 2020). On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm. In H. D. Iii & A. Singh (Eds.), Proceedings of the 37th International Conference on Machine Learning (pp. 7673–7682). Retrieved from https://proceedings.mlr.press/v119/pham20a.html", "limitations": "See above.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The manuscript is well-written and it was easy to understand.", "weaknesses": "I have several concerns about the contributions and the assumptions of this work:\n- First of all, I think the contribution of this work is rather limited since the only benefit is to avoid the tuning of the parameter $\\varepsilon$. In my experience, tuning $\\varepsilon$ is usually not a big issue since Sinkhorn is rather robust.  In addition, the work does not discuss any computational tradeoff in doing so. Tuning the parameter $\\varepsilon$ only affects the accuracy of EOT, which is rather minor since the main bottleneck of large-scale Optimal Transport is the number of support $n$.\n- Secondly, this work uses a lot of strong assumptions to establish theoretical guarantees. For example, it assumes Euclidean cost, convex, and compact condition and the inverse mapping has at least three continuous derivatives. These assumptions are very restrictive and would limit the applicability of the theoretical analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your review and for your comments. Our response follows.\n\n> **I think the contribution of this work is rather limited since the only benefit is to avoid the tuning of the parameter $\\varepsilon$.**\n\nIn light of your comment, we will delete the sentence *“Setting $\\varepsilon$ can be difficult, … and bias”* in our abstract, so that readers do not get confused on what constitutes our contribution.\n\nOur paper is not “only about avoiding tuning $\\varepsilon$” in Sinkhorn. While this was the main focus of other recent works, such as [1, 2], **L. 59**, this is not our goal.\n\nOur solver blends McCann type interpolation, local Sinkhorn computations and entropic map estimation to yield both a coupling and map estimator (**L.12~18**). While easing the burden to pick $\\varepsilon$, our solver does a lot more than that, as listed in *Contributions*, **L.65~76**, and demonstrated in Fig. 5 (computations, highlighting very different outputs compared to Sinkhorn), or Table 1 (out-of-sample prediction). \n\nPlease also take a look at our new experimental results [**Fig. A, Tab. B**] on the entire CIFAR-10 dataset.\n\n\n> **In my experience, tuning is usually not a big issue since Sinkhorn is rather robust.**\n\nUse cases for Sinkhorn vary a lot, and we do trust that you may not have run into such problems in the past. In practice, however, EOT solvers are widely used, for many tasks, impacting many user profiles. \n\nIn the simplest cases, e.g. when computing **couplings on normalized data** (e.g. embeddings on the sphere), we also observe that Sinkhorn's outputs can be fairly stable across various $\\varepsilon$ values. On **real** data sources (such as the genomics point clouds), this is not necessarily the case. $\\varepsilon$ has a considerable impact on stability and compute time.\n\n\n$\\varepsilon$ also impacts disproportionately the predictive performance for map estimation. Not tuning it often results in poor outcomes, which is why *all* Sinkhorn baselines in our draft use cross-validation. To convince you, we have added a line with \"untuned Sinkhorn\" (i.e. using the default given by OTT-JAX) to our benchmarks [see **Tab. D** in the PDF], showing that it performs significantly worse than cross-validated Sinkhorn and ProgOT.\nThank you for this suggestion.\n\n\n> **In addition, the work does not discuss any computational tradeoff in doing so.** \n\nWe did target this tradeoff explicitly in Figure 5, see **L.314**. We have updated appendix A with more experimental results to highlight that tradeoff.\n\n> **Tuning the parameter only affects the accuracy of EOT,...**\n\nCan you clarify what you mean by accuracy of EOT? Do you mean proximity to the unregularized solution?\n\n> **...which is rather minor since the main bottleneck of large-scale Optimal Transport is the number of support $n$**\n\nEach Sinkhorn iteration has indeed a cost in $n^2$. However, **the number of iterations varies tremendously with $\\varepsilon$**, e.g. from 10 iterations for large $\\varepsilon$ to 10k for small ones, as shown e.g. in our paper, and [**Fig. F & Fig. G**].\n\nTo our knowledge, the only linear time RegOT solver is the low-rank (LRSink) approximation from [3]. That solver does not yield a map estimator, as it does not solve the dual OT problem. Exploring hybridizing ProgOT with LRSink is an interesting direction.\n\n> **...this work uses a lot of strong assumptions to establish theoretical guarantees. For example, it assumes Euclidean cost, convex, and compact condition and the inverse mapping has at least three continuous derivatives. These assumptions are very restrictive and would limit the applicability of the theoretical analysis.**\n\nTo prove guarantees for estimators, the strength of assumptions should naturally be commensurate with the difficulty of the task. \n\nProving theoretical recovery of OT maps is hard, and to our knowledge, the entropic map estimator [4] that we build upon is the **only** tractable estimator that has such guarantees (**L. 126~129**). We follow their assumptions, and we do not see any alternatives now. The fact that the inverse mapping must have three continuous derivatives is a strong assumption, but we believe this is a fundamental issue, and not an artefact or our analysis. Other approaches (e.g., [5, 6, 7]) share many of the same regularity assumptions as we do. \n\nNote that many OT map estimators are routinely used without any theoretical guarantees (e.g. all neural-OT approaches, flow-matching like formulations), even in biological settings [8]. Aside from the entropic map, ProgOT is the only other tractable map estimator that comes with convergence guarantees.\n\nTo summarize, apart from the entropic map estimator, ProgOT is the only currently known tractable OT map estimator with guarantees, and we show that it performs better in practice.\n\n\n> **How would $\\nu_{\\min}, \\nu_{\\max}$ affect the theoretical analysis?**\n\n$\\nu_{\\min}$ and $\\nu_{\\max}$ appear as constants in the bounds, as established by [4], which are just meant to be multiplicative constants in the bounds. The role is analogous in the other papers mentioned above \n\n> **I would recommend the authors to consider quadratic regularized OT/UOT formulations [1], [2] and entropic-UOT [3]. In addition, I would recommend the authors to cite and discuss the following works as well.**\n\nThis is an interesting research direction. However, extending our work to the settings you mention (quadratic regularization/ unbalanced settings) is far from being straightforward. Crucially, there is no equivalent of an “entropic potential map”, that can be extended out of sample, for the dual solutions outputted from quadratic solvers (first proposed in [DPR16]), nor for the unbalanced case.\n\nWe have added references that you mentioned to the conclusion section, to highlight possible future work around ProgOT.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "EhtKJfWRoE", "reviewer": "Reviewer_Z1nb", "summary": "The choice of appropriate entropic trade-off term is one of the main headaches for finding maps between data distributions with sample access when considering Optimal Transport (OT) with entropic regularization. While the selection of sufficiently small regularization terms leads to unstable learning, the picking of large ones causes biased solutions. Considering the EOT problem with Euclidean quadratic cost, the authors propose to divide the entropic problem to small sub-problems between intermediate distributions with their own regularization terms. The final entropic map between initial data distributions is a composition of maps from sub-problems. The authors provide the theoretical guarantees that the constructed map does not differ much from the true one. Moreover, the authors offer the algorithm for selecting appropriate regularization terms for each step as well as demonstrate the method’s performance on synthetic data and single-cell experiments.", "questions": "- From intuition’s point of view, it is understandable that a sub-problem is easier and well-conditioned, than the initial problem. Anyway, could you provide a theoretical understanding why it is so? It seems that this fact depends on the step-size of the McCann interpolator.\n\n- You probably use entropic OT formulation with Euclidean quadratic cost due to the convenient use of linear McCann interpolation between distributions. Is there generalization of the method for arbitrary transport cost function?\n\n- Does the method’s algorithm need a pre-trained OT map between initial distributions that approximately matches one distribution to another?\n\n- Is there intuition or methodology of picking a more appropriate step-size for McCann interpolator at step k?\n\n- Why does the speed of convergence of the proposed algorithm compare to the speed of the Sinkhorn algorithm of EOT problem between initial distributions? Could you provide a theoretical explanation and plot of convergence, at least in a synthetic Gaussian example?\n\n- In accordance with the developed methodology, the geodesic curves between initial distributions should be straight. Could you provide practical evidence that obtained OT curves are really straight, at least synthetic Gaussian experiments?\n\nI suggest the authors to test their algorithm on benchmark with available ground-truth OT maps for bigger values of $d$.\n\n**References.**\n\nA. Korotin, L. Li, A. Genevay, J. M. Solomon, A. Filippov, and E. Burnaev. Do neural optimal transport solvers work? A continuous Wasserstein-2 benchmark. Advances in neural information processing systems, 2021.", "limitations": "In accordance with the theorem 3 of the main text, the main limitation of the method is scalability. However, it seems that the authors did not mention this fact.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper provides a new methodology for solving OT problems with entropy regularization which does not combine the ideas of prior works. The paper is well-written, well-organized and clear. The storyline is perfect and understandable through the paper. \n\n- Since the method alleviates the tuning process for entropic regularization terms and allows to avoid unstable learning as well as finding primitive OT maps, this work might be interesting to the ML community. I am sure that other researchers might apply this methodology for Flow matching or Schroedinger bridge methods that aim to build interpolation from one dataset to another one.\n\n- The ProgOT alleviates the tuning process for entropic regularization terms, thereby it does not suffer from unstable learning as well as primitive solutions.", "weaknesses": "Although the paper offers a unique theoretical approach, a crucial shortcoming of this method is **scalability**. Indeed, the estimation, which is provided by Theorem 3, demonstrates poor scalability of the method since the rate of convergence to the true OT maps depends on the dimension $d$. From my point of view, this issue of the approach could not be ignored since it severely limits its practical usability.\n\nThis my concern is supported by the  set of experiments considered in the paper. It seems that the authors avoid empirical evaluation of the method in high-dimensional problems (which might be useful for mitigating the concern). The performance of the method is tested only in low-dimensional ($d<=64$) synthetic data experiments using the benchmark with available ground-truth OT maps (Korotin et al, 2021), as well as biological experiments with no sufficiently large dimensionality ($d<=256$) of data. I am especially confused by the fact that the authors do not test their approach using the provided benchmark pairs for $d$ larger than 64 and at the same time apply the method to an important biological task in case of $d=256$. However, I believe that prior to adapting the algorithm to the tasks from single-cell biology (especially as important as prediction of cancer cells responses to drug perturbation), it should be thoroughly tested on synthetic tasks.\n\n**Overall**, although the proposed approach is as fast as the Sinkhorn algorithm and demonstrates better performance in low dimensional problems, there is no guarantees and understanding of the method’s behavior in high dimensional problems as well as there are no comparisons with other EOT solvers in this case. This is the main reason of my current score. However, I am open to adjust the score based on the authors' answers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for your review, and for the many thought provoking questions you have asked. We did our best to answer all of them.\n\n> **a crucial shortcoming of this method is scalability. […] severely limits its practical usability.**\n\nThere might be a confusion: Theorem 3 proves a theoretical recovery guarantee for ProgOT as a map estimator. This is a worst-case statistical analysis. This worst-case study has little impact on prediction performance on real tasks (e.g. Table 1), and no impact on computational aspects (Figure 5) or coupling estimation, which are all practically governed by other factors ($K, \\alpha_k, \\varepsilon_k, \\tau_k$ for instance).\n\nTheorem 3 guarantees the soundness of the ProgOT map. In contrast, **none** of the NN-based methods (e.g. those compared in the [Korotin et al. 21] benchmark) have a theoretical statistical guarantee. The fact that ProgOT has such a guarantee (even if pessimistic) cannot therefore be seen as a weakness. Their estimation is hard, because of non-convexity. In contrast, ProgOT is recovered as a sequence of convex problems.\n\n> **the authors do not test their approach using the provided benchmark pairs for $d$ larger than 64 […] it should be thoroughly tested on synthetic tasks.**\n\t\nThanks for this suggestion. To demonstrate the scalability of ProgOT to larger $n$ and larger $d$, we have added **two new experiments**: MSE error on the Korotin benchmark for $d=128, 256$ (as you requested) and a new large scale experiment on *all* $n=60.000$ grayscale images ($d=1024$) of the CIFAR10 dataset.\n\n**GMM Experiment**: We use the code from the [Korotin et al., 2021] benchmark to create high-$d$ synthetic problems, reaching the maximum dimension available $d=256$. **Tab. D** in the uploaded PDF presents the results. ProgOT (with default settings) dominates other baselines. This is remarkable, because: (1) this benchmark favors Neural OT solvers by design: the ground truth transport is itself generated as the gradient of an ICNN. (2) running ProgOT takes about 3 minutes, while Neural solvers take *10-20X longer* (about 30’ for ICNN and 60’ for the Monge Gap), all on a A100 GPU.\n\n**CIFAR10 Experiment**: We designed a novel benchmark task, to match images and their blurred versions. As described in our general response, the ground truth OT coupling should be the diagonal. We only compare ProgOT to Sinkhorn (ICNN and Monge Gaps do not return couplings). See *Fig. A& Tab. B* in the uploaded PDF for the results and refer to the general response for more details.\n\nNote that a sample size $n=60,000$ is, to our knowledge, unheard of for entropic OT for such high-$d$ (we overwhelmingly see $n\\leq 10,000$ in papers). We achieve this through a tight implementation of ProgOT within OTT-JAX, and JAX, getting the immediate benefits of sharding across GPUs.\n\n> **You use entropic OT formulation with Euclidean quadratic cost due to […] Is there generalization of the method for arbitrary transport cost function?**\n\nThis generalization is already in the paper: we defined ProgOT using arbitrary translation invariant (TI) costs, $c(x,y) = h(x-y)$ (see **L.81**). **Alg. 2 and 3** both use $h$ explicitly. We use a a $p$-norm (with $p=1.5$) in the bottom of Fig. 5. We only need $h$ to be $\\tfrac12\\|\\cdot\\|^2$ when proving map guarantees in Section 3.2, see **L.184**.\n\nTo complete the empirical overview on general costs, **we have added a new table to the appendix A**, which benchmarks map estimation with general costs (similar to Table 1).\n\n> **it is understandable that a sub-problem is easier and well-conditioned, than the initial problem […] could you provide a theoretical understanding why it is so?**\n\nBecause the intermediate problems lie along the geodesic path, the distance/transport costs for them is smaller, **L. 154~163**. Calling Sinkhorn for problems with smaller transport costs is easier/more stable w.r.t. $\\varepsilon$. This also allows warm-starting to recycle solutions from the previous step (**L.3** in **Algo. 2**.)\n\n> **Does the method need a pre-trained OT map?**\n\nNo, ProgOT runs off-the-shelf using source/target data (see L.229, inputs to **Algo.2**) and is deterministic, as it relies on sub-Sinkhorn calls. Could you please clarify your question?\n\n> **Is there intuition or methodology of picking a more appropriate step-size for McCann interpolator at step $k$?**\n\nWe investigated this thoroughly. We found out that scheduling the step-size does not seem to play a key role, neither theoretically nor empirically (see Fig 7 in the Appendix). We stick to the simplest choice, the constant-speed, a.k.a. linear, schedule.\n\n> **Why does the speed of convergence of the proposed algorithm compare to the speed of the Sinkhorn algorithm of EOT[…]**\n\nThe convergence rates are comparable because ProgOT uses calls to local Sinkhorn problems as a subroutine. Theorem 3 guarantees theoretically that by using ProgOT we remain statistically sample efficient, while gaining a computational edge. Theorem 3 is illustrated empirically in Figure 4 (A), where we plot the MSE for various $n$, using the ground-truth OT map from [Korotin et al., 2021].\n\n> **Could you provide practical evidence that obtained OT curves are really straight, at least synthetic Gaussian experiments?**\n\nUsing **Algo.2**, curves can be either obtained by sampling straight lines from the final coupling matrix $\\mathbf{P}$ returned by ProgOT, or using the map (Alg. 3) out of sample. Sampling from coupling yields straight lines, by construction. This might be used in, e.g., flow-matching estimation. The \"straightness\" of lines sampled with Alg. 3 will hinge on the various $\\varepsilon_i$, as represented schematically in Fig. 3.\n\n> **I suggest the authors to test their algorithm on benchmark with available ground-truth OT maps for bigger values of $d$**\n\nMany thanks for this suggestion. We hope our new results [**Tab. D**] assuage your concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "idoCFdJNY8", "reviewer": "Reviewer_mHCn", "summary": "This paper introduces a new class of entropic optimal transport (EOT) solvers called PROGOT. This work aims to address the challenges of selecting entropic regularization strength $\\epsilon$ for original EOT. As we know $\\epsilon$ is significant to the performance of EOT like computation time and convergence rate. PROGOT utilizes a progressive framework that interpolates the whole entropic transportation process into multiple steps by using dynamic OT formulations. This work proposed algorithms to set the parameters throughout the interpolation process. As it claimed, the new framework enhances the robustness and performance of EOT, and avoid the headache of tuning the value of $\\epsilon$. Experimental results show that PROGOT surpasses classic EOT in both synthetic and real dataset.", "questions": "- Line 141: do you mean $S^{(0)}$ instead of $S^{(1)}$ in the equation $\\mu^{(1)}=S^{(1)}\\mu$? The same issue in line 143.\n- Figure 5, why does PROGOT with larger $K$ value has larger cost and appear closer to the original Sinkhorn point? Please correct my intuition if wrong: PROGOT with smaller $K$ runs fewer interpolation steps, leading to fewer number of iterations and closer to the original Sinkhorn cost. When $K=0$, PROGOT should be the same as EOT. In figure 5, the number of iteration results align with my intuition, but the \"distances\" to Sinkhorn results looks inconsistent.", "limitations": "None", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The major contribution of this work is the proposal of a novel framework to estimate entropic transport plans. This framework addressed the challenge of selecting an appropriate entropic regularization strength $\\varepsilon$ in traditional EOT algorithms.\n- The mathematical notation and theoretical analysis are clear and easy to follow. And the theoretical analysis looks good to me.\n- The convergence of PROGOT map is supported both in theoretical proof and experimental results.\n- The methodology for selecting hyper-parameters (regularization and threshold schedule) and the corresponding justifications are well provided.", "weaknesses": "- While the effort in setting the hyper-parameters and the justifications is acknowledged, from a broader perspective, to replace the selection of $\\varepsilon$, PROGOT introduces a series of new hyper-parameters (step/regularization schedules and the length of the schedules, $K$) and theoretic assumptions on the inputs. One may find this less favorable, as it replaces one hyper-parameter with several others. \n- Some of the experiments would be benefit from including of the real OT cost (or map) as a reference. For example, adding the real OT map in Figure 1 and adding the real OT cost data point in Figure 5 will be helpful. \n- The gradient of PROGOT is missed, which is important for machine learning application.\n- In Figure 5, could you provide the actual number of iterations instead of just indicating different marker sizes? For example, comparing the size of markers for $\\beta=0.08$ vs. $K=8$ is difficult. Additionally, since all the sub-figures use the same legend, it should be clarified whether they run the same number of iterations for the same configuration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your careful review, we thank you for your positive score, encouraging comments and questions. We did our best to answer them.\n\n> **… One may find this less favorable, as it replaces one hyper-parameter with several others.**\n\nWhile we certainly agree with this assessment, one message we want to convey is that having to choose a single variable $\\varepsilon$ for everything (Sinkhorn, map estimation) holds too much sway over EOT.  A \"good\" $\\varepsilon$ can be difficult to nail right, notably for map estimation, and most practitioners want to bypass this. As it stands, this is likely one of the factors that limits usage of EOT.\n\nOur goal is to \"divide and conquer\" $\\varepsilon$ too, translating it into simpler quantities. In practice, a user *would only need to only choose the number of steps $K$*, according to the means of compute available to them, since we provide automatic routines for choosing the rest of the hyper-parameters: the step schedule is linear (a.k.a constant-speed), the threshold scaling log-linear, and the epsilon schedule is tuned as in Algo. 4.\n\n> **The gradient of PROGOT is missed, which is important for machine learning application.**\n\nDifferentiability is an important point. We would like to provide a satisfactory answer by splitting the various meanings this could take in the context of ProgOT. \n\nProgOT can return three objects that a user may want to differentiate:\n- “Wasserstein-like” OT transport cost between two point-clouds, equal to $t(\\mathbf{X},\\mathbf{Y}):=\\langle \\mathbf{P}, [h(x_i-y_j)]_ij\\rangle$ where $\\mathbf{P}$ is the solution outputted by Algo.2. This is the quantity displayed in the `x-axis` of Fig. 5 (as well as in our new plots). We believe this is the quantity you have in mind when you mention a gradient. Differentiating this quantity w.r.t. $\\mathbf{x}, \\mathbf{y}$ is, in fact, very easy, since it amounts to using a “Danskin-like” linearization, i.e. use the chain rule to differentiate $t$ w.r.t. $\\mathbf{X}$ or $\\mathbf{Y}$ (or any other parameter), while keeping $\\mathbf{P}$ fixed. We illustrate this in **Figure H** of the uploaded pdf, and also provided a `colab` link to the AC, who can share it with you.\n- The Jacobian of the coupling matrix $\\mathbf{P}$ returned by Algo.2. This can be achieved by differentiating chained Sinkhorn iterations. This is the strategy used, in a different context, when differentiating the optimal coupling obtained in Gromov-Wasserstein for instance. It is costly, but doable.\n- The Jacobian of the ProgOT map (that returned by Algo. 3) at a given point $x$, w.r.t. Parameters. This is the same as above, and will require a chained differentiation of successive OT solutions.\n\nHence, it is very easy to use ProgOT to obtain descent directions to minimize OT-costs. Differentiating higher-order objects (couplings / map estimator) is more involved.\n\n\n> **In Figure 5, could you provide the actual number of iterations instead of just indicating different marker sizes?** \n\nWe have updated the figures in the paper to show the number of iterations, see Figure **F** for the updated version of Figure 5 in the paper. We have also added 4 new figures to the appendix in new data configurations. **Fig G** in the upload PDF shows an example.  To answer your question, We have prepared **Table E** with the exact number of iterations corresponding to Figure F, (the updated variant of Figure 5).\n\n> **Line 141: do you mean $S^{(0)}$ instead of $S^{1)}$ in the equation $\\mu^{(1)}=S^{(1)}\\mu$? The same issue in line 143.\n\nThank you for pointing this typo out, we have updated the text.\n\n>**Figure 5, why does PROGOT with larger $K$ value has larger cost and appear closer to the original Sinkhorn point? Please correct my intuition if wrong: PROGOT with smaller runs fewer interpolation steps, leading to fewer number of iterations and closer to the original Sinkhorn cost. When $K=0$, PROGOT should be the same as EOT. In figure 5, the number of iteration results align with my intuition, but the \"distances\" to Sinkhorn results looks inconsistent.**\n\nEven if $K$ is small, the cost values might be different if the regularization parameter is different. We visualize Sinkhorn using different values $\\beta$, and in many cases (c.f. **Fig F** and **Fig G**) ProgOT with $K=2$ lies close to a Sinkhorn instance. Perhaps your intuition is better reflected in the updated plots for this experiment.\n\nWe also highlight that the two axes of cost and entropy should be viewed together. One possible scenario is that by choosing a larger $K$, we target a smaller $\\varepsilon$ in the last iteration, and gaining sharper entropy, at the price of having a larger cost.\n\nSome of the results in Figure 5 in our initial draft were slightly off because we were using a maximum number of iterations when implementing Sinkhorn and ProgOT. While the overall message of these plots have not changed, this error sometimes contradicted with our claim that all methods were converging to target accuracy (**L.322**). We now set that maximum iterations to infinity, and therefore all solutions converge correctly within the desired threshold, irrespective of the number of iterations. See **Fig F** and **Fig G** in the uploaded PDF for an example.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "iPob1t2DHk", "reviewer": "Reviewer_9mT1", "summary": "In this paper, the authors propose a new entropic optimal transport solver as an alternative to the commonly used Sinkhorn algorithm named ProgOT. This solver has three main properties:\n\n(i) It is less sensitive to the choice of the entropy-regularized parameter than the Sinkhorn algorithm;\n\n(ii) When computing couplings between point-clouds, the runtime of ProgOT is no longer than the Sinkhorn algorithm;\n\n(iii) The resulting optimal transport map estimator is consistent and stable.\n\nThe authors provide both theoretical and empirical results to demonstrate the above claims.", "questions": "1. Are there any chances that the assumption (A.2) can be reduced?\n\n2. What is the computational complexity of the ProgOT algorithm (Algorithm 2)?\n\n3. Could the authors please explain more clearly why line 3 in Algorithm 2 helps improve the runtime?\n\n4. Are there any instructions on how to choose the number of steps $K$?", "limitations": "The limitations have not been discussed in the paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed ProgOT is a new entropic optimal transport solver built based on the McCann-type interpolation.\n\n2. The theoretical results in the paper are sound, and provided with rigorous proofs.\n\n3. The paper is well written and organized. The background section is very useful, particularly for readers who are not familiar with optimal transport.", "weaknesses": "1. Since the optimal transport map $T_0$ is unknown in practice, the assumption (A.2), which says that the inverse map $T^{-1}_0$ has at least three continuous derivatives, is quite strong. \n\n2. In lines 33-36, when comparing the efficiency of the Sinkhorn algorithm to the linear programs for solving the EOT problems, the authors should make it more explicit by stating the computational complexity of both algorithms.\n\n3. In lines 238-239, when initializing the entropy-regularized parameter $\\varepsilon_0$, the authors should at least briefly present the intuition for setting it to be the average of the values in the cost matrix between sources and targets. \n\n4. Minor issues: There are some undefined notations and grammatical mistakes:\n\n- In line 141, I think the term $S^{(1)}$ should be $S^{(0)}$.\n\n- In the inequality between lines 196-197, the notation $\\lesssim_{\\log(n),k}$ has not been defined yet.\n\n- In line 167, the notation $\\alpha(k)$ has not been defined. Is it a constant depending on $k$?\n\n- In the assumption (A.3), what does the notation $D$ stand for?\n\n- In line 204, $\\mu^{(k)}$ is corresponds a location --> grammatical mistake.\n\n5. The authors should add a discussion about the limitations of the proposed method as well as future directions. For instance, can we generalize the ProgOT so that it applies for entropic unbalanced optimal transport. I believe that such discussion would make the paper more complete.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our work. We have fixed the notation and typos issues, thank you for pointing them out.\n\n\n> **Since the optimal transport map is unknown in practice, the assumption (A.2), which says that the inverse map has at least three continuous derivatives, is quite strong. Are there any chances that the assumption (A.2) can be reduced?**\n\nWe highlight that requiring up to second-order derivatives on the inverse map is a standard requirement, see [1, 2, 3, 4]. The additional derivative comes from the assumptions in [4], which is, to date, the only work that demonstrates that the entropic OT map can statistically estimate the OT map. It is largely believed that **the assumption is not superfluous**, and that the extra derivative is required to make this statistically viable [5]. Though there are restricted scenarios under which the 3rd derivative is unnecessary, e.g., if $\\nu$ is a discrete measure [5]. One can hope that ultimately these results can be bridged, but for now this is outside the scope of this work.\n\n\n> **What is the computational complexity of the ProgOT algorithm (Algorithm 2)? In lines 33-36, when comparing the efficiency of the Sinkhorn algorithm to the linear programs for solving the EOT problems, the authors should make it more explicit by stating the computational complexity of both algorithms.**\n\nThe exact computational complexity depends on the problem geometry, since $\\varepsilon_i$ is data-driven.\n\nWe can specify the worst-case computational complexity of ProgOT compared to Sinkhorn. The algorithm runs the sinkhorn subroutine $K$ times, with a decreasing sequence of regularizations $\\varepsilon_i$. ProgOT also performs linear interpolations, yielding a total of  $\\mathcal{O}(nd + \\sum_{i=1}^K C_{\\mathrm{Sink}}(\\varepsilon_i))$ operations, where $n$ is the size of the input domain, $d$ the dimension and $C_{\\mathrm{Sink}}(\\varepsilon_i)$ denotes the worst-case complexity of Sinkhorn with regularization $\\varepsilon_i$. We *added a paragraph to Appendix A* discussing this matter.\n\n\n> **In lines 238-239, when initializing the entropy-regularized parameter $\\varepsilon$, the authors should at least briefly present the intuition for setting it to be the average of the values in the cost matrix between sources and targets.**\n\nWe have *updated that paragraph and included a brief explanation*. This choice is not ours, but largely followed by the community. To avoid simple scaling effects (such as multiplying all features by a constant), cost matrices in entropic OT are rescaled, typically by their maximum value (as commonly done in POT [7]) or by mean-value (as done in OTT [6]). This rescaling is often absorbed into $\\varepsilon$, and helps balance out cost/entropy in EOT optimization.\n\n\n> **The authors should add a discussion about the limitations of the proposed method as well as future directions. For instance, can we generalize the ProgOT so that it applies for entropic unbalanced optimal transport. I believe that such discussion would make the paper more complete.**\n\nThis is indeed an interesting research direction. *We have updated the conclusions section* with a discussion on directions such as unbalanced OT. We note that such extension is far from being straightforward. Crucially, there is no equivalent of an “entropic potential map”, or McCann interpolation, that can be extended out of sample. \n\n> **Could the authors please explain more clearly why line 3 in Algorithm 2 helps improve the runtime?**\n\nLine 3 implements a warm-start for the Sinkhorn algorithm. The learned potentials at iterate $k$ might still be relevant for the next Sinkhorn subroutine. Our update is essentially a back-of-the-hand envelope: since the distances between point clouds decrease at each iteration by roughly a factor of $(1-\\alpha_k)$, we decrease the dual potentials by a similar amount. We observe that this simple update works well across various costs. Warm-starting Sinkhorn is an active area of work [6, 7, 8, 9].\n\n\n> **Are there any instructions on how to choose the number of steps?**\n\nIntuitively, the number of steps should correlate positively with the distance between source and target, but for now, we do not have explicit instructions for this.\n\nA simple solution can be choosing $K$ depending on the available compute. Choosing large $K$ typically does not affect performance negatively, but smaller $K$ may create a less stable algorithm (as we get back to Sinkhorn). In our experiments we have used K = 4, 8, 16, and noticed diminishing returns in increasing $K$. \n\n## References\n---\n[1] Hutter and Rigollet. \"Minimax estimation of optimal transport maps\", (2021)\n\n[2] Manole, et al. \"Plugin estimation of smooth optimal transport maps\", (2021)\n\n[3] Muzellec, et al. \"Near-optimal estimation of smooth transport maps with kernel sums-of-squares\", (2021)\n\n[4] Pooladian, Aram-Alexandre, and Jonathan Niles-Weed. \"Entropic estimation of optimal transport maps.\" arXiv preprint (2021)\n\n[5] Pooladian et al. “Minimax estimation of discontinuous optimal transport maps: The semi-discrete case” ICML (2023)\n\n[6] Cuturi, Marco, et al. \"Optimal transport tools (ott): A jax toolbox for all things wasserstein.\" arXiv preprint (2022)\n\n[7] Flamary, Rémi, et al. \"Pot: Python optimal transport.\" Journal of Machine Learning Research (2021)\n\n[8] Amos, B., Cohen, S., Luise, G., & Redko, I. Meta optimal transport. arXiv preprint (2022)\n\n[9] Thornton, James, and Marco Cuturi. \"Rethinking initialization of the sinkhorn algorithm.\" AISTATS (2023)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We would like to thank all reviewers for taking the time to read about our new method. We take pride in seeing many appreciative comments, notably from reviewers `Z1nb`, `mHCn` and `9mT1`:
> *I am sure that other researchers might apply this methodology for Flow matching or Schroedinger bridge methods that aim to build interpolation from one dataset to another one*. 

> *The paper provides a new methodology for solving OT problems with entropy regularization which does not combine the ideas of prior works.*

> *The storyline is perfect and understandable through the paper*

> *The theoretical results in the paper are sound, and provided with rigorous proofs.*

> *The paper is well written and organized. The background section is very useful, particularly for readers who are not familiar with optimal transport.*

We enjoyed very much, during this past week, doing our best to come up with clarifications, given separately for each reviewer below.

A common thread found in Reviewers' `Agip` and `Z1nb` comments was about scalability and applicability. We have added various experiments in the pdf (notably higher-$d$ results from the [Korotin et al] benchmark, or a concrete example showing the importance of tuning $\varepsilon$ for vanilla Sinkhorn), but we would like to highlight in our general response a new experiment, on real data, that we designed to assuage specific concerns on scalability.

We wanted to provide an experiment that was large $n$, large $d$, challenging enough to be of interest, on real data to stay relevant, and for which some form of ground truth was known.

We propose an OT problem on the *entire* (grayscaled) **CIFAR10** dataset, that is, $n=60,000$ images, each of size $d=32\times32=1024$.

*The transport task we propose is the following*: The isotropic Gaussian blur operation on an image is a linear operator. Following notations from the textbook **[Peyre & Cuturi, 2019, Remark 4.17]**, for $N\times N$ sized-images (here $N=32$), writing 
$$K=\left[\exp\left(-\frac{(i-j)^2}{N^2\sigma}\right)\right]\_{ij},$$ 
for $i,j\leq N$, the Gaussian blur operation takes an image stored in matrix form as $U\in\mathbb{R}^{N\times N}$ and transforms it into 
$$L(U) = KUK\in\mathbb{R}^{N\times N}.$$

**Crucially, the Gaussian blur linear map is *positive-definite*.**

**Proof**: Let $U_1, \dots, U_s$ be $s$ images, and form the kernel matrix $$A_{ij} = \langle U_i, L(U_j) \rangle.$$
Each entry can be re-written as $\langle U_i, L(U_j) \rangle = \langle U_i, K U_j K \rangle = \langle K U_i, K U_j \rangle$. This is a dot-product matrix (of all elements $KU_i$), and is therefore always psd. This proves that $L$ is a positive-definite linear operator.

As a result, following Brenier's theorem, $L$ is necessarily a $\ell_2^2$ OT Monge map in $R^{N\times N}$, from any family of images onto their blurred counterparts ($L$ is the gradient of $h(U)=\tfrac12\langle U, L(U) \rangle$, which is a convex potential, therefore an OT map, see **[Santambrogio 2015 S.1.3.1]**). 

Practically, if one solves an assignment problem (with $\ell_2^2$ cost) between a family of $n$ images and their blurred $n$ counterparts, the optimal assignment is necessarily that which maps an image to itself, blurred, regardless of the value of $\sigma$ ⎯ the optimal permutation is the identity.

We create these two blurred datasets using a Gaussian kernel with $\sigma=2$ and $\sigma=4$ (visualized in **Figure A** of the uploaded PDF). We then use `ProgOT` and Sinkhorn to match the blurred dataset to the original CIFAR10 (de-blurring). We test various $\varepsilon$ and schedules, and, as always in our experiments, focus on relevant regimes for both methods.

As explained above, the ground-truth (unregularized) OT coupling matrix $\boldsymbol P^\star$ should be the diagonal matrix $q \boldsymbol{I}$, with $q=1/60000$. 

We evaluate the performance of an OT solver by checking how close the trace of the recovered coupling $\mathrm{Tr}(\hat{\boldsymbol P})$ is to $1.0$, or with the KL divergence to the ground truth, that is: $\mathrm{KL}(\boldsymbol P^\star||\hat{\boldsymbol P}) = \log q -\tfrac1q\sum_{i} \log (\hat{\boldsymbol P}_{ii})$. **Table A** in the uploaded PDF shows the performance of `ProgOT` and Sinkhorn, along with the number of iterations needed to achieve this performance. Both Algorithms scale well and show high accuracy, while requiring a similar amount of computation.

We highlight that at this scale, simply storing the cost or coupling matrices would require `28.8Gb`. The experiment has to happen across multiple GPUs. Thanks to its integration in JAX, and OTT-JAX, `ProgOT` supports sharding by simply changing ~3 lines of code. The algorithms scales seamlessly and each run takes about 15 minutes, on a single node of 8 A100 GPUs.

We hope that this experiment sets a convincing example on how `ProgOT` scales to much larger (in $n$ and $d$) problems than considered previously. This ties with the comment from Reviewer `Z1nb` on guided generation or flow-matching, which could benefit from our implementation to handle very large batches.

We would be happy to answer all follow up questions.

the authors

Meta-review: For entropic optimal transport (EOT) problem, the paper proposes a dynamic solvers which breaks down the original difficult problem to a set of simpler problems via time discretization. 
Numerics demonstrate the efficiency of the proposed method.

From the reviewers' comment and the author-reviewer discussions, below is a summary of the reviewers concerns and how they are addressed by the authors 
 - **Reviewer Agip** 1) "only contribution is avoid tuning $\varepsilon$"; (the authors argue that their contribution is not about avoiding tuning $\varepsilon$, and their contributions include a schme that can recover both Kantorovitch couplings and Monge map, good practical performance and map estimation, etc) 2) large-scale problems (the authors provide numerical experiments on large-scale problem showing that the scalability is not an issue);  3) "strong assumptions" (the authors argue that they are standard assumptions). **Author-Reviewer discussion** The reviewer commented that the concerns are addressed. 
 - **Reviewer Z1nb** the main concern is about scalabilirty, and the authors' rebuttal addressed this problem, and this is acknowledged by the reviewer. **Author-Reviewer discussion** The reviewer commented that the concerns in the reviews are addressed. But *I still have some concerns regarding the soundness and practical usability of the proposed approach*.
 - **Reviewer mHCn** hyper-parameters introduced. The authors argued that they only parameter needs to be tune is $K$, which they provide some discussions on how to choose it, though not perfect. **Author-Reviewer discussion** The reviewer commented that the concerns in the reviews are addressed.
 - **Reviewer 9mT1** 1) strong assumption as need to know "optimal transport map $T_0$" (the authors argue that "requiring up to second-order derivatives on the inverse map is a standard requirement"); 2) more explicit comparison against Sinkhorn algorithm in terms of complexity (the authors provide some derivations). **Author-Reviewer discussion** The reviewer commented that the concerns in the reviews are addressed.

In general, the reviewers are satisfied with the authors' reponses. 

During the AC-Reviewer discussion, Reviewers Agip, Z1nb and 9mT1 join the discussion, and the main concerns of them are
 - the contribution is limited.
 - Reviewer 9mT1 still thinks that requiring "optimal transport map $T_0$" is a strong assumption. 

Given the overall reviews and discussions, I think the submission worth a chance, hence recommending accept.
Decision: Accept (poster)