GEODESIC SINKHORN FOR FAST AND ACCURATE OPTIMAL TRANSPORT ON MANIFOLDS

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\left(n^2\right)$ 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\mathrm{l}\mathrm{o}\mathrm{g}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.

1. INTRODUCTION

Optimal Transport (OT) distances or Wasserstein distances are computed by lifting ground distances between points to distances between measures. This distance is computed relative to a ground distance on the support of the distributions, making it more informative than distances based only on a pointwise comparison of the densities. However, to compute the Wasserstein, one needs to find the optimal transport plan from the source distribution to a target distribution; this is a linear programming problem requiring $O\left(n^3\mathrm{l}\mathrm{o}\mathrm{g}n\right)$ for discrete distributions of size $n$ [1].

An efficient modification of the optimal transport problem is to consider entropy-regularized transportation. This formulation is solved with the Sinkhorn algorithm [2] by iteratively rescaling a Gaussian kernel based on the distance matrix. It is equivalent to the Schrödinger Bridge problem, for which similar algorithms were developed [3]–[5]. In the discrete case, it requires $O\left(n^2\right)$ for distributions of size $n$, since it relies on matrix-vector products. Furthermore, this formulation allows for fast computation of the discrete barycenter with fixed support (the average distributions w.r.t. the Sinkhorn distance). An important drawback of the Sinkhorn algorithm is the necessity of storing and multiplying the pairwise distance matrix with a vector.

Additionally, the ground distance is commonly chosen as the Euclidean distance. The Euclidean distance is often sub-optimal for high-dimensional datasets over larger distances according to the manifold hypothesis, which says observations lie near a low dimensional (curved) manifold in high dimensional space [6]. For higher dimensional datasets assumed to be sampled from a lower dimensional manifold, using a distance closer to the manifold for OT has shown promising results [7]–[10].

In this work, we present Geodesic Sinkhorn¹ a Sinkhorn-based method for fast optimal transport with a heat-geodesic ground distance. Our method is based on the geometry of the dataset constructed from a common graph and uses the heat kernel on the graph to defined a heat-geodesic distance. Key to this approach, we will never need to construct or operate on an $n\times n$ distance matrix, and we will only use the sparse Laplacian matrix and sparse matrix-vector products. For sparse graphs, this can be used for $O(n\mathrm{l}\mathrm{o}\mathrm{g}n)$ computation of the Sinkhorn distance with a manifold ground distance, improving on the $O\left(n^2\right)$ implementations based on dense matrices.

Increasing the state-of-the art efficiency in Sinkhorn computation opens us up to being able to perform complex operations on large groups of datasets. In particular, we consider interpolating between datasets and show that using our heat-geodesic distance improves the interpolation accuracy compared to OT with Euclidean distance. The barycenter corresponds to the average distribution of a set of distributions. Our method allows for finer-grained barycenters on a data manifold, which motivates us to define a novel notion of dissimilarity between families of distributions called barycentric distance.

We apply the barycentric distance to single cell data from patient-derived cancer organoids (PDOs) to assess the effect of treatments (such as drugs and chemotherapy). Here we have one set of PDOs from control conditions, and another set that are treated. The treatment effect is thus the distance between these barycenters. In addition, we use Geodesic Sinkhorn’s barycenter to compare the effect from one family of distributions to another.

Our main contributions include: (1) A new method for computing optimal transport distances on a manifold called Geodesic Sinkhorn, which is highly efficient in time and memory. (2) Defining the barycentric distance; a novel distance between families of distributions, and showing its utility in deriving treatment effect from control and treated patient samples.

2. RELATED WORK

Geodesic Sinkhorn is related to prior work linking the entropy-regularized optimal transport problem triangular mesh with the heat operator [8], [11], but using different graph filtering techniques. These approaches approximate the application of the heat kernel to a vector by discretizing the heat equation and solving systems of linear equations. This technique was used in different contexts, either with the cotangent Laplacian [8] or to learn a ground metric [12]. Solving these systems for each Sinkhorn iteration can be done efficiently with a sparse Cholesky decomposition. However, this method’s efficiency depends mainly on the efficiency of the Cholesky decomposition which can be slow depending on the sparsity pattern is $O\left(n^3\right)$ for an $n\times n$ matrix, and necessitates solving $2K$ systems of linear equations per Sinkhorn iteration, where $K$ is the number of sub-steps in the backward Euler discretization.

