この文章について

Wasserstein GANの全体の解説や実装・実験については他に記事があるので割愛させていただきます.
損失関数およびニューラルネットワークで近似する部分についてのみ記述しています.
厳密な証明をつけるわけではありません. 測度論が必要になるからです.
線形計画法を知っていることが前提です. 記法が分かっていれば大丈夫です.
先行している類似記事に, Wasserstein GAN と Kantorovich-Rubinstein 双対性（和訳）があります. 本記事はこの記事やこの引用元などに記載されている内容を自分なりに整理したものになります.

本記事では, 論文中で用いられているKantorovich-Rubinstein双対性
$$
\inf_{\gamma \in \Pi(p,q)}
\int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y
= \sup_{f : {\rm 1-Lip}} \int_{\mathbb{R}^n} f(x) (p(x)-q(x)) {\rm d}x
$$
の意味を, 線形最適化を通して理解することを目指します.

Wasserstein距離の提案

論文中では, Wasserstein距離を導入しています：
$$
W(p,q) := \inf_{\gamma \in \Pi(p,q)}
\int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y
$$
ここで
$$
\Pi(p,q) := \Bigl\{
\gamma : \mathbb{R}^n \times \mathbb{R}^n \to \mathbb{R} \Bigm|
\int_{\mathbb{R}^n} \gamma(x,y) {\rm d}x = q(y),
\int_{\mathbb{R}^n} \gamma(x,y) {\rm d}y = p(x),
\gamma \ge 0
\Bigr\}
$$
は, 周辺化すると$p,q$になるような$\mathbb{R}^{n} \times \mathbb{R}^{n}$上の確率密度関数の全体を指します.

イメージとしては,

場所$x$に土が密度$p(x)$で置いてある. これを密度$q(x)$にしたい.
$\gamma(x,y)$で$x$から$y$に運ぶ土の量を表すことにする. $\gamma(x,x)$は$x$に残す土の量を表すことになる.
- $y$に運んでくる土の総量は$$\int_{\mathbb{R}^n} \gamma(x,y) {\rm d}x = q(y)$$である.
- $x$から持ち出す土の総量は$$\int_{\mathbb{R}^n} \gamma(x,y) {\rm d}y = p(x)$$である.
運ぶ際には距離に応じたコストがかかる.
総コストの下限を以って$p$と$q$の距離と定める.

です.

$W(p,q) :=\inf_{\gamma \in \Pi(p,q)} \int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y$で2つの確率密度関数$p,q$の距離を定めている.

Wasserstein距離の計算

論文中では,
$$
W(p,q) = \sup_{f \in C_1} \int_{\mathbb{R}^n} (p(x)-q(x)) \ f(x) {\rm d}x
= \sup_{f \in C_1} \left( E_{x \sim p} [f(x)] - E_{x \sim q} [f(x)] \right)\
$$
が成り立つことを使っています. ただし
$$
C_1 := \Bigl\{ f:\mathbb{R} \to \mathbb{R} \Bigm| f(x) - f(y) \le \|x-y\| \ (\forall x,y \in \mathbb{R}^n) \Bigr\}
$$

とおいています. 損失関数を定義する際には,

$p_{\rm data}$はデータを生成している真の分布（に関する確率密度関数）を表す
Generatorを$G$で表す
$G$によって確率密度関数$p_{\rm known}$（正規分布など）が$p_G$に変換される

とき,

$$
\underset{G}{\text{minimize }} \sup_{f \in C_1} \left( E_{x \sim p_{\rm data}} [f(x)] - E_{z \sim p_{\rm known}} [f(G(z))] \right)\
$$

の解$G$がデータを最もよく生成する, と考えています.

このとき必要となるのは,

$E_{x \sim p_{\rm data}} [f(x)] - E_{z \sim p_{\rm known}} [f(G(z))]$を最大化するような$f \in C_1$
Generator $G$

の2つです. これらをニューラルネットワークで近似します.

