Efficient and Accurate Optimal Transport with Mirror Descent and Conjugate Gradients
Authors: Mete Kemertas
Allan Douglas Jepson
Amir-massoud Farahmand
Venue: Accepted by TMLR ()

Abstract: We propose Mirror Descent Optimal Transport (MDOT), a novel method for solving discrete optimal transport (OT) problems with high precision, by unifying temperature annealing in entropic-regularized OT (EOT) with mirror descent techniques. In this framework, temperature annealing produces a sequence of EOT dual problems, whose solution gradually gets closer to the solution of the original OT problem. We solve each problem efficiently using a GPU-parallel nonlinear conjugate gradients algorithm (PNCG) that outperforms traditional Sinkhorn iterations under weak regularization. Moreover, our investigation also reveals that the theoretical convergence rate of Sinkhorn iterations can exceed existing non-asymptotic bounds when its stopping criterion is tuned in a manner analogous to MDOT.

Our comprehensive ablation studies of MDOT-PNCG affirm its robustness across a wide range of algorithmic parameters. Benchmarking on 24 problem sets of size $n=4096$ in a GPU environment demonstrate that our method attains high-precision, feasible solutions significantly faster than a representative set of existing OT solvers—including accelerated gradient methods and advanced Sinkhorn variants—in both wall-clock time and number of operations. Empirical convergence rates range between $O(n^2 \varepsilon^{-1/4})$ and $O(n^2 \varepsilon^{-1})$, where $\varepsilon$ is the optimality gap. For problem sizes up to $n=16\,384$, the empirical runtime scales as $\widetilde{O}(n^2)$ for moderate precision and as $\widetilde{O}(n^{5/2})$ at worst for high precision. These findings establish MDOT-PNCG as a compelling alternative to current OT solvers, particularly in challenging weak-regularization regimes.

Published in Transactions on Machine Learning Research (06/2025)
Efficient and Accurate Optimal Transport with Mirror Descent and Conjugate Gradients
Mete Kemertas
kemertas@os.torontor
University of Toronto, Department of Computer Science
Vector Institute
Allan D. Jepson
jepson@cs.torontoocem edu
University of Toronto, Deportment of Computer Science
Amir-massoud Farahmand
famamandd@mila..oom quebec
Polytechnique Montréal
Mila Quebec Al Institute
University of Toronto, Deportment of Computer Science
Reviewed on OpenReview: https: // openrevee.. formm? id- Fd=FFFqrDeF8
Abstract
We propose Mirror Descent Oplimal Transport (MDOT); a novel method for solving discrete optimal transport (OT) problems with high precision, by unifying temperature annealing in entropic-rogullarized OT (EOT) with mirror descent tcchniques. In this framework, temper- ature annealing produces sequence of EOT dual problems, whose solution gradually gets closer to the solution of the original OT problem. We solve each problem efficiently using GPU-parallel nonlinear conjugate gradients algorithm (PNCG) that outperforms traditional Sinkhorn iterations under weak regularization. Moreover, our investigation also reveals that the theoretical convergence rate of Sinkhorn iterations can exceed existing non-asymptotic bounds when its stopping criterion tuned in manner analogous to MDOT.
Our comprehensive ablation studies of MDOT-PNCG affirm its robustness across wide range of algorithmic parameters. Benchmarking 24 problem sets of size 4096 in GPU environment demonstrate that our method attains high-procision, feasible solutions significantly faster than a representative set of existing OT solvers including accelerated gradient methods and advanced Sinkhorn variants in both wall-clock time and number of operations. Empirical convergence rates range between 1(44 and O(?2==1), where e is the optimality gap. For problem sizes up to n 18484, the empirical runtime scales Ónn²) for moderate precision and as O(n5/2) at worst for high precision. These findings establish MDOT-PNCG compelling alternative to current OT solvers, particularly in challenging weak-regularizztion regimes.
1 INTRODUCTION
When a statistical distance required for an event space equipped with metric, optimal transport (OT) distances, such the Wasserstein metric, provide intuitive means account for the inherent structure of the metric space. Consequently, fast, scalable, and accurate computation of OT distances is major problem encountered in various scientific ficlds. Example application arcas include point cloud registration (Shen et al., 2021), color transfer (Pitie et al., 2005; Ferradans el al., 2014; Rabin et al., 2014), shape matching (Feydy et al., 2017), texture mixing (Ferradans et al., 2013; Bonneel et al., 2015) and meshing (Digne et al., 2014) in computer vision and graphics, quantum mochanics (L6onard, 2012), astronomy (Frisch et al., 2002; Levy et al., 2021) and quantum chemistry (Bokanowski Grébert, 1996) in physics, and generative modeling (Gulrajani et al., 2017: 7; Genevay et al., 2018), reinforcement. learning (Ferns et al., 2004; Dadashi et al., 2021),

---

Published in Transactions on Machine Learning Research (06/2025)
and neural architecture search (Kandasamy et al., 2018) in machine learning. Exact solvers for the discrete OT problem encounter significant, computational hurdles in high dimensions, with theoretical complexity ((n and practical complexity Ö(n²) (Lee & Sidford, 2014; Pele & Werman, 2009).
Entropic regularization, as pioneered by Cuturi (2013), has mitigated challenges in scalability by regularizing the classical problem, thereby allowing approximate solutions in Ö(n²) time via the -uuhorn-KKopp (SK) matrix scaling algorithm. This advancement, together with GPU parallelization, has yielded substantial speed improvements, making it several orders of magnitude faster than conventional CPU-based solvers (e.g.: linear programming) in high dimensions (Peyré et al., 2019). However, these methods necossitate a delicate balance between regularization strength and convergence speed, a trade-off that can compromise the precision of the solution. Despite significant progress in recent years, many state-of-the-arl solvers still struggle to strike a better trade-off than aggressively tuned Sinkhorn iterations in practice (Dvurechensky et al., 2018; Jambulapati et al., 2019; Lin et al., 2019). Although they offer superior theoretical guarantees, their practical performance is often less compelling, particularly in lerms of speed and scalability. Existing algorithms either suffer from high computational complexity or do not take advantage of modern hardware capabilities, such as GPU parallelization (Tang et al., 2024). To understand and combat these challenges, we make the following contributions:
1. We empirically show that in GPU environment the decades-old Sinkhorn-Knopp algorithm for OT can still outperform many theoretically grounded recent OT algorithms in practice, especially when tuned with a seemingly unconventional stopping criterion formula proposed here (Fig. 5).
2. We introduce mirror descent optimal transport (MDOT), a method which generalizes temperature annealing in entropic OT (EOT) (Schmitzer, 2019; Feydy, 2020), and connects temperature annealing to mirror descent (Alg. 1).
3. We introduce an instantiation of MDOT that empirically improves speed and robustness to temper- ature (regularization strength) decay rate compared to e-scaling of Schmitzer (2019) (Fig. 2).
4. We show that MDOT can compute high precision, feasible solutions and its performance can be boosted by adopting a specialized GPU-parallel conjugate gradients (CG) algorithm developed here (Alg- 2); this method is highly competitive in practice, as we show empirically (Figs. 5, 6, 10-19).
The remainder of this paper is organized as follows. In the next section, we introduce our notation and the necessary background, followed by related work in Sec. 3. In Sec. 4.1-4.2, we introduce the MDOT framework and establish its connection to temperature annealing strategies, and make some practical recommendations. In Sec. 4.3, we introduce the non-linear CG algorithm to be used within MDOT as an alternative to SK. In Sec. 5, we benchmark various algorithms on upsampled MNIST (n = 4096) under L1 and squared L2 costs, and a color transfer problem in terms of wall-clock time, and further study the operation count dependence of the proposed algorithm on problem size n. Lastly, we present concluding remarks in Sec. 6.
2 Background
Here, we present our notation, the basics of EOT, and the necessary background on mirror descent and CG.
Notation and Definitions. We consider discrete OT, where the event space is finite with parti- cles and A, C 15 is the (n-1)-simplex. The row sum of an x matrix P is r(P) = P1 and the column sum is c(P)-
1 Given marginals r, An, the transportation polytope is written as u(r,c)= (P E Raxa (r(P)r,,ct c(()))-()) c). Division, exp and log over vectors or matrices are element- wise. Vectors in g are column vectors, and (æ, y) denotes the concatenation of and y. Vector and Frobenius inner products alike are given by (,-)). Hadamard product is given by A O B. An n x n diagonal matrix with æ E R" along the diagonal is written as D(x), and the vector formed by the diagonal entries of a matrix Q is diag(Q). Log reductions over the rows and columns of X E R are given by (X), = log( ( 1x) and (X) = log ( (xp(XT]1). The Shannon entropy of r E An is denoted H(r)- = (r. logr) with the convention that 0 loo = Under the same convention, the KL divergence (Ru(rlr') (r, log(r/'))) + (r' owr-rr1 1) for r, r' E Rg given r absolutely continuous with respect to r'.
2

---