3. PRELIMINARIES

In this section, we start by reviewing the basics of OT and the Wasserstein distance, as well the Sinkhorn distance. Then we review two notions fundamental to our method; the heat equation on a graph and the Chebyshev approximation of the heat kernel.

3.1. Wasserstein Distance

In the following, we assume that all distributions admit a density or a probability mass function, and we use the same notation for both. Let $\mu ,\nu$ be two probability distributions on a measurable space $𝒳\subseteq {\mathbb{R}}^d$ with metric $d(\cdot ,\cdot )$, let $\mathrm{\Pi }(\mu ,\nu )$ be the set of joint probability distributions $\pi$ on the space $𝒳\times 𝒳$ where, for any measurable subset $\omega \subset 𝒳,\pi (\omega \times 𝒳)=\mu \left(\omega \right)$ and $\pi (𝒳\times \omega )=\nu \left(\omega \right)$. The $p$-Wasserstein distance is defined as:

$$W_p\left(\mu ,\nu \right)≔{\left(\underset{\pi \in \mathrm{\Pi }\left(\mu ,\nu \right)}{\mathrm{i}\mathrm{n}\mathrm{f}}{\int }_{{𝒳}^2}d(x,y{)}^p\mathrm{d}\pi (x,y)\right)}^{1/p}.$$ (1)

In the following, we consider $p=2$. An exact algorithm based on linear programming can solve this problem in $O\left(n^3\mathrm{l}\mathrm{o}\mathrm{g}n\right)$ time for discrete distributions of size $n$.

3.2. Sinkhorn Distances

The Kullback-Leibler (KL) divergence between $\pi$ and some strictly positive $K$ on $𝒳\times 𝒳$ is defined as

$$D_{\mathrm{K}\mathrm{L}}\left(\pi ∣K\right)≔{\int }_{{𝒳}^2}\left(\text{ln}\frac{\pi \left(x,y\right)}{K\left(x,y\right)}-1\right)\mathrm{d}\pi \left(x,y\right).$$ (2)

The Sinkhorn distance² is a relaxation of equation 1 where the infimum is over all coupling in $\left\{\pi \in \mathrm{\Pi }(\mu ,\nu )\mid D_{\mathrm{K}\mathrm{L}}(\pi \mid \mu \times \right.\nu )\le \xi \}$ for $\xi >-1$. Introduced in [13], the optimization of this distance can be solved by considering the entropyregularized transport

$$W_{d,\lambda }^2(\mu ,\nu )≔{\left(\underset{\pi \in \mathrm{\Pi }(\mu ,\nu )}{\mathrm{i}\mathrm{n}\mathrm{f}}{\int }_{{𝒳}^2}d(x,y{)}^2\mathrm{d}\pi (x,y)-\lambda H(\pi )\right)}^{1/2},$$ (3)

