SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport
Machine Learning
Summary
The authors address two main problems with a popular algorithm for entropic optimal transport (EOT), which is slow on large datasets and tends to spread mass unnecessarily at moderate regularization. They introduce SinkSLOT, a new method that uses a smarter way to sparsify computations by slicing the transport plan, making each iteration much faster and still accurate. They prove SinkSLOT converges reliably and produces a mathematically sound objective without needing corrections. Their experiments show SinkSLOT is faster than existing methods and works well in applications involving gradient flows.
Authors
Ian Hsieh, Soumya Snigdha Kundu, Tom Vercauteren, Reuben Dorent
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\geq10^4$). 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.