SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport

SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport Ian Hsieh1,∗ , Soumya Snigdha Kundu2 , Tom Vercauteren2,† , Reuben Dorent1,† 1

Sorbonne Université, Institut du Cerveau - Paris Brain Institute - ICM, CNRS, Inria, Inserm, AP-HP, Hôpital de la Pitié Salpêtrière, 75013 Paris, France 2 King’s College London, UK

arXiv:2608.28262v1 [cs.LG] 28 Aug 2026

∗

Corresponding author: ian-yee-yang.hsieh@inria.fr. † Equal contribution.

Abstract Entropic optimal transport (EOT) has been shown to offer a computationally tractable approximation to exact optimal transport. However, the standard Sinkhorn-Knopp algorithm has two main limitations. First, given discrete measures with N points, each iteration requires O(N 2 ) operations, which restricts its use on large-scale datasets (e.g. N ≥ 104 ). Second, it uses the independent coupling as a reference measure for regularisation. This assigns mass to high-cost transport edges at moderate regularisation strengths. We propose SinkSLOT, which addresses both limitations by putting forth the expected sliced lifted transport plan as a natural way to sparsify the Gibbs kernel with a non-independent prior coupling. We prove that: 1) SinkSLOT converges; 2) with L slices, each resulting sparse Sinkhorn iteration costs O(LN ); and 3) the resulting objective is a divergence requiring no debiasing. Experiments on synthetic benchmarks show that SinkSLOT delivers substantial speedups over state-of-the-art dense and sparse EOT methods. We also demonstrate the applicability of the proposed divergence in a gradient flow experiment. The code is publicly available at https://github.com/cai4cai/SinkSLOT.

Introduction Optimal transport (OT) and the associated Wasserstein distance (WD) have become fundamental tools in machine learning and beyond [Peyré and Cuturi, 2019, Villani, 2009]. However, computing exact OT requires solving a linear programme with O(N 3 log N ) complexity for discrete measures supported on N points, limiting its applicability to large-scale problems. This computational cost is particularly acute when WD is used as a loss function within an optimisation pipeline, such as training a generative neural network, where its gradient must be computed at every optimisation step. Entropic OT (EOT) [Cuturi, 2013] offers a computationally tractable approximation for the OT plan by adding entropy regularisation. The resulting strictly convex objective admits a unique regularised transport plan, computable via the Sinkhorn-Knopp algorithm [Sinkhorn and Knopp, 1967]. For two N -point measures, standard dense Sinkhorn iterations require O(N 2 ) time. The resulting runtime is thus substantially lower than exact OT but still computationally demanding for large datasets. Sliced optimal transport (SOT) [Bonneel et al., 2015] provides a computationally attractive alternative to EOT, reducing runtime to O(N log N ) per projection through 1D sorting. However, it can underperform exact OT as a distance measure [Rowland et al., 2019]. The lack of a transport plan in the original space further limits its utility when an explicit coupling is needed. Recently, slice-based approaches that recover a transport plan by lifting the 1D matchings back to the original space have been proposed [Liu et al., 2025]. Recently, Nguyen [2026] merges these two frameworks by replacing the reference coupling of EOT with the lifted SOT plan, introducing sliced-regularised optimal transport (SROT). SROT provides more 1

Inputs

Stage 1: Sparse SOT Construction (a) Sample L directions θ ∼ Uniform(Sd−1 )

Stage 2: Sparse Sinkhorn

θ3

Prohibitive pairwise cost: O(N ×M )

θi

Output (a) Sparse coupling P⋆ε

Kε = PSOT ⊙exp(−C/ε) Inherits sparsity of PSOT (b) Log-domain Sinkhorn loop φ-update

O(N log N ) per slice

(c) Aggregate the per-slice plans N L X 1 PSOT = L PSOT θℓ M

µ ν source target N points M points

(a) Form sparse kernel M

θ1 θ2

N

(b) 1D OT by sorting

φi ← log ai − LSEj∈Si(log Kij + ψj ) ψj ← log bj − LSEi∈S j(log Kij + φi ) ψ-update repeat until ∥marginal error∥ < τ

ℓ=1

S = |S| ≤ L(N +M −1)

O(S) per iteration

(b) SinkSLOT divergence SLOTε (µ, ν) = ⟨C, P⋆ε ⟩
+ ε·KL P⋆ε PSOT

✓ No debiasing needed ✓ Gradient and Hessian

Figure 1: Overview of the SinkSLOT pipeline. Inputs: two discrete measures µ and ν; the naïve dense coupling has N × M nonzeros. Stage 1: L random directions are sampled and 1D OT problems are solved by sorting in O(N log N ) per slice; the resulting plans are aggregated into a sparse lifted plan PSOT with at most L(N + M − 1) nonzeros. Stage 2: a sparse Gibbs kernel Kε = PSOT ⊙ exp(−C/ε) inherits this sparsity, reducing each log-domain Sinkhorn iteration from O(N M ) to O(L(N + M )). Output: the sparse coupling P⋆ε and the SinkSLOT divergence SLOTε , which requires no debiasing. accurate approximations of the true OT plan than EOT under the same level of regularisation, and significantly outperforms SOT. Both SROT and EOT are solved via Sinkhorn-type algorithms, providing efficient approximations of the OT plan at reduced cost compared to exact solvers. However, the computational complexity of SROT remains O(N 2 ) per Sinkhorn iteration due to the dense kernel, substantially higher than the O(N log N ) complexity of SOT. The primary objective of Nguyen [2026] is thus to improve coupling accuracy over EOT, rather than to retain the computational efficiency of SOT. In this work, we build upon Nguyen [2026] and show that the sparsity of the SOT plan can be leveraged to formulate an efficient sparse Sinkhorn-type algorithm. Rather than operating on the full N × N coupling matrix, the Sinkhorn iterations are restricted to the sparse support defined by the SOT plan, reducing periteration complexity from O(N 2 ) to O(LN ), where L is the number of projection slices. As a result, the per-iteration cost remains substantially below EOT, making the method significantly more scalable than dense Sinkhorn while retaining the improved coupling accuracy of SROT. Our contributions are threefold. 1. We formalize SinkSLOT, the sparse version of Nguyen [2026], and prove key properties, including convergence and uniqueness of the solution. 2. We show that the SinkSLOT objective is a divergence without requiring EOT debiasing terms and admits an analytical gradient, making it suitable for optimisation and learning tasks. 3. Experiments on synthetic benchmarks demonstrate that SinkSLOT achieves substantial speedups over state-of-the-art dense and sparse EOT methods.

Background and related work In this section, we cover the necessary background on discrete OT for this paper. Readers are referred to Peyré and Cuturi [2019] for a comprehensive review, and to Peyré [2025] for a treatment more geared toward machine learning. Nguyen [2025] serves a similar purpose for the sliced problem. P PM Optimal Transport (OT) Given two discrete probability measures µ = N i=1 ai δxi and ν = j=1 bj δyj , N M d with positive unitary weight vectors a ∈ ∆ , b ∈ ∆ and support points xi , yj ∈ R , and given a ground ×M cost matrix C ∈ RN with entries Cij = c(xi , yj ), typically Cij = ∥xi − yj ∥p , p ≥ 1, the OT +

2

problem [Villani, 2009, Peyré and Cuturi, 2019] seeks a cost-minimising coupling: P OT ∈ arg min ⟨C, P ⟩,

(1)

P ∈Π(a,b)

 ×M where Π(a, b) = P ∈ RN | P 1M = a, P ⊤ 1N = b is the set of admissible couplings (joint + measures) with marginals a and b. When Cij = ∥xi − yj ∥p with p ≥ 1, the optimal coupling P OT defines
1/p the p-Wasserstein distance Wp (µ, ν) = ⟨C, P OT ⟩ . Problem (1) is a linear programme whose best known general-purpose solvers require O(N M (N + M ) log(N + M )) arithmetic operations [Peyré and Cuturi, 2019], limiting its use as a loss function within large-scale optimisation pipelines. Entropic Optimal Transport (EOT) To alleviate this cost, Cuturi [2013] popularised an entropy-regularised relaxation of (1):
PεEOT = arg min ⟨C, P ⟩ + ε KL P ∥ a ⊗ b , (2) P ∈Π(a,b)

where ε > 0 is a regularisation strength, a⊗b denotes the independent coupling with entries (a⊗b)ij = ai bj , P P and KL(P ∥Q) = ij Pij log Qijij − Pij + Qij is the generalised Kullback-Leibler divergence. The unique solution has the form PεEOT = diag(u) K diag(v), where Kij = ai bj exp(−Cij /ε) is the Gibbs kernel and (u, v) are positive scaling vectors found by the Sinkhorn-Knopp algorithm: uℓ+1 = a ⊘ (Kvℓ ),

vℓ+1 = b ⊘ (K ⊤ uℓ+1 ),

(3)

where ⊘ is elementwise division. The solution defines a discrepancy functional: OTε (µ, ν) = ⟨C, PεEOT ⟩ + ε KL(PεEOT ∥a ⊗ b).

(4)

Despite its widespread adoption, EOT has three well-known limitations. First, although, for structured C matrices, streaming backends help avoid the O(N M ) memory to store K at the cost of extra online computations [Feydy et al., 2019, Ye et al., 2026], each Sinkhorn iteration requires two dense matrix-vector products, resulting in O(N M ) cost per iteration. This quadratic scaling is prohibitive for very large point clouds (e.g. N, M ≥ 104 ). Second, regularising toward the cost-agnostic independent coupling a ⊗ b biases the plan toward dense and diffuse matchings that potentially assign non-negligible mass to higher-cost transport edges, particularly for moderate values of ε. Finally, the EOT cost is biased in the sense that OTε (µ, µ) ̸= 0 in general, so that using it directly as a loss requires an additional debiasing step to obtain the Sinkhorn divergence [Feydy et al., 2019]. Sparsification techniques for OT Several methods reduce the complexity of OT by imposing sparsity or low-dimensional structure on the transport plan. We refer to Khamis et al. [2024], Ouyang et al. [2026] for broad surveys. Blondel et al. [2018] induce sparse transport plans through squared-L2 and group-lasso regularisation, while Liu et al. [2023] constrain each column to contain at most k nonzeros. Low-rank alternatives include factored couplings [Forrow et al., 2019, Scetbon et al., 2021, Scetbon and Cuturi, 2022] and low-rank-plus-sparse decompositions [Liu et al., 2021]. These methods require specialised optimisation procedures rather than standard Sinkhorn scaling on a fixed kernel. More closely related to our work, sparse-kernel methods reduce the O(N M ) cost of Sinkhorn iterations. Schmitzer [2019] adaptively truncate the kernel using the current dual potentials, Gasteiger et al. [2021] retain local interactions via locality-sensitive hashing, and Li et al. [2023] pre-build a sparse kernel by importance sampling so each Sinkhorn iteration costs O(|S|) for support S. However, the latter requires O(N M ) preprocessing to construct their sparse kernel. Furthermore, these methods all use the independent coupling as the reference measure, so the value of the nonzero entries do not contain geometric information. 3

