One-dimensional Wasserstein Distance

给定两个定义在 $\mathbb{R}$ 上的概率测度 ${\mu }_1,{\mu }_2$，对应的累积分布函数分别为 $F_1(x),F_2(x)$，如何计算概率测度 ${\mu }_1,{\mu }_2$之间的 Wasserstein Distance 呢？维基百科给出的计算：

乍一看很难理解，为什么可以这么算？

先看第一个式子，它是扫描 $q$ 进行积分的，$q=F_1(x_1)=F_2(x_2)$，即，扫描至当前 $q$ 时，两个累计概分布函数值是一样的，都等于 $q$，如下图：

此时 $x_1=F_1^{-1}(q),x_2=F_2^{-1}(q)$，想把 $dq$ 的 mass (绿色长条) 从 $x_1$ 运输至 $x_2$，距离为 $|x_1-x_2|=|F_2^{-1}(q)-F_2^{-1}(q)|$，那么这次移动代价为 $|F_2^{-1}(q)-F_2^{-1}(q)|dq$，扫描完整个 $q\in [0,1]$，其结果就是第一个积分式。

再看第二个式子，它是扫描 $x$ 进行积分的，如下图，扫描至 $x$ 处时，两个概率分布的累积分布函数值分别为 $F_1(x),F_2(x)$，对应着已扫描过的空间 $(-\infty ,x)$ 的概率测度。

$x$ 在 $-\infty$ 时，测度都为 $0$，向右移动一点点到 $x$ 处，测度差为 $|F_1(x)-F_2(x)|$，若想运输 mass 以使扫描过的测度相同（多推少挖），则需要的 mass 量为 $|F_1(x)-F_2(x)|$，从哪运输呢？当然是从距离最近的 $dx$ 处运来（或推走），代价为 $|F_1(x)-F_2(x)|dx$。

归纳法，假设红线左侧的测度已经被填一样了，现在到了另一个 $x$ 处，刚才就是在 $x$ 处挖走的 mass，现在 $x$ 处的 mass 差就是 $|F_1(x)-F_2(x)|$（还可以从右侧想，右侧的测度分文没动，左边又一样，这个差可不就是$|F_1(x)-F_2(x)|$嘛）。依然从其右侧一点点挖，以使扫描过的部分测度一样，挖的量为 $|F_1(x)-F_2(x)|$，距离为 $dx$，代价为 $|F_1(x)-F_2(x)|dx$。如此下去，直至 $+\infty$。就得到了第二个积分式。
【多推少挖。其实不必担心从右边一点点处挖的时候不够挖，你可以假设后面的即使不够，它也从右边的右边挖过了，再不然假设从右端向左扫描，向左推挖，总能符合实际意义】

但是，这两个式子如何划等号呢？它们的直观解释差别还挺大的，前者是时刻保持两者扫描过的 $q$ 一致，后者是时刻保持位置 $x$ 一致。积分的变量代换也不好使。以后再说吧！

PyTorch 实现

连续的 One-dimensional Distribution 之间的 Wasserstein 距离是不好实现的, 因为的公式是积分. 对于均匀的离散分布, 或者在连续分布的离散采样, 即 $\alpha =\frac{1}{n}{\sum }_{i=1}^n{\delta }_{x_i},\beta =\frac{1}{n}{\sum }_{i=1}^n{\delta }_{y_i}$, 其中 $x_1\le \cdots \le x_n$ and $y_1\le \cdots \le y_n$. 则可以按公式: $$W_p(\alpha ,\beta {)}^p=\frac{1}{n}\sum _{i=1}^n|x_i-y_i{|}^p$$ 计算. 代码:

def one_dimensional_Wasserstein(x: torch.Tensor, y: torch.Tensor, p: float):
	"""
	:param x: n 个 samples
	:param y: n 个 samples
	:param p: the exponent value in the formulation of Wasserstein distance
	:return: one_dimensional_Wasserstein^p
	"""
	distances = torch.abs(  # 排序, 相减
		torch.sort(x)[0] - torch.sort(y)[0]
	)
	wasserstein_distance_p = torch.pow(distances, p).mean()
	return torch.pow(wasserstein_distance_p, 1.0 / p)

可以看到, 计算过程含有排序, 那么当其作为损失函数时, 梯度能进行正常的反向传播吗? 答案是肯定的. 试验:

import torch

x = torch.tensor([1.0, 2.0, 3.0])
y = torch.tensor([3.3, 2.2, 1.1], requires_grad=True)
print(torch.sort(y))

optimozer = torch.optim.SGD(params=[y], lr=0.01)

for _ in range(20):
	sorted_y = torch.sort(y)[0]
	loss = torch.norm(x - sorted_y)
	optimozer.zero_grad()
	loss.backward()
	optimozer.step()
	print(y)

#################### output #######################
torch.return_types.sort(
values=tensor([1.1000, 2.2000, 3.3000], grad_fn=<SortBackward0>),
indices=tensor([2, 1, 0]))
tensor([3.2759, 2.1840, 1.0920], requires_grad=True)
tensor([3.2519, 2.1679, 1.0840], requires_grad=True)
tensor([3.2278, 2.1519, 1.0759], requires_grad=True)
tensor([3.2038, 2.1359, 1.0679], requires_grad=True)
tensor([3.1797, 2.1198, 1.0599], requires_grad=True)
tensor([3.1557, 2.1038, 1.0519], requires_grad=True)
tensor([3.1316, 2.0877, 1.0439], requires_grad=True)
tensor([3.1076, 2.0717, 1.0359], requires_grad=True)
tensor([3.0835, 2.0557, 1.0278], requires_grad=True)
tensor([3.0595, 2.0396, 1.0198], requires_grad=True)

可以看到, y 在逐渐靠近 x, 说明梯度正常传播了.

问: 为何是 $\frac{1}{n}$?

在上面的公式 $$\begin{array}{rl}\alpha & =\frac{1}{n}\sum _{i=1}^n{\delta }_{x_i}\\ \beta & =\frac{1}{n}\sum _{i=1}^n{\delta }_{y_i}\end{array}$$ 中, 为何对连续分布的离散采样成了均匀分布? 这个好理解, 对于概率密度高的点 $x_0$, 它出现在采样结果中的概率是高的, $10$ 个采样点中, 它可能出现了三次哦, 而概率密度低的点, 可能采样很多点才出现一次.

对于 $$\begin{array}{rl}{W}_p(\alpha ,\beta {)}^p& =\frac{1}{n}\sum _{i=1}^n|x_i-y_i{|}^p\end{array}$$ 不就是维基百科中积分公式的离散版吗! 但为何是 $\frac{1}{n}$? 因为 $q=F(x)$ 是服从均匀分布 $U(0,1)$ 的!