冒頭にも書きましたが, 本記事では, 等式
$$
\inf_{\gamma \in \Pi(p,q)}
\int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y
= \sup_{f \in C_1} \int_{\mathbb{R}^n} (p(x)-q(x))\ f(x) {\rm d}x
$$
が成り立つ理由を, 線形最適化を通して理解することを目指します.

定理$W(p,q)=\sup_{f \in C_1} \int_{\mathbb{R}^n} (p(x)-q(x)) \ f(x) {\rm d}x$を用いて計算・設計を行なっている.

離散化して考える

そのままだと積分の下限などを考えることになるので, 代わりに「離散的な世界」で考えることにします.
ここでは,

集合$A \subset \mathbb{R}^n$は有限集合$\{x_1,\ldots,x_N\} \subset A$に,
関数$\ \ f:A \to \mathbb{R}$は点群$y_1=f(x_1),\ldots,y_N=f(x_N)$に,
$A$上の積分$\int_A f(x){\rm d}x$は和$\sum_{i=1}^N y_i$に,
$A$上の確率密度関数$p:A \to \mathbb{R}$は確率関数$p_1,\ldots,p_N$に

読み替えることにします. ここでは以下の図に従って考えていくことにします.

Step 1. 左辺の離散化

左辺
$$
\inf_{\gamma \in \Pi(p,q)} \int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y
$$

に対応する離散化を考えましょう. まず$\inf$の中身については, 離散化したものは
$$
\sum_{i=1}^N \sum_{j=1}^N \|x_i-x_j\| \gamma(x_i,x_j) \ \ (x_1,\ldots,x_N \in \mathbb{R}^n)
$$
となります. ここで
$$
d_{ij} := \|x_i - x_j\| \\
\gamma_{ij} := \gamma(x_i, x_j)
$$
とおくことにすれば,
$$
\sum_{i=1}^N \sum_{j=1}^N \|x_i-x_j\| \gamma(x_i,x_j) = \sum_{i=1}^N \sum_{j=1}^N d_{ij} \gamma_{ij}
$$
です.
また, 制約である$\gamma \in \Pi(p,q)$は,
$$
\Pi(p,q) := \Bigl\{
\gamma : \mathbb{R}^n \times \mathbb{R}^n \to \mathbb{R} \Bigm|
\int_{\mathbb{R}^n} \gamma(x,y) {\rm d}x = q(y),
\int_{\mathbb{R}^n} \gamma(x,y) {\rm d}y = p(x),
\gamma \ge 0
\Bigr\}
$$
でしたから, $\gamma$を離散化した$\gamma_{ij}$たちは
$$
\sum_{i=1}^N \gamma_{ij} = q_j ,\sum_{i=1}^N \gamma_{ij} = p_i, \gamma_{ij} \ge 0
$$
を満たすことにします. ここで$p_i,q_j$は$p,q$をそれぞれ離散化したもので,

$p,q$が確率密度関数であることに対応して
$$
\sum_{i=1}^N p_i = \sum_{j=1}^N q_j = 1
$$
が成り立つものとします.

以上から,
$$
\inf_{\gamma \in \Pi(p,q)}\int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y
$$
は, 離散化した場合には, 線形最適化問題
$$
\begin{aligned}
\min_{\gamma_{ij}} &\sum_{i=1}^N \sum_{j=1}^N d_{ij} \gamma_{ij}\\
{\rm subject \ to} &\sum_{i=1}^N \gamma_{ij} = q_j \ \ (j=1,\ldots,N)\\
&\sum_{i=1}^N \gamma_{ij} = p_i \ \ (i=1,\ldots,N)\\
&\gamma_{ij} \ge 0 \ \ (i,j=1,\ldots,N)
\end{aligned}
$$
に対応することが分かります.
目的関数は積分そのものに, 制約は$\gamma \in \Pi(p,q)$であることに対応しています.

$\inf_{\gamma \in \Pi(p,q)} \int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y$に対応する線形最適化問題が得られた.

