PyTorch 實戰：計算 Wasserstein 距離

PyTorch 實戰：計算 Wasserstein 距離

2019-09-23 18:42:56 

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

 

最優傳輸理論及 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（如上所示）、二、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