最優(yōu)傳輸理論及 Wasserstein 距離是很多讀者都希望了解的基礎(chǔ)，本文主要通過簡單案例展示了它們的基本思想，并通過 PyTorch 介紹如何實(shí)戰(zhàn) W 距離。

機(jī)器學(xué)習(xí)中的許多問題都涉及到令兩個(gè)分布盡可能接近的思想，例如在 GAN 中令生成器分布接近判別器分布就能偽造出逼真的圖像。但是 KL 散度等分布的度量方法有很多局限性，本文則介紹了 Wasserstein 距離及 Sinkhorn 迭代方法，它們 GAN 及眾多任務(wù)上都展示了杰出的性能。

在簡單的情況下，我們假設(shè)從未知數(shù)據(jù)分布 p(x) 中觀測到一些隨機(jī)變量 x（例如，貓的圖片），我們想要找到一個(gè)模型 q(x|θ)（例如一個(gè)神經(jīng)網(wǎng)絡(luò)）能作為 p(x) 的一個(gè)很好的近似。如果 p 和 q 的分布很相近，那么就表明我們的模型已經(jīng)學(xué)習(xí)到如何識(shí)別貓。

因?yàn)?KL 散度可以度量兩個(gè)分布的距離，所以只需要最小化 KL(q‖p) 就可以了?？梢宰C明，最小化 KL(q‖p) 等價(jià)于最小化一個(gè)負(fù)對數(shù)似然，這樣的做法在我們訓(xùn)練一個(gè)分類器時(shí)很常見。例如，對于變分自編碼器來說，我們希望后驗(yàn)分布能夠接近于某種先驗(yàn)分布，這也是我們通過最小化它們之間的 KL 散度來實(shí)現(xiàn)的。

盡管 KL 散度有很廣泛的應(yīng)用，在某些情況下，KL 散度則會(huì)失效。不妨考慮一下如下圖所示的離散分布：

KL 散度假設(shè)這兩個(gè)分布共享相同的支撐集（也就是說，它們被定義在同一個(gè)點(diǎn)集上）。因此，我們不能為上面的例子計(jì)算 KL 散度。由于這一個(gè)限制和其他計(jì)算方面的因素促使研究人員尋找一種更適合于計(jì)算兩個(gè)分布之間差異的方法。

在本文中，作者將：

• 簡單介紹最優(yōu)傳輸問題

• 將 Sinkhorn 迭代描述為對解求近似

• 使用 PyTorch 計(jì)算 Sinkhorn 距離

• 描述用于計(jì)算 mini-batch 之間的距離的對該實(shí)現(xiàn)的擴(kuò)展

移動(dòng)概率質(zhì)量函數(shù)

我們不妨把離散的概率分布想象成空間中分散的點(diǎn)的質(zhì)量。我們可以觀測這些帶質(zhì)量的點(diǎn)從一個(gè)分布移動(dòng)到另一個(gè)分布需要做多少功，如下圖所示：

接著，我們可以定義另一個(gè)度量標(biāo)準(zhǔn)，用以衡量移動(dòng)做所有點(diǎn)所需要做的功。要想將這個(gè)直觀的概念形式化定義下來，首先，我們可以通過引入一個(gè)耦合矩陣 P（coupling matrix），它表示要從 p(x) 支撐集中的一個(gè)點(diǎn)上到 q(x) 支撐集中的一個(gè)點(diǎn)需要分配多少概率質(zhì)量。對于均勻分布，我們規(guī)定每個(gè)點(diǎn)都具有 1/4 的概率質(zhì)量。如果我們將本例支撐集中的點(diǎn)從左到右排列，我們可以將上述的耦合矩陣寫作：

也就是說，p(x) 支撐集中點(diǎn) 1 的質(zhì)量被分配給了 q(x) 支撐集中的點(diǎn) 4，p(x) 支撐集中點(diǎn) 2 的質(zhì)量被分配給了 q(x) 支撐集中的點(diǎn) 3，以此類推，如上圖中的箭頭所示。

