Accelerating Sinkhorn algorithm with sparse Newton iterations
Authors: Xun Tang
Michael Shavlovsky
Holakou Rahmanian
Elisa Tardini
Kiran Koshy Thekumparampil
Tesi Xiao
Lexing Ying
Venue: ICLR 2024 poster (2024)
Keywords: Optimal transport
Convex optimization
Quasi-Newton methods
Non-asymptotic analysis
Extremal combinatorics
Primary area: optimization
TL;DR: A quasi-Newton method with a sparse approximation of a Hessian matrix greatly accelerates entropic optimal transport in both iteration count and in runtime. Numerical analysis is provided which corroborates the algorithm.

Abstract: Computing the optimal transport distance between statistical distributions is a fundamental task in machine learning. One remarkable recent advancement is entropic regularization and the Sinkhorn algorithm, which utilizes only matrix scaling and guarantees an approximated solution with near-linear runtime. Despite the success of the Sinkhorn algorithm, its runtime may still be slow due to the potentially large number of iterations needed for convergence. To achieve possibly super-exponential convergence, we introduce Sinkhorn-Newton-Sparse (SNS), an extension to the Sinkhorn algorithm, by introducing early stopping for the matrix scaling steps and a second stage featuring a Newton-type subroutine. Adopting the variational viewpoint that the Sinkhorn algorithm maximizes a concave Lyapunov potential, we offer the insight that the Hessian matrix of the potential function is approximately sparse. Sparsification of the Hessian results in a fast $O(n^2)$ per-iteration complexity, the same as the Sinkhorn algorithm.  In terms of total iteration count, we observe that the SNS algorithm converges orders of magnitude faster across a wide range of practical cases, including optimal transportation between empirical distributions and calculating the Wasserstein $W_1, W_2$ distance of discretized continuous densities. The empirical performance is corroborated by a rigorous bound on the approximate sparsity of the Hessian matrix.

Published as a conference paper at ICLR 2024
ACCELERATING SINKHORN GLRIITH WITH SPARSE NEWTON ITERATIONS
Xun Tang', Michael Shavlovsky', Holakou Rahmanian', Elisa Tardini!, Kiran Koshy Thekumparampil; Tesi Xiao', Lexing Ying!
Staaford University
AAmazon
(xutang, lexing lex.nng@staaffrd. edu
(shavlov, holakou, ettardin, kkt, tesixiao) ama com com
ABSTRACT
Computing the optimal transport distance between statistical distributions a fun- damental task in machine learning. One remarkable recent advancement is en- tropic regularization and the Sinkhorn algorithm, which utilizes only matrix scal- approximated solution with near-linear runtime. Despite
ing and guarantees the success of the Sinkhorn algorithm, its runtime may still be slow due to the po- tentially large number of iterations needed for convergence. To achieve possibly super-exponential convergence, we present Sinkhorn- -eetin--eett--ttteeeeeeenee -parse (SNS), an extension the Sinkhorn algorithm, by introducing early stopping for the matrix scaling steps and a second stage featuring a Newton-type subroutine. Adopting the variational viewpoint that the Sinkhorn algorithm maximizes a concave Lya- punov potential, we offer the insight that the Hessian matrix of the potential func- tion is approximately sparse. Sparsification of the Hessian results a fast (n?) per-iteration complexity, the same the Sinkhorn algorithm. In terms of total iteration count, we observe that the SNS algorithm converges orders of magni- tude faster across wide range of practical cases, including optimal transportation between empirical distributions and calculating the Wasserstein W1, W2 distance of discretized densities. The empirical performance corroborated by rigorous bound on the approximate sparsity of the Hessian matrix.
INTRODUCTION
Optimal transport (OT) calculates the best transportation plan from an ensemble of sources to targets
Villani et al. 2009 Linial et al. 1998 PPeyré et al.) 2017) (117)am) and is becoming increasingly an important task machine learning (Sandler & Lindenbaum Jitkritum al.| [2016] Arjovsky Arjossss al. [017] Salimans et al.) 2018 Genevay et al.) 2018 Chen et al. 2020 Fatras et al. 2021). In this work, we
focus on optimal transportation problem with entropic regularization:
C. C-P+ MP+-MCg -P+-)(()) 11
min P.P.
(1)
where 7) 0 is the entropy regularization parameter, C' E Rxn is the cost matrix, r E R are respectively the source and target density, and H(P) E Pij log is the entropy of P. The insight of using the Sinkhorn algorithm that mntropy-regularte
optimal transport is equivalent to
an instance of matrix scaling (Linial et al, 1998; Cuturi 2013 Garg et al. 2020):
Find diagonal matrix X, Y so that P = X exp( Y satisfies P1 -r. P
The Sinkhorn algorithm Yule 1912) alternates between scaling the rows and columns matrix to target vector, and its convergence property first proved in [Sinkhorn (1964). Theoretical results show that the Sinkhorn algorithm converges at a relatively slow rate. While the Sinkhorn al-
gorithm satisfies exponential convergence (Franklin & Lorenz 1989 Carlier 2022). its best proven The symbol exp(-r/)) denotes entry-wise exponential.

---

Published as a conference paper at ICLR 2024
exponential convergence rate is often too close to one for practical use (see Section for detailed discussions) and in practice it behaves more like a polynomially converging method (Altschuler
et al. 2017: Ghosal & Nutz 2022). Therefore, to reach within sub-optimality gap e. the tightest bound of iteration complexity in practice tends to be ((poly(e-)) in early stages.
We introduce new algorithm that greatly accelerates convergence by reducing the required iteration counts. Utilizing a well-established variational perspective for the Sinkhorn algorithm (Aftschuler al. 2017), we consider the Lyapunov potential f: IR" R" IR associated with the entropic optimal transport problem in equation
f(r,y):= := :o.,))=--- 17 Σ
exp(1(-c0 + 34 93) 1) Σ PaR + Σ
In particular, the discussion in Section 3 shows that solving equation I is equivalent to obtaining (x",y") = arg maxx.y [(x,, v). We emphasize that f is concuve, allowing routine convex optimiza- tion techniques to used. Under this framework, the matrix scaling step the Sinkhorn algorithm can be seen an alternating maximization algorithm
arg max, fzz v), y f- arg max, (2,
The formula in equation 3 provides surprisingly clear guidance and justification for our approach to accelerate the Sinkhorn algorithm. First, one can only achieve substantially reduced iteration count by jointly optimizing (2r, y). Second, no first-order method can be used, as they achieve polynomial convergence at best, reaching an iteration complexity of O(6=1) in the case of gradient descent for a sub-optimality of (Boyd & Vandenberghe 2004). O(c-1/2) in the case of accelerated In conclusion, one can only hope to achieve better convergence
gradient descent [Nesterov] with second-order methods, which enjoy super-exponential convergence (Boyd & Vandenberghe 2004). The use of Newton's 's method for the Sinkhorn algorithm has been introduced in [Baaer et al. (217)) However, even single Newton step has an O(n²) cost, which violates the goal of having near-lineal time algorithm with O(n²) total complexity. This naturally leads to the question:
Is there an algorithm that reaches the optimality gap with an iteration complexity of Newton's 's algorithm and an O(n²) per-iteration complexity of the Sinkhorn algorithm?
Practical Newton's algorithm via Hessian sparsification We answer the question in the affirma- tive by introducing the Munkhorn-Newton-Spee (SNS) algorithm, which achieves fast convergence and O(n?) per-iteration complexity. The main technical novelty is solve the Hessian system (V²f)u = f in Newton's algorithm efficiently with sparse approximation. In particular, we show that the Hessian matrix of the Lyapunov potential is approximately sparse in the following sense:
Definition (Sparsity and approximate sparsity) Let I|-llo denote the lo norm The sparsity of matrix M E mxx is defined by r(M) = [|M]]o Furthermore, a matrix M E Rmx is (X, sparse if there exists matrix M so that T(M) < and || -MIl Se.
In Section 5] we prove rigorous bound on the approximate sparsity of the Hessian matrix. We highlight our result by providing an informal version of our sparsity analysis:
Theorem. (Informal version Theorem 7 Assume minp PT 1=c C-P admits a unique solution. Then, if t. 17 are sufficiently large, the Hessian matrix after Sinkhorn matrix 3 12n" exp (-pn) + ))-spasse for some parameter p,q.
scaling step is 2n
As v2 approximately sparse, one can approximate v with a relaxed target sparsity = O(1/n) 3. The cost for solving the sparsified linear system thus reduces to O()x²) O(n²).
Contributions The contribution of this paper is threefold. First, we point out the sparsification technique, and the resultant SNS algorithm reduces the per-iteration cost of Newton step to On??).
2

