论文阅读——《Wasserstein GAN》《Improved Training of Wasserstein GANs》

论文阅读之 Wasserstein GAN 和 Improved Training of Wasserstein GANs

本博客大部分内容参考了这两篇博客: 再读WGAN（链接已经失效）和令人拍案叫绝的Wasserstein GAN, 自己添加了或者删除了一些东西, 以及做了一些修改.

基础知识：

f-Divergence

原始GAN采用的是JS divergence来衡量两个分布之间的距离。事实上有一个统一的模式来衡量两个分布间的距离，它就是f-divergence。
假设 $P$和 $Q$是两个分布。 $p(x)$和 $q(x)$是对应样本 $x$的概率，则：

$D_f(P||Q)=\int_xq(x)f(\frac{p(x)}{q(x)})dx$
就是f-Divergence, 其中 $f$需要满足: 1. 是凸函数, 2. $f(1)=0$.
可以看到对任意的 $x$都有 $p(x)=q(x)$, 则 $D_f(P||Q)=\int_xq(x)f(\frac{p(x)}{q(x)})dx=0$, 即两个分布相同的时候f-Divergence等于0.

想要让f-divergence能表示两个分布的距离, 不但要求当两个分布相同时距离是0, 还需要保证这个0是f-Divergence能取到的最小值, 证明如下:
由于 $f(x)$是凸函数，则有下面的不等式：
$D_f(P||Q)=\int_xq(x)f(\frac{p(x)}{q(x)})dx \ge f(\int_xq(x)\frac{p(x)}{q(x)}dx)=f(1)=0$

Fenchel Conjugate

每个凸函数 $f$都具备一个共轭函数 $f^∗$：
$f^∗(t)=\max_{x\in dom(f)}\{xt−f(x)\}$
关于Fenchel Conjugate有两个性质：

1. 所有的凸函数f都有一个conjugate函数 $f^∗$;
2. $((f^∗)^∗)=f$.

和GAN的联系

假设 $f^*$是 $f$的Fenchel Conjugate函数, 则由上面的式子可以得到:
$f(x)=\max_{t\in dom(f^*)}\{tx-f^*(t)\}$
将这个 $f$作为上面f-Divergence中的 $f$, 代入f-Divergence的式子得到:
$D_f(P||Q)=\int_xq(x)f(\frac{p(x)}{q(x)})dx=\int_xq(x)(\max_{t\in dom(f∗)}\{\frac{p(x)}{q(x)}t−f^∗(t)\})dx$
此时我们如果构建一个函数 $D(x)\in dom(f^*)$, 即输入 $x$, 输出在 $f^*$的定义域中, 这样我们就能有 $D(x)$代替 $t$, 但是这种替换对原来的公式并不是等价的. 因为我们所能用 $D(x)$找到的 $t$并不是那个能够让 $f$最大的那个 $t$. 所以我们替换之后构造的函数永远要小于等于f-Divergence：
$D_f(P||Q)\ge\int_xq(x)(\frac{p(x)}{q(x)}D(x)−f^∗(D(x)))dx\\=\int_xp(x)D(x)dx−\int_xq(x)f^∗(D(x))dx$
这就相当于我们找到了一个 $D_f(P||Q)$的下界。接下来，如果我们能找到一个让上面公式不等号右边最大的 $D$, 那么如果我们在公式中采用了这个找到的 $D$, 那就可以去逼近 $D_f(P||Q)$, 即：
$D_f(P||Q)\approx \max_D\int_xp(x)D(x)dx-\int_xq(x)f^∗(D(x))dx\\=\max_D\{\mathbb{E}_{x\sim P}[D(x)]−\mathbb{E}_{x\sim Q}[f^∗(D(x))]\}$
上面公式的第二行是将前面的概率积分变成了期望. 然而我们在工程上没法真的去求期望, 所以一般的做法就是分别从 $P$和 $Q$中去抽样数据.
假设现在我们的 $P$是 $P_data$， $Q$是 $P_G$，则公式变成：
$D_f(P_{data}||P_G)=\max_D{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[f^∗(D(x))]}$
但实际上想要将f-Divergence用在GAN中是不可能的，因为我们并不知道 $P_{data}$和 $P_G$的表达式. 所以我们这一路的公式推导, 最终得到的公式只需要我们简单地从 $P_{data}$和 $P_G$中采样就可以计算得到f-Divergence. 而不再需要知道这两个分布的表达式即可计算.

