Wasserstein distance in Optimal transport

Wasserstein Distance

Optimal transport

1. Notations

​ Consider two probability measures µ and ν defined on measure spaces X and Y . In most applications X and Y are subsets of ${\mathbb{R}}^d$ and µ and ν have density functions which we denote by $I_0$ and $I_1$, $d\mu (x)=I_0(x)dx$ and $dv(x)=I_1(x)dx$, (originally representing the height of a pile of soil/sand and the depth of an excavation).

2. Monge’s formulation

​ Monge’s optimal transportation problem is to find a measurable map f : X → Y that pushes µ onto ν and minimizes the following objective function,
$$M(\mu ,v)=in{f}_{f\in MP}{\int }_{x}c(x,f(x))d\mu (x)$$
Where $c:X\times Y\to {\mathbb{R}}^{+}$, is the cost functional, and M P = { f : X → Y ∣ f # μ = v } MP=\{f:X\rightarrow Y|f_\#\mu=v\} MP={f:X→Y∣f#​μ=v} represents the pushforward of measure µ and is characterized as, ${\int }_{f^{-1}(A)}d\mu (x)={\int }_{A}d\nu (y)$ for any measurable $A\subset Y$.

​ Simply put, the Monge formulation of the problem seeks the best pushforward map that rearranges measure µ into measure ν while minimizing a specific cost function.

Drawback:

1. Nonlinear with respect to f(x)
2. For certain measures the Monge’s formulation of the optimal transport problem is illposed; in the sense that there is no transport map to rearrange one measure to another. For instance, consider the case where µ is a Dirac mass while ν is not.

3. Kantorovich’s formulation

​ Kantorovich’s formulation alleviates this problem by finding the optimal transport plan as opposed to the transport map. Kantorovich formulated the transportation problem by optimizing over transportation plans, where a transport plan is a probability measure $\gamma \in P(X\times Y)$ with marginals µ and ν. The quantity $\gamma (A,B)$ tells us how much ‘mass’ in set A is being moved to set B. Let $\Gamma (\mu ,v)$ be the set of all such plans. Kantorovich’s formulation can then be written as,
$$K(\mu ,\nu )=\underset{\gamma \in \Gamma (\mu ,\nu )}{\min }{\int }_{X\times Y}c(x,y)d\gamma (x,y)$$
Note that unlike the Monge problem, in Kantorovich’s formulation the objective function and the constraints are linear with respect to $\gamma (x,y)$. Moreover, Kantorovich’s formulation is in the form of a convex optimization problem.

The Kantorovich problem is especially interesting in a discrete setting, that is for probability measures of the form $\mu ={\sum }_{i=1}^M{p}_i{\delta }_{x_i}$ and $v={\sum }_{j=1}^N{q}_j{\delta }_{x_j}$, where ${\delta }_{x_i}$ is a dirac measure centered at $x_i$. the Kantorovich problem can be written as,
$$\begin{array}{rl}K(\mu ,\nu )=& \underset{\gamma }{\min }\sum _i\sum _{j}c\left(x_i,y_j\right){\gamma }_{ij}\\ \text{ s.t. }& \sum _j{\gamma }_{ij}=p_i,\sum _i{\gamma }_{ij}=q_j\\ & {\gamma }_{ij}\ge 0,i=1,\dots ,M,j=1,\dots ,N\end{array}$$
where ${\gamma }_{ij}$ identifies how much of the mass particle $m_i$ at $x_i$ needs to be moved to $y_i$.

4. Wasserstein Distance

​ Let $\Omega$ be a subset of ${\mathbb{R}}^d$ on which the measures we consider are defined. In most applications $\Omega$ is the domain where the signal is defined and thus bounded. Let $P_p(\Omega )$ be the set of Borel probability measures on $\Omega$, with finite p’th moment, that is the set of probability measures $\mu$ on ${\mathbb{R}}^d$ such that ${\int }_{\Omega }|x{|}^{p}d\mu (x)<\infty$ The p-Wasserstein distance(metric), $W_p$ , for $p>=1$ on $P_p(\Omega )$ is then defined as using the Kantorovich’s formulation with the cost function $c(x,y)=|x-y{|}^p$, for $\mu$ and $v$ in $P_p(\Omega )$ ,
$$W_p(\mu ,\nu )={\left(\underset{\gamma \in \Gamma (\mu ,\nu )}{\inf }{\int }_{\Omega \times \Omega }|x-y{|}^{p}d\gamma (x,y)\right)}^{\frac{1}{p}}$$
For any $p\ge 0$, $W_p$ is a metric on $P_p(\Omega )$.The metric space $(P_p(\Omega ),W_p)$is referred to as the p-Wasserstein space.

Note that the p-Wasserstein metric can equivalently be defined using the dual Kantorovich problem,
$$W_p(\mu ,\nu )={\left(\underset{\varphi }{\sup }\left\{{\int }_{\Omega }\varphi (x)d\mu (x)-{\int }_{\Omega }{\varphi }^c(y)d\nu (y)\right\}\right)}^{\frac{1}{p}}$$
Where ϕ c ( y ) = inf ⁡ x { ϕ ( x ) − ∣ x − y ∣ p } \phi^{c}(y)=\inf _{x}\left\{\phi(x)-|x-y|^{p}\right\} ϕc(y)=infx​{ϕ(x)−∣x−y∣p}. For the specific case of p = 1 the p-Wasserstein metricis also known as the Monge–Rubinstein metric, or the earth mover distance.

Advantage:

1. 能够很自然地度量离散分布和连续分布之间的距离；
2. 不仅给出了距离的度量，而且给出如何把一个分布变换为另一分布的方案；
3. 能够连续地把一个分布变换为另一个分布，在此同时，能够保持分布自身的几何形态特征；(Geodesic) and (Barycenters)

Disadvantage:

1. 计算复杂度高，大多数情况下没有闭式解（除了一维和高斯分布）

The p-Wasserstein metric for one-dimensional probability measures is specifically interesting due to its simple and unique characterization.
$$\begin{array}{rl}{W}_p(\mu ,\nu )& ={\left({\int }_X|x,F_{\nu }^{-1}\left(F_{\mu }(x)\right){|}^{p}d\mu (x)\right)}^{\frac{1}{p}}\\ & ={\left({\int }_0^1|F_{\mu }^{-1}(z),F_{\nu }^{-1}(z){|}^{p}dz\right)}^{\frac{1}{p}}\end{array}$$
The closed-form solution of the p-Wasserstein distance in one dimension is an attractive property, as it alleviates the need for optimization. This property was employed in the Sliced Wasserstein distance as defined below.

5. Sliced Wasserstein distance

1. Radom Transform

简单来说，Radom Transform 阐述了一幅分布与其在各个角度下投影的具体表示。

将分布$f(x,y)$映射到$f(\alpha ,s)$.下图也是一个radom transfrom例子:

2. Sliced Wasserstein distance

   ​ The idea behind the sliced p-Wasserstein distance is to first, obtain a family of one-dimensional representations for a higher-dimensional probability distribution through linear projections (via the Radon transform), and then, calculate the distance between two input distributions as a functional on the p-Wasserstein distance of their one-dimensional representations. The sliced p-Wasserstein distance between $I_{\mu }$ and $I_v$ is then formally defined as:
   $$S{W}_p\left(I_{\mu },I_{\nu }\right)={\left({\int }_{{\mathbb{S}}^{d-1}}{W}_p^p\left(\mathcal{R}{I}_{\mu }(.,\theta ),\mathcal{R}{I}_{\nu }(.,\theta )\right)d\theta \right)}^{\frac{1}{p}}$$
   This is indeed a distance function as it satisfifies positive definiteness, symmetry and the triangle inequality.

3. Maximum sliced p-Wasserstein (max-SW)
   $$\max -S{W}_p\left(I_{\mu },I_{\nu }\right)=\underset{\theta \in {\mathbb{S}}^{d-1}}{\max }{W}_p\left(\mathcal{R}{I}_{\mu }(\cdot ,\theta ),\mathcal{R}{I}_{\nu }(\cdot ,\theta )\right)$$

6. Generalized Sliced-Wasserstein Distances

​ The GSW distance is obtained using the same procedure as for the SW distance, except that here, the one-dimensional representations are acquired through nonlinear projections.

1. Generalized Radon Transform
   
   $$\mathcal{G}I(t,\theta )={\int }_{{\mathbb{R}}^d}I(x)\delta (t-g(x,\theta ))dx$$

2. Generalized Sliced-Wasserstein and Maximum Generalized Sliced-Wasserstein Distances

   ​ Following the defifinition of the SW distance, we define the generalized sliced p-Wasserstein distance using the generalized Radon transform as:
   $$GS{W}_p\left(I_{\mu },I_{\nu }\right)={\left({\int }_{{\Omega }_{\theta }}{W}_p^p\left(\mathcal{G}{I}_{\mu }(\cdot ,\theta ),\mathcal{G}{I}_{\nu }(\cdot ,\theta )\right)d\theta \right)}^{\frac{1}{p}}$$
   Maximum Generalized Sliced-Wasserstein Distances:
   $$\max -GS{W}_p\left(I_{\mu },I_{\nu }\right)=\underset{\theta \in {\Omega }_{\theta }}{\max }{W}_p\left(\mathcal{G}{I}_{\mu }(\cdot ,\theta ),{\mathcal{G}}_{\nu }(\cdot ,\theta )\right)$$

3. Algorithm