---

Published as a conference paper at ICLR 2024
While related idea has been discussed with the name of the mankhom-Newtorr algorithm in Brauer et al. (2017). the (nn)) complexity of Sinkhorn- Neatooon-eetonn makes the algorithm prohibitively expen- sive to work with. Second, we give non non-asymptotic analysis, showing that can expect sparse Hessian generically. Third, we fully adapt our sparsity argument to the case of non-unique solutions. We introduce novel argument that shows O (n)) sparsity. Moreover, the provided numerical analysis directly treats the case of non-unique solutions.
Notation For N. we denote [n] :- {.,...: n). We use shorthand for several matrix operations for the sake of notational compactness. The operation between matrices is defined by C Σij=1 CiyPij. For a matrix M, the notation log M stands for entry-wise logarithm, and similarly exp(MM denotes entry-wise exponential. The symbol 1 stands for the all-one vector in R Finally, use the symbol ||1 to denote the entry-wise l norm, i.c. [|M == ||vec[M)]] Eig /mig).
RELATED LITERATURE
Convergence of Sinkhorn We give an overview of the convergence of the Sinkhorn algorithm and discuss the iteration complexity to reach within an error tolerance a on the marginal KL divergence, which is defined by C(P) = [L(r[|P1) + [L(c|IPT]) The Sinkhorn algorithm has a sub-
optimality gap of O ((1-e (1-e- (11-e-21100)>>)) after t steps Carlier 2022), which implies an iteration
complexity of O( log(1/a)) The analysis by Altschuler et al. (2017) proves an iteration complexity of O(a and has been refined to ((a 1172) in special cases (Ghosal & Nutz, 2022)
We refer the readers to |Peyré et al. (2017): Carlier] (2022) for more comprehensive review of the convergence of the Sinkhorn algorithm. The analysis in (Mallasto et al., 2020 Weed (018)
considers the convergence of the entropic optimal transport solution to the true transport plan, which further connects the Sinkhorn solution to the ground-truth optimal transport plan.
Sparsification in Sinkhorn Sparsification techniques have been extensively explored to improve the efficiency of the Sinkhorn algorithm. Sparsification has been considered during matrix scal- ing (Li et al., 2023), and a 2-norm or group lasso penalty function has been considered to boost sparsity in the transport plan (Blondel et al. 2018). Unlike the methods discussed, which modify the optimization task, our work retains the original formulation while introducing enhancement via sparsification inside the optimization step. Our approach uses Newton's algorithm, and the sparsity is applied to the Hessian matrix for the Lyapunov function. The numerical analysis in our work lends further support sparse transportation plan in the entropic case.
Acceleration for Sinkhorn There is considerable research interest in speeding up the runtime of the Sinkhorn algorithm. There are a few strategies related to our work, including randomized or greedy row/column scaling (Genevay et al.) 2016; Altschuler et al. 2017), dynamic schedule of entropic regularization (Chen et al.) 2023). These works complement our work, as both strategies be naturally used in our Newton-type algorithm. Although both techniques can potentially boost the convergence speed of SNS, we focus on the case with fixed entropy strength and joint optimization over all variables to show that sparse approximation to the Hessian can reach rapid convergence on its own. Another notable line of acceleration technique considers acceleration using an approximation of the kernel K exp(-C'r) (Deriche 1993 Solomon et al. 2015 Bonneel
et al., 2016; Altschuler et al., 2019 Scetbon & Cuturi 2020 Huguet et al. 2023 or by exploring
low-rankness of the transport plan Scetbonn et al. 2021
Variational methods in Entropic OT Furthermore, a stream of research focuses on jointly op- timizing the potential / using accelerated first-order method (Dvurechensky et al., 2018: Thibault et al. 2 Kemertas al. 2023). The SNS algorithm can be extended by replacing Sinkhorn ma- trix scaling steps with iteration schemes based on first-order methods. As variational interpretation
of the Sinkhorn algorithm, it can also be viewed as mirror descent (Mishchenko) 2019 Léger 202
OT in Machine Learning A rich body of literature exists the applications of optimal transport in various Mmchine-learning domains. Research such as Kolouri et al.] (2017); /Vaeer et al. (2018); (00tt)OOneto
Genevay et al. (201);] Luise et al. (2018); Oneto al. 2020): Huynh al. (2020) studies statistical
3

---