Reference coupling beyond a⊗b Non-product reference measures have been considered in the Schrödingerbridge literature [Léonard, 2014], though they have received little attention in static EOT. Freulon et al. [2025] study EOT with a Gaussian reference coupling, but restrict their analysis to Gaussian marginals. Another framework to efficiently approximate OT is sliced optimal transport (SOT) [Bonneel et al., 2015], which averages 1D Wasserstein distances over random projections, each solved via sorting, giving O(L(N + M ) log(N + M )) cost for L directions. To produce a transport plan in the original space, Muzellec and Cuturi [2019], Liu et al. [2025], Tanguy et al. [2025] propose lifting the 1D optimal matchings back to the original support points. For each direction θ ∈ Sd−1 , let PθSOT ∈ Π(a, b) denote the coupling that solves OT inP the projected 1D space. Averaging over L directions yields the expected sliced transport plan: 1 SOT SOT . By construction, P SOT ∈ Π(a, b) is sparse. Each 1D plan contributes at most P = L L ℓ=1 Pθℓ
N +M −1 nonzero entries [Brualdi, 2006, Theorem 8.1.2], so |S| = O L(N +M ) where S = supp(P SOT ), which is much smaller than N M when L ≪ min(N, M ). Nguyen [2026] observed that the sliced plan P SOT is a more informative prior than the independent coupling and proposed sliced-regularised OT (SROT):
SROT Pε,γ = arg min ⟨C, P ⟩ + ε KL P ∥ PγSOT , (5) P ∈Π(a,b)

where PγSOT = (1 − γ)P SOT + γ a ⊗ b is smoothed with γ > 0 to guarantee full support. SROT yields couplings closer to the true OT plan than EOT under the same regularisation level, but the smoothing destroys the sparsity of P SOT . Each Sinkhorn iteration still costs O(N M ), the same as EOT. Similar to (4), the optimal solution leads to a biased discrepancy functional SROTε,γ (µ, ν).

SinkSLOT: Sinkhorn via Sparse Lifted OT Problem formulation We propose to solve the SROT problem (5) with the unsmoothed, i.e. γ = 0, expected slice transport reference plan P SOT :
Pε⋆ = arg min ⟨C, P ⟩ + ε KL P ∥ P SOT , (6) P ∈Π(a,b)

Absorbing the transport cost into the KL divergence, we obtain the equivalent formulation
Pε⋆ = arg min KL P ∥ Kε ,

(7)

P ∈Π(a,b)

where Kε = P SOT ⊙ exp(−C/ε) is the sparse Gibbs kernel. The SinkSLOT discrepancy associated with a SOT plan is SLOTε (µ, ν) = ⟨C, Pε⋆ ⟩ + ε KL(Pε⋆ ∥P SOT ). (8) Setting γ = 0 confers key advantages over the smoothed formulation. Assuming no overlapping projected samples (zero probability under sufficient numerical precision), each 1D optimal matching contributes at most N + M − 1 nonzero entries to its corresponding PθSOT . Consequently, the support of the aggregated lifted plan, S := supp(P SOT ), satisfies |S| ≤ L(N + M − 1). Since the Gibbs kernel Kε inherits the sparsity pattern
of P SOT , the cost of each Sinkhorn iteration is reduced from O(N M ) to O(|S|), and to O L(N + M ) in the worst case. To avoid infinite costs, optimising the KL divergence forces supp(P ) ⊆ S for any feasible coupling P with finite objective value. Hence, the minimisation in (7) is implicitly restricted to couplings supported on S.

4

Algorithm 1: SinkSLOT pseudo-code M Require: Source {(xi , ai )}N i=1 , target {(yj , bj )}j=1 , L projections, regularisation ε > 0, and tolerance τ >0 Ensure: Sparse coupling Pε⋆ and discrepancy SLOTε (µ, ν)

1: Sample θ1 , . . . , θL

i.i.d.

∼ U(Sd−1 )

2: for ℓ = 1, . . . , L do

Project: x̂i = θℓ⊤ xi and ŷj = θℓ⊤ yj 4: Solve the 1D OT problem by sorting → PθSOT ℓ 5: end for P L 1 SOT 6: P SOT = L ℓ=1 Pθℓ 7: Compute cost Cij = c(xi , yj ), ∀(i, j) ∈ supp(P SOT ) SOT − C /ε 8: Form sparse log-kernel: log(Kε )ij = log Pij ij (0) (0) 9: Initialise φ = 0N and ψ = 0M 10: repeat (ℓ+1) (ℓ)
= log ai − LSEj∈Si log(Kε )ij + ψj , ∀i 11: φi (ℓ+1) (ℓ+1)
12: ψj = log bj − LSEi∈S j log(Kε )ij + φi , ∀j 13: until marginal violation < τ 14: Pε⋆ = diag(eφ )Kε diag(eψ ) 15: Compute SLOTε (µ, ν) from Pε⋆ 16: return Pε⋆ and SLOTε (µ, ν) 3:

{sparse, nnz ≤ L(N + M )}

Existence, uniqueness, and convergence of sparse Sinkhorn We first establish that (6) admits a unique solution of diagonal-scaling form, then show that the sparse Sinkhorn algorithm converges to it. Proposition 1 (Unique diagonal scaling). Let Kε = P SOT ⊙ exp(−C/ε) be the sparse Gibbs kernel with support S = supp(P SOT ). The SinkSLOT problem (7) has a unique solution of the form Pε⋆ = diag(u) Kε diag(v),

(9)

M for positive scaling vectors u ∈ RN >0 and v ∈ R>0 .

The proof is in Appendix A.1. Proposition 2 (Convergence of sparse Sinkhorn). Let (uℓ , vℓ ) be defined by vℓ+1 = b ⊘ (Kε⊤ uℓ+1 ),

uℓ+1 = a ⊘ (Kε vℓ ),

(10)

initialised at u0 = v0 = 1. P ℓ = diag(uℓ ) Kε diag(vℓ ) converges to the unique solution Pε⋆ of (6). The proof is in Appendix A.2. SinkSLOT pipeline, computational complexity Figure 1 and Algorithm 1 summarise SinkSLOT. The total cost decomposes into (i) construction of the sparse SOT plan and (ii) sparse Sinkhorn iterations. Stage 1 corresponds
to the SOT construction. For each of the L directions, projecting N + M points costs O (N + M )d , while sorting the projected source and target points costs O(N log N + M log M ). Constructing the resulting 1D transport plan requires O(N +M ) operations and produces at most N +M −1
nonzero entries. Accumulating the L sparse plans therefore costs O L(N + M ) . Hence, the total cost of

this stage is O L (N + M )d + N log N + M log M . 5

Table 1: Computational complexity of the reference coupling and each Sinkhorn iteration. Here, S = |S| ≤ L(N + M − 1) denotes the number of nonzeros in the SOT plan, and we assume N = M for simplicity. Method EOT [Cuturi, 2013] SROT [Nguyen, 2026] SinkSLOT (ours)

Reference coupling

Reference coupling cost

Sinkhorn iter.

a⊗b (1 − γ)P SOT + γ a ⊗ b P SOT

O(N 2 ) O(LN d + LN log N + N 2 ) O(LN d + LN log N )

O(N 2 ) O(N 2 ) O(S)

In Stage 2, the sparse Sinkhorn is applied. Each iteration of (10) performs two sparse matrix-vector products on Kε , which has |S| ≤ L(N + M ) nonzero
entries, costing O(|S|). For T iterations until convergence, the Sinkhorn stage costs O T · L(N + M ) . The cost entries Cij are computed only for (i, j) ∈ S at O(d) per entry, adding O(|S| d). For small ε, the entries of Kε can span many orders of magnitude. We work entirely in the log domain, C defining φℓ = log uℓ , ψ ℓ = log vℓ and log(Kε )ij = log PijSOT − εij . The Sinkhorn updates then read   (ℓ+1) (ℓ) φi = log ai − LSEj∈Si log(Kε )ij + ψj , (11)   (ℓ+1) (ℓ+1) , (12) ψj = log bj − LSEi∈S j log(Kε )ij + φi where Si = {j : (i, j) ∈ S}, S j = {i : (i, j) ∈ S}, log-sum-exp (LSE) uses an online max-shift trick for stability [Ye et al., 2026, Milakov and Gimelshein, 2018, Nowozin, 2016].

The total complexity is O L (N + M )(d + T ) + N log N + M log M . In terms of memory, only the sparse kernel (O(|S|)), the scaling vectors (O(N + M )), and the cost entries on S (O(|S|)) need to be stored. This sparsity is intrinsic to our formulation. Even without matrix-free backends such as KeOps [Feydy et al., 2019] or fused Triton kernels [Ye et al., 2026], the full N × M cost or kernel matrices are never formed. SinkSLOT as a divergence A practical advantage of the unsmoothed formulation is that the SinkSLOT functional (8) is already a well-behaved discrepancy without debiasing. Recall that with standard entropic regularisation, OTε (µ, µ) > 0 in general. Feydy et al. [2019] proposed to subtract self-transport terms to obtain a proper divergence: Sε (α, β) = OTε (α, β) −

1 1 OTε (α, α) − OTε (β, β) 2 2

