FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport

2026-10-01

AI Summary

This study addresses the memory and computational bottlenecks encountered in all-pairs evaluation for large-scale Entropic Optimal Transport (EOT). We propose a two-stage solver that first obtains an approximate solution through coarse-grained centroid iterations, followed by block-sparse fine-grained refinement to enhance accuracy. The core innovation lies in introducing a coupled centroid promotion and block-sparse filtering mechanism, integrated with Morton encoding to enable Tensor Core-fused acceleration, thereby significantly reducing redundant computation. This approach achieves high-precision solutions for hundred-million-particle cosmological simulation data within 2.5 hours on a single A100 GPU, establishing a new record for the largest scale among comparable methods.

Abstract

Streaming GPU solvers for entropic optimal transport (EOT), such as FlashSinkhorn, avoid storing the dense kernel but still evaluate all $n\times 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\times10^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

Problem

• Entropic Optimal Transport
• Large-scale discrete EOT
• GPU solver
• Computational efficiency

Innovation

• Entropic Optimal Transport
• Block-Sparse
• Sinkhorn Algorithm
• Tensor-Core
• Morton Order