A Vectorial Approach to Unbalanced Optimal Mass Transport

The original formulation of OMT due to Monge [15, 16] may be expressed as follows:

$$\inf_{T}\{\int_{E}c(x,T(x))\rho_{0}(x)dx\ |\ T_{\#}\rho_{0}=\rho_{1}\},$$

(1)

where $c(x,y)$ is the cost of moving unit mass from $x$ to $y$, $T$ is the transport map, and $\rho_{0},\rho_{1}$ are two probability distributions defined on $E$ a subdomain of $\mathbb{R}^{n}$. $T_{\#}$ denotes the push forward of $T$.

As pioneered by Leonid Kantorovich [15, 16], the Monge formulation of OMT may be relaxed replacing transport maps $T$ by couplings $\pi$:

$$\inf_{\pi\in\Pi(\rho_{0},\rho_{1})}\int_{E}c(x,y)\pi(dx,dy),$$

(2)

where $\Pi(\rho_{0},\rho_{1})$ denotes the set of all the couplings between $\rho_{0}$ and $\rho_{1}$ (joint distributions whose two marginal distributions are $\rho_{0}$ and $\rho_{1}$). One may show that for $c(x,y)=\|x-y\|^{2}$ (square of distance function), the Kantorovich and Monge formulations are equivalent; see [15, 16] and the references therein.

Moreover for $c(x,y)=||x-y||^{2}$, the specific infimum is called Wasserstein 2 distance ($\mathcal{W}_{2}$). Benamou and Brenier pointed out that this can be written in an equivalent computational fluid dynamic formulation:

	$\displaystyle\mathcal{W}_{2}(\rho_{0},\rho_{1})^{2}=$	$\displaystyle\inf_{\rho,v}\int_{0}^{1}\int_{E}\rho(t,x)||v(t,x)||^{2}dxdt$		(3a)
		$\displaystyle\frac{\partial\rho}{\partial t}+\nabla\cdot(\rho v)=0$		(3b)
		$\displaystyle\rho(0,\cdot)=\rho_{0},\rho(1,\cdot)=\rho_{1}.$		(3c)

Thus, the latter reformulates OMT as one of minimizing a kinetic energy functional subject to a continuity constraint. The continuity constraint states that the change of density at each point over time is due to the flux of its neighborhood.

The Benamou-Brenier optimal solution may be related to the the original one of Monge. Indeed, if we have the optimal transport plan $T$ in (1), the interpolation can be expressed as $\rho(t,\cdot)=((t\cdot T+(1-t)\cdot id)_{\#}\rho_{0}$. For computational purposes, the Benamou-Brenier model is usually written as a convex optimization problem in momentum form ($p=\rho v$):

	$\displaystyle\mathcal{W}_{2}(\rho_{0},\rho_{1})^{2}=$	$\displaystyle\inf_{\rho,p}\int_{0}^{1}\int_{E}\frac{p(t,x)^{2}}{\rho(t,x)}dxdt$		(4a)
		$\displaystyle\frac{\partial\rho}{\partial t}+\nabla\cdot p=0$		(4b)
		$\displaystyle\rho(0,\cdot)=\rho_{0},\rho(1,\cdot)=\rho_{1}.$		(4c)

In the standard formulation, $\rho_{0},\rho_{1}$ need to be balanced, i.e., have the same total mass. We enforced this by simply taking them to be probability distributions ($\int_{E}\rho_{0}(x)dx=\int_{E}\rho_{1}(x)dx=1$). In the unbalanced case, total mass is not required to be preserved. This is very important for a number of applications in which mass may be created or destroyed.

There are many ways to extend the original setting to the unbalanced case [10]. In the present work, we will just treat the case in which a source term is added under the Fisher-Rao smooth formulation [9]. Namely,

	$\displaystyle\mathcal{W}_{FR}(\rho_{0},\rho_{1})^{2}=$	$\displaystyle\inf_{\rho,p}\int_{0}^{1}\int_{E}\frac{p(t,x)^{2}}{\rho(t,x)}+\gamma\frac{s(t,x)^{2}}{\rho(t,x)}dxdt$		(5a)
		$\displaystyle\frac{\partial\rho}{\partial t}+\nabla\cdot p=s$		(5b)
		$\displaystyle\rho(0,\cdot)=\rho_{0},\rho(1,\cdot)=\rho_{1}.$		(5c)

There is an extra source term $s$ in the continuity equation. The change of density at each point over time now is not only due to the flux of its neighborhood, but also a source. Mass can be created or destroyed at any point and time. The parameter $\gamma$ is a weight to control how much source one wants to use.

Vector-valued OMT considers vector-valued distributions instead of only the classical scalar ones [6]. With more than one channel, the change of mass is not only due to the flux of its neighborhood but also the redistribution within the given channels. Namely, we have

	$\displaystyle\mathcal{W}_{V}(\rho_{0},\rho_{1})^{2}=$	$\displaystyle\inf_{\rho,p,u}\int_{0}^{1}\int_{E}p^{T}{\rm diag}(\rho)^{-1}p+\gamma u^{T}[{\rm diag}(\mathbb{F}^{T}_{2}\rho)^{-1}+{\rm diag}(\mathbb{F}^{T}_{1}\rho)^{-1}]u\ dxdt$		(6a)
		$\displaystyle\frac{\partial\rho}{\partial t}+\nabla_{x}\cdot p-\nabla_{\mathcal{F}}^{*}u=0$		(6b)
		$\displaystyle\rho(0,\cdot)=\rho_{0},\quad\rho(1,\cdot)=\rho_{1}.$		(6c)

As mentioned above, there are two kinds of flows for the vector-valued density $\rho$. The first is the spatial flux $p$ and its corresponding divergence is denoted as $\nabla_{x}\cdot$, and the second is the flux $u$ between channels, whose discrete divergence is denoted by $\nabla_{\mathcal{F}}^{*}$. $p$ and $u$ are both vectors. We note that $\mathbb{F}_{1}$ is the all-1 part of the incident matrix $\mathbb{F}$ (sources of all the edges) and $\mathbb{F}_{2}=\mathbb{F}_{1}-\mathbb{F}$ (sinks of all the edges). As $u$ is defined on every edge and the density is only defined on nodes, ${\rm diag}(\mathbb{F}^{T}_{2}\rho)^{-1}+{\rm diag}(\mathbb{F}^{T}_{1}\rho)^{-1}$ assigns a density to each edge using the density of two end points so that we can compute the kinetic energy for the flows on edges.

Vector-valued OMT can be considered as a general OMT problem on $E\times F$, where $E$ is the space on which the vector-valued density is defined and $F$ is a connected graph. With this understanding, we can rewrite the energy functional in the following equivalent form:

$$\mathcal{W}_{V}(\rho_{0},\rho_{1})^{2}=\inf_{\rho,p,u}\int_{0}^{1}\int_{E}\sum_{c\in Nodes(F)}\frac{p(t,x,c)^{2}}{\rho(t,x,c)}dxdt+\gamma\int_{0}^{1}\int_{E}\sum_{e\in Edges(F)}\frac{u(t,x,e)^{2}}{\tilde{\rho}(t,x,e)}dxdt,$$

(7)

where $\tilde{\rho}$ is density of edge.

In this section, we write down the unbalanced vector-valued version of OMT. Namely,

	$\displaystyle\mathcal{W}_{VS}(\rho_{0},\rho_{1})^{2}=$	$\displaystyle\inf_{\rho,p,u}\int_{0}^{1}\int_{E}p^{T}{\rm diag}(\rho)^{-1}p+\gamma u^{T}[{\rm diag}(\mathbb{F}^{T}_{2}\rho)^{-1}+{\rm diag}(\mathbb{F}^{T}_{1}\rho)^{-1}]u+\eta s^{T}{\rm diag}(\rho)^{-1}s\ dxdt$		(8a)
		$\displaystyle\frac{\partial\rho}{\partial t}+\nabla_{x}\cdot p-\nabla_{\mathcal{F}}^{*}u=s$		(8b)
		$\displaystyle\rho(0,\cdot)=\rho_{0},\quad\rho(1,\cdot)=\rho_{1}.$		(8c)

Note that the source term is also a vector. So we need to write the energy functional for this vector term. Regarding the continuity equation, there are three possibilities for the change of mass over time: the flux of its neighborhood, flows between channels and source.

Here the input source and target are regarded as two mass distributions on the subdomain $E$. Total mass may or may not be preserved. We add an extra source layer which is parallel to original space. For each point in $E$, there is an edge connecting with source layer. As now we have an extra layer, we can put the difference of mass into the new layer so that the total mass of the 2-vector new structure is preserved. The difference of mass can just be distributed uniformly to source layer or put to some area of interest for further specific applications.

Now we view the flow between channels as source term. Consider the continuity equations of (5b) and (6b):

	$\displaystyle\frac{\partial\rho}{\partial t}+\nabla\cdot p=$	$\displaystyle s$
	$\displaystyle\frac{\partial\rho}{\partial t}+\nabla_{x}\cdot p=$	$\displaystyle\nabla_{\mathcal{F}}^{*}u.$

We can use $\nabla_{\mathcal{F}}^{*}u$ as the source term $s$. Because of the special graph structure (there is only one edge), we have:

$$\nabla_{\mathcal{F}}^{*}u=u=s.$$

(9)

Now consider the second term in (7). Since there is only one edge, we can just omit the summation:

$$\gamma\int_{0}^{1}\int_{E}\sum_{e\in Edges(F)}\frac{u(t,x,e)^{2}}{\tilde{\rho}(t,x,e)}dxdt=\gamma\int_{0}^{1}\int_{E}\frac{u(t,x)^{2}}{\tilde{\rho}(t,x)}dxdt.$$

(10)

If we further take the density of each edge $\tilde{\rho}(t,x)=\rho(t,x)$, this term is exactly the same as the second term in the energy functional in unbalanced OMT setting.

Now we consider the first term of the energy functional:

$$\int_{0}^{1}\int_{E}\sum_{c\in Nodes(F)}\frac{p(t,x,c)^{2}}{\rho(t,x,c)}dxdt=\int_{0}^{1}\int_{E}\frac{p(t,x,c_{1})^{2}}{\rho(t,x,c_{1})}+\frac{p(t,x,c_{2})^{2}}{\rho(t,x,c_{2})}dxdt.$$

(11)

There are only two channels, i.e., $c_{1}$ denotes the original channel and $c_{2}$ denotes the source layer. Notice that the integral with respect to $c_{1}$ term is exactly the first term of (5). We want to get rid of the $c_{2}$ integral. This can be done by introducing a small weight parameter for that layer.

In the original setting, one treats the kinetic energy of different layers in an identical manner. But by simply introducing a weight, we can treat each layer differently:

$$\mathcal{W}_{V^{\prime}}(\rho_{0},\rho_{1})^{2}=\inf_{\rho,p,u}\int_{0}^{1}\int_{E}\sum_{c\in Nodes(F)}w(c)\frac{p(t,x,c)^{2}}{\rho(t,x,c)}dxdt+\gamma\int_{0}^{1}\int_{E}\sum_{e\in Edges(F)}\frac{u(t,x,e)^{2}}{\tilde{\rho}(t,x,e)}dxdt,$$

(12)

where $w(c)$ is the weighting parameters. We can choose different values for different channels. For another form of vector OMT:

$$\mathcal{W}_{V^{\prime}}(\rho_{0},\rho_{1})^{2}=\inf_{\rho,p,u}\int_{0}^{1}\int_{E}p^{T}{\rm diag}(w){\rm diag}(\rho)^{-1}p+\gamma u^{T}[{\rm diag}(\mathbb{F}^{T}_{2}\rho)^{-1}+{\rm diag}(\mathbb{F}^{T}_{1}\rho)^{-1}]u\ dxdt.$$

(13)

Under this setting, if we take $w(c_{1})=1$ and $w(c_{2})$ very small, then the integral in the source layer is very small in that energy functional. Clearly, unbalanced OMT is almost equivalent to the extra source layer vector OMT with the latter set of weight parameters.

By introducing the source layer, there is an edge between two channels for which we use the flow on that edge as source, and thus there is an extra integral on that new layer. We add very small weight parameters so that two energy functionals are almost identical. See Figure 1.

It is quite straightforward to implement unbalanced OMT from vector-valued OMT code. We need to only alter a few lines of code to make it work for unbalanced case:

1. 1.

   Initialization: From the input, first construct two extended structures for vector OMT. The first layer is the original input and the second layer contains the mass difference. If the starting density distribution has more total mass, then the mass difference is added to the second layer of target density distribution. If the starting density distribution has less total mass, then the mass difference is added to the second layer of starting density distribution itself.

2. 2.

   Set weighting parameters: Employing the code for the energy functional, just add weighting parameters to the corresponding layers.

In the vector-valued case, the input source and target are two $n$-vector valued distributions. Total mass may or may not be preserved. We add an extra source layer which is connected to each of the existing layers. As before, we can put the difference of mass into the new layer so that the total mass of the $(n+1)$-vector new structures is preserved.

We call the incident matrix for those newly added edges $\mathcal{G}$, which is a submatrix of the new large incident matrix $\tilde{\mathcal{F}}$. We just split $\tilde{\mathcal{F}}$ for illustration purpose. They can be easily pieced together as in original vector setting:

$$\nabla_{\tilde{\mathcal{F}}}^{*}\left[\begin{array}[]{c}u\\ v\end{array}\right]=\left[\nabla_{\mathcal{F}}^{*}\ \nabla_{\mathcal{G}}^{*}\right]\left[\begin{array}[]{c}u\\ v\end{array}\right]=\nabla_{\mathcal{F}}^{*}u+\nabla_{\mathcal{G}}^{*}v,$$

(14)

where $u$ is the flow within the existing edges and $v$ is the flow within the newly added edges.

Now we consider the continuity equations of unbalanced vector OMT and vector OMT with a new layer:

	$\displaystyle\frac{\partial\rho}{\partial t}+\nabla\cdot p-\nabla_{\mathcal{F}}^{*}u=$	$\displaystyle s$
	$\displaystyle\frac{\partial\rho}{\partial t}+\nabla_{x}\cdot p-\nabla_{\mathcal{F}}^{*}u=$	$\displaystyle\nabla_{\mathcal{G}}^{*}v.$

We can use $\nabla_{\mathcal{G}}^{*}v$ as the source $s$. Now similar to the scalar case:

$$\nabla_{\mathcal{G}}^{*}v=v=s.$$

(15)

The energy functional now becomes:

	$\displaystyle\mathcal{W}_{V^{\prime\prime}}(\rho_{0},\rho_{1})^{2}=\inf_{\rho,p,u,v}\int_{0}^{1}\int_{E}\sum_{c\in Nodes(\tilde{\mathcal{F}})}\frac{p(t,x,c)^{2}}{\rho(t,x,c)}dxdt+$	$\displaystyle\gamma\int_{0}^{1}\int_{E}\sum_{e\in Edges(\mathcal{F})}\frac{u(t,x,e)^{2}}{\tilde{\rho}(t,x,e)}dxdt$		(16)
	$\displaystyle+$	$\displaystyle\eta\int_{0}^{1}\int_{E}\sum_{e\in Edges(\mathcal{G})}\frac{v(t,x,e)^{2}}{\tilde{\rho}(t,x,e)}dxdt.$

With the above relationship (15), it is easy to see that the last term exactly fits the original source energy term (8a).

Again there is a new part of kinetic energy due to the newly added source layer. We can introduce weight parameters to make the extra term almost disappear. Same as adding weight technique for different layers, we can put $\gamma$ and $\eta$ together as a weight matrix for different edges so that the final form is more concise:

	$\displaystyle\mathcal{W}_{V^{\prime\prime}}(\rho_{0},\rho_{1})^{2}=$	$\displaystyle\inf_{\rho,p,\tilde{u}}\int_{0}^{1}\int_{E}p^{T}{\rm diag}(w_{1}){\rm diag}(\rho)^{-1}p+\tilde{u}^{T}{\rm diag}(w_{2}){\rm diag}(\tilde{\rho})^{-1}\tilde{u}\ dxdt$		(17a)
		$\displaystyle\frac{\partial\rho}{\partial t}+\nabla_{x}\cdot p-\nabla_{\tilde{\mathcal{F}}}^{*}\tilde{u}=0$		(17b)
		$\displaystyle\rho(0,\cdot)=\rho_{0},\quad\rho(1,\cdot)=\rho_{1},$		(17c)

where ${\rm diag}(w_{1})$ denotes the weighting matrix for different layers and ${\rm diag}(w_{2})$ denotes the weighting matrix for the edges we add (weight $=\eta$) and existing edges (weight$=\gamma$). $\tilde{u}=\left[\begin{array}[]{c}u\\ v\end{array}\right]$ is all the flows that between channels within the new large graph $\tilde{\mathcal{F}}$ and $\tilde{\rho}$ is the density assigned to each edge.

By introducing a source layer, there is an edge between the source layer and each existing channel. We use the flows on those edges as sources. As before, we add very small weight parameters for the source layer so that two energy functionals are almost identical. See Figure 2.

As before, it is also very easy to implement unbalanced vector OMT from vector OMT code.

1. 1.

   Initialization: From the input, first construct the extended structures for vector OMT. Add an extra layer. Connect that new layer with each of other layers. Put the difference of total mass into that newly added layer.

2. 2.

   Set weight parameters: Employing the code for the energy functional, add weighting parameters to corresponding layer and corresponding edges (existing edges and newly added edges).

With the simple change of the original vector-valued OMT, we can use the same code for unbalanced vector OMT case without even touching the main structure of the original code.

We tested our approach on color image data. While the interpolations of density are color, the source layer is still gray scale.