Published as a conference paper at ICLR 2024
learning with optimal transport. Studies like Genevay et al. (2017); Bousquet al. (2017) : Sanjabi eet al. (0118) Deshpande et al. (2019); Lei et al. (2019): Patrini et al. (2020); Onken et al. (2021)
utilize optimal transport distance metric for improving the performance and robustness of un-
supervised learning algorithms. The Sinkhorn a algorithm, featured in works such as Fernando et al.
(2013): Redko et al. (2017); Courtyy al. (2017); Alvarez-Melis al. (2018); Nguyen et al. (2022)
2021]: Xu et al. (2022): Turrisi et al. (2022), is commonly used in robust domain adaptation, partic-
ularly in addressing distribution shift. Additionally, various approaches employ neural networks to
learn the optimal transport map, as seen in (Seguy et al. 2017 Damodaran et al. 2018 Makkuva
al. 2020 Mokrov et al., 2021 Daniels al. 2021).
VARIATIONAL. FORM OF THE SINKHORN ALGORITHM
This section summarizes the variational form of the Sinkhorn algorithm, mathematical representa- tion crucial for understanding our proposed algorithm's theoretical underpinnings. As pointed out, the Sinkhorn algorithm performs alternating maximization for the Lyapunov potential. By introduc- ing the Lagrangian variable and using the minimax theorem (for a detailed derivation, see Appendix we formulate the associated primal-dual dual problem to equation Das:
max min L(P, x. v) P log P + CC-P- x 1..--xx-PP--001 (P)-r) - "
a.y
The Lyapunov function f. described in equation comes from eliminating P' (see Appendix A):
f(x,)) 3) min L(P, v.
Maximizing over is equivalent to solving the problem defined in equation | As consequence of the minimax theorem, (x*, 3²* = arg max,, f(x,z) y) effectively solves equation as the following equation shows:
C 1 H(() P exp (n(-C+=11++ P.P1=r,P1ic
arg min
11
Let P be defined as exp (n(-C + I1 ly 1y -)), where serves the intermediate matrix corre- sponding to dual variables r, y. As f is concave, the first-order condition is equivalent optimality. Upon direct calculation, one has
fz, = Σ Pake (x. = Σ PRS- &
As consequence, maximizing x with fixed y corresponds to scaling the of P so that P1 Likewise, maximizing y with fixed corresponds to scaling the column of matrix P so that P c. Thus, the Sinkhorn matrix scaling algorithm corresponds ascent approach to the Lyapunov function, as illustrated in equation 3 For the reader's 's convenience, we write down the first and second derivatives of the Lyapunov function f:
alternating coordinate
= 1 v2 = [diag(P)) (1)) PT
P diag(P 1)
ves(s,y)= - Pl, -c-(T,
(4)
4 MAIN ALGORITHM
We introduce Algorithm herein referred to as Sinkhom--Newoo-Spaass (SNS). This algorithm extends the Sinkhorn algorithm with second stage featuring Newton-type subroutine. In short, Algorithm [ starts with running the Sinkhorn algorithm for N₁ steps, then switches to sparsified Newton algorithm for fast convergence. Algorithm I employs a straightforward thresholding rule for the sparsification step. Specifically, any entry in the Hessian matrix smaller than constant is truncated to zero, and the resulting sparsified matrix stored in sparse data structure. The trun- cation procedure preserves symmetry and diagonal dominance of the Hessian matrix (as a simple consequence of equation 44). which justifies the use of conjugate gradient for linear system solving (Golub & Van Loan 2013). The obtained search direction Az is approximation to the exact New- ton step search direction, i.e., the solution to the linear system (V² = -Vf. In other words, removing sparse approximation will recover the Newton algorithm, and the Newton algorithm with- out sparsification considered Appendix [) In Appendix [ we introduce additional techniques used for improving the numerical stability of SNS.
4

---

Published as a conference paper at ICLR 2024
Algorithm 1 Sinkhomn-Newoon-Spans (SNS)
Require: f, E R", M/nii E RR" ,,, N2, 0
1: # Sinkhorn stage
2: (z.z) f (Xiitt Minit)
Initialize dual variable
3: while < N1 do
4: P - exp (m(-C + x1 T + 1y 1)
5:
+- + (log(r) - log (1)) /n
6:
P - exp (n(-C x1 + 1y y -1)
7:
y (log(c) log( (T (1))
8:
++itl
9: end while
10: # Newton stage
< (x,y)
12: while < Nı + N₂ do
13:
M (
Truncate with threshold p
14:
A: f Conjugate_Gradientt Gradient - -V((2))
>> Solve sparse linear system Line search for step size
a Line search 2, Az)
16:
2 2-2+00 +
i+i+1
18:
end while
19: Output dual variables (x, v) +
Complexity analysis of Algorithm Inside the Newton stage, the conjugate gradient method in- volves O(n) left multiplications by M := Sparsify (z) Furthermore, left multiplication by M involves (MT))) arithmetic operations, where T(.) the sparsity defined in Definition Thus, obtaining Az is of complexity ((T).)n²). To maintain an upper bound for per-iteration complexity, practice, one sets target sparsity À and picks dynamically to be the [Am21 entry of f² which ensures r(M) < X. For the forthcoming numerical experiments, whenever transport problem has unique solutions, we have found setting X = 2/n suffices for a conver- gence performance quite indistinguishable from a full Newton step, which is corroborated by the approximate sparsity results mentioned in Theorem
Necessity of Sinkhorn steps We remark that the Sinkhorn stage in SNS has two purposes. The first purpose is warm initialization. Using the Sinkhorn steps, we bring the intermediate dual variable (x,, v) closer to the optimizer that the subsequent Newton step has good convergence speed. Secondly, this proximity ensures that the intermediate matrix P satisfies approximate sparsity. While the number Sinkhorn steps is written as fixed parameter Algorithm I one can alternatively switch to the Newton stage dynamically. In particular, we switch to sparsified Newton step when the two following conditions hold: First, the intermediate matrix should admit a good sparse approximation, and secondly, the Newton step should improve the Lyapunov potential more than the Sinkhorn algorithm. Our analysis in Section 5 demonstrates that it requires at most (1/22) Sinkhorn steps for the sparsification to be within an error of e.
5 SPARSITY ANALYSIS OF SNS
This section gives a complete theoretical analysis of the approximate sparsity throughout the SNS algorithm: Theorem [ Tanalyzes the approximate sparsity of the Hessian after the Ni Sinkhorn step; Theorempaanal analyzes the approximate sparsity within the N₂ Newton steps and proves monotonic im- provement the approximate sparsity guarantee. Three symbols are heavily referenced throughout this analysis, which we list out for the reader's convenience:
Pt, The result of Sinkhorn algorithm after iterations.
F: The set of optimal transport plan the original problem
P The eetropy-regulariied optimal solution,
5

---

Published as a conference paper at ICLR 2024
where F, P. satisfy the following equation:
C HPP) P.Pl
=
C.P, := P.P
arg min
arg min
(5)
Approximate sparsity in the Sinkhorn stage We first prove the main theorem on the approximate sparsity of P..p which in particular gives an approximate sparsity bound for the Hessian matrix when one initiates the sparsified Newton step in Algorithml I We generalize allow for the setting where potentially multiple solutions exist to the optimal transport problem. Definition 2 lists the main concepts the analysis. For concepts such as polyhedron, face, and vertex solution, the readers can consult Cook et al. (1998) for detailed definitions.
Definition 2. Define P as the feasible set polyhedron, i.e. P := (PI =
- -rr,, c, 0}.
The symbol V denotes the set of vertices of P. The symbol O stands for the of optimal vertex solution, iee.
arg min C. PEV
The symbol A denotes the vertex optimality gap
A min C min P.
QEV-O
We use F (O) denote the optimal face the optimal transport problem, and (FF) = maxMer (M) is defined to be the sparsity of F. We define a distance function to F by dist(F, P) arg
Pll1-
We move on to prove the theorem for approximate sparsity of Pt. Theorem 1. Assume |r]] = ||c||] 1, and let A be as Definition There exists constant q, l
such that, for 1j > A
142logn f > t one has
dist(F, Pan) S 6n exp (-nA) V
Vt
Therefore, for X = T(F), it follows that Pt.g is (X* -spars, whereby
V Vt
6t, : 6n² exp (-nA)
As result, the Hessian matrix in Algorithm I after N1 Sinkhorn steps is 22 2n 2 spparse.
Define the subset S C gaxnn that C' E CC is a cost matrix for which minp.pier,PTI C P has non-unique solutions. Then, S of Lebesgue measure zero, and so the optimal transport minp,P1 pammr, PT C P has unique solution generically. The generic condition C S leads to X= < 2/n. 1f one further assumes = 11, then X* =1/n.
Proof. It suffices to construct sparse approximation to Pt., that matches the requirement in Definition [ We choose the sparse approximation to the matrix Pa E F which satisfies p arg min MlI. By triangle inequality, one has
- <e:= [Pt,, - Pf | dist(F,
Thus Pr,, (T(F), c)ssarrse..
sparse. Existence of sparse approximation to the Hessian matrix V is shown by the following construction:
diag Pt Pt,y (P*)
P
M
diag( (PT.
One can directly count that the number nonzero entries of M is upper bounded by 2n + 2 Moreover, direct computation shows || - v fl| Pr, - P=1|1 I1 < 2e. Thus, we show that the Hessian matrix is (2' 2n 2e -spaaree
Thus, the proof reduces to proving ct.- By Corollary 9 in Weed 2018 if follows
1-2log
dist(FF, 5 2n exp (-7A 1) Gn exp (-nA)-

---

Published as a conference paper at ICLR 2024
By Pinsker inequality (Pinsker 1964) and Theorem 4.4 in Ghosal & Nutz 2022 there exists stants q. 1 such that, for any t >> 11, one has
Pen P; (L(Pp,))P;)) KL(P,
and thus ||Pt. P7 I Vi
Combining equation 7 and equation 8 one has Et,ys as desired.
One has C' E E R if and only if there exist two distinct permutation matrices Pt. P2 so that PI = C. P2. For each Pi, P2 the condition that C' = C P2 is on a subset of measure zero on R As there are only a finite number of choices for PI, P2, it follows that S is a finite union of sets of measure zero, and thus in particular S is of measure zero. Thus, one has C' & S generically.
When mill:: 1-c C P has a unique solution P it must be an extremal point the poly- hedron P. which has 2n - 1 non-zero entries (Peyré et al. 2017). and therefore r(F) r(P") 2n-1 n² I 2. In the case where r = c = 1, follows from Birkhoff-von Neumann theorem that n'
nP is permutation matrix, and therefore TFF) = TPP)) = n'
As Theorem [ shows, taking a target sparsity good sparse approximation, which leads to O(n²) per-iteration cost. It is worth pointing out that exp (-Ar) term ér,y shows the appealing property that the matrix Pt,, has a better sparse approximation guarantee in the challenging case where 7) is very large.
O(1/n) 3n/2 Algorithm leads
Approximate sparsity in the Newton stage We consider next the approximate sparsity inside the N₂ Newton loops. We show that approximate sparsity is monotonically improving as the Newton step converges to the optimal solution:
(Ik, wa) denote the dual variable at the k-th Newton step, and let eq max; f(=) 1003-(le f=xx be the sub-optimality gap for zk. For the sake of normalizing the trunsport plan formed by zk. take Uk.+ arg y (I v) and define Pe to be the transport plan formed after column normalization, ie..
Theorem 2. Let zhe
Ps = exp (n( (n(-C++ 1 (Jke,) -1) --1)
Then, for the same constant in Theorem sparse, where
for Cr < and 1
1+2loRn A
the matrix Pa is (X*, GR)-
SK = 6n exp (-n +
The proof is similar to that of Theorem and we defer it to Appendix B We remark that the EK term Theorem 2 becomes insignificant e decreases at a super-exponential rate. Moreover, the statement in Theorem 2 would have shown monotonically improving approximate sparsity if one were to add a column scaling step inside the Newton inner loop in Algorithm []] For clarity, we do not include any matrix scaling in the Newton stage of Algorithm [ However, combining different techniques in the second step might aid convergence and is worth investigating.
1/4
Sparsity under non-uniqueness We explore sparsity guarantees when the optimal transport prob- lem lacks unique solution. Although these cases are rare in practice, we point out the somewhat surprising result: There exist conditions for which one can derive sparsity properties of the optimal face using tools from extremal combinatorics 2011]. On high level, the optimality of p* F forbids certain substructures from forming in an associated bipartite graph, which in turn gives upper bound the number of nonzero entries of p*. We defer the proof to Appendix B
bipartite graphs (Erdōs & Spencer] 1974| Juknal
Theorem 3. For a cost matrix C' = lcylii== E Rux", suppose that oi + cet = Ce3 + Citi if and
only if = i' or = j.. Then one has ((F) < 1+o(1)
7

---

Published as a conference paper at ICLR 2024
Sinkhom
Sinkhorn -oweorn--ee------- Sparse
Switch to Newton
SNS terminates
Convergence in Iterations
Convergence in wall time
-10
-10
Optimality
-35
20
-25
-10
-30
150
200
250
Iterations
Seconds
(a) Entropic random linear assignment
Convergence iterations
Convergence
-5
-5
-10
-10
Log Optimality
-15
-15
-20
-20
-25
50
100
150
200
250
0.9
2.5
5.0
75
10.0
12.
15.0
17.5
Iterations
Secands
(b) Optimal transport on the MNIST dataset under transportation cost ||>
vll
Convergence in iterations
Convergence in wall time
Gap
-5
-10
-30
-20
-20
-25
-25
-30
600
Iterations
Seconds
(c) Optimal transport on the MNIST dataset under transportation cost || - w|li
Figure I: Performance comparison between Algorithm and the Sinkhorn algorithm.
NUMERICAL RESULT
We conduct numerical experiments to compare the original Sinkhorn and SNS algorithms. We use the following settings: We obtained the runtime data on a 2021 Macbook Pro laptop with an Apple MI Pro chip. The lincar system solving is done through the conjugate gradient step as mentioned in Algorithm To set a challenging case, use entropic regularization with r7 1200 throughout the experiments. We refer the reader to Appendix F for the performance of SNS under different entropy parameter 1)- We defer a comparison of SNS with the Sinkhom-Newtt -ewton algorithm (Brauer ett al. (017)) to Appendix D where we show that removing sparsification from SNS results in prohibitively expensive algorithm due to its (1²²) runtime complexity. Additionally, we perform experiments on the numerical performance of quasi-Newton methods (Noceaal & Wright 1999).
In the first numerical test, consider the random assignment problem with entropic regularization (Mézard & Parisi 1987 Stecle 1997; Aldous 2001), considered a hard instance of optimal trans- port. The cost matrix C Gylle gazz with 500 is generated by cjj ~ Unif((0, 1]). The
source and target vectors are c= r 11. We run Algorithm with N, = 20 and a target sparsity of X = 2/n. Figure Talshows that SNS drastically outperforms Sinkhorn in iteration and runtime.
8

---

Published as a conference paper at ICLR 2024
Table 1: Performance comparison between SNS and Sinkhorn. Both algorithms are run until they reach machine accuracy.
Case
Method
Stage
Time (s)
Iterations
Sinkhorn
0.12
20
SNS
Newton
0.22
Random
Total
0.34
29
Sinkhorn
Total
233.36
56 199
Sinkhorn
0.17
20
SNS
Newton
2.16
33
MNIST L2
Total
2.33
53
Sinkhorn
Total
18.84
2041
Sinkhorn
6.28
700
SNS
Newton
5.94
77
MNIST LI
Total
12.22
777
Sinkhorn
Total
45.75
5748
In the second numerical experiment, similar to the experiment setting in in(Cuturi] (013), we consider the more practical case of optimal transport on the MNIST dataset. In particular, two images are respectively converted ized to sum to 1. The entry corresponding to the (11, i2)-t pixel is conceptualized as the point (11/28, 22 28) E R2, and the transport cost is the Euclidean distance cost || u11?. Similarly,
vector of intensities
the 28 x 28 pixel grid, which are then normal-
we pick N1 = 20 in Algorithm with a target sparsity of À = 2/nz Figure
lb shows a similar
performance boost the Sinkhorn algorithm.
As the approximate sparsity analysis underlying Section 5 mainly focuses on the situation of unique solutions, is of interest test problems with many optimal transport solutions, it is case where the SNS algorithm might potentially break in practice. For this purpose, we consider the MNIST example under the l transport cost || -yll. which known have non-unique solutions due to the lack of convexity in the |-I| norm (Villani et al. 2009). As the Sinkhorn algorithm converges quite slowly, we pick N₁ 700 before switching the Newton stage. Choosing target sparsity of 15/n is sufficient for convergence to the ground truth, which shows that SNS runs quite well even without uniqueness guarantee.
In Table T we benchmark the runtime for Sinkhorn versus SNS to reach machine accuracy, which shows that SNS has an overall runtime advantage for convergence. In particular, we list the per- formance of SNS in the Newton stage, which shows that early stopping of Sinkhorn matrix scaling steps and switching to the Newton stage results an order of magnitude speedup in iteration counts. While we cannot prove the conjectured super-exponential convergence, the low iteration count in the Newton stage shows strong numerical support.
7 CONCLUSION
We propose the Sinkhom--Newton-Sparss algorithm, demonstrating its empirical super-exponential convergence at O(n²) per-iteration cost through numerical validation. We prove several novel bounds on approximate sparsity underlying the algorithm. For problems with non-unique solutions, we elucidate novel relationship between approximate sparsity and extremal combinatories. We contend that this new method significantly advances the field of high-precision computational opti- mal transport and complements the existing Sinkhorn algorithm.
For future work, it may be interesting to study how the approximation accuracy of the sparsification step affects the algorithm's 's convergence and how to devise more sophisticated sparse approximation techniques beyond simple thresholding. Moreover, it is exciting direction to incorporate existing optimal transport techniques with the SNS algorithm, including greedy row/column optimization, dynamic regularization scheduling, and hardware-bai based parallel accelerations. Theoretically, alytic properties on the Lyapunov potential might provide more insight into the region where the sparsified Newton algorithm achieves the super-exponential convergence empirically observe.

---

Published as a conference paper at ICLR 2024
REFERENCES
David J Aldous. The ( (2) limit in the random assignment problem. Random Structures & Algo- rithms, 1844):311-418, 2001.
Jason Altschuler, Jonathan Kiles-Weed and Philippe Rigollet. Near-linear time approximation al- gorithms for optimal transport via sinkhorn iteration. Advances in neural information processing systems, 30, 2017.
Jason Altschuler, Francis Bach, Alessandro Rudi, and Jonathan Niles- Weedd Massively scalable sinkhorn distances via the nyström method. Advances in neural information processing systems, 32, 2019.
David Alvarez-Melis, Tommi Jaakkola, and Stefanie Jegelka, Structured optimal transport. In International conference artificial intelligence and statistics, pp. 1771-1780. PMLR, 2018,
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative ad adversarial networks. In International conference on machine learning, pp. 214-223. PMLR, 2017.
Mathicu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In Inter- national conference artificial intelligence and statistics, pp. 880-889. PMLR, 2018.
Nicolas Bonneel, Gabriel Peyré, and Marco Cuturi. Wasserstein barycentric coordinates: histogram regression using optimal transport. ACM Trans. Graph., 35(4):71-1, 2016.
Olivier Bousquet, Sylvain Gelly, Ilya Tolstikhin, Carl-Johann Simon-Gabrics and Bernhard Schoelkopf. From optimal transport to generative modeling: the vegan cookbook. arXiv preprint arXiv:1705.07642, 2017.
Stephen P Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Christoph Brauer, Christian Clason, Dirk Lorenz, and Benedikt Wirth. A sinkhom-- newoon method for entropic optimal transport. arXiv preprint arXiv: 17X0:v/710006635, 2017.
Emmanuel Candès, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? Journal of the ACM (JACM), 58(3):1337, 2011.
Guillaume Carlier. On the linear convergence of the multimarginal sinkhorn algorithm. SIAM
Journal on Optimization, 32(2):786-794, 2022.
Jingbang Chen, Li Chen, Yang P Liu, Richard Peng. and Arvind Ramaswami. Exponential conver- gence of sinkhorn under regularization scheduling. In SIAM Conference on Applied and Compu- tational Discrete Algorithms (ACDA23), pp. 180-188. SIAM, 2023.
Liqun Chen, Zhe Gan, Yu Cheng, Linjie Li, Lawrence Carin, and Jingjing Liu. Graph optimal transport for cross-domain aligament. In International Conference Machine Learning, pp- 1542-1553. PMLR, 2020.
William J Cook, William H Cunningham, William R Pulleyblank, and Alexander Schrijver. Combi- natorial optimisation. Wiley-Interscience Series in Discrete Mathematics and Optimization, USA, 1:998, 1998.
Nicolas Courty, Rémi Flamary, Amaury Habrard, and Alain Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. Advances neural information processing systems, 30, 2017.
Marco Cuturi. Sinkhor distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Bharath Bhushan Damodaran, Benjamin Kellenberger, Rémi Flamary, Devis Tuia, and Nicolas Courty. Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation. In Proceedings of the European conference on computer vision (ECCV), pp. 447-463. 2018.
Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems, 34:12955-11298 12965, 2021.
10

---

Published as a conference paper at ICLR 2024
Rachid Deriche. Recursively implementating the Gaussian and its derivatives. PhD thesis, INRIA, 1993.
Ishan Deshpande, Yuan-Time Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander Schwing. Max-sliced wasserstein distance and its use for gans. In Proceedings of the IEEE/CVF Conference Computer Vision and Putlern Recognition, pp. 10648-10656, 2019.
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn's 's algorithm. Interna- lional conference machine learning, pp. 1367-1376. PMLR, 2018.
Paul Erdös and Joel Spencer. Probabilistic methods in combinatorics, (No Title), 1974.
Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty, Unbalanced minibatch optimal transport; applications domain adaptation. In International Conference on Machine Learning. pp. 3186-3197. PMLR, 2021.
Basura Fernando, Amaury Habrard, Marc Sebban, and Tinne Tuytelaars. Unsupervised visual do- main adaptation using subspace alignment. In Proceedings of the IEEE International Conference Computer Viston (ICCV), December 2013.
Joel Franklin and Jens Lorenz. On the scaling of multidimensional matrices. Linear Algebra and its applications, 11::717-735, 1989.
Ankit Garg, Leonid Gurvits, Rafael Oliveira, and Avi Wigderson. Operator scaling: theory and applications. Foundations of Computationnll Mathematics, 20(2):223-291 2020.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. Advances in neural information processing systems, 29, 2016.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Gan and vae from optimal transport point of view. arXiv preprint arXiv: 0177000000.7 2017.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn di- vergences, In International Conference on Artificial Intelligence and Statistics, pp. 1608-1617.
PMLR, 2018.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complex- ity of sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pp. 1574-1583. PMLR, 2019.
Promit Ghosal and Marcel Nutz. On the convergence rate of sinkhorn's algorithm. arXiv preprint arXiv:2212.06000. 2022.
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 2013.
Guillaume Huguet, Alexander Tong, María Ramos Zapatero, Christopher J Tape, Guy Wolf, and Smita Krishnaswamy. Geodesic sinkhorn for fast and accurate optimal transport on manifolds. ArXiv, 2023.
Viet Huynh, He Zhao, and Dinh Phung. Otlda: geometry-aware optimal transport approach for topic modeling. Advances in Neural Information Processing Systems, 33:18573-18582, 2020.
Wittawat Jitkrittum. Zoltán Szabó, Kaeper P Chwialkowski, and Arthur Gretton. Interpretable dis- tribution features with maximum testing power. Advances in Neural Information Processing Systems, 29, 2016.
Stasys Jukna. Extremal combinatorics: with applications in computer science, volume 571. Springer, 2011.
Mete Kemertas, Allan D Jepson, and Amir-massoud Farahmand. Efficient and accurate optimal transport with mirror descent and conjugate gradients. arXiv preprint NXXi:2307.00807 2023.
11

---

Published as a conference paper at ICLR 2024
Soheil Kolouri, Se Rim Park, Matthew Thorpe, Dejan Slepcev, and Gustavo K. Rohde. Optimal mass transport: Signal processing and machine- applications. IEEE signal processing magazine, 34(4):43-55) 202--59,11111 2017.
Flavien Léger. gradient descent perspective on sinkhorn. Applied Mathematics Optimization, 84(2):843-1855, 2021.
Na Lei, Kehua Su, Li Cui, Shing-Tung Yau, and Xianfeng David Gu. A geometric view of optimal transportation and generative model. Computer Aided Geometric Design, 68:1-21, 2019.
Mengyu Li, Jun Yu, Tao Li, and Cheng Meng. Importance sparsification for sinkhorn algorithm. arXiv preprint orXiv:2306.06500 2023.
Nathan Linial, Alex Samorodnissky, and Avi Wigderson. A deterministic strongly polynomial al- gorithm for matrix scaling and approximate permanents. Proceedings of the thirtieth annual ACM symposium on Theory of computing, pp. 644-652, 1998.
Giulia Luise, Alessandro Rudi, Massimiliano Pontil, and Carlo Ciliberto. Differential properties of sinkhorn approximation for learning with wasserstein distance. Advances in Neural Information Processing Systems, 31, 2018.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672- 6681. PMLR, 2020.
Anton Mallasto, Augusto Gerolin, and Ha Quang Minh. Entropy-reeullriized 2-wasserstein distance between gaussian measures. arXiv preprint arXiv:2006.03416 2020.
Marc Mézard and Giorgio Parisi. On the solution of the random link matching problems. Journal de Physique, 1551:144, 1987.
Konstantin Mishchenko. Sinkhorn algorithm as special case of stochastic mirror descent. arXiv preprint urXiv:1909.066 2019.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Bur- naev. Large-scale wasserstein gradient flows. Advances in Neural Information Processing Sys- tems, 34:15243-15556, 2021.
Yurii Evgen' evich Nesterov. A method of solving convex programming problem with convergence rate of olk bbggg))) bigr). Doklady Akademii Nauk, volume 269, pp. 543-547. Russian Academy of Sciences, 1983.
Khai Nguyen, Dang Nguyen, Tung Pham, Nhat Ho, et al. Improving mini-batch optimal transport via partial transportation. In International Conference on Machine Learning, pp. 16656-16690. PMLR, 2022.
Tuan Nguyen, Trung Le, He Zhao, Quan Hung Tran, Truyen Nguyen, and Dinh Phung. Most: Multi-source domain adaptation via optimal transport for student-teacher learning. In Uncertainty in Artificial Intelligence, pp. 225-235. PMLR, 2021.
Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999.
Luca Oneto, Michele Donini, Giulia Luise. Carlo Ciliberto, Andreas Maurer, and Massimiliano Pontil. Exploiting mmnd and sinkhom divergences for fair and transferable representation learning. Advances in Neural Information Processing Systems, 33:1 5360--33-,,,, 2020.
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continu- normalizing flows via optimal transport. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 9223-9223 2021.
Giorgio Patrini, Rianne Van den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Uncertainty in Artificial Intelligence, pp. 733-743. PMLR, 2020.
12

---

Published as a conference paper at ICLR 2024
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport. Center for Research in Eco- nomics and Statistics Working Papers, (2017-86), 2017.
Mark S Pinsker. Information and information stability of random variables and processes, Holden- Day, 1964.
levgen Redko, Amaury Habrard, and Mare Sebban. Theoretical analysis of domain adaptation with optimal transport. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2017, Skopje, Macedonia, September 18-22, 28-22, 2017, Proceedings, Part II 10, pp. 737-753. Springer, 2017.
Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas, Improving gans using optimal transport. arXiv preprint 1803. 2018.
Roman Sandler and Michael Lindenbaum. Nonnegative matrix factorization with earth mover's s dis- tance metric for image analysis. IEEE Transactions on Pattern Analysis and Machine Intelligence, 3308):1590-1602, 2011.
Maziar Sanjabi, Jimmy Ba, Meisam Razaviyayn, and Jason D Lee. On the convergence and ro- bustness of training gans with regularized optimal transport. Advances Neural Information Processing Systems, 31, 2018.
Meyer Scetbon and Marco Cuturi. Linear time sinkhorn divergences using positive features. Ad- vances in Neural Information Processing Systems, 33:1 11133-3343333 2020.
Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pp. 9344 9354. PMLR, 2021.
Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv:1711.02283 2017.
Richard Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matri- ces. The annals of mathematical statistics, 35(2):876 879, 1964.
Justin Solomon, Fernando De Goes, Gabriel Peyré, Mareo Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional wasserstein distances: Efficient optimal transporta- tion on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1-11 2015.
J Michael Steele. Probability theory and combinatorial optimization. SIAM, 1997.
Alexis Thibault, Lénaïc Chizat, Charles Dossal, and Nicolas Papadakis. Overrelaxed sinkhorn- knopp algorithm for regularized optimal transport. Algorithms, 14(5):143, 2021.
Rosanna Turrisi, Rémi Flamary, Alain Rakotomamonjy, and Massimiliano Pontil. Multi-source domain adaptation via weighted joint distributions optimal transport. In Uncertainty in Artificial Intelligence, pp. 1970-1980 PMLR, 2022.
Titouan Vayer, Lactitia Chapel, Rémi Flamary, Romain Tavenard, and Nicolas Courty. Optimal transport for structured data with application graphs. arXiv preprint arXiv:1805.09111 2018.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Jonathan Weed. An explicit analysis of the entropic penalty in linear programming. In Conference On Learning Theory, pp. 1841-1855. PMLR, 2018
Yingxue Xu, Guihua Wen, Yang Hu, and Pei Yang. Unsupervised domain adaptation via deep hierarchical optimal transport. arXiv preprint arXiv:2211.11442 2022.
G Udny Yule. On the methods of measuring association between two attributes. Journal of the Royal Statistical Sociery, 75(6):579-63 652,, 1912.
13

---

Published as a conference paper at ICLR 2024
A EQUIVALENCE OF PRIMAL AND PRIMAL-DUAL FORM
We now show that the primal form equation 1 can be obtained from the primal-dual form by eliminating the dual variables.
Proposition Define
L(P, = P P - - - c),
and then the following equation holds:
max min L(P,z,y)
min
P log ..P
(9)
J.
P.Pl=r,P' 1=c
Moreover, for the Lyapunov potential function equation one has
f(x.y) = min L(P, x,
(10)
Proof. The use of Lagrange multiplier implies the following equality:
P log P PPC.P.
min max L(P. x, y) P.Pl=r,P 2,9
min
As L is concave in and convex in P. one can invoke the minimax theorem to interchange the operations of maximization and minimization. Therefore,
max L(P.x, Z.y
min P.PI-r,PTi-CT
P. lo. P C.P
(11)
In terms of entries, one writes L(P, v) as follows:
max min
Σ Pij log
CosPes
Pis Pis r4))
Σ Uj Σ Pij - -g).
(12)
21:9f Pij
We then solve the inner min problem explicitly by taking the derivative of Pı to zero, from which one obtains
Pi exp(n(- cij + x, UU) 1).
Plugging in the formula for Pij, one has
min L(P,r,)) P
= 1 2j
1 Σ
exp(n(-cc) + + m₁) 1) + Σ rits + Σ Cjys = f(x,z).
B PROOF THEOREM 2 AND THEOREM
We present the proof follows:
Proof. (of Theorem 2 Same as Theorem 1 we use the proof strategy that the approximate sparsity guarantee monotonically improves as Ps converges Py By Lemma 3.2 in Ghosal & Nutz (2022) and Pinsker's inequality, for 04k PP.1), one has
P 111 < (L(Px||))) + (L(P))|PP) < min (ou,, / Vor)
By Lemma in Altschuler et al. (017)) one has
oa 5 max (=) [( )),))) S max ((z) [[r, Ed::
where the second inequality comes from the definition of yk... From the proof in Theorem there exists P. E F such that
IP - 1* < 6n exp (-7A) min ((x) (44) ((x)))
Therefore, the statement in the theorem specializes to the more relevant where 6 1.
n
14

---

Published as a conference paper at ICLR 2024
Proof. (of Theorem
We begin by constructing the bipartite graph associated with an optimal transport plan P Let A, B be two index sets with An B = 9 and |A| [B] For optimal transport plan ps E R=*= we define its associated bipartite graph Gp. = Ep-) where (i.j) Ep- if and only if the (i, entry of P* is non-zero.
Suppose that there exists i, E A and j. E B with i " such that P is nonzero on the (f,5), (i, ). (i', j)), (t', 3") entries. Then, one consider perturbation E to P Specifically, E is +1 on entries (i, (V', j'), and is -1 on entries (i, (i', .)). and is zero everywhere else. As E1 E 1= 0, for sufficiently small 6, it follows that P eE is still feasible.
We note that one must have C E = 0, otherwise one would contradict the optimality of Pt, but this would mean Gy + Cyyy Cij' + Cgj. which contradicts the assumption on C. Thus, f cannot be simultaneously nonzero on the (i,j', (i', j), (i', j") entries. By the definition of Ep., it would mean that Gp- cannot simultaneously contain the edges (i.j), (i,j'). (i), (t',3"). Thus, in terms of graph-theoretic properties, we have shown that Gp. is C2.2 free, whereby C2.2 is the 2 x 2 clique. Then, by the C2 free property Theorem 2.10 in (Jukna 2011 shows Ep. cannot have more than nyn+ + edges. Thus T(P*) I 1-o(1) as desired.
v
C SINKHORN-NEWTON-SSAAN
FOR AUGMENTED LYAPUNOV FUNCTION
Since the Lyapunov function f satisfies f(z, v) = f(x + y1. - 11 for any scalar y, this implies that f has a degenerate direction of u = practice an augmented Lyapunov potential -1
1
Thus, to maintain numerical stability, we use in
fang (z, u) = f(x, v)
-
Σ x Σ 31) 2
1
Switching the augmented Lyapunov potential does not change the task, as (x* v/*) arg max fang (2. y) is simply the unique maximizer of f which satisfies Σ, क = E, Vr. More-
over, as f v -vuT, one can adapt the Hessian approximation
superposition of rank-1 1 and sparse matrix (Candès et al. 2011), which will likewise lead to O(n²) complexity.
We introduce Algorithm
2 a variation of Algorithm
[ This altered version uses the augmented
Lyapunov potential faug to accommodate the degenerate direction
fang (z, v) = f(z,8) 9)
2 Σ 35)
As the discussion in Section 4 shows, optimizing for faug is equivalent to optimizing for f. Algo- rithm differs from the original algorithm in two respects. First, the initialization of is through the output of the Sinkhorn stage after projection into the orthogonal complement Newton stage, one obtains the Hessian approximation term M with the superposition of sparse ma- trix Sparsify(V2f(2> p) and a rank-1 matrix wo Importantly, for À = r(pparsify [$)a)ssffyV""(((),,,) p)), the cost of left multiplication of M is O(An²) +((n). As the O(n) term is dominated by the (())²) term, the conjugate gradient step still has O(An²) scaling, Overall, the computational complexity of Algorithm is nearly identical to Algorithm]
v. Second, in the
15

---

Published as a conference paper at ICLR 2024
Algorithm 2 Sinkhomn-Newion-Spans (SNS) with augmented Lyapunov potential
Require: Sang: Vinit. E RR M/nit
M/ini E RR,,,, p, 0
1: # Sinkhorn stage
2: + -1
Initialize degenerate direction
3: (a, - (Zinit, Winit)
> Initialize dual variable
4: while N1 do
5:
P x exp (n(-C + æ1 T + 1y ) -1)
6:
I - x + (log(r) - log P1)) /n
7:
P - exp + x1 + +1y -1)
8:
- y (log(c) log( (S((P1)) 1)) /n
9:
i+i+1
10: end while
# Newton stage
12: 2 < Proj,- ((x, v))
D Project into non-degenerate direction of
13: while < Nı + N₂ do
14: M < Sparsify(V p) vv
Truncate with threshold p.
15:
Á: f- Conjugate_ Gradient( M, V Saug (=))
b. Solve sparse linear system Line search for step size
16:
Line_search(faar 2, Az)
17:
2 + ad>
18:
i+i+1
19:
end while
Output dual variables (z,
D WNKKORN-NEWTON WITHOUT SPARSITY
In this section, we show that the Sinkhorn- -eekoorn-Newton Newton algorithm (Brauer al. 2017) without account- ing for Hessian sparsity would be prohibitively slower than the SNS Algorithm. We remark that the Sinkhorn-Newton algorithm can be obtained from SNS by removing the Sparsify step in Algorithm In this case, one solves for the descent direction by directly using the Hessian, i.e., changing to
-
v(()).
Theoretically, this leads to On per- per-iteration complexity, which is considerably costlier than our best-scenario complexity of (n?)) under ((1/n) sparsity.
To empirically verify the impracticality of this method, we repeat the entropic random linear assign- ment problem experiment Section with 2000, N₁ 20 and r 5000. Table 2 summarizes our findings. As expected, we observe that the mankhor--Newton method is significantly slower than SNS, especially in terms per-iteration complexity. For larger ni, the Sinkhorn-Newic algorithm will be even more unfavorable.
Table 2:
Performance comparison between SNS and Sinkhorn-Newton during the Newton stage.
Both algorithms are run until they reach machine accuracy.
Method
Time (s)
Iterations
Time per iteration (s)
SNS
3.26
11
0.30
Sinkhom-Newton
118.56
10
11.86
E COMPARISON BETWEEN SINKHORN-NEWTON- QUASI-NEWTON METHODS
WITH
This section presents the result of quasi-Newton algorithms (Ncedal & Wright 1999) applied to entropic optimal transport problems. We show that, while being a reasonable proposal for solv- ing entropic optimal transport with second-order information, traditional quasi-Newton algorithms
16

---

Published as a conference paper at ICLR 2024
work poorly in practice. In short, a Quasi-Newton algorithm can be obtained from SNS by placing the Hessian approximation step in Algorithm [ Specifically, instead of sparse approxima- tion, a -qasi--Nettoo method approximates the Hessian M through the history of gradient informa- tion. In particular, we consider the -eaeed---lleceee-oodfaabbbbb miiteyd-emory Proyden-Fletchess Fletcher-GGoldfarbbbb-e- Sooldfarb-Shanno (L-BFGS) algorithm, which two of the most widely used quasi-Newton methods.
(BFGS) algorithm and the
We repeat the three experiment settings in Section 6 and the result is shown in Figure 2 a fair comparison, the quasi-Newton candidate algorithms are given the same Sinkhorn initializa- tion as in the SNS algorithm. As the plot shows, quasi- seasi-Newtoon algorithms do not show significant improvement over the Sinkhorn algorithm. Moreover, in the two experiments based on the MNIST dataset, both quasi-Newton candidates perform worse than the Sinkhorn algorithm in terms of run- time efficiency. As the iteration complexity of the quasi-Newton candidates does not exhibit the nu- merical super-exponential convergence shown in SNS, we conclude that noisy Hessian estimation from gradient history accumulation is inferior to direct sparse approximation the true Hessian matrix.
To ensure
F SINKHORN-NEWTON-SPAEEN SPARSE UNDER DIFFERENT ENTROPY REGULARIZATION PARAMETER
In this section, we show acceleration of SNS over the Sinkhorn algorithm under wider range for the entropy regularization parameter 7). In particular, we focus on the setting of MNIST image under I and l2 costs. Importantly, it is common practice pixel distance is used to form the cost matrix. For example, the Sarth-Mover distance (EMD) considered in (Altschuler et al.) 2017) is defined by
dpisel ((i, f). (t', (() = -i 2'| + Lj J'I.
In our work, a pixel (i,j) embedded to the point (8/28.j/28) E R2 before the distance function is applied. Thus, this text uses the normalized distance function
1' = 28
j
28
As d 28 dpixel. our choice of n 1200 in Section 6 is equivalent to choosing 1) = 1200/28 s 42. As the range used in Altschuler et al. (2017) is 17 [1.9]. our n is similar to the range of entropy regularization commonly used. To show the performance of SNS is robust under different 7. we benchmark the performance of SNS under an extended practical choice of
28k for 3, 5, 7, 9, 11]. For the Sinkhorn stage warm initialization, we take N1 = 10k + 100 for each 7) 28k, and the target sparsity is taken to be X 15/n. In Table 3 we show the performance of SNS compared with the Sinkhorn algorithm to reach machine accuracy. One can see that the SNS algorithm consistently outperforms the Sinkhorn algorithm, and the improvement is more significant under larger choices of 17.
For further validation, in Figure 3a we plot the Wasserstein W1 transport distance for the entropy regularized optimal solution n under different v), which shows indeed that v) zs 150 is sufficient for the practical goal of obtaining transport plan with relatively low transport cost.
For the case of squared l2 distance, as the squared 12 distance is scaled by a factor of 576, we bench- mark the performance of SNS under 7) 576k for = {1, 5, 7,9, 11]. J. As can be seen in Figure 3, one needs 1) P 1000 to reach within 1 1% accuracy of the ground-truth transport cost, which is why the range of 17 considered a reasonable choice. For the Sinkhorn stage warm initialization, we take Nı = 10k for each 7) 576k, and the target sparsity is taken to be X = 4/n. In Table we show the performance of SNS compared with the Sinkhorn algorithm to reach machine accuracy. One can see that the SNS algorithm consistently outperforms the Sinkhorn algorithm, and the improvement is more significant under larger choices of r). Moreover, for the case of r) = 576, even though the presence of entropy regularization is strong, the numerical result shows that À = 4/r nz in the sparse Hessian approximation is sufficient to reach machine accuracy rapidly.
17

---

Published as a conference paper at ICLR 2024
Sinkhorn
BFGS
End of Sinkharn stage terminates
Sinxhon--NewoonSparre
Convergence iterations
Convergence wall
Log Optimallty Gap
10
-15
-35
-20
-20
-25
-25
-30
-30
-15
50
100
150
200
250
10
Iterations
Secands
(a) Entropic random linear assignment
Convergence iterations
Convergence wall
-10
-10
Log Optimality
-15
-15
-20
-00
-25
-25
-10
-30
50
100
150
200
250
10
Iterations
Secands
(b) Optimal transport on the MNIST dataset under transportation cost ||
vl1
Convergence Iterations
Convergence in wall
Log Optimallty Gap
-10
-10
-15
-35
-20
-20
-30
-30
-35
-35
400
600
1000
20
20
25
30
Iterations
Seconds
(c) Optimal transport on the MNIST dataset under transportation cost || vll
Figure 2: Performance of Quasi-Newton methods, compared against the algorithm and the Sinkhorn algorithm.
18

---

Published as a conference paper at ICLR 2024
Table 3: Performance comparison between SNS and Sinkhorn for different under the transporta- tion cost ||.- Both algorithms are run until they reach machine accuracy. The time and iteration of the SNS method refers to the combined time and iteration of the two stages combined.
Entropy
Method
Time (s)
Iterations
SNS
1.45
110
7) = 28
Sinkhorn
1.96
173
SNS
5.32
147
7 84
Sinkhorn
9.63
899
SNS
5.41
167
77 140
Sinkhorn
14.53
1399
SNS
6.10
189
196
Sinkhorn
15.56
1499
SNS
7.78
216
17 252
Sinkhorn
17.81
1699
SNS
8.08
236
77 = 308
Sinkhorn
19.76
1899
Table 4: Performance comparison between SNS and Sinkhorn for different under the transporta- tion cost || yl. Both algorithms are run until they reach machine accuracy. The time and iteration of the SNS method refers to the combined time and iteration of the two stages combined.
Entropy
Method
Time (s)
Iterations
SNS Sinkhorn
3.40
33
7 = 576
9.95
946
SNS
5.10
64
1728
Sinkhorn
32.92
3072
SNS
7.08
96
17 2880
Sinkhorn
62.96
6083
SNS
9.70
134
4032
Sinkhorn
108.24
10 596
SNS 12.83
177
17 5184
Sinkhorn
166.97
16 299
SNS
22.30
259
6336
Sinkhorn
248.43
23 498
19

---

Published as a conference paper at ICLR 2024
Transport cost w.r.t. entropy parameter
0.185
0.180
W, Cost
0.175
0.170
0.165
50
100
150
200
250
300
(a) MNIST dataset under transportation cost ||
y|1
Transport cost w.r.t. entropy parameter
0.1650
0.1645
0.1640
W, Cost
0.1635
0.1630
0.1620
1000
5000
6000
n
(b) MNIST dataset under transportation cost ||.
y|1
Figure Optimal transport cost of the obtained entropic regularized solution for different r).
20


Reviews:
[{"review_id": "AoLK3MDTi6", "reviewer": "Reviewer_ZDaz", "summary": "Despite the success of the Sinkhorn algorithm, its runtime may still be slow due to the potentially large number of iterations needed for convergence. To achieve possibly super-exponential convergence, we introduce Sinkhorn-Newton-Sparse (SNS), an extension to the Sinkhorn algorithm, by introducing early stopping for the matrix scaling steps and a second stage featuring a Newton-type subroutine.", "questions": "No", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Sparsification of the Hessian results in a fast per-iteration complexity, the same as the Sinkhorn algorithm. In terms of total iteration count, we observe that the SNS algorithm converges orders of magnitude faster across a wide range of practical cases, including optimal transportation between empirical distributions and calculating the Wasserstein distance of discretized continuous densities. The empirical performance is corroborated by a rigorous bound on the approximate sparsity of the Hessian matrix.", "weaknesses": "1. The theoretical results in this paper seem to be simple corollaries of existing resluts. \nFor example, Eq. (7) which is crucial to the proof of Theorem 1, follows Corollary 9 of (Weed, 2018). \n2. Theorem 1 requires that $\\eta$ is large enough. However, I believe that a large $\\eta$ will bring something negative. \nOtherwise, why don't people minimize the original problem without the regularization? Thus, I think the author should remark the condition that $\\eta$ is large enough of Theorem 1 and clarify how the value of $\\eta$ affect the convergence rate.\n3. The paper claims that ``to achieve possibly super-exponential convergence, we introduce Sinkhorn-Newton-Sparse (SNS)''. However, this paper does not provide any sound convergence analysis of SNS.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8t0R8i4z5X", "reviewer": "Reviewer_39Qf", "summary": "This paper introduces the sinkhorn-newton-sparse algorithm to solve the regularised optimal transport problem as described in eq (1). This improves on the previous work by Brauer et al. to the case of sparse transport matrices. The authors define the notion of sparsity as in Definition 1 and provide the rate of convergence of the modified sinkhorn algorithm in theorem 1. The authors then provide the approximate sparsity rate for the newton step in theorem 2. The paper is closed with numerical examples.", "questions": "Can the authors explain the notion of sparsity introduced in this paper? It seems to me that $\\tau$ is just the proportion of non-zero elements. However, is it usual to approximate $M$ by a sparse $\\bar M$?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "There is a notable reduction in the number of iterations and the time to convergence while solving a problem using Algorithm 1. It also seems that the definition of sparsity plays a key role in the analysis.", "weaknesses": "One notable omission I felt was the lack of guarantees for computational complexity. Although the algorithm performs well numerically, there is no theoretical backing for that.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M68iKCS7I0", "reviewer": "Reviewer_CUDR", "summary": "This paper proposes a method for the fast numerical solution of the Kantorovich problem with entropic regularization. The submission is closely related to the cited work of Brauer, Clason, Lorenz and Wirth that explored the use of Newton iterations for this problem. The main contributions are two-fold:\n1. The observation (that the authors prove) that after several Sinkhorn iterations, the Hessian matrix becomes close to sparse.\n2. An algorithm based on this observation that combines several iterations of Sinkhorn followed by Newton iterations with a sparsified Hessian. The Hessian is sparsified by zeroing-out small entries.", "questions": "* For the numerical section, I would like to see numerical results for the same data sets with several regularization strengths (e.g. \\eta=10,100,1000). Additionally, I think it is important to compare the proposed method to other current methods for accelerating OT with entropic regularization. In particular to \"Massively scalable Sinkhorn distances via the Nystrom method\" by Altschuler, Bach, Rudi, and Niles-Weed (2019). \n\n* Have you considered any quasi-Newton methods? Will these be applicable in the context of your paper?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper focuses on the important and timely problem of efficiently computing regularized optimal transport.\n* The authors clearly explain the potential advantages of second-order methods while noting the inefficiency of classic Newton iterations. This motivates their approach of using a sparse approximation to the Hessian.\n* The convergence of the Hessian to a nearly-sparse matrix throughout the steps of the algorithm is proved (in both the Sinkhorn and Newton stages).\n* The method is clearly explained and the results should be easy to reproduce.", "weaknesses": "* A potential red flag is that, in the numerical results, only a single entropic regularization coefficient is considered (\\eta=1200). I suspect that this is a weak regularization term which results in rather sparse couplings and is therefore favorable to their method. It is OK if the method is not beneficial in the strong-regularization regime. However, this needs to be acknowledged in the paper and not swept under the rug.\n\n* The method is not compared to any other current methods aside from the classical methods on which it is based (Sinkhorn and Newton).\n\nThese two weaknesses taken together are the main reason for my \"marginally below\" score. If they are properly addressed I will increase my rating of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "neX5nnBFgG", "reviewer": "Reviewer_orpT", "summary": "The authors focus on accelerating the Sinkhorn algorithm for optimal transport (OT). The Sinkhorn algorithm is an iterative matrix scaling algorithm used to solve the entropy-reularized OT problem. Its per-iteration complexity is $O(n^2)$ for discrete distributions of size $n$. Its convergence rate behaves like a polynomial converging method, which can be slow. Other optimization methods like second-order methods enjoy super-exponential convergence, but a greater per-itereration cost $O(n^3)$ (e.g. Newton). Here they find a link between the Sinkhorn iterations and the sparsity of the Hessian of the Lyapunov potential. Once the Hessian is sparse, the per-iteration cost can reach $O(n^2)$, hence the same per-iteration cost as Sinkhron while benefiting from the faster convergence rate of the second-order methods. The authors propose to combine the two methods, i.e. Newton iteration once the Hessian is sparse, which improves the convergence rate. They provide quantitative results to validate the algorithm.", "questions": "Questions and minor comments.\n\n- I believe that to compute the optimality gap (e.g. Fig.1) you need a ground truth (i.e. the minimizer), in this case how is it computed ? Is it the OT plan without entropy regularization ?\n- I might be missing something, but I don't see the target sparsity $\\lambda$ in Alg.1.\n- How do you choose the number of Sinkhorn steps $N_1$? Would it be possible to switch to the Newton stage once the Hessian is sparse enough ? \n- The related work section is great, but maybe you can add other references that speed up Sinkhorn using a factorization of the ground cost [1,3] or Chebyshev polynomials with a sparse graph [2]. \n- The results in Tab.1 are on the number of iterations until reaching machine accuracy. They don't inform us on the quality of the approximation. It could be interesting to compare the accuracy of the plan in the case where we have a closed form (e.g. entropic OT between Gaussian distributions [4]). \n- The authors could add a brief review of Newton’s algorithm since it is an important piece of the paper. In particular, on the importance of solving the Hessian system. \n- In eq. 1, you need to define $\\eta$ and specify if it should be greater than zero.\n- Having a preliminary or background section could help. Parts of the introduction could be in this section, and definition such as the optimality gap could be added in that section.\n- In section 2 \"Convergence of Sinkhorn\", what is $\\alpha$ ? Should it be $\\eta$ ?\n\n[1] Scetbon, Meyer, Marco Cuturi, and Gabriel Peyré. \"Low-rank Sinkhorn factorization.\" International Conference on Machine Learning. PMLR, 2021.\n\n[2] G. Huguet, A. Tong, M. R. Zapatero, C. J. Tape, G. Wolf and S. Krishnaswamy, \"Geodesic Sinkhorn For Fast and Accurate Optimal Transport on Manifolds,\" 2023 IEEE 33rd International Workshop on Machine Learning for Signal Processing (MLSP).\n\n[3] Scetbon, Meyer, and Marco Cuturi. \"Linear time Sinkhorn divergences using positive features.\" Advances in Neural Information Processing Systems 33 (2020): 13468-13480.\n\n[4] Mallasto, A., Gerolin, A., & Minh, H. Q. (2022). Entropy-regularized 2-Wasserstein distance between Gaussian measures. Information Geometry, 5(1), 289-323.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The overall presentation of the manuscript is good, and the method is clear.\n- The method accelerates the convergence of the Sinkhorn algorithm without increasing the per-iteration cost.  \n- The method is theoretically motivated by theorem 1; we know that after some Sinkhorn steps the Hessian can be sparse, hence the per-iteration cost of Newton's methods is lower.", "weaknesses": "- I find the evaluation a bit limited. In particular, it would be great to see more results for different regularization parameter, and include standard deviation in the tables. \n- It would be great to have comparison with other methods than Sinkhorn, although I agree that most methods can be combined with SNS, it would be beneficial to compare their convergence.\n- Have you tried removing the Sinkhorn stage  ($N_1=0$) ? I would like to understand how much of the sparsity is due to Sinkhorn (Thm1) vs the `Sparsify` step in Alg.1, and its influence on the convergence.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)