[ICML 2017] Wasserstein Generative Adversarial Networks (WGAN)

Introduction

冷知识: Wasserstein 中的 $W$ 念 $v$ (因为是俄语)

• 在原始 GAN 中，discriminator 衡量的是两个分布间的 JS 散度，而作者指出了使用 JS 散度会导致训练不稳定。为了解决上述问题，作者使用 Wasserstein Distance 来衡量两个概率分布间的距离，相应地提出了 WGAN 来提高 GAN 模型训练的稳定性，避免 mode collapse，提供可用于 debug 和超参搜索的有意义的训练曲线

Wasserstein Distance (Earth Mover’s Distance)

Wasserstein Distance

Earth Mover’s Distance

• Considering one distribution $P$ as a pile of earth, and another distribution $Q$ as the target.
  
• There many possible “moving plans”. Using the “moving plan” with the smallest average distance to define the earth mover’s distance.
  

---

Formal Definition

• A “moving plan” is a matrix. The value of the element is the amount of earth from one position to another.
  Wasserstein Distance 即为 :
  $$\mathcal{W}[p,q]=\underset{\gamma \in \Pi [p,q]}{\min }\sum _{x_p,x_q}\gamma (x_p,x_q)d(x_p,x_q)$$其中，$\gamma$ 表示运输方案，$\gamma \in \Pi [p,q]$ 表示 $\gamma$ 是一个联合分布，它的边缘分布就是原来的 $p$ 和 $q$，即
  $$\sum _x\gamma (x,y)=p(y),\sum _y\gamma (x,y)=p(x)$$$\gamma (x,y)$ 为从 $x$ 运到 $y$ 的土量，$d(x,y)$ 为成本函数，代表从 $x$ 运输到 $y$ 的成本. 常用的 $d$ 是基于 $l$ 范数衍生出来的，比如
  $$\Vert \mathbf{x}-\mathbf{y}{\Vert }_1,\Vert \mathbf{x}-\mathbf{y}{\Vert }_2,\Vert \mathbf{x}-\mathbf{y}{\Vert }_2^2$$
• 推广到连续概率分布，Wasserstein Distance 为
  $$\mathcal{W}[p,q]=\underset{\gamma \in \Pi [p,q]}{\inf }\iint \gamma (\mathbf{x},\mathbf{y})d(\mathbf{x},\mathbf{y})d\mathbf{x}d\mathbf{y}$$其中，$\gamma (\mathbf{x},\mathbf{y})$ 表示要从 $x$ 处搬 $\gamma (x,y)dx$ 那么多的东西到 $y$ 处，$\inf$ 表示下确界

Why Earth Mover’s Distance?

Different Distances between two distributions

• The Total Variation (TV) distance
  其中，$\mathcal{X}$ 为图像空间 $[0,1{]}^d$，$\Sigma$ 为 $\mathcal{X}$ 的所有 Borel subsets 的集合
• The Kullback-Leibler (KL) divergence
  
• The Jensen-Shannon (JS) divergence
  其中，${\mathbb{P}}_m=({\mathbb{P}}_r+{\mathbb{P}}_g)/2$
• The Earth-Mover (EM) distance or Wasserstein-1
  其中，$\Pi \left({\mathbb{P}}_r,{\mathbb{P}}_g\right)$ 为所有联合分布 $\gamma (x,y)$ 的集合，$\gamma (x,y)$ 的边缘分布为 ${\mathbb{P}}_r,{\mathbb{P}}_g$

---

假设数据集分布为 ${\mathbb{P}}_0$，模型学得的数据分布为 ${\mathbb{P}}_{\theta }$. 下面作者给出了具体的例子来说明在某些情况下，使用 EM 距离作为概率分布间的距离度量能够使得 ${\mathbb{P}}_{\theta }$ 收敛，而使用其余距离度量则不能收敛

• 在上述例子中，只有 EM 距离能收敛的原因就在于，当两个概率分布不重合时，只有 EM 距离是连续的并能为模型参数提供有效的梯度信息
  

---

$W({\mathbb{P}}_r,{\mathbb{P}}_{\theta })$ is a continuous loss function on $\theta$ under mild assumptions

• 这种连续的性质对于实际的图像生成任务是十分重要的。首先，假设数据集样本分布为 $P_{data}$，模型生成的样本分布为 $P_G$，则在大部分情况下，这两个分布都是没有交集的。这是因为图像是高维空间中的低维流形，两个不同的图像分布往往交集很小，同时我们往往是通过采样的方式来估计分布，采样不充分也会导致采样出的分布没有交集。而以原始 GAN 使用的 JS 散度为例，当两个分布不重合时，JS 散度就为 $\log 2$，当刚开始训练 GAN 时， $P_G$ 和 $P_{data}$ 一直都没有重合，因此 JS divergence 一直为 $\log 2$，使得 Generater 的 loss function $L(D)$ 一直为 0，训练难以收敛
  相比之下，EM 距离就好在它在衡量两个低维流形上分布间的距离时更加敏感，它处处连续，几乎处处可导且导数不为 0
  

---

作者进一步用以下定理证明了使用 EM 距离进行训练的优越性

Wasserstein GAN

WGAN (Weight Clipping)