Step 2. 右辺の離散化

次に右辺
次に
$$
\sup_{f \in C_1} \int_{\mathbb{R}^n} (p(x)-q(x)) \ f(x) {\rm d}x
$$
に対応する離散化を考えましょう. 1.と同様に, $\sup$の中身に対応するものは
$$
\sum_{i=1}^N (p_i-q_i)f_i
$$
となります. また, 制約である$f \in C_1$は,
$$
C_1 := \{ f:\mathbb{R} \to \mathbb{R} \mid f(x) - f(y) \le \|x-y\| \ (\forall x,y \in \mathbb{R}^n) \} \\
d_{ij} = \|x_i - y_j\|
$$
に注意すれば
$$
f_i - f_j \le d_{ij} \ \ (i,j=1,\ldots,N)
$$
に対応します. 以上から,
$$
\sup_{f \in C_1} \int_{\mathbb{R}^n} f(x) (p(x)-q(x)) {\rm d}x
$$
は, 離散化した場合には, 線形最適化問題
$$
\begin{aligned}
\max_{f_i} &\sum_{i=1}^N (p_i - q_i) f_i \\
{\rm subject \ to} \ \ &f_i - f_j \le d_{ij} \ \ (i,j=1,\ldots,N) \\
\end{aligned}
$$
に対応することが分かります.

$\sup_{f \in C_1} \int_{\mathbb{R}^n} (p(x)-q(x)) \ f(x) {\rm d}x$に対応する線形最適化問題が得られた.

Step 3. 双対問題

Step 1,2で得られた2つの線形最適化問題の関係を調べましょう.

線形最適化問題には双対性と呼ばれるものがあり,
一般に以下の２つの最適値が一致することが知られています.
$\boldsymbol{a}_1,\ldots,\boldsymbol{a}_K,\boldsymbol{c} \in \mathbb{R}^m$と$\boldsymbol{b} \in \mathbb{R}^n$に対して,

$$
\begin{aligned}
\min_{\boldsymbol{x} \in \mathbb{R}^m} \ &\boldsymbol{c}^{\top} \boldsymbol{x} \\
{\rm subject \ to}\ &\boldsymbol{a}_k^{\top} \boldsymbol{x} = b_k \ \ (k=1,\ldots,n) \\
& \boldsymbol{x} \ge \boldsymbol{0} \\
\end{aligned}
$$

$$
\begin{aligned}
\max_{\boldsymbol{y} \in \mathbb{R}^n} \ &\boldsymbol{b}^{\top} \boldsymbol{y} \\
{\rm subject \ to}\ &\sum_{k=1}^n y_k \boldsymbol{a}_k \le \boldsymbol{c} \ \ \\
\end{aligned}
$$

一方の問題に対するもう一方の問題のことを, その問題の双対問題とよびます.
証明が読みたい方はここなどを参照すると良いでしょう（Google検索で一番上にあったpdfを載せています）.

この事実を用ると, Step 1.で導出した問題
$$
\begin{aligned}
\min_{\gamma_{ij}} &\sum_{i=1}^N \sum_{j=1}^N d_{ij} \gamma_{ij}\\
{\rm subject \ to} &\sum_{i=1}^N \gamma_{ij} = q_j \ \ (j=1,\ldots,N)\\
&\sum_{i=1}^N \gamma_{ij} = p_i \ \ (i=1,\ldots,N)\\
&\gamma_{ij} \ge 0 \ \ (i,j=1,\ldots,N)
\end{aligned}
$$
の双対問題は
$$
\begin{aligned}
\max_{f_i,h_j} &\sum_{i=1}^N p_i f_i + \sum_{j=1}^N q_j h_j \\
{\rm subject \ to} \ \ &f_i + h_j \le d_{ij} \ \ (i,j=1,\ldots,N)
\end{aligned}
$$
であることがわかります. Step 2.で得られた線形最適化問題

