Matricial Wasserstein-1 Distance

TL;DR

This paper introduces a matrix-valued extension of the Wasserstein 1-metric, enabling efficient computation and unbalanced mass transport interpretation for matrix probability densities.

Contribution

It develops a novel matrix analogue of the Wasserstein 1-metric using duality theory, expanding the applicability of optimal transport to matrix-valued data.

Findings

Provides a duality-based formulation for matrix Wasserstein-1 distance

Enables easier computation of matrix optimal transport distances

Offers an unbalanced interpretation of mass transport for matrices

Abstract

In this note, we propose an extension of the Wasserstein 1-metric ($W_1$) for matrix probability densities, matrix-valued density measures, and an unbalanced interpretation of mass transport. The key is using duality theory, in particular, a "dual of the dual" formulation of $W_1$. This matrix analogue of the Earth Mover's Distance has several attractive features including ease of computation.

Tables

Table 1. TABLE I: Distances between power spectra

	${\rho }_0,{\rho }_1$	${\rho }_1,{\rho }_2$	${\rho }_0,{\rho }_2$
${\beta }_1=10,{\beta }_2=1$	77.85	77.76	137.36
${\beta }_1=1,{\beta }_2=1$	249.40	162.03	199.78
${\beta }_1=1,{\beta }_2=10$	210.93	110.25	113.46

Equations

$$W_1({\rho }_0,{\rho }_1):=\underset{\pi \in \Pi ({\rho }_0,{\rho }_1)}{\inf }{\int }_{{\mathbb{R}}^m\times {\mathbb{R}}^m}\Vert x-y\Vert \pi (dx,dy),$$

$$W_1({\rho }_0,{\rho }_1)$$

$$\text{mbox}subjectto\Vert f{\Vert }_{\mathrm{Lip}}\le 1\},$$

$$W_1({\rho }_0,{\rho }_1)$$

$$\text{mbox}subjectto\Vert {\nabla }_{x}f\Vert \le 1\},$$

$$W_1({\rho }_0,{\rho }_1)$$

$$\text{mbox}subjectto{\rho }_0-{\rho }_1+{\nabla }_x\cdot u=0\},$$

$$W_1({\rho }_0,{\rho }_1)$$

$$\text{mbox}subjectto\frac{\partial \rho (t,x)}{\partial t}+{\nabla }_x\cdot u(t,x)=0,$$

$$\text{mbox}and\rho (0,x)={\rho }_0(x),{\rho }_1(1,x)={\rho }_1(x)\},$$

$$\mathcal{D}$$

$${\mathcal{D}}_{+}$$

$$T_{\rho }=\{\delta \in \mathcal{H}\mid \mathrm{tr}(\delta )=0\},$$

$$⟨X,Y⟩=\mathrm{tr}(X^{\ast }Y),$$

$$⟨X,Y⟩=\sum _{k=1}^N\mathrm{tr}(X_k^{\ast }{Y}_k).$$

$$XY=\left[\begin{array}[]{c}X_{1}\\ \vdots\\ X_{N}\end{array}\right]Y:=\left[\begin{array}[]{c}X_{1}Y\\ \vdots\\ X_{N}Y\end{array}\right],$$

$$YX=Y\left[\begin{array}[]{c}X_{1}\\ \vdots\\ X_{N}\end{array}\right]:=\left[\begin{array}[]{c}YX_{1}\\ \vdots\\ YX_{N}\end{array}\right].$$

$$\nabla_{L}:{\mathcal{H}}\rightarrow{{\mathcal{S}}}^{N},~{}~{}X\mapsto\left[\begin{array}[]{c}L_{1}X-XL_{1}\\ \vdots\\ L_{N}X-XL_{N}\end{array}\right]$$

$$\nabla_{L}^{*}:{{\mathcal{S}}}^{N}\rightarrow{\mathcal{H}},~{}~{}Y=\left[\begin{array}[]{c}Y_{1}\\ \vdots\\ Y_{N}\end{array}\right]\mapsto\sum_{k}^{N}L_{k}Y_{k}-Y_{k}L_{k},$$

$$⟨{\nabla }_{L}X,Y⟩=⟨X,{\nabla }_L^{\ast }Y⟩.$$

$$W_1({\rho }_0,{\rho }_1):=\underset{f\in \mathcal{H}}{\sup }\left\{\mathrm{tr}[f({\rho }_0-{\rho }_1)]\mid \Vert {\nabla }_{L}f\Vert \le 1\right\}.$$

