Wasserstein Distance

Wasserstein Distance は、二つの probability distribution の距離を測る指標です。直感的には、ある distribution の形をした土の山を、別の distribution の形へ変形するために必要な最小のコストとして理解できます。

このため、Wasserstein Distance は Earth Mover's Distance（EM Distance）とも呼ばれます。

Earth Mover's Distance の直感

Earth Mover's Distance では、コストは「動かした土の量」と「動かした距離」の積として考えます。

$$\text{cost} = \text{moved amount} \times \text{moving distance}$$

たとえば、discrete な domain に四つの位置があり、二つの distribution $P$ と $Q$ が同じ総量の土を持っているとします。

$$\begin{aligned} & P_1 = 3,\ P_2 = 2,\ P_3 = 1,\ P_4 = 4 \\ & Q_1 = 1,\ Q_2 = 2,\ Q_3 = 4,\ Q_4 = 3 \end{aligned}$$

$P$ を $Q$ と同じ形に変えるには、余っている場所から足りない場所へ土を移動します。

画像出典: Lilian Weng, “From GAN to WGAN”。$P$ の土の山を $Q$ に一致させるために、どのように土を移動するかが示されています。

$P_i$ と $Q_i$ を一致させるために必要な差分を $\delta_i$ と書き、$\delta_{i+1} = \delta_i + P_i - Q_i$ とすると、この例では次のようになります。

$$\begin{aligned} \delta_0 &= 0 \\ \delta_1 &= 0 + 3 - 1 = 2 \\ \delta_2 &= 2 + 2 - 2 = 2 \\ \delta_3 &= 2 + 1 - 4 = -1 \\ \delta_4 &= -1 + 4 - 3 = 0 \end{aligned}$$

したがって、Earth Mover's Distance は次のように計算できます。

$$W = \sum_i |\delta_i| = 5$$

Continuous distribution での定義

Continuous な probability distribution の場合、Wasserstein Distance は次のように定義されます。

$$W(p_r, p_g) = \inf_{\gamma \sim \Pi(p_r, p_g)} \mathbb{E}_{(x,y) \sim \gamma}\left[\|x - y\|\right]$$

ここで、$\Pi(p_r, p_g)$ は、$p_r$ と $p_g$ の間で考えられるすべての joint probability distribution の集合です。一つの $\gamma \in \Pi(p_r, p_g)$ は、点 $x$ から点 $y$ へどれだけの土を運ぶかを表す transport plan として解釈できます。

$\inf$ は infimum、つまり下限を意味します。したがって、この定義は、すべての transport plan の中で期待コストが最小になるものを選ぶことを意味します。

KL Divergence や JS Divergence との違い

Wasserstein Distance の大きな利点は、二つの distribution の support が重なっていない場合でも、距離を意味のある滑らかな値として表せることです。

単純な例として、二つの distribution $P$ と $Q$ を考えます。

$$\forall (x,y) \in P,\ x = 0 \text{ and } y \sim U(0,1)$$ $$\forall (x,y) \in Q,\ x = \theta,\ 0 \leq \theta \leq 1 \text{ and } y \sim U(0,1)$$

画像出典: Lilian Weng, “From GAN to WGAN”。$\theta \neq 0$ のとき、$P$ と $Q$ は overlap しません。

$\theta \neq 0$ のとき、$P$ と $Q$ は disjoint になります。この場合、各距離は次のようになります。

$$\begin{aligned} D_{KL}(P \,\|\, Q) &= +\infty \\ D_{KL}(Q \,\|\, P) &= +\infty \\ D_{JS}(P, Q) &= \log 2 \\ W(P, Q) &= |\theta| \end{aligned}$$

一方で、$\theta = 0$ のときには、二つの distribution は完全に overlap します。

$$D_{KL}(P \,\|\, Q) = D_{KL}(Q \,\|\, P) = D_{JS}(P, Q) = 0$$ $$W(P, Q) = 0 = |\theta|$$

KL Divergence は、二つの distribution が disjoint であるときに無限大になります。Jensen-Shannon Divergence は、$\theta = 0$ の点で急に値が変わり、滑らかではありません。一方で、Wasserstein Distance は $|\theta|$ のように滑らかに変化します。

この滑らかさは、gradient descent による training にとって非常に重要です。

GAN で重要になる理由

GAN では、real data distribution $p_r$ と generated distribution $p_g$ の support が、高次元空間の中で重なりにくいことがあります。この状況では、JS Divergence に基づく training signal が不安定になりやすくなります。

Wasserstein Distance は、support が disjoint であっても distance を滑らかに与えるため、Wasserstein GAN の loss として使われます。

関連ページ

• KL Divergence
• Jensen-Shannon Divergence
• Wasserstein GAN