The smoothed SROTε,γ similarly requires debiasing. In contrast, we show that SinkSLOT bypasses these issues entirely. Proposition 3 (SLOTε is a divergence). Assume Cij = ∥xi − yj ∥p for some p ≥ 1, and that the support points of µ and ν are in general position.1 Then, for a fixed set of projection directions, SLOTε satisfies: 1. Non-negativity. SLOTε (µ, ν) ≥ 0. 2. Identity of indiscernibles. SLOTε (µ, ν) = 0 ⇔ µ = ν. 3. Symmetry. SLOTε (µ, ν) = SLOTε (ν, µ). In particular, SLOTε is a divergence in the sense of Feydy et al. [2019], with no debiasing required. The proof is in Appendix A.3. 1

General position ensures that, for almost every projection direction, projected coordinates are distinct, so the 1D OT matching is uniquely determined.

6

Linear convergence rate Having established convergence of the sparse Sinkhorn iterates (Proposition 2), we now show that convergence is linear (geometric) and give an explicit expression for the asymptotic contraction factor. Definition 1 (Support graph). The bipartite support graph GS = (U ∪ V, S) has row vertices U = {1, . . . , N }, column vertices V = {1, . . . , M }, and edge set S = supp(P SOT ). Proposition 4 (Spectral gap). Define the normalised plan matrix A = diag(a)−1/2 Pε⋆ diag(b)−1/2 ∈ RN ×M ,

(13)

and let σ1 ≥ σ2 ≥ · · · ≥ 0 denote its singular values. Then: √ √ 1. σ1 = 1, with left and right singular vectors a and b, respectively. 2. σ2 < 1 if and only if GS is connected. Proposition 5 (Linear convergence rate). When GS is connected, the sparse Sinkhorn iterates P ℓ = diag(uℓ ) Kε diag(vℓ ) converge linearly: 1/ℓ

lim sup P ℓ − Pε⋆ 1 ℓ→∞

≤ σ22 ,

(14)

where σ2 is the second singular value of A from Proposition 4. In particular, the number of iterations to achieve ∥P ℓ − Pε⋆ ∥1 ≤ δ is at most   log(1/δ) T = O , (15) log(1/σ22 )
giving a total arithmetic cost of O L(N +M ) · log(1/δ) / log(1/σ22 ) . If GS is disconnected, the scaling problem decomposes into independent subproblems on the connected components [Kalantari et al., 2008], and the above applies to each with its own second singular value. The proofs are in Appendix A.4. The convergence rate σ22 depends on ε through the optimal plan Pε⋆ . As ε → 0, Pε⋆ converges to the KL-selected minimiser of the support-restricted OT problem (Proposition 6 below), and σ2 → 1: convergence degenerates, a well-known difficulty of solving near-linear programmes with Sinkhorn. As ε → ∞, Pε⋆ → P SOT and σ2 is bounded away from 1 (for connected GS ), giving fast convergence. This mirrors the standard trade-off in entropic OT between approximation quality (small ε) and computational efficiency (large ε). Limiting behaviour as ε → 0 and ε → ∞ The regularisation strength ε interpolates between two natural extremes: the sliced plan P SOT (large ε) and a support-restricted exact OT solution (small ε). Proposition 6 (Limiting behaviour). Let ΠS = {P ∈ Π(a, b) : supp(P ) ⊆ S} denote the set of couplings supported on S.
1. Small-ε limit. As ε → 0+ , Pε⋆ converges to the unique coupling Pb that minimises KL P ∥ P SOT among all solutions of the support-restricted OT problem. 2. Large-ε limit. As ε → +∞, Pε⋆ → P SOT .

(16)

The proof is in Appendix A.5. Proposition 6 reveals that SinkSLOT interpolates continuously between the two interpretable endpoints. In practice, a moderate ε offers the best trade-off: a coupling close to true OT with rapid convergence. 7

Gradient and Hessian of the functional When the SLOTε divergence is used as a loss function for optimising the source positions X, we require its gradient ∇X SLOTε and, for second-order methods, Hessian-vector products T A = (∇2X SLOTε ) A. We derive both under the practical choice of treating the SOT reference plan as a fixed (detached) prior, and discuss when this is justified. Should we differentiate through P SOT ? The divergence SLOTε (µ, ν) depends on the source positions X through two channels: (i) the cost matrix Cij = ∥xi − yj ∥2 and (ii) the SOT reference plan P SOT (X, Y ), which itself depends on X through the 1D projections and rank-based matchings. By the envelope theorem (Danskin’s theorem), at any point X where the projected rank orders are locally constant (which holds for Lebesgue-almost every X), the gradient decomposes as ! X Pij⋆ ∂PijSOT d SLOTε X ⋆ 1 − SOT , (17) = Pkj ∇xk Ckj + ε dxk ∂xk Pij j (i,j)∈S {z } | | {z } (I) direct / Danskin

(II) reference-plan term

where P ⋆ is the SinkSLOT solution, dropping the explicit reference to ε for notational conciseness. In the classical EOT setting [Cuturi, 2013, Feydy et al., 2019], the reference coupling is the independent product a ⊗ b, which does not depend on X, so term (II) vanishes identically and the gradient reduces to term (I) alone. In SinkSLOT, however, P SOT depends on X through one-dimensional sorting operations, making term (II) nonzero in principle. We advocate treating P SOT as a fixed prior (i.e., stop-gradient / detach) and using only term (I), for the following reasons. Term (II) is small for moderate ε. By Proposition 6, P ⋆ → P SOT as ε → ∞, so Pij⋆ /PijSOT → 1 and term (II) vanishes. At moderate ε, the ratio stays close to 1, especially on the high-mass entries of P SOT that dominate the sum. Conversely, as ε → 0, P ⋆ approaches the support-restricted OT (Proposition 6), whose gradient depends on the support set S (a combinatorial object) rather than the values of P SOT . P SOT is locally constant. Each 1D OT plan PθSOT is determined by the rank order of the projected points ⊤ θ xi ; this rank order is a piecewise-constant function of X, introducing discontinuities in ∂P SOT /∂X. Although relaxations exist, they add complexity and lose the sparsity advantage. Note that at differentiable points (Lebesgue-almost everywhere), ∂P SOT /∂X = 0, so term (II) vanishes exactly, making stop-gradient not merely practical but formally correct. This choice is consistent with EOT practice. The standard approach in EOT-based learning [Feydy et al., 2019, Cuturi, 2013] is to apply Danskin’s theorem and differentiate only through C. Under the stop-gradient convention, the gradient and Hessian of SLOTε have exactly the same analytic form as those of EOT [Feydy et al., 2019, Ye et al., 2026], with the dense plan P EOT replaced everywhere by the sparse plan P ⋆ . All computational gains therefore come from the sparsity of P ⋆ : every transport-vector product costs O(|S|) instead of O(N M ). First-order gradient Proposition 7 (Gradient of SinkSLOT). Assume the squared Euclidean cost Ckj = ∥xk − yj ∥2 and treat P SOT as fixed. Then the gradient of SLOTε with respect to the source positions is X
⋆ ∇xk SLOTε = 2 Pkj (xk − yj ) = 2 ak xk − Tε (xk ) , (18) j∈Sk

8

where Tε (xk ) = a1k

P

⋆ j∈Sk Pkj yj is the barycentric projection and Sk = {j : (k, j) ∈ S}. In matrix form,

∇X SLOTε = 2 diag(a) X − Tε (X) , −1

Tε (X) = diag(a)

⋆

P Y.

(19) (20)

The proof is in Appendix A.6. Computational cost The sparse matrix-vector product P ⋆ Y costs O(|S| d) = O(L(N +M ) d), compared with O(N M d) for the dense EOT gradient. No sorting or projection is required at differentiation time. Second-order Hessian and Hessian-vector products Following the derivation of Ye et al. [2026, Appendix C] and Li et al. [2025], we derive the Hessian ∇2X SLOTε under stop-gradient in Appendix A.7. Since the analytic structure is identical to EOT, the Hessian-vector product (∇2X SLOTε ) A can be computed via sparse linear solves involving the sensitivity matrix H ⋆ ∈ R(N +M )×(N +M ) , which inherits the sparsity of P ⋆ . Each Hessian-vector product therefore costs O(|S|) per conjugate-gradient iteration, compared with O(N M ) for dense EOT.

Experiments To assess the effectiveness of the proposed method, we evaluate SinkSLOT both as a transport solver and as a loss function via its divergence. We benchmark against FlashSinkhorn [Ye et al., 2026], the most computationally efficient GPU implementation of dense Sinkhorn available; SROT [Nguyen, 2026], which uses a smoothed sliced reference coupling but operates on the full dense kernel; and a sparse Sinkhorn approach, Spar-Sink [Li et al., 2023], which sparsifies the Gibbs kernel by selecting entries via importance sampling. Experiments are run on an NVIDIA H100-80GB. Implementation details and full results are in Appendix B and C.

Synthetic Benchmarks Setup We consider five datasets with different geometric structures and dimensionalities, following previous experiments [Nguyen, 2026, Ye et al., 2026]: (i) half-moons (d=2), a pair of interleaved crescent-shaped point clouds; (ii) 8-Gaussians (d=2), a mixture of eight isotropic Gaussians; (iii) two-rings (d=2), where µ and ν are supported on concentric circles with different radii (iv) Gaussian d=3; and (v) Gaussian d=64. We sample N = M = 104 points with weights drawn uniformly from [0.1, 1.1], normalised to sum to one. All methods use identical early stopping based on marginal violation (details in Appendix B.1). Accuracy is measured by the relative transport cost gap (%), ∆ := 100 × (⟨C, P ⟩ − ⟨C, P OT ⟩) / ⟨C, P OT ⟩, where P OT is the exact OT plan from linear programming. SinkSLOT matches dense Sinkhorn accuracy at lower computational cost First, we assess whether sparsifying the Gibbs kernel with sliced lifted plans degrades plan quality. We sweep over regularisation strength ε (all methods), slice count L (SROT and SinkSLOT), and subsampling budget S0 (Spar-Sink). For each method, we record the minimum runtime required to reach a relative cost gap below three thresholds: ∆ ≤ 1%, 5%, and 10%. Table 2 reports SinkSLOT’s speedup over each competitor across all five datasets, along with SinkSLOT’s total runtime. Three conclusions emerge. First, Spar-Sink fails to reach any gap threshold on the Gaussian datasets (d=3; d=64) and cannot attain the ≤ 5% gap on 8-Gaussians. Where Spar-Sink reaches the accuracy threshold, SinkSLOT is 2–45× faster. Second, SROT is the slowest method overall: SinkSLOT achieves 9

