Importance Sparsification for Sinkhorn Algorithm

TL;DR

This paper introduces Spar-Sink, a sparsification method that accelerates the Sinkhorn algorithm for optimal transport problems by reducing computational complexity while maintaining accuracy.

Contribution

The paper proposes a novel importance sparsification technique, Spar-Sink, that significantly speeds up Sinkhorn iterations with theoretical guarantees and practical effectiveness.

Findings

Spar-Sink reduces computational cost from O(n^2) to near O(n) per iteration.

Theoretical proof of estimator consistency under mild conditions.

Empirical results show Spar-Sink outperforms competitors in accuracy and speed.

Abstract

Sinkhorn algorithm has been used pervasively to approximate the solution to optimal transport (OT) and unbalanced optimal transport (UOT) problems. However, its practical application is limited due to the high computational complexity. To alleviate the computational burden, we propose a novel importance sparsification method, called Spar-Sink, to efficiently approximate entropy-regularized OT and UOT solutions. Specifically, our method employs natural upper bounds for unknown optimal transport plans to establish effective sampling probabilities, and constructs a sparse kernel matrix to accelerate Sinkhorn iterations, reducing the computational cost of each iteration from $O(n^2)$ to $\widetilde{O}(n)$ for a sample of size $n$. Theoretically, we show the proposed estimators for the regularized OT and UOT problems are consistent under mild regularity conditions. Experiments on various synthetic data demonstrate Spar-Sink outperforms mainstream competitors in terms of both estimation error and speed. A real-world echocardiogram data analysis shows Spar-Sink can effectively estimate and visualize cardiac cycles, from which one can identify heart failure and arrhythmia. To evaluate the numerical accuracy of cardiac cycle prediction, we consider the task of predicting the end-systole time point using the end-diastole one. Results show Spar-Sink performs as well as the classical Sinkhorn algorithm, requiring significantly less computational time.