作者：Daniel Daza
機(jī)器之心編譯
最優(yōu)傳輸理論及 Wasserstein 距離是很多讀者都希望了解的基礎(chǔ)，本文主要通過(guò)簡(jiǎn)單案例展示了它們的基本思想，并通過(guò) PyTorch 介紹如何實(shí)戰(zhàn) W 距離。
機(jī)器學(xué)習(xí)中的許多問(wèn)題都涉及到令兩個(gè)分布盡可能接近的思想，例如在 GAN 中令生成器分布接近判別器分布就能偽造出逼真的圖像。但是 KL 散度等分布的度量方法有很多局限性，本文則介紹了 Wasserstein 距離及 Sinkhorn 迭代方法，它們 GAN 及眾多任務(wù)上都展示了杰出的性能。
在簡(jiǎn)單的情況下，我們假設(shè)從未知數(shù)據(jù)分布 p(x) 中觀測(cè)到一些隨機(jī)變量 x（例如，貓的圖片），我們想要找到一個(gè)模型 q(x|θ)（例如一個(gè)神經(jīng)網(wǎng)絡(luò)）能作為 p(x) 的一個(gè)很好的近似。如果 p 和 q 的分布很相近，那么就表明我們的模型已經(jīng)學(xué)習(xí)到如何識(shí)別貓。
因?yàn)?KL 散度可以度量?jī)蓚€(gè)分布的距離，所以只需要最小化 KL(q‖p) 就可以了?？梢宰C明，最小化 KL(q‖p) 等價(jià)于最小化一個(gè)負(fù)對(duì)數(shù)似然，這樣的做法在我們訓(xùn)練一個(gè)分類(lèi)器時(shí)很常見(jiàn)。例如，對(duì)于變分自編碼器來(lái)說(shuō)，我們希望后驗(yàn)分布能夠接近于某種先驗(yàn)分布，這也是我們通過(guò)最小化它們之間的 KL 散度來(lái)實(shí)現(xiàn)的。
盡管 KL 散度有很廣泛的應(yīng)用，在某些情況下，KL 散度則會(huì)失效。不妨考慮一下如下圖所示的離散分布：
KL 散度假設(shè)這兩個(gè)分布共享相同的支撐集（也就是說(shuō)，它們被定義在同一個(gè)點(diǎn)集上）。因此，我們不能為上面的例子計(jì)算 KL 散度。由于這一個(gè)限制和其他計(jì)算方面的因素促使研究人員尋找一種更適合于計(jì)算兩個(gè)分布之間差異的方法。
在本文中，作者將：
簡(jiǎn)單介紹最優(yōu)傳輸問(wèn)題
將 Sinkhorn 迭代描述為對(duì)解求近似
使用 PyTorch 計(jì)算 Sinkhorn 距離
描述用于計(jì)算 mini-batch 之間的距離的對(duì)該實(shí)現(xiàn)的擴(kuò)展
移動(dòng)概率質(zhì)量函數(shù)
我們不妨把離散的概率分布想象成空間中分散的點(diǎn)的質(zhì)量。我們可以觀測(cè)這些帶質(zhì)量的點(diǎn)從一個(gè)分布移動(dòng)到另一個(gè)分布需要做多少功，如下圖所示：
接著，我們可以定義另一個(gè)度量標(biāo)準(zhǔn)，用以衡量移動(dòng)做所有點(diǎn)所需要做的功。要想將這個(gè)直觀的概念形式化定義下來(lái)，首先，我們可以通過(guò)引入一個(gè)耦合矩陣 P（coupling matrix），它表示要從 p(x) 支撐集中的一個(gè)點(diǎn)上到 q(x) 支撐集中的一個(gè)點(diǎn)需要分配多少概率質(zhì)量。對(duì)于均勻分布，我們規(guī)定每個(gè)點(diǎn)都具有 1/4 的概率質(zhì)量。如果我們將本例支撐集中的點(diǎn)從左到右排列，我們可以將上述的耦合矩陣寫(xiě)作：
也就是說(shuō)，p(x) 支撐集中點(diǎn) 1 的質(zhì)量被分配給了 q(x) 支撐集中的點(diǎn) 4，p(x) 支撐集中點(diǎn) 2 的質(zhì)量被分配給了 q(x) 支撐集中的點(diǎn) 3，以此類(lèi)推，如上圖中的箭頭所示。
為了算出質(zhì)量分配的過(guò)程需要做多少功，我們將引入第二個(gè)矩陣：距離矩陣。該矩陣中的每個(gè)元素 C_ij 表示將 p(x) 支撐集中的點(diǎn)移動(dòng)到 q(x) 支撐集中的點(diǎn)上的成本。點(diǎn)與點(diǎn)之間的歐幾里得距離是定義這種成本的一種方式，它也被稱(chēng)為「ground distance」。如果我們假設(shè) p(x) 的支撐集和 q(x) 的支撐集分別為 {1,2,3,4} 和 {5,6,7,8}，成本矩陣即為：
根據(jù)上述定義，總的成本可以通過(guò) P 和 C 之間的 Frobenius 內(nèi)積來(lái)計(jì)算：
你可能已經(jīng)注意到了，實(shí)際上有很多種方法可以把點(diǎn)從一個(gè)支撐集移動(dòng)到另一個(gè)支撐集中，每一種方式都會(huì)得到不同的成本。上面給出的只是一個(gè)示例，但是我們感興趣的是最終能夠讓成本較小的分配方式。這就是兩個(gè)離散分布之間的「最優(yōu)傳輸」問(wèn)題，該問(wèn)題的解是所有耦合矩陣上的最低成本 L_C。
由于不是所有矩陣都是有效的耦合矩陣，最后一個(gè)條件會(huì)引入了一個(gè)約束。對(duì)于一個(gè)耦合矩陣來(lái)說(shuō)，其所有列都必須要加到帶有 q(x) 概率質(zhì)量的向量中。在本例中，該向量包含 4 個(gè)值為 1/4 的元素。更一般地，我們可以將兩個(gè)向量分別記為 a 和 b，因此最有運(yùn)輸問(wèn)題可以被寫(xiě)作：
當(dāng)距離矩陣基于一個(gè)有效的距離函數(shù)構(gòu)建時(shí)，最小成本即為我們所說(shuō)的「Wasserstein 距離」。
關(guān)于該問(wèn)題的解以及將其擴(kuò)展到連續(xù)概率分布中還有大量問(wèn)題需要解決。如果想要獲取更正式、更容易理解的解釋?zhuān)x者可以參閱 Gabriel Peyré 和 Marco Cuturi 編寫(xiě)的「Computational Optimal Transport」一書(shū)，此書(shū)也是本文寫(xiě)作的主要參考來(lái)源之一。
這里的基本設(shè)定是，我們已經(jīng)把求兩個(gè)分布之間距離的問(wèn)題定義為求最優(yōu)耦合矩陣的問(wèn)題。事實(shí)證明，我們可以通過(guò)一個(gè)小的修改讓我們以迭代和可微分的方式解決這個(gè)問(wèn)題，這將讓我們可以很好地使用深度學(xué)習(xí)自動(dòng)微分機(jī)制完成該工作。
熵正則化和 Sinkhorn 迭代
首先，我們將一個(gè)矩陣的熵定義如下：
正如信息論中概率分布的熵一樣，一個(gè)熵較低的矩陣將會(huì)更稀疏，它的大部分非零值集中在幾個(gè)點(diǎn)周?chē)Ｏ喾?，一個(gè)具有高熵的矩陣將會(huì)更平滑，其最大熵是在均勻分布的情況下獲得的。我們可以將正則化系數(shù) ε 引入最優(yōu)傳輸問(wèn)題，從而得到更平滑的耦合矩陣：
通過(guò)增大 ε，最終得到的耦合矩陣將會(huì)變得更加平滑；而當(dāng) ε 趨近于零時(shí)，耦合矩陣會(huì)更加稀疏，同時(shí)最終的解會(huì)更加趨近于原始最優(yōu)運(yùn)輸問(wèn)題。
通過(guò)引入這種熵正則化，該問(wèn)題變成了一個(gè)凸優(yōu)化問(wèn)題，并且可 以通過(guò)使用「Sinkhorn iteration」求解。解可以被寫(xiě)作 P=diag(u)Kdiag(v)，在迭代過(guò)程中交替更新 u 和 v：
其中 K 是一個(gè)用 C 計(jì)算的核矩陣（kernel matrix）。由于這些迭代過(guò)程是在對(duì)原始問(wèn)題的正則化版本求解，因此對(duì)應(yīng)產(chǎn)生的 Wasserstein 距離有時(shí)被稱(chēng)為 Sinkhorn 距離。該迭代過(guò)程會(huì)形成一個(gè)線性操作的序列，因此對(duì)于深度學(xué)習(xí)模型，通過(guò)這些迭代進(jìn)行反向傳播是非常簡(jiǎn)單的。
通過(guò) PyTorch 實(shí)現(xiàn) Sinkhorn 迭代
為了提升 Sinkhorn 迭代的收斂性和穩(wěn)定性，還可以加入其它的步驟。我們可以在 GitHub 上找到 Gabriel Peyre 完成的詳細(xì)實(shí)現(xiàn)。
項(xiàng)目鏈接：https://github.com/gpeyre/SinkhornAutoDiff。
讓我們先用一個(gè)簡(jiǎn)單的例子來(lái)測(cè)試一下，現(xiàn)在我們將研究二維空間（而不是上面的一維空間）中的離散均勻分布。在這種情況下，我們將在平面上移動(dòng)概率質(zhì)量。讓我們首先定義兩個(gè)簡(jiǎn)單的分布：
%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();
我們很容易看出，最優(yōu)傳輸對(duì)應(yīng)于將 p(x) 支撐集中的每個(gè)點(diǎn)分配到 q(x) 支撐集上的點(diǎn)。對(duì)于所有的點(diǎn)來(lái)說(shuō)，距離都是 1，同時(shí)由于分布是均勻的，每點(diǎn)移動(dòng)的概率質(zhì)量是 1/5。因此，Wasserstein 距離是 5×1/5= 1?，F(xiàn)在我們用 Sinkhorn 迭代來(lái)計(jì)算這個(gè)距離：
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
結(jié)果正如我們所計(jì)算的那樣，距離為 1?，F(xiàn)在，讓我們查看一下「Sinkhorn( )」方法返回的矩陣，其中 P 是計(jì)算出的耦合矩陣，C 是距離矩陣。距離矩陣如下圖所示：
plt.imshow(C)
plt.title('Distance matrix')
plt.colorbar();
plt.imshow(C)plt.title('Distance matrix')plt.colorbar();
元素「C[0, 0]」說(shuō)明了將（0,0）點(diǎn)的質(zhì)量移動(dòng)到（0,1）所需要的成本 1 是如何產(chǎn)生的。在該行的另一端，元素「C[0, 4]」包含了將點(diǎn)（0,0）的質(zhì)量移動(dòng)到點(diǎn)（4,1）所需要的成本，這個(gè)成本是整個(gè)矩陣中最大的：
由于我們?yōu)榫嚯x矩陣使用的是平方后的 ?2 范數(shù)，計(jì)算結(jié)果如上所示?，F(xiàn)在，讓我們看看計(jì)算出的耦合矩陣吧：
plt.imshow(P)
plt.title('Coupling matrix');
plt.imshow(P)plt.title('Coupling matrix');
該圖很好地向我們展示了算法是如何有效地發(fā)現(xiàn)最優(yōu)耦合，它與我們前面確定的耦合矩陣是相同的。到目前為止，我們使用了 0.1 的正則化系數(shù)。如果將該值增加到 1 會(huì)怎樣？
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
正如我們前面討論過(guò)的，加大 ε 有增大耦合矩陣熵的作用。接下來(lái)，我們看看 P 是如何變得更加平滑的。但是，這樣做也會(huì)為計(jì)算出的距離帶來(lái)一個(gè)不好的影響，導(dǎo)致對(duì) 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)
讓我們?cè)谝粋€(gè)更有趣的分布（Moons 數(shù)據(jù)集）上完成這項(xiàng)工作。
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 距離
在深度學(xué)習(xí)中，我們通常對(duì)使用 mini-batch 來(lái)加速計(jì)算十分感興趣。我們也可以通過(guò)使用額外的批處理維度修改 Sinkhorn 迭代來(lái)滿足該設(shè)定。將此更改添加到具體實(shí)現(xiàn)中后，我們可以在一個(gè) mini-batch 中計(jì)算多個(gè)分布的 Sinkhorn 距離。下面我們將通過(guò)另一個(gè)容易被驗(yàn)證的例子說(shuō)明這一點(diǎn)。
代碼：https://github.com/dfdazac/wassdistance/blob/master/layers.py
我們將計(jì)算包含 5 個(gè)支撐點(diǎn)的 4 對(duì)均勻分布的 Sinkhorn 距離，它們垂直地被 1（如上所示）、2、3 和 4 個(gè)單元分隔開(kāi)。這樣，它們之間的 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])
這樣做確實(shí)有效！同時(shí)，也請(qǐng)注意，現(xiàn)在 P 和 C 為 3 維張量，它包含 mini-batch 中每對(duì)分布的耦合矩陣和距離矩陣：
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])
結(jié)語(yǔ)
分布之間的 Wasserstein 距離及其通過(guò) Sinkhorn 迭代實(shí)現(xiàn)的計(jì)算方法為我們帶來(lái)了許多可能性。該框架不僅提供了對(duì) KL 散度等距離的替代方法，而且在建模過(guò)程中提供了更大的靈活性，我們不再被迫要選擇特定的參數(shù)分布。這些迭代過(guò)程可以在 GPU 上高效地執(zhí)行，并且是完全可微分的，這使得它對(duì)于深度學(xué)習(xí)來(lái)說(shuō)是一個(gè)很好的選擇。這些優(yōu)點(diǎn)在機(jī)器學(xué)習(xí)領(lǐng)域的最新研究中得到了充分的利用（如自編碼器和距離嵌入），使其在該領(lǐng)域的應(yīng)用前景更加廣闊。
聯(lián)系客服