Table 2: SinkSLOT speedup over competing methods at each transport cost gap threshold. “—” indicates that the method did not reach the threshold at convergence (or after 10k Sinkhorn iterations). N = 104 . d = 2 unless otherwise specified. H100-80GB. Half-moons

8-Gaussians

Gauss d=3

Two-rings

Gauss d=64

≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10%

Cost gap threshold:

SinkSLOT speedup (×): FlashSinkhorn [Ye et al., 2026] 13.1 8.4 8.4 17.5 19.4 16.7 36.4 14.7 8.2 SROT [Nguyen, 2026] 124.9 118.6 118.6 103.0 160.2 119.5 122.8 131.3 131.3 Spar-Sink [Li et al., 2023] 4.8 2.5 2.5 — — 37.4 45.5 5.5 3.3

4.3 29.5 —

Runtime (ms): SinkSLOT (ours)

1577.9 351.0 165.7 358.9 147.9 87.2

16.5 106

7.0

7.0

145.2 18.5

Gaussian, d=3, =0.01

11.0

7.1

Gaussian, d=64, =0.1

SROT OOM

SROT OOM

105

5.1

5.1

9.7 62.1 —

9.6 89.7 —

2.8 6.7 15.2 27.1 — —

11.4 35.1 —

Gaussian, d-scaling, N=M=10k, =0.1

Time (ms)

104 103 102 101 5k

SinkSLOT (ours), L=1024

10k

20k

N, M

30k

50k 5k

SinkSLOT, L=4096

10k

20k

N, M

30k

SROT (Nguyen 2026), L=1024

50k

4

8

16 32 64

SROT, L=4096

d

256

1024

FlashSinkhorn (Ye et al. 2026)

Figure 2: Forward-pass runtime against (N, M ) and d. Error bars indicate the standard error over five random seeds. 15–160× speedups across all datasets and gap levels. Third, and most importantly, SinkSLOT consistently outperforms FlashSinkhorn by 3 to 36×, with the largest gains on the low-dimensional benchmarks. Even on the challenging d=64 Gaussian setting, SinkSLOT is 3–11× faster. These results show that SinkSLOT can reach the same accuracy as dense methods at a fraction of the computation time. While SinkSLOT primarly targets a reduction of the computation cost, results related to peak memory usage are provided in Appendix C.3. Impact of the number of slices L The number of slices L determines the accuracy–speed trade-off. Results in Appendix C.1 show the Pareto front for SinkSLOT at various L alongside FlashSinkhorn. First, increasing L improves the achievable accuracy but increases computation time, consistent with the O(L(N + M )) per-iteration cost. Second, for a given L, decreasing ε improves accuracy only up to a plateau that reflects the gap between the support-restricted OT solution and the true OT plan; only enriching the sparse support by increasing L can lower this floor. Third, despite this trade-off, SinkSLOT consistently reaches the same accuracy as FlashSinkhorn at lower runtime for moderate gap levels. Very high precision (∆ ≤ 1%) requires large L, but such precision is rarely needed in practice due to the prohibitive computational cost in this regime. Scalability SinkSLOT’s O(L(N + M )) per-iteration cost is asymptotically lower than FlashSinkhorn’s O(N M ) for fixed L. Figure 2 confirms this: SinkSLOT consistently outperforms FlashSinkhorn across both d=3 and d=64, with speedups of 2–11× at L=4096 and 4–36× at L=1024. The advantage also persists across dimensions (right panel): at fixed N =104 , SinkSLOT remains faster for all d up to 1024, even at L=4096. 10

SOT (Bonneel et al. 2015)

step 5

step 10

step 20

step 30

step 50

W22 = 1.0122

W22 = 0.7010

W22 = 0.4482

W22 = 0.1159

W22 = 0.0053

W22 = 0.0009

W22 = 1.0122

W22 = 0.4101

W22 = 0.1569

W22 = 0.0190

W22 = 0.0027

W22 = 0.0004

W22 = 1.0122

W22 = 0.4267

W22 = 0.2209

W22 = 0.0528

W22 = 0.0068

W22 = 0.0004

W22 = 1.0122

W22 = 0.3540

W22 = 0.1238

W22 = 0.0154

W22 = 0.0022

W22 = 0.0005

SinkSLOT (ours)

SROT (Nguyen 2026)

EOT (Feydy et al. 2019)

step 0

Figure 3: Blob-to-crescent gradient flows. N = M = 1000. Regularisation strength ε = 0.01. L = 100 slices. Performance is evaluated using the exact squared 2-Wasserstein distance W22 ; bold indicates the lowest value at each intermediate step.

Gradient Flows We now assess whether the SinkSLOT divergence SLOTε (8) is effective as a loss function for optimisation. We run the blob-to-crescent gradient flow experiment of EOT [Feydy et al., 2019] and SROT [Nguyen, 2026] with N = M = 1000, and additionally report results for SOT [Bonneel et al., 2015] and SLOTε (8). At step t, source points are updated: X (t+1) = X (t) − η n ∇X D(X (t) , Y ) with η = 0.05. D denotes either the Sliced Wasserstein Distance [Bonneel et al., 2015] or the corresponding divergence for EOT, SROT, and SLOTε . All methods use the squared Euclidean cost Cij = ∥xi − yj ∥2 . For the divergences, we set ε = 0.01. L = 100 projection vectors are sampled whenever a sliced plan is involved (all except EOT), and the SROT smoothing parameter γ in (5) is set to 10−8 . We evaluate all methods using the exact squared 2-Wasserstein distance, W22 (1). Figure 3 shows the results after 50 Euler steps. While all four methods attain similar final W22 values, SLOTε transports the source points toward the target more rapidly during the earlier steps. Quantitatively, SLOTε achieves the lowest W22 and lingers substantially below SROT through step 30. Qualitatively, at step 5, SLOTε already rearranges the source cloud into the target crescent shape, whereas the initial blob shape remains clearly visible for the compared methods. The fact that SLOTε converges without debiasing, while EOT requires the Sinkhorn divergence correction [Feydy et al., 2019], is consistent with Proposition 3. Overall, the informative gradients and efficient forward computation make SLOTε divergence a promising loss function for large-scale optimisation.

Conclusion We introduced SinkSLOT, a sparse EOT method that replaces the dense independent reference coupling with an expected lifted sliced OT plan, reducing per-iteration cost from O(N M ) to O(L(N + M )). We proved convergence of the sparse Sinkhorn iterations, characterised limiting behaviour as ε varies, and derived first11

and second-order derivatives. Experiments show that SinkSLOT achieves accuracy comparable to dense and hardware-aware baselines while substantially reducing computation time. The number of slices L is a tunable parameter: too few slices restrict the transport polytope, while too many reduce the sparsity advantage. Promising extensions include ε-scaling [Kosowsky and Yuille, 1994], non-uniform slice weighting [Liu et al., 2025], warm starts [Peyré and Cuturi, 2019], and data-dependent projection directions [Nguyen et al., 2021]. Overall, experiments show that SinkSLOT delivers substantial speedups over state-of-the-art dense and sparse EOT methods as a transport solver, and that the resulting divergence is a promising loss function, as illustrated on a gradient flow task.

Acknowledgments This work was supported by ANR through the “Investissements d’avenir” program (ANR-10-IAIHU-06, ANR-19-P3IA-0001, PRAIRIE 3IA Institute), the “France 2030” program (ANR-23-IACL-0008, PRAIRIEPSAI). The ARAMIS Lab is affiliated with DIM C-BRAINS, funded by the Conseil Régional d’Ile-deFrance. This work was performed using HPC resources from GENCI-IDRIS (Grant 2026-AD011017987). RD received Marie Skłodowska-Curie grant No. 101154248 (SafeREG). SSK acknowledges support from the UKRI MRC Doctoral Training Partnership at King’s College London (MR/W006820/1). TV is co-founder and shareholder of Hypervision Surgical.

References Dimitri P Bertsekas. Nonlinear programming. Journal of the Operational Research Society, 48(3):334–334, 1997. Mathieu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In International conference on artificial intelligence and statistics, pages 880–889. PMLR, 2018. Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 51(1):22–45, 2015. Richard A Brualdi. Combinatorial matrix classes, volume 13. Cambridge University Press, 2006. Imre Csiszár. I-divergence geometry of probability distributions and minimization problems. The annals of probability, pages 146–158, 1975. Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013. Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memoryefficient exact attention with io-awareness. In Advances in Neural Information Processing Systems, 2022. Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pages 2681–2690. PMLR, 2019. Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Zahdi Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. Pot: Python optimal transport. Journal of machine learning research, 2021.

12

Aden Forrow, Jan-Christian Hütter, Mor Nitzan, Philippe Rigollet, Geoffrey Schiebinger, and Jonathan Weed. Statistical optimal transport via factored couplings. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2454–2465. PMLR, 2019. Paul Freulon, Nikitas Georgakis, and Victor Panaretos. Entropic optimal transport beyond product reference couplings: the gaussian case on euclidean space. arXiv preprint arXiv:2507.01709, 2025. Johannes Gasteiger, Marten Lienen, and Stephan Günnemann. Scalable optimal transport in high dimensions for graph distances, embedding alignment, and more. In International Conference on Machine Learning, pages 5616–5627. PMLR, 2021. Bahman Kalantari, Isabella Lari, Federica Ricca, and Bruno Simeone. On the complexity of general matrix scaling and entropy minimization via the RAS algorithm. Mathematical Programming, 112(2):371–401, 2008. Abdelwahed Khamis, Russell Tsuchida, Mohamed Tarek, Vivien Rolland, and Lars Petersson. Scalable optimal transport methods in machine learning: A contemporary survey. IEEE transactions on pattern analysis and machine intelligence, 2024. Philip A Knight. The sinkhorn–knopp algorithm: convergence and applications. SIAM Journal on Matrix Analysis and Applications, 30(1):261–275, 2008. Jeffrey J Kosowsky and Alan L Yuille. The invisible hand algorithm: Solving the assignment problem with statistical physics. Neural networks, 7(3):477–490, 1994. Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems-Series A, 34(4):1533–1574, 2014. Mengyu Li, Jun Yu, Tao Li, and Cheng Meng. Importance sparsification for sinkhorn algorithm. Journal of Machine Learning Research, 24(247):1–44, 2023. Xingjie Li, Fei Lu, Molei Tao, and Felix X-F Ye. Robust first-and second-order differentiation for regularized optimal transport. SIAM Journal on Scientific Computing, 47(3):C630–C654, 2025. Tianlin Liu, Joan Puigcerver, and Mathieu Blondel. Sparsity-constrained optimal transport. In International Conference on Learning Representations, 2023. Weijie Liu, Chao Zhang, Nenggan Zheng, and Hui Qian. Approximating optimal transport via low-rank and sparse factorization. arXiv preprint arXiv:2111.06546, 2021. Xinran Liu, Rocio Diaz Martin, Yikun Bai, Ashkan Shahbazi, Matthew Thorpe, Akram Aldroubi, and Soheil Kolouri. Expected sliced transport plans. In International Conference on Learning Representations, volume 2025, pages 60948–60971, 2025. Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867, 2018. Boris Muzellec and Marco Cuturi. Subspace detours: Building transport plans that are optimal on subspace projections. Advances in Neural Information Processing Systems, 32, 2019. Khai Nguyen. An introduction to sliced optimal transport: foundations, advances, extensions, and applications. Foundations and Trends® in Computer Graphics and Vision, 17(3-4):171–391, 2025. Khai Nguyen. Sliced-regularized optimal transport. arXiv preprint arXiv:2604.23944, 2026. 13

Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Distributional sliced-wasserstein and applications to generative modeling. In International Conference on Learning Representations, 2021. Sebastian Nowozin. Streaming log-sum-exp computation. streaming-log-sum-exp-computation.html, 2016. Blog post.

