FINOM: Fast Sinkhorn on Non-uniform Meshes

Abstract

A linear-complexity algorithm for computing the Wasserstein-1 distance on non-uniform meshes is proposed. This work extends the fast Sinkhorn algorithms from [Q. Liao et al., Commun. Math. Sci., 20(2022)] and [Q. Liao et al., J. Sci. Comput., 98 (2024)] to non-uniform meshes. In those prior works, a distinctive collinear structure of the kernel matrix on uniform meshes was identified, enabling \(O(N)\) acceleration via dynamic programming. While non-uniform meshes are prevalent in practical applications like computational fluid dynamics and finance, their lack of collinearity has hindered direct acceleration. In this paper, we introduce the concept of a ``dividing index'', which partitions the kernel matrix into two blocks. We demonstrate that each block exhibits a quasi-collinear property, a generalization of the structure found in uniform meshes. Leveraging this insight, we develop \textbf{F}ast S\textbf{I}nkhorn algorithm on \textbf{NO}n-uniform \textbf{M}eshes (\textbf{FINOM}), a dynamic programming approach that reduces the per-iteration complexity of the Sinkhorn algorithm from \(O(N^2)\) to \(O(N)\). Extensive numerical experiments on 1D and 2D problems confirm these improvements, achieving speed-ups of several orders of magnitude while maintaining accuracy.