where we define the entropy of a coupling $\pi$ as $H\left(\pi \right):=-\int \mathrm{l}\mathrm{n}\pi (x,y)\mathrm{d}\pi (x,y)$, and $\lambda >0$. This formulation converges to the Wasserstein distance as $\lambda \to 0$, and can be solved with the Sinkhorn algorithm with complexity of the order $O\left(n^2/ϵ\right)$ for discrete distributions of size $n$ [13]. In the discrete case, the transport matrix $\pi$ admits the form $\mathrm{d}\mathrm{i}\mathrm{a}\mathrm{g}\left(\mathit{v}\right)K_{\lambda }\mathrm{d}\mathrm{i}\mathrm{a}\mathrm{g}\left(\mathit{w}\right)$, where $\mathit{v},\mathit{w}$ are vectors of size $n$. The Sinkhorn algorithm iteratively updates the vectors as $(\mathit{v},\mathit{w})\leftarrow \left(\mu ./K_{\lambda }\mathit{w},\nu ./K_{\lambda }^{\mathrm{\text{'}}}\mathit{v}\right)$, where $K_{\lambda }:=e^{-d(x,y{)}^2/\lambda }$.

Following [8], using the kernel $K_{\lambda }$ gives an alternative interpretation of the Sinkhorn distance as

$$W_{d,\lambda }^2\left(\mu ,\nu \right)={\lambda }^{1/2}{\left(1+\underset{\pi \in \mathrm{\Pi }\left(\mu ,\nu \right)}{\mathrm{m}\mathrm{i}\mathrm{n}}{D}_{\mathrm{K}\mathrm{L}}\left(\pi \mid K_{\lambda }\right)\right)}^{1/2}.$$ (4)

The problem in equation 3 is strictly convex and continuous yielding a unique minimizer. In the discrete case, this leads to an algorithm for the entropy-regularized Wasserstein distance based on the Sinkhorn algorithm enforcing the marginal constraints on the kernel $K_{\lambda }$ while minimizing the distance as quantified by $D_{\mathrm{K}\mathrm{L}}$.

The underlying metric $d(\cdot ,\cdot )$ is generally unknown, thus the kernel $K_{\lambda }$ cannot be evaluated. The authors of [8] proposed to approximate $K_{\lambda }$ with the heat kernel ${\mathscr{H}}_t(x,y)$ on $𝒳$. According to Varadhan’s formula [14], the geodesic distance on a manifold can be recovered from the heat transfer at small timescales as

$$d(x,y{)}^2=\underset{t\to 0^{+}}{\mathrm{l}\mathrm{i}\mathrm{m}}-4t\text{ln}{\mathscr{H}}_t\left(x,y\right).$$ (5)

Hence, motivating the use of the heat-geodesic distance $d_{\mathscr{H}}^2(x,y):=-4t\mathrm{l}\mathrm{n}{\mathscr{H}}_t(x,y)$, with associated kernel $K_{\lambda }(x,y)=$ ${\mathscr{H}}_{\lambda /4}(x,y)$. Interestingly, Sinkhorn-based methods admit an efficient algorithm to solve the barycenter problem which we present next.

3.3. Interpolation with discrete support

By constraining the support to a set $𝒳$ (or a graph), we can efficiently interpolate between more than two distributions. The barycenter problem [1], [8], [15] generalizes the notion of average between points to an average between distributions. For a set of $m$ distributions $\left\{{\mu }_1,\dots ,{\mu }_m\right\}$ supported on $𝒳$, the objective is to find a distribution minimizing the average distance

$${\mu }^{*}≔\underset{\mu \in 𝒫\left(𝒳\right)}{\mathrm{a}\mathrm{r}\mathrm{g}\mathrm{m}\mathrm{i}\mathrm{n}}\sum _{i=1}^m{\alpha }_i{W}_d^p{\left(\mu ,{\mu }_i\right)}^p,$$

where $𝒫\left(𝒳\right)$ denotes the space of probability distributions supported on $𝒳$, and $\left\{{\alpha }_1,\dots ,{\alpha }_m\right\}$ are non-negative weights. Finding the barycenter is a challenging optimization problem, however the barycenter for Sinkhorn-based methods admits an efficient computation. It involves updating $m$ vectors ${\mathit{v}}_i,{\mathit{w}}_i$, which define a transport plan from ${\mu }_i$ to the barycenter ${\mu }^{*}$. The support of the barycenter is constrained to $𝒳$, for most Sinkhorn-based methods the size of $𝒳$ needs to be small for computational reason. Our method does not suffer from such a limitation. Hence, we can consider barycenter with greater expressivity, and interpolate between large sets of distributions.

3.4. Heat Diffusion on a Graph

Consider an undirected graph $𝒢=(V,E)$ with a set $V$ of $n$ vertices and a set of edges $E$, and its weighted adjacency matrix $\mathit{A}$ with non-negative edge weights, and the diagonal degree matrix $\mathit{D}$, where ${\mathit{D}}_{ii}:={\sum }_k{\mathit{A}}_{ik}$. We define the combinatorial Laplacian as $\mathit{L}:=\mathit{D}-\mathit{A}$, for any function $f:V\to \mathbb{R}$ we have $\left(\mathit{L}f\right)\left(v\right)={\sum }_u{a}_{u,v}\left(f\right(v)-f(u\left)\right)$. The combinatorial Laplacian is a symmetric positive semi-definite matrix, and has an eigendecomposition $\mathit{L}=\mathrm{\Psi }\mathrm{\Lambda }{\mathrm{\Psi }}^T$ with orthonormal eigenvectors $\mathrm{\Psi }$ and diagonal eigenvalue matrix $\mathrm{\Lambda }=\mathrm{d}\mathrm{i}\mathrm{a}\mathrm{g}\left({\lambda }_1,{\lambda }_2,\dots ,{\lambda }_n\right)$, such that $0\le {\lambda }_1\le {\lambda }_2\le \cdots \le {\lambda }_n$. The combinatorial Laplacian is a natural extension of the negative of the Laplacian operator to a graph. For a signal ${\mathit{f}}_0\in {\mathbb{R}}^n$ on $𝒢$, the diffusion of ${\mathit{f}}_0$ on the graph evolves according to the heat equation

$$\frac{d}{dt}\mathit{f}\left(t\right)+\mathit{L}\mathit{f}\left(t\right)=\mathbf{0}\mathit{\text{, s.t. }}\hspace{1em}\mathit{f}\left(0\right)={\mathit{f}}_0\hspace{1em}t\in {\mathbb{R}}^{+}.$$

The heat kernel solves this ODE, it is defined by the matrix exponential ${\mathit{H}}_t:=e^{-t\mathit{L}}$. By orthogonality of the eigenvectors of $\mathit{L}$, we can write ${\mathit{H}}_t=\mathrm{\Psi }{e}^{-t\mathrm{\Lambda }}{\mathrm{\Psi }}^T$ and $\mathit{f}\left(t\right)={\mathit{H}}_t{\mathit{f}}_0$. Computing ${\mathit{H}}_t$ by eigendecomposition would require $O\left(n^3\right)$ operations. Recall that, for the Sinkhorn algorithm, we are only concerned with the application of the heat operator ${\mathit{H}}_t$ on a signal $\mathit{f}\in {\mathbb{R}}^n$. For larger diffusion time, the heat kernel converges to its eigenvector associated to the lowest eigenvalues of the Laplacian, hence, intuitively, the heat kernel corresponds to a low-pass filter. In Geodesic Sinkhorn, we use Chebyshev polynomials [16], [17] to approximate the application of the heat operator to a signal. For a short timescale $t$, using the heat kernel accounts for using the geodesic distance as ground distance in the entropy-regularized OT formulation equation 3.

3.5. Chebyshev Polynomials

Polynomial sequences are often used to approximate functions or operator. With Chebyshev polynomials, we can approximate the application of the matrix exponential ${\mathit{H}}_t=e^{-t\mathit{L}}$ to a signal $\mathit{f}$ on the graph. An attractive property of Chebyshev polynomials is that the approximation error decays exponentially with the maximum degree $K$. They are defined by the recursive relation ${\left\{T_k\right\}}_{k\in \mathbb{N}}$ with $T_0\left(y\right)=0,T_1\left(y\right)=y$ and $T_k\left(y\right)=2y{T}_{k-1}\left(y\right)-T_{k-2}\left(y\right)$ for $k\ge 2$. On [−1, 1] these polynomials are orthogonal w.r.t. the weight $(1-y{)}^{-1/2}$, and can be used to express the operator ${\mathit{H}}_t$. Assuming the largest eigenvalue ${\lambda }_n\le 2$, we can write

$${\mathit{H}}_t=\frac{b_0}{2}+\sum _{k=1}^{\mathrm{\infty }}{b}_k{T}_k(\mathit{L}-\mathit{I}\mathit{d}),$$

where the $K+1$ scalar coefficient $\left\{b_k\right\}$ depend on time and can be evaluated with the Bessel function. The approximation of ${\mathit{H}}_t$ is based on the first $K$ term of the series which we note $p_K(\mathit{L},t)$. It results in $K$ matrix-vector products which can be efficient since, in general, $\mathit{L}$ is a sparse matrix. On a $m$-nearest neighbor graph, this can be $O(Kmn/\lambda )$, where $\lambda$ is a regularization parameter. Chebyshev polynomials admits interesting theoretical properties and are known to converge faster than other polynomials [17], [18]. The choice of the parameter $K$ is related to the number of neighbors or the connectivity of the graph. For small diffusion time, hence only diffusing in a local neighborhood, the approximation is accurate even with a small $K$. As the diffusion time increases, $K$ has to increase in order to consider a larger neighborhood around a node. For OT, we consider small diffusion time, and we found that our results were stable for all $K$ greater than 10.

4. GEODESIC SINKHORN DISTANCES

We define the Geodesic Sinkhorn distance between any signals or distributions on a graph $𝒢$ by the entropy-regularized OT with the heat kernel ${\mathit{H}}_t$ on the graph. This construction is also valid between any point cloud datasets. In that case, for $m$ datasets $\left\{{\mathrm{X}}_1,\dots ,{\mathrm{X}}_{\mathrm{m}}\right\}$ sampled from a set of distributions $\left\{{\mu }_1,\dots ,{\mu }_m\right\}$, we construct a common graph using an affinity kernel on the $m$ datasets and compare two samples by taking the distance between two indicator functions on the graph. We approximate the heat kernel ${\mathit{H}}_t$ with Chebyshev polynomials $p_K(\mathit{L},l)$ of order $K$. In Algorithm 1, we present the main steps to evaluate the Geodesic Sinkhorn. It is based upon Sinkhorn iterations [2], [13], where ⊘ and ⊙ denote respectively the elementwise division and multiplication. Note that, as opposed to the usual Sinkhorn algorithm, we never have to store a dense $n\times n$ distance matrix, but only the usually sparse graph Laplacian.

Definition 4.1. The Geodesic Sinkhorn distance between two distributions $\mu ,\nu$ on a graph $𝒢$ is

$$W_{{\mathit{H}}_t}(\mu ,\nu )≔4t^{1/2}{\left(1+\underset{\pi \in \mathrm{\Pi }(\mu ,\nu )}{\mathrm{m}\mathrm{i}\mathrm{n}}{D}_{\mathrm{K}\mathrm{L}}\left(\pi \mid {\mathit{H}}_t\right)\right)}^{1/2}.$$

In the following proposition, we find the ground distance implicitly used in the optimal transport defined by Geodesic Sinkhorn. We use ≃ for the equivalence relation between distances.

Proposition 4.2. There exists a maximum Chebyshev polynomial degree $K$ such that the ground distance in Geodesic Sinkhorn is equivalent to the one based ${\mathit{H}}_t$

$$-4t\mathrm{l}\mathrm{o}\mathrm{g}{\left(p_K(\mathit{L},t)\right)}_{ij}\simeq -4t\mathrm{l}\mathrm{o}\mathrm{g}{\left({\mathit{H}}_t\right)}_{ij}$$

In particular, the Wasserstein distances with these ground distances are equivalent.

Proof. Because the approximation error decreases exponentially in $K$ [17], we have that for any $ϵ>0$ sufficiently small there exist $K_i$ such that ${\left(H_t\right)}_{ij}-ϵ<{\left(p_K\right)}_{ij}<{\left(H_t\right)}_{ij}+$ $ϵ$. Choose $K$ such that this is true for all vertices $K:=\mathrm{m}\mathrm{a}\mathrm{x}\left\{K_1\dots ,K_n\right\}$. We define

$$c≔\underset{ij}{\mathrm{m}\mathrm{i}\mathrm{n}}\frac{{\left({\mathit{H}}_t\right)}_{ij}-ϵ}{{\left({\mathit{H}}_t\right)}_{ij}}\text{ and }C≔\underset{ij}{\mathrm{m}\mathrm{a}\mathrm{x}}\frac{{\left({\mathit{H}}_t\right)}_{ij}+ϵ}{{\left({\mathit{H}}_t\right)}_{ij}}$$

and we have the equivalence between the distances since

$$c{\left({\mathit{H}}_t\right)}_{ij}\le {\left(p_K(\mathit{L},t)\right)}_{ij}\le C{\left({\mathit{H}}_t\right)}_{ij}\text{ for all }ij\in \left[n\right],$$

and since the logarithm is a monotonic function. □

In [8], [12], using the Euler implicit discretization results in a ground cost of the form $-ϵ\mathrm{l}\mathrm{n}\left({\left(\mathbf{I}\mathbf{d}-\frac{ϵ}{4K}\mathit{L}\right)}^{-K}\right)$, where Id is the identity matrix, and can be seen as another approximation for the matrix exponential.

The efficiency of Geodesic Sinkhorn improves the notion of barycenter as it is possible to consider much larger graph $𝒢$, thus a finer grained support of the barycenter. This leads us to define a novel distance between families of distributions.

Definition 4.3. For two finite families of distributions $𝒯$ and $𝒞$ supported on $𝒢$, we define the barycentric distance between the families $𝒯,𝒞$ as

$$\gamma (𝒯,𝒞)≔W_{{\mathit{H}}_t}\left({\mu }_{𝒯}^{*},{\mu }_{𝒞}^{*}\right)$$

where ${\mu }_{𝒯}^{*},{\mu }_{𝒞}^{*}$ are respectively the barycenters of $𝒯$ and $𝒞$.

The previous definition is valid for any distances between distributions or barycenters. However, OT barycenters are known to be more informative than others [15]. We will further explore this comparison in our experiments. We use it to distinguish between two groups in a medical setting where a set of patients received a treatment (defining the family $𝒯$), and another set acts as a control family $𝒞$. Following this idea, we define a notion of effect between two families.

Definition 4.4. For two family of distributions $𝒯$ and $𝒞$ supported on $𝒢$, define the Expected Barycenter Effect of $𝒯$ as

$$\tau \left(𝒯\right)≔{\mathbb{E}}_{{\mu }_{𝒯}^{*}}\left(Y_t\right)-{\mathbb{E}}_{{\mu }_{𝒞}^{*}}\left(Y_c\right),$$

where ${\mu }_{𝒯}^{*},{\mu }_{𝒞}^{*}$ are respectively the barycenters of $𝒯$ and $𝒞$, and the features $Y_c~{\mu }_{𝒞}^{*}$ and $Y_t~{\mu }_{𝒯}^{*}$.

Note that we compute the average on the family of distributions instead of the average on their support, hence we evaluate their expectations in a closed form. This definition also extends to a conditional equivalent where families of distribution can be subdivided with discrete covariate variables. When the barycenters are computed with the total variation, this definition is equivalent to the naive Average Treatment Effect(ATE) [19]; i.e. difference of empirical means.

5. RESULTS

We demonstrate the accuracy and efficiency of the Geodesic Sinkhorn distance on two tasks: (1) Nearest-Wasserstein-neighbor calculation on simulated data with manifold structure similar to the setup of [10]; (2) A newly defined Barycentric distance between families of distributions computed to quantify the effect of a treatment on patient-derived organoids. In Appendix A.1, we present additional results on time series interpolation.

5.1. Nearest-Wasserstein-neighbor distributions

In this experiment, we compare our method with Sinkhorn [13], and LR Sinkhorn [20], both algorithms with Euclidean and squared Euclidean ground distance, with DiffusionEMD [21], and Sinkorn with Euler approximation of the heat filter. We created 15 Gaussian distributions sampled randomly on a swiss roll dataset, and sampled 10k observations from each distribution. We rotated the observations in 10 dimensions. We consider a k-nearest neighbors task on these distributions. We evaluate the methods with the ground truth, since we know the exact geodesic distance on the manifold. In Tab. 1, we report the average and standard deviation over 10 seeds of the Spearman and Pearson correlations to the ground truth, and the runtime in seconds with and without the computation of the graph. Our method is the most accurate while being much faster than other Sinkhorn-based methods.

Table 1.

KNN task for 15 distributions, best score highlighted is bold. Geodesic Sinkhorn is the most accurate, while being faster than other Sinkhorn-based methods.

Method	SpearmanR	PearsonR	P@5	Time(s) no graph	Time(s)
Diffusion EMD	0.62±0.097	0.736±0.023	0.66±0.072	2.845±0.135	7.877±0.531
Sinkhorn $W1$	0.387±0.044	0.523±0.036	0.471±0.028	112.406±0.206	112.406±0.206
Sinkhorn $W2$	0.411±0.036	0.485±0.027	0.492±0.053	133.686±5.234	133.686±5.234
LR Sinkhorn $W1$	−0.31±0.07	−0.131±0.086	0.237±0.037	578.631±107.82	578.631±107.82
LR Sinkhorn $W2$	0.366±0.048	0.379±0.051	0.447±0.023	204.191±3.656	204.191±3.656
Euler Sinkhorn	0.776±0.061	0.718±0.009	0.728±0.072	449.752±42.985	455.059±43.083
Geodesic Sinkhorn	0.847±0.023	0.754±0.016	0.833±0.034	10.176±1.249	16.682±1.705

5.2. Barycentric distance

We test if we can identify a linear treatment effect with the Expected Barycenter Effect (EBE). In this experiment, we create a control family of distributions $𝒞$ of ten standard Gaussian distributions. The treatment group consists of nine Gaussian distributions $𝒩\left(\mathrm{5,1}\right)$, and one outlier centered at different means. For each distribution, we sample 500 observations, and reproduce the experiment over ten seeds. In Tab. 2, we report the EBE and its standard deviation with the Geodesic Sinkhorn, the Total Variation (TV) distance, and Sinkhorn. Since the TV only compares the mean, it is sensitive to the outlier, whereas our method can identify the true treatment effect.

Table 2.

Expected Barycenter Effect (EBE) with one outlier distribution centered at −60,−30,0, or 5. Comparison using the barycenter from Sinkhorn, total variation, or Geodesic Sinkhorn. Values closer to the real treatment effect of 5 are better.

Outlier	EBE Geo Sinkhorn	EBE Sinkhorn	EBE TV
−60	5.016±0.226	−0.103±0.005	−1.429±0.144
−30	5.053±0.196	0.355±0.049	1.571±0.144
0	4.917±0.315	4.954±0.157	4.571±0.144
No outlier	5.059±0.159	5.054±0.16	5.071±0.144

5.3. Single-cell signaling data

We use single-cell signaling data produced by mass cytometry (MC) for a screening study to compare the treatment effect of different chemotherapies on 10 colorectal cancer (CRC) patient-derived organoids (PDOs) [22]. These PDOs can be grouped into chemoresistant PDOs, that show little-to-no effect when treated with chemotherapies; and chemosensitive PDOs, that present strong shifts in their phenotypes upon treatment. The observations include single-cell data information on the cell cycle and signaling changes upon treatment of PDOs with different CRC treatments at a range of concentrations. In Fig. 1, we present the barycentric distances matrices between treatments a) and between four concentrations of treatment SN-38 (S) b). In both cases, the control groups corresponds to AH and DMSO, the two rightmost columns. We compare the distance matrices between Sinkhorn (left) and our method (right). Our method provide a finer distinction between treatments (Fig. 1 top) and concentrations (Fig. 1 bottom), especially for the chemosensitive group. As observed in [22], chemosensitive PDOs show little-to-no response to lower concentrations of SN-38 (S1), but their phenotype shifts very strongly upon treatment with higher concentrations (S2, S3, and S4) (Fig. 1 b). When comparing combinations of different treatments (Fig. 1 a), Geodesic Sinkhorn better resolves the difference between SN-38 (S) alone and in combination with Cetuximab (C), showing that S is the main agent creating the treatment effect and the combination with C does not resolve in a synergistic effect [22]. Note that we only consider the relative magnitude of the distances, since the two algorithm use different ground distances.

Fig. 1.

a) Barycentric distances matrices for the Sinkhorn algorithm (left) and our method Geodesic Sinkhorn (right). b) Barycentric distances matrices between doses of treatment SN-38, for four concentrations S1 S2 S3 S4. Control groups correspond to AH and DSMO. Geodesic Sinkhorn provides a clearer distinction between treatments, and doses.