https://www.nowozin.net/sebastian/blog/

James M Ortega and Werner C Rheinboldt. Iterative solution of nonlinear equations in several variables. SIAM, 2000. Xiaxue Ouyang, Hao Zheng, Haoxian Liang, Jingyi Zhang, Yixuan Qiu, Cheng Meng, and Mengyu Li. Sparsification techniques for large-scale optimal transport problems. Wiley Interdisciplinary Reviews: Computational Statistics, 18(1):e70056, 2026. Gabriel Peyré. Optimal transport for machine learners. arXiv preprint arXiv:2505.06589, 2025. Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Now Foundations and Trends, 2019. Mark Rowland, Jiri Hron, Yunhao Tang, Krzysztof Choromanski, Tamas Sarlos, and Adrian Weller. Orthogonal estimation of wasserstein distances. In The 22nd International conference on artificial intelligence and statistics, pages 186–195. PMLR, 2019. Meyer Scetbon and Marco Cuturi. Low-rank optimal transport: Approximation, statistics and debiasing. Advances in Neural Information Processing Systems, 35:6802–6814, 2022. Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank sinkhorn factorization. In International Conference on Machine Learning, pages 9344–9354. PMLR, 2021. Bernhard Schmitzer. Stabilized sparse scaling algorithms for entropy regularized transport problems. SIAM Journal on Scientific Computing, 41(3):A1443–A1481, 2019. Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343–348, 1967. George W Soules. The rate of convergence of sinkhorn balancing. Linear algebra and its applications, 150: 3–40, 1991. Eloi Tanguy, Laetitia Chapel, and Julie Delon. arXiv:2508.01243, 2025.

Sliced optimal transport plans.

arXiv preprint

Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009. Felix X-F Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, and Davis Wertheimer. Flashsinkhorn: Io-aware entropic optimal transport on gpu. In Forty-third International Conference on Machine Learning, 2026.

14

A.

Proofs

A.1

Proof of Proposition 1 (unique diagonal scaling)

Proof. Since P SOT ∈ Π(a, b) is supported on S, the feasible set Π(a, b) ∩ {P : supp(P ) ⊆ S} is nonempty. The KL divergence is strictly convex in its first argument over this set, so the minimiser Pε⋆ exists and is unique. Writing the Lagrangian with multipliers fi (row marginals) and gj (column marginals), the first-order P condition for (i, j) ∈ S gives log (Kεij)ij − fi − gj = 0, hence Pij = efi (Kε )ij egj . Defining ui = efi and vj = egj yields the claimed form; for (i, j) ∈ / S, Pij⋆ = 0 by the support restriction. We additionally verify scalability in the sense of Kalantari et al. [2008]: a nonnegative matrix K is (a, b)-scalable if and only if there exists a matrix with the same zero pattern satisfying both marginal constraints. Since P SOT shares the zero pattern of Kε and lies in Π(a, b), the condition is met.

A.2

Proof of Proposition 2 (convergence of sparse Sinkhorn)

×M Proof. We apply Csiszár [1975, Theorem 3.2]. Matrices P ∈ RN supported on S are identified with + probability distributions on the finite set S, and Kε plays the role of the reference measure R in Csiszár [1975]. Define two linear constraint families

Ec = {P ≥ 0 : P ⊤ 1 = b},

Er = {P ≥ 0 : P 1 = a},

so that Π(a, b) = Er ∩ Ec . Each is a linear family in the sense of Csiszár [1975, Section 3, case (A)], since the constraints are finite linear expectations. The domain condition of Csiszár [1975, Theorem 3.2] requires a distribution in Er ∩ Ec supported on S; this is satisfied by P SOT . We identify each half-iteration of (10) with the I-projection (KL-projection) of the current iterate onto Er or Ec . Starting from P 0 = Kε = diag(1) Kε diag(1), the I-projection of P ℓ onto Er preserves the ℓ diagonal-scaling structure: the first-order optimality conditions give Peijℓ+1 = uℓ+1 i (Kε )ij vj , and imposing P eℓ+1 = ai yields uℓ+1 = a ⊘ (Kε vℓ ). Similarly, projecting onto Ec gives vℓ+1 = b ⊘ (Kε⊤ uℓ+1 ). j Pij By Csiszár [1975, Theorem 3.2], the alternating I-projections converge to the I-projection of Kε onto Er ∩ Ec = Π(a, b), which is Pε⋆ .

A.3

Proof of Proposition 3 (SLOTε is a divergence)

Proof. Non-negativity Since C ≥ 0 and Pε⋆ ≥ 0, we have ⟨C, Pε⋆ ⟩ ≥ 0. By Gibbs’ inequality, KL(Pε⋆ ∥P SOT ) ≥ 0. Hence SLOTε (µ, ν) ≥ 0. Identity of indiscernibles (⇐) When µ = ν, projecting identical point clouds onto any direction θ yields identical sorted positions, so the rank-based 1D matching maps each point to itself. Thus P SOT(µ,µ) = diag(a). Since Cii =
∥xi − xi ∥p = 0, the candidate plan P = diag(a) achieves ⟨C, diag(a)⟩ = 0 and KL diag(a)∥ diag(a) = 0, giving SLOTε (µ, µ) = 0. (⇒) If SLOTε (µ, ν) = 0, then both ⟨C, Pε⋆ ⟩ = 0 and KL(Pε⋆ ∥P SOT ) = 0. The latter implies ⋆ Pε = P SOT . The former requires Cij = 0, hence xi = yj , for all (i, j) in the support of P SOT . Since P SOT ∈ Π(a, b), every source point xi sends mass only to target points yj = xi , and every target point yj

15

receives mass only from source points xi = yj . For any location z ∈ Rd , X X X µ({z}) = ai = PijSOT =

i:xi =z

i:xi =z

X X

PijSOT =

j

i:xi =z j:yj =z

X

bj = ν({z}),

j:yj =z

so µ = ν. Symmetry For any direction θ, the 1D rank-based optimal matching is symmetric: swapping source and
⊤ target reverses the matching bijection. Hence P SOT(ν,µ) = P SOT(µ,ν) . Since Cij = ∥xi − yj ∥p = ∥yj − xi ∥p , the cost matrix for the reversed problem is C ⊤ . The Gibbs kernel for the reversed problem is therefore Kε⊤ , and the optimal plan is (Pε⋆ )⊤ . Because KL is invariant under simultaneous transposition of both arguments:
SLOTε (ν, µ) = ⟨C ⊤ , (Pε⋆ )⊤ ⟩ + ε KL (Pε⋆ )⊤ ∥(P SOT )⊤ = SLOTε (µ, ν).

A.4

Proof of Proposition 5 (linear convergence rate)

Proof. The argument combines the classical Jacobian analysis of Sinkhorn iterates [Soules, 1991, Knight, 2008] with a connectivity characterisation specific to the sparse support S. Step 1: Jacobian and its spectrum Write one full Sinkhorn cycle as a map Ψ : RM → RM on the log-scaling vectors P ψ = log v (cf. equations (11)-(12)). A direct calculation using the softmax identity x k ∂xk LSE(x) = e / l exl shows that the Jacobian of Ψ at the fixed point ψ ∗ satisfies J

jj ′

X

=

i: (i,j),(i,j ′ )∈S

Pij⋆ Pij⋆ ′ ai bj

,

(21)

i.e. J = Db−1 (Pε⋆ )⊤ Da−1 Pε⋆ ; see Knight [2008, Theorem 3.1] for the general (possibly sparse) case. The 1/2 −1/2 similarity transformation Je = Db J Db = A⊤ A shows that the eigenvalues of J are the squared singular values σ12 ≥ σ22 ≥ · · · of the normalised plan matrix A defined in (13). Step 2: Leading singular value

A direct check confirms σ1 = 1: √ √ −1/2 ⋆ −1/2 Pε 1 = Da a = a, A b = Da √ √ −1/2 A⊤ a = Db (Pε⋆ )⊤ 1 = b. −1/2 √

The corresponding eigenvector of J is Db

b = 1, which spans the gauge direction ψ 7→ ψ + c 1.

Step 3: Connectivity implies σ2 < 1 The entry (A⊤ A)jj ′ is positive if and only if columns j and j ′ share a row neighbour in GS . This adjacency defines the column intersection graph H. When GS is connected, so is H: any path cj −ri1 −cj1 −ri2 −· · ·−cj ′ in GS maps to a walk j−j1 −· · ·−j ′ in H, since each consecutive column pair shares a row neighbour. Hence A⊤A is irreducible, and the Perron-Frobenius theorem implies that its leading eigenvalue 1 is simple, giving σ22 < 1. Conversely, if GS has q ≥ 2 connected components, A decomposes into q independent blocks each contributing a unit singular value, so σ2 = 1. 16

