Wasserstein Distance

https://zhuanlan.zhihu.com/p/58506295

今天来学一个数学知识，Wasserstein Distance。如果听说过WGAN的话，里面的W就是代表Wasserstein。

参考资料

Optimal Transport and Wasserstein Distance (slides)

特色

常见的有很多衡量概率分布差异的度量方式，比如total variation（TRPO推导里面有用到），还有经常被用到的KL散度。相比于这些度量方式，Wasserstein距离有如下一些好处。

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

过程

1. 其他距离度量的缺陷

首先注意到KL散度不是距离度量，它不满足对称性。常见的距离度量有

这些距离度量没法衡量离散分布和连续分布之间的距离：假设 $p$ 是均匀分布 $U[0, 1]$ 的概率密度， $q$ 是离散均匀分布 $\{0, 1/N, \cdots, 1\}$ 的概率密度。其total variation等于1，即完全不相似，但是凭感觉上来说，它们两个是很相似的。在Wasserstein距离度量下，它们的距离为 $1/N$ ，这看起来就比较合理了。
这些距离都忽略了概率分布之间的几何特性：它们几乎都有一个共同的特征，那就是都是对应点的概率密度函数相比较，这会忽略其几何特性。比如图1中，左边的分布应该离中间的分布更近，而中间的离右边的更远，但是其他度量无法反应这个特性，但Wasserstein距离可以。

基于Wasserstein距离可以找出Wasserstein平均（Wasserstein barycenter），相比于欧式平均（Eulidean average）来说，它更能够描述其形态特征，如图2所示。

Wasserstein距离不仅告诉两个分布之间的距离，而且能够告诉我们它们具体如何不一样，即如何从一个分布转化为另一个分布。如图3所示，Wasserstein能够告诉我们每一份probability density的转移方案。

这个转化过程还可以做成一个连续的过程，可以把A分布连续转化为B分布，并且这个转化过程是能够保持其几何特征的，如下面两图所示。

2. Wasserstein距离定义

Wasserstein距离的起源是optimal transport problem，把概率分布想象成一堆石子，如何移动一堆石子，通过最小的累积移动距离把它堆成另外一个目标形状，这就是optimal transport所关心的问题。

首先，要能完成这个操作，先要确保本来的这一堆石子的总质量要和目标石子堆总质量一样；考虑到概率分布的归一化条件，这一点是自然被满足的。

其次，我们暂时假设石子都是很小的，无限可分的；毕竟如果一个大石块要求堆成两座山仅仅通过移动肯定没法做到。（后面会有另外的定义方式来把大石块“劈开”）

假设地面上 $\mathcal{X} = \mathbb{R}^2$ 堆了一些石子，石子的分布我们用 $\mu: \mathcal{X} \to \mathbb{R}$ 来表示，采取这样的表示方法对于地面上的任意一块面积 $A \subseteq \mathcal{X}$ ， $\mu(A)$ 表示这块面积上放置了质量为多少的石子。同样的我们可以定义目标石子堆的分布 $\nu$ 。定义一个输运方案 $T:\mathcal{X} \to \mathcal{X}$ 把现有的石子堆变成目标石子堆。 $T(A)=B$ 表示把原来放在A处的石子都运到B处放好，类似地可以定义反函数 $T^{-1}(B)=A$ 。该输运方案成立需要满足 $\nu(B) = \mu(T^{-1}(B)), \ \forall B\subseteq\mathcal{X}$ ，即任意位置的石子通过输运过后都刚好满足分布 $\mu$ 的要求。这也可以写为 $T\#\mu = \nu$ 。

由此，两堆石子之间的距离可以被定义成把一堆石子挪动成另外一堆所需要的最小输运成本

$W_p(\mu, \nu) = \left( \inf_{T: T\#\mu=\nu} \int_{\mathcal{X}} ||x-T(x)||^p \mu(x) dx \right)^{1/p}$

遇到石子不可分的情况上述的定义就不再保证能存在可行的输运方案了，比如 $\mathcal{X}=\mathbb{R}$ ， $\mu(x) = \delta(x)$ ， $\nu(x) = \dfrac{1}{2}\delta(x-1)+\dfrac{1}{2}\delta(x+1)$ ，这样位于 $x=0$ 的质量为 $1$ 的“大石头”又不能劈开，那应该运到哪个位置呢？下面引入一个新的定义来解决这个问题。

下面的这个方案就是认为每个位置的石子都能够分开并且按照比例输运到不同的位置，先写出其定义

$W_p(\mu, \nu) = \left( \inf_{\gamma\in\Gamma(\mu, \nu)} \int_{\mathcal{X}\times\mathcal{X}} ||x-y||^p d\gamma(x,y) \right)^{1/p}$

其中 $\gamma$ 是一个联合概率分布，称coupling，它要求其边缘分布刚好是 $\mu$ 和 $\nu$ ，即 $\gamma(A\times\mathcal{X})=\mu(A)$ ， $\gamma(\mathcal{X}\times B)=\nu(B)$ ；如下图所示，它能够表示分布 $\mu$ 上的某个位置的质量被拆开（红竖线），然后再按照权重被分配给目标分布（红箭头）。

要注意到，图中画的这种联合概率分布是很难实际成为真实的optimal coupling的，一般来讲optimal coupling都会比较稀疏。排除少数不可分的情况，大多数的optimal coupling都是稀疏的，即 $\gamma(A\times B)=\mu(A\cap T^{-1}(B))$ 。

3. 对偶形式

最小化输运成本也可以写成对偶形式

注意到可以利用后面的约束关系推到类似前面定义的形式，它有弱对偶（weak duality）性质，即这个式子的最大值不超过Wasserstein距离定义中的最小值。

现在来考虑 $p=2$ 的形式，考虑到

做替换 $\varphi(x) = ||x||^2 - \phi(x)$ ， $\Psi(y) = ||y||^2 - \psi(y)$ ，可以得到如下形式

现在考虑把 $\Phi(y)$ 替换掉，这样就只有一个变量，在最优情况下有 $\Phi = \varphi^*$ 以及 $\varphi=\varphi^{**}$ ，其中

是 $\varphi(x)$ 的 Fenchel conjugate function。这样

优化问题就转变为了unconstrained问题， $\inf_\varphi \mathbb{E} \left[\varphi(X) +\varphi^*(Y)\right]$ ；
$\varphi=\varphi^{**}$ 表明 $\varphi$ 是一个凸函数；
$y=\nabla \varphi(x)$ ，注意到这刚好是输运函数 $y=T(x)$ 的形式，即输运函数为 $\varphi$ 的梯度；

另外注意到，Wasserstein距离不仅可以对两个分布来计算，也可以对从两个分布里面采样得到的样本来计算，如下示意图。

4. 从一个分布渐变到另一个分布（Geodesics）

找一个分布的路径 $c(t), t\in[0,1]$ 从 $\mu$ 到 $\nu$ ，其中 $c(0)=\mu$ ， $c(1)=\nu$ ，如果对于任意的 $t$ ， $W_p(c(0),c(t))+W_p(c(t),c(1))=W_p(c(0),c(1))$ ，那么这条路径就是最短路径。举例来说就是本文的图4。