$$\nabla_{L}f=\left[\begin{array}[]{c}L_{1}f-fL_{1}\\ \vdots\\ L_{N}f-fL_{N}\end{array}\right].$$

$$d_D({\rho }_0,{\rho }_1)=\underset{f\in \mathcal{H}}{\sup }\left\{\mathrm{tr}[f({\rho }_0-{\rho }_1)]\mid \Vert [D,f]\Vert \le 1\right\}.$$

$${\hat{W}}_1({\rho }_0,{\rho }_1)=\underset{u\in {\mathcal{S}}^N}{\inf }\left\{\Vert u{\Vert }_{\ast }\mid {\rho }_0-{\rho }_1-{\nabla }_L^{\ast }u=0\right\},$$

$$W_1({\rho }_0,{\rho }_1)={\hat{W}}_1({\rho }_0,{\rho }_1).$$

$$\Vert u{\Vert }_{\ast }=\underset{g\in {\mathcal{S}}^N,\Vert g\Vert \le 1}{\sup }⟨u,g⟩.$$

$${\hat{W}}_1({\rho }_0,{\rho }_1)$$

$$W_1({\rho }_0,{\rho }_2)\le W_1({\rho }_0,{\rho }_1)+W_1({\rho }_1,{\rho }_2).$$

$${\rho }_0-{\rho }_2-{\nabla }_L^{\ast }(u_1+u_2)=0.$$

$$W_1({\rho }_0,{\rho }_2)\le \Vert u_1+u_2{\Vert }_{\ast }\le \Vert u_1{\Vert }_{\ast }+\Vert u_2{\Vert }_{\ast },$$

$${\rho }_0-{\rho }_1-{\nabla }_L^{\ast }u-v=0.$$

$$V_1({\rho }_0,{\rho }_1)=\underset{\begin{array}{c}u\in {\mathcal{S}}^N\\ v\in \mathcal{H}\end{array}}{\inf }\left\{\Vert u{\Vert }_{\ast }+\alpha \Vert v{\Vert }_{\ast }\mid {\rho }_0-{\rho }_1-{\nabla }_L^{\ast }u-v=0\right\}.$$

$$W_1(\mu ,\nu )+\alpha \Vert {\rho }_0-\mu {\Vert }_{\ast }+\alpha \Vert {\rho }_1-\nu {\Vert }_{\ast }.$$

Full text

I Introduction

Optimal mass transport (OMT) has proven to be a powerful methodology for numerous problems in physics, probability, information theory, fluid mechanics, econometrics, systems and control, computer vision, and signal/image processing [14, 26, 1, 30, 29, 20]. Developments along purely controls-related issues ensued when it was recognized that mass transport may be naturally reformulated as a stochastic control problem; see [15, 21, 16, 11, 5, 6, 7, 8, 9] and the references therein.

Historically, the problem of OMT [26, 30] began with the question of minimizing the effort of transporting one distribution to another, typically with a cost proportional to the Euclidean distance between starting and ending points of the mass being transported. However, the control-theoretic reformulation [1] which was at the root of the aforementioned developments was based on the choice of a quadratic cost. The quadratic cost allowed the interpretation of the transport effort as an action integral and gave rise to a Riemannian structure on the space of distributions [18, 13, 25]. The originality in our present work is two-fold. First, we formulate the transport problem with an $L_{1}$ cost in a similar manner, as a control problem with an $L_{1}$-path cost functional, and secondly, we develop theory for shaping flows of matrix-valued distributions which is a non-trivial generalization of classical OMT.

The relevance of OMT on flows of matrix-valued distributions was already recognized in [22, 23] and was cast as a control problem as well, albeit in a quadratic-cost setting. At that point, interest in the geometry of matrix-valued distributions stemmed from applications to spectral analysis of vector-valued time series (see [22] and the references therein). Yet soon it became aparent that flows of matrix-valued distributions represent evolution of quantum systems. In fact, there has been a burst of activity in applying ideas of quantum mechanics to OMT of matrix-valued densities as well as, utilizing an OMT framework to study the dynamics of quantum systems: three groups [3, 4, 19] independently and simultaneously developed quantum mechanical frameworks for defining a Wasserstein-2 distance on matrix-valued densities (normalized to have trace 1), via a variational formalism generalizing the work of [1]. We note that [3, 4, 19] develop matrix-valued generalizations of the Wasserstein 2-metric ($W_{2}$) and explore the Riemannian-like structure for studying the entropic flows of quantum states.