6. CONCLUSION

In this work, we considered the use of OT for graphs and large datasets in high dimensions potentially sampled from a lower dimensional manifold. We proposed Geodesic Sinkhorn, a fast implementation of the Sinkhorn algorithm using the graph Laplacian and Chebyshev polynomials. Our method is well adapted for large and high dimensional datasets as it is defined with a geodesic ground distance, which takes into account the underlying geometry of the data, and requires less computation time and less memory. On a synthetic dataset, we showed that Geodesic Sinkhorn is much faster than other Sinkhorn-based methods while being more accurate. With the Wasserstein barycenter, we defined the barycentric distance to compare entire families of distributions, and the expected barycenter effect, then applied both methods to a large PDO drug screen dataset.

A. APPENDIX

A.1. Time series Interpolation

To evaluate Geodesic Sinkhorn’s performance on inferring dynamics, we test its performance on a task for time series interpolation. In this setting the most used datasets are Embryoid Body [1], and WOT [2]. We curated ten scRNA-seq datasets; WOT [2], Clark [3], Embryoid Body [1], and seven more from the 2022 multimodel single-cell integration challeng³ to test our method. The observations are the gene expression of single cells from a distribution evolving through time. The Waddington-OT dataset (WOT) has 38 timpoints of a developing stem cell population over 18 days collected roughly every 6–12 hours. This is the most densely sampled dataset in time. The Embryoid Body dataset is a single-cell RNA seq dataset of developing embryoid bodies from 0–30 days with 5 datasets collected over time. The Clark dataset contains 12 samples over 9 unique timepoints of a developing mouse retina. Finally, the NeurIPS 2022 data contains four donors with single-cell transcriptomic data collected over 4 timepoints for 3 donors with 10X-multiome and 4 combined cite-seq (in the publically released training data) leading to 7 additional time series with 4 timepoints.