接下来，假如我们想要去寻找一个能让这个距离最小的 $P_G$，则这个G应该是:
$G^∗=\arg \min_GD_f(P_{data}||P_G)\\=\arg \min_G\max_D{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[f^∗(D(x))]}\\=\arg\min_G\max_DV(G,D)$

WGAN

前面我们介绍了使用f-Divergence来将“距离”定义到一个统一框架之中的方法. 而Fenchel Conjugate则将这个f-Divergence与GAN联系在一起。这么做的目的在于, 我们只要能找到一个符合f-Divergence要求的函数，就能产生一个距离的度量, 从而定义一种不同的GAN。

原生GAN的问题

对于原生的GAN来说, 选择特定的度量函数之后, 会导致目标函数变成生成分布与真是分布的JS divergence. 但是这个divergence有很多问题. 比如说一个最严重的问题就是当两个分布之间完全没有重叠时, 分布间距离的大小并不会直接反映在divergence上. 这对基于迭代的优化算法是个致命问题.

根据原始GAN定义的判别器loss, 我们可以得到最优判别器的形式; 而在最优判别器的下, 我们可以把原始GAN定义的生成器loss等价变换为最小化真实分布 $P_r$与生成分布 $P_g$之间的JS散度. 我们越训练判别器, 它就越接近最优, 最小化生成器的loss也就会越近似于最小化 $P_r$和 $P_g$之间的JS散度.

问题就出在这个JS散度上. 我们会希望如果两个分布之间越接近它们的JS散度越小, 我们通过优化JS散度就能将 $P_g$“拉向” $P_r$, 最终以假乱真. 这个希望在两个分布有所重叠的时候是成立的, 但是如果两个分布完全没有重叠的部分, 或者它们重叠的部分可忽略(下面解释什么叫可忽略), 它们的JS散度是多少呢?
答案是 $log2$, 因为对于任意一个x只有四种可能:
$P_1(x) = 0且P_2(x) = 0$
$P_1(x) \neq 0且P_2(x) \neq 0$
$P_1(x) = 0且P_2(x) \neq 0$
$P_1(x) \neq 0且P_2(x) = 0$
第一种对计算JS散度无贡献, 第二种情况由于重叠部分可忽略所以贡献也为0, 第三种情况对公式7右边第一个项的贡献是 $\log \frac{P_2}{\frac{1}{2}(P_2 + 0)} = \log 2$, 第四种情况与之类似, 所以最终 $JS(P_1||P_2) = \log 2$.

换句话说, 无论 $P_r$跟 $P_g$是远在天边, 还是近在眼前, 只要它们俩没有一点重叠或者重叠部分可忽略, JS散度就固定是常数 $\log 2$, 而这对于梯度下降方法意味着——梯度为0! 此时对于最优判别器来说, 生成器肯定是得不到一丁点梯度信息的; 即使对于接近最优的判别器来说, 生成器也有很大机会面临梯度消失的问题.

但是 $P_r$与 $P_g$不重叠或重叠部分可忽略的可能性有多大? 不严谨的答案是: 非常大. 比较严谨的答案是：当 $P_r$与 $P_g$的支撑集(support)是高维空间中的低维流形(manifold)时, $P_r$与 $P_g$重叠部分测度(measure) 为0的概率为1.

论文里面针对这个也做了相应的实验:

这个图的大致意思是作者用不同的网络结构做了实验, 发现随着迭代次数的增加, 生成图像的质量在提高, 但是真实数据和生成数据之间的JS散度却增加了或者是为一个常数, 这就从实验上验证了原来的GAN利用JS散度来学习的问题.

Earth Mover’s Distance

假设我们有下面的两个分布:

如何将P上的内容“匀一匀”得到Q呢? 比如说把最高的哪一条分开一部分分到其他地方? 这或许是一种解决方案:

但是显然除此之外还有很多种方法, 例如:

既然移动的方法有很多种, 如果每一种都表示了一种代价, 那么显然有“好”方法, 就会有“坏”方法. 假设我们衡量移动方法好坏的总代价是“移动的数量”X“移动的距离”. 那这两个移动的方案肯定是能分出优劣的.当我们用分布Q上不同颜色的色块对应分布P的相应位置, 就可以将最好的移动方案画成下面这个样子:

我们可以将这个变化画成一个矩阵:

对于每一个移动方案 $\gamma$, 都能有这样一个矩阵. 矩阵的每一行表示分布 $P$的一个特定位置. 该行中的每一列表示需要将该行的内容移动到分布 $Q$对应位置的数量. 即矩阵中的一个元素 $(x_p,x_q)$表示从 $P(x_p)$移动到 $Q(x_q)$的数量。

