Wasserstein GAN and Beyond

Posted on
GAN Deep Learning
thumbnail

$$ \DeclareMathOperator{\E}{\mathbb{E}} $$

本文是关于GAN[1]训练稳定性的几篇文章的笔记。 主要包括(但不仅限于)文献[2]分析了原生GAN[1]训练不稳定的原因在于其优化目标等价于 最小化真实数据分布$p_r$和生成数据分布$p_g$之间的J-S散度 $D_{JS}(p_r || p_g)$。 文献[3]提出使用连续性和数值稳定性更好的Wasserstein distance来代替原生GAN中的J-S散度。 文献[4]提出对判别器的参数采用“谱归一化” (spectral normalization,这里翻译成谱归一化很蹩脚), 从而使判别器满足lipschitz continuity。 因为在[2]简单的使用参数剪裁(weight clip)来限制判别器的参数,这样做过于粗暴而且会影响模型收敛。

在GAN原文[1]中整个模型(包括G和D)的优化目标为:

\begin{equation} \begin{split} \min_G \max_D V(D, G) &= \E_{x \sim p_r} \left[\log D(x)\right] + \E_{z \sim p_z(z)} \left[\log(1 - D(G(z)))\right] \\ &= \E_{x \sim p_r} \left[\log D(x)\right] + \E_{x \sim p_z{z}} \left[\log(1 - D(G(z)))\right] \\ &= \E_{x \sim p_r} \left[\log D(x)\right] + \E_{x \sim p_g(x)} \left[\log(1 - D(x))\right] \\ &= \int_x \left[p_r(x)\log D(x) + p_g(x)\log(1-D(x))\right]dx \end{split} \label{eq:gan-loss} \end{equation}

这里解释一下$\ref{eq:gan-loss}$式的含义: 对$G$而言最好的状态就是$D(G(Z))=1$,也就是让$D$觉得G生成的图片是真实的, 此时$\ref{eq:gan-loss}$式中右边 $\log(1-D(G(z))) = \log{0} = -\infty$; 对$D$而言就是$D(x) = 1, x \sim p_r$,以及$D(G(z)) = 0$,此时$\ref{eq:gan-loss}$式为$\log{1} + \log{1} = 0$。 因此整个$\ref{eq:gan-loss}$式在$G$上求损失的最小值,在$D$上求损失的最大值。

事实上$\ref{eq:gan-loss}$式可以拆分成$D$和$G$各自的两个损失函数:

  • 判别器D的损失
\begin{equation} L(D, g_{\theta}) = \mathbb{E}_{x \sim p_r}\left[\log D(x)\right] + \mathbb{E}_{x \sim p_g}\left[1 - \log D(x)\right] \label{eq:loss-d} \end{equation}
  • 生成器G的损失:
\begin{equation} L(G, g_{\theta}) = 2D_{JS}(\mathbb{P}_r || \mathbb{P}_g) - 2\log2 \label{eq:loss-g} \end{equation}

这里G的损失写成了真实数据的分布$\mathbb{P}_r$和生成数据分布$\mathbb{P}_g$之间的JS散度的形式, 后面我们将给出证明。

典型的GAN的训练过程是根据$\ref{eq:loss-d}$式迭代训练D,然后根据$\ref{eq:loss-g}$式训练G。 当训练D时,我们根据$\ref{eq:gan-loss}$式很容易得到最佳的判别器$D^*$:

\begin{equation} \partial D(x) = \frac{p_r}{D(x)} + p_g \frac{-1}{1-D(x)} = 0 \Rightarrow D^* = \frac{p_r}{p_r+p_g} \label{eq:optimal-d} \end{equation}

这里有一个重要结论:当D是optimal的时候,优化$G$实际上是在优化$p_r$和$p_g$的J-S散度。 将$\ref{eq:optimal-d}$式的结论带入$\ref{eq:gan-loss}$式,得到:

\begin{equation} \begin{split} \min_G \max_D V(D^*, G) &= \min_G V(D^*, G) \\ &= \int_x \left[ p_r(x)\log D^*(x) + p_g(x) \log (1 - D^*(x)) \right] \\ &= \int_x \left[ p_r(x)\log \frac{p_r(x)}{p_r(x)+p_g(x)} + p_g(x) \log (\frac{p_g(x)}{p_r(x)+p_g(x)}) \right] \end{split}\label{eq:loss-optimal-d} \end{equation}