• EM 距离里有一个下确界比较难以处理，因此现在的问题就是如何去估算出 EM 距离。Kantorovich-Rubinstein duality 给出了估算方法：
  $$W\left({\mathbb{P}}_r,{\mathbb{P}}_{\theta }\right)=\underset{\Vert f{\Vert }_L\le 1}{\sup }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_{\theta }}[f(x)]$$其中，$f:\mathcal{X}\to \mathbb{R}$ 为所有满足 1-Lipschitz 性质的函数。如果 $f$ 仅满足 $K$-Lipschitz 性质，则有
  $$K\cdot W\left({\mathbb{P}}_r,{\mathbb{P}}_{\theta }\right)=\underset{\Vert f{\Vert }_L\le K}{\sup }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_{\theta }}[f(x)]$$其中，$K$-Lipschitz 是指
  $$\Vert f(x_1)-f(x_2)\Vert \le K\Vert x_1-x_2\Vert$$($K=1$ for “1 −Lipschitz” $\Rightarrow$ $||f(x_1)-f(x_2)||\le ||x_1-x_2||$) 也就是说，Lipschitz 性质是要求 $f$ 足够光滑，当输入发生改变时，输出发生的改变不会太大 (function does not change fast). 试想如果没有 Lipschitz 约束，由于 real data 和 generated data 之间很少会有 overlap，那么 $f$ (i.e., discriminator) 会使 real data 对应的值趋近于 $\infty$，generated data 对应的值趋近于 $-\infty$，因此训练永远都不会收敛。因此，WGAN 的优化目标为
  
• 那么我们如何保证 $f$ (i.e., discriminator) 满足 $K$-Lipschitz 呢？WGAN 采用的方法是在每次梯度更新后进行 weight clip，保证权重 $w$ 里的每个值都介于 $c$ 和 $-c$ 之间 (if $w>c$, $w=c$; if $w<-c$, $w=-c$)。直观上看，$w$ 里的值都比较小，因此当输入变化时，输出的变化总是有限的

WGAN 优化器的选择：作者发现 WGAN 在使用 Adam (with ${\beta }_1>0$) 等 momentum based optimizer 或是采用大学习率时训练就会变得不稳定。对此，作者的解释是 “Since the loss for the critic is nonstationary, momentum based methods seemed to perform worse. We identified momentum as a potential cause because, as the loss blew up and samples got worse, the cosine between the Adam step and the gradient usually turned negative. The only places where this cosine was negative was in these situations of instability.” 因此，作者采用 RMSProp 优化器，该优化器在 nonstationary problems 上表现较好

---

Limitation of Weight Clipping

• (1) We only ensure that $\Rightarrow$ $||f(x_1)-f(x_2)||\le K||x_1-x_2||$ for some $K$. 最后得出来的 wasserstein distance 时原来的 $K$ 倍，并且这个 $K$ 是未知的
• (2) Do not truly find function $f$ maximizing the function: 有可能 $w$ 不满足 weight clipping 的条件，但也可以使 $f$ 满足 1-Lipschitz 限制；也就是说，weight clipping 只覆盖了满足 1-Lipschitz 限制的所有函数的一个 subspace

Two main benefits of WGAN

Meaningful loss metric

• 由于 discriminator 衡量的是 wasserstein distance，因此 discriminator 的损失函数可以直观地反映 generator 是否收敛及其生成质量
• 如果不使用 WGAN，由于 real data 和 generated data 之间通常没有 overlap，因此 JS divergence 经常保持在 $\log 2\approx 0.69$，无法通过 loss function 来判断 real data 和 generated data 的相似程度
  而使用 WGAN 后，就可以用 wasserstein distance 来衡量 GAN 训练的好坏了！但需要注意的是，该损失函数并不能直接量化 GAN 的性能，这是因为我们只能约束 discriminator 满足 $K$-Lipschitz 性质，因此最终的损失函数其实是 $K$ 倍的 wasserstein distance. 此外，实际中的 discriminator 也不是最优的，该损失函数也只是 wasserstein distance 的一个近似
  

---

Improved stability

• 在设计 GAN 的网络结构时，我们需要合理地平衡 generator 和 discriminator 的性能来保证训练的稳定性，一旦 discriminator 的性能过强，能完美区分生成数据和真实数据，JS 散度就会变成一个常数，损失函数无法给 generator 提供有效的梯度信息，训练就会停滞。通常的解决方法有 (1) 削弱 discriminator 的能力: 通过给它添加 Dropout、减少更新次数，让它无法 overfitting；但 Discriminator 能力过弱也有问题，它就无法衡量 JS divergence 了… (2) Add Noise (Noises decay over time): Add some artificial noise to the inputs of discriminator; Make the labels noisy for the discriminator (train discriminator 时，给 input / label 加 noise) $\Rightarrow$ $P_{data}(x)$ and $P_G(x)$ have some overlap. Discriminator cannot perfectly separate real and generated data. 而由于 WGAN 采用 wasserstein distance，因此我们不用特别考虑 generator 和 discriminator 的平衡，可以大胆地训练 discriminator 到最优，discriminator 性能越好，衡量的 wasserstein distance 就越准确，给 generator 提供的梯度信息也就越有意义
• 此外，作者还观察到 WGAN 对 generator 模型架构的选择更具鲁棒性，能极大地避免了 GAN 中经常出现的 mode collapse
  

References

• Martin Arjovsky, Soumith Chintala, Léon Bottou. Wasserstein Generative Adversarial Networks. Proceedings of the 34th International Conference on Machine Learning, PMLR 70:214-223, 2017.
• 李宏毅 2021 ML 课程