想要算一算Wasserstein距离?这里有一份PyTorch实战

最优传输理论及 Wasserstein 距离是很多读者都希望了解的基础,本文主要通过简单案例展示了它们的基本思想,并通过 PyTorch 介绍如何实战 W 距离。

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

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

因为 KL 散度可以度量两个分布的距离,所以只需要最小化 KL(q‖p) 就可以了。可以证明,最小化 KL(q‖p) 等价于最小化一个负对数似然,这样的做法在我们训练一个分类器时很常见。例如,对于变分自编码器来说,我们希望后验分布能够接近于某种先验分布,这也是我们通过最小化它们之间的 KL 散度来实现的。

尽管 KL 散度有很广泛的应用,在某些情况下,KL 散度则会失效。不妨考虑一下如下图所示的离散分布:

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(如上所示)、2、3 和 4 个单元分隔开。这样,它们之间的 Wasserstein 距离将分别为 1、4、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

工程GANKL散度PyTorchWasserstein 距离
8
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

范数技术

范数(norm),是具有“长度”概念的函数。在线性代数、泛函分析及相关的数学领域,是一个函数,其为向量空间内的所有向量赋予非零的正长度或大小。半范数反而可以为非零的向量赋予零长度。

机器学习技术

机器学习是人工智能的一个分支,是一门多领域交叉学科,涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论,机器学习与推断统计学联系尤为密切,也被称为统计学习理论。算法设计方面,机器学习理论关注可以实现的,行之有效的学习算法。

变分自编码器技术

变分自编码器可用于对先验数据分布进行建模。从名字上就可以看出,它包括两部分:编码器和解码器。编码器将数据分布的高级特征映射到数据的低级表征,低级表征叫作本征向量(latent vector)。解码器吸收数据的低级表征,然后输出同样数据的高级表征。变分编码器是自动编码器的升级版本,其结构跟自动编码器是类似的,也由编码器和解码器构成。在自动编码器中,需要输入一张图片,然后将一张图片编码之后得到一个隐含向量,这比原始方法的随机取一个随机噪声更好,因为这包含着原图片的信息,然后隐含向量解码得到与原图片对应的照片。但是这样其实并不能任意生成图片,因为没有办法自己去构造隐藏向量,所以它需要通过一张图片输入编码才知道得到的隐含向量是什么,这时就可以通过变分自动编码器来解决这个问题。解决办法就是在编码过程给它增加一些限制,迫使其生成的隐含向量能够粗略的遵循一个标准正态分布,这就是其与一般的自动编码器最大的不同。这样生成一张新图片就比较容易,只需要给它一个标准正态分布的随机隐含向量,这样通过解码器就能够生成想要的图片,而不需要给它一张原始图片先编码。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

概率分布技术

概率分布(probability distribution)或简称分布,是概率论的一个概念。广义地,它指称随机变量的概率性质--当我们说概率空间中的两个随机变量具有同样的分布(或同分布)时,我们是无法用概率来区别它们的。

收敛技术

在数学,计算机科学和逻辑学中,收敛指的是不同的变换序列在有限的时间内达到一个结论(变换终止),并且得出的结论是独立于达到它的路径(他们是融合的)。 通俗来说,收敛通常是指在训练期间达到的一种状态,即经过一定次数的迭代之后,训练损失和验证损失在每次迭代中的变化都非常小或根本没有变化。也就是说,如果采用当前数据进行额外的训练将无法改进模型,模型即达到收敛状态。在深度学习中,损失值有时会在最终下降之前的多次迭代中保持不变或几乎保持不变,暂时形成收敛的假象。

凸优化技术

凸优化,或叫做凸最优化,凸最小化,是数学最优化的一个子领域,研究定义于凸集中的凸函数最小化的问题。凸优化在某种意义上说较一般情形的数学最优化问题要简单,譬如在凸优化中局部最优值必定是全局最优值。凸函数的凸性使得凸分析中的有力工具在最优化问题中得以应用,如次导数等。 凸优化应用于很多学科领域,诸如自动控制系统,信号处理,通讯和网络,电子电路设计,数据分析和建模,统计学(最优化设计),以及金融。在近来运算能力提高和最优化理论发展的背景下,一般的凸优化已经接近简单的线性规划一样直捷易行。许多最优化问题都可以转化成凸优化(凸最小化)问题,例如求凹函数f最大值的问题就等同于求凸函数 -f最小值的问题。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

神经网络技术

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

欧几里得距离技术

在数学中,欧几里得距离或欧几里得度量是欧几里得空间中两点间“普通”(即直线)距离。 使用这个距离,欧氏空间成为度量空间。

正则化技术

当模型的复杂度增大时,训练误差会逐渐减小并趋向于0;而测试误差会先减小,达到最小值后又增大。当选择的模型复杂度过大时,过拟合现象就会发生。这样,在学习时就要防止过拟合。进行最优模型的选择,即选择复杂度适当的模型,以达到使测试误差最小的学习目的。

信息论技术

信息论是在信息可以量度的基础上,研究有效地和可靠地传递信息的科学,它涉及信息量度、信息特性、信息传输速率、信道容量、干扰对信息传输的影响等方面的知识。通常把上述范围的信息论称为狭义的信息论,又因为它的创始人是香农,故又称为香农信息论。

暂无评论
暂无评论~