而对于方案 $\gamma$我们可以定义一个平均移动距离(Average distance of a plan $\gamma$):
$B(\gamma)=\sum_{x_p, x_q}\gamma(x_p, x_q)||x_p−x_q||$
而Earth Mover’s Distance就是指所有方案中平均移动距离最小的那个方案：
$W(P,Q)=\min_{\gamma\in\Pi}B(\gamma)$
其中 $\Pi$是所有可能的方案.
作者在论文中也通过实验说明了EM距离比JS散度更能反应连个分布间的距离:

上面两幅图说明在训练过程中, 随着EM距离的减小, 生成图片的质量越来越高; 而下面这幅图b表示EM距离为一个常数, 生成的图片也就很差.

EM距离与GAN结合

回忆一下f-Divergence：
$D_f(P_{data}||P_G)=\max_D{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[f^∗(D(x))]}$
而WGAN的文章中写到，EM距离也可以类似f-Divergence，用一个式子表示出来：
$W(P_{data},P_G)=\max_{D\in1-Lipschitz}{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[D(x)]}$
公式中1-Lipschitz表示了一个函数集. 当 $f$是一个Lipschitz函数时，它应该受到以下约束:
$||f(x_1)−f(x_2)||\le K||x_1−x_2||$
当K=1时，这个函数就是1-Lipschitz函数。
为什么要限制生成器D时1-Lipschitz函数呢
假设我们现在有两个一维的分布, $x_1$和 $x_2$的距离是 $d$, 显然他们之间的EM距离也是 $d$:

此时如果我们想要去优化
$W(P_{data},P_G)=\max_{D\in1-Lipschitz}\{{\mathbb{E}_{x∼P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[D(x)]}$
只需要让 $D(x_1)=+\infty$，而让 $D(x_2)=−\infty$就可以了. 也就是说，如果不加上1-Lipschitz的限制的话, 只需要让判别器判断 $P_{data}$时大小是正无穷, 判断PG时是负无穷就足够了. 这样的判别器可能会导致训练起来非常困难: 判别器区分能力太强, 很难驱使生成器让生成分布适应数据分布.

这个时候我们加上了这个限制, 也就是说 $||D(x_1)−D(x_2)||\le||x_1−x_2||=d$. 此时如果我们想要满足上面的优化目标的话, 就可以让 $D(x_1)=k+d$, 让 $D(x_2)=k$. 其中 $k$具体是什么无所谓, 关键是我们通过 $d$将判别器在不同分布上的结果限制在了一个较小的范围中. 传统的GAN所使用的判别器是一个最终经过sigmoid输出的神经网络, 它的输出曲线肯定是一个S型. 在真实分布附近是1, 在生成分布附近是0. 而现在我们对判别器施加了这个限制, 同时不再在最后一层使用sigmoid, 它有可能是任何形状的线段, 只要能让 $D(x_1)−D(x_2)\le d$即可. 如下图所示:

这样做的好处显而易见. 传统GAN的判别器是有饱和区的, 而现在的GAN如果是一条直线, 那就能在训练过程中无差别地提供一个有意义的梯度.说了这么多, WGAN主要的变化在这两点:

1. 不要用sigmoid输出;
2. 换成受限的1-Lipschitz来实现一个类似sigmoid的“范围限制”功能.

1-Lipschitz限制应该如何施加?
文章中所用的方法是截断权重.

WGAN的伪代码如下:

• 初始化D的 $θ_d$和G的 $θ_g$
• 在每一个训练循环进行：
  • 从数据分布 $P_{data(x)}$中采样m个样本 $\{x_1,x_2,…,x_m\}$
  • 从先验噪声分布 $P_{prior(z)}$中采样m个样本 $\{z_1,z_2,…,z_m\}$
  • 将这些噪声样本输入生成器G，得到生成样本 $\{\tilde{x}_1,\tilde{x}_2,…,\tilde{x}_m\},\tilde{x}_i=G(z_i)$
  • 更新判别器的参数θd，即最大化：
    • $\tilde{V} =\frac{1}{m}∑^m_{i=1}D(x_i)−\frac{1}{m}∑^m_{i=1}D(\tilde{x}_i)$
    • $θ_d\leftarrowθ_d+\eta \nabla \tilde{V} (θ_d)$
    • 更新参数后，截断参数
• 从先验噪声分布 $P_{prior(z)}$中再采样m个样本 $\{z_1,z_2,…,z_m\}$
• 更新生成器的参数θg，即最小化：
  • $\tilde{V} =\sout{\frac{1}{m}\sum^{m}_{i=1}logD(x_i)}−\frac{1}{m}∑^m_{i=1}D(G(z_i))$
  • $θ_g\leftarrowθ_g−\eta \nabla \tilde{V}(θ_g)$

尤其需要注意的是, 判别器的输出不再需用sigmoid函数了! 并且需要训练k次判别器, 然后只训练一次生成器.

Gradient Penalty

WGA中试通过权重裁剪来施加1-Lipschitz限制的, 但是做样做会有一些问题, 比如通过权重裁剪并不能保证判别器满足1-Lipschitz限制, 而是K-Lipschitz限制, 而且权重裁剪很可能讲那些满足1-Lipschitz限制的函数给剪掉. 而WGAN GP提出了通过施加梯度惩罚来施加1-Lipschitz限制的条件.

Improved WGAN 的文章中提到, 1-Lipschitz函数有一个特性: 一个可微函数是1-Lipschitz函数时, 当且仅当它的梯度的norm将永远小于等于1, 即:
$D\in1-Lipschitz\Leftrightarrow||∇xD(x)||\le1 for\ all\ x$
有了这个理论, 我们就可以改变我们的目标函数了.
原来我们优化目标是:
$W(P_{data},P_G)=\max_{D\in1-Lipschitz}\{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x∼P_G}[D(x)]\}$
此时WGAN的优化目标是在1-Lipschitz中挑一个函数作为判别器D.
而Improved WGAN则是这样:
$W(P_{data},P_G)=\max_D\{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[D(x)]−\lambda\int_xmax(0,||\nabla xD(x)||−1)dx\}$
也就是说, 现在我们寻找判别器的函数集不再是1-Lipschitz中的函数了, 而是任意函数. 但是后面增加了一项惩罚项. 这个惩罚项就能够让选中的判别器函数倾向于是一个“对输入梯度为1的函数”. 这样也能实现类似weight clipping的效果. 但与之前遇到的问题一样, 积分在实现的时候无法计算, 所以我们用采样的方法去计算这个惩罚项, 即:
$W(P_{data},P_G)=\max_D\{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[D(x)]−\lambda\mathbb{E}_{x\sim P_{penalty}}[max(0,||\nabla xD(x)||−1)]\}$
也就是说, 在训练过程中, 我们更倾向于得到一个判别器D, 它能对从 $P_penalty$中采样得到的每一个 $x$都能 $||\nabla xD(x)||\le1$,. 涉及到采样, 那就要关系到如何采样. 而首先是从哪里采? 即 $P_penalty$是什么?
Improved WGAN设计了一个特别的 $P_penalty$. 它的产生过程如下:

1. 从Pdata中采样一个点
2. 从PG中采样一个点
3. 将这两个点连线
4. 在连线之上在采样得到一个点，就是一个从 $P_penalty$采样的一个点.

重复上面的过程就能不断采样得到 $x\sim P_{penalty}$。最终得到下图中的蓝色区域就可以看作是 $P_{penalty}$:

也就是说, 我们采样的范围不是整个 $x$, 只是P_G和P_{data}中间的空间中的一部分.

再更进一步, Improved WGAN真正做的事是这样:
$W(P_{data},P_G)=\max_D\{\mathbb{E}_{x\sim P_{data}}[D(x)]−\mathbb{E}_{x\sim P_G}[D(x)]−\lambda\mathbb{E}_{x\sim P_{penalty}}[(||\nabla xD(x)||−1)^2]\}$
这个惩罚项的目的是让梯度尽可能趋向于等于1. 即当梯度大于1或小于1时都会受到惩罚. 而原来的惩罚项仅仅在梯度大于1时受到惩罚而已. 这样做是有好处的, 就像我们在SVM中强调最大类间距离一样, 虽然有多个可以将数据区分开的分类面, 但我们希望找到不但能区分数据, 还能让区分距离最大的那个分类面. 这里这样做的目的是由于可能存在多个判别器, 我们想要找到的那个判别器应该有一个“最好的形状”. 一个“好”的判别器应该在 $P_{data}$附近是尽可能大, 要在 $P_G$附近尽可能小. 也就是说处于 $P_{data}$和 $P_G$之间的 $P_{penalty}$区域应该有一个比较“陡峭”的梯度. 但是这个陡峭程度是有限制的, 这个限制就是1.