PyTorch 实战:计算 Wasserstein 距离

PyTorch 实战:计算 Wasserstein 距离

2019-09-23 18:42:56 html

This blog is copied from: https://mp.weixin.qq.com/s/nTUKYNxdiPK3xdOoSXvTJQ git

 

最优传输理论及 Wasserstein 距离是不少读者都但愿了解的基础,本文主要经过简单案例展现了它们的基本思想,并经过 PyTorch 介绍如何实战 W 距离。github

 

机器学习中的许多问题都涉及到令两个分布尽量接近的思想,例如在 GAN 中令生成器分布接近判别器分布就能伪造出逼真的图像。可是 KL 散度等分布的度量方法有不少局限性,本文则介绍了 Wasserstein 距离及 Sinkhorn 迭代方法,它们 GAN 及众多任务上都展现了杰出的性能。算法

 

在简单的状况下,咱们假设从未知数据分布 p(x) 中观测到一些随机变量 x(例如,猫的图片),咱们想要找到一个模型 q(x|θ)(例如一个神经网络)能做为 p(x) 的一个很好的近似。若是 p 和 q 的分布很相近,那么就代表咱们的模型已经学习到如何识别猫。网络

 

由于 KL 散度能够度量两个分布的距离,因此只须要最小化 KL(q‖p) 就能够了。能够证实,最小化 KL(q‖p) 等价于最小化一个负对数似然,这样的作法在咱们训练一个分类器时很常见。例如,对于变分自编码器来讲,咱们但愿后验分布可以接近于某种先验分布,这也是咱们经过最小化它们之间的 KL 散度来实现的。框架

 

尽管 KL 散度有很普遍的应用,在某些状况下,KL 散度则会失效。不妨考虑一下以下图所示的离散分布:dom

 

 

KL 散度假设这两个分布共享相同的支撑集(也就是说,它们被定义在同一个点集上)。所以,咱们不能为上面的例子计算 KL 散度。因为这一个限制和其余计算方面的因素促使研究人员寻找一种更适合于计算两个分布之间差别的方法。机器学习

 

在本文中,做者将:函数

 

  • 简单介绍最优传输问题性能

  • 将 Sinkhorn 迭代描述为对解求近似

  • 使用 PyTorch 计算 Sinkhorn 距离

  • 描述用于计算 mini-batch 之间的距离的对该实现的扩展

 

移动几率质量函数

 

咱们不妨把离散的几率分布想象成空间中分散的点的质量。咱们能够观测这些带质量的点从一个分布移动到另外一个分布须要作多少功,以下图所示:

 

 

接着,咱们能够定义另外一个度量标准,用以衡量移动作全部点所须要作的功。要想将这个直观的概念形式化定义下来,首先,咱们能够经过引入一个耦合矩阵 P(coupling matrix),它表示要从 p(x) 支撑集中的一个点上到 q(x) 支撑集中的一个点须要分配多少几率质量。对于均匀分布,咱们规定每一个点都具备 1/4 的几率质量。若是咱们将本例支撑集中的点从左到右排列,咱们能够将上述的耦合矩阵写做:

 

 

也就是说,p(x) 支撑集中点 1 的质量被分配给了 q(x) 支撑集中的点 4,p(x) 支撑集中点 2 的质量被分配给了 q(x) 支撑集中的点 3,以此类推,如上图中的箭头所示。

 

为了算出质量分配的过程须要作多少功,咱们将引入第二个矩阵:距离矩阵。该矩阵中的每一个元素 C_ij 表示将 p(x) 支撑集中的点移动到 q(x) 支撑集中的点上的成本。点与点之间的欧几里得距离是定义这种成本的一种方式,它也被称为「ground distance」。若是咱们假设 p(x) 的支撑集和 q(x) 的支撑集分别为 {1,2,3,4} 和 {5,6,7,8},成本矩阵即为:

 

 

根据上述定义,总的成本能够经过 P 和 C 之间的 Frobenius 内积来计算:

 

 

你可能已经注意到了,实际上有不少种方法能够把点从一个支撑集移动到另外一个支撑集中,每一种方式都会获得不一样的成本。上面给出的只是一个示例,可是咱们感兴趣的是最终可以让成本较小的分配方式。这就是两个离散分布之间的「最优传输」问题,该问题的解是全部耦合矩阵上的最低成本 L_C。

 

因为不是全部矩阵都是有效的耦合矩阵,最后一个条件会引入了一个约束。对于一个耦合矩阵来讲,其全部列都必需要加到带有 q(x) 几率质量的向量中。在本例中,该向量包含 4 个值为 1/4 的元素。更通常地,咱们能够将两个向量分别记为 a 和 b,所以最有运输问题能够被写做:

 

 