$$
\begin{aligned}
\max_{f_i} &\sum_{i=1}^N (p_i - q_i) f_i \\
{\rm subject \ to} \ \ &f_i - f_j \le d_{ij} \ \ (i,j=1,\ldots,N) \\
\end{aligned}
$$

に非常に似た問題が出てきました.

$\inf_{\gamma \in \Pi(p,q)} \int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y$に対応する線形最適化問題を, 双対問題によって変形した結果,
$\sup_{f \in C_1} \int_{\mathbb{R}^n} (p(x)-q(x)) \ f(x) {\rm d}x$に対応する問題に形が似通ったものが得られた.

Step 4. 2つの問題が等価であること

最後に左辺を離散化した問題の双対問題

$$
\begin{aligned}
\max_{f_i,h_j} &\sum_{i=1}^N p_i f_i + \sum_{j=1}^N q_j h_j \\
{\rm subject \ to} \ \ &f_i + h_j \le d_{ij} \ \ (i,j=1,\ldots,N)
\end{aligned}
$$

と, 右辺を離散化した問題

$$
\begin{aligned}
\max_{f_i} &\sum_{i=1}^N (p_i - q_i) f_i \\
{\rm subject \ to} \ \ &f_i - f_j \le d_{ij} \ \ (i,j=1,\ldots,N) \\
\end{aligned}
$$

が等価な問題であることを示します.
これで$\inf_{\gamma \in \Pi(p,q)} \int_{\mathbb{R}^n \times \mathbb{R}^n} \|x-y\| \gamma(x,y) {\rm d}x {\rm d}y=\sup_{f \in C_1} \int_{\mathbb{R}^n} (p(x)-q(x)) \ f(x) {\rm d}x$を,
離散化したもの同士で確認することができます.

さて, 等価性を示すためには, 最適解において$f_i + h_i =0$であることを示せばよいでしょう.
いま$p_i,q_j \ge 0$なので, 目的関数は$f_i,h_j$が大きいほど大きくなります.
$f_i,h_j$たちにある制約は$f_i + h_j \le d_{ij}$の形の制約だけなので, 最適解においては
$$
f_i = \min_j (d_{ij} - h_j) \\
h_j = \min_i (d_{ij} - f_i)
$$がそれぞれ成り立ちます. したがって,
$$
\begin{aligned}
f_i &= \min_j(d_{ij} - h_j) \\
&= \min_j(d_{ij} - \min_k(d_{jk} - f_k))
\end{aligned}
$$ここで$i,k=1,\ldots,N$に対して, $j$が$f_i=d_{ij}-h_j$を満たすとすると,

$$
\begin{aligned}
f_k - f_i &= f_k - d_{ij} + h_j \\
&\le f_k - d_{ij} + (d_{jk} - f_k) \\
&= d_{jk} - d_{ij} \\
&\le d_{ik}
\end{aligned}
$$が得られます. ここで最後の評価は$d_{ik}$に関する三角不等式
$$
d_{jk} = \|x_j-x_k\| \le \|x_j-x_i\| + \|x_i-x_k\| = d_{ij}＋d_{ik}
$$を用いています. $k=1,\ldots,N$のすべてで$f_k-f_i \le d_{ik}$が成り立つので,
$$
-f_i \le \min_k(d_{ik} - f_k) = h_i
$$すなわち$f_i+h_i \ge 0$です.
一方, 問題中の制約式$f_i+h_j \le d_{ij}$において$j=i$とすると,
$d_{ii}=0$なので$f_i + h_i \le 0$を得ます.
以上より最適解においては$f_i + h_i = 0$となることが示されました.

左辺を離散化した問題は, 右辺を離散化した問題と双対の関係にある.

まとめ

上記の4step（＋前半の文章）にしたがって,

Kantorovich-Rubinstein双対性は、離散化した場合には線形最適化の双対性に相当すること.
WGANの設計にはこの双対性が使われていること.

を確認しました.

参考文献

本内容は主に元論文およびOptimal transport, Old and newを参考にしました.