Published in Transactions on Machine Learning Research (06/2025)
2.1 Optimal Transport
Given a cost matrix C' E [0, where Cij is the transportation cost between the 7th and jth particles, we study the EOT problem:
minimize PEu(r, (P,C)- -
(1)
7
where ? 0. Here, the regularization weight --1 is called temperature. The Lagrangian of (1) is strictly convex in P, which renders the solution (7) unique. (7) converges to solution of the unregularized OT problem as y + oo and admits the following form (Cuturi, 2013):
(u, 7;7) = exp(u: + b/ CGg),
(2)
where u,v E R", An optimal pair (u, v) minimizes the following convex dual problem (Lin et al., 2019):
gu 7, 7, = > Pis(u, 0; - (u, r) (v,c). c), u.u R.
minimize
(3)
if
where Vug rrr) -r and c The gradient norm naturally mcasures the constraint violation; L norm is lypically used to monitor convergence as it relates to the total variation metric between probability distributions (Altschuler et al., 2017). The Sinkhorn-Knopp (SK) algorithm (see Alg. 4 in Appx. A) can be used to minimize g with guaranteed convergence to dual-optimal variables as the number of iterations k + oo (Sinkhorn & Knopp, 1967; Sinkhorn, 1967; Franklin & Lorenz, 1989; Knight, 2008). Dvurechensky et al. (2018) showed that SK can be used to compute solution P E U(r, c) satisfying (P - P* CC) S with complexity Oen"/="), where P* an optimal solution of the unregularized OT problem. In particular, one first minimizes the dual objective until the L1 norm of its gradient is below a prescribed threshold, then applics the rounding algorithm of Altschuler et al. (2017) the infeasible plan given by (2) to obtain P E U(r,c) c) with an upper bound on the primal cost increase.
2.2 Mirror Descent
Consider a constrained convex minimization problem minper f(P), where the convex objective f f: is defined over some domain D, and the feasible set F S D, c.g., in OT, we take (P,C), and D = RRz"" Mirror descent (MD) method of Nemirovski & Yudin (1983) generalizes projected gradient descent (GD) by first choosing a strictly convex function h : D + R called the márror map, which induces a Bregman divergence (PIQ) 2 O between points P, Q E D as the difference between h(P) and its first
R
order approximation around Q:
(a(P)Q) = h(P) ((Q) - (VM()) P Q).
(4)
Then, the variable p(O) is updated as follows (Bubeck, 2015):
p(++1)= Vh-1 (Vh(P())) - A (((((())))
(5a)
p(t-1) min Dj (P(((tt))
(5b)
PeFnD
where (0) > is the step size. For h(P) = /22 projected GD is recovered as a special case since Vh(P) = P and Da
|PP /2 For certain problem geometries (e.g., if D is the simplex), choosing different h provably accelerates convergence. The idea is that in (5a), Vh: D + D* maps primal variables P to a dual space D', where applying a gradient update better aligns with the local curvature of the underlying geometry (Bubeck, 2015). Once the updated point is mapped back to primal space D via -h-1: D > D. (5b) performs a Breyman projection onto the feasible set F. Plugging (5a) into (5b) and rearranging yields an equivalent form with another interpretation (Beck & Teboulle, 2003):
p(++1) = arg (V PEFnD
P)
P)+ Aro (PIP(()))
(6)
Each MD step (6) seeks an update jointly maximizing the inner product with the steepest descent direction,
V p ((P(t),, subject to (i) a penalty for diverying from the current point p(t) and (ii) feasibility constraints.
A common choice is the negative entropy mirror map h(P)) (P, log P) in domain = R which yields Da, = Dwt. Throughout this work, we only use this mirror map." Hence, we write (5) explicitly
SSee Di Marino Gerolin (2020) for other diverggnce-regullarieddd optimal transport problems (using mirror maps besides negative Shannon entropy). They consider single step of (6), while we focus on multiple steps.
3

---

Published in Transactions on Machine Learning Research (06/2025)
for this case, where Vh(P) 1 + log P and Vh (RR exp(R 1).
p(+11 = p(t) о exp( - (8) (P()))
(7a)
p(t+1) arg arg DRL. (P)p(())))) PEFND
(7b)
Hero, we ro-weight p(t) entrywise by a factor of expl- A followed by a KL projection onto F, e.g,, if is the probability simplex, (7b) is a simple renormalization P + P/||PIl- We leverage the properties of this mirror map for the case F U((,c) c) design MDOT in Sec. 4.
3 Related Work
Accelcration of approximate OT solvers has been a focus of machine learning rescarch since the seminal work of Cuturi (2013). For instance, Altschuler et al. (2017) proposed the Greenkhorn algorithm, which greedily selects individual rows or columns to scale at given step and requires fewer row/column updates than SK to converge, but performs poorly due to low GPU utilization unless n is extremely large. Dvurochensky et al. (2018) proposed an adaptive primal-dual accelerated gradient descent (APDAGD) algorithm. Lin et al. (2019) later proposed adaptive primal-dual accelerated mirror descent (APDAMD) with theoretical guarantees. Lin et al. (2019) showed APDAMD to outperform APDAGD in terms of number of iterations, but not SK. Further, these tests only covered a high relative error regime (>50%); we investigate broader scope down 10- error in Section 5. Modest gains over SK in terms of number of iterations in the same regime were later obtained by Lin et al. (2022) via an accelerated alternating minimization (AAM) algorithm similar to that of Guminov et al. (2021). Notably, APDAMD applies mirror descent to the dual (3) of the EOT problem, while we apply it to the primal of the unregularized OT problem, i.e., problem (1) as Y + oo.
Application mirror descent to the primal of the OT problem has also been considered. Yang & Toh (2022) propose a more general inexact mirror descent algorithm (iBPPA), which they apply to the OT problem. Our approach differs from theirs in several ways, most notably how we decide to terminate the Bregman projection inner loop to solve (7b) inexactly (details in Section 4.2). Recently, Ballu & Berthet (2023) introduced Mirror Sinkhorn (MSK), which also takes gradient steps in the dual space in (7a), but instead of approximately projecting onto U(r, c) as in (7b) (as we do here), they alternately project onto U(: e) and u(r,-) via Sinkhorn updates, satisfying only half of the marginal constraints at a time. Our experiments in Soc. 5 suggest that this approach is efficient only in the low precision regime. Furthermore, MSK requires maintaining a running average of the transport plan at each iteration, precluding a straightforward ((n) memory implementation. In contrast, all algorithms presented here admit O(n) memory implementations,, assuming individual cost matrix entries can be computed on-the-fly in O(1) time. Xie et al. (2020) previously proposed an algorithm (IPOT) similar to MSK with fixed, even number of Sinkhorn updates (usually 2) following temperature updates. Alg. 3.5 of Feydy (2020) is also similar to these algorithms in spirit and is discussed thoroughly in Sec. 5.2. As discussed in detail in Sec. 4.1, well-known e-scaling strategies are also closely related (Kosowsky & Yuille, 1994; Schmitzer, 2019). Similar ideas have also been applied to the Wasserstein barycenter problem (Gramfort et al., 2015; Xie et al., 2020), which is left outside the scope of this work.
An alternative line of acceleration research focuses on multi-scale strategies, which employ clustering or grid- based methods to solve a series of coarse-to-fie OT problems and are sometimes combined with e-scaling (Schmitzer, 2016; 2019; Feydy, 2020). These are known to provide performance gains when the marginals are defined over well-chustered particles or in low-dimensional event spaces (Peyré et al., 2019). Lastly, in a similar spirit to our use of non-linear CG here, curvature-aware convex optimization techniques such as L-BEGS have also been considered for OT, e.g., Mérigot (2011); Blondel et al. (2018); however, scalability, precision and better performance than SK on GPUs has not been demonstratod simultancously our knowledge. Tang al. (2024) recently adopted Newton's 's method with Hessian sparsification to efficiently use second order information, but their key sparsification strategy is maximally utilized only on CPUs.
4 A Mirror Descent Framework for Optimal Transport
4.1 Temperature Annealing Mirror Descent
The OT objective (P, C²) has a constant gradient Vp C) C. Given step sizes (1) at time t 2 0, we obtain mirror descent. iterates (using the negative entropy mirror map) from (7)
4

---

Published in Transactions on Machine Learning Research (06/2025)
p((+)) arg min DKL (l|P() o expl exp(-
(8)
PeU(r,c)
Comparing the conditioning term p(t) o cxp( above with the Lagrangian closed form solution the EOT problem in (2), namely P(u, 0;7) = 11u1++1v" +1v -7C), suggests considering an MD step starting with p(t) of this same form. Specifically, we find the following result (proved in Appx. A.1).
Lemma 4.1. Given any initial u, t E R", 7 20020 0 and 4, > 0, we have
P(u", ; An) aig DKL. ()) Ped(r,c)
(9)
where u', v² E arg min g + r,.r.
That is, if we begin an MD step for the OT problem with any p(t) P(u, oy 117(1) (i.e., in the Lagrangian form but not necessarily optimal for the EOT problem), then the next MD iterate is the unique solution to
the EOT dual problem at -((+1) = 70 + (t) namely p((-1) = (*((((+11). Therefore, by induction, if we begin the MD iteration (8) for the OT problem with such an initial p()), then the MD iterates p(t) trace solutions of the EOT problems at increasing values of 711).
For initialization, recall that the FOT solution at 7=0 is the maximum entropy coupling, rc. lim,-o- arg mitpeetrr.) (P, M Thus, can initialize > = 0 and p(o) P(log log c) re Alternatively, in view of Lemma 4.1, it would suffice to use p(o) P(u, v; 0) for any u. E IR" at 70) = 0. That is, p(0) can be any rank-1 matrix in R50".
Finally, note that for the purpose of computing P(t-1) in (8), Lemma 4.1 ensures that any values u()) and wl0 can serve to provide a suitable p(() = P(u (a0))((0 (0)) Therefore, when implementing the MD iteration in (8) we do not need to converge to exact optimality each step (see Fig. 1). We formalize this discussion and provide further insight in the following proposition.
Proposition 4.2. Let (0) = and (((+1) = 700 + +4) which together imply ((+1)) = E' A!) Suppose p(0) E R850" is rank-1 and p(t) are computed via (8) for t 2 0. The following are true.
1. p(++)) PM (y(t+1)), i.e., the solution of the EOT problem (1) at 7((+1),
Given any u. R", p(++)) arg miurru) DKL Dxx(PlP(u,
3. (p - P',C)),) C) < Hmin (r, e)/ where p argmin (P, (P,C) C) and Hmmm(r = min (H(r),
4. (p(t)-((tt) p(ttt((())))))))))))))))))) ,) L (DRL (P(0)p((pp(tH)) + DKL. D
for all t 0.
A detailed proof is deferred to Appx. A.2. Next we comment on each statement in the order presented.
1. proves the equivalence of MD iterates to sequence of solutions to EOT problems with decreasing temperature, and therefore connects MD to temperature annealing.
2. shows that MD iterates p(++1) derived from a rank-1 initialization can be written independently of the previous solution p(t) (cf. 8), but rather only as function of ((+1), This means that we need not start from exact minimizer of the previous KL projection problem to reach the correct solution of the new problem. Given this result and Lemma 4.1, MD reduces in the dual space to a numerical continuation method (Allgower & Georg, 2003), in which we numerically trace the curve u² (7), v (7):²
(t) (+)) as arg min g(u, (0,
(10)
u.VER"
where each problem is initialized with some guess u', v E Rd. This is the key idea behind MDOT.
3. bounds the primal optimality gap at a given step t 2 1 in terms of the entropies of the marginals r, c. Since min (r),, Hc)))= < log n for r, c on the (n-1)-simplex, this tighter bound than the standard upper bound 7 log n used in prior work (Altschuler et al., 2017; Dvurechensky et al., 2018; Lin et al., 2019).*
"Slight abuse of terminology: in fact, there is a space of optimal curves, since g(u, 7) = u- 7) for any finite d. 3hhis log n term appears the time complexity varions algorithms, but is hidden in O-notation.
5

---

Published in Transactions on Machine Learning Research (06/2025)
Visualization
Space
q'n
n'n
9'1
9'
9%
W1
z- Id, v1
2 (u,
Path of Minima
MDOT iberates
Minimization.c of
Post Difference Extrapolation
Warm-start Guess
Figure In each iteration of Alg. 1, MDOT approximately minimizes v(u,v; as in (10) to compute duals vo(),v()) (i.e., "MDOT iterates"), starting initially with ? and increasing y by factor of q > per iteration until final value ye With incrcasing Y (weaker entropic regularization), the convex dual objectives defined over 2 become increasingly ill-conditioned (illustrated here in 1D by parabolas with increasingly high curvature) and therefore harder to solve. At each step, MDOT uses approximate solutions of prior problems to produce an initial guess for the next problem (left), and generates increasingly accurate warm-start guesses to mitigate the difficulty posed by ill-conditioning (right).
4. shows the one-step reduction in transport cost with equality; this is in contrast to the more standard analysis of MD where the improvement is bounded with an inequality.
4.2 A mirror descent method for optimal transport: MDOT
Using insights derived in the previous section, we construct the MDOT method shown in Alg. 1. As illustrated in Fig. 1, MDOT minimizes a sequence of FOT dual objectives g(>(")). At each step, MDOT uses prior approximate minimizers (t-kk), l(t-k) of (t)(t=)) for E [0, K] to extrapolate better initial guesses (warm-starting) for the next problem at + (sce Fig. 1, where K depicted).
The MDOT loop is summarized with the following steps:
A Algorithm MDOT 2, n. 1, 1)
L3: Choose stopping criterion εа for [[Vg(,, c)lli.
t 1, + min(1, 7t)
L4:
Obtain smoothed marginals F, ē to avoid zero or infinitesimal entries.
while True do
ea + H (r,, c)
L5: Rank-1 initial guess with positive C.C.
4:
(F, + (1 { (r.c) 44 12m
L6: Solve (10) given initial u', v', i.e. minimize g(r,5, ē) until (V S <a/2. By the triangle inequality, this implies [Vg[,, e) S Ea-
5:
if fft=11 then u', v' + log logë
10)) x(1) + Given initialization u',
6:
minimize g(7, F, until ||Vll] S sa/2
L7: Exit if y reached user input Y,, otherwise continue.
7:
7 = 7r then break
L8: Increase y for the next iteration.
A, + (g - 1),, + (-mmn(( + Ay, 71)
L9: Using previous (** (kk-))vvee (t-A) for k 2 0, extrapolate a warm-start, guess for the new KL projection (or, ROT dual) problem of minimizing (t+(1))), ,,, r.c).
u','' + WarmSttat(uu)) (t))............................
10: l 1+1+1
11: end while
In L13, approxination of p* (*())) is rounded via Altschuler et al. (2017) onto U(r, e); see Alg. in Appx. A.
12: P + exp(u + 1, (0) CC
13: Output +- (o.
Using (i) the guarantee that ||Vg(T, r. c)ll S Hmin(r, c//,,, (ii) the third statement of Prop. 4.2, and (iii
Lemma 7 of Altschuler et al. (2017), we obtain guarantee on the quality of the solution.

---

Published in Transactions on Machine Learning Research (06/2025)
Remark 4.3. 1f 2 5mmm (r, /2s, the output P E c)) of Alg. 1 satisfies (P-P", 0y S Ō())
The proof follows from a special case of Prop. 4.5, which is deferred to Sec. 4.2.2. Next, we discuss key MDOT steps listed above, how they differ from existing work, and tradc-offs of input parameters.
Stopping criterion (L3) and input p. Several prior methods such as IPOT (Xie et al., 2020), Alg. 3.5 of Feydy (2020) and MSK (Ballu & Berthet, 2023) that use MD ideas do not enforce any constraints on |Voll Rather, they only take a few Sinkhorn steps after Y is increased. Feydy (2020) takes increasingly bigger updates A = (q- (q-1)7 as we do in L8, such that when y is large, few Sinkhorn steps are insufficient to prevent divergence from the optimal curve u (1), (7)).. This effectively caps their precision as we show in Sec. 5. IPOT takes fixed steps (typically A, = 1), which helps trace the curve u² (1),, (1) closely, but roduces the regularization weight > more slowly; see Appx. G for further discussion and empirical data. On the other hand, MSK takes increasingly smaller A!" = O(1/VD), such that its convergence rate suffers (shown in Soe. 5). All three rely on Sinkhorn updates, while MDOT can benefit. from any (potentially better) convex optimization algorithus for minimizing g. The practical approach of Yang & Toh (2022) also uses a fixed 4, and Sinkhorn iteration. As the stopping criterion, they measure DKL DKL (Round r, c) after each Sinkhorn step to satisfy an inexactness condition for the KL projection, which facilitates their theoretical analysis. This incurs some O(r²) overhead per inner loop iteration. In contrast, we only check the IA norm of the dual gradient as in L6 (which can be evaluated in O(n) time during optimization), since by Proposition 4.2 maintaining the structure P(u, v; throughout ensures that P (1t) can be recovered thooretically, regardless of how incxact provious KL projections were.
Another line of work runs (some version of) a single iteration of MDOT (Altschuler et al., 2017; Dvurechensky et al., 2018; Lin et al., 2019). In these works, &d x (log n2/7,, while we adapt it according to the entropy of r, e using the third statement of Prop. 4.2. In Appx. C, the use of Hmin (r, c) instead of log n is shown to yield speedups of order log H/Hain (r, c). Our use of p 2 1 is also novel and controls how closely the curve u" (1)), 00 (7) is traced. We study theoretical benefits and practical trade-offs in Sec. 4.2.2.
Marginal smoothing (L4). Our "smoothing" of the target marginals T and e (mixing in the uniform distribution) follows prior work by Dvurechensky et al. (2018) and Lin et al. (2019). This step helps provide convergence guarantees for certain choices of minimization algorithms for the dual g in L6. Since the mixing weight used is proportional to P-, it gradually decreases with the temperature in our case. This scheme smoothes marginals more aggressively in earlier iterations of MDOT when y is small. In Appx. D, we empirically show the performance benefits of this variable smoothing.
EOT dual minimization (L6). The algorithm for minimizing 9(7, r. c) is left unspecified here for gener- alily; for example, Sinkhorn iteration can be used. In Sec. 4.3, we introduce a new algorithm (PNCG).
Temperature decay and input q (L8). L8 effectively sets -((+1) = () which follows prior work (Schmitzer, 2019; Feydy, 2020). With larger MD steps (higher q), the extrapolation of duals for warm- starting is over a longer range, which degrades quality and poses a tradc-off, investigated in Sec. 4.2.1.
Warm-starting duals (L9). Our explicit warm-starting approach in Sec. 4.2.1 is novel to our knowledge and was left unspecified in Alg. for generality. We show that e-scaling as implemented by Schmitzer (2019) and Feydy (2020) amount to an implicit warm-starting, which is shown to be less effective than ours.
4.2.1 Warm- Strm--tarring KL Projections
Assume that at each prior temperature value, we obtained the dual-optimal u ())()), o (7(()) without error. How should u, e be initialized for 2(++17 A simple, memory-efficiett approach is to consider a Taylor expansion around recent 7 to predict u u=("""(( (+(++))), v² (7((1))) Letting 2 = (u, v) to reduce clutter:
0z* z (++++)) = z () + By (7 (1(1)) - +...
(11)
As we cannot compute 0z" /ay analytically, we use a numerical approximation (backward finite differencing) 0z x (0 (2) 2" (7 (t-))) / (1-1) and keep only the first two terms in (11):
x ()(((+2) = x ())) A (1-1) 2 (0) - x ()((-11))
(12)
7

---

Published in Transactions on Machine Learning Research (06/2025)
gradient L norm at initialization vs. y
Suboptimality gap vs, number O(n") operations
Sinkham projection
PNCG projection
q
10"
4 2
10-
10-2
10
10
10
cost
10
10-3
10-5
MDOT escaling
10-1
10
10-
2
24
2
119g 19g*1||
10a
10
10
10-3
10
10
L3 cost
10-5
10-5
10-
10
10-1
2'
23
22
10 107
10
10
10
101
10
10
10
10
Inverse temperature Y
Number of operations
Figure 2: Comparison of the MDOT warm-start proposed in Sec. 4.2.1 to e-scaling. Curves show the median over 36 upsampled-MNIIs problems (n=4096) under La (top) and L3 (bottom) distance costs (see Sec. 5 for details). In all experiments, 1.5 and 7 2 For the Lı cost, Y 2 and for the L3 cost, 7t 21..
In contrast, the E-scaling approach of Schmitzer (2019) and Feydy (2020) maintains reparamatrized dual variables ž == x/4 as the temperature decays. Rewriting (2) in terms & reveals that E-scaling amounts to prodicting 2" (t+1)) (y(t+1) // (1)=) i.c., simply scaling the dual variables instoad of modelling the trajectory of 2z with a Taylor approximation, which we argue is better approach.
In Fig. 2, present empirical study with varying step sizes Dy (g 1)y by ablating q, where the advantage of (12) over the E-scaling warm-start of Schmitzer (2019) is demonstrated. On the left, MDOT warm-start initializes each dual problem closer to the solution than the e-scaling approach. The quality of initial guesses increase markedly with decreasing temperature (left); at high temperatures dual problems are initialized very close to the solution with the gradient norm just small multiple of the target Ea- contrast, the e-scaling warm-start stays relatively fixed. For the same decay rate q, this translates about 10x gains in convergence speed precision (mid-right). The performance gap widens for slow temperature decay (lower q), as MDOT benefits from reduced Taylor approximation errors given smaller step sizes Ay,
4.2.2 KL Projection Stopping Criteria
Our assignment Eα X T for some p 2 1 in L4 of MDOT departs from the conventional wisdom of choosing Ea x (Altschuler et al., 2017: Dvurechensky et al., 2018: Lin et al., 2019). Here, we provide justification for this departure. Consider first the exed-temperature problem (3) for simplicity. Building the results of Cominetti & Martín (1994), Weed (2018) showed in his Prop. 4 and Thm. 5 that there is both uniform bound (7) P*, C) < logn/r (slow rate), and a fast asymptotic rate O(exp(-7K) which takes over for large enough 7, where the constant K > 0 is -ebendedd Taking these as a starting point, the following remark generalizes the third statement of Prop. 4.2.
Remark 4.4. For any constant p E 10) and OT problem given by (r, c, C), there exists a no such that for any 7 2 70o we have (>) _p, CP* < Hmin (r, c)//P.
8

---

Published in Transactions on Machine Learning Research (06/2025)
Suboptimality gap vs. number of O(n²² operations for Sinkhorn and MDOT
Sinkhorn
MDOT (q=2"%)
10a4
10
PNCG
Sinkhorn
10
10
10
10
10-2
p=
1-25
10
p= 1.75
1-1.5
10
2.0
10-7
P 2-0.0
10:
10a
10a
10a
10a
10a
Number of LogSumExp reductions
Figure 3: Ablation of stopping criterion parameter (Sec. 4.2.2) for the SK algorithm (left) and MDOT with Sinkhorn and PNCG as KL projectors (right). The SK algorithm (left) is called by running MDOT (Alg. 1) with " = Y where higher precision is achieved by increasing Mr. Results show the median over 36 random problems from the upsampled MNIST dataset (n = 4096) with the L, cost.
That is, below some temperature % stronger bound Hmin (r, c)y for some p replaces the uniform bound Hmin (r, et Thus, the SK algorithm (see Alg. 4 in the Appx.) can be tuned (via the p parameter in Alg. 1) to enjoy rate substantially better than Onn² log n/E²) given by Dvurechensky et al. (2018).
Proposition 4.5. Sinkhorn iteration, as instantiated by calling Alg. 1 (L6) with p E [1, o0) and sufficiently large ? = "/r = 5/5min (r, /2e, returns a plan P EE((, c) sulisfying (P- P*, PO,C) S a O(2)) in at most
0 ( n² H Hmin (r, 1/p
p=1
arithmetic operations.
(13)
This result consistent with the empirical findings of Jambulapati et al. (2019), who noted "The [tuned] Sinkhorn algorithm converged at. rates much faster than the predicted e rale on all experiments, outper- forming all other methods, which we believe merits further investigation." We believe Prop. 4.5 shods some light on this phenomenon, and further present an ablation of p in Fig. 3 for SK and MDOT algorithms.
For the SK algorithm, Fig. 3 (left) verifies the insight derived from Prop. 4.5. The choice p 1 is better at low precision, but the trend gradually shifts in favor of higher with (sufficiently) higher k. That is, for sufficiently low temperature 7- it is advantageous to reduce the gradient norm error tolerance, from Hmin/y to Hmin/7" for > 1. In contrast, MDOT is more robust to the p parameter in the high precision regime (right). Moreover, the use of PNCG projections (Alg. 2) for KI. projections MDOT (L6 of Alg. 1) provides speedup of 2 3x over Sinkhorn projections. PNCG is introduced and discussod next.
4.3 Preconditioned Non-linear Conjugate Gradients for KL Projections
SK converges more slowly at low temperatures (Kosowsky & Yuille, 1994). For faster alternative, we develop Alg. 2 based non-linear CG (NCG) methods (Fletcher & Reeves, 1964; Nocedal & Wright, 2006). which now briefly review. Given objective g. NCG takes descent directions (0)) (g(2(0)) and ( + vy(0)) + ,(k) (k-1), and iterates z(++1) + (k) + (k)) where (k) is the step size. Optimal (k) has a closed-form solution for quadratics, but for general lineiinear objectives, line search is necessary to find suitable step sizes (k). Many formulas for computing (Kk) exist; for quadratic objectives, they are equivalent and guarantee convergence at most n' iterations, where n'' S n is the number of distinct eigenvalues of V2g. Further, the objective decreases faster if eigenvalues are tightly clustered (Stiefel, 1958; Kaniel, 1966; Nocedal & Wright, 2006). For example, the Hestenes-Stiefel formula sets (Nocedal & Wright, 2006):
g(4) (k) (y(k-1), (k) (Vg(1) Vg 1), (x:k1)) 1)}
(14)
9

---

Published in Transactions on Machine Learning Research (06/2025)
A practical way to further improve the convergence rate of CG methods is via preconditioning. By making change of variables z = M- given some symmetric positive-defiite matrix M, one reduces the condition number of the problem or tightens the clustering of eigenvalues for improved convergence (ideally, M- ~ V)g-1). We refer the reader to Hager & Zhang (2006b) for further details on CG mothods.
For the EOT problem, recall the 1st and 2nd order derivatives of the dual objective g in (3) at z = (u,
-
(D(r(P)) P D(c(P))) )
Vg= (r(P)-r, c(P)-c), v²g =
(15)
A typical choice of preconditioner M, known to be effective for diagonally-dominant matrices (Golub & Van Loan, 2013), is the diagonal approximation of the Hessian, which yields the following descent direction:
$ D(diag((V' (V2 'Vg 1000(110 (r(())' )(P)) 12n- -12n.
(16)
((IP))((P))
Observe, however, that if at any point in the optimization r(P) or c(P) has infinitesimal entries, numerical instabilities may occur when evaluating 8. We propose using the Sinkhorn direction, 8, in place of ø:
r s log r(P)
log C )
(17)
This has the benefit that it can be evaluated via mumerically stable LogSumExp reductions, e.g., see lines 2-3 of Alg. 2. The Sinkhorn direction can be understood as the result of an alternative diagonal preconditioner, namely M- = D Vg)), since 8 = -M- Furthermore, for any sub-optimal (u, v), we have
(8, Vg)
DKL. (r())) + DRL. (ri.(r/r(P)) (r/r(P)) + DRL. DR..(c(P))))) (c(P)lc) (ca((cl((P))
0,
and therefore is also a descent direction. Empirically we find that this Sinkhorn preconditioner results in improved numerical stability. Finally, note that near the solution (for r Rs r(P) and c c(P)) we have
) log r(P) log c(P)
C
~ r(P) c(P) 12
(18)
where we have used log 1-1 for ~ 1. Therefore, near the solution, the Sinkhorn direction s approaches the direction å obtained using the common preconditioner from the diagonal of the Hessian.
Plugging the preconditioner = = D gg/)) into the preconditioned Hestenes-Stiie Fletcher, 1996), we take g(k) in L7 of Alg. 2:
Algorithm 2 PNCG 7, C, C,r,c,, Ea)
formula (Al-Baali
1: (u, v) < x, p- 02., 8 + 0
2: log (ogr(P): + u + 1,," - -C)
g00) (Vg(*) (g (k-1) (())
(19)
3:
LSE.
7C)
(V Vg 1) 11)-1
4: Vg - (r(P) (eP))-rc((P rrc(P)) -c)
5: while |Vgll1 > Ed do
where (k)) is the Sinkhorn direction as in (17), 300 = O, (0) Ozn- Observe that (b))) above simply replaces g term in the numerator of (14).
s llgr-llog (oggr(P, l log c(P))
p + + ++33p 3p
See (19)
if (p. Vg) 2 0 then
We defer details of the line search in L11 of Alg. 2 to Appx. B, but note that by design, the proposed line search only carries out the same form of LogSumExp reductions as the log-domain stabilized SK algorithm
9:
p s > Reset CG if not a descent dir. end if
10:
11:
o, log (P)), log c( + LineSearch v.
12:
(u,v) 4- (u, v) + ap
(Alg. 4 in the Appx A), so that its output is reused when evaluating the Sinkhorn direction next iteration (see L11 of Alg, 2). This also allows for fair comparison of the two algorithms' performance. Indeed, Fig, 4 plots the runtime of the two algorithms in terms of LogSumExp evaluations; PNCG outshines SK empirically, especially at lower temperatures (higher 7). Further, to see whether the added numerical stability of the newly proposed Sinkhorn preconditioner comes at a performance trade-off, we implement an alternative stabilization scheme for the diagonal Hessian preconditioner. In particular, for this alternative we use s only if the voctor //r(PP)c/c(())))) has any entries outside the range [0.01, 100) and otherwise assign s given by (16) in L6 of Alg. 2. The results shown in Fig. suggest that, on the contrary, the Sinkhorn pre- conditioner also provides performance benefit over the diagonal Hessian in addition to numerical stability.
c) 15: Output 2 2 (
13:
Vg (r(P) - (r(P))r,c(P)
(17) at the
14:
end while
10

---

Published in Transactions on Machine Learning Research (06/2025)
Gradient norm (scaled) vs. number of O(n²) operations (upsampled MNIST dataset)
L cost
L3 cost
10
Sinkhorn
2
PNCG: Sinkhar
10c
PNCG-Dlagene (Stabil (zed)
10a
10c
10
100
10a
10a
10a 10a
10t
10a
10a
Number of Exp reductions
Figure 4: Comparison of KL projection algorithms (used in L7 of Alg. 1) over the upsampled MNIST dataset with L1 (left) and L3 (right) costs. Algorithms evaluated are the SK algorithm, the newly proposed PNCG given Alg. 2, and a variant (soc text). In all 36 problems, n 4096. p = 1.5 and g = Each curve shows the convergence behavior, at one specific lemperature 1/7, in terms of the median nunber of LogSumExp reductions (x-axis) until gradient norm (y-axis) reaches below target dual gradient norm Ea(?).
5 EXPERIMENTS
In this section, we first detail the MNIST experimental setup in Figs. 2-4. Then, we describe an additional color transfer task we use for benchmarking. Next, performance evaluations in terms of precision vs. wall- clock time are discussed given the results over 4 scts of problems shown in Fig. 5. In Appx. H, we add 20 more problem sets from the DOTmark benchmark of Schrieber et al. (2017) showing similar results both in terms of wall-clock time and operation counts. Lastly, the dependence
problem size n is investigated in Fig. 6. All experiments were performed on an NVIDIA GeForce RTX 2080 Ti GPU with 64-bit precision. Exact costs (P*, C) are evaluated using the implementation of the CPU-based algorithm of Bonneel et al. (2011) from the Python Optimal Transport (POT) library (Flamary et al., 2021), which was run on an Intel Xeon Silver 4110 (2.10GHzz) CPU.
5.1 Experimental Setup
Upsampled MNIST. In line with prior work (Cuturi, 2013; Altschuler et al., 2017; Lin et al., 2022; Tang et al., 2024), we first consider the MNIST dataset, where each pixel represents an event and each image a probability distribution. Unlike prior work, we form higher dimensional problems by upsampling the original 28 x 28 images to be 64 x 64 (with bilinear interpolation) so that n = 4096. Cost matrices are constructed by mcasuring the L1 or squared L2 distances between pixel locations on a 2D grid, and dividing all entries by the maximum distance value so that all entries of C lie in [). The probability of each pixel obtained by flattening the pixel intensity matrices
is proportional to its intensity value; marginals r, and subsequent L1 normalization. To select random problems, without replacement, and compute the OT distances between the first and second halves of the samples. Our selection of n = 4096 is driven by the objective of conducting a large number of tests per con configuration to ensure statistically significant results, rather than by any inherent limitations the algorithm. In fact, our MDOT code supports the use of on-the-fly CUDA kernels to evaluate entries of the cost matrix on the go using the PyKeOps package (Charlier et al., 2021). In this case, MDOT leaves an O(n) memory footprint (with both Sinkhorn and PNCG projections) rather than O(n²); it has been verified to scale to much larger problems (n & 100, 000).
sample 2m images from the dataset
Color Transfer. For the color transfer problem, each image is viewed point cloud in RGB space (pixel locations carry no importance). Cost matrices C' are constructed by mcasuring the L1 or L3 distances between pixels in RGB space and dividing all entries by the maximum distance. Marginals r, e are taken to be uniform over An With the help of GPT-4, we prompt DALL-E 2 to generate 20 vibrant and colorful images with intricate details or patterns. To match the dimensionality of the upsampled MNIST problem set, we downsample the original 1024 1024 images 64 64 so that 4096. Once again, cost matrix entries normalized to lie in [0, 1].
11

---

Published in Transactions on Machine Learning Research (06/2025)
Optimality gap vs. wall-clock time
upsampled MNIST, L1 cost
upsampled MNIST, L3 cost
10 -
10-2
10
10-3
10-4
MDOT-PACC
NDOT-Sindom
10
Sinesom 1-11
pmmudapati et (2019)
10-1
10
MSK
Sisnom (2:1:11
10-3
10-7
APONED
Peyay No.
10"4
10
10
10-7
color transfer. cost
color transfer, L3 cost
10-1
10-4
10-2
10-1
10
10
10
10-
10-4
10-7
10-1
10-4
10-1
10-2
10-1
10 10a
10-1
10-1
10a
10c
Wall-clock time (5)
Figure 5: Wall-clock time error benchmarking over the upsampled MNIST (top) and color transfer (bottom) problems using Li (left) and L3 (right) distances cost functions. Each marker shows the median time to converge (over 18 random problems) for each algorithm at a given hyperparameter setting, which controls the precision level, and the error PP-P",Ch C) after rounding the output of the algorithm onto U(r, c) with the exception of Alg. 3.5 Feydy (2020); see text. Vertical dashed lines show the median time taken by the CPU-based exact solver with 80% confidence intervals.
5.2 Wall-clock Time Comparisons With Prior Work
In Fig. 5, we present wall-clock time benchmarking MDOT (with both Sinkhorn and PNCG projections) against existing GPU-parallel algorithms on the upsampled MNIST and color transfer problems. All bench- mark methods were implemented in PyTorch and run on the GPU. For MDOT, we use q=21/3,p 1.5 and 7= in all experiments. For the closely related Mirror Sinkhorn (MSK) algorithm of Ballu & Berthet (2023) the variable step size schedule prescribed by their Thm. 3.3 is used in our implementation. For Alg, 3.5 of Feydy (2020), we decay temperature at a rate 0.7-1, which interpolates their fast (q=0.5-1) and safe (g=0.9-1) settings. For AAM (Guminov et al., 2021), Mirror Prox Sherman Optimized (Jambulapati et al., 2019) and APDAGD (Dvurechensky et al., 2018), each implementation closely follows an open-source NumPy implementation. Our PyTorch implementation was verified to produce identical results to the pub- licly available NumPy code. We additionally attempted comparison with APDAMD (Lin et al., 2019) and PDASMD (Luo et al., 2023), but observed extremely long convergence times for 4096 and omitted the results. For further details on the implementation of benchmark methods, see Appx. F.
While MDOT optimizes (3) to satisfy convergence criterion following each temperature decrease, Alg. 3.5 of Feydy (2020) performs single (symmetrized) Sinkhorn update instead, i.e., it does not minimize the sequence dual objectives sufficiently despite taking increasingly large gradient steps in the dual space (cf. 5a-5b). This causes an accumulation of projection errors and results the algorithm hitting precision wall. Their debiasing option for estimating the OT distance via Sinkhorn divergences (introduced by Ramdas al. (2017)) fares slightly better and is used here to comprise stronger baseline, albeit this approach does not find member of U(r, c), which may be strict requirement in some applications. MSK also runs single row www/column scaling update after a temperature decrease, but takes increasingly smaller steps and maintains a rumning average of transport plans to ensure convergence. It performs well at low precision, but shrinking step sizes slow it down, so that it exhibits O(n2= convergence behavior. Sinkhorn iteration (log-domain the Appx. A) benefits substantially from sctting p 1.5 rather than p 1 at sufficiently low temperatures for L2 costs (see also See. 4.2.2). APDAGD underperforms SK with p = and AAM performs similarly to it. Mirror Prox Sherman Optimized of Jambulapati et al. (2019) overtakes
stabilized, see Alg.
12

---

Published in Transactions on Machine Learning Research (06/2025)
Number of O(n²) operations problem size for different datasets and cost functions
upsampled MNIST, L, cost
upsampled MNIST, L3 cost
10a
Inw. temperature
103
215
10a
Number of O()²)) operations
10a
10a
color transfer. L1 cost
color transfer. L3 cost
10a
10a
10a
10a
10a
102 10a
10
104
10a
10'
Problem size n
Figure 6: Problem size dependence of MDOT-PNCC convergence over the upsampled MNIST (top) and color transfer (boltom) problems using L, (left) and L2 (right) distance cost functions with E {26, 2 212 215). Each marker displays the median over 20 problems and shaded areas show 75% confidence intervals. Dashed lines show (n) 5/2 for visual comparison, with selected to intersect the purple curve the lowest n.
SK (p= 1.0) one case only (top-left) in the high precision range. Mcanwhile, MDOT-Sinkhorn enjoys laster convergence than the more competitive SK (p = 1.5) owing to warm-started temperature annealing, especially in the high precision range (near the solution). MDOT-PNCG is the quickest to converge in all cases. The performance gap with its close second, MDOT-Sinkhorn, grows with higher precision.
5.3 Empirical Dependence Problem Size of MDOT-PNCG
Our last set of experiments investigates the practical dependence of MDOT-PNCG the problem size Over the same problem sets as Fig. 5, change n from 36 to 16, 384 by up- down-sampling images. The values are selected to be approximately equally spaced the behavior of MDOT-PNCG for a range of final temperature values "i E {26 29, 21,,222 2153. At medium precision (green and blue), we find that the algorithm behaves no worse than ((n)) in practice as implicd by the flatness of the curves. As seen visually, with higher precision (roughly 5-decimals) with or E {212, 215), the proposed GPU-parallel algorithm behaves roughly as O(n5/2) at worst and even better for some of the problems in practice (see the reference line in Fig. 6). These should be compared to the 5(2 theoretical and (n²) practical complexity of CPU-based exact solvers (Pele & Werman, 2009; Lee & Sidford, 2014). Indeed, Figure 5 suggests that the CPU-based solver offers a strictly better precision-time trade-off beyond an optimality gap of around 10 10-5-10-5 for this value of n = 4096 and this particular CPU-GPU setup; however, in Table of the Appendix, we show that the trade-off skews favor of MDOT-PNCG when n is increasod.
logarithmic scale. In Fig. 6, we plot
6 CONCLUSION
In this work, we first presented general procedure, MDOT, for computing OT distances with high precision and described its relation to well-known temperature annealing strategy (e-scaling). MDOT employs a novel warm-starting of the sequence of EOT dual problems encountered in temperature annealing, which empirically shown to be highly effective compared to existing approaches. In addition, specialized non-lincar CG algorithm was developed as an alternative to Sinkhorn itcration and was shown to be more effective al low lemperatures (under weak regularization). Over 24 different. problem sets, the combined MDOT-PNCG algorithm outperforms aggressively tuned Sinkhorn iteration and many other recent baselines in terms convergence of the primal suboptimality gap measured wall-clock time. The algorithm was also shown to behave well with respect to the problem size. Interesting directions for future research include the theoretical convergence behavior of PNCG, bounds the gradient norm of warm-started init initialization,
13

---

Published in Transactions on Machine Learning Research (06/2025)
better warm-starting methods, and adaptive dual problem stopping criteria. The development of faster KL projection algorithms and adaptive temperature decay schedules are also of interest; see also our recent work in this direction (Kemertas et al., 2025).
Acknowledgements
Amir-masooud Farahmand acknowledges the support of the Natural Sciences and Engineering Research Council of Canada (NSERC) through the Discovery Grant program (2021-03701). Mete Kemertas acknowl- edges the support of NSERC through the CGS-D scholarship. Resources used in preparing this research were provided, in part, by the Provincc of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institule.
14

---

Published in Transactions on Machine Learning Research (06/2025)
References
Mehiddin Al-Baali and Robert Fletcher. On the order of convergence of preconditioned nonlinear conjugate gradient methods. SIAM Journal on Scientific Computing, 17(3):658 665, 1996.
Eugene L Allgower and Kurt Georg, Introduction to numerical continuation methods. SIAM, 2003.
Jason Altschuler, Jonathan Niles-Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. Advances in neural information processing systems, 30, 2017.
Marin Ballu and Quentin Berthet. Mirror Sinkhorn: Fast online optimization on transport polytopes. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 1595-1613. PMLR, 23 23-29 Jul 2023. URL https://proceedings mlr.
preas/v202/ban ha1222//baall2aaaa html.
Amir Beck and Mare Teboulle. Mirror descent. and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167-11) 175, 2003. URL https://ww.seeeneeddreee sciencedirect.
Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. Proceed- ings of the Twenty-Fiigg First International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pp. 880 889. PMLR, 09 11 Apr 2018. URL https: / //proceedings.nt mmr. mlr. prres//84/blondee
html.
Olivier Bokanowski and Benoit Grébert. Deformations density functions molecular quantum chemistry.
Journal of Mathematical Physics, 37(4::1553 1573, 1996.
Nicolas Bonneel, Michiel van de Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. In Proceedings of the 2011 SIGGRAPH Asia Conference, SA '11, New York, NY, USA, 2011. Association for Computing Machinery.
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasserstein barycen- ters of measures. Journal of Mathematical Imaging and Vision, 51:22-45, 2015.
Stephen Boyd and Lieven Vandenberghe. Conver optimization. Cambridge university press, 2004.
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends0 in Machine Learning, 8(3-4):231-333 2015.
Benjamin Charlier, Joan Foydy, Joan Alexis Glaunes, François- David Collin, and Ghislain Durif. Kernel operations on the gpu, with autodiff, without memory overflows. The Journal of Machine Learning Research, (2(1))3457-3444 3462, 2021.
Roberto Cominetti and J San Martín. Asymptotic analysis of the exponential penalty trajectory linear programming. Mathematical Programming, 67:169 187, 1994.
Thomas M Cover. Flements of information theory. John Wiley & Sons, 1999.
Marco Cuturi. Sinkhoru distances: Lightspeed compulation of optimal transport. Advances in neural information processing systems, 26, 2013.
Robert Dadashi, Leonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal Wasserstein imitation
learning. In International Conference Learning Representations, 2021. URL https://openrevvee: net//forumidd--t-ss22:gg
Simone Di Marino and Augusto Gerolin. Optimal transport losses and Sinkhorn algorithm with general convex regularization. arXiv preprint arXio:2007.00976 2020.
Julie Digne, David Cohen-Steiner, Pierre Alliez, Fernando De Goes, and Mathicu Desbrun. Feature- preserving surface reconstruction and simplification from defect-laden point sets. Journal of mathematical imaging and vision, 48:369-33 382, 2014.
15

---

Published in Transactions on Machine Learning Research (06/2025)
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Com- plexity by accelerated gradient descent is better than by Sinkhorn's algorithm. In International conference on machine learning, pp. 1367- 1376. PMLR, 2018.
Stephan Eckstein and Marcel Nutz. Quantitative stability of regularized optimal transport and convergence of sinkhorn's algorithm. SIAM Journal on Mathematical Analysis, 54(6):5922-55945 2022.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite Markov decision processes. In Pro- ceedings of the 20th Conference on Uncertainty in Artificial Intelligence, UAI '04, pp. 162-169, Arlington, Virginia, USA, 2004. AUAT Press.
Sira Ferradans, Gui-Song Xia, Gabriel Peyré, and Jean-François Aujol. Static and dynamic texture mixing using optimal transport. In Scale Space and Variational Methods in Computer Vision, pp. 137-148, Berlin, Heidelberg, 2013. Springer Berlin Heidelberg.
Sira Ferradans, Nicolas Papadakis, Gabriel Peyré, and Jean-François Aujol. Regularized discrete optimal transport. SIAM Journal on Imaging Sciences, 7(3):1853 1882, 2014.
Jean Fcydy. Geometric data analysis, beyond convolutions. PhD thesis, Université Paris-Saclay, 2020.
Jean Feydy, Benjamin Charlier, François-Xavea Vialard, and Gabriel Peyré. Optimal transport for diffco- morphic registration. In Medical Image Computing and Computer Assisted Intervention- MICCAL: 20th International Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part I 20, pp. 291-299. Springer, 2017.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Cham- bon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. POT: Python optimal transport. The Journal of Machine Learning Research, 2251:3357733578 2021.
Reeves Fletcher and Colin M Reeves. Function minimization by conjugate gradients. The computer journal, 7(2):149-11 154, 1964.
Joel Franklin and Jens Lorenz. On the scaling of multidimensional matrices. Linear Algebra and its appli- calions, 114:717-735, 1989.
Uriel Frisch, Sabino Matarrese, Roya Mohayaee, and Andrei Sobolevski. A reconstruction of the initial conditions of the universe by optimal mass transportation. Nature, 417(6886)) 262, 2002.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608-1617. PMLR, 2018
Gene H. Golub and Charles F. Van Loan, Matriz computations. JHU press, 4th edition, 2013.
A. Gramfort, G. Peyré, and M. Cuturi. Fast optimal transport averaging of neuroimaging data. In Sebastien Oursolin, Daniel C. Alexander, Carl-Fredrik Westin, and M. Jorge Cardoso (ods.), Information Processing in Medical Imaging, pp. 261 272, Cham, 2015. Springer International Publishing. ISBN 978-3-319-19992-4.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of Wasserstein GANs. Advances in neural information processing systems, 30, 2017.
Sergey Guminov, Pavel Dvurechensky, Nazari Tupitsa, and Alexander Gasnikov. On combination of alternating minimization and Nesterov's momentum. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8886-3898. PMLR, 18-24 Jul 2021. URL https: https://prroeeedigss , mlr ppres/vv3ss https///11/g/mmmn11111 gww.gguuua..ai..mmmmmoo html.
William W Hager and Hongchao Zhang. CG DESCENT, conjugate gradient method with guaranteed descent. ACM Transactions on Mathematical Software (TOMS), 32(1):113-137, 2006a.
William W Hager and Hongchao Zhang. A survey of nonlinear conjugate gradient methods. Pacific journal of Oplimization, 2(1):35 58, 2006b.
16

---

Published in Transactions on Machine Learning Research (06/2025)
Arun Jambulapati, Aaron Sidford, and Kevin Tian. A direct Õ(1/c) iteration parallel algorithm for optimal transport. In Advances in Newral Information Processing Systems, volume 32. Cur- ran Associates, Inc., 2019. URL Paper. pdf.
wwwerips.cc/ppaeeeeeeeeeeee cc/paper http/pperrfiee//pppppeeeeeeee ftles//apper/219//tii
024266996CC12222988888884442
Kirthevasan Kandasamy, Willie Neiswanger, Jeff Schneider, Barnabas Poczos, and Eric P. Xing. Neural architecture search with Bayesian optimisation and optimal transport. Advances in neural information processing systems, 31, 2018.
Shmuul Kanicl. Estimates for some computational techniques in linear algebra. Mathematics of Computation, 20(95):369 378, 1966.
Mete Kemertas, Amir- massoud Farahmand, and Allan Douglas Jepson. A truncated newton mothod for optimal transport. In The Thirleenth International Conference on Learning Representations (ICLR), 2025.
Philip Knight. The Sinkhorn Knopp algorithm: convergence and applications. SIAM Journal on Mutria Analysis and Applications, 30(1):261-255 225, 2008.
Jeffrey J Kosowsky and Alan L Yuille. The invisible hand algorithmr Solving the assignment problem with statistical physics. Neural networks, 7(3):477-499) 490, 1994.
Yin Tat Lee and Aaron Sidford. Path linding methods for linear programming: Solving linear programs o (vrank) iterations and faster algorithms for maximum flow. In 2014 IEEE 55th Annual Symposium on Foundations of Computer Science, pp. 424-433. IEEE, 2014.
Christian Léonard. From the Schrödinger problem the Monge Kantorovich problem. Journal of Functional Analysis, 262 (0)44418779-192))) 2012.
Bruno Levy, Roya Mohayaee, and Sebastian von Hausegger. A fast semidiscrete optimal transport algorithm for unique reconstruction of the early Universe. Monthly Notices of the Royal Astronomical Society, 506 (1):1165 1185, 06 2021.
Tianyi Lin, Nhat Ho, and Michael I. Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3982-3991 PMLR, 09-15 Jun 2019. URL http::::////oceeedippe/ .mw.. wtw.ar.resss//////// /wt.////////...1/// html.
Tianyi Lin, Nhat Ho, and Michael I. Jordan. On the efficiency of entropic regularized algorithms for optimal transport. Journal of Machine Learning Research, 23(137):1 2022. URL /ttt::/jjmr....e waa.ar.rrrroo.oeeoeee.eee orggpaperss v23/20-277 777html.
Yiling Luo, Yiling Xie, and Xiaoming Huo. Improved rate of first order algorithms for entropic optimal transport. In Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. 2750.2 PMLR, 25-27 27 Apr 2023. URL https://proceedings mlr wtwppp/sos//1102000/0 html.
Quentin Mérigot. multiscale approach to optimal transport. Computer Graphics Forum, 30(3):1583-1553 1592,
2011.
Arkadi Nemirovski and Dmitry Yudin. Problem complexity and method efficiency in optimization. John Wiley & Sons, 1983.
Jorge Nocedal and Stephen J. Wright. Numerical Optimization. Springer, New York, NY, USA, 2e edition, 2006.
Ofir Pele and Michael Werman. Fast and robust Earth Mover's distances. In 2009 IEEE 12th international conference on computer vision, pp. 460-467. TEFE, 2009.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications data science. Foundalions and Trendso in Machine Learning, 11(556):355-600 607, 2019.
17

---

Published in Transactions on Machine Learning Research (06/2025)
François Pitie, Anil Kokaram, and Rozenn Dahyot. N-dimensional probability density function transfer and its application to color transfer. In Tenth IEFF International Conference on Computer Vision (ICCV'05) Volume I, volume 2, pp. 1434 1439 Vol. 2, 2005.
Julien Rabin, Sira Ferradans, and Nicolas Papadakis. Adaptive color transfer with relaxed optimal transport. In 2014 IEEE international conference on image processing (ICIP), pp. 4852 4856. IEEE, 2014.
Aaditya Ramdas, Nicolás García Trillos, and Marco Cuturi. On Wasserstein two-sample testing and related families of nonparametric tests. Entropy, 19(2):47, 2017.
Bernhard Schmitzer. A sparse multiscale algorithm for dense optimal transport. Journal of Mathematical Imaging and Vision, 56:238 259, 2016.
Bernhard Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM 2019.
Journal on Scientific Computing, 41 A3)3A
Jörn Schrieber, Dominic Schuhmacher, and Carsten Gottschlich. Dotmark - benchmark for discrete optimal transport. IEEF Access, 5:271 -282, 2017.
Zhengyang Shen, Jean Feydy, Peirong Liu, Ariel H Curiale, Ruben San Jose Estepar, Raul San Jose Estepar, and Marc Niethammer. Accurate point cloud registration with robust optimal trans- port. In Advances in Neural Information Processing Systems, volume 34, pp. 5373 5389, Cur- ran Associates, Inc., 2021. URL 1boff58....004494411119966666111
wnw.prooreii.nn....r............ htppai/..errrpp....ppppe..... httpr//s///pppppp..eee////////// wtt.f//s////ppe///////aa pdf.
Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathemalical Monthly, 74(4):402-45) 1967.
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathemalics, (2):3343 348, 1967.
Eduard L Stiefel. Kernel polynomials in linear algebra and their numerical applications. Not. Bur. Standards Appl. Math. Ser, 49:1-22, 1958.
Xun Tang, Michael Shavlovsky, Holakou Rahmanian, Elisa Tardini, Kiran Koshy Thekumparampil, Tesi Xiao, and Lexing Ying. Accelerating Sinkhorn algorithm with sparse Newton iterations. Tn The Twelfth International Conference on Learning Representations, 2024. URL http:pp:///peeeeeeeeeeeteeeeeeee Kuj5gVp5GQ.
Jonathan Weed. An explicit analysis of the entropic penalty in linear programming. In Conference On Learning Theory, pp. 1841-1855. PMLR, 2018.
Philip Wolfe. Convergence conditions for ascent methods. STAM review, 11(2):226-233 1969.
Philip Wolfe. Convergence conditions for ascent methods. ii: Some corrections. SIAM review, 13(2):185-188, 1971.
Yujia Xie, Xiangfeng Wang, Ruijia Wang, and Hongyuan Zha. A fast proximal point method for computing
exact Wasserstein distance. In Proceedinggs of The 35th Uncertainty in Artificial Intelligence Conference, volume 115 of Proceedings of Machine Learning Research, pp. 433 453. PMLR, 22 25 Jul 2020. URL attps://proceeding: mlr. ppwss/v115/xie200000 html.
Lei Yang and Kim-Chuan Toh. Bregman proximal point algorithm revisited: A new inexact version and its inertial variant. SIAM Journal on Optimization, 32(3):1523 1554, 2022.
G Zoutendijjk. Nonlinear programming: a numerical survey. SIAM Journal on Control, 4(1):194 210, 1966.
18

---

Published in Transactions on Machine Learning Research (06/2025)
Appendix
The Appendix is organized as follows:
Appendix A provides the proofs of theoretical results.
Appendix B presents the line search routine used for Alg, after primer on line search.
Appendix C investigates our use of Hmia (r, c) when deciding the KL projection stopping criteria in contrast to the more standard use of log n.
Appendix ablates the design decision MDOT (Alg. 1) to vary the KL projection stopping criterion based on the present temperature, and shows its performance benefit.
Appendix E compares MDOT-PNCG the network simplex solver in higher dimensions (n 16384) Ihan the experiments in the main text (n 4096), and shows belter scaling.
Appendix F discusses implementation details for baseline algorithmss
Appendix provides experiments vs. the IPOT algorithm of Xie et al. (2020).
Appendix provides additional benchmarking on 20 problem sets from the DOTmark benchmark (Schrieber et al., 2017).
A Proofs
Here, we provide proofs for the theoretical results in the main text; Lemma 4.1 in Appx. A.1, Proposition 4.2 in Appx. A.2 and Proposition 4.5 in Appx. A.3.
A.1 Proof of Lemma 4.1
Lemma 4.1. Given any initial u, v E R", 7 0 and A, > 0, we have
P(u", 0 3 A,) arg min DKL (PIP(u, w;7) (xpp(-A,C)). Pel((n.))
(9)
where u*, o² aarg g(7 + Ty,
Proof. First, emphasize the following observation, from which Lemma 4.1 follows immediately using P(u, v;)) - P(u, u, A,).
Observation A.1. Given any incliall u, u E Ran and ((u,w;1) defined as (2), we have
(*( arg min DKL (PIP(u,)))) P(u", "ou","""") ,7);), where u², o0 E arg min g(r,r,)). rell(r.e)
(20)
In words, (i) minimizing the objective g in (3) given any initial u, R" amounts to KL projection of P(u, 7;7) onto U(r,c) c). and (ii) the set of matrices P(u, 7;7) = exp(u1 + Io -CC} all have the same KT. projection in U(r, c). The intuition here is that since the elements u, of u, v simply rescale the jth row and 5th column of P(u, v; y), and the unique projection onto U(r,o c) is the optimal scaling (with r(P) = r and c(P) = c), the specific initial scaling of rows and columns the conditioning matrix P(u, u; is irrelevant.
We now start the formal proof by deriving the relationship
P(u", v²;r) arg (PP((,((uu)))))) + K]. Pel((r.e)
(21)
given any initial u, u E R and u", x² E arg min g(u, 7: Here we have introduced constant K which, of course, does not effect the arg min. We will choose simplify the derivation below. Note that K depend the given quantities, such u, v, and on, but not the unknown P. Given the
19

---

Published in Transactions on Machine Learning Research (06/2025)
definition of DKL in Sec. 2 for un-normalized P E R50", observe that
DRt. Drt.(P|P(u,, (( 0;7)) = (P, log P - log P(u, v; + (P(u, 0; - P.1)
(P. P- u1 10 T +사C) + (P(u, - P.1)
(P, IP.-1+ 7C + log P) - r ()) - (0, c()) + [| ;; Mlh
As shown below, it is convenient to use K = (u, r) + (u,c) - || P(u, ;: 7)1,, for which we find
DEL (PIP(,, 0;4) + K = (P, 12.-14 + C + log P) - (u, r(P) - r) - (p) c).
(22)
We represent the equality constraints on P. that is, r(P) and c(P) = c, using Lagrange multipliers a and / to form the Lagrangian
C(P. 8) = [DKi.
0; (PIPu,, + log P) - (u + a, r(p) -r) - (v- B, c
+ KK] (()) r) - (B,c(P) - c)
=
(23)
where we have used DKL. + as in (22). Recall that the Lagrange dual function is given by (Boyd & Vandenberghe, 2004):
inf C(P, a. B)
PE
where R55"" is the domain of the KL divergence objective in (21). If strong duality holds, we have (P(P(()))) + K].
sup inf a) = min P
(24)
a.BER PER
In this case, we conclude that strong duality holds since Slater's condition is satisficd, i.c., there exists a strictly feasible P; namely, the independence coupling Te.. See Ch. 5.2.3 of Boyd & Vandenberghe (2004).
Now, to find the point-wise minimum of L(P, or, 9) with respect to P, we require:
dC = -Gi + log Pij un 0: 1/3 33 0.
aP.j
Therefore the minimizer P must have the form
P = P(u + a, B:7) exp(a; + u; + B 01- -Cig)-
(25)
for some a and B. To climinate the variable P from inf PERYX" L(P, ax (9), we plug the form above into the Lagrangian. The first term on the right hand side of (23) is then
(P*, 1P.-14 +1+ 70
log PT) (P*, - + 70) (P*, (u (u 11 + 1(1 -
-B)= B)T -0,,
-PP" 1) + (u a. (v B.
Using this in (23), we find
C(P*, az B) = -Pp
, 1) + +
=- E exp + B,)
(wy
-O.) + (t B,c).
(26)
where and are the only unknowns. Without loss of generality we can change variables in both p Plu+a, a, u + B:7) and [(P", a. B) to a' a + u and B' 9 u, in which case the Lagrangian is equal to
C(P(dl.
Σ exp(a + + Cin) + (a'rr) r) + (B',c) c)
= =-g(d',8
(27)
20

---

Published in Transactions on Machine Learning Research (06/2025)
where g(a', (1:7) is the dual objective for the EOT problem in (3). By strong duality and (24) we conclude that the solution is given by p = P((x", ;);,,, where ,", , satisfy:
gr mag (a, a,
o.BeR"
arg mammingga g 8; a.BER"
That is, the solution is identical to the optimal dual solution in (3), and therefore Observation A.1 holds.
Piu, Lemma 4.1 follows immediately from Observation A.1 using P(u, v; exp(-A, P(u, v17 + A,).
A.2 Proof of Proposition 4.2
Proposition 4.2. Let 70)) = and (((+1) - -00 + A which together imply 7 (1+2) Ee- A Suppose p(0) 1250" is rank-1 and p(t) are computed vio (8) for t 0. The following are true.
p(++1) P" (y(t(2)) i.e., the solution of the EOT problem (1) at (((+1).
2. Given any u.u RR" p(++)) arg peu((r)).)) (PL(P)P(u, P(u, ( ((+1))).
3. (P()) - C', s Hmin (r, e)/7(0), where Pa E arg aag pau C) and Hmin = min (H(r), H(c)).
4. p(t+)), 0) N. (pl0)
(p(0)p((+1) p(t+))) + DRL (p((+))p(()) for all t 2 0.
Proof. We prove each of the four statements in order.
A.2.1 Proof of the 1st statement.
Here, we first provide a step by step derivation for the following expression for p(t++) for t 0:
p(++1) = arg min DEL (PIP(u(>?1)) o (;))); = (+()()
Peu(r.c)
where (0) = by construction, and (*()) = P(u" (7), for u (1), v² (7) E arg gin g(7), i.e., solutions of the EOT primal (1) and dual (3) at y. Starting with a rank-1 positive matrix p(0) FēT P(log F, log ē; 0) given vectors F, with positive entries, we obtain using MD updates (8) and Lemma 4.1:
p()) = arg min DKL C A())) Pel(r,C)
log a
argmin DKL (PIP(log log (og Pellin.))
(Since 70 and 7((+1) = (0) + A by construction.)
P" (?("") 0² (1))))))
(By Observation A.1.) (By definition.)
= (').
Repeatedly using the same and continuing the iteration:
p()) arg min DKL (PIP(u (" (0)) PeU((..c)
p(3) = arg min DEL (PIP(ri() u (2)))) = p*))
Pel((..))
(28)
p(++) = arg min DEL (,"",, = =P*(((0t2)) pel((..e)
Hence, each iterate p(++)) (+()((+)) for t 220. 0.
21

---

Published in Transactions on Machine Learning Research (06/2025)
A.2.2
Proof of the 2nd statement.
The result follows immediately by applying Observation A.1 to cach iterate in (28) to replace u' (1), w² (7) terms by arbitrary u, E IR..
A.2.3 Proof of the 3rd statement. First, write the following helper lemma.
Lemma A.2 (A mirror descent bound for linear objectives). Given linear objective function f(P) = (P, C),
initial point p(o) F, an optimal point p and any T > 0, sequence [P(N) obtained via (6) satisfies:
-(P())- - - P((
(29)
Σ A
Proof. Recall the definition of p(++1) from mirror descent iterates in (5a):
p(++1) vi-1
For any PED. D,
f(p((+1) f(P) (V)(P p((t)) P)
(since is linear)
A) (Vi(P()) vh(Be-1), p(tt)) - PP)
(due to (5a))
AJ (Vi(Pio) - vh(p(tt1) p(ttz) - P)
(by Lemina 4.1 in Bubeck (2015))
A (Da(PIP(t) Dn(P(p(ttt Du p(tt)t))pp()) (by Eq. 4.1 in Bubeck (2015))
(Da(PIP(t)- -D((PI)ttt))
(since D₁ 2 0)
A
which implies
(F(p(tt))))(P)) ( DA, ((pptl)))))).
The above incquality proves monotonic improvement in each step t once we take P = P(t), Letting P= P* taking telescopic sum and dividing both sides by IT-1 (s) arrive at:
IT-1 A (
(P D (P
1--1 (3))
ΣT A11)
(P(
Σ Als)
which implies (29) since improvement. is monotonie and the first term the LHS is convex combination of objective values.
By Lemma A.2, for p(o) E Eu(r, we have
Dr (P (P(O)
(p(0) P". S
Σ1-1 (!)
Given 7=)(1) -1-11 A9)
remains to show that D₁ (2x (P'pp())) & Hmin e).
22

---

Published in Transactions on Machine Learning Research (06/2025)
Recall that for the negative entropy h(x) Σ, F, log have Dn(xly) DKL (x Suppose we take p(0) = re:
DRL(P"IPLU)
= Σ
Pij (log P. - logricy)
17
Σ
Pij (log Pt logri ri - llo
-H(P)) Σ logra Epj Σ log Pi Poj
Elogr:
Pūj
-H((P) Σ r; 1;oog,; Σ Cj log
(since p pr e c))
H(r) + H(c) - H(P*)
max ((r), (c)) + min(((r), - H(P")
$ min(H(r). (()).
The last incquality holds since H(P) 2 H(r) and H(P) 2 H(c) for any P ((r,c) (Cover, 1999), which logether imply H(P) (ax(H(r),HH()))))
Proof of the 4th statement. First, note that given h(P) = Lis Pag log P.j have Vh(P)y 1 + log Po and VA (Q)as exp(Qu - 1). Then, given the definition of p((+ from mirror descent iterates in (5a):
p((+1) = VI- (Vi(P(0)- - ((V(())())
exp p(() - -)
= =exp (u" (1))))11 + 11 (T)1)) - (7 + (()).). (given u' (7(1),, 0² (x((1)) E arg min g(u, v; (t))) exp (u" (7 1 + 10 (7(1)) - (()
(30)
In the third equality, we used the known closod-form expression (2) expand p(t).
In the special case that the feasible set J = U(r, c)2
(p(t), C) - (p(++)) C)
-V]((P()) p(t) p((+1) -p
(since (P,C) is linear, Vpf(P) = C.)
1 (Vh(p(!)- vn(Ptt1), p(t) p((+1))
(due to (5a))
1000
(Vh(P(() (M(p(++1), p(o) -p((t2)
(see below)
(Dx(P(0))ptt)) Dr. (m (p(l-)))(())
(by definition of the Bregman divergence as in (4))
A
(Dk((P())pttt))) + DKL ( (p(+)))p())
40
To why the third equality holds, observe that p++1 P exp(ă; + +0 for some optimal update vectors â²" 6² E R given the closed-forms (2) and (30). Then, for any P. P' E (((rrc)).
(Vh(pitt)))PP P- PP)
E (1+ + log pt-1 + at + ()) P!)))
=(V((Pe+)) , P- pr) + Σ iū Σ (P.) (Psy-PPP) -P,))+ + Σ 07 [(,, Ply))
-(Vi(D(H)) (ū" - r) + (0* c c)
(since P, P' E EU(r, e) by construction)
-V P-P
23

---

Published in Transactions on Machine Learning Research (06/2025)
Algorithm 4 Sinkhorn(z, C, r, c,
Algorithm 3 rr, (Altschuler et al., 2017)
1: (u, U) t 2
1: X +- ((x) with a rrr((P)) A1
2: log r(P) + u + (1, yC)
F+ XP
3: while [Vgl] = |r r-- r(P) || Ea do
Y Dyy) with y c/c(FF) A1
logr- logr( (P)
F' F FY
V log (ul, -C)
5: etr, t r(F),, erte + C - c(F")
6: logr(P) t u +
-0)
6: Output G t- F' + err, / ||err,l1
7: end while
8: Output 2 t (u, v)
A.3 Proof of Proposition 4.5
In the remainder of this section, the L norm [P]|, of a matrix denotes the I norm of the vectorized form of the matrix, and not the L1 matrix norm.
First we state the following lemma, which is a simple combination of Lemmas and by Weed (2018).
Lemma A.3 (Entropy increase from mixing (Weed, 2018)). Lel T1,72, T3 E An and (1 E)ri + ers, where E E (0, 1]. We have,
H(r2) (L))))(()) + +(11 - log < H(r1) + + (1+ + log
(31)
Next, we provide simple proof for Remark 4.4
Remark 4.4. For any constant p E (1) and OT problem given by (r. C)2 there crists a > 0 such that for any 2 "0o, we have (P) - Px C2 < Hmin (r,,))/// c)/7P.
Proof. Recall from Thm. 5 of Weed (2018) that the quantity (P" - - P',C) decays at an exponential rate with increasing y for sufficiently large 7. Since the exponential function exp{->K} decays more quickly than - for any constant K and finite P, we conclude that there exists some constant % > 0 such that
(p' - C) (p))
& Hmin(r, c)/7"
(32)
for all optimal transport problems given by r, c, C provided that y 2 70-
Proposition 4.5. Sinkhorn iteration, instantiated by calling Alg. 1 (L6) with p E [1, 00) and sufficiently large ?1- 7/- min (r, (r,,e)/2 /2s, relurns plan PE Uu(r c) salisfying (P C*, + Ö(s2) in at most
O Hmin
Hmin (r, c) 1/p
E
pil P
arithmetic operations.
(13)
T - CC} after the termination of the main loop (before rounding in L13) of Alg. 1, which takes place after a single outer loop iteration in this setting, since 7i = "r by construction. Since B is the output of Sinkhorn iteration (Alg. 4), it lies on the simplex, do its row and column marginals (specifically, Furthermore, B is the unique optimizer of the EOT problem over U(r', c') due to Prop. 4.2 and the fact that it has the form Byj = exp(u; + Uj - :i
Proof. Let B E U(r',c' d') be the transport plan P(u, v) = exp(u1 + le
have c' c(B) č E A, from Alg. 4).
1 H(P) 7f
== arg min mig (P,C) PeM(r', .c)
(33)
Sinkhorn iteration returns a solution u, v such that
| - rB)], + c < 6/12
|Voll1
|r r(B) + -c-llccl30 c(B)ll
r Fll, + [[--r(B]]l + ||c ē|1 + |2 c(B)I),
(triangle incquality)
| - Fl], +| č|, + +
24

---

Published in Transactions on Machine Learning Research (06/2025)
Then, given mixing weights e'/4 in L5 of Alg. 1:
VVgll S
(34)
Now, we make the following dcfinitions:
B = Round B, r, c), the rounding of B onto u(r,c) via Alg. 2 of Altschuler et al. (2017), returned
by our Alg. 1,
B* E pru .t) C), an optimal plan in the feasible set u(r', cc),
P E pentr (P. C), an optimal plan the feasible set ((r,c)) c).
We have that,
=B--BCC- nnxi + (B--B",)+ (B"
PP,
(since B, B E Anxn.)
$ 2 2 B + (BB
(Hölder's ineq., given Ci. E [0, Vi,j E [n])
< [[Va|| + (B BB)B',C) + (B' PB')
(by Lemma of Altschuler et al. (2017))
< [[gg11 + Hmin (r',c) + (B' P",C)
(given (32-33), assuming Yr sufficiently large)
Se'+ Hmin(r' (F) (B-P*,C), P',
(35)
where B is any transport plan in u(r', c'). We take B to be the "shadow" of Pt in the sense of Definition 3.1 Eckstein & Nutz (2022), under the discrete metric. In other words, letting
B= arg min 1P
rEU((",e))
and noting that the 1-Wasserstein distance under the discrete metric is equal to the total variation (TV) distance, the first equation in Lemma 3.2 of Eckstein & Nutz (2022) yields the equality (x) below:
1l = =TV(B,P") (+) TV(r,r)) + TV(C, = VVlli
(36)
Then, continuing from (35);
Hmin (T', + - (B-P",CY
(B-P*,CC p' C- S
= + Amin(r',c)
+ (B - (BP'", p'
2 1nxa)
(since B, P E Aaxn.)
5 + Hmin (T',
+
B 1-ppl
pell
c 2 1nx
V
3
+ Amin(r',e)
5 =
(given (34-36))
2
7
$ c + (min(r,
3
+ E'
(1 log(n///)))
(by Lemma A.3)
2
of
op
5/ (r, c)
+ Ō(75 20)
(since e' = Hmin(r.c) in L4 of Alg. 1)
27F
7
1/p
= =
(since op 55lmmn (r,c) c2/3))
by construction)
The computational complexity of the algorithm follows simply from the same line of reasoning as Thm. 1 and Thm. 2 of Dvurechensky et al. (2018). In particular, they show that Sinkhorn iteration converges in O(R/e")
25

---

Published in Transactions on Machine Learning Research (06/2025)
steps, where R O(1)) (H (r, ( //////// in our case. The complexity result O Hmin (T, //// )
follows since g' O(Hmin(r: c)n; ") = (), and each Sinkhorn step costs (n²²).
B An Efficient Line Search Algorithm
In Section 4.3, we developed the PNCG algorithm, which required a line search procedure. Here, we develop the line search used in our implementation, following short background on relevant aspects of line search in numerical optimization.
B.1 Background: Line Search
Given a descent direction p(*) E R,, i.c., a direction that satisfies (p v)(x()))) & O, line search algorithms aim to find an appropriate step size a, where x ((+11) + xx() +apl(). Perhaps the most well-known of desirable properties that a step size a should satisfy at any given optimization step are the Wolfe conditions (Wolfe, 1969; 1971). Given d(a)= f(x(k) + (p)))))
Φ(a) -00) S (
(37a)
a s'(a) 2 (2)
(37b)
where 0 < c1 < <22 < 1 and (37a) and (37b) are known as the sufficiend decrease and curvalure condi- tions respectively (Nocedal & Wright, 2006). It is well-known that given step sizes satisfying the Wolfe conditions and descent directions pik) that are not nearly orthogonal to the steepest descent directions (*)), line search methods ensure convergence of the gradient norms to zero (Zoutendijk, 1966; Wolfe, 1969; 1971). Instead of satisfying (37), some algorithms or theoretical analyses consider evact line search, where a² E arg min ((a), which has unique closed-form solution for quadratic objectives with posi- live definite Hessian. However, a rule of thumb for general non-linear objectives is to not spend too much time finding a" (Nocedal & Wright, 2006). Hager & Zhang (2006a) proposed approrimate Wolfe conditions, derived by replacing the σ(a) and (00) terms in (37a) with g(a) and q(0), where g is a quadratic interpolant of & such that q(0)- = (00), d((0)- d' d0) and da(a) = +'(a):
- 12c-1)o''00 2 d'(a) 2 02 (0).
(38)
A key advantage of replacing (37) by (38) stems from the fact that one only needs to evaluate o' rather than both & and ø' to check whether the conditions are satisfied, thereby halving the amount of computation necessary per iteration in cases where their evaluation has similar computational cost.
Bisection is a line search strategy with convergence guarantees when the objective is convex. One simply maintains a bracket [ato: ani), where &' ((no) < 0 and (a(m)) > 0, and recursively considers their average and updates either endpoint of the bracket given the sign of ø' ((ani + o2.)
B.2 PNCG Line Search
To perform line search in PNCG (Alg. 2), we adopt a hybrid strategy combining bisection and the secant method to find ag satisfying approximate Wolfe conditions (38). Given alo abi, the secant method computes the minimizer of a quadratic interpolant q that satisfics (((a)) = o' (aw) and (a(m)) = ø'(ana) as follows:
0166 (ama) - On (ano)
O'oe
d'(ona) ('(o1o
(39)
Thanks to the convexity of the objective g, by ensuring d((aw) < 0 and ('(om) > 0 with simple algorithmic checks, we can guarantee that on < Osec < Ohi- Thus, the updated bracket is guaranteed to be smaller once we replace either of a% or ahi by asec for the next bracket given the sign of d'(ase). If & behaves like a quadratic inside the bracket, the sccant method converges very quickly, but convergence can be arbitrarily slow otherwise. For this reason, we simply average the bisection estimate and asec for a less aggressive but more reliable line search that still converges quickly, i..., Ohybrid 0.50 + 0.5(Oni 022)
26

---

Published in Transactions on Machine Learning Research (06/2025)
Sinkhorn iteration using Hmn logn (4g cost)
Sinkhorn iteration using Hain vs. log (L) cost)
Hea 0. 0.03logn
Hime - 0.03logn
1.4
10l0gg
Roc 0. 0.30
Haa 0.30lo
Hi... 0.30logn
9
Hov 0.00kogn
199
15
14-
11
124
d'e:Hoally
8' Vogony
E-lognly
32 of LogSumExp reductions
Number LegiuumExp reductions
Figure 7: We control the problem parameter Hmiu(r,c) over synthetically sampled OT problems (r,c, C) and show that using the entropy-aware stopping criterion can yield substantial performance gains, with the gap between the two approaches growing proportionally to the gap between Hmiu (r, e) and log n. The experiments carried out use ? E (24,25... .24) to control precision. We sample 18 problems for each y value and plot the median along both axes. The results are consistent between L1 (left) and L3 (right) costs.
Evaluation of o' has computational complexity Onn)) as does a single step of Sinkhorn's algorithm (given by the two LogSumExp reductions seen in L5-6 of Alg. 4):
(ax))= = (Pu,, (Paa) - r) + (po: c(Pa) - c),
(40)
where (Pu. Pe) is the descent direction. Since evaluating ø' requires the computation of r r(Pa) and c(Pa) for the new matrix Pa == exp((u+ s + 1(v + ap.) -C),, the last step of the line search readily carries out the LogSumExp reductions necessary for computing the Sinkhorn direction in the next step of PNCG (see L11 of Alg. 2). Observe also that at the next PNCG iteration, d'(0) can also be computed in O(n) time rather than On?²) since ((Po). are already known from the last line search step of the previous PNCG iteration. With these important implementation details in place, we find that the average number of ø' evaluations necessary to find that satisfies (38) is lypically between 1.5 - 2.5 for the PNCG algorithm. While the approach outlined here is easy to implement (including as batch process) and works well in practice, better line scarch mothods may further benefit Alg. 2.
C Entropy-aware Stopping Criteria on the Dual Objective Gradient Norm
Here, we show the effect of choosing Hmin(r, c) over the weaker bound log in L4 of Alg. 1, where the stopping critcrion E' is selected. To control the problem setting H (r, c), we construct synhetic problems by randomily sampling r from the simplex via a Dirichlet distribution constructed to meet a larget entropy level (r) as a fraction of the maximum possible entropy log n. The column marginal e is simply taken to be the uniform distribution I,/n, so that Hmin e) = (r). Cost matrices are constructed by sampling n points R from multivariate normal distribution and assigning Cip- = || zyl"; for {1,2], before entrywise division by max;j Ci to ensure Cij E [0,1].
Fig. 7 illustrates the effect of this choice by ranging Hrin(r,)) c)/ logn E {0.03, 0.1,0.33.33 Towards the RHS of the plots, we observe an improvement in precision roughly proportional to log logn/ H/Hain (r, c) for the same number of operations, which agree with our complexity result O(n² H Hun (r, c)) for p = in (32) vs. the O(n² logn/-3) result by Dvurechensky et al. (2018).
D Variable vs. Fixed Smoothing of the Marginals in MDOT
As discussod in Sec. 4.2, MDOT smoothes the marginals r, e (by mixing in the uniform distribution) with a weighting factor that tracks the temperature. Since MDOT anneals the temperature, this means that the smoothing weight is higher in earlier iterations of MDOT. In particular, the mixture weight gradually
27

---

Published in Transactions on Machine Learning Research (06/2025)
Error number of iterations (Sinkhorn)
Error vs. number iterations (PNCG)
q=22
2-2
4-2
17a4
00-1
19
15.
smosihin smoothing
smoothing
Number
Number of interrations
12'
Figure 8: The variable smoothing scheme has almost no impact on the convergence behavior of MDOT- Sinkhorn (left) for both settings g = (rapid temperature decay) and q = 21/8 (slow temperature decay). MDOT-PNCG (right) enjoys a speedup of nearly 2x under rapid decay and a more modest speedup under slow decay. All curves show the median over 36 sample problems from the MNIST dataset (L1 cost).
decays from Mmin (r, c)r,e)// to Hmin( 4)/47/ given input parameter p 2 1. Here, we study the effect of this design choice as it influences the convergence of two KL projection algorithms used in L6 of Alg. 1: Sinkhorn iteration and the newly proposed PNCG algorithm (Alg. 2). The approach is benchmarked against a bascline that fixes the smoothing weight at Hmim (r, c)/ 4n, all throughout instead. For these experiments, we fix p 1.5 following our experimental setup in Sec. 5. Fig. 8 shows that while MDOT-Sinkhorn is largely unaffected by this design choice, MDOT-PNCC enjoys notable speedup from variable smoothing. We thus conclude that the approach provides a performance benefit.
E Comparison with CPU-based Solver for Higher n
In Figure 5, we displayed via a vertical line the time taken on a CPU by the network simplex solver from the Python Optimal Transport package of Flamary et al. (2021). Clearly, a faster CPU would benefit the network simplex and a faster GPU would benefit. the MDOT family of algorithms, which makes them difficult to compare in a fair setup. However, given the known Ó(n²) complexity of the network simplex and the Ö(n?) - ((2255) empirical dependence of MDOT-PNCG seen in Fig. 6, we suspect that the precision-speed tradc-off posed by MDOT-PNCG improves with increasing n. In Table 1, we provide a comparison of the two algorithms, with several values of "r used for MDOT-PNCG and increased from 4, 096 to 16, 384. We observe the same trend across both In and La cost functions on the MNIST dataset; for a comparable level of relative error achieved by MDOT-PM. PNCG (at a fixed 74), the speedup over the network simplex solver is better for n = 16, 384. We expect the trend to continue with higher n, and note that MDOT stands to benefit from ongoing rapid developments in GPU hardware innovation, as well as faster projection algorithms than PNCG; see also our concurrent work in this direction showing substantial speedups (Kemertas et al., 2025).
F Details of Baseline Algorithm Implementations
Here, we provide details and sources on the implementation of various algorithms shown in Fig. 5. Our implementations of other algorithms will be open-sourced for transparency.
Mirror Prox Sherman Optimized (Jambulapati et al., 2019). For this algorithm, the source code is originated in the NumPy code at this repository. The owner of the repository notes that this NumPy implementation is based on a Julia implementation by the original authors, which was provided in a private exchange. The code used in this paper is a PyTorch adaptation of the NumPy code and has been verified to produce identical output as the NumPy version over multiple problems. The algorithm was called with entropy factor parameter set the default 2.75 in all experiments. The number of iterations for the algorithm was varied from 2 to 215 to achieve different levels of precision.
28

---

Published in Transactions on Machine Learning Research (06/2025)
APDAGD (Dvurechensky et al., 2018). For APDAGD, similar strategy was used, except with this code repository. A PyTTrrch version of the original NumPy code was written and verified to produce identical output. For different levels of precision, the parameter of the algorithm was varied from 2 to 2 For smaller E, non-convergence was observed.
AAM (Guminov et al., 2021). The implementation based on NumPy code by the original authors at this repository. Py'orrc version was verified to produce identical output for GPU execution. The e parameter varied from 2 to 2-10, 10, For smaller E, numerical errors
encountered.
Feydy, Alg. 3.5 (Feydy, 2020). The implementation is based on the algorithm presented in the original work. For different levels of precision, the number of total iterations was varied from 2 to 212, Beyond the upper bound, numerical errors were observed. As it produced betler estimates than the alternative, the algorithm was called with debiasing turned on; hence, the error (P P C) was instead measured in absolute value as (P- P", P) for this algorithm only. Scaling ratio was to an intermediate 0.7, which is between the listed 0.5 (fast) and 0.9 (safe) settings.
Sinkhorn (Cuturi, 2013). A log-domain stabilized implementation was varied from 2 to 2 for L1 distance cost and to 2 for L3 distance cost. Stopping criteria were given by our formula in L4 of Alg. 1, and the results obtained by calling Alg. 1 with 7/1 so that the algorithm terminates after a single KT. projection via SK iteration.
used. For different precision levels,
Mirror Sinkhorn (MSK) (Ballu & Berthet, 2023). The implementation is based on the algorithum pre- sented in the original paper. For different levels of precision, the number of total iterations was varied from 2 to 2
Table 1: Comparison of the exact Network Simplex solver with MDOT PNCG method on MNIST transport problems under L, and L2 costs. Relative error is computed 100 (P P C) (P'.)P.) CC)
n 4, 096
16, 384
Cost Fn. Algorithm
RelErr %
Time (s)
Speedup
RelErr %
Time (s)
Speedup
Net. Simplex
0.0000
5.45
1.00
236.46
1.00
28
16.556
0.18
30.55
22.580
1.72
137.24
L1
MDOT-PNCG
29
0.167
2.28
2.39
0.585
26.94
8.78
212
0.002
11.85
0.46
0.002
318.58
0.74
Net. Simplex
0.000
11.74
1.00
728.40
1.00
2a
26.877
0.52
22.66
23.827
5.53
131.83
MDOT- PNCG
212
3.166
3.85
3,04
3.372
38.12
19.11
215
0.044
39.38
0.30
0.303
294.12
2,48
29

---

Published in Transactions on Machine Learning Research (06/2025)
cost
4 cost
V007-5 1O0T-INCC
Optimaliyy 0eT 1--
0611
3441
Number operations
Figure 9: Optimality gap (P- P", C) of rounded (feasible) plans vs. number of O(n²) operations for TPOT
(Xie al., 2020) and MDOT (with Sinkhorn and PNCG for KL projections, using p = 1.5, q 21/3 as in Section 5). Both methods are run up to final - 2 on the upsampled MNIST dataset (n = 4096) with I (left) and L3 (right) cost. functions. MDOT reaches the larget "r satislying its stopping criterion [Vg(())) = (()-P) in fewer operations than IPOT takes until termination, so we contimue the KL projection at or until the same number of operations as IPOT is reached, without increasing y further (dashed lines).
G Similarities and Differences with IPOT
In this section, we discuss the similarities and differences between MDOT and the Inexact Proximal point method for exact Optimal Transport (IPOT) algorithm of Xie al. (2020) in more detail. In IPOT, authors propose to update (in our notation) the y parameter by fixed increments and typically choose A = for all / 2 0, while we look A, = 1 for a hyperparameter q > 1. Following each temperature update, they run a fixed number of L column scaling (Sinkhorn) updates and recommend L = 1, whereas MDOT requires that the KL projection is continued until the dual objective gradient norm reaches below a threshold al each 7. A small number of L Sinkhorn updates clearly does not amount to an "exact" KL projection onto the feasible set ()r,c), but Xie et al. (2020) show that under some conditions there exists some finite T. that guarantees linear convergence of (P(u, v; 7), C²2 to the (P*,C), On the other hand, (i) MDOT naturally adapts the number of optimization updates to the difficulty of the problem and the strictness of the stopping criterion at Y, and (ii) it stands to benefit, from polentially faster convex optimization algorithms besides Sinkhorn iteration, e.g., our PNCG approach proposed in Sec. 4 or any other approach that may be developed in the future. A similarity between the approaches is that IPOT also includes an implicit warm-start of the dual variable, which can be considered a special case of our warm-start proposed in Section 4.2.1, where A = for all t.
In Figure 9, we compare IPOT with MDOT (p=1.5, q=21/3 (01////,=2)))) -=) on the upsampled MNIST dataset. For a fair comparison, we followed the conventions of Xie et al. (2020) and implemented all algorithms by computing row/column sums of P(u,,;; explicitly rather than via LogSum Exp reductions in Algorithms 2 and 4.* Each matrix-vector product, row /colllmmm sum, row /wo/ucuum scaling of matrix and entry-wise operations on matrices is counted as one operation of cost O(n²). Since MDOT satisfies its stopping criterion (reaches and satisfies [|Vg()))] S Hrin(r, c)/--P) up to 10x more quickly, we continue minimizing g(u, 0; c) using the respective KL projection algorithm (Sinkhorn or PNCG) until the same number of O(n²) operations as IPOT are executed (namely, 7r).
Overall, the methods behave similarly for the L1 cost, while MDOT is up to 10x faster for the L3 cost.
These results also reveal an interesting contrast between the L1 and L3 cost functions. For the L1 cost the optimality gap rapidly approaches 0 as the KL projection becomes more precise (sce dashed lines), even
IIn this implementation, we evaluate the Sinkhorn direction (17) for PNCG by adding small constant =10-30 -00 each entry of rP) and c(P) for numerical stability.
30

---

Published in Transactions on Machine Learning Research (06/2025)
though "I is fixed. This suggests that the entropic gap, namely (1) -P C)2 is small at a for these MNIST problems, and the optimality gap is dominated by the inexactness of the projection onto U(r, c). We believe that the fast rate O(exp(->X3) of Weed (2018) discussed in Section 4.2.2 is active here for the entropic gap, although this requires further study. On the other hand for the L3 cost, continuing to iterate on the projection error to improve the approximation of (*(r)) does not reduce the optimality gap any further, which implics the entropic gap (P*( C))) still dominant at y these problems.
31

---

Published in Transactions on Machine Learning Research (06/2025)
H Additional Benchmarking on DOTmark
Figs. 10-19 add further benchmarking 10 more datasets from the DOTmmark benchmark of Schrieber et al. (2017), which include various kinds of randomly generated images, classical test images and real data from microscopy. Each dataset contains 45 unique pairs of marginals (r,c) c) obtained from pixel values. The cost matrix is constructed from distances in 2D pixel locations; we evaluate on both L1 and Lg distance costs for n 4096 following our sctup in Sec.
Besides clock time, additionally plot here the total number of Onn")) operations for each algorithm, e-g.; matrix-vector products, row/column sums of matrices, vector outer products, element-wise operations matrices. We count primitive operations for consistency across algorithms; counting higher-level function call such as the number of gradient evaluations would be unfair due to inherent differences in the design of various algorithms. For instance, some require costly line search between gradient evaluations. These plots show that operation counts of the baseline algorithms follow similar trends to wall-clock time and no algorithm unfairly advantaged via. low-level optimizations.
For each of 20 problem sets (10 image datasets 2 cost functions), 20 out of 45 problems are sampled without replacement. The wall-clock time plots for the respective cost functions (L1 and L3) follow similar trends as Fig. 5. In addition to the median, we also include 75% confidence intervals along both axes, which show that MDOT is generally robust.
32

---

Published in Transactions on Machine Learning Research (06/2025)
CauchyDensity
10
10
1
MDOT Sinkhern
L1
MOOT-PACG
cost
10-4
janbulaçati
10
Sinehom (p-1.00
Sinkhern (pm1:5)
10
MSK
Feysy
AAM
Optimality gap
10
L3
10
cost
10
Number of O(n²) operations
Wall-clock time (s)
Figure 10: CauchyDensity problem with L, (top) and L (bottom) costs, showing excess cost (error) vs. number of Onn²) operations (left) and wall-clock time (right).
Classiclmages
10
10
10
L1
MOOT-SeKhern
cost
jambuiscat
Sinchom (pst.00
Sinchern (p=1.5)
MSK
10
APOAGO
Paysy
10-1
AMM
Optimality gap
10
10
L3 cost
10
Number of O(n?) operations
Wall-clock time (s)
Figure ClassicImage problem with Ln (top) and 12 (bottom) costs, showing excess cost. (error) vs. number of O(n²) operations (left) and wall-clock time (right).
33

---

Published in Transactions on Machine Learning Research (06/2025)
GRFsmooth
10
L1
MOOT-Sekheen
cost
10
jonbulacat
Sinehom 10
10
Sinkhern (pm1.))
MSK
10
Feysy
10
AAM
Optimality gap (P
10-2
L3 cost
10
10-
Number of O(n²) operations
Wall-clock time (s)
Figure 12: GRFSmooth problem with L, (lop) and L3 (boltom) costs, showing excess cost (error) vs. number of O(n²) operations (left) and wall-clock time (right).
GRFmoderate
10
10
L1
MOOT-SeKhern
cost
jambuiscat
Sinchor (pst
10-4
Sinchern (p=1.5)
MSK
10
APDAGO
Paysy
AMM
10
Optimality gap
10-4
10
L3
10-4
cost
10
10th
Number of O(n?) operations
Wall-clock time (s)
Figure GRFModerate problem with Tn (top) and I. (bottom) costs, showing excess cost. (error) vs. number of O(n²) operations (left) and wall-clock time (right).
34

---

Published in Transactions on Machine Learning Research (06/2025)
GRFrough
10
10
L1
MOOT-Sekheen
cost
10
jonbulspat
Sinehom 1p-1.00
Sinthern (pm1.))
MSK
Feyey Alg,
AAM
Optimality gap
10
L3 cost
10-4
10-4
1.7
Number of O(n²) operations
Wall-clock time (s)
Figure 14: GRFRough problem with La (top) and L3 (boltom) costs, showing excess cost (error) vs. number of O(n²) operations (left) and wall-clock time (right).
LogGRF
10-4
10
L1
KDOT-SeKhern
cost
10-4
jambulsaat
Sinchor (pst.00
10-5
Sinchern (p-1.5)
MSK
APDAGO
Paysy
10-7
Optimality gap
10-1
10
L3
cost
10"
Number of O(n?) operations
Wall-clock time (s)
Figure 15: LogGRF problem with Ia (top) and T32 (bottom) costs, showing excess cost (error) vs. number of O(n²) operations (left) and wall-clock time (right).
35

---

Published in Transactions on Machine Learning Research (06/2025)
LogitGRF
10
10
L1
MOOT-Sekhern
cost
10
jonbulacat
Sinehom 1p-1.00
10
Sinkhern (pm1.))
MSK
Feysy
AMM
10-1
Optimality gap
10
L3 cost
10
10a5
10
10a
10
Number of O(n²) operations
Wall-clock time (s)
Figure 16: tog problem with La (top) and L3 (boltom) costs, showing excess cost (error) vs. number of O(n²) operations (left) and wall-clock time (right).
Microscopylmages
10a4
10
L1
KDDT-SeKhern
cost
jambulsaat
Sinchor (pst
10-1
Sinchern (p=1.5) 5)
MSK
APDAGO
10
Paysy
AMM
Optimality gap
10-4
10-1
L3 cost
10
10-5
Number of O(n?) operations
Wall-clock time (s)
Figure 17: MicroscopyImage problem with L, (top) and E. (bottom) costs, showing excess cost (error) vs. number of O(n²) operations (left) and wall-clock time (right).
36

---

Published in Transactions on Machine Learning Research (06/2025)
Shapes
10
10
L1
MOOT-Sekhern
cost
jonbulapat
Sinehom 1p-1.00 0)
10
Sinkhern (pm1.))
MSK
10
Feysy
AAM
Optimality gap
10
L3 cost
10
10-
10a
10a2
Number of O(n²) operations
Wall-clock time (s)
Figure 18: Shape problem with L, (lop) and L3 (bottom) costs, showing excess cost (error) vs. number of (n?²) operations (left) and wall-clock time (right).
WhiteNoise
10
L1
10-4
MOOT-SeKhern
cost
10
jambuiscat
Sinchor (pst
10
Sinchorn (p-1.5) 5)
MSK
APOAGO
10
Paysy
Optimality gap
10
L3
cost
10-5
10-1
Number of O(n?) operations
Wall-clock time (s)
Figure 19: WhiteNoise problem with Ln (top) and L& (bottom) costs, showing excess cost (error) vs. number of O(n²) operations (left) and wall-clock time (right).
37


Reviews:
[{"review_id": "B8ttYxqwJz", "reviewer": "Reviewer_Xten", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The empirical performance of the method is strong.\n* I felt the writing was good, and communicated the basic idea of the method quite well.\n* There appear to be well-reasoned motivations behind their tweaks to the method.\n\nWeaknesses:\n* Most evidence is empirical for the superiority of their method and there is not much in the way of proven guarantees on improved performance.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The work proposes a method called Mirror Descent Optimal Transport (MDOT), which empirically produces superior performance to existing OT algorithms in the discrete setting. It works by progressively decreasing the regularization for entropic OT problems in a fashion motivated by mirror descent applied to the primal form of the problem. Several practically impactful tweaks to existing methods are presented:\n* a different tolerance $\\varepsilon_d$ from [Dvuruchensky et al. 18], [Lin et al. 19], and [Altschuler et al. 17]\n* a warm start to the Bregman projection different from that of [Schmitzer 19] and [Feydy 20]\n* use of a pre-conditioned nonlinear conjugate gradients method for Bregman projections instead of the standard Sinkhorn-Knopp iterations\n\nThey demonstrate on a set of problems from upsampled MNIST dataset and a color transfer task and show that their method empirically outperforms existing methods in terms of LogSumExp calls and wall-clock time, especially when higher precision is required.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I do not have too many requests, as I was not able to check the article in great detail. One small note I have is that it would be helpful to place some lines on Fig. 5 (e.g., $n^{2.5}$) to guide the estimations referenced in the text. It's hard to see otherwise (especially with varied scaling on the horizontal and vertical axes).", "broader_impact_concerns": "No concerns here.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "CLgAwmFhUJ", "reviewer": "Reviewer_7ogQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strength \n\nThe authors propose a GPU-parallel nonlinear conjugate gradients algorithm to efficiently solve the subproblem, comparable to existing algorithms such as Sinkhorn iterations. The effectiveness and efficientcy of the proposed method are demonstrated on various datasets.\n\n# Weaknesses \n\nThe paper is not very easy to follow. In particular, the connection between the main algorithm, i.e., Algorithm 1, and the mirror descent framework is not clear based on the current presentation. The explanation lacks a detailed step-by-step derivation that would help readers understand how the algorithm is formulated from the mirror descent principles. Additionally, the paper cites many references when discussing the main algorithm, but the information is presented in an unstructured manner, making it difficult to discern the key contributions and innovations.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces Mirror Descent Optimal Transport, which unifies temperature annealing in entropic-regularized OT with mirror descent techniques, for solving the optimal transport problem approximately.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address the above-mentioned and the following issues. \n\nTo improve clarity, the authors should provide a more comprehensive and structured derivation of the algorithm from the mirror descent framework. This should include a clear explanation of the theoretical underpinnings, assumptions, and any approximations made. Furthermore, the authors should ensure that the context is self-contained, allowing readers to follow the development of the algorithm without needing to consult numerous external sources.\n\nMoreover, while the idea of using mirror descent (or Bregman/Entropic proximal methods) for solving optimal transport is not novel (see e.g., [1]), the paper should better highlight the advanced techniques and innovations in the proposed method compared to existing methods. This could include a detailed comparison of the proposed algorithm's performance, complexity, and convergence properties relative to other state-of-the-art methods. By doing so, the authors can more effectively demonstrate the significance and impact of their contributions.\n\n[1] Xie, Y., Wang, X., Wang, R., and Zha, H. A fast proximal point method for computing exact Wasserstein distance. In Uncertainty in artificial intelligence, pp. 433–453. PMLR, 2020.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Z3nr6Owwoq", "reviewer": "Reviewer_yvDj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n* The work is generally easy to read.\n* The performance of the proposed algorithm is compelling and is compared with many other algorithms. There is also a suitable selection of experimental verifications of the scaling properties of the proposed method.\n\n**Weaknesses**\n* Using \"exponentiated gradient descent\" (mirror descent with entropy mirror map) in place of regular gradient steps has been done in [1] in the context of computing Wasserstein barycenters. Please add this as a reference and possible other applications of mirror descent in the context of optimal transport in the literature review.\n* The introduction to mirror descent in Sec 2.2 is short and confusing. In particular, the $P$ notation and the new introduction of $f$ which is only defined at the end of this section. Please rewrite in a more general manner.\n* (Critical) Prop 4.1 seems sketchy, as the statement is very surprising, in particular that the single MD step can work to exactly solve the regularized OT problem from any rank-1 initialization. I believe the issue is in the recursion step (below Eq. 23, there is sum over $t$, up to $t$...). I think it is fine to simply remove the erroneous claims, including Prop 4.2. It is worth checking Prop 4.3 as well as it appears not to refer to the initialization (again, surprising). Prop 4.4 seems plausible but I did not check it as the statement does not look useful.\n* Limited theoretical properties of the proposed algorithm, e.g. convergence rates. It also does not link into the proposed algorithmic components, which are plentiful and motivated but poorly linked to content earlier in the paper.\n* How is $P^*$ computed in Figure 4? \n* Using the entropic mirror descent has the distinct disadvantage of not being able to reach the boundary of $\\mathbb{R}_{\\ge 0}$ in the mirror descent step. How does the proposed method compare with competing methods when the true transportation plan lies on the boundary, like if the cost matrix has zero entries off the diagonal? This appears to be a possible numerical instability problem with the proposed Sinkhorn direction Eq. 18.\n* Why is the unregularized objective being compared e.g. in Figure 2 when the target problem is the regularized objective? How do the regularized objectives compare? It would be interesting to see how the regularized OT objective compares instead of the unregularized primal, as the competing methods solve the entropic OT problem. Perhaps the slow convergence of the other methods is due to solving an approximate problem, in which case, the proposed MDOT problem has a clear advantage.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes Mirror Descent Optimal Transport for discrete OT, which applies mirror descent steps on the unregularized primal objective instead of the dual regularized objective. Instead of a projection given by Sinkhorn iterations, a Bregman projection is used instead, computed approximately using a conjugate gradients method. Experiments demonstrate significantly improved memory footprint and faster convergence compared to existing OT methods, as well as experimental support for the scaling properties.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Critical**\n* Address weaknesses, especially proof of Prop 4. \n\n**Strengthening**\n* (Fig 3) Change the caption or the figure to point out that the y-axis is in terms of the threshold $\\epsilon_d$.\n* Big-O notation is used inconsistently as $\\tilde{O}$ earlier and $O$ later (e.g. section 5).\n* Add short takeaways to the captions of each figure.\n\n[1] Gramfort, A., Peyré, G., & Cuturi, M. (2015). Fast optimal transport averaging of neuroimaging data. \n[2] Bonet, C., Nadjahi, K., Sejourne, T., Fatras, K., & Courty, N. (2024). Slicing Unbalanced Optimal Transport.", "broader_impact_concerns": "Usage of generative AI models to generate data for color transfer experiment.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper was appreciated by the reviewers that still had some concerns in particular about the proof of Prop 4.1,  the positioning wrt existing OT solvers and the formulation as a Mirror Descent. 

The responses and the new revision have convinced all three reviewer that the paper is OK for an Accept and I concur. But a few things remain to be done in the final version:
- The authors promised to add to Figure 5 lines illustrating the stated "empirical speed" which is indeed necessary because it is not clear form the Figure.
- The authors stated in their reply to yvDj that the focus in Fig 3 is to solve exact OT. This is an interesting reply and the proposed method indeed seem better at that than competitors. But SOTA solvers for exact OT are not entropic/Sinkhorn based, they are network simplex solvers. The paper need to discuss and compare to the exact OT solver (they already use the one from POT) in term of complexity and computational time. For instance $10^4$ logsumexp reduction of complexity $N^2$ for $N=4096$ might not compare favorably to the network flow whose complexity which is $N^3log(N)$ (runs in 2 seconds on my laptop at this size). Admittedly comparison between GPU and CPU solver is a bit unfair toward CPU but would also illustrate the point if the network flow is competitive.  Note that whatever methods is faster is not really important but since the focus is exact OT loss. But this discussion/paragraph needs to be in the paper is order to bring concrete numbers for future users of the solver especially since the "empirical" complexity stated in the abstract are very interesting but potentially misleading for practical problems. For instance in Figure 5 the authors must add a black vertical line with the wall-clock time of the exact solver.
Decision: Accept (poster)