给定两个定义在 R \mathbb{R} R 上的概率测度 μ 1 , μ 2 \mu_1, \mu_2 μ1,μ2,对应的累积分布函数分别为 F 1 ( x ) , F 2 ( x ) F_1(x), F_2(x) F1(x),F2(x),如何计算概率测度 μ 1 , μ 2 \mu_1, \mu_2 μ1,μ2之间的 Wasserstein Distance 呢?维基百科给出的计算:
乍一看很难理解,为什么可以这么算?
先看第一个式子,它是扫描 q q q 进行积分的, q = F 1 ( x 1 ) = F 2 ( x 2 ) q = F_1(x_1) = F_2(x_2) q=F1(x1)=F2(x2),即,扫描至当前 q q q 时,两个累计概分布函数值是一样的,都等于 q q q,如下图:
此时 x 1 = F 1 − 1 ( q ) , x 2 = F 2 − 1 ( q ) x_1 = F^{-1}_1(q), x_2 = F^{-1}_2(q) x1=F1−1(q),x2=F2−1(q),想把 d q dq dq 的 mass (绿色长条) 从 x 1 x_1 x1 运输至 x 2 x_2 x2,距离为 ∣ x 1 − x 2 ∣ = ∣ F 2 − 1 ( q ) − F 2 − 1 ( q ) ∣ |x_1 - x_2| = |F^{-1}_2(q) - F^{-1}_2(q)| ∣x1−x2∣=∣F2−1(q)−F2−1(q)∣,那么这次移动代价为 ∣ F 2 − 1 ( q ) − F 2 − 1 ( q ) ∣ d q |F^{-1}_2(q) - F^{-1}_2(q)|dq ∣F2−1(q)−F2−1(q)∣dq,扫描完整个 q ∈ [ 0 , 1 ] q \in [0,1] q∈[0,1],其结果就是第一个积分式。
再看第二个式子,它是扫描 x x x 进行积分的,如下图,扫描至 x x x 处时,两个概率分布的累积分布函数值分别为 F 1 ( x ) , F 2 ( x ) F_1(x), F_2(x) F1(x),F2(x),对应着已扫描过的空间 ( − ∞ , x ) (-\infin, x) (−∞,x) 的概率测度。
x x x 在 − ∞ -\infin −∞ 时,测度都为 0 0 0,向右移动一点点到 x x x 处,测度差为 ∣ F 1 ( x ) − F 2 ( x ) ∣ |F_1(x) - F_2(x)| ∣F1(x)−F2(x)∣,若想运输 mass 以使扫描过的测度相同(多推少挖),则需要的 mass 量为 ∣ F 1 ( x ) − F 2 ( x ) ∣ |F_1(x) - F_2(x)| ∣F1(x)−F2(x)∣,从哪运输呢?当然是从距离最近的 d x dx dx 处运来(或推走),代价为 ∣ F 1 ( x ) − F 2 ( x ) ∣ d x |F_1(x) - F_2(x)|dx ∣F1(x)−F2(x)∣dx。
归纳法,假设红线左侧的测度已经被填一样了,现在到了另一个
x
x
x 处,刚才就是在
x
x
x 处挖走的 mass,现在
x
x
x 处的 mass 差就是
∣
F
1
(
x
)
−
F
2
(
x
)
∣
|F_1(x) - F_2(x)|
∣F1(x)−F2(x)∣(还可以从右侧想,右侧的测度分文没动,左边又一样,这个差可不就是
∣
F
1
(
x
)
−
F
2
(
x
)
∣
|F_1(x) - F_2(x)|
∣F1(x)−F2(x)∣嘛)。依然从其右侧一点点挖,以使扫描过的部分测度一样,挖的量为
∣
F
1
(
x
)
−
F
2
(
x
)
∣
|F_1(x) - F_2(x)|
∣F1(x)−F2(x)∣,距离为
d
x
dx
dx,代价为
∣
F
1
(
x
)
−
F
2
(
x
)
∣
d
x
|F_1(x) - F_2(x)|dx
∣F1(x)−F2(x)∣dx。如此下去,直至
+
∞
+\infin
+∞。就得到了第二个积分式。
【多推少挖。其实不必担心从右边一点点处挖的时候不够挖,你可以假设后面的即使不够,它也从右边的右边挖过了,再不然假设从右端向左扫描,向左推挖,总能符合实际意义】
但是,这两个式子如何划等号呢?它们的直观解释差别还挺大的,前者是时刻保持两者扫描过的 q q q 一致,后者是时刻保持位置 x x x 一致。积分的变量代换也不好使。以后再说吧!
PyTorch 实现
连续的 One-dimensional Distribution 之间的 Wasserstein 距离是不好实现的, 因为的公式是积分. 对于均匀的离散分布, 或者在连续分布的离散采样, 即 α = 1 n ∑ i = 1 n δ x i , β = 1 n ∑ i = 1 n δ y i \alpha = \frac{1}{n}\sum_{i=1}^{n} \delta_{x_i}, ~ \beta = \frac{1}{n}\sum_{i=1}^{n} \delta_{y_i} α=n1∑i=1nδxi, β=n1∑i=1nδyi, 其中 x 1 ≤ ⋯ ≤ x n x_1 \le \cdots \le x_n x1≤⋯≤xn and y 1 ≤ ⋯ ≤ y n y_1 \le \cdots \le y_n y1≤⋯≤yn. 则可以按公式: W p ( α , β ) p = 1 n ∑ i = 1 n ∣ x i − y i ∣ p W_p(\alpha, \beta)^p = \frac{1}{n}\sum_{i=1}^{n}|x_i - y_i|^p Wp(α,β)p=n1i=1∑n∣xi−yi∣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
, 说明梯度正常传播了.
问: 为何是 1 n \frac{1}{n} n1?
在上面的公式 α = 1 n ∑ i = 1 n δ x i β = 1 n ∑ i = 1 n δ y i \begin{aligned} \alpha &= \frac{1}{n}\sum_{i=1}^{n} \delta_{x_i} \\ \beta &= \frac{1}{n}\sum_{i=1}^{n} \delta_{y_i} \end{aligned} αβ=n1i=1∑nδxi=n1i=1∑nδyi 中, 为何对连续分布的离散采样成了均匀分布? 这个好理解, 对于概率密度高的点 x 0 x_0 x0, 它出现在采样结果中的概率是高的, 10 10 10 个采样点中, 它可能出现了三次哦, 而概率密度低的点, 可能采样很多点才出现一次.
对于 W p ( α , β ) p = 1 n ∑ i = 1 n ∣ x i − y i ∣ p \begin{aligned} W_p(\alpha, \beta)^p &= \frac{1}{n}\sum_{i=1}^{n}|x_i - y_i|^p \end{aligned} Wp(α,β)p=n1i=1∑n∣xi−yi∣p 不就是维基百科中积分公式的离散版吗! 但为何是 1 n \frac{1}{n} n1? 因为 q = F ( x ) q=F(x) q=F(x) 是服从均匀分布 U ( 0 , 1 ) U(0,1) U(0,1) 的!