Hierarchical Low-Rank Approximation of Regularized Wasserstein Distance

Let ${\mathcal{X}}$ and ${\mathcal{Y}}$ be two measurable subsets of the Euclidean space ${\mathbb{R}}^{d}$, with $d\in{\mathbb{N}}$. Let further ${\bf f}$ and ${\bf g}$ be two $n$-dimensional probability vectors, i.e. two vectors in the probability simplex $\Sigma_{n}:=\{{\bf f}\in{\mathbb{R}}_{+}^{n}:\,\,{\bf f}^{\top}\mathds{1}_{n}=1\}$, defined on ${\mathcal{X}}$ and ${\mathcal{Y}}$, respectively. Here, $\mathds{1}_{n}$ is the $n$-dimensional vector of ones. The two probability vectors ${\bf f}$ and ${\bf g}$ are assumed to be given at two sets of $n$ discrete points $\{{\bf x}_{1},\dotsc,{\bf x}_{n}\}\subset{\mathcal{X}}$ and $\{{\bf y}_{1},\dotsc,{\bf y}_{n}\}\subset{\mathcal{Y}}$ in ${\mathbb{R}}^{d}$, respectively. Figure 1 shows a schematic representation of probability vectors in $d=1$ and $d=2$ dimensions.

We denote by $U({\bf f},{\bf g})$ the transport polytope of ${\bf f}$ and ${\bf g}$, i.e. the set of all nonnegative $n\times n$ matrices with row and column sums equal to ${\bf f}$ and ${\bf g}$, respectively, that is,

Each matrix $P=[P_{ij}]\in{\mathbb{R}}_{+}^{n\times n}$ in the transport polytope $U({\bf f},{\bf g})$ is a transport (or coupling) matrix that encodes a transport plan, that is, each element $P_{ij}$ of $P$ describes the amount of mass to be transported from a source point ${\bf x}_{i}\in{\mathcal{X}}$ to a traget point ${\bf y}_{j}\in{\mathcal{Y}}$, with $i,j=1,\dotsc,n$. The two constraints $P\mathds{1}_{n}={\bf f}$ and $P^{\top}\mathds{1}_{n}={\bf g}$ are necessary for a transport plan $P$ to be admissible; all the mass taken from a point ${\bf x}_{i}$, i.e. $\sum_{j=1}^{n}P_{ij}$, must be equal to the mass at point ${\bf x}_{i}$, i.e. ${\bf f}_{i}$, and all the mass transported to a point ${\bf y}_{j}$, i.e. $\sum_{i=1}^{n}P_{ij}$, must be equal to the mass at point ${\bf y}_{j}$, i.e. ${\bf g}_{j}$. We say that ${\bf f}$ and ${\bf g}$ are the marginals of $P$.

Let further $c:{\mathcal{X}}\times{\mathcal{Y}}\rightarrow{\mathbb{R}}_{+}$ be a non-negative cost function defined on ${\mathcal{X}}\times{\mathcal{Y}}$. For any pair of points $({\bf x},{\bf y})\in{\mathcal{X}}\times{\mathcal{Y}}$, the value $c({\bf x},{\bf y})$ represents the cost of transporting one unit of mass from a source point ${\bf x}\in{\mathcal{X}}$ to a target point ${\bf y}\in{\mathcal{Y}}$. Correspondingly, we consider the cost matrix

The optimal transport problem can now be formulated as follows. We first note that for any given transport plan $P\in U({\bf f},{\bf g})$, the total transport cost associated with $P$ is given by the Frobenius inner product $\langle P,C\rangle=\sum_{i,j}P_{ij}C_{ij}$. Kantorovich’s optimal transport problem then aims at minimizing the total transport cost over all admissible transport plans. It reads

where $T_{C}({\bf f},{\bf g})$ is the optimal total cost of transporting ${\bf f}$ onto ${\bf g}$.

In some applications the main goal of solving Kantorovich’s optimal transport problem (1) is to find the optimal transport plan, i.e. the transport matrix $P$ that minimizes the total cost. In some other applications the goal is to directly find the optimal cost $T_{C}$ without the need to explicitly form the optimal transport matrix. An example of the latter application appears in the construction of loss functions in optimization problems, where loss functions are built based on metrics induced by the optimal cost $T_{C}$. An important type of such metrics is Wasserstein metric.

For the optimal cost $T_{C}$ in (1) to induce a metric, we need the cost matrix $C\in{\mathbb{R}}_{+}^{n\times n}$ to be a metric matrix belonging to the cone of distance matrices [17], i.e. we need $C\in{\mathcal{M}}^{n\times n}$, where

Throughout this paper, we always assume that the cost matrix $C$ is a metric matrix: $C\in{\mathcal{M}}^{n\times n}$. An important class of metric matrices is driven by cost functions of the form $c({\bf x},{\bf y})=d({\bf x},{\bf y})^{p}$, where $p\in[1,\infty)$, and $d$ is a distance function (or metric). In this case the cost matrix reads