Thus, in our present note, we develop a natural extension of the Wasserstein 1-metric to matrix-valued densities and matrix-valued measures. Our point of view is somewhat different from the earlier works on matricial Wassserstein-2 metrics. We mainly use duality theory [12, 30]. Further, we do not employ the Benamou and Brenier [1] control formulation of OMT, but rather the Kantorovich-Rubinstein duality. This new scheme is computationally more attractive and, moreover, it is especially appealing when specialized to weighted graphs (discrete spaces) that are sparse (few edges), as is the case for many real-world networks [27, 28, 31].

The present paper is structured as follows. Section II is a quick review of several different formulations of Wasserstein-1 distance in the scalar setting. Using the quantum gradient operator defined in Section III, we generalize the Wasserstein-1 metric to the space of density matrices in Section IV. The case where the two marginal matrices have different traces is discussed in Section V. We finally extend the framework to deal with matrix-valued densities in Section VI, which may find applications in multivariate spectra analysis as well as comparing stable multi-inputs multi-outputs (MIMO) systems. The paper concludes with an academic example in Section VII.

II Optimal mass transport

We begin with duality theory, explained for scalar densities, upon which our matricial generalization of the Wasserstein-1 metric is based.

Given two probability densities $\rho_{0}$ and $\rho_{1}$ on ${\mathbb{R}}^{m}$, the Wasserstein-1 distance between them is

[TABLE]

where $\Pi(\rho_{0},\rho_{1})$ denotes the set of couplings between $\rho_{0}$ and $\rho_{1}$. The Wasserstein-1 distance has a dual formulation via the following result due to Kantorovich and Rubinstein [12, 26, 30]:

[TABLE]

where $\|f\|_{\rm Lip}$ denotes the Lipschitz constant. When $f$ is differentiable, $\|f\|_{\rm Lip}=\|\nabla_{x}f\|$. It follows that,

[TABLE]

Starting from (3), by once again considering the dual, we readily obtain the very important reformulation

[TABLE]

where the (Lagrange) optimization variable $u$ now represents flux. Alternatively, this can be written as the control-optimization problem in the Benamou-Brenier style [1]

[TABLE]

This “dual of the dual” formulation turns the Kantorovich and Rubinstein into a control problem to determine a suitable velocity (control vector) $u$. We remark that from a computational standpoint, when applied to discrete spaces (graphs), this formulation leads to a very substantial computational benefit in the case of sparse graphs; this is due to the fact that (1) involves solving systems of the order of the square of the number of nodes, while equation (4), solving systems of the order of the number of edges.

III Gradient on space of Hermitian matrices

We closely follow the treatment in [4]. In particular, we will need a notion of gradient on the space of Hermitian matrices and its dual, i.e. the divergence.

Denote by ${\mathcal{H}}$ and ${\mathcal{S}}$ the set of $n\times n$ Hermitian and skew-Hermitian matrices, respectively. We will assume that all of our matrices are of fixed size $n\times n$. Next, we denote the space of block-column vectors consisting of $N$ elements in ${\mathcal{S}}$ and ${\mathcal{H}}$ as ${\mathcal{S}}^{N}$ and ${\mathcal{H}}^{N}$, respectively. We also let ${\mathcal{H}}_{+}$ and ${\mathcal{H}}_{++}$ denote the cones of nonnegative and positive-definite matrices, respectively, and

[TABLE]

We note that the tangent space of ${\mathcal{D}}_{+},$ at any $\rho\in{\mathcal{D}}+$ is given by

[TABLE]

and we use the standard notion of inner product, namely

[TABLE]

for both ${\mathcal{H}}$ and ${\mathcal{S}}$. For $X,Y\in{\mathcal{H}}^{N}$ (${\mathcal{S}}^{N}$),

[TABLE]

Given $X=[X_{1}^{*},\cdots,X_{N}^{*}]^{*}\in{\mathcal{H}}^{N}$ (${\mathcal{S}}^{N}$), $Y\in{\mathcal{H}}$ (${\mathcal{S}}$), set

[TABLE]

and

[TABLE]

For a given $L\in{\mathcal{H}}^{N}$ we define

[TABLE]

to be the gradient operator. By analogy with the ordinary multivariable calculus, we refer to its dual with respect to the Hilbert-Schmidt inner product as the (negative) divergence operator, and this is

[TABLE]

i.e., $\nabla_{L}^{*}$ is defined by means of the identity

[TABLE]

A standing assumption throughout, is that the null space of $\nabla_{L}$, denoted by ${\rm ker}(\nabla_{L})$, contains only scalar multiples of the identity matrix.