The goal is to interpolate the distribution between two timepoints. The number of timepoints in each dataset range from 4 to 40. Here for a dataset with $T$ single-cell distributions $\left\{{\mu }_1,{\mu }_2,\dots ,{\mu }_T\right\}$ over time for $t\in [2\dots T-1]$ we compute the exact Euclidean 2-Wasserstein distance between the interpolated distribution ${\hat{\mu }}_t$ at time $t$ and the ground truth distribution ${\mu }_t,W_2\left({\hat{\mu }}_t,{\mu }_t\right)$. Since we interpolate between two distributions we used the McCann interpolant as its support is ${\mathbb{R}}^d$. We compare our Geodesic Sinkhorn interpolation with either the Sinkhorn Mccann interpolant with Euclidean ground distance ($L^2$ Sinkhorn), Sinkhorn Euler Mccann with the Euler heat approximation in Tab. 3 Sinkhorn with Euler approximation ran out of memory on the Clark and WOT datasets. We see that across all 10 the Geodesic Sinkhorn interpolation with the Mccann interpolant outperforms all other methods, hence showcasing the importance of the heat-geodesic distance and our kernel approximation. We also compare Sinkhron Euler and Geodesic Sinkhorn for different nearest neighbors graphs in Tab. 4 and Tab. 5, where Geodesic Sinkhorn outperforms Sinkhorn Euler on most datasets.

Table 3.

Time series interpolation task comparing mean and standard deviation $(\mu \pm \sigma )$ across 5 seeds the 2-Wasserstein metric averaged across time $(1/T-2){\sum }_{t\in [2\dots T-1]}{W}_2\left({\hat{\mu }}_t,{\mu }_t\right)$ for 10 single-cell timeseries datasets. Sinkhron Euler ran out of memory on two datasets. Lower is better, best performance on each dataset is bold.

Dataset	$L^2$ Sinkhorn	Sinkhron Euler	Geo Sinkhorn (ours)
Cite Donor0	48.545 ± 0.057	46.254 ± 3.192	44.440 ± 0.108
Cite Donor1	48.220 ± 0.055	45.897 ± 3.254	44.165 ± 0.103
Cite Donor2	50.281 ± 0.016	47.773 ± 3.958	45.673 ± 0.092
Cite Donor3	49.339 ± 0.081	46.565 ± 3.553	45.022 ± 0.146
Clark	13.500 ± 0.003	–	13.288 ± 0.008
EB	12.415 ± 0.008	12.298 ± 0.140	12.133 ± 0.011
Multiome Donor0	56.648 ± 0.048	55.373 ± 7.234	53.431 ± 0.077
Multiome Donor1	54.028 ± 0.126	52.396 ± 4.394	50.238 ± 0.022
Multiome Donor2	58.798 ± 0.155	57.182 ± 5.511	55.041 ± 0.058
WOT	8.096 ± 0.003	–	7.397 ± 0.106

Table 4.

Time series interpolation task comparing mean of the 2-Wasserstein for a KNN graph with 50 neighbors. Lower is better, best score is bold.

Dataset	Sinkhorn Euler	Geodesic Sinkhorn
Cite Donor0	48.507	46.850
Cite Donor1	48.207	46.883
Cite Donor2	50.383	49.176
Cite Donor3	49.210	47.646
Clark	13.506	13.378
EB	12.409	12.394
Multiome Donor0	56.676	55.095
Multiome Donor1	54.028	53.952
Multiome Donor2	58.821	57.187
WOT	8.070	8.279

Table 5.

Time series interpolation task comparing mean of the 2-Wasserstein for a KNN graph with 100 neighbors. Lower is better, best score is bold.

Dataset	Sinkhorn Euler	Geodesic Sinkhorn
Cite Donor0	48.573	47.102
Cite Donor1	48.197	47.373
Cite Donor2	50.298	49.517
Cite Donor3	49.290	48.433
Clark	13.500	13.396
EB	12.415	12.416
Multiome Donor0	56.708	55.562
Multiome Donor1	54.063	54.016
Multiome Donor2	58.802	57.732
WOT	8.096	8.253