当距离矩阵基于一个有效的距离函数构建时,最小成本即为咱们所说的「Wasserstein 距离」。

 

关于该问题的解以及将其扩展到连续几率分布中还有大量问题须要解决。若是想要获取更正式、更容易理解的解释,读者能够参阅 Gabriel Peyré 和 Marco Cuturi 编写的「Computational Optimal Transport」一书,此书也是本文写做的主要参考来源之一。

 

这里的基本设定是,咱们已经把求两个分布之间距离的问题定义为求最优耦合矩阵的问题。事实证实,咱们能够经过一个小的修改让咱们以迭代和可微分的方式解决这个问题,这将让咱们能够很好地使用深度学习自动微分机制完成该工做。

 

熵正则化和 Sinkhorn 迭代

 

首先,咱们将一个矩阵的熵定义以下:

 

 

正如信息论中几率分布的熵同样,一个熵较低的矩阵将会更稀疏,它的大部分非零值集中在几个点周围。相反,一个具备高熵的矩阵将会更平滑,其最大熵是在均匀分布的状况下得到的。咱们能够将正则化系数 ε 引入最优传输问题,从而获得更平滑的耦合矩阵:

 

 

经过增大 ε,最终获得的耦合矩阵将会变得更加平滑;而当 ε 趋近于零时,耦合矩阵会更加稀疏,同时最终的解会更加趋近于原始最优运输问题。

 

经过引入这种熵正则化,该问题变成了一个凸优化问题,而且可 以经过使用「Sinkhorn iteration」求解。解能够被写做 P=diag(u)Kdiag(v),在迭代过程当中交替更新 u 和 v:

 

 

其中 K 是一个用 C 计算的核矩阵(kernel matrix)。因为这些迭代过程是在对原始问题的正则化版本求解,所以对应产生的 Wasserstein 距离有时被称为 Sinkhorn 距离。该迭代过程会造成一个线性操做的序列,所以对于深度学习模型,经过这些迭代进行反向传播是很是简单的。

 

经过 PyTorch 实现 Sinkhorn 迭代

 

为了提高 Sinkhorn 迭代的收敛性和稳定性,还能够加入其它的步骤。咱们能够在 GitHub 上找到 Gabriel Peyre 完成的详细实现。

 

项目连接:https://github.com/gpeyre/SinkhornAutoDiff。

 

让咱们先用一个简单的例子来测试一下,如今咱们将研究二维空间(而不是上面的一维空间)中的离散均匀分布。在这种状况下,咱们将在平面上移动几率质量。让咱们首先定义两个简单的分布:

 

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
np.random.seed(42)

n_points = 5
a = np.array([[i, 0] for i in range(n_points)])
b = np.array([[i, 1] for i in range(n_points)])

plt.figure(figsize=(6, 3))
plt.scatter(a[:, 0], a[:, 1], label='supp($p(x)$)')
plt.scatter(b[:, 0], b[:, 1], label='supp($q(x)$)')
plt.legend();

 

 

咱们很容易看出,最优传输对应于将 p(x) 支撑集中的每一个点分配到 q(x) 支撑集上的点。对于全部的点来讲,距离都是 1,同时因为分布是均匀的,每点移动的几率质量是 1/5。所以,Wasserstein 距离是 5×1/5= 1。如今咱们用 Sinkhorn 迭代来计算这个距离:

 

import torch
from layers import SinkhornDistance

x = torch.tensor(a, dtype=torch.float)
y = torch.tensor(b, dtype=torch.float)

sinkhorn = SinkhornDistance(eps=0.1, max_iter=100, reduction=None)
dist, P, C = sinkhorn(x, y)
print("Sinkhorn distance: {:.3f}".format(dist.item()))

————————————————————————————————————————————————
Sinkhorn distance: 1.000

 

结果正如咱们所计算的那样,距离为 1。如今,让咱们查看一下「Sinkhorn( )」方法返回的矩阵,其中 P 是计算出的耦合矩阵,C 是距离矩阵。距离矩阵以下图所示:

 

plt.imshow(C)
plt.title('Distance matrix')
plt.colorbar();
plt.imshow(C)plt.title('Distance matrix')plt.colorbar();

 

元素「C[0, 0]」说明了将(0,0)点的质量移动到(0,1)所须要的成本 1 是如何产生的。在该行的另外一端,元素「C[0, 4]」包含了将点(0,0)的质量移动到点(4,1)所须要的成本,这个成本是整个矩阵中最大的:

 

 

因为咱们为距离矩阵使用的是平方后的 ℓ2 范数,计算结果如上所示。如今,让咱们看看计算出的耦合矩阵吧:

 

