SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport
Authors: Ian Hsieh, Soumya Snigdha Kundu, Tom Vercauteren, Reuben Dorent
Organizations: 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 · King’s College London, UK
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(N2) 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.
Figures & tables
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(NlogN) 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(NM) to O(L(N+M)) . Output: the sparse coupling Pε⋆ and the SinkSLOT divergence SLOTε , which requires no debiasing.
Method
Reference coupling
Reference coupling cost
Sinkhorn iter.
EOT ( Cuturi, 2013 )
a⊗b
O(N2)
O(N2)
SROT ( Nguyen, 2026 )
(1−γ)PSOT+γa⊗b
O(LNd+LNlogN+N2)
O(N2)
SinkSLOT (ours)
PSOT
O(LNd+LNlogN)
O(S)
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.
Half-moons
8-Gaussians
Two-rings
Gauss d=3
Gauss d=64
Cost gap threshold:
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
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
4.3
9.7
9.6
2.8
6.7
11.4
SROT Nguyen (2026)
124.9
118.6
118.6
103.0
160.2
119.5
122.8
131.3
131.3
29.5
62.1
89.7
15.2
27.1
35.1
Spar-Sink Li et al. (2023)
4.8
2.5
2.5
—
—
37.4
45.5
5.5
3.3
—
—
—
—
—
—
Runtime (ms):
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.
Figure 2: Forward-pass runtime against (N,M) and d . Error bars indicate the standard error over five random seeds.
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.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
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.
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 .
Figure 6: Pareto fronts for the accuracy–speed trade-off on the Gaussian d=64 dataset. FlashSinkhorn vs SinkSLOT at four slice counts.
Half-moons
8-Gaussians
Two-rings
Gauss d=3
Gauss d=64
Method
≤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
250
150
90
90
10000
3240
1430
430
520
990
FlashSink Ye et al. [2026]
260
70
70
3060
430
220
310
90
50
8220
4090
1920
430
430
430
SROT Nguyen [2026]
450
140
140
3210
470
260
150
100
100
7320
3290
1550
1230
870
700
Spar-Sink Li et al. [2023]
260
70
70
—
—
1010
654
96
50
—
—
—
—
—
—
Appendix
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
8-Gaussians
Two-rings
Gauss d=3
Gauss d=64
Method
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
≤1%
≤5%
≤10%
Peak memory (MB):
SinkSLOT (ours)
820
820
820
896
820
820
820
820
820
2580
1799
1281
5050
3007
1302
FlashSink Ye et al. [2026]
1785
1785
1785
1785
1785
1785
1785
1785
1785
1785
1785
1785
1806
1806
1806
SROT Nguyen [2026]
2359
2359
2359
2359
2359
2359
2359
2359
2359
2359
2359
2359
2380
2380
2380
SparSink Li et al. [2023]
1258
1258
1258
—
—
1258
1258
1258
1258
—
—
—
—
—
—
Appendix
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.
Figure 7: Forward-pass peak memory against (N,M) and d . Error bars indicate the standard error over five random seeds.
step
analytic
FD (frozen)
rel. err.
∣ΔF∣/F
0
4.431×10−4
4.431×10−4
1.1×10−9
1.6×10−6
10
−2.906×10−4
−2.906×10−4
8.7×10−5
1.3×10−5
25
−6.904×10−5
−6.903×10−5
6.6×10−4
8.6×10−5
40
−7.862×10−6
−7.853×10−6
2.3×10−3
3.2×10−4
50
−6.683×10−6
−6.696×10−6
3.8×10−3
1.8×10−4
Appendix
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 PSOT frozen; ∣ΔF∣/F is the relative change in the value caused by rebuilding PSOT at the perturbed point, i.e. the only trace term (II) leaves.
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.
We propose a new regularized optimal transport (OT) formulation, termed sliced-regularized optimal transport (SROT). Unlike entropic OT (EOT), which regularizes the transport plan toward an independent coupling, SROT regularizes it toward a smoothened sliced OT (SOT) plan. To the best of our knowledge, SROT is the first approach to leverage a version of SOT plan as a reference to improve classical OT. We provide a formal definition of SROT, derive its dual formulation, and provide a post-Bayesian interpretation of SROT. We then develop a Sinkhorn-style algorithm for efficient computation, retaining the same scalability advantages as EOT. By incorporating a scalable SOT plan as a prior, SROT yields more accurate approximations of the exact OT plan than EOT under the same level of regularization. Moreover, the resulting transport plan improves upon the reference SOT plan itself. We further introduce the corresponding OT divergence induced by SROT, named SROT divergence, and analyze its topological and computational properties. Finally, we validate our approach through experiments on synthetic datasets and color transfer tasks, demonstrating that SROT is better than both EOT and SOT in approximating exact OT. Additional experiments on gradient flows further highlight the advantages of SROT divergence.
Streaming GPU solvers for entropic optimal transport (EOT), such as FlashSinkhorn, avoid storing the dense kernel but still evaluate all n×m point pairs in every Sinkhorn iteration. We present FlashSinkhorn~2 (FS2), a solver for squared-Euclidean cost on low-dimensional point clouds that solves large discrete EOT problems to a prescribed marginal residual on a single GPU by coupling two stages. A coarse stage solves on cell centroids, lifts the potentials to every point and, when a sampled marginal check rejects the lift, continues on the centroids, replacing most point-level updates. A block-sparse fine stage then removes the centroid error that coarse updates cannot. Its Morton-ordered blocks support screening and fused tensor-core execution, and a threshold set by the block masses bounds each omitted tile's contribution to every row and column. On synthetic benchmarks, FS2 reaches the target residual on all 32 problems and GeomLoss multiscale on 10. On one A100, FS2 solves discrete EOT between two 1.34×108-particle measures from a cosmological N-body simulation, at an entropic blur equal to the mean interparticle distance, to an all-particle marginal residual below 0.01 in under 2.5 hours. To our knowledge, it is the largest discrete EOT problem solved to this accuracy within hours. For reproducibility, we release an open-source implementation at https://github.com/ot-triton-lab/flash-sinkhorn
Felix X. -F. Ye, Yu Chin Fabian Lim, Naigang Wang +1
Department of Mathematics & Statistics, University at Albany, Albany, NY, USA · IBM T. J. Watson Research Center, Yorktown Heights, NY, USA
Entropic regularized optimal transport (OT) via the Sinkhorn algorithm has become a fundamental tool in machine learning, yet existing implementations either suffer from numerical instability for small regularization parameters or incur significant overhead from deep learning frameworks. We present FastSinkhorn, a lightweight, native CUDA implementation of the log-domain Sinkhorn algorithm that combines warp-level shuffle reductions with shared-memory tiling to achieve high GPU utilization without sacrificing numerical stability. Our solver operates entirely in the log-domain, enabling robust computation for regularization parameters as small as epsilon = 10^{-4} where standard-domain methods fail. On dense OT problems with n = m = 8192, our implementation achieves 12x speedup over the widely-used POT library and 5.9x speedup over GPU-accelerated PyTorch baselines, while consuming only 256 MB of GPU memory. We validate our solver on image color transfer, 3D point cloud matching, and convergence analysis, demonstrating that native CUDA kernels with careful numerical treatment provide a practical and efficient foundation for large-scale optimal transport computation.