本文海量参考 郑华滨:令人拍案叫绝的Wasserstein GAN,部分地方加入了自己的理解与解释,但是还有一些细节未完善,原文非常推荐。

原始 GAN 的问题

原始 GAN 的目标是: $$ \min_{G} \max_{D} V(D, G)= \mathbb{E}_{\boldsymbol{x} \sim p_{\text{data}}}[\log D(\boldsymbol{x})] + \mathbb{E}_{\boldsymbol{z} \sim p_{\boldsymbol{z}}} [\log (1-D(G(\boldsymbol{z})))] $$ 其中,生成器 $G$ 的损失项: $$ \begin{aligned} &\text{(1):} &\mathbb{E}_{\boldsymbol{x} \sim p_g} [\log (1-D(\boldsymbol{x}))] \\ \end{aligned} $$ 后来作者又提出了改进的损失项以增大梯度信息: $$ \begin{aligned} &\text{(2):} &\mathbb{E}_{\boldsymbol{x} \sim p_g} [- \log (D(\boldsymbol{x}))] \end{aligned} $$ 其中二者都存在着一些问题,以下一一解释。

log(1-D(x)) 的问题

判别器越好,生成器梯度消失越严重。

对于第一种形式的 $G$ 损失,GAN 的训练目标相当于 在最优判别器 $D$ 下最小化 $p_\text{data}$ 与 $p_g$ 的 JS 散度

我们来研究一下 JS 散度: $$ \begin{aligned} D_{\text{JS}}(P||Q) &= \frac{1}{2}D_{\text{KL}}(P||\frac{P+Q}{2}) + \frac{1}{2}D_{\text{KL}}(Q||\frac{P+Q}{2}) \\ D_{\text{KL}}(P||Q) &= \int_xp(x)(\log\frac{p(x)}{q(x)}) \end{aligned} $$ 对于任意一个变量 $x$,在两个分布 $P$、$Q$ 上只有 4 种可能:

  1. $P(x) = 0$ 且 $Q(x) = 0$,对 JS 散度的贡献为 0;
  2. $P(x) = 0$ 且 $Q(x) \ne 0$,对 JS 散度的贡献为式子第二项 $\frac{1}{2}D_{\text{KL}}(Q||\frac{0+Q}{2}) = \log 2$;
  3. $P(x) \ne 0$ 且 $Q(x) = 0$,对 JS 散度的贡献为式子第一项 $\frac{1}{2}D_{\text{KL}}(P||\frac{P+0}{2}) = \log 2$;
  4. $P(x) \ne 0$ 且 $Q(x) \ne 0$,两个分布有重叠,是 JS 散度的主要贡献。

因此,如果两个分布 $P$ 跟 $Q$ 没有重叠会怎样?会导致它们的 JS 散度 $\operatorname{JS} (P||Q) = \log 2$,此时的梯度恒等于 0

那么在 GAN 的场景中,数据分布 $p_\text{data}$ 与 $p_g$ 没有重叠的概率其实是非常大的。在实践中,生成器 $G$ 往往从一个低维分布 $p_\boldsymbol{z}$(例如 100 维)中采样随机噪声 $\boldsymbol{z}$ 作为输入,再经过深度神经网络生成高维样本(例如 $64 \times 64 = 4096$ 维),这意味着生成数据的分布 $p_g$ 本质上是 4096 维的高维空间中的一个 100 维的低维 流形(manifold,例如三维空间中的一个曲面是一个二维流形,因为它只有两个方向的自由度),然而真实数据分布 $p_\text{data}$ 却是实实在在分布在这 4096 高维空间的,此时 $p_g$ 与 $p_\text{data}$ 的重叠部分的 测度(measure,可以理解为「超体积」)无限接近于 0,从而 JS 散度为常数,生成器 $G$ 的梯度为 0,梯度消失

事实上,我们也可以直观上从高维空间角度进行理解。$p_\text{data}$ 是 4096 维空间分布,$p_g$ 是该空间内的 100 维流形,判别器 $D$ 很容易彻底将它们分开,对几乎所有真实样本给出常数概率 1,几乎所有生成样本给出常数概率 0,常数求导得到 $G$ 的梯度为 0,梯度消失

总结一下:

  • 判别器训练得太好,生成器梯度消失,训练不了;
  • 判别器训练得不好,生成器梯度不准,四处乱跑;
  • 我太难了。

-log(D(x)) 的问题

等价于最小化一个不合理的距离衡量,导致梯度不稳定、Mode Collapse。

使用 $- \log (D(\boldsymbol{x}))$ 代替 $\log (1-D(\boldsymbol{x}))$ 之后,经过一系列等价变换(见参考资料),GAN 的最小化目标等价于: $$ \min_G \operatorname{KL}(p_g || p_\text{data}) - 2 \operatorname{JS}(p_\text{data} || p_g) $$ 这个目标存在两个严重问题:

问题一

对于两个分布 $p_g$ 与 $p_\text{data}$,这个目标同时最小化它们的 KL 散度、最大化它们的 JS 散度,而这是两个相反的目标,将会导致梯度很不稳定。

问题二:

KL 散度是不对称的,以 $\operatorname{KL}(p_g || p_\text{data})$ 为例:

  • 当 $p_g(\boldsymbol{x}) \rightarrow 0$ 而 $p_\text{data}(\boldsymbol{x}) \rightarrow 1$ 时(生成器没能生成真实样本;缺乏多样性),$p_g\log\frac {p_g(\boldsymbol{x})} {p_\text{data}(\boldsymbol{x})} \rightarrow 0$,对 KL 散度贡献趋近 0;
  • 当 $p_g(\boldsymbol{x}) \rightarrow 1$ 而 $p_\text{data}(\boldsymbol{x}) \rightarrow 0$ 时(生成器生成了不真实样本,缺乏准确性),$p_g\log\frac {p_g(\boldsymbol{x})} {p_\text{data}(\boldsymbol{x})} \rightarrow +\infty$,对 KL 散度贡献趋近于正无穷。

以上两者都是我们应该惩罚的,如果我们用该损失函数,那么前者的惩罚极其微小,后者的惩罚巨大。这样导致的后果就是,生成器倾向于生成准确但没有多样性的样本,这就是我们常说的 Mode Collapse(模式坍缩)。

总结一下:

  • 优化两个矛盾的目标,梯度不稳定;
  • 容易造成 Mode Collapse,生成样本多样性差。

WGAN

WGAN 正是为了解决以上提到的原始 GAN 的诸多问题,其核心在于引入了 Wasserstein 距离并且成功解决了一些求解问题,理论部分比较复杂(原作者不愧是数学出身),但是实现起来却不难。下面先按照理论依次解释,最后给出具体做法。

Wasserstein 距离

Wasserstein 距离 又叫 推土机距离(Earch-Mover's Distance),定义如下: $$ W(p_\text{data}, p_g) = \inf_{\gamma\sim\Pi(p_\text{data},p_g)} \mathbb{E}_{(\boldsymbol{x},\boldsymbol{y})_\sim\gamma} [||\boldsymbol{x} - \boldsymbol{y}||] $$

其中 $\Pi(p_\text{data},p_g)$ 是 $p_\text{data}$ 与 $p_g$ 组合起来的所有可能的联合分布的集合。对于每一个可能的联合分布 $\gamma$,我们可以从中采样 $(x, y) \sim \gamma$ 得到真实样本 $x$、生成样本 $y$,并算出这对样本之间的距离 $||x-y||$,从而可以计算该联合分布下样本对之间的距离的期望 $\mathbb{E}_{(x,y)_\sim\gamma} [||x-y||]$。在所有可能的联合分布中能够对该期望值取到的下界,就是 Wasserstein 距离

之所以被称为 推土机距离,是因为该距离在直观上可以理解为将一堆土 $p_\text{data}$ 推到 $p_g$ 的位置所要花费的最小代价,其中的 $\gamma$ 就是推土的方案。

与 KL 散度、JS 散度相比,即使两个分布之间没有重叠,Wasserstein 距离依然能够很好地反应它们之间的差异,并且该差异是平滑的。反应到基于梯度下降法的训练中,Wasserstein 距离能够在 KL、JS 散度无法提供梯度的时候提供有意义的、平滑的梯度

Wasserstein GAN

如何能够将 Wasserstein 距离应用到 GAN 的训练中,首先要解决的就是公式中的 $\inf_{\gamma\sim\Pi(p_\text{data},p_g)}$ 项,这里没办法直接求解,因此作者基于 Kantorovich-Rubinstein Duality 将其转换为如下形式(证明参见原论文): $$ W(p_\text{data}, p_g) = \frac{1}{K} \sup_{||f||_l\le K} \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f(\boldsymbol{x})] $$

其中 $||f||_L$ 表示符合 Lipschitz 连续(Lipschitz Continuity)的连续函数 $f$ 的 Lipschitz 常数。Lipschitz 连续是一个函数光滑性的条件,它限制了函数改变的速度(即斜率、导数、梯度)不大于某个常数 $K$(该常数视具体函数而定),即: $$ |f(x_1) - f(x_2)| \le K |x_1 - x_2| $$ 因此上式的意思就是要求函数 $f$ 在 Lipschitz 常数 $||f||_L \le K$ 的条件下,对所有可能满足条件的 $f$ 取到 $\mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f(\boldsymbol{x})]$ 的上界,最后除以 $K$。