plt.imshow(P)
plt.title('Coupling matrix');
plt.imshow(P)plt.title('Coupling matrix');

 

该图很好地向咱们展现了算法是如何有效地发现最优耦合,它与咱们前面肯定的耦合矩阵是相同的。到目前为止,咱们使用了 0.1 的正则化系数。若是将该值增长到 1 会怎样?

 

sinkhorn = SinkhornDistance(eps=1, max_iter=100, reduction=None)
dist, P, C = sinkhorn(x, y)
print("Sinkhorn distance: {:.3f}".format(dist.item()))
plt.imshow(P);

————————————————————————————————————————————————
Sinkhorn distance: 1.408

 

正如咱们前面讨论过的,加大 ε 有增大耦合矩阵熵的做用。接下来,咱们看看 P 是如何变得更加平滑的。可是,这样作也会为计算出的距离带来一个很差的影响,致使对 Wasserstein 距离的近似效果变差。

 

可视化支撑集的空间分配也颇有意思:

 

def show_assignments(a, b, P):    
    norm_P = P/P.max()
    for i in range(a.shape[0]):
        for j in range(b.shape[0]):
            plt.arrow(a[i, 0], a[i, 1], b[j, 0]-a[i, 0], b[j, 1]-a[i, 1],
                     alpha=norm_P[i,j].item())
    plt.title('Assignments')
    plt.scatter(a[:, 0], a[:, 1])
    plt.scatter(b[:, 0], b[:, 1])
    plt.axis('off')

show_assignments(a, b, P)

 

让咱们在一个更有趣的分布(Moons 数据集)上完成这项工做。

 

from sklearn.datasets import make_moons

X, Y = make_moons(n_samples = 30)
a = X[Y==0]
b = X[Y==1]

x = torch.tensor(a, dtype=torch.float)
y = torch.tensor(b, dtype=torch.float)

sinkhorn = SinkhornDistance(eps=0.1, max_iter=100, reduction=None)
dist, P, C = sinkhorn(x, y)
print("Sinkhorn distance: {:.3f}".format(dist.item()))
show_assignments(a, b, P)

——————————————————————————————————————————
Sinkhorn distance: 1.714

Mini-batch 上的 Sinkhorn 距离

 

在深度学习中,咱们一般对使用 mini-batch 来加速计算十分感兴趣。咱们也能够经过使用额外的批处理维度修改 Sinkhorn 迭代来知足该设定。将此更改添加到具体实现中后,咱们能够在一个 mini-batch 中计算多个分布的 Sinkhorn 距离。下面咱们将经过另外一个容易被验证的例子说明这一点。

 

代码:https://github.com/dfdazac/wassdistance/blob/master/layers.py

 

咱们将计算包含 5 个支撑点的 4 对均匀分布的 Sinkhorn 距离,它们垂直地被 1(如上所示)、二、3 和 4 个单元分隔开。这样,它们之间的 Wasserstein 距离将分别为 一、四、9 和 16。

 

n = 5
batch_size = 4
a = np.array([[[i, 0] for i in range(n)] for b in range(batch_size)])
b = np.array([[[i, b + 1] for i in range(n)] for b in range(batch_size)])

# Wrap with torch tensors
x = torch.tensor(a, dtype=torch.float)
y = torch.tensor(b, dtype=torch.float)

sinkhorn = SinkhornDistance(eps=0.1, max_iter=100, reduction=None)
dist, P, C = sinkhorn(x, y)
print("Sinkhorn distances: ", dist)

——————————————————————————————————————————
Sinkhorn distances:  tensor([ 1.0001,  4.0001,  9.0000, 16.0000])

 

这样作确实有效!同时,也请注意,如今 P 和 C 为 3 维张量,它包含 mini-batch 中每对分布的耦合矩阵和距离矩阵:

 

print('P.shape = {}'.format(P.shape))
print('C.shape = {}'.format(C.shape))

——————————————————————————————————————————
P.shape = torch.Size([4, 5, 5])
C.shape = torch.Size([4, 5, 5])

 

结语

 

分布之间的 Wasserstein 距离及其经过 Sinkhorn 迭代实现的计算方法为咱们带来了许多可能性。该框架不只提供了对 KL 散度等距离的替代方法,并且在建模过程当中提供了更大的灵活性,咱们再也不被迫要选择特定的参数分布。这些迭代过程能够在 GPU 上高效地执行,而且是彻底可微分的,这使得它对于深度学习来讲是一个很好的选择。这些优势在机器学习领域的最新研究中获得了充分的利用(如自编码器和距离嵌入),使其在该领域的应用前景更加广阔。

 

原文连接:https://dfdazac.github.io/sinkhorn.html

相关文章
相关标签/搜索