FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport
Abstract
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 \textbf{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
Related
- The optimal-transport cartogram: world population as a Brenier map
- The gravitational Olbers’ paradox and the entropic accelerating universe
- On entropic cosmology, late-time acceleration and the quantum bounce
- Theoretical consistency and observational constraints on entropic cosmology
- Entropic Chaplygin-Gas cosmology and late-universe tension diagnostics
BibTeX
@misc{ye2026flashsinkhorn,
title = {{FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport}},
author = {Felix X. -F. Ye and Yu Chin Fabian Lim and Naigang Wang and Davis Wertheimer},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.02395},
eprint = {2610.02395},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.02395},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport
AU - Felix X. -F. Ye
AU - Yu Chin Fabian Lim
AU - Naigang Wang
AU - Davis Wertheimer
PY - 2026
DA - 2026/10/01
PB - arXiv
DO - 10.48550/arxiv.2610.02395
UR - https://arxiv.org/abs/2610.02395
AB - 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 \textbf{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×10^8-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
ER -