IV Wassertein-1 distance for density matrices

In this section, we show that both (3) and (4) have natural counterparts for probability density matrices, i.e. matrices in ${\mathcal{D}}$. This set-up obviously works for matrices in ${\mathcal{H}}_{+}$ of equal trace.

We treat (3) as our starting definition and define the $W_{1}$ distance in the space of density matrices as

[TABLE]

Here $\|\cdot\|$ is the operator norm. The above is well-defined since by assumption, the null space of $\nabla_{L}$ is spanned by the identity matrix $I$. As above, we have that

[TABLE]

This should be compared to the Connes spectral distance [10], which is given by

[TABLE]

It is not difficult to see that the dual of (11) is

[TABLE]

which is the counterpart of (4). Here $\|\cdot\|_{*}$ denotes the nuclear norm [2]. In particular, we have the following theorems.

Theorem 1

Notation as above. Then

[TABLE]

Proof:

We start from (12) and use the fact that

[TABLE]

It follows

[TABLE]

This implies that (11) and (12) are dual to each other. Since both of them are strictly feasible, the duality gap is zero. Therefore $W_{1}(\rho_{0},\rho_{1})=\hat{W}_{1}(\rho_{0},\rho_{1})$. ∎

Theorem 2

The $W_{1}$ distance defined as in (11) is a metric on the space of density matrices ${\mathcal{D}}$.

Proof:

Obviously $W_{1}(\rho_{0},\rho_{1})\geq 0$ holds with equality if and only if $\rho_{0}=\rho_{1}$. The symmetric property that $W_{1}(\rho_{0},\rho_{1})=W_{1}(\rho_{1},\rho_{0})$ is also clear from the definition. Here we prove the triangle inequality. That is, for any $\rho_{0},\rho_{1},\rho_{2}\in{\mathcal{D}})$, we have

[TABLE]

It is easier to see this from the dual formulation (12). Let $u_{1}$, $u_{2}$ be the optimal fluxes for $(\rho_{0},\rho_{1})$ and $(\rho_{1},\rho_{2})$ respectively. Then $u_{1}+u_{2}$ is a feasible flux for $(\rho_{0},\rho_{2})$, namely,

[TABLE]

It follows that

[TABLE]

which completes the proof. ∎

V Wassertein-1 distance: the unbalanced case

In this section, we extend the definition of Wasserstein-1 distance to the space nonnegative matrices ${\mathcal{H}}_{+},$ i.e., we remove the constraint of both matrices having equal traces. Compare also with some very interesting recent work [17] on fast computational methods for $W_{1}$ in the unbalanced scalar case.

In order to compare matrices of unequal trace we relax the constraint in (12), which forces $\operatorname{tr}(\rho_{0})=\operatorname{tr}(\rho_{1})$, by introducing a “source” term $v\in{\mathcal{H}}$. That is, we replace our continuity equation (12) with

[TABLE]

With this added source, we define a Wasserstein-1 distance in ${\mathcal{H}}_{+}$ as follows. Given $\rho_{0},\rho_{1}\in{\mathcal{H}}_{+}$, we define

[TABLE]

Here $\alpha>0$ measures the relative significance between $u$ and $v$.

Another natural way to compare $\rho_{0},\rho_{1}\in{\mathcal{H}}_{+}$ is by finding $\mu,\nu\in{\mathcal{H}}_{+}$ having equal trace that are close to $\rho_{0},\rho_{1}$ in some norm (here taken to be the nuclear norm), as well as close to one another. More specifically, we seek $\mu,\nu$ to minimize

[TABLE]

Putting the two terms together we obtain the following definition of Wasserstein-1 distance

[TABLE]

It turns out these two relaxations of $W_{1}$ are in fact equivalent.

Theorem 3

With notation and assumptions as above,

[TABLE]

Proof:

Clearly, $\hat{V}_{1}(\rho_{0},\rho_{1})\geq V_{1}(\rho_{0},\rho_{1})$. On the other hand, let $u,v$ be a minimizer of (14), and $v=v_{1}-v_{0}$ with $v_{0},v_{1}\in{\mathcal{H}}_{+}$, i.e., $v_{0},v_{1}$ are the negative and positive parts of $v$ respectively, then $\mu=\rho_{0}+v_{0},\nu=\rho_{1}+v_{1}$ together with $u$ is a feasible solution to (16). With this solution,

[TABLE]

which implies that $\hat{V}_{1}(\rho_{0},\rho_{1})\leq V_{1}(\rho_{0},\rho_{1})$. This completes the proof. ∎

