Wassertein距离详解

该文是我投稿于PaperWeekly的一篇文章，如果大家有什么问题想要跟我讨论，可以留言跟我一起交流。

引言

   WGAN的横空出世引出了一个更好度量两个概率分布差异的指标即Wassertein距离（或叫做推土机距离），它主要优势就在于该距离具有连续性的特质。TV散度和JS散度的缺点在于这两个距离不具有连续性，这会导致在神经网络参数优化的过程中梯度会消失，KL散度则是因为该距离不具有对称性即$KL(p(x),p(y))\ne KL(p(y),p(x))$。本文会从Monge问题开始详细介绍关于Wassertein距离的一些相关背景和证明，最后给出了一个实例有助于理解Wassertein距离这个概念。

Monge问题(最优传输映射)

   一般情况下，假定$X$和$Y$是完备、可分的度量空间，例如欧式空间的子集$\Omega \subset {\mathbb{R}}^n$，通常是紧集。$\mathcal{P}(X)$代表$X$上所有概率测度构成的空间。

问题1（Monge问题或最优传输问题）给定两个概率测度$\mu \in \mathcal{P}(X)$，$\nu \in \mathcal{P}(Y)$和一个代价函数$c$：$X\times Y\to [0,+\infty ]$，求 $$(MP)\inf \left\{M(T):={\int }_{X}c(x,T(x){\}}_1\mu (x):T_{\#}\mu =\nu \right\},$$
其中由映射$T$诱导的推前测度$T_{\#}\mu$定义为：
$$\left(T_{\#}\mu \right)(A):=\mu \left(T^{-1}(A)\right)$$
这里$A\subset X$是任意可测集合，映射空间
$$\Sigma (\mu ,\nu ):=\left\{T:X\to Y\mid T_{\#}\mu =\nu \right\}$$

如下图示，Monge问题其实是在找一个最优传输映射$T$，映射$T$将空间$X$映射为空间$Y$，$A$是空间$Y$中的区域，$T^{-1}(A)$是$A$在$X$中原像。需要满足的条件是对于$X$空间中区域$T^{-1}(A)$的测度要与$Y$空间中区域$A$的测度要相等。根据下图通俗易懂的理解是面积$T^{-1}(A)$在$X$ 中的占比与面积$A$ 在$Y$中的相等。

Kantorovich问题(最优传输方案)

   对于Monge问题，由于需要求解一个映射，实际上只能允许多对一或者一对一(映射的定义)，Kantorovich
将传输映射放松成传输方案，允许一对多，从而将原来问题进行简化。]{}假定有一个联合概率分布$\gamma \in \mathcal{P}(X\times Y)$，其边际概率分布分别等于$\mu$
和$\nu$，${\left({\pi }_x\right)}_{\#}\gamma =\mu$，${\left({\pi }_y\right)}_{\#}\gamma =\nu$，这里${\pi }_x:X\times Y\to X$
和${\pi }_y:X\times Y\to Y$是投影映射：${\pi }_x(x,y)=x,{\pi }_y(x,y)=y$。

问题2（Kantorovich问题或最优传输方案问题）给定两个概率测度$\mu \in \mathcal{P}(X)$，$\nu \in \mathcal{P}(Y)$和一个代价函数$c$：$X\times Y\to [0,+\infty ]$，求解
$$(KP)\inf \left\{K(\gamma ):={\int }_{X\times Y}c(x,y)d\gamma (x,y):\gamma \in \Pi (\mu ,\nu )\right\},$$
其中联合概率测度$\gamma$属于传输方案空间
$$\Pi (\mu ,\nu ):=\left\{\gamma \in \mathcal{P}(X\times Y):{\left({\pi }_x\right)}_{\#}\gamma =\mu ,{\left({\pi }_y\right)}_{\#}\gamma =\nu \right\}.$$
如果一个最优传输映射为$T:X\times Y$，则最优传输方案$\gamma$有如下形式：
$$\gamma =(id,T{)}_{\#}\mu$$
给定离散的$\mu$和$\nu$，可能不存在任何传输映射：
$$\mu =\delta \left(x-x_0\right),\nu =\frac{1}{n}\sum _{i=1}^n\delta \left(y-y_i\right)$$
但传输方案总是存在的，例如$\mu \otimes \nu \in \Pi (\mu ,\nu )$。所以说Kantorovich问题是Monge问题的一个弱化和放松。如下图所示，图中左半部分是映射，映射的要求是$X$空间中的元素有且只有一个$Y$空间中的元素与之对应。图中的有右半部分是方案，方案可以是多多的形式。

Wassertein距离

Wassertein距离定义

   在深度学习中利用最优传输代价作为统计散度来反应两个概率密度之间的离散程度。主要用来表示两个概率分布的不相似程度，类似的统计散度有比较知名的KL散度，JS散度，还有一种度量概率分布很好方式就是利用Wassertein距离。如果将Section 3的Kantorovich问题的公式（2）里的代价函数$c(x,y)$换成欧式距离：$$c(x,y)=\Vert x-y{\Vert }_2^2$$ 就得到了2-Wassertein距离： $$(WD-2)\inf \left\{K(\gamma ):={\int }_{X\times Y}\Vert x-y{\Vert }_2^{2}d\gamma (x,y):\gamma \in \Pi (\mu ,\nu )\right\},$$
类似的，对于任意正整数$k$，k-Wasserstein距离可以定义为： $$(WD-k)\inf \left\{K(\gamma ):={\int }_{X\times Y}\Vert x-y{\Vert }_k^{k}d\gamma (x,y):\gamma \in \Pi (\mu ,\nu )\right\},$$
其中，$\Vert {\Vert }_k^k$是$L_k$范数，即$L_k=\sqrt[k]{{\sum }_{i=1}^n{x}_i^k}$，$\mathbf{x}=\left(x_1,x_2,\cdots ,x_n\right)$。
我们最熟知的推土机距离（Earth-Mover）为1-Wassertein距离： $$(WD-1)\inf \left\{K(\gamma ):={\int }_{X\times Y}\Vert x-y{\Vert }_1^{1}d\gamma (x,y):\gamma \in \Pi (\mu ,\nu )\right\},$$将公式（5）写成期望的形式为： $$W\left({\mathbb{P}}_r,{\mathbb{P}}_g\right)=\underset{\gamma \in \Pi \left({\mathbb{P}}_r,{\mathbb{P}}_g\right)}{\inf }{\mathbb{E}}_{(x,y)\sim \gamma }[\Vert x-y\Vert ]$$其中$\Pi \left({\mathbb{P}}_r,{\mathbb{P}}_g\right)$表示边缘分布${\mathbb{P}}_r$和${\mathbb{P}}_g$的所有联合分布$\gamma (x,y)$的集喝。$\gamma (x,y)$可以理解为将分布${\mathbb{P}}_r$转换成分布${\mathbb{P}}_g$所消耗的成本的方案。
我们熟知的WGAN的对抗目标损失其实是公式（6）的对偶形式，具体公式如下所示：
$$W\left({\mathbb{P}}_r,{\mathbb{P}}_{\theta }\right)=\underset{\Vert f{\Vert }_L\le 1}{\sup }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_{\theta }}[f(x)]$$
其中上确界在所有1-Lipschitz$f:\mathcal{X}\to \mathbb{R}$的上面。如果将$\Vert f{\Vert }_L\le 1$替换为$\Vert f{\Vert }_L\le K$，则会有一个参数化的函数${\left\{f_w\right\}}_{w\in \mathcal{W}}$，进而可以求解问题：
$$\underset{w\in \mathcal{W}}{\max }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}\left[f_w(x)\right]-{\mathbb{E}}_{z\sim p(z)}[f_w\left(g_{\theta }(z)\right]$$
其中，$w$为神经网络的参数。

Wassertein距离连续性

   Wassertein距离衡量分布的一个最大的好处就在于它距离度量的连续性。
证明：假定$\theta$和${\theta }^{\prime }$是${\mathbb{R}}^d$中的两个参数向量，其中联合的分布$\left(g_{\theta }(Z),g_{{\theta }^{\prime }}(Z)\right)$的集合为
$\gamma \in \Pi \left({\mathbb{P}}_{\theta },{\mathbb{P}}_{{\theta }^{\prime }}\right)$。由公式（6）的Wassertein距离的定义可知： $$\begin{array}{rl}W\left({\mathbb{P}}_{\theta },{\mathbb{P}}_{{\theta }^{\prime }}\right)& \le {\int }_{\mathcal{X}\times \mathcal{X}}\Vert x-y\Vert \mathrm{d}\gamma \\ & ={\mathbb{E}}_{(x,y)\sim \gamma }[\Vert x-y\Vert ]\\ & ={\mathbb{E}}_z\left[\Vert g_{\theta }(z)-g_{{\theta }^{\prime }}(z)\Vert \right]\end{array}$$
如果$g$在$\theta$中是连续的，那么$g_{\theta }(z){\to }_{\theta \to {\theta }^{\prime }}{g}_{{\theta }^{\prime }}(z)$，则有$\Vert g_{\theta }-g_{{\theta }^{\prime }}\Vert \to 0$，由于$\mathcal{X}$是紧的，其中任意两个元素的距离必须由某个常数$M$一致有界，因此对于所有$\theta$和$z$，$\Vert g_{\theta }(z)-g_{{\theta }^{\prime }}(z)\Vert \le M$一致。根据有界收敛定理，则有： $$W\left({\mathbb{P}}_{\theta },{\mathbb{P}}_{{\theta }^{\prime }}\right)\le {\mathbb{E}}_z\left[\Vert g_{\theta }(z)-g_{{\theta }^{\prime }}(z)\Vert \right]{\to }_{\theta \to {\theta }^{\prime }}0$$最终则有： $$\left|W\left({\mathbb{P}}_r,{\mathbb{P}}_{\theta }\right)-W\left({\mathbb{P}}_r,{\mathbb{P}}_{{\theta }^{\prime }}\right)\right|\le W\left({\mathbb{P}}_{\theta },{\mathbb{P}}_{{\theta }^{\prime }}\right){\to }_{\theta \to {\theta }^{\prime }}0$$
$W\left({\mathbb{P}}_r,{\mathbb{P}}_{\theta }\right)$连续性证明完毕。

Wassertein距离实例

   如下图所示考虑怎样将下面左图中的（实线）方块区域搬到右图中的（虚线）方块区域中，并且综合所移动距离最短。我们可以把这里的移动距离（或者传输距离）堪称是位置之间的距离差。比如左边的方块1在位置4，如果搬到右边的位置11，那么距离就是7。

   对于上图的问题，我们的一个直观的感受就是有很多种的搬运方块的方案。搬运方案一如下图所示，我们可以把位置1的方块A，放到位置7处，剩下两个方块B和C放到位置8中；位置2处的方块D和E放到位置9处；位置3的方块F放到位置9处，剩下的两个方块H和G放到位置10处；位置4的方块放到位置11处。搬运方案一总体的运输距离为：
$$6+7+7+7+7+6+7+7+7=61$$

   搬运方案二如下图所示，我们可以把位置1的方块A，放到位置11处，剩下两个方块B和C放到位置10中；位置2处的方块D和E放到位置9处；位置3的方块H放到位置[9]{}处，剩下的两个方块F和G放到位置8处；位置4的方块放到位置7处。搬运方案二总体的运输距离为：
$$10+9+9+7+7+6+5+5+3=61$$

   上面两种运输方案的运输距离是一样的，运输方案不一样，距离也会不一样。不同的运输方案也会有不同的距离。最优传输方案就是遍历所有的传输方案找出传输代价最小的传输方案。