為了算出質(zhì)量分配的過程需要做多少功，我們將引入第二個(gè)矩陣：距離矩陣。該矩陣中的每個(gè)元素 C_ij 表示將 p(x) 支撐集中的點(diǎn)移動(dòng)到 q(x) 支撐集中的點(diǎn)上的成本。點(diǎn)與點(diǎn)之間的歐幾里得距離是定義這種成本的一種方式，它也被稱為「ground distance」。如果我們假設(shè) p(x) 的支撐集和 q(x) 的支撐集分別為 {1,2,3,4} 和 {5,6,7,8}，成本矩陣即為：

根據(jù)上述定義，總的成本可以通過 P 和 C 之間的 Frobenius 內(nèi)積來計(jì)算：

你可能已經(jīng)注意到了，實(shí)際上有很多種方法可以把點(diǎn)從一個(gè)支撐集移動(dòng)到另一個(gè)支撐集中，每一種方式都會(huì)得到不同的成本。上面給出的只是一個(gè)示例，但是我們感興趣的是最終能夠讓成本較小的分配方式。這就是兩個(gè)離散分布之間的「最優(yōu)傳輸」問題，該問題的解是所有耦合矩陣上的最低成本 L_C。

由于不是所有矩陣都是有效的耦合矩陣，最后一個(gè)條件會(huì)引入了一個(gè)約束。對于一個(gè)耦合矩陣來說，其所有列都必須要加到帶有 q(x) 概率質(zhì)量的向量中。在本例中，該向量包含 4 個(gè)值為 1/4 的元素。更一般地，我們可以將兩個(gè)向量分別記為 a 和 b，因此最有運(yùn)輸問題可以被寫作：

當(dāng)距離矩陣基于一個(gè)有效的距離函數(shù)構(gòu)建時(shí)，最小成本即為我們所說的「Wasserstein 距離」。

關(guān)于該問題的解以及將其擴(kuò)展到連續(xù)概率分布中還有大量問題需要解決。如果想要獲取更正式、更容易理解的解釋，讀者可以參閱 Gabriel Peyré 和 Marco Cuturi 編寫的「Computational Optimal Transport」一書，此書也是本文寫作的主要參考來源之一。

這里的基本設(shè)定是，我們已經(jīng)把求兩個(gè)分布之間距離的問題定義為求最優(yōu)耦合矩陣的問題。事實(shí)證明，我們可以通過一個(gè)小的修改讓我們以迭代和可微分的方式解決這個(gè)問題，這將讓我們可以很好地使用深度學(xué)習(xí)自動(dòng)微分機(jī)制完成該工作。

熵正則化和 Sinkhorn 迭代

首先，我們將一個(gè)矩陣的熵定義如下：

正如信息論中概率分布的熵一樣，一個(gè)熵較低的矩陣將會(huì)更稀疏，它的大部分非零值集中在幾個(gè)點(diǎn)周圍。相反，一個(gè)具有高熵的矩陣將會(huì)更平滑，其最大熵是在均勻分布的情況下獲得的。我們可以將正則化系數(shù) ε 引入最優(yōu)傳輸問題，從而得到更平滑的耦合矩陣：

通過增大 ε，最終得到的耦合矩陣將會(huì)變得更加平滑；而當(dāng) ε 趨近于零時(shí)，耦合矩陣會(huì)更加稀疏，同時(shí)最終的解會(huì)更加趨近于原始最優(yōu)運(yùn)輸問題。

通過引入這種熵正則化，該問題變成了一個(gè)凸優(yōu)化問題，并且可 以通過使用「Sinkhorn iteration」求解。解可以被寫作 P=diag(u)Kdiag(v)，在迭代過程中交替更新 u 和 v：