Theorem 4

The formula (14) defines a metric on ${\mathcal{H}}_{+}$.

Proof:

The proof follows exactly the same lines as in Theorem 2. ∎

Using the technique of Lagrangian multipliers one can deduce the dual formulation of (12) and establish the following:

Theorem 5

Notation as above. Then

[TABLE]

Proof:

Straight calculation gives

[TABLE]

This together with the strong duality completes the proof. ∎

VI Wasserstein-1 distance for matrix-valued densities

With little effort we are able to generalize the definition of Wasserstein-1 distance to the space of matrix-valued densities. Examples of matrix-valued densities include power spectra of multivariate time series, stress tensors, diffusion tensors and so on, and hence our motivation in considering matrix-valued distribution on possibly more than a one dimensional spatial coordinates.

Given two matrix-valued densities $\rho_{0},\rho_{1}$ satisfying

[TABLE]

we can define their Wasserstein-1 distance as

[TABLE]

or through its dual

[TABLE]

For more general densities where condition (19) may not be valid, we define

[TABLE]

or, equivalently,

[TABLE]

One can introduce positive coefficients $\beta_{1}>0,\beta_{2}>0$ to trade-off the relative importance of $u_{1}$ and $u_{2}$ in establishing correspondence between the two distributions as follows:

[TABLE]

or, equivalently,

[TABLE]

VII Example

We use our framework to compare power spectra of multivariate time series (in discrete time). Evidently, the distance between two power spectra induces a distance between corresponding linear modeling filters and, thereby, can be used to compare (stable) MIMO systems [22].

Consider the three power spectra as shown in Figure 1 (in different colors). What is shown in the three subplots are power spectra of two time series (in subplots (a) and (c)) and their cross-spectrum (in subplot (b)) as functions of time (the phase of the cross spectra are not shown). Thus, the three different colors represent the three different matrix-valued power spectra given by:

[TABLE]

where

[TABLE]

The distances between the each pair for different $\beta_{1},\beta_{2}$ values, $\alpha=1$, and the choice $L=[L_{1},\ L_{2}]$ with

[TABLE]

are tabulated in Table I. We observe that when the penalty on the rotation part is large ($\beta_{1}>>\beta_{2}$), we have $V_{1}(\rho_{0},\rho_{2})>V_{1}(\rho_{0},\rho_{1})$ and $V_{1}(\rho_{0},\rho_{2})>V_{1}(\rho_{2},\rho_{1})$. On the other hand, when the penalty on translation is large relative to the cost of rotation ($\beta_{1}<<\beta_{2}$), we have $V_{1}(\rho_{0},\rho_{1})>V_{1}(\rho_{0},\rho_{2})$ and $V_{1}(\rho_{0},\rho_{1})>V_{1}(\rho_{1},\rho_{2})$. These findings are in agreement with the intuition when observing the relative frequency directionality of power in the three spectra. More specifically, $\rho_{1}$ requires a significant drift in directionality before we can match it with the other two, while this is less important when comparing $\rho_{0}$ and $\rho_{2}$. For this latter case, it is the actual frequency where the power resides that distinguishes the two while the directionality is more in agreement.

What this example underscores is the ability of the metric to be tailored to applications where we need to trade off and compromise, in a principled way, between two vastly different features of matrix-valued distributions, i.e., spatial location versus directionality of the “intensity.” What was achieved in this paper is the construction of a suitable and easily computable metric that can be utilized for this purpose.

VIII Future research

We introduced generalization of the scalar $W_{1}$ distance to matrices and matrix-valued measures. This new metric, $W_{1}$, is computationally simpler and more attractive than earlier metrics, based on quadratic cost criteria. In fact, our “dual of the dual” formulation makes the metric especially attractive when comparing matrix-valued data on a discrete space (graph, network).

We note that the Wasserstein 1-metric has been used as a tool in defining curvature [24] and in analyzing the robustness of complex networks derived from scalar-valued data [27, 28]. The formalism presented in the current work, suggests alternative notions of curvature and robustness when the nodes of a network carry matrix-valued data, e.g., in diffusion tensor imaging. We plan to pursue such issues in future work.

Acknowledgements

This project was supported by AFOSR grants (FA9550-15-1-0045 and FA9550-17-1-0435), grants from the National Center for Research Resources (P41- RR-013218) and the National Institute of Biomedical Imaging and Bioengineering (P41-EB-015902), National Science Foundation (NSF), and National Institutes of Health (P30-CA-008748 and 1U24CA18092401A1).