下面我们会介绍KL散度(Kullback–Leibler divergence) 和J-S散度(Jensen–Shannon divergence), 并且证明$\ref{eq:loss-optimal-d}$式的优化目标就是$\ref{eq:loss-g}$式中J-S散度的形式。


K-L 散度和 J-S 散度

K-L散度是两个概率分布$P, Q$之间的距离的一种度量:

\begin{equation} D_{KL}(P, Q) = \int_x p(x) \cdot \log\frac{p(x)}{q(x)}\label{eq:kld} \end{equation}

根据$\ref{eq:kld}$式,KL散度是不对称的 ($D_{KL}(P, Q)\neq D_{KL}(Q, P)$),实际上不满足成为度量的条件。 而 JS散度是基于 K-L 散度的一种改进,满足对称性:

\begin{equation} D_{JS}(P, Q) = \frac{1}{2}D_{KL}(P, M) + \frac{1}{2}D_{KL}(Q, M) = D_{JS}(Q, P) \end{equation}

其中$M=\frac{1}{2}(P + Q)$,也即$m(x) = \frac{1}{2}(p(x) + q(x))$

JS(KL)散度的问题:当分布没有交集

KL散度有一个缺陷,就是当两个分布$P$和$Q$没有分布的时候,会得到无定义的散度值。 例如下图中的两个高斯分布,没有相交的部分(或者相交的部分其中一个分布的概率值特别小),这样算出来的散度值会趋向于无穷。 由于 J-S 散度的计算基于K-L散度,因此也继承了这个缺陷。

上图中的两个高斯分布就没有重合的地方(没有交集),这会导致计算K-L(J-S)散度的时候 有一部分点$\frac{p(x)}{q(x)}$中的分母$q(x)$为零,导致整个散度无定义。


原生的GAN实际上在优化$D_{JS}(p_r, p_g)$

这里我们证明原生的GAN实际上在优化数据分布$p_r$和生成数据的分布$p_g$之间的J-S散度。 当$D$达到最优判别器 $D^*$ 时,GAN的优化目标从$\displaystyle\min_G \max_D V(D, G)$变成了$\displaystyle\min_G V(D^*, G)$ 我们继续展开Eq.\ref{eq:loss-optimal-d}

\begin{equation} \begin{split} D_{JS}(p_r, p_g) &= \frac{1}{2}D_{KL}(p_r, \frac{p_r+p_g}{2}) + \frac{1}{2}D_{KL}(p_g, \frac{p_r+p_g}{2}) \\ &= \frac{1}{2} \left[ \int_x p_r(x)\log\frac{2p_r(x)}{p_r(x)+p_g(x)} + \int_x p_g(x)\log\frac{2p_g(x)}{p_r(x)+p_g(x)}\right] \\ &= \frac{1}{2}\left[ \log4 + \int_x p_r(x)\log\frac{p_r(x)}{p_r(x)+p_g(x)} + \int_x p_g(x)\log\frac{p_g(x)}{p_r(x)+p_g(x)} \right] \\ &= \frac{1}{2}\left[ \log4 + V(D, G)\right] \end{split}\label{eq:djs} \end{equation}

接合$\ref{eq:loss-optimal-d}$式,我们得到当$D^*$固定时我们的优化目标为:

\begin{equation} \begin{split} \min_G V(D^*, G) &= \int_x \left[ p_r(x)\log D^*(x) + p_g(x) \log (1 - D^*(x)) \right] \\ &= \int_x \left[ p_r(x)\log \frac{p_r(x)}{p_r(x)+p_g(x)} + p_g(x) \log (\frac{p_g(x)}{p_r(x)+p_g(x)}) \right] \\ &= D_{JS}(p_r, p_g) - \log4 \end{split} \end{equation}

前面提到过 J-S 散度数值不稳定的缺陷,因此当D已经达到最优,正在训练G的时候,会使得G得到的梯度不平滑,造成G难以训练。


Wasserstein Distance