其中 K 是一個(gè)用 C 計(jì)算的核矩陣（kernel matrix）。由于這些迭代過程是在對原始問題的正則化版本求解，因此對應(yīng)產(chǎn)生的 Wasserstein 距離有時(shí)被稱為 Sinkhorn 距離。該迭代過程會(huì)形成一個(gè)線性操作的序列，因此對于深度學(xué)習(xí)模型，通過這些迭代進(jìn)行反向傳播是非常簡單的。

通過 PyTorch 實(shí)現(xiàn) Sinkhorn 迭代

為了提升 Sinkhorn 迭代的收斂性和穩(wěn)定性，還可以加入其它的步驟。我們可以在 GitHub 上找到 Gabriel Peyre 完成的詳細(xì)實(shí)現(xiàn)。

項(xiàng)目鏈接：https://github.com/gpeyre/SinkhornAutoDiff。

讓我們先用一個(gè)簡單的例子來測試一下，現(xiàn)在我們將研究二維空間（而不是上面的一維空間）中的離散均勻分布。在這種情況下，我們將在平面上移動(dòng)概率質(zhì)量。讓我們首先定義兩個(gè)簡單的分布：

我們很容易看出，最優(yōu)傳輸對應(yīng)于將 p(x) 支撐集中的每個(gè)點(diǎn)分配到 q(x) 支撐集上的點(diǎn)。對于所有的點(diǎn)來說，距離都是 1，同時(shí)由于分布是均勻的，每點(diǎn)移動(dòng)的概率質(zhì)量是 1/5。因此，Wasserstein 距離是 5×1/5= 1。現(xiàn)在我們用 Sinkhorn 迭代來計(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 是距離矩陣。距離矩陣如下圖所示：

元素「C[0, 0]」說明了將（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ì)怎樣？

正如我們前面討論過的，加大 ε 有增大耦合矩陣熵的作用。接下來，我們看看 P 是如何變得更加平滑的。但是，這樣做也會(huì)為計(jì)算出的距離帶來一個(gè)不好的影響，導(dǎo)致對 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)

讓我們在一個(gè)更有趣的分布（Moons 數(shù)據(jù)集）上完成這項(xiàng)工作。

Mini-batch 上的 Sinkhorn 距離

在深度學(xué)習(xí)中，我們通常對使用 mini-batch 來加速計(jì)算十分感興趣。我們也可以通過使用額外的批處理維度修改 Sinkhorn 迭代來滿足該設(shè)定。將此更改添加到具體實(shí)現(xiàn)中后，我們可以在一個(gè) mini-batch 中計(jì)算多個(gè)分布的 Sinkhorn 距離。下面我們將通過另一個(gè)容易被驗(yàn)證的例子說明這一點(diǎn)。

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

我們將計(jì)算包含 5 個(gè)支撐點(diǎn)的 4 對均勻分布的 Sinkhorn 距離，它們垂直地被 1（如上所示）、2、3 和 4 個(gè)單元分隔開。這樣，它們之間的 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í)，也請注意，現(xiàn)在 P 和 C 為 3 維張量，它包含 mini-batch 中每對分布的耦合矩陣和距離矩陣：

結(jié)語

分布之間的 Wasserstein 距離及其通過 Sinkhorn 迭代實(shí)現(xiàn)的計(jì)算方法為我們帶來了許多可能性。該框架不僅提供了對 KL 散度等距離的替代方法，而且在建模過程中提供了更大的靈活性，我們不再被迫要選擇特定的參數(shù)分布。這些迭代過程可以在 GPU 上高效地執(zhí)行，并且是完全可微分的，這使得它對于深度學(xué)習(xí)來說是一個(gè)很好的選擇。這些優(yōu)點(diǎn)在機(jī)器學(xué)習(xí)領(lǐng)域的最新研究中得到了充分的利用（如自編碼器和距離嵌入），使其在該領(lǐng)域的應(yīng)用前景更加廣闊。