特别地,函数 $f$ 可以用一组参数 $w$ 来定义(即 $f_w$),因此上式又被近似为求解以下形式: $$ W(p_\text{data}, p_g) \approx \frac{1}{K} \max_{w:|f_w|_L \le K} \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f_w(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})] $$ 在深度学习中,我们可以把用带参数 $w$ 的神经网络来表示 $f_w$,由于神经网络的拟合能力足够强大,因此这样定义出来的 $f_w$ 足以近似上式要求的 $\sup_{||f||_l\le K}$。

最后,我们要记得还有一个 Lipschitz 连续的约束,即 $|f_w|_L \le K$。因为我们并不关心具体的常数 $K$ 是多少,它造成的效果只是改变了梯度的总体大小,而不会影响梯度的方向,因此作者提出只需要限制神经网络 $f_w$ 的所有参数 $w_i$ 不超过某个范围 $[-c, c]$,此时关于输入样本 $\boldsymbol{x}$ 的导数 $\frac{\partial f_w}{\partial \boldsymbol{x}}$ 也不会超过此范围,因此肯定存在这样的常数 $K$ 使得 $f_w$ 满足 Lipschitz 连续。

综上,我们构造一个判别器网络 $f_w$,在限制其参数 $w$ 不超过某个范围的条件下,最大化 $$ L = \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f_w(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})] $$ 此时,$L$ 就近似真实分布 $p_\text{data}$ 与生成分布 $p_g$ 之间的 Wasserstein 距离,当 $L$ 越小时,表示真实分布 $p_\text{data}$ 与生成分布 $p_g$ 的距离越小,GAN 训练得越好。我们的优化目标为: $$ \min_{G} \max_{D} L = \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f_w(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})] $$ 总结一下此时的损失函数:

  • 判别器 $D$:$- \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f_w(\boldsymbol{x})] + \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})]$;
  • 生成器 $G$:$- \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})]$。

从原始 GAN 到 WGAN

从原始 GAN 到 WGAN 的 4 点改进

由此,我们对比一下原始 GAN 价值函数与 Wasserstein 距离的优化目标: $$ \begin{aligned} \min_{G} \max_{D} V(D, G) &= \mathbb{E}_{\boldsymbol{x} \sim p_\text{data}}[\log D(\boldsymbol{x})] + \mathbb{E}_{\boldsymbol{x} \sim p_g} [\log (1-D(\boldsymbol{x}))] \\ \min_{G} \max_{D} L &= \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f_w(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})] \end{aligned} $$ 这里要注意,原始判别器 $D$ 是分类任务,基于 Wasserstein 距离的判别器 $D$ 则是回归任务。因此我们做如下改动:

  1. 判别器 $D$ 的最后一层去掉 Sigmoid 激活函数;
  2. 判别器 $D$ 和生成器 $G$ 的损失都不取 $\log$;
  3. 每次更新判别器 $D$ 之后把每个参数的绝对值截断到不超过固定常数 $。c$

除了以上三点是从理论中得到的,作者还提出从实践中得到的第 4 点改进:

  1. 不使用基于动量的优化算法(如 Momentum、Adam),推荐 RMSProp 或 SGD。

作者发现如果使用Adam,判别器的 loss 有时候会崩掉,当它崩掉时,Adam给出的更新方向与梯度方向夹角的cos值就变成负数,更新方向与梯度方向南辕北辙,这意味着判别器的loss梯度是不稳定的,所以不适合用Adam这类基于动量的优化算法。作者改用RMSProp之后,问题就解决了,因为RMSProp适合梯度不稳定的情况。

WGAN-GP

WGAN 为了保证 Lipschitz 连续,采用了 weight clipping 的方法,但是这样会导致模型的建模能力弱化、梯度爆炸或消失等问题,最终也会产生低质量样本、难以收敛等。

因此 WGAN-GP 提出为 WGAN 的损失函数加入 梯度惩罚(Gradient Penalty,GP)。具体来说,原始 WGAN 的目标 $L$ 如下: $$ L = \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f_w(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})] $$ 而 WGAN-GP 的目标如下: $$ L = \underbrace{ \mathbb{E}_{\boldsymbol{x}\sim p_\text{data}}[f_w(\boldsymbol{x})] - \mathbb{E}_{\boldsymbol{x}\sim p_g }[f_w(\boldsymbol{x})] }_{Original} + \underbrace{ \lambda \mathbb{E}_{\hat{\boldsymbol{x}}\sim p_\hat{\boldsymbol{x}} } [(||\nabla_\hat{\boldsymbol{x}} D(\hat{\boldsymbol{x}})||_2 - 1)^2] }_{Gradient~Penalty} $$ 其 GP 的逻辑是:当且仅当一个可微函数的梯度范数(Gradient Norm)在任意处都不超过 1 时,该函数满足 1-Lipschitz 条件。公式中用于计算 GP 的样本 $\hat{\boldsymbol{x}}$ 是真实样本 $p_\text{data}$ 与生成样本 $p_g$ 的线性插值。

参考