Geodesic Sinkhorn for Fast and Accurate Optimal Transport on Manifolds

TL;DR

This paper introduces Geodesic Sinkhorn, a fast and accurate optimal transport method on manifolds that leverages heat kernel diffusion, reducing computational complexity from quadratic to near-linear, and is applied to high-dimensional biological data analysis.

Contribution

The paper presents Geodesic Sinkhorn, a novel optimal transport algorithm on manifolds using heat kernel approximation, significantly improving computational efficiency over existing methods.

Findings

Achieves $O(n \, \log n)$ computation complexity.

Effectively computes barycenters of high-dimensional single-cell data.

Identifies optimal transport paths related to treatment effects.

Abstract

Efficient computation of optimal transport distance between distributions is of growing importance in data science. Sinkhorn-based methods are currently the state-of-the-art for such computations, but require $O(n^2)$ computations. In addition, Sinkhorn-based methods commonly use an Euclidean ground distance between datapoints. However, with the prevalence of manifold structured scientific data, it is often desirable to consider geodesic ground distance. Here, we tackle both issues by proposing Geodesic Sinkhorn -- based on diffusing a heat kernel on a manifold graph. Notably, Geodesic Sinkhorn requires only $O(n\log n)$ computation, as we approximate the heat kernel with Chebyshev polynomials based on the sparse graph Laplacian. We apply our method to the computation of barycenters of several distributions of high dimensional single cell data from patient samples undergoing chemotherapy. In particular, we define the barycentric distance as the distance between two such barycenters. Using this definition, we identify an optimal transport distance and path associated with the effect of treatment on cellular data.