Step 4: Linear convergence of the coupling The log-domain map Ψ satisfies the gauge equivariance Ψ(ψ + c1) = Ψ(ψ) + c1, so the eigenvalue 1 of J corresponds to the gauge direction and does not affect ℓ ℓ the coupling P ℓ = diag(eφ ) Kε diag(eψ ). On the quotient space RM /span(1) the induced map has a unique fixed point whose Jacobian has spectral radius σ22 < 1. By the standard convergence theorem for differentiable fixed-point iterations [Ortega and Rheinboldt, 2000, Theorem 4.1], 1/ℓ

lim sup P ℓ − Pε⋆ 1 ℓ→∞

≤ σ22 .

The iteration count (15) and the per-iteration cost O(|S|) = O(L(N +M )) follow immediately.

A.5

Proof of Proposition 6 (limiting behaviour)

Proof. Part 1: ε → 0+ We first note that KL(P ∥P SOT ) is finite and continuous on ΠS : for any (i, j) ∈ S we have PijSOT > 0, and the map x 7→ x log(x/q) is continuous on [0, ∞) for q > 0 (with the convention 0 log 0 = 0). Since ΠS is compact (a closed bounded subset of R|S| ), the KL divergence is bounded: supP ∈ΠS KL(P ∥P SOT ) =: Kmax < ∞. Let εk → 0+ and write Pk = Pε⋆k . By compactness, (Pk ) has a convergent subsequence (still denoted Pk ) with Pk → P̄ ∈ ΠS . P̄ is a cost minimiser. The optimality of Pk yields, for any P ∈ ΠS : ⟨C, Pk ⟩ + εk KL(Pk ∥P SOT ) ≤ ⟨C, P ⟩ + εk KL(P ∥P SOT ).

(22)

Since KL(Pk ∥P SOT ) ≥ 0: ⟨C, Pk ⟩ ≤ ⟨C, P ⟩ + εk Kmax . Taking k → ∞: ⟨C, P̄ ⟩ ≤ ⟨C, P ⟩ for all P ∈ ΠS . Hence P̄ ∈ P ∗ := arg minP ∈ΠS ⟨C, P ⟩. P̄ has minimum KL among cost minimisers. Let OTS = ⟨C, P̄ ⟩ and let Pb = arg min{KL(P ∥P SOT ) : P ∈ P ∗ } (unique by strict convexity of KL and convexity of P ∗ ). Setting P = Pb in (22) and using ⟨C, Pk ⟩ ≥ OTS : ⟨C, Pb⟩ − ⟨C, Pk ⟩ KL(Pk ∥P SOT ) ≤ KL(Pb∥P SOT ) + εk SOT b ≤ KL(P ∥P ),

(23) (24)

where the last inequality uses ⟨C, Pk ⟩ ≥ ⟨C, Pb⟩ = OTS . By continuity of KL on ΠS : KL(P̄ ∥P SOT ) ≤ lim inf KL(Pk ∥P SOT ) ≤ KL(Pb∥P SOT ). k

Since P̄ ∈ P ∗ and Pb minimises KL over P ∗ , we also have KL(P̄ ∥P SOT ) ≥ KL(Pb∥P SOT ). Hence equality holds, and by the uniqueness of the KL minimiser over the convex set P ∗ , P̄ = Pb. Since every convergent subsequence has the same limit, the entire family converges: Pε⋆ → Pb as ε → 0+ . Part 2: ε → +∞ Dividing the optimality condition (22) by ε and setting P = P SOT (for which KL(P SOT ∥P SOT ) = 0): ⟨C, P SOT ⟩ ⟨C, Pε⋆ ⟩ + KL(Pε⋆ ∥P SOT ) ≤ . ε ε 17

Since ⟨C, Pε⋆ ⟩ ≥ 0:

⟨C, P SOT ⟩ KL(Pε⋆ ∥P SOT ) ≤ . ε p By Pinsker’s inequality, ∥P − Q∥1 ≤ 2 KL(P ∥Q), so r 2 ⟨C, P SOT ⟩ ε→∞ ⋆ SOT −−−−→ 0. ∥Pε − P ∥1 ≤ ε

A.6

Proof of Proposition 7 (gradient of SLOTε )

Proof. By Danskin’s theorem [Bertsekas, 1997], since P ⋆ is the unique minimiser of a strictly convex programme in P (Proposition 1), X d SLOTε X ⋆ ∂Ckj ⋆ = Pkj = Pkj · 2(xk,t − yj,t ), dxk,t ∂xk,t j

j∈Sk

⋆ = 0 for (k, j) ∈ / S. Using where the sum is restricted to Sk because Pkj form (18).

A.7

⋆ j Pkj = ak yields the barycentric

P

Proof of Proposition 8 (Hessian of SLOTε )

Definition 2 (Sensitivity matrix). Define the sensitivity matrix   diag(a) P⋆ ⋆ H = ∈ R(N +M )×(N +M ) . (P ⋆ )⊤ diag(b)

(25)

Under assumption that the bipartite support graph GS is connected, H ⋆ is symmetric positive semidefinite with a simple zero eigenvalue (null vector (1N , −1M )), and its Moore-Penrose pseudoinverse H ⋆† is welldefined. Remark 1. In the dense EOT case, H ⋆ is a dense (N +M ) × (N +M ) matrix because P ⋆ > 0 everywhere. In SinkSLOT, H ⋆ inherits the sparsity of P ⋆ : its off-diagonal blocks have |S| nonzeros, so the matrix-vector product H ⋆ w costs O(|S| + N + M ) rather than O(N M ). ⋆ , and R ∈ R(N +M )×N ×d Definition 3 (Auxiliary tensors). Define B ∈ RN ×d×M by Bktj = 2(xk,t −yj,t ) Pkj slice-wise (for each t ∈ [d]) as
! diag B:,t,: 1M R:,:,t = . (B:,t,: )⊤

Proposition 8 (Hessian of SinkSLOT). Under the stop-gradient convention and squared Euclidean cost, the Hessian T = ∇2X SLOTε ∈ RN ×d×N ×d decomposes as N +M

1 X Tktsl = Rikt Hij⋆† Rjsl + Ektsl , ε

(26)

i,j=1

where E is block-diagonal across source points: Ek,:,s,: = 0 for k ̸= s, and Ek,:,k,: = 2ak Id −

4 X ⋆ Pkj (xk − yj )(xk − yj )⊤ , ε

(27)

j∈Sk

where each point xk , yj ∈ Rd×1 is treated as a column vector, so (xk − yj )(xk − yj )⊤ ∈ Rd×d is the outer product. 18

Half Moons

8 Gaussians

Two Rings

source target

Figure 4: Visualisation of SinkSLOT transport plan for the three 2D datasets. Only 120 of the 104 points are shown for readability. Regularisation strength ε = 0.01. L = 1024 slices. Proof. Under the stop-gradient convention, ∂Pij⋆ /∂xk,t = (Pij⋆ /ε)(∂fi /∂xk,t + ∂gj /∂xk,t − ∂Cij /∂xk,t ), since the only x-dependence in log Pij⋆ = log PijSOT + (fi + gj − Cij )/ε comes from Cij and the dual potentials (f, g) (with P SOT treated as constant). This is the same structure as the dense EOT case (where log Pij⋆ = log ai +log bj +(fi +gj −Cij )/ε and log ai , log bj are constant), so the remainder of the derivation (differentiating the marginal constraints, solving the linear system via H ⋆† , and collecting terms) is identical to Ye et al. [2026, Theorem 7 / Appendix C].

B.

Implementation details

B.1

Synthetic Benchmark Experiments

Shared protocol All methods solve the same problem instances and are held to the same stopping rule, so runtimes are directly comparable. Except in the scalability experiment in Figure 2, each dataset contains N = M = 104 points. The weights are drawn uniformly from [0.1, 1.1] and normalised to sum to one. The three 2D datasets are visualised in Figure 4. All methods are run in float32 on the same NVIDIA H100-80GB GPU. Iterations terminate when the maximum marginal violation satisfies   max max (P 1)i − ai , max (P ⊤ 1)j − bj < 10−6 , i

j

or when the common limit of 104 Sinkhorn iterations is reached. This limit is increased to 2 × 104 for the Gaussian dataset with d = 64 and for the scalability experiment. P Every output plan is subject to a feasibility check, including two criteria: (i) mass conservation, | ij Pij − 1| < 10−3 and (ii) a marginal violation below 10−4 . All 840 output plans produced in the benchmark satisfy both criteria. Hyperparameters We sweep eight values of the entropic regularisation parameter ε for every method, eight slice counts L for SROT and SinkSLOT, and four sampling budgets s for Spar-Sink. Since the effective entropic smoothing is governed by the ratio Cij /ε, ε should be selected relative to the scale of the ground cost [Peyré and Cuturi, 2019]. We therefore sample eight log-spaced values in [0.001, 0.1] for the low-dimensional datasets and in [0.1, 1] for the Gaussian dataset with d = 64. 19

For the slice-based methods, we use L ∈ {32, 64, 128, . . . , 4096} on the low-dimensional datasets and 2L on Gaussian d = 64. For Spar-Sink, we follow the experimental set of Li et al. [2023] and use s ∈ {5, 10, 15, 20} × s0 (N ), where s0 (N ) = 10−3 N log4 (N ). For N = M = 104 , these values correspond to expected sampling densities s/N 2 of approximately 0.36%, 0.72%, 1.08%, 1.44%. For each method and cost-gap threshold, Table 2 reports the minimum runtime achieved across the hyperparameter combinations considered. The reference plan P OT is the exact solution of the unregularised linear programme. Baseline implementation SinkSLOT’s implementation is described in Appendix B.2. For the baselines, we start from the authors’ released code and adapt it to our benchmarking stack. The released implementations target a different software stack and report different stopping and timing conventions, so running them as published would confound the method with its implementation. Adapting both holds the GPU, precision, stopping rule and timing instrumentation fixed across all methods, leaving the algorithm as the only difference. The changes are confined to that harness: the update equations, the reference coupling and the sampling scheme are the authors’ own, and we link the original code in each case. FlashSinkhorn We use the authors’ official FlashSinkhorn implementation2 . Its solver is left intact apart from the instrumentation required to apply the shared marginal-violation stopping criterion instead of its native potential-change criterion. We use the alternating backend, rather than the symmetric backend, which updates both dual potentials simultaneously using half-steps. We set allow_tf32=False to ensure that computations are performed in float32, consistent with all other methods. SROT Adapted from the author’s code [Nguyen, 2026]3 . The reference plan is the smoothed sliced plan PγSOT = (1 − γ)P SOT + γ a ⊗ b of (5). Consistent with the author’s choice, we set γ = 10−8 . Spar-Sink Adapted from the authors’ code [Li et al., 2023]4 . Following their equations p (7) and (9), entries are kept independently with inclusion probability qij = min(1, s pij ) where pij ∝ ai bj , and survivors are rescaled by 1/qij so that the sampled kernel is an unbiased estimator for K. The budget s bounds the expected number of nonzero entries, so the realised count varies across runs. The values of s included in the sweep are specified above.

