Accelerating the Sinkhorn Algorithm for Sparse Multi-Marginal Optimal Transport via Fast Fourier Transforms

Abstract

We consider the numerical solution of the discrete multi-marginal optimal transport (MOT) by means of the Sinkhorn algorithm. In general, the Sinkhorn algorithm suffers from the curse of dimensionality with respect to the number of marginals. If the MOT cost function decouples according to a tree or circle, its complexity is linear in the number of marginal measures. In this case, we speed up the convolution with the radial kernel required in the Sinkhorn algorithm via non-uniform fast Fourier methods. Each step of the proposed accelerated Sinkhorn algorithm with a tree-structured cost function has a complexity of $\mathcal{O}\left(KN\right)$ instead of the classical $\mathcal{O}\left(K{N}^2\right)$ for straightforward matrix–vector operations, where K is the number of marginals and each marginal measure is supported on, at most, N points. In the case of a circle-structured cost function, the complexity improves from $\mathcal{O}\left(K{N}^3\right)$ to $\mathcal{O}\left(K{N}^2\right)$. This is confirmed through numerical experiments.

1. Introduction

The optimal transport (OT) problem is an optimization problem that deals with the search for an optimal map (plan) that moves masses between two or more measures at low cost [1,2]. OT appears in a wide range of applications such as image and signal processing [3,4,5,6,7], economics [8,9], finance [10,11], and physics [12,13]. The OT problem was first introduced in 1781 by Monge. His objective was to find a map between two probability measures ${\mu }^1,{\mu }^2$ on ${\mathbb{R}}^d$ that transports ${\mu }^1$ to ${\mu }^2$ with minimal cost, where the cost function describes the cost of transporting a mass between two points in ${\mathbb{R}}^d$. However, such maps do not always exist, so that Kantorovich [14] relaxed the problem in 1942 by looking for a transport plan with two prescribed marginals ${\mu }^1$ and ${\mu }^2$ that minimizes a certain cost functional.

Several authors have generalized the formulation to multi-marginal optimal transport (MOT) [15,16,17], where more than two marginal measures are given. For the given probability measures ${\mu }^k$ on ${\Omega }^k\subset {\mathbb{R}}^d$, $k=1,\cdots ,K$, an optimal transport plan $\pi$ is defined as a solution of the MOT problem

$${\displaystyle \underset{\pi \in \Pi ({\mu }^1,\dots ,{\mu }^k)}{min}}{\int }_{{\Omega }^1\times \cdots \times {\Omega }^K}c(x^1,\dots ,x^K)\mathrm{d}\pi (x^1,\dots ,x^K)$$

where $\Pi ({\mu }^1,\cdots ,{\mu }^K)$ is the convex set of all joint probability measures $\pi$ whose marginals are ${\mu }^k$, and $c:{\Omega }^1\times \cdots \times {\Omega }^K\to \mathbb{R}$ is the cost function.

Since the numerical computation of a transport plan is difficult in general, a regularization term such as the entropy [13,18], Kullback–Leibler divergence [19], general f-divergence [20], or $L^2$-regularization [21,22] can be added to make the problem strictly convex. Different approaches such as the Sinkhorn algorithm [1,13], stochastic gradient descent [23], the Gauss–Seidel method [22], or proximal splitting [24] have been used to iteratively determine a minimizing sequence of the MOT problem.

However, the problem suffers from the curse of dimensionality as the complexity grows exponentially with the number K of marginal measures. One way to circumvent this lies in incorporating additional sparsity assumptions into the model. Polynomial-time algorithms to solve certain sparse MOT problems have been studied in [13,25,26]. We will assume that the cost function decouples according to a graph, where the nodes correspond to the marginals and the cost function is the sum of functions that depend only on two variables which correspond to two marginals connected by an edge of the graph. For example, the circle-structured Euclidean cost function reads as

$$c(x^1,\cdots ,x^K)={\parallel x^1-x^2\parallel }_2^2+\cdots +{\parallel x^{K-1}-x^K\parallel }_2^2+{\parallel x^K-x^1\parallel }_2^2,x^1,\cdots ,x^K\in {\mathbb{R}}^d.$$

MOT problems with graph-structured cost functions with a constant tree-width can be solved with the Sinkhorn algorithm in polynomial time [27]. In [25], polynomial-time algorithms were presented for the MOT problem and its regularized counterpart for the cases of a graph structure and a set-optimization structure, as well as a low-rank and sparsely structured cost function. Another sparsity assumption lies on the transport plan to be thinly spread, which is, e.g., the case for the $L^2$-regularized problem [21,22].

1.1. Our Contributions

In the present paper, we study the discrete, entropy-regularized MOT problem with a tree-structured [3,13] or a circle-structured cost function, where all measures are supported on a finite number of points (atoms) in ${\mathbb{R}}^d$. Then, the computational time of the Sinkhorn algorithm [28,29,30], which iteratively determines a sequence converging to the solution, depends linearly on the number K of input measures. If the numbers of atoms is large, however, the Sinkhorn algorithm still requires considerable computational time and memory, which mainly comes from computing a discrete convolution, i.e., a matrix–vector product with a kernel matrix.

This is significantly improved by Fourier-based fast summation methods [31,32]. The key idea is the approximation of the kernel function using a Fourier series, which enables the application of the non-uniform fast Fourier transform (NFFT). Although such fast summation methods are frequently used in different applications such as electrostatic particle interaction [33], tomography [34], image segmentation [35], and, very recently, also OT with two marginals [36], they have not been utilized for MOT so far. Furthermore, a method for accelerating the Sinkhorn algorithm for Wasserstein barycenters via computing the convolution with a different kernel, namely the heat kernel, was discussed in [37].

Our main contribution is the combination of the fast summation method with the sparsity of the tree- or circle-structured cost function in the MOT problem for accelerating the Sinkhorn algorithm. Each iteration step has a complexity of $\mathcal{O}\left(KN\right)$ for a tree- and $\mathcal{O}\left(K{N}^2\right)$ for a circle-structured cost function, compared to $\mathcal{O}\left(K{N}^2\right)$ and $\mathcal{O}\left(K{N}^3\right)$, respectively, with the straightforward matrix–vector operations, where N is an upper bound of the number of atoms for each of the K marginal measures. Our numerical tests with both tree- and circle-structured cost functions confirm a considerable acceleration, while the accuracy stays almost the same. A different acceleration of the Sinkhorn algorithm via low-rank approximation for tree-structured cost yields the same asymptotic complexity [38]. We note that MOT problems with tree-structured cost functions are used for the computation of Wasserstein barycenters [4], and with circle-structured cost for computing Euler flows [39].

1.2. Outline of the Paper

Section 2 introduces the notation. In Section 3, we focus on the discrete MOT problem with squared Euclidean norm cost functions and the numerical solution of the corresponding entropy-regularized problem, using the Sinkhorn algorithm. We investigate sparsely structured cost functions that decouple according to a tree or circle in Section 4. Then, the complexity of the Sinkhorn algorithm depends only linearly on K. Section 5 describes a fast summation method for further accelerating the Sinkhorn algorithm. Finally, in Section 6, we verify the practical performance by applying it to generalized Euler flows and for finding generalized Wasserstein barycenters. We compare the computational time of the proposed Sinkhorn algorithm based on the NFFT, with the algorithm based on direct matrix multiplication.

2. Notation

Let $K\in \mathbb{N}$ and $\mathit{n}=(n_1,\cdots ,n_k)\in {\mathbb{N}}^K$. We set $\left[K\right]:= 1,\cdots ,K$ and consider K finite sets

$${\Omega }^k=\left\{{\mathit{x}}_{i_k}^k:i_k\in \left[n_k\right]\right\}\subset {\mathbb{R}}^d,k\in \left[K\right],$$

consisting of points ${\mathit{x}}_{i_k}^k\in {\mathbb{R}}^d$, which are called atoms. We set $\Omega := {\Omega }^1\times \cdots \times {\Omega }^K$. Additionally, we define the index set

$$\mathit{I}:= \left\{\mathit{i}=(i_1,\cdots ,i_K):i_k\in \left[n_k\right],k\in \left[K\right]\right\}$$

and the set of K-dimensional matrices (tensors)

$${\mathbb{R}}^{\mathit{i}}:= {\mathbb{R}}^{i_1}\times \cdots \times {\mathbb{R}}^{i_K},\mathit{i}\in \mathit{I}.$$

Let $\mathbb{P}\left({\Omega }^k\right)$ denote the set of probability measures on ${\Omega }^k$. In this paper, we consider K discrete probability measures, also called marginal measures, ${\mu }^k\in \mathbb{P}\left({\Omega }^k\right)$, given by

$${\mu }^k=\sum _{i_k=1}^{n_k}{\mu }_{i_k}^k{\delta }_{{\mathit{x}}_{i_k}^k},k\in \left[K\right],$$

where the probabilities satisfy

$${\mu }_{i_k}^k\ge 0,\sum _{i_k=1}^{n_k}{\mu }_{i_k}^k=1\mathrm{for}\mathrm{all}{i}_k\in \left[n_k\right],k\in \left[K\right],$$

and, for all $A\subset {\Omega }^k$, the Dirac measure is given by

$${\delta }_{x_{i_k}^k}\left(A\right):= \left\{\begin{array}{cc}1\hfill & \mathrm{if}{x}_{i_k}^k\in A,\hfill \\ 0\hfill & \mathrm{otherwise}.\hfill \end{array}\right.$$

For $\mathit{G},\mathit{H}\in {\mathbb{R}}^{\mathit{n}}$, we denote their component-wise (Hadamard) product by

$$\mathit{G}\odot \mathit{H}{\left(G_{\mathit{i}}{H}_{\mathit{i}}\right)}_{\mathit{i}\in \mathit{I}}\in {\mathbb{R}}^{\mathit{n}},$$

and similarly, their component-wise division by ⊘, as well as the Frobenius inner product

$${\langle \mathit{G},\mathit{H}\rangle }_{\mathrm{F}}{\displaystyle \sum _{\mathit{i}\in \mathit{I}}}{G}_{\mathit{i}}{H}_{\mathit{i}}\in \mathbb{R}.$$

The tensor product (Kronecker product) of $\mathit{u},\mathit{v}\in {\mathbb{R}}^m$ is denoted by $\mathit{u}\otimes \mathit{v}\in {\mathbb{R}}^{m\times m}$. Analogues can be defined for tensors of different size.

3. Multi-Marginal Optimal Transport

In the following, we consider the discrete multi-marginal optimal transport (MOT) between K marginal measures ${\mu }^k\in \mathbb{P}\left({\Omega }^k\right)$, $k\in \left[K\right]$. We define the set of admissible transport plans by

$$\Pi ({\mu }^1,\dots ,{\mu }^k):= \Pi \in {\mathbb{R}}_{\ge 0}^{\mathit{n}}:P_k(\Pi )={\mu }^k\mathrm{for}\mathrm{all}k\in \left[K\right],$$

(1)

where the k-th marginal projection is defined as

$$P_k(\Pi ):= {\displaystyle \sum _{\ell \in \left[K\right]\backslash k}}{\displaystyle \sum _{i_{\ell }\in \left[n_{\ell }\right]}}{\Pi }_{i_1,\cdots ,i_K}\in {\mathbb{R}}^{n_k}.$$

(2)

For a cost function $c:\Omega \to {\mathbb{R}}_{\ge 0}$ and the samples ${\mathit{x}}_{\mathit{i}}=\left(x_{i_1}^1,\dots ,x_{i_K}^K\right)$, $\mathit{i}\in \mathit{I}$, we define the respective cost matrix

$$\mathit{C}:= {\left[{\mathit{C}}_{\mathit{i}}\right]}_{\mathit{i}\in \mathit{I}}={\left[c\left({\mathit{x}}_{\mathit{i}}\right)\right]}_{\mathit{i}\in \mathit{I}}={\left[c({\mathit{x}}_{i_1}^1,\cdots ,{\mathit{x}}_{i_K}^K)\right]}_{(i_1,\cdots ,i_K)\in \mathit{I}}\in {\mathbb{R}}_{\ge 0}^{\mathit{n}}.$$

(3)

The discrete MOT problem reads

$${\displaystyle \underset{\Pi \in \Pi ({\mu }^1,\dots ,{\mu }^k)}{min}}{\langle \Pi ,\mathit{C}\rangle }_{\mathrm{F}},$$

(4)

whose solution $\Pi \in {\mathbb{R}}_{\ge 0}^{\mathit{n}}$ is called optimal plan.

Entropy Regularization

As the MOT problem (4) is numerically unfeasible, we consider for $\eta >0$ the entropy-regularized multi-marginal optimal transport (${\mathrm{MOT}}_{\eta }$) problem

$${\displaystyle \underset{\Pi \in \Pi ({\mu }^1,\dots ,{\mu }^K)}{min}}{\langle \Pi ,\mathit{C}\rangle }_{\mathrm{F}}+\eta {\langle \Pi ,log\Pi -{\mathit{1}}_{\mathit{n}}\rangle }_{\mathrm{F}}={\displaystyle \underset{\Pi \in \Pi ({\mu }^1,\dots ,{\mu }^k)}{min}}{\displaystyle \sum _{\mathit{i}\in \mathit{I}}}{\Pi }_{\mathit{i}}{\mathit{C}}_{\mathit{i}}+\eta {\displaystyle \sum _{\mathit{i}\in \mathit{I}}}{\Pi }_{\mathit{i}}\left(log{\Pi }_{\mathit{i}}-1\right),$$

(5)

which is a convex optimization problem. It is possible to numerically deduce the optimal transport plan $\widehat{\Pi }$ of (5) from the solution of the corresponding Lagrangian dual problem. The following theorem is a special case of [3] for a constant entropy function.

Theorem 1.

The Lagrangian dual formulation of the discrete ${\mathit{MOT}}_{\eta }$ problem (5) states

$${\displaystyle \underset{\begin{array}{c}{\varphi }^k\in {\mathbb{R}}_{\ge 0}^{n_k},k\in \left[K\right]\end{array}}{sup}}\mathcal{S}\left({\varphi }^1,\cdots ,{\varphi }^K\right):= {\displaystyle \underset{\begin{array}{c}{\varphi }^k\in {\mathbb{R}}_{\ge 0}^{n_k},k\in \left[K\right]\end{array}}{sup}}\eta {\displaystyle \sum _{k\in \left[K\right]}}{\displaystyle \sum _{j\in \left[n_k\right]}}{\mu }_j^{k}log{\varphi }_j^k-\eta {\displaystyle \sum _{\mathit{i}\in \mathit{I}}}{\mathcal{K}}_{\mathit{i}}{\mathbf{\Phi }}_{\mathit{i}},$$

(6)

where the kernel matrix $\mathcal{K}\in {\mathbb{R}}^{\mathit{n}}$ is defined by

$${\mathcal{K}}_{\mathit{i}}:= exp\left(-\frac{{\mathit{C}}_{\mathit{i}}}{\eta }\right),\mathit{i}\in \mathit{I},$$

and the dual tensor $\mathbf{\Phi }={⨂}_{k=1}^K{\varphi }^k$ by

$${\mathbf{\Phi }}_{\mathit{i}}:= \prod _{k=1}^K{\varphi }_{i_k}^k,\mathit{i}\in \mathit{I}.$$

The functional $\mathcal{S}:{\mathbb{R}}_{\ge 0}^{\mathit{n}}\to \mathbb{R}$ is called the Sinkhorn function.

The solutions of the dual and the primal problem are generally not equal. The solution of (6) is generally a lower bound to the solution of the primal problem (5). Equality holds if the cost function c is lower semi-continuous; i.e., $lim\; infc\left(\mathit{x}\right)\ge c\left({\mathit{x}}_{\mathit{0}}\right)$ as $\mathit{x}\to {\mathit{x}}_{\mathit{0}}$ for every ${\mathit{x}}_{\mathit{0}}\in {\mathbb{R}}^{\mathit{n}}$, or Borel measurable and bounded [40]. This is obviously the case for the squared Euclidean norm cost function

$$c:{\mathbb{R}}^{\mathit{n}}\to [0,\infty ),c\left(\mathit{x}\right)={\displaystyle \sum _{\begin{array}{c}{k}_1,k_2\in \left[K\right]\\ k_1\ne k_2\end{array}}}{\parallel x_{i_{k_1}}^{k_1}-x_{i_{k_2}}^{k_2}\parallel }_2^2$$

that we study here.

Proposition 1

([41]). An optimal plan of the ${\mathit{MOT}}_{\eta }$ problem (5) is given by

$$\widehat{\Pi }=\mathcal{K}\odot \widehat{\mathbf{\Phi }},$$

where $\widehat{\mathbf{\Phi }}={⨂}_{k\in \left[K\right]}{\widehat{\varphi }}^k$ and ${\widehat{\varphi }}^k,$$k\in \left[K\right]$ are the optimal solutions of dual problem (6).

A sequence converging to the optimal dual vectors ${\widehat{\varphi }}^k,$$k\in \left[K\right]$ in (6) can be iteratively determined by the Sinkhorn algorithm [13,42] presented in Algorithm 1, where we note that line 4 is obtained by deriving the Sinkhorn function $\mathcal{S}$ with respect to ${\varphi }^k,$$k\in \left[K\right].$ The complexity of the algorithm mainly comes from the computation of the marginal $P_k(\mathcal{K}\odot \mathbf{\Phi })$, where the projection $P_k$ is defined in (2). In general, the number of operations depends exponentially on K.

Algorithm 1 Sinkhorn iterations for the ${\mathrm{MOT}}_{\eta }$ problem

1: Input: Initial values ${\left({\varphi }^k\right)}^{\left(0\right)}\in {\mathbb{R}}^{n_k},$$k\in \left[K\right],$ regularization parameter $\eta >0,$ threshold $\delta >0$ 2: Set $r\leftarrow 0$ 3: do 4:     for $k=1,\dots ,K$ do 5:         Compute ${\left(\tilde{\mathbf{\Phi }}\right)}_k^{(r+1)}:= {⨂}_{\ell \in [k-1]}{\left({\varphi }^{\ell }\right)}^{(r+1)}\otimes {⨂}_{\ell \in \left[K\right]\backslash [k-1]}{\left({\varphi }^{\ell }\right)}^{\left(r\right)}$ 6:         Compute dual vectors $$\begin{array}{cc}\hfill & {\left({\varphi }^k\right)}^{(r+1)}:= \left({\mu }^k\odot {\left({\varphi }^k\right)}^{\left(r\right)}\right)\oslash P_k\left(\mathcal{K}\odot {\left(\tilde{\mathbf{\Phi }}\right)}_k^{(r+1)}\right)\hfill \end{array}$$ 7:         Increment $r\leftarrow r+1$ 8:     end for 9: while$\left|\mathcal{S}\left({\left({\varphi }^1\right)}^{\left(r\right)},\cdots ,{\left({\varphi }^K\right)}^{\left(r\right)}\right)-\mathcal{S}\left({\left({\varphi }^1\right)}^{(r-1)},\cdots ,{\left({\varphi }^K\right)}^{(r-1)}\right)\right|\ge \delta$ 10: return optimal plan $\widehat{\Pi }=\mathcal{K}\odot \widehat{\mathbf{\Phi }}$ and $\widehat{\mathbf{\Phi }}={⨂}_{k\in \left[K\right]}{\left({\varphi }^k\right)}^{(r+1)}$

4. Sparse Cost Functions

In this section, we take a look at sparsely structured cost functions, for which the Sinkhorn algorithm becomes much faster and we overcome the curse of dimensionality. Let $G=\left(\mathcal{V},\mathcal{E}\right)$ be an undirected graph with vertices $\mathcal{V}$ and edges $\mathcal{E}$. We say that the ${\mathrm{MOT}}_{\eta }$ problem has the graph G structure if $\mathcal{V}=\left[K\right]$ and the cost matrix (3) decouples according to

$${\mathit{C}}_{\mathit{i}}={\displaystyle \sum _{\{k_1,k_2\}\in \mathcal{E}}}\parallel x_{i_{k_1}}^{k_1}-x_{i_{k_2}}^{k_2}{\parallel }_2^2,\mathit{i}=(i_1,\cdots ,i_K)\in \mathit{I}.$$

This implies that the kernel $\mathcal{K}\in {\mathbb{R}}_{\ge 0}^{\mathit{n}}$ satisfies

$${\mathcal{K}}_{\mathit{i}}=exp\left(-\frac{{\mathit{C}}_{\mathit{i}}}{\eta }\right)={\displaystyle \prod _{\{k_1,k_2\}\in \mathcal{E}}}{\mathcal{K}}_{i_{k_1},i_{k_2}}^{(k_1,k_2)},\mathit{i}\in \mathit{I},$$

(7)

where the kernel matrix ${\mathcal{K}}^{(k_1,k_2)}\in {\mathbb{R}}_{\ge 0}^{n_{k_1}\times n_{k_2}}$ for $\{k_1,k_2\}\in \mathcal{E}$ is given by

$${\mathcal{K}}_{i_{k_1},i_{k_2}}^{(k_1,k_2)}:= exp\left(-\frac{1}{\eta }\parallel {\mathit{x}}_{i_{k_1}}^{k_1}-{\mathit{x}}_{i_{k_2}}^{k_2}{\parallel }_2^2\right).$$

(8)

We use the indices $k\in \mathcal{V}=\left[K\right]$ to identify the marginal measures ${\mu }^k$ in the rest of the paper.

The discrete, dual formulation (6) of the ${\mathrm{MOT}}_{\eta }$ problem has the same form independently of the structure of the graph G; only the marginals $P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)$ differ. If the graph G is complete, i.e., each of the two vertices in $\mathcal{V}$ are connected with an edge, then the computational complexity of the Sinkhorn Algorithm 1 depends exponentially on the number K of marginal measures (vertices). For larger values of K, it is practically impossible to numerically compute an optimal plan for the ${\mathrm{MOT}}_{\eta }$ problem. We consider two sparsity assumptions of the ${\mathrm{MOT}}_{\eta }$ problem, each of them yielding that the Sinkhorn algorithm has a linear complexity in the number K of the nodes. It was shown in [25] that ${\mathrm{MOT}}_{\eta }$ problems with graphically structured cost functions of constant tree-width can be implemented in polynomial time. This is the case for the tree and circle, whose tree-widths are 1 and 2, respectively. In the following, we give an explicit scheme to efficiently compute the Sinkhorn iterations for the tree and circle structures.

4.1. Tree Structure

We consider the ${\mathrm{MOT}}_{\eta }$ problem with the structure of a tree $\left(\mathcal{V},\mathcal{E}\right)$, which is a connected and circle-free graph with $\left|\mathcal{E}\right|=\left|\mathcal{V}\right|-1.$ We define the (non-empty) neighbour set ${\mathcal{N}}_k$ of $k\in \mathcal{V}$ as the set of all nodes $\ell \in \mathcal{V}$, such that $\{k,\ell \}\in \mathcal{E}.$ Furthermore, we denote by $\mathcal{L}\{k\in \mathcal{V}:|{\mathcal{N}}_k|=1\}$ the set of all leaves of the tree.

We call $1\in \mathcal{V}$ the root of the tree. For every $k\in \mathcal{V}$, there is a unique path between k and the root. For $k\in \mathcal{V}\backslash 1$, we define the parent $p\left(k\right)$ as the node in ${\mathcal{N}}_k$ such that $p\left(k\right)$ lies in the path between k and the root 1. The root has no parent. Without loss of generality, we assume that $p\left(k\right)<k$ holds for all $k\in \mathcal{V}\backslash 1$. We define the set of children ${\mathcal{C}}_k\{\ell \in {\mathcal{N}}_k:\ell >k\}$. Then, we can derive a recursive formula for the k-th marginal of the tensor $\mathcal{K}\odot \mathbf{\Phi }\in {\mathbb{R}}^{\mathit{n}}.$

Theorem 2.

Let $(\mathcal{V},\mathcal{E})$ be a tree with leaves $\mathcal{L}$, and let c be the ${\mathit{MOT}}_{\eta }$ cost function associated. Furthermore, let $k\in \mathcal{V}$ be an arbitrary node. Then, the k-th marginal of the transport plan $\mathcal{K}\odot \mathbf{\Phi }$ is given by

$$P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)={\varphi }^k\odot {\displaystyle \underset{\ell \in {\mathcal{N}}_k}{⨀}}{\alpha }^{(k,\ell )},$$

where the vectors ${\alpha }^{(k,\ell )}\in {\mathbb{R}}^{n_k}$ are recursively defined as

$${\alpha }^{(k,\ell )}=\left\{\begin{array}{cc}{\mathcal{K}}^{(k,\ell )}{\varphi }^{\ell }\hfill & \mathit{if}\ell \in \mathcal{L},\hfill \\ {\mathcal{K}}^{(k,\ell )}\left({\varphi }^{\ell }\odot {\displaystyle \underset{t\in {\mathcal{N}}_{\ell }\backslash k}{⨀}}{\alpha }^{(\ell ,t)}\right)\hfill & \mathit{otherwise}.\hfill \end{array}\right.$$

A proof of Theorem 2 can be found in [13] (Theorem 3.2). The main idea is to split the rooted tree at the node k into $|{\mathcal{N}}_k|$ subtrees. Therefore, the kernel matrix holds

$${\mathcal{K}}_{\mathit{i}}={\displaystyle \prod _{\{k_1,k_2\}\in \mathcal{E}}}{\mathcal{K}}_{i_{k_1},i_{k_2}}^{(k_1,k_2)}={\displaystyle \prod _{\ell \in {\mathcal{N}}_k}}{\displaystyle \prod _{t\in {\mathcal{D}}_{\ell }}}{\mathcal{K}}_{i_{p\left(t\right)},i_t}^{\left(p\right(t),t)},\mathit{i}\in \mathit{I},$$

where ${\mathcal{D}}_{\ell }$ is the set of descendants of $\ell ,$ i.e., the nodes $t\in \mathcal{V}$ such that ℓ lies in the path between k and $t.$ Inserting ${\mathcal{K}}_{\mathit{i}}$ into $P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)$ and recalling the definition of $P_k$ in (2) yields the result.

In order to efficiently perform the Sinkhorn algorithm, we compute iteratively for $\ell =K,\cdots ,2$ the vectors ${\beta }_{\ell }:= {\alpha }^{\left(p\right(\ell ),\ell )}\in {\mathbb{R}}^{n_{p(\ell )}}.$ From Theorem 2, we obtain

$$\begin{array}{cc}\hfill & {\beta }_{\ell }=\left\{\begin{array}{cc}{\mathcal{K}}^{\left(p\right(\ell ),\ell )}{\varphi }^{\ell }\hfill & \mathrm{if}\ell \in \mathcal{L},\hfill \\ {\mathcal{K}}^{\left(p\right(\ell ),\ell )}\left({\varphi }^{\ell }\odot {\displaystyle \underset{t\in {\mathcal{C}}_{\ell }}{⨀}}{\beta }_t\right)\hfill & \mathrm{otherwise}.\hfill \end{array}\right.\hfill \end{array}$$

Since we assumed that $p\left(k\right)<k$, the computation of ${\beta }_{\ell }$ requires only ${\beta }_t$ for $t>\ell$. Similarly, the vectors ${\gamma }_{\ell }:= {\alpha }^{(\ell ,p(\ell \left)\right)}\in {\mathbb{R}}^{n_{\ell }}$ for $\ell >1$ and ${\gamma }_1:= {\mathit{1}}_{n_1}$ can be iteratively computed by    

$${\gamma }_{\ell }={\left({\mathcal{K}}^{\left(p\right(\ell ),\ell )}\right)}^{⊺}\left({\varphi }^{p(\ell )}\odot {\gamma }_{p(\ell )}\odot {\displaystyle \underset{t\in {\mathcal{C}}_{p(\ell )}\backslash \ell }{⨀}}{\beta }_t\right),\ell =2,\cdots ,K,$$

where the vectors ${\beta }_t\in {\mathbb{R}}^{n_{p(\ell )}}$ are assumed to be known from above. Then, the marginals are

$$P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)={\varphi }^k\odot {\gamma }_k\odot {\beta }^{{\odot }_{C_k}},k\in \left[K\right],$$

where for any subset $\mathcal{U}\subseteq {\mathcal{C}}_k,$ we define

$${\mathbb{R}}^{n_k}\ni {\beta }^{{\odot }_{\mathcal{U}}}:= \left\{\begin{array}{cc}{\mathit{1}}_{n_k}\hfill & \mathrm{if}\mathcal{U}=\varnothing ,\hfill \\ {\displaystyle {⨀}_{t\in \mathcal{U}}}{\beta }_t\hfill & \mathrm{otherwise}.\hfill \end{array}\right.$$

The resulting procedure is summarized in Algorithm 2.

Algorithm 2 Sinkhorn algorithm for tree structure

1: Input: Tree $T(\mathcal{V},\mathcal{E})$ with leaves $\mathcal{L}$ and root $1\in \mathcal{V}=\left[K\right]$, initialization ${\left({\varphi }^k\right)}^{\left(0\right)}$, $k\in \left[K\right],$ parameters $\eta ,\delta >0$ 2: Initialize $r\leftarrow 0$ 3: for $k=K,\cdots ,2$  do 4:     ${\displaystyle {\beta }_k^{\left(0\right)}:= \left\{\begin{array}{ll}{\mathcal{K}}^{\left(p\right(k),k)}{\left({\varphi }^k\right)}^{\left(0\right)}& \mathrm{if}k\in \mathcal{L},\\ {\mathcal{K}}^{\left(p\right(k),k)}\left({\left({\varphi }^k\right)}^{\left(0\right)}\odot {\left({\beta }^{{\odot }_{{\mathcal{C}}_k}}\right)}^{\left(0\right)}\right)& \mathrm{otherwise}\end{array}\right.}$ 5: end for 6: do 7:     for $k=1,\cdots ,K$ do 8:         ${\displaystyle {\gamma }_k^{\left(r\right)}:= \left\{\begin{array}{ll}\mathit{1}& \mathrm{if}k=1,\\ {\left({\mathcal{K}}^{\left(p\right(k),k)}\right)}^{⊺}\left({\left({\varphi }^{p\left(k\right)}\right)}^{(r+1)}\odot {\gamma }_{p\left(k\right)}^{\left(r\right)}\odot {\left({\beta }^{{\odot }_{{\mathcal{C}}_{p\left(k\right)}\backslash k}}\right)}^{\left(r\right)}\right)& \mathrm{otherwise}\end{array}\right.}$ 9:         Compute the dual vector ${\left({\varphi }^k\right)}^{(r+1)}:= {\mu }^k\oslash \left({\gamma }_k^{\left(r\right)}\odot {\left({\beta }^{{\odot }_{{\mathcal{C}}_k}}\right)}^{\left(r\right)}\right)$ 10:     end for 11:     for $k=K,\cdots ,2$ do 12:         Compute ${\beta }_k^{(r+1)}$ according to step 4. 13:     end for 14:     Set ${\displaystyle {\mathcal{S}}^{\left(r\right)}:= \eta \left(\sum _{k=1}^K{\left({\mu }^k\right)}^{⊺}log{\left({\varphi }^k\right)}^{\left(r\right)}-{\left({\mu }^1\oslash {\left({\varphi }^1\right)}^{(r+1)}\right)}^{⊺}{\left({\varphi }^1\right)}^{\left(r\right)}\right)}$ 15:     Increment $r\leftarrow r+1$ 16: while $|{\mathcal{S}}^{\left(r\right)}-{\mathcal{S}}^{(r-1)}|\ge \delta$ 17: return optimal plan $\widehat{\Pi }=\mathcal{K}\odot \widehat{\mathbf{\Phi }}$, where $\widehat{\mathbf{\Phi }}={⨂}_{k\in \left[K\right]}{\left({\varphi }^k\right)}^{\left(r\right)}$

4.2. Circle Structure

We consider the ${\mathrm{MOT}}_{\eta }$ problem where the graph $(\mathcal{V},\mathcal{E})$ is a circle. We assume for each $k\in \mathcal{V}=\left[K\right]$ that $\{k,k+1\}\in \mathcal{E}$ is an edge of the circle, where we set $k+1=1$ if $k=K$ and $k-1=K,$ if $k=1.$ Thus, we can define the distance between two nodes $k_1,k_2\in \mathcal{V}$ as

$$d(k_1,k_2):= \left\{\begin{array}{cc}{k}_2-k_1\hfill & \mathrm{if}{k}_2\ge k_1,\hfill \\ K-k_1+k_2\hfill & \mathrm{otherwise}.\hfill \end{array}\right.$$

Theorem 3.

Let $(\mathcal{V},\mathcal{E})$ be a circle and c be the MOT cost function associated. Furthermore, let $k\in \mathcal{V}$ be an arbitrary node. The k-th marginal of the transport plan $\mathcal{K}\odot \mathbf{\Phi }$ is given by

$$P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)=\left(\left({\varphi }^k\odot {\mathcal{K}}^{(k,k+1)}\right)\odot {\left({\varphi }^{k+1}\odot {\alpha }^{(k+1,k)}\right)}^{⊺}\right){\mathit{1}}_{k+1},$$

where the matrices ${\alpha }^{(\ell ,t)}\in {\mathbb{R}}^{n_{\ell }\times n_t}$, $\ell ,t\in \mathcal{V}$ recursively satisfy

$${\alpha }^{(\ell ,t)}=\left\{\begin{array}{cc}{\mathcal{K}}^{(\ell ,t)}\hfill & \mathrm{if}\mathrm{d}(\ell ,t)=1,\hfill \\ {\mathcal{K}}^{(\ell ,\ell +1)}\left({\varphi }^{\ell +1}\odot {\alpha }^{(\ell +1,t)}\right)\hfill & \mathrm{otherwise},\hfill \end{array}\right.$$

(9)

and we set $k+1=1$ if $k=K$ and $k-1=N,$ if $k=1.$

Proof. 

The kernel $\mathcal{K}$ can be decomposed as (7). Let $k\in \left[K\right]$ and $i_k\in \left[n_k\right]$. It holds then, that

$$\begin{array}{ll}{\left[P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)\right]}_{i_k}& ={\displaystyle \sum _{\ell \in \mathcal{V}\backslash k}}{\displaystyle \sum _{i_{\ell }\in \left[n_{\ell }\right]}}{\mathcal{K}}_{\mathit{i}}\odot {\mathbf{\Phi }}_{\mathit{i}}={\varphi }_{i_k}^k{\displaystyle \sum _{\ell \in \mathcal{V}\backslash k}}{\displaystyle \sum _{i_{\ell }\in \left[n_{\ell }\right]}}{\displaystyle \prod _{\{k_1,k_2\}\in \mathcal{E}}}{\mathcal{K}}_{i_{k_1},i_{k_2}}^{(k_1,k_2)}{\displaystyle \prod _{j\in \mathcal{V}\backslash k}}{\varphi }_{i_j}^j\\ & ={\varphi }_{i_k}^k{\displaystyle \sum _{i_1}}{\varphi }_{i_1}^1{\displaystyle \sum _{i_2}}{\mathcal{K}}_{i_1,i_2}^{(1,2)}{\varphi }_{i_2}^2\cdots {\displaystyle \sum _{i_{k-2}}}{\mathcal{K}}_{i_{k-3},i_{k-2}}^{(k-3,k-2)}{\varphi }_{i_{k-2}}^{k-2}{\displaystyle \sum _{i_{k-1}}}{\mathcal{K}}_{i_{k-2},i_{k-1}}^{(k-2,k-1)}{\varphi }_{i_{k-1}}^{k-1}{\mathcal{K}}_{i_{k-1},i_k}^{(k-1,k)}\\ & {\displaystyle \sum _{i_{k+1}}}{\mathcal{K}}_{i_k,i_{k+1}}^{(k,k+1)}{\varphi }_{i_{k+1}}^{k+1}\cdots {\displaystyle \sum _{i_K}}{\mathcal{K}}_{i_{K-1},i_K}^{(K-1,K)}{\varphi }_{i_K}^K{\mathcal{K}}_{i_K,i_1}^{(K,1)}\\ & ={\varphi }_{i_k}^k{\displaystyle \sum _{i_{k+1}}}{\mathcal{K}}_{i_k,i_{k+1}}^{(k,k+1)}{\varphi }_{i_{k+1}}^{k+1}\cdots {\displaystyle \sum _{i_K}}{\mathcal{K}}_{i_{K-1},i_K}^{(K-1,K)}{\varphi }_{i_K}^K{\displaystyle \sum _{i_1}}{\mathcal{K}}_{i_K,i_1}^{(K,1)}{\varphi }_{i_1}^1{\displaystyle \sum _{i_2}}{\mathcal{K}}_{i_1,i_2}^{(1,2)}{\varphi }_{i_2}^2\cdots \\ & {\displaystyle \sum _{i_{k-2}}}{\mathcal{K}}_{i_{k-3},i_{k+2}}^{(k-3,k-2)}{\varphi }_{i_{k-2}}^{k-2}{\displaystyle \sum _{i_{k-1}}}{\mathcal{K}}_{i_{k-2},i_{k-1}}^{(k-2,k-1)}{\varphi }_{i_{k-1}}^{k-1}{\mathcal{K}}_{i_{k-1},i_k}^{(k-1,k)}.\end{array}$$

Setting ${\psi }^{\ell }:= {\varphi }^{(k+\ell -1)\mathrm{mod}K}$ and ${\tilde{\mathcal{K}}}^{(j,\ell +1)}: = {\mathcal{K}}^{\left((k+\ell -1)\mathrm{mod}K,(k+\ell )\mathrm{mod}K\right)}$ for every $\ell \in \left[K\right]$, we obtain

$$\begin{array}{ll}{\left[P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)\right]}_{i_k}& ={\varphi }_{i_k}^k{\displaystyle \sum _{j_2}}{\tilde{\mathcal{K}}}_{i_k,j_2}^{(1,2)}{\psi }_{j_2}^2\cdots {\displaystyle \sum _{j_{K-1}}}{\tilde{\mathcal{K}}}_{j_{K-2},j_{K-1}}^{(K-2,K-1)}{\psi }_{j_{K-1}}^{K-1}{\displaystyle \sum _{j_K}}{\tilde{\mathcal{K}}}_{j_{K-1},j_K}^{(K-1,K)}{\psi }_{j_K}^K{\tilde{\mathcal{K}}}_{j_K,i_k}^{(K,1)}\\ & ={\left[P_1\left(\tilde{\mathcal{K}}\odot \tilde{\mathbf{\Phi }}\right)\right]}_{i_k}.\end{array}$$

We define

$${\tilde{\alpha }}^{(K,1)}:= {\tilde{\mathcal{K}}}^{(K,1)}={\alpha }^{(k-1,k)}\in {\mathbb{R}}^{n_{k-1}\times n_k}$$

and recursively for all $\ell =K-1,\cdots ,1$, the matrix ${\tilde{\alpha }}^{(\ell ,1)}\in {\mathbb{R}}^{\left((k+\ell -1)\mathrm{mod}K\right)\times n_k}$ by

$${\tilde{\alpha }}_{j_l,i_k}^{(\ell ,1)}:= {\left[{\tilde{\mathcal{K}}}^{(\ell ,\ell +1)}\left({\psi }^{\ell +1}\odot {\tilde{\alpha }}^{(\ell +1,1)}\right)\right]}_{j_l,i_k}={\displaystyle \sum _{j_{\ell +1}}}{\tilde{\mathcal{K}}}_{j_{\ell },j_{\ell +1}}^{(\ell ,\ell +1)}{\psi }_{j_{\ell +1}}^{\ell +1}{\tilde{\alpha }}_{j_{\ell +1},i_k}^{(\ell +1,1)}.$$

Inserting this into the marginal yields

$$\begin{array}{ll}{\left[P_1\left(\tilde{\mathcal{K}}\odot \tilde{\mathbf{\Phi }}\right)\right]}_{i_k}& ={\varphi }_{i_k}^k{\displaystyle \sum _{j_2}}{\tilde{\mathcal{K}}}_{i_k,j_2}^{(1,2)}{\psi }_{j_2}^2\cdots {\displaystyle \sum _{j_{K-2}}}{\tilde{\mathcal{K}}}_{j_{K-2},j_{K-1}}^{(K-2,K-1)}{\psi }_{j_{K-1}}^{K-1}{\displaystyle \sum _{j_{K-1}}}{\tilde{\mathcal{K}}}_{j_{K-2},j_{K-1}}^{(K-2,K-1)}{\psi }_{j_{K-1}}^{K-1}{\tilde{\alpha }}_{j_{K-1},i_k}^{(K-1,1)}\\ & ={\psi }_{i_k}^1{\displaystyle \sum _{j_2}}{\tilde{\mathcal{K}}}_{i_k,j_2}^{(1,2)}{\psi }_{j_2}^2{\tilde{\alpha }}_{j_2,i_k}^{(2,1)}={\psi }_{i_k}^1{[\tilde{\mathcal{K}}}_{i_k,j_2}^{(1,2)}{]}_{j_2\in \left[n_{k+1}\right]}^{⊺}\left({\psi }^2\odot {\left[{\tilde{\alpha }}_{j_2,i_k}^{(2,1)}\right]}_{j_2\in \left[n_{k+1}\right]}\right).\end{array}$$

Henceforth,

$$P_1\left(\tilde{\mathcal{K}}\odot \tilde{\mathbf{\Phi }}\right)=\left(\left({\psi }^1\odot {\tilde{\mathcal{K}}}^{(1,2)}\right)\odot {\left({\psi }^2\odot {\tilde{\alpha }}^{(2,1)}\right)}^{⊺}\right)\mathit{1}.$$

Finally, defining ${\tilde{\alpha }}^{(\ell ,1)}={\alpha }^{\left(\right(k+\ell )\mathrm{mod}K,k)}$ yields the assumption.    □

To efficiently compute the marginal optimal transport as for the tree structure, we choose $k=1$ as the starting point and decompose the marginal into two matrices, which can be computed recursively as follows. For matrices $\mathit{G},\mathit{H}\in {\mathbb{R}}^{n\times m},$ we define the inner product with respect to the second dimension by

$$\langle \mathit{G},\mathit{H}\rangle := {\left(\sum _{j=1}^m{G}_{ij}{H}_{ij}\right)}_{i\in \left[n\right]}\in {\mathbb{R}}^n.$$

Theorem 4.

Under the assumptions of Theorem 3, we have

$$P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)=\left\{\begin{array}{cc}{\varphi }^1\odot 〈{\mathcal{K}}^{(1,2)},{\left({\varphi }^2\odot {\alpha }^{(2,1)}\right)}^{⊺}〉,\hfill & k=1,\hfill \\ {\varphi }^k\odot 〈{\alpha }^{(k,1)},{\left({\varphi }^1\odot {\lambda }^{(1,k)}\right)}^{⊺}〉,\hfill & k=2,\cdots ,K,\hfill \end{array}\right.$$

where ${\alpha }^{(k,1)}$ is given in (9) and

$${\lambda }^{(1,k)}:= \left\{\begin{array}{cc}{\mathcal{K}}^{(1,2)},\hfill & k=2,\hfill \\ {\lambda }^{(1,k-1)}\left({\varphi }^{(k-1)}\odot {\mathcal{K}}^{(k-1,k)}\right),\hfill & k=3,\cdots ,K.\hfill \end{array}\right.$$

(10)

Proof. 

Let $k\in \left[K\right]$ and $i_k\in \left[n_k\right]$. For $k=1$, the assertion follows from Theorem 3. For $k>2$, we have

$$\begin{array}{cc}\hfill {\left[P_k\left(\mathcal{K}\odot \mathbf{\Phi }\right)\right]}_{i_k}& ={\varphi }_{i_k}^k{\displaystyle \sum _{i_1}}{\varphi }_{i_1}^1{\displaystyle \sum _{i_2}}{\mathcal{K}}_{i_1,i_2}^{(1,2)}{\varphi }_{i_2}^2\cdots {\displaystyle \sum _{i_{k-1}}}{\mathcal{K}}_{i_{k-2},i_{k-1}}^{(k-2,k-1)}{\varphi }_{i_{k-1}}^{k-1}{\mathcal{K}}_{i_{k-1},i_k}^{(k-1,k)}\hfill \end{array}$$

$$\begin{array}{cc}\hfill & {\displaystyle \sum _{i_{k+1}}}{\mathcal{K}}_{i_k,i_{k+1}}^{(k,k+1)}{\varphi }_{i_{k+1}}^{k+1}\cdots {\displaystyle \sum _{i_K}}{\mathcal{K}}_{i_{K-1},i_K}^{(K-1,K)}{\varphi }_{i_K}^K{\mathcal{K}}_{i_K,i_1}^{(K,1)}\hfill \end{array}$$

$$\begin{array}{cc}\hfill & ={\varphi }_{i_k}^k{\displaystyle \sum _{i_1}}{\varphi }_{i_1}^1\underset{\mathrm{I}:= }{\underbrace{\sum _{i_2}{\mathcal{K}}_{i_1,i_2}^{(1,2)}{\varphi }_{i_2}^2\cdots \sum _{i_{k-1}}{\mathcal{K}}_{i_{k-2},i_{k-1}}^{(k-2,k-1)}{\varphi }_{i_{k-1}}^{k-1}{\mathcal{K}}_{i_{k-1},i_k}^{(k-1,k)}}}{\alpha }_{i_k,i_1}^{(k,1)}.\hfill \end{array}$$

Furthermore, the term $\mathrm{I}$ can be rewritten as

$$\begin{array}{c}\hfill \mathrm{I}={\displaystyle \sum _{i_{k-1}}}{\mathcal{K}}_{i_{k-1},i_k}^{(k-1,k)}{\varphi }_{i_{k-1}}^{k-1}{\displaystyle \sum _{i_{k-2}}}{\mathcal{K}}_{i_{k-2},i_{k-1}}^{(k-2,k-1)}{\varphi }_{i_{k-2}}^{k-2}\cdots {\displaystyle \sum _{i_2}}{\mathcal{K}}_{i_1,i_2}^{(1,2)}{\varphi }_{i_2}^2{\mathcal{K}}_{i_2,i_3}^{(2,3)}\end{array}$$

By (10), we have

$${\lambda }_{i_1,i_3}^{(1,3)}={\displaystyle \sum _{i_2}}{\mathcal{K}}_{i_1,i_2}^{(1,2)}{\varphi }_{i_2}^2{\mathcal{K}}_{i_2,i_3}^{(2,3)},i_1\in \left[n_1\right],i_3\in \left[n_3\right].$$

This implies that

$$\begin{array}{ll}\mathrm{I}& ={\displaystyle \sum _{i_{k-1}}}{\mathcal{K}}_{i_{k-1},i_k}^{(k-1,k)}{\varphi }_{i_{k-1}}^{k-1}{\displaystyle \sum _{i_{k-2}}}{\mathcal{K}}_{i_{k-2},i_{k-1}}^{(k-2,k-1)}{\varphi }_{i_{k-2}}^{k-2}\cdots {\displaystyle \sum _{i_3}}{\mathcal{K}}_{i_3,i_4}^{(3,4)}{\varphi }_{i_3}^3{\lambda }_{i_1,i_3}^{(1,3)},\\ & =\cdots ={\displaystyle \sum _{i_{k-1}}}{\mathcal{K}}_{i_{k-1},i_k}^{(k-1,k)}{\varphi }_{i_{k-1}}^{k-1}{\lambda }_{i_1,i_{k-1}}^{(1,k-1)}={\left[{\lambda }^{(1,k-1)}\left({\varphi }^{k-1}\odot {\mathcal{K}}^{(k-1,k)}\right)\right]}_{i_1,i_k}={\lambda }_{i_1,i_k}^{(1,k)}.\end{array}$$

Hence, the hypothesis is true for all $k>2$. The case $k=2$ can be proven using the same procedure.    □

In order to efficiently compute a maximizing sequence of the dual ${\mathrm{MOT}}_{\eta }$ problem (12), we set for $k\in \left[K\right]$ the dual matrices

$$\begin{array}{c}\hfill {\beta }_k:= {\alpha }^{(k,1)}\in {\mathbb{R}}^{n_k\times n_1},\end{array}$$

$$\begin{array}{c}\hfill {\gamma }_k:= {\lambda }^{(k,1)}\in {\mathbb{R}}^{n_1\times n_k},\end{array}$$

given in Theorems 3 and 4, respectively. The method is shown in Algorithm 3.

Algorithm 3 Sinkhorn algorithm for circle structure

1: Input: Initialization ${\left({\varphi }^k\right)}^{\left(0\right)},k\in \left[K\right],$ parameters $\eta ,\delta >0$ 2: Initialize $r\leftarrow 0$ 3: for $k=K,\cdots ,2$do 4:     ${\displaystyle {\left({\beta }_k\right)}^{\left(0\right)}:= \left\{\begin{array}{ll}{\mathcal{K}}^{(K,1)}& \mathrm{if}k=K,\\ {\mathcal{K}}^{(k,k+1)}\left({\left({\varphi }^{k+1}\right)}^{\left(0\right)}\odot {\left({\beta }_{k+1}\right)}^{\left(0\right)}\right)& \mathrm{otherwise}\end{array}\right.}$ 5: end for 6: do 7:     ${\displaystyle {\left({\varphi }^1\right)}^{(r+1)}:= {\mu }^1\oslash 〈{\mathcal{K}}^{(1,2)},{\left({\left({\varphi }^2\right)}^{\left(r\right)}\odot {\left({\beta }_2\right)}^{\left(r\right)}\right)}^{⊺}〉}$ 8:     for $k=2,\cdots ,K$ do 9:         ${\displaystyle {\left({\gamma }_k\right)}^{\left(r\right)}:= \left\{\begin{array}{ll}{\mathcal{K}}^{(1,2)}& \mathrm{if}k=2,\\ {\left({\gamma }_{k-1}\right)}^{\left(r\right)}\left({\left({\varphi }^{k-1}\right)}^{(r+1)}\odot {\mathcal{K}}^{(k-1,k)}\right)& \mathrm{otherwise}\end{array}\right.}$ 10:         Compute dual vector ${\displaystyle {\left({\varphi }^k\right)}^{(r+1)}:= {\mu }^k\oslash 〈{\left({\beta }_k\right)}^{\left(r\right)},{\left({\left({\varphi }^1\right)}^{(r+1)}\odot {\left({\gamma }_k\right)}^{\left(r\right)}\right)}^{⊺}〉}$ 11:     end for 12:     for $k=K,\cdots ,2$ do 13:         Compute ${\left({\beta }_k\right)}^{(r+1)}$ according to step 4 14:     end for 15:     Compute ${\displaystyle {\mathcal{S}}^{\left(r\right)}:= \eta \left(\sum _{k=1}^K{\left({\mu }^k\right)}^{⊺}log{\left({\varphi }^k\right)}^{\left(r\right)}-{\left({\mu }^1\oslash {\left({\varphi }^1\right)}^{(r+1)}\right)}^{⊺}{\left({\varphi }^1\right)}^{\left(r\right)}\right)}$ 16:     Increment $r\leftarrow r+1$ 17: while $|S^{\left(r\right)}-S^{(r-1)}|\ge \delta$ 18: return Optimal plan $\widehat{\Pi }=\mathcal{K}\odot \widehat{\mathbf{\Phi }}$, where $\widehat{\mathbf{\Phi }}={⨂}_{k\in \left[K\right]}{\left({\varphi }^k\right)}^{(r+1)}$

The tree-structured and circle-structured ${\mathrm{MOT}}_{\eta }$ problems both have a sparse cost function that considerably improves the computational complexity of the Sinkhorn algorithm. In each iteration step, Algorithm 2 requires only $2(K-1)$ matrix–vector products, which have a complexity of $\mathcal{O}\left(N^2\right)$, where $N:= {\parallel \mathit{n}\parallel }_{\infty }$, and Algorithm 3 requires $2(K-1)$ matrix–matrix products, which have a complexity of $\mathcal{O}\left(N^3\right)$. This can be considerably improved by employing fast Fourier techniques, as we will see in the next section.

5. Non-Uniform Discrete Fourier Transforms

The main computational cost of the Sinkhorn algorithm comes from the matrix–vector product with the kernel matrix (8). Let $k,\ell \in \left[K\right]$ and $\alpha \in {\mathbb{R}}^{n_{\ell }}$. We briefly describe a fast summation method for the computation of $\beta ={\mathcal{K}}^{(k,\ell )}\alpha$, i.e.,

$${\beta }_{i_k}=\sum _{i_{\ell }=1}^{n_{\ell }}{\alpha }_{i_{\ell }}exp\left(-\frac{1}{\eta }{\parallel x_{i_k}^k-x_{i_{\ell }}^{\ell }\parallel }_2^2\right),i_k\in \left[n_k\right].$$

(11)

We refer to [32] for a detailed derivation and error estimates. The main idea is to approximate the kernel function

$$\kappa \left(x\right):= exp\left(-\frac{1}{\eta }{x}^2\right),x\in \mathbb{R},$$

(12)

using a Fourier series. In order to ensure the fast convergence of the Fourier series, we extend $\kappa$ to a periodic function of certain smoothness $p\in \mathbb{N}$. Let ${\epsilon }_{\mathrm{B}}>0$ and $\tau >{\epsilon }_{\mathrm{B}}+max\left\{\parallel x_{i_k}^k-x_{i_{\ell }}^{\ell }{\parallel }_2^2:i_k\in \left[n_k\right],i_{\ell }\in \left[n_{\ell }\right]\right\}$. For $x\in \mathbb{R}$, we define the regularized kernel

$${\kappa }_{\mathrm{R}}\left(x\right):= \left\{\begin{array}{cc}\kappa \left(x\right),\hfill & \left|x\right|\le \tau -{\epsilon }_{\mathrm{B}},\hfill \\ {\kappa }_{\mathrm{B}}\left(x\right),\hfill & \tau -{\epsilon }_{\mathrm{B}}<\left|x\right|\le \tau ,\hfill \\ {\kappa }_{\mathrm{B}}\left(\tau \right),\hfill & \left|x\right|>\tau ,\hfill \end{array}\right.$$

(13)

where ${\kappa }_{\mathrm{B}}$ is a polynomial of degree p that fulfills the Herminte interpolation conditions ${\kappa }_{\mathrm{B}}^{\left(r\right)}(\tau -{\epsilon }_{\mathrm{B}})={\kappa }^{\left(r\right)}(\tau -{\epsilon }_{\mathrm{B}})$ for $r=0,\cdots ,p-1,$ and ${\kappa }_{\mathrm{B}}^{\left(r\right)}\left(\tau \right)=0$ for $r=1,\cdots ,p-1$; see Figure 1.

Then, we define a $2\tau$-periodic function on ${\mathbb{R}}^d$ by

$$\tilde{\kappa }\left(x\right):= {\kappa }_{\mathrm{R}}(\parallel x\parallel ),x\in {[-\tau ,\tau )}^d.$$

(14)

By construction, $\tilde{\kappa }$ is p times continuously differentiable and we have $\tilde{\kappa }(x_{i_k}^k-x_{i_{\ell }}^{\ell })=exp(-\frac{1}{\eta }{\parallel x_{i_k}^k-x_{i_{\ell }}^{\ell }\parallel }_2^2)$ for all $i_k\in \left[n_k\right],i_{\ell }\in \left[n_{\ell }\right]$.

Let $M\in \mathbb{N}$. We approximate $\tilde{\kappa }$ using the $2\tau$-periodic Fourier expansion of degree $2M$,

$$\tilde{\kappa }\left(x\right)\approx {\displaystyle \sum _{m\in {\{-M,\cdots ,M-1\}}^d}}\widehat{\kappa }\left(m\right)exp\left(\mathrm{i}\frac{\pi }{\tau }{m}^{⊺}x\right),x\in {\mathbb{R}}^d,$$

with the discrete Fourier coefficients $\widehat{\kappa }\left(m\right)\in \mathbb{C}$, which can be efficiently approximated by the fast Fourier transform (FFT)

$$\widehat{\kappa }\left(m\right):= \frac{1}{{\left(2M\right)}^d}{\displaystyle \sum _{x\in \frac{\tau }{M}{\{-M,\cdots ,M-1\}}^d}}\tilde{\kappa }\left(x\right)exp(-\mathrm{i}\frac{\pi }{\tau }{m}^{\top }x),m\in {\{-M,\cdots ,M-1\}}^d.$$

(15)

This yields an approximation of (11) by

$$\begin{array}{c}{\beta }_{i_k}=\sum _{j=1}^{n_{\ell }}\tilde{\kappa }(x_{i_k}^k-x_{i_{\ell }}^{\ell }){\alpha }_{i_{\ell }}\approx \sum _{i_{\ell }=1}^{n_{\ell }}{\displaystyle \sum _{m\in {\{-M,\cdots ,M-1\}}^d}}\widehat{\kappa }\left(m\right)exp\left(\mathrm{i}\frac{\pi }{\tau }{m}^{⊺}(x_{i_k}^k-x_{i_{\ell }}^{\ell })\right){\alpha }_{i_{\ell }}\hfill \\ \hfill ={\displaystyle \sum _{m\in {\{-M,\cdots ,M-1\}}^d}}\widehat{\kappa }\left(m\right)\left(\sum _{i_{\ell }=1}^{n_{\ell }}{\alpha }_{i_{\ell }}exp\left(-\mathrm{i}\frac{\pi }{\tau }{m}^{⊺}{x}_{i_{\ell }}^{\ell }\right)\right)exp\left(\mathrm{i}\frac{\pi }{\tau }{m}^{⊺}{x}_{i_k}^k\right).\end{array}$$

(16)

The non-uniform discrete Fourier transform (NDFT) of $\widehat{\kappa }:= {\left[\widehat{\kappa }\left(m\right)\right]}_{m\in {\{-M,\cdots ,M-1\}}^d}$ at the nodes ${\Omega }^k\subset {\mathbb{R}}^d$ is defined by

$${\left[{\mathit{F}}_k\widehat{\kappa }\right]}_{i_{\ell }}:= {\displaystyle \sum _{m\in {\{-M,\cdots ,M-1\}}^d}}\widehat{\kappa }\left(m\right)exp\left(-\mathrm{i}\frac{\pi }{\tau }{m}^{⊺}{x}_{i_{\ell }}^k\right),i_{\ell }\in \left[n_{\ell }\right],$$

(17)

and the adjoint NDFT of $\alpha \in {\mathbb{R}}^{n_{\ell }}$ on the set ${\Omega }^{\ell }$ is given by

$${\left[{\mathit{F}}_{\ell }^{*}\alpha \right]}_m:= \sum _{i_{\ell }=1}^{n_{\ell }}{\alpha }_{i_{\ell }}exp\left(-\mathrm{i}\frac{\pi }{\tau }{m}^{⊺}{x}_{i_{\ell }}^{\ell }\right),m\in {\{-M,\cdots ,M-1\}}^d,$$

(18)

cf. [43] (Section 7). Therefore, the approximation (16) can be written as

$$\beta ={\mathcal{K}}^{(k,\ell )}\alpha \approx {\mathit{F}}_k(\widehat{\kappa }\odot {\mathit{F}}_{\ell }^{*}\alpha ).$$

(19)

The procedure is summarized in Algorithm 4.

Algorithm 4 NFFT-based fast summation

1: Input: Vector $\alpha \in {\mathbb{R}}^{n_{\ell }}$, kernel function $\kappa$ in (12), parameters ${\epsilon }_{\mathrm{B}}>0$, $M\in \mathbb{N}$ 2: precomputation 3:     Compute the regularized kernel ${\kappa }_{\mathrm{R}}$ by (13) 4:     Compute the periodized kernel $\tilde{\kappa }$ by (14) 5:     Compute the discrete Fourier coefficients $\widehat{\kappa }\left(m\right)$, $m\in {\{-M,\cdots ,M-1\}}^d$, by an FFT, see (15) 6: end 7: Compute the adjoint NDFT ${\mathit{F}}_{\ell }\alpha$, see (18) 8: Compute the pointwise product $\widehat{\beta }:= \widehat{\kappa }\odot {\mathit{F}}_{\ell }\alpha$ 9: Compute the NDFT $\beta := {\mathit{F}}_k\widehat{\beta }$, see (17) 10: return   $\beta \approx {\mathcal{K}}^{(k,\ell )}\alpha$

There are fast algorithms, known as non-uniform fast Fourier transform (NFFT), allowing for the computation of an NDFT (17) and its adjoint (18) in $\mathcal{O}(M^{d}logM+N)$ steps up to arbitrary numeric precision; see e.g., [44,45] and [43] (Section 7), where $N={\parallel \mathit{n}\parallel }_{\infty }$. Note that the direct implementation of (16) requires $\mathcal{O}\left(M^{d}N\right)$ operations. We call the Sinkhorn algorithm where the matrix–vector multiplication is performed via (19) the NFFT-Sinkhorn algorithm. If we fix the Fourier expansion degree M, which is possible because of the smoothness of $\kappa$, then we end up at a numerical complexity of $\mathcal{O}\left(KN\right)$ for each iteration step of the NFFT-Sinkhorn algorithm for trees. In the case of a circle (Algorithm 3), we can apply the fast summation column by column for the matrix–matrix product with ${\mathcal{K}}^{(k,k+1)}$, yielding a complexity of $\mathcal{O}\left(K{N}^2\right)$ for each iteration step.

6. Numerical Examples

We illustrate the results from Section 5 concerning the Sinkhorn algorithm and its accelerated version, the NFFT-Sinkhorn. First, we investigate the effect of parameter choices in some artificial examples. Then, we look at the one-dimensional Euler flow problem of incompressible fluids and the fixed support barycenter problem of images (the code for our examples is available at https://github.com/fatima0111/NFFT-Sinkhorn, accessed on 8 August 2022). All computations were performed on an eight-core Intel Core i7-10700 CPU with 32 GB memory. For computing the NFFT of Section 5, we rely on the implementation [46].

6.1. Uniformly Distributed Points

We consider the ${\mathrm{MOT}}_{\eta }$ problem for uniform measures ${\mu }^k$ of uniformly distributed points on $\Omega =[-1/2,1/2]$. We chose the entropy regularization parameter $\eta =0.1$, so that a boundary regularization (13) for the fast summation method is necessary. For the tree-structured cost function, we set the boundary regularization ${\epsilon }_{\mathrm{B}}=1/16$, the Fourier expansion degree $M=156,$ and the smoothness parameters $p=3$ (see Section 5). In Figure 2 left, we see the linear dependence of the computational time on the number K of marginals. For a growing number N of points, the NFFT-Sinkhorn algorithm, which requires $\mathcal{O}\left(KN\right)$ steps, clearly outperforms the standard method, which requires $\mathcal{O}\left(K{N}^2\right)$ steps (see Figure 2 right).

In Figure 3, we show the computation times for the circle-structured cost function, with the parameters $\eta =0.1$, ${\epsilon }_{\mathrm{B}}=3/32$, $M=2000$, and $p=3$. As the Sinkhorn iteration requires matrix–matrix products, it is more costly than for the tree-structured cost function, and so we used a lower number of points N. The advantage of the NFFT-Sinkhorn is smaller than for the tree, but still considerable. We point out here that the fast summation method is applied column by column to the matrix–matrix product, and the Fourier expansion degree M is larger.

Finally, we investigate how the approximation error between the Sinkhorn and NFFT-Sinkhorn algorithm, $|{\mathcal{S}}^{\left(r\right)}-{\tilde{S}}^{\left(r\right)}|$, depends on the entropy regularization parameter $\eta$ and the Fourier expansion degree M at a fixed iteration $r=10$, where ${\mathcal{S}}^{\left(r\right)}$ denotes the evaluation with the Sinkhorn algorithm and ${\tilde{\mathcal{S}}}^{\left(r\right)}$ its evaluation with the NFFT-Sinkhorn algorithm. Since the time differences for different M in the one-dimensional case are very small, we consider the two-dimensional uniform marginal measures for the ${\mathrm{MOT}}_{\eta }$ problems with tree- or circle-structured cost functions. In Figure 4 and Figure 5, we see that for smaller $\eta$, we need a larger expansion degree M to achieve a good accuracy. The error stagnates at a certain level and does not decrease anymore for increasing M. This could be improved by increasing the approximation parameter p and the cutoff parameter of the NFFT, cf. [43] (Section 7). The parameter choice methods for the NFFT-based summation were discussed in [47]. For an appropriately chosen M, the NFFT-Sinkhorn is usually much faster than the Sinkhorn algorithm. However, for very small $\eta$, the kernel function (12) is concentrated over a small interval, and therefore, a simple truncation of the sum (11) might be beneficial to the NFFT approximation.

6.2. Fixed-Support Wasserstein Barycenter for General Trees

Let $\mathcal{T}=\left(\mathcal{V},\mathcal{E}\right)$ be a tree with set of leaves $\mathcal{L}$; see Section 4.1. For $k\in \mathcal{L}$, let measures ${\mu }^k$ and weights ${\tilde{w}}_k\in [0,1]$ be given that satisfy ${\sum }_{k\in \mathcal{L}}{\tilde{w}}_k=1$. For any edge $e\in \mathcal{E}$, we set

$$w_e:= \left\{\begin{array}{cc}1\hfill & \mathrm{if}\left|e\cap \mathcal{L}\right|\ne 1,\hfill \\ {\tilde{w}}_k\hfill & \mathrm{if}e\cap \mathcal{L}=k.\hfill \end{array}\right.$$

The generalized barycenters are the minimizers ${\mu }^k$, $k\in \mathcal{V}\backslash \mathcal{L}$, of

$${\displaystyle \underset{{\mu }^k\in \mathbb{P}\left({\Omega }^k\right),k\in \mathcal{V}\backslash \mathcal{L}}{inf}}{\displaystyle \sum _{e=e_1,e_2\in \mathcal{E}}}{w}_e{\mathrm{W}}_2^2\left({\mu }^{e_1},{\mu }^{e_2}\right),$$

(20)

where ${\mathrm{W}}_2^2({\mu }^{e_1},{\mu }^{e_2})$ is the squared Wasserstein distance [48,49] between the measures ${\mu }^{e_1}$ and ${\mu }^{e_2}.$ The well-known Wasserstein barycenter problem [50,51] is a special case of (20), where the tree is star-shaped and the barycenter corresponds to the unique internal node. We consider the fixed-support barycenter problem [26,52], where the nodes ${\mathit{x}}^k$, $k\in \mathcal{V}\backslash \mathcal{L}$ are also given, so that we need to optimize (20) only for ${\mu }_{i_k}^k$, $k\in \mathcal{V}\backslash \mathcal{L}$. This yields an MOT problem with the tree-structured cost

$$C_{\mathit{i}}={\displaystyle \sum _{e=e_1,e_2\in \mathcal{E}}}{w}_e{\parallel {\mathit{x}}_{i_{e_1}}^{e_1}-{\mathit{x}}_{i_{e_1}}^{e_2}\parallel }_2^2,\mathit{i}\in \mathit{I},$$

where the marginal constraints of (1) are only set for the known measures ${\mu }^k$, $k\in \mathcal{L}$ (see [53]). This barycenter problem can be solved approximately using a modification of the Sinkhorn Algorithm 2, in which we replace line 12 by

$${\left({\varphi }^k\right)}^{(r+1)}=\left\{\begin{array}{cc}\mathit{1}\hfill & \mathrm{if}k\in \mathcal{L},\hfill \\ {\mu }^k\oslash \left({\gamma }_k^{\left(r\right)}\odot {\left({\beta }^{{\odot }_{{\mathcal{C}}_k}}\right)}^{\left(r\right)}\right)\hfill & \mathrm{otherwise}.\hfill \end{array}\right.$$

We test our algorithm with a tree consisting of $K=7$ nodes; see Figure 6. The four given marginals ${\mu }^k$, $k\in \mathcal{L}$, are dithered images in ${\mathbb{R}}^2$ with uniform weights ${\mu }_{i_k}^k=1/N$. As support points of the barycenters ${\mu }^k$, $k\in \mathcal{V}\backslash \mathcal{L}$, we take the union over all support points ${\mathit{x}}_{i_k}^k$ of all four input measures ${\mu }^k,$$k\in \mathcal{L}.$ Furthermore, we use the barycenter weights ${\tilde{w}}^k=1/4$. The given images and the computed barycenters are show in Figure 7, where we executed $r=150$ iterations of the Sinkhorn algorithm and its accelerated version. We chose the regularization parameter $\eta =5·{10}^{-3}$ and the fast summation parameters $M=156$, $p=3.$

6.3. Generalized Euler Flows

We consider the motion of N particles of an incompressible fluid, e.g., water, in a bounded domain $\Omega$ in discrete time steps $t_k(k-1)/(K-1)$, $k\in \left[K\right]$. We assume that we know the function $\sigma :\Omega \to \Omega$, which connects N initial positions $x_{i_1}\in \Omega$ of particles with their final positions $x_{i_K}\in \Omega$. At each time step $t_k$, we know an image of the particle distribution, which is described by the discrete marginal measure ${\mu }^k$, $k\in \left[K\right]$, with uniform weights ${\mu }_{i_k}^k=1/N$. We want to find out how the single particles move, i.e., their trajectories. Due to the least-action principle, this problem can be formulated as the MOT problem (4) with the circle-structured cost

$$C_{\mathit{i}}={\parallel x_{i_K}^K-\sigma \left(x_{k_1}^1\right)\parallel }_2^2+\sum _{k=1}^{K-1}{\parallel x_{i_{k+1}}^{k+1}-x_{i_k}^k\parallel }_2^2,\mathit{i}\in \mathit{I},$$

see [55,56,57]. Then, the pair marginal

$${\Pi }_{1,k}:= {\displaystyle \sum _{\ell \in \left[K\right]\backslash 1,k}}{\displaystyle \sum _{i_{\ell }\in \left[n_{\ell }\right]}}{\Pi }_{i_1,\cdots ,i_K}$$

of the optimal plan $\Pi$ provides the (discrete) probability that a particle that was initially at position $x_{i_1}\in \Omega$ is in position $x_{i_k}\in \Omega$ at time $t_k$, $k=2,\cdots ,K-1$. The one-dimensional problem has been studied by several authors [25,26,39], where the particles are assumed to be on a grid. Here, we consider the case where the positions are uniformly distributed. We draw $N=400$ uniformly distributed points on $\Omega =[0,1]$. We use $K=5$ marginal constraints and the entropy regularization parameter $\eta =0.05$. Figure 8 and Figure 9 display the probability matrix ${\Pi }_{1,k}$ describing the motion of the particles from initial time $t_1=0$ to time $t_k$, where we use $r=50$ iterations for both the NFFT-Sinkhorn and the Sinkhorn algorithms, and two different connection functions $\sigma$.

7. Conclusions

We have proposed the NFFT-Sinkhorn algorithm to solve the ${\mathrm{MOT}}_{\eta }$ problem efficiently. Assuming that the cost function of the multi-marginal optimal transport decouples according to a tree or a circle, we obtain a linear complexity in K. The complexity of the algorithm with respect to the numbers $n_k,$$k\in \left[K\right]$, of atoms of the discrete marginal measures is further improved using the non-uniform fast Fourier transform. This results in a considerable acceleration in our numerical experiments compared to the usual Sinkhorn algorithm. The tree-structured ${\mathrm{MOT}}_{\eta }$ problem gives a much better numerical complexity than the circle-structured ${\mathrm{MOT}}_{\eta }$ problem due to the fact that in the latter case, matrix–matrix products are required in Algorithm 3 instead of just the matrix–vector products of Algorithm 2.