Intuitively, this choice implies that the cost of transporting one unit of mass from a source point ${\bf x}\in{\mathcal{X}}$ to a target point ${\bf y}\in{\mathcal{Y}}$ corresponds to the distance between the two points. The optimal cost $T_{C}({\bf f},{\bf g})$ in (1) corresponding to the cost matrix (2) induces a metric known as the Wasserstein distance of order $p$ [16, 17],

Due to its desirable convexity and convergence features, Wasserstein distance is considered as an important measure of dissimilarity in many applications, ranging from computer vision and machine learning to Seismic and Bayesian inversion; see e.g. [13, 5, 18, 11]. Unfortunately, the numerical computation of Wasserstein metric is in general infeasible in high dimensions and is mainly limited to low dimensions; see e.g. [13]. For instance, in the case $c({\bf x},{\bf y})=||{\bf x}-{\bf y}||^{2}$, the quadratic Wasserstein distance between two continuous probability density functions $f:{\mathcal{X}}\rightarrow{\mathbb{R}}+$ and $g:{\mathcal{Y}}\rightarrow{\mathbb{R}}+$ is given by [2]

where $M$ is the optimal map. In one dimension, i.e. when ${\mathcal{X}},{\mathcal{Y}}\subset{\mathbb{R}}$, the optimal map is given by $M=G^{-1}\circ F$, where $F$ and $G$ are the cumulative distribution functions corresponding to $f$ and $g$, respectively. The discrete $W_{2}$ metric between the corresponding probability vectors ${\bf f}=[f({\bf x}_{i})]$ and ${\bf g}=[g({\bf y}_{i})]$ can then be efficiently approximated by a quadrature rule,

where $G^{-1}$ may be computed by interpolation and an efficient search algorithm, such as binary search. The overall cost of computing the discrete $W_{2}$ metric in one dimension will then be ${\mathcal{O}}(n\,\log n)$. In multiple dimensions, i.e. when ${\mathcal{X}},{\mathcal{Y}}\subset{\mathbb{R}}^{d}$ with $d\geq 2$, however, the computation of the optimal map $M$, and hence the $W_{2}$ metric, may be either expensive or infeasible. The map is indeed given by the gradient of a convex function that is the solution to a constrained nonlinear Monge-Ampère equation. We refer to [5] for more discussion on the finite difference approximation of the Monge-Ampère equation in two dimensions.

One approach to computing Kantorovich’s optimal transport problem (1) is to regularize the original problem and seek approximate solutions to the regularized problem. One possibility is to add an entropic penalty term to the total transport cost and arrive at the regularized problem

where $\lambda>0$ is a regularization parameter, and $H(P)$ is the discrete entropy of transport matrix,

The regularized problem (4)-(5) with $C\in{\mathcal{M}}^{n\times n}$ has a unique solution, say $P_{\lambda}$ [3]. Indeed, the existence and uniqueness follows from the boundedness of $U({\bf f},{\bf g})$ and the strict convexity of the negative entropy. On the one hand, as $\lambda$ increases, the unique solution $P_{\lambda}$ converges to the solution with maximum entropy within the set of all optimal solutions of the original Kantorovich’s problem; see Proposition 4.1 in [13],

In particular, we have $T_{C}^{\lambda}({\bf f},{\bf g})\stackrel{{\scriptstyle\lambda\rightarrow\infty}}{{\longrightarrow}}T_{C}({\bf f},{\bf g})$. On the other hand, as $\lambda$ decreases, the unique solution $P_{\lambda}$ converges to the transport matrix with maximum entropy between the two marginals ${\bf f}$ and ${\bf g}$, i.e. the outer product of ${\bf f}$ and ${\bf g}$; see Proposition 4.1 in [13],

It is to be noted that in this latter case, as $\lambda\rightarrow 0$, the optimal solution $P_{\lambda}$ becomes less and less sparse (or smoother), i.e. with more and more entries larger than a prescribed small threshold.

With $P_{\lambda}$ being the optimal solution to the regularized problem (4)-(5) corresponding to the cost matrix (2), the Sinkhorn divergence of order $p$ between ${\bf f}$ and ${\bf g}$ is defined as

Sinkhorn divergence (6) has several important properties. First, it satisfies all axioms for a metric with the exception of identity of indiscernibles (or the coincidence axiom), that is, it satisfies non-negativity, symmetry, and triangle inequality axioms [3]. Second, it is convex and smooth with respect to the input probability vectors and can be differentiated using automatic differentiation [13]. Third, as discussed below, it can be computed by a simple alternating minimization algorithm, known as Sinkhorn’s algorithm, where each iteration involves two matrix-vector products. Consequently, Sinkhorn divergence may serve as a feasible alternative to classical Waaserstein distance, particularly when probability measures are supported on multi-dimensional metric spaces.

To make the assumptions precise, we first introduce a special Kronecker sum of matrices.

It is to be noted that the all-ones Kronecker sum is different from the common Kronecker sum of two matrices in which identity matrices $I_{p}$ and $I_{q}$ replace all-ones matrices $J_{p}$ and $J_{q}$. This new operation facilitates working with elementwise matrix operations. We also note that the two-term sum in Definition 1 can be recursively extended to define all-ones Kronecker sums with arbitrary number of terms, thanks to the following associative property inherited from matrix algebra,

We next consider kernel matrices generated by kernel functions with a special type of regularity, known as “asymptotic smoothness” [7].

We now make the following two assumptions.

• A1.

  The cost matrix $C\in{\mathcal{M}}^{n\times n}$ can be decomposed into a $d$-term all-ones Kronecker sum

  $$C=C_{d}\oplus\dotsi\oplus C_{1},\qquad C_{k}\in{\mathcal{M}}^{n_{k}\times n_{k}},\ \ \ k=1,\dotsc,d,\ \ \ n=\prod_{k=1}^{d}n_{k},$$

  (16)

  where each matrix $C_{k}\in{\mathcal{M}}^{n_{k}\times n_{k}}$ is generated by a metric $c_{k}:{\mathbb{R}}\times{\mathbb{R}}\rightarrow{\mathbb{R}}_{+}$ on the real line.

• A2.

  The kernel matices $Q^{(k)}=\exp[-\lambda C^{(k)}]$ and $\hat{Q}^{(k)}=C^{(k)}\odot Q^{(k)}$, with $k=1,\dotsc,d$, generated by the kernels $\kappa_{k}(x,y)=\exp(-\lambda\,c_{k}(x,y))$ and $\hat{\kappa}_{k}(x,y)=c_{k}(x,y)\,\kappa_{k}(x,y)$, where $(x,y)\in{\mathbb{R}}^{2}$, are asymptotically smooth.

It is to be noted that by assumption A2 all $2\,d$ kernels $\kappa_{k}$ and $\hat{\kappa}_{k}$, with $k=1,\dotsc,d$, satisfy similar estimates of form (15) with possibly different sets of constants $(c_{0},\alpha,\beta,s)$.

We will first show that assumption A1 holds for $L^{p}$ cost functions. Let the two sets of discrete points $\{{\bf x}_{i}\}_{i=1}^{n}\in{\mathcal{X}}$ and $\{{\bf y}_{j}\}_{j=1}^{n}\in{\mathcal{Y}}$, at which the two probability vectors $({\bf f},{\bf g})$ are defined, be given on two regular grids in ${\mathbb{R}}^{d}$. We note that if the original probability vectors are given on irregular grids, we may find their values on regular grids by interpolation with a cost ${\mathcal{O}}(n)$ that would not change the near-linearity of our overall target cost. Set $I_{n}:=\{1,\dotsc,n\}$. Then, to each pair of indices $i\in I_{n}$ and $j\in I_{n}$ of the cost matrix $C=[c({\bf x}_{i},{\bf y}_{j})]$, we can assign a pair of $d$-tuples $(i_{1},\dotsc,i_{d})$ and $(j_{1},\dotsc,j_{d})$ in the Cartesian product of $d$ finite sets $I_{n_{1}}\times\dotsc\times I_{n_{d}}$, where $n=\prod_{k=1}^{d}n_{k}$. Consequently, the additive representation of $c$ in (17) will imply (16).

It is also easy to see that assumption A2 holds for $L^{p}$ cost functions, with corresponding kernels,

Indeed, these kernels are asymptotically smooth due to the exponential decay of $\exp(-\lambda\,|x-y|^{p})$ as $|x-y|$ increases. For instance, in the particular case when $p=2$, the estimate (15) holds for both $\kappa$ and $\hat{\kappa}$ with $c_{0}=1$, $\alpha=2$, $\beta=0$, and $s=0$. To show this, we note that

where $H_{m}(z)$ is the Hermite polynomial of degree $m$, satisfying the following inequality [8],

Setting $z:=\sqrt{\lambda}(x-y)\neq 0$, and noting that $(2^{m}\,m!)^{1/2}\leq 2^{m}\,m!$ for $m\in{\mathbb{N}}$, we obtain

Hence we arrive at

Similarly, we can show that the same estimate holds for $|\partial_{x}^{m}(x-y)^{2}\exp(-\lambda\,(x-y)^{2})|$.

As noted in Section 2.2, the main computational bottleneck of Sinkhorn’s algorithm is the matrix-vector multiplications by kernels $Q$ and $Q^{\top}$ in each iteration (12) and by kernel $\hat{Q}$ in (14). We present a strategy that reduces the complexity of each Sinkhorn iteration to ${\mathcal{O}}(n\,\log^{2}n)$, enabling Sinkhorn’s algorithm to achieve a near-linear overall complexity ${\mathcal{O}}(n\,\log^{3}n)$. The proposed strategy takes two steps, following two observations made based on assumptions A1 and A2.