B.2

Performance Improvements from Fusion Strategies

Table 1 shows that the number of arithmetic operations per Sinkhorn iteration scales as O(|S|). In practice, however, wall-clock time scales as |S| times a constant factor that depends on implementation details, and a naive implementation can inflate that constant well beyond what the asymptotic bound suggests. For this reason, we wrote our kernels in Triton, taking inspiration from FlashSinkhorn’s approach. They read the support S through a compressed sparse layout we build ourselves, rather than through a general-purpose sparse library. The reason is that the row and column half-steps are segmented log-sum-exp reductions, not matrix-vector products. No sparse Basic Linear Algebra Subprograms (BLAS) library exposes that primitive. Two fusions carry the implementation, one in each stage of the pipeline. 2

https://github.com/ot-triton-lab/flash-sinkhorn https://github.com/khainb/SROT 4 https://github.com/Mengyu8042/Spar-Sink 3

20

Fusing the sparse cost evaluation Stage 1 needs Cij only on S. Written directly, (X[rows]−Y [cols])2 .sum(1) allocates three (|S|, d) intermediates before reducing them. We fuse the gather, difference, square and reduction into one Triton pass that accumulates over d in registers and writes only the |S| costs, evaluating 26M entries in 3.1 ms and matching the dense computation to a relative error of 2.8 × 10−7 . Fusing the segmented log-sum-exp Each half-step of (11)-(12) is a segmented LSE over S. We store the row half-step in CSR and the column half-step in CSC so both reductions are contiguous, and fuse each into a kernel that sweeps a row in tiles of width B, carrying a running maximum m and accumulator s. This is the streaming log-sum-exp recurrence of Nowozin [2016], in the generalised, parallel-reduction form of Milakov and Gimelshein [2018]:
P ′ ′ m′ = max m, maxt xt , s′ = s em−m + t ext −m , where xt enumerates log(Kε )ij + ψj over the tile’s entries j. This never materialises those values and reads each entry once per half-step rather than twice [Ye et al., 2026]. This is the same online-softmax trick FlashAttention [Dao et al., 2022] uses to fuse attention without materialising the score matrix, and that FlashSinkhorn [Ye et al., 2026] uses for its own dense Sinkhorn kernels; we apply it to the segmented, per-row case the sparse support requires. Folding log ai − LSE into the same kernel halves the launches per iteration, from four to two. One program owns one row and the tile width is set to B = 2⌈log2 (|S|/N )⌉ , the mean row length rounded up to a power of two and clamped to [32, 1024].

B.3

Gradient Flows

We follow the convention of Nguyen [2026], Feydy et al. [2019] as closely as possible. N = M = 1000 uniformly weighted source and target points are sampled from two density images (blob → crescent), using density-weighted pixel sampling with sub-pixel jitter, ±(0.5/h) × noise, where h is the image height in pixels, to avoid samples lying exactly on the discrete pixel grid. We then perform 50 Euler steps, as described in the main text. The four methods are evaluated using the exact squared 2-Wasserstein distance W22 : • SOT: We compute the exact analytical gradient using automatic differentiation through ot.sliced_ wasserstein_distance from the POT library [Flamary et al., 2021], using L = 100 slices. Since 1D OT has a closed-form solution based on sorting, no Sinkhorn iterations or entropic approximation are involved, and the results are therefore independent of ε. • EOT: Rather than applying the marginal-violation stopping criterion used in the benchmark experiments, we run dense Sinkhorn for a fixed 1000 iterations and use full backpropagation to compute the gradients. We set ε = 0.01. • SROT: In addition to the EOT setup, we use L = 100 slices and set the smoothing parameter γ in (5) to 10−8 , as proposed in Nguyen [2026]. As in our approach, the authors’ implementation treats the prior coupling as a stop-gradient constant. • SinkSLOT (ours): No backpropagation is required; analytical gradients are computed using Equation (18). We run 1000 fixed iterations with ε = 0.01 and L = 100. In Appendix C.4, we validate the substitution of our analytical gradient formula (18) for full backpropagation through P SOT , in which the prior coupling is treated as differentiable with respect to the source points.

21

cost gap from OT cost (%)

Half-moons, d=2, N=M=10k

8-Gaussians, d=2, N=M=10k

1

10

0.1

1

0.01 101

102

101

103

time (ms)

cost gap from OT cost (%)

Two-rings, d=2, N=M=10k

102

103

time (ms)

Gaussian, d=3, N=M=10k 100

1

10 1

0.1 101

102

time (ms)

FlashSinkhorn (Ye et al. 2026) =0.100

=0.052

0.1

103

101

102

time (ms)

method (color) SinkSLOT (ours), L=128 SinkSLOT (ours), L=1024 (shape)

=0.027

=0.014

=0.007

=0.004

104

103

SinkSLOT (ours), L=4096 =0.002

=0.001

Figure 5: Pareto fronts for the accuracy-speed trade-off on the four low-dimensional datasets. FlashSinkhorn vs SinkSLOT at three slice counts: L = 128, 1024, and 4096. N = M = 104 .

C.

Additional results

C.1

Impact of L

To understand the accuracy-speed trade-off one faces when varying regularisation strength ε and slice count L, Figure 5 plots the cost gap from exact OT cost against runtime for FlashSinkhorn and SinkSLOT on each of the four low-dimensional datasets. The plot uses the eight values of ε swept in the benchmark experiment and three slice counts: 128 (low), 1024 (medium), 4096 (high). Since the results needed for this plot are readily available from the benchmark experiment, no additional runs were performed. For the Gaussian d = 64 dataset, we run additional experiments using larger ε to show the trade-off at higher cost gaps in Figure 6. Figure 5 shows that, in low-dimensional settings, SinkSLOT achieves a better accuracy-speed trade-off than FlashSinkhorn. In other words, given a target cost gap, we can find a combination of hyperparameters (L, ε) for which SinkSLOT is faster. Similarly, given a runtime budget, we can find (L, ε) for which SinkSLOT achieves a lower cost gap. As with many optimisation methods, the question is then how to select (L, ε) to achieve the optimal trade-off. There is no simple answer, and the appropriate choice depends strongly on the dataset. For tworings, a low slice count of L = 128 appears sufficient, as its curve lies entirely southwest of the other three 22

d=64, N=M=10k

cost gap from OT cost (%)

method (color) FlashSinkhorn (Ye et al. 2026) SinkSLOT (ours), L=1024 SinkSLOT (ours), L=2048 SinkSLOT (ours), L=4096 SinkSLOT (ours), L=8192

10

1 (shape)

=20.000 =11.808 =6.971 =4.116 =2.430 =1.435 =0.847 =0.500

102

time (ms)

103

Figure 6: Pareto fronts for the accuracy–speed trade-off on the Gaussian d = 64 dataset. FlashSinkhorn vs SinkSLOT at four slice counts. Table 3: Iterations run by the SAME configuration selected in the speedup table 2 (the fastest wall-clock run reaching each transport cost-gap threshold). “—” means no feasible run reached the threshold. N = 104 . d = 2 unless otherwise specified. H100-80GB. Half-moons Method

8-Gaussians

Two-rings

Gauss d=3

Gauss d=64

≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10%

Iterations of the fastest-time (speedup-table) run: SinkSLOT (ours) 420 140 140 3180 460 FlashSink [Ye et al., 2026] 260 70 70 3060 430 SROT [Nguyen, 2026] 450 140 140 3210 470 Spar-Sink [Li et al., 2023] 260 70 70 — —

250 220 260 1010

150 310 150 654

90 90 100 96

90 50 100 50

10000 3240 1430 430 520 8220 4090 1920 430 430 7320 3290 1550 1230 870 — — — — —

990 430 700 —

curves. Decreasing ε reduces the cost gap while increasing runtime. For the other three datasets, the overall Pareto front for SinkSLOT comprises points obtained using multiple slice counts. The Gaussian dataset with d = 3 illustrates this behaviour particularly well. At L = 128, runtime is low, and decreasing ε initially reduces the cost gap. Beyond a certain point, however, further reduction in ε increases runtime without improving the quality of the output plan. In contrast to two-rings, increasing the slice count to L = 1024 is more effective. At large ε, the cost gap is comparable to L = 128, but decreasing ε has a greater effect on the cost gap at this slice count. The same reasoning applies when increasing the slice count from L = 1024 to L = 4096 when the objective is to achieve a cost gap below 1%. The same conclusion holds in the high-dimensional setting shown in Figure 6, even though the advantage over FlashSinkhorn is less pronounced. For a given cost-gap-runtime trade-off achieved by FlashSinkhorn, meticulously selected combinations (L, ε) can achieve both a lower cost gap and a shorter runtime, but this set is likely to be small. In practice, the gains in accuracy and speed over FlashSinkhorn may not compensate for the time investment on experimentation to identify an appropriate hyperparameter range for a given distribution. This limitation is not unique to SinkSLOT. Rather, it is a broader challenge for slice-based OT methods, as the approximation ability of lifted plans deteriorates when dimension increases [Tanguy et al., 2025].

23

Table 4: Lowest peak GPU memory (MB) among feasible runs reaching each transport gap threshold. “—” means no feasible run reached the threshold. N = 104 . d = 2 unless otherwise specified. H100-80GB. Half-moons

8-Gaussians

Gauss d=3

Two-rings

Gauss d=64

≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10% ≤1% ≤5% ≤10%

Method

Peak memory (MB): SinkSLOT (ours) 820 FlashSink [Ye et al., 2026] 1785 SROT [Nguyen, 2026] 2359 SparSink [Li et al., 2023] 1258