既然GAN训练不稳定是由JS散度造成的,那么有没有其他的概率分布的距离度量,能够避免数值不稳定的问题呢? 答案当然是有的,就是Wasserstein Distance。 Wasserstein Distance解释起来比较麻烦,我首先从一个例子开始。

假设有五个城市生产鞋子,产能分别是$A_1, … A_5$; 另外有五个城市需要从生产鞋子的五个城市购买鞋子,需求量分别是$B_1, … B_5$。 假设$\sum_{i=1}^5 A_i = \sum_{i=1}^5 B_i$,也就是产能刚好满足需求。 假设 $\mathbf{G} \in \mathbb{R}^{5 \times 5}$ 是一个 $5\times 5$ 的矩阵,$g_{i,j}$ 表示从 $A_i$ 运输一双鞋子到$B_j$的成本。 假设有运输方案 $\mathbf{T} \in \mathbb{R}^{5 \times 5}$ 其中 $T_{i,j}$ 表示从 $A_i$ 运输到 $B_j$ 鞋子的数量。 很显然我们有

\begin{equation} \begin{split} \sum_i \mathbf{T}_{i, j} = B_j \\ \sum_j \mathbf{T}_{i, j} = A_i \end{split}\label{eq:constrain} \end{equation}

用数学语言表示,就是“联合概率分布$\mathbf{T}$的边缘分布(marginal distribution)分别是$A$和$B$”。 不难得到这个方案的总运输成本为:

\begin{equation} \mathcal{L} = \sum_{i, j}\mathbf{G}_{i, j}\cdot\mathbf{T}_{i, j} \end{equation}

好了,基本模型描述完毕了,现在求一个最佳的从$A_1, .. A_5$运输鞋子到$B_1, …, B_5$的运输方案 $\mathbf{T}^* $,能够拥有最小的运输成本$\mathcal{L}^*$:。

\begin{equation} \begin{split} \mathcal{L}^* &= \inf_{T \in \prod (A, B)} \sum_{i, j}\mathbf{G}_{i, j}\cdot\mathbf{T}_{i, j} \\ &= \inf_{T \in \prod (A, B)} \ <\mathbf{G}, \mathbf{T}>_F \end{split} \label{eq:wasserstein} \end{equation}

$\ref{eq:wasserstein}$就是Wasserstein distance的定义,其中$T \in \prod (A, B)$表示 “所有边缘分布分别是$A$和$B$的分布的集合”。 这样的问题又叫最优传输问题(optimal transport)[6],离散情况下的Wasserstein Distance 可以形象的比喻为把土堆从一个地方移动到另一个地方的最小传输损耗,因此又称为 Earth Mover’s Distance (EMD)。

一般情况下我们可以用$i, j$两点的欧氏距离(Euclidean distance)作为两点的ground-metric, 并将A和B归一化为概率分布($A = \frac{A}{\sum A}, B = \frac{B}{\sum B}$, 也即$\sum_{i, j} T_{i, j}=1$), $\ref{eq:wasserstein}$式又可以写为:

\begin{equation} \begin{split} \mathcal{L}^* &= EMD(A, B) \\ &= \inf_{T \in \prod (A, B)} \ \mathbb{E}_{(i, j) \sim T} \mathbf{G}_{i, j} \mathbf{T}_{i, j} \\ &= \inf_{T \in \prod (A, B)} \ \mathbb{E}_{(i, j) \sim T} ||i - j||^2_2 \end{split} \label{eq:wasserstein1} \end{equation}

$\ref{eq:wasserstein}$式中的问题是一个标准的线性规划(Linear programming)问题: 在线性约束下($\ref{eq:constrain}$式)求解线性方程($\ref{eq:wasserstein}$式)的最值。

虽然 Wasserstein Distance 比 J-S 散度更适合做概率分布的度量,但是$\ref{eq:wasserstein}$式并 不是那么容易的求得结果。 在我们的例子里 $A$ 和 $B$ 都是一维分布,而且只有五种状态。 在现实应用中图像基本都有万以上个像素点,每个像素有3个颜色值,每个值有256种状态。 $\ref{eq:wasserstein}$式的求解会随着维度增加呈指数增长。 因此在每一次迭代中都求解$\ref{eq:wasserstein}$式显然是无法做到的。

而且,我们并不需要这个最优的传输 $\mathbf{T}^* $,我们只对最优传输对应的传输成本 $L^* $ 感兴趣。 同时,为了能够通过 $L^* = EMD(p_r, p_g)$ 训练生成器 $G$,我们还需要求得 $\frac{\partial EMD(p_r, p_g)}{\partial p_g}$。

下面一节会介绍 K-R 对偶 (Kantorovich-Rubinstein Duality),通过K-R对偶可将$\ref{eq:wasserstein}$式的问题 转换为更容易求解的问题。

到目前为止我们讲的一维分布的EMD能够很好的通过例子来理解。接下来在具体应用中,图像数据的维度有上万, 因此 $p_r$ 和 $p_g$ 都是维度非常高的概率分布。

总结一下我自己一开始理解上出现的误差:

  • 我们要计算的是真实数据的分布 $p_r$ 和生成数据的分布 $p_g$ 之间的Wasserstein Distance;
  • $p_r$ 和 $p_g$ 是多元概率分布,其随机变量个数等于图像像素个数;
  • 每幅图像本身只是从 $p_r$ (或者是 $p_g$) 采样出来的一个数据点;
  • 不要因为图像是二维的就把** $p_r$,$p_g$ 想象成二维分布。


    K-R 对偶 (Kantorovich-Rubinstein Duality)

    事实上每一个线性规划问题都有对应的对偶问题 (dual problem)。 Kantorovich-Rubinstein Duality 说:对于$\ref{eq:wasserstein}$式这样的线性规划问题,当 $G_{i, j} = ||i - j||$ 是欧氏距离 (Euclidean distance)的时候,等价于以下对偶问题:

\begin{equation} EMD(p_r, p_g) = \sup_{\lVert f\rVert_L < 1} \mathbb{E}_{x \sim p_r} \ f(x) - \mathbb{E}_{x \sim p_g} \ f(x) \label{eq:k-r-dual} \end{equation}

其中$\lVert f\rVert_L < 1$ 表示函数 $f(x)$ 满足“Lipschitz连续性” (Lipschitz continuity): 对任意的$x \neq y$, 有$\frac{f(y) - f(y)}{x - y} <= K$。 简而言之,Lipschitz 连续性要求函数 $f(x)$ 在任意局部定义域内的坡度是有限的。 注意这里并不要求f(x)可导,因此不可导函数 $f(x) = |x|$ 是满足 lipschitz 连续性的; 而可导函数 $f(x) = \frac{1}{x}$ 并不满足Lipschitz连续性。

$\ref{eq:k-r-dual}$式说明,在 $G_{i, j} = \lVert i - j \rVert$ 的条件下, 求$\ref{eq:wasserstein}$式中的最优传输代价,等价于求使得$\ref{eq:k-r-dual}$ 式取最小值并且满足 Lipschitz 连续性的 $f(x)$。

关于K-R对偶的证明,可以参考这篇博客: Wasserstein GAN and the Kantorovich-Rubinstein Duality

Tobe completed


参考资料


参考文献

[1] Goodfellow, Ian, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. “Generative adversarial nets.” In Advances in neural information processing systems, pp. 2672-2680. 2014. [PDF]

[2] Arjovsky, Martin, and Léon Bottou. “Towards principled methods for training generative adversarial networks.” arXiv preprint arXiv:1701.04862 (2017). [PDF]

[3] Arjovsky, Martin, Soumith Chintala, and Léon Bottou. “Wasserstein generative adversarial networks.” In International Conference on Machine Learning, pp. 214-223. 2017. [PDF]

[4] Miyato, Takeru, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. “Spectral normalization for generative adversarial networks.” arXiv preprint arXiv:1802.05957 (2018). [PDF]

[5] Gulrajani, Ishaan, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C. Courville. “Improved training of wasserstein gans.” In Advances in Neural Information Processing Systems, pp. 5767-5777. 2017. [PDF]

[6] Villani, Cédric. Optimal transport: old and new. Vol. 338. Springer Science & Business Media, 2008. [PDF]