820 820 896 820 820 820 820 820 2580 1799 1281 5050 3007 1302 1785 1785 1785 1785 1785 1785 1785 1785 1785 1785 1785 1806 1806 1806 2359 2359 2359 2359 2359 2359 2359 2359 2359 2359 2359 2380 2380 2380 1258 1258 — — 1258 1258 1258 1258 — — — — — —

Gaussian, d=3, =0.01

Gaussian, d-scaling, N=M=10k, =0.1

Gaussian, d=64, =0.1

Peak memory (MB)

104

103 5k

10k

SinkSLOT (ours), L=1024

N, M

20k

30k

SinkSLOT, L=4096

50k 5k

10k

N, M

20k

30k

SROT (Nguyen 2026), L=1024

50k 4

8

16 32 64

SROT, L=4096

d

256

1024

FlashSinkhorn (Ye et al. 2026)

Figure 7: Forward-pass peak memory against (N, M ) and d. Error bars indicate the standard error over five random seeds.

C.2

Convergence speed

Table 3 reports the number of iterations required for convergence for each fastest run reported in Table 2. No additional runs were performed to produce this table. Here, one iteration consists of sequentially updating the two dual potentials, corresponding to Steps 11 and 12 of Algorithm 1. In all selected runs except one (SinkSLOT on Gaussian d = 3), Sinkhorn converged within 104 iterations. SinkSLOT required a number of iterations comparable to that of its dense counterpart, SROT.

C.3

Memory cost

Table 2 and Figure 2 report the shortest runtime achievable at each predefined cost gap threshold. In practice, however, computational efficiency is constrained not only by runtime but also by hardware memory requirements. We therefore investigate the lowest peak memory usage achievable by each method at each cost gap threshold. Table 4 reports our findings for fixed N = M = 104 on the same datasets, at the three cost gap thresholds of 1%, 5% and 10%. Figure 7 reports the peak memory usage as N, M increase from 5 × 103 to 5 × 104 for Gaussian d = 3, 64, as well as while increasing d up to 1024. Although the fused Triton kernels described in Appendix B.2 keep the Sinkhorn iteration itself memoryefficient, the memory reported here often comes from storing the sparse support S of the prior coupling P SOT , whose size grows with the slice count L. For example, in Table 4, on Gaussian d=64 at the 1% threshold, SinkSLOT needs L=8192 slices to bring the cost gap under 1%. At N =M =104 , this many slices covers 55.7% of all N × M pairs. Since the sparse support must store row and column indices in addition to the corresponding values, its memory footprint can exceed that of a dense N × M matrix. This explains why SinkSLOT’s peak memory of 5050 MB exceeds SROT’s flat, fully dense 2380 MB in this setting. Nevertheless, in the high-sample-size setting (N =5 × 104 ) shown in Figure 7, SinkSLOT’s memory usage remains manageable even with 4096 slices, whereas SROT runs out of memory at all slice counts because 24

its dense matrix grows quadratically with N .

C.4

Stop-gradient experiments

We verify the two claims underlying Equation (17) numerically: that term (I) is the gradient of SLOTε under the stop-gradient convention, and that term (II) is negligible. We also quantify a third, distinct quantity that is easily conflated with term (II): the sensitivity of the Sinkhorn solve itself. Protocol All measurements are taken along the blob-to-crescent flow of Appendix B.3, with N = M = 1000, L = 100 projections, ε as stated, 600 inner Sinkhorn iterations and 50 gradient steps at η = 0.05. The trajectory is driven by term (I) throughout, since that is the method under study. At each step the reference plan P SOT is rebuilt from the current X, exactly as the method does, and is then shared by every estimator evaluated at that step. The projection directions are drawn from a fixed seed, so across a perturbation only the rank orders they induce can change, never the directions themselves. Term (I) is computed in two independent ways: in closed form as 2 diag(a)(X − Tε (X)) with no automatic differentiation, and by automatic differentiation through the cost with the prior coupling held constant. The two agree to a relative 10−4 at every step, the discrepancy tracking the marginal violation maxi |ai /ri − 1| of the truncated solve. Reported values are accumulated in float32 on GPU, whose atomic reductions are not order-deterministic, so all numbers below should be read to two or three significant figures; the finite-difference study is carried out in float64. Term (I) is exact, and term (II) is a jump rather than a derivative We compare the analytic directional derivative along random unit directions v against [SLOTε (X + hv) − SLOTε (X − hv)]/2h, with P SOT either frozen at the base point or rebuilt at each perturbed point so that rank flips are free to occur. With the reference plan frozen, the finite difference matches term (I) to a relative 1.1 × 10−9 at the first step (Table 5); the growth of the reported error along the flow reflects the derivative itself decaying to ∼ 7 × 10−6 against a fixed finite-difference noise floor, not a degradation of the identity. With the reference plan rebuilt, the divided difference does not converge as h → 0: at ε = 0.01 and the first step it takes the values 4.7 × 10−3 , 1.5 × 10−2 , 5.8 × 10−2 and 2.4 × 10−1 for h = 10−3 , . . . , 10−6 , against an analytic value of 1.76 × 10−3 . This O(h−1 ) growth is the signature of a jump discontinuity, not of a missing gradient term, and it is consistent with P SOT being piecewise constant: at h = 10−7 , where no rank flip occurs between the two evaluations, the rebuilt difference returns to the analytic value. Term (II) is therefore exactly zero wherever it is defined, and at the rank-flip boundaries SLOTε is discontinuous rather than differentiable. Those discontinuities are small. Rebuilding P SOT moves roughly 0.3% of the support entries and changes the value by a relative 1.6 × 10−6 at the first step, rising to at most 3.2 × 10−4 at the last, where the value itself has fallen by two orders of magnitude; the jump shrinks with h (4.8 × 10−4 at h = 10−3 to 3.3 × 10−5 at h = 10−6 , at step 50) and is systematically negative, since the rebuilt plan is optimal at the perturbed point. Stop-gradient on P SOT is thus justified in practice, the justification being piecewise-constancy plus small jumps rather than differentiability with a vanishing derivative. The solver term is negligible until convergence A separate question is what is lost by treating the solve as exact. Differentiating ⟨C, P ⟩ through all inner Sinkhorn iterations yields a gradient that differs from term (I) by a residual. This residual is not term (II), since the reference plan is frozen in both arms, but the sensitivity of the optimal plan, which the envelope identity for ⟨C, P ⟩ + ε KL(P ∥P SOT ) absorbs: unrolling that regularised objective instead recovers term (I) to a cosine similarity of 0.999992.

25

step

analytic

FD (frozen)

rel. err.

|∆F |/F

0 10 25 40 50

4.431 × 10−4 −2.906 × 10−4 −6.904 × 10−5 −7.862 × 10−6 −6.683 × 10−6

4.431 × 10−4 −2.906 × 10−4 −6.903 × 10−5 −7.853 × 10−6 −6.696 × 10−6

1.1 × 10−9 8.7 × 10−5 6.6 × 10−4 2.3 × 10−3 3.8 × 10−3

1.6 × 10−6 1.3 × 10−5 8.6 × 10−5 3.2 × 10−4 1.8 × 10−4

Table 5: Finite-difference check of term (I) along the flow, ε = 0.01, h = 10−4 , six random directions per point, float64. “rel. err.” compares the analytic directional derivative to the central difference with P SOT frozen; |∆F |/F is the relative change in the value caused by rebuilding P SOT at the perturbed point, i.e. the only trace term (II) leaves. Figure 8 tracks both norms along the flow. The residual is non-zero everywhere, but its magnitude is nearly constant, varying by under 13% from step 10 onwards, while term (I) decays by a factor of 102 as the flow converges. Their ratio therefore grows from 0.029 at the first step to 1.12 at the last, the two crossing at step 47. The transport cost dates the crossover: ⟨P, C⟩ falls by nearly three orders of magnitude and is within a factor of two of its final value by step 27, where the two gradients still agree to a cosine similarity of 0.982. All of the disagreement accrues afterwards, over steps in which the objective is no longer moving. The same three-part picture holds across ε ∈ {0.003, 0.01, 0.03, 0.1}, with the residual’s floor, and hence the crossover step, set by ε; at ε = 0.003 even 600 inner iterations do not converge the solve, so that column mixes the effect with truncation error and is marked accordingly. The exact W2 distance, computed by dense EMD and so independent of ε and L, plateaus at the same step as ⟨P, C⟩ in every case, confirming that the flow has genuinely arrived rather than that the regularised objective has saturated. The analytic gradient is within a few percent of the fully-unrolled one throughout the phase in which the flow transports mass, and departs from it only once the objective has settled at its O(ε) floor, where the descent direction no longer matters. Unrolling the solve is least defensible where it is usually assumed to be safest, and its O(k) stored activations buy nothing.

26

(I) envelope gradient vs (II) reference-plan gradient along the flow

gradient norm

= 0.003 *

= 0.01

N = 1000, L = 100, 600 inner iterations = 0.03

= 0.1

10 2 10 3

transport cost P, C

100 10 1 10 2

Wasserstein distance W2

10 3 100

10 1

0

20 40 gradient step

0

20 40 gradient step

(I) envelope gradient (analytical)

0

20 40 gradient step

(I) envelope gradient (numerical)

0

20 40 gradient step

(II) reference-plan gradient

Figure 8: Envelope gradient (I) and the solver residual along the blob-to-crescent flow, one column per ε, N = M = 1000, L = 100, 600 inner iterations. Top: term (I) computed analytically (closed form) and numerically (automatic differentiation with the prior coupling held constant), indistinguishable at this scale, together with the residual. Middle: the plan’s transport cost ⟨P, C⟩. Bottom: the exact 2-Wasserstein distance W2 to the target. The residual is roughly constant while term (I) decays by two orders of magnitude, so it becomes comparable only after the transport cost has plateaued at its ε-dependent floor. ∗ marks a column whose solve is not converged at 600 iterations.

27

Related documents

• SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport#792283
• SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport#788283
• Sobolev Gradient Ascent for Optimal Transport: Barycenter Optimization and Convergence Analysis#153063
• An exact and fast solution of the inverse Regularized Optimal Transport problem#821231
• TBR: Transport-Based Rendering with Deposition Strokes for Inverse Graphics#999314