Wasserstein GAN (WGAN)

这是 Wasserstein GAN 的实现。

最初的GAN损耗基于实际分布${\mathbb{P}}_r$和生成的分布之间的Jensen-Shannon（JS）差异${\mathbb{P}}_g$。Wasserstein GAN 基于这些分布之间的 Earth Mover 距离。

$$W({\mathbb{P}}_r,{\mathbb{P}}_g)=\underset{\gamma \in \Pi ({\mathbb{P}}_r,{\mathbb{P}}_g)}{\inf }{\mathbb{E}}_{(x,y)\sim \gamma }\Vert x-y\Vert$$

$\Pi ({\mathbb{P}}_r,{\mathbb{P}}_g)$是所有联合分布的集合，其边际概率为$\gamma (x,y)$。

${\mathbb{E}}_{(x,y)\sim \gamma }\Vert x-y\Vert$是给定关节分布的地球移动器距离（$x$$y$也是概率）。

因此，等$W({\mathbb{P}}_r,{\mathbb{P}}_g)$于实际分布${\mathbb{P}}_r$和生成的分布之间任何关节分布的最小地球移动器距离${\mathbb{P}}_g$。

本文表明，Jensen-Shannon（JS）背离和其他衡量两个概率分布之间差异的度量并不平滑。因此，如果我们对其中一个概率分布（参数化）进行梯度下降，它将不会收敛。

基于坎托罗维奇-鲁宾斯坦二元性，$$W({\mathbb{P}}_r,{\mathbb{P}}_g)=\underset{\Vert f{\Vert }_L\le 1}{\sup }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_g}[f(x)]$$

所有$\Vert f{\Vert }_L\le 1$的 1-Lipschitz 函数都在哪里。

也就是说，它等于所有 1-Lipschitz 函数$${\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_g}[f(x)]$$之间的最大差异。

对于$K$-Lipschitz 函数，$$W({\mathbb{P}}_r,{\mathbb{P}}_g)=\underset{\Vert f{\Vert }_L\le K}{\sup }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[\frac{1}{K}f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_g}[\frac{1}{K}f(x)]$$

如果所有$K$-Lipschitz 函数都可以表示$f$为参数化了$f_w$哪里$w\in \mathcal{W}$，

$$K\cdot W({\mathbb{P}}_r,{\mathbb{P}}_g)=\underset{w\in \mathcal{W}}{\max }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f_w(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_g}[f_w(x)]$$

如果$({\mathbb{P}}_g)$由生成器表示$$g_{\theta }(z)$$$z$并且来自已知分布$z\sim p(z)$，

$$K\cdot W({\mathbb{P}}_r,{\mathbb{P}}_{\theta })=\underset{w\in \mathcal{W}}{\max }{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f_w(x)]-{\mathbb{E}}_{z\sim p(z)}[f_w(g_{\theta }(z))]$$

现在为了收敛$g_{\theta }$，${\mathbb{P}}_r$我们可以通过梯度下降$\theta$来最小化上述公式。

同样，我们可以${\max }_{w\in \mathcal{W}}$通过上升来找到$w$，同时保持$K$界限。保持$K$界限的一种方法是裁剪神经网络中定义范围内的$f$裁剪的所有权重。

以下是在一个简单的 MNIST 生成实验中尝试此操作的代码。

87import torch.utils.data
88from torch.nn import functional as F
89
90from labml_helpers.module import Module

鉴别器丢失

我们想找到最大化$w$$${\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f_w(x)]-{\mathbb{E}}_{z\sim p(z)}[f_w(g_{\theta }(z))]$$，所以我们最小化，$$-\frac{1}{m}\sum _{i=1}^m{f}_w\left(x^{(i)}\right)+\frac{1}{m}\sum _{i=1}^m{f}_w(g_{\theta }(z^{(i)}))$$

93class DiscriminatorLoss(Module):

• f_real 是$f_w(x)$

f_fake 是$f_w(g_{\theta }(z))$

这将返回带有 and 亏损的 a 元组$f_w(x)$$f_w(g_{\theta }(z))$，稍后会添加这些元组。它们分开存放以进行日志记录。

104    def forward(self, f_real: torch.Tensor, f_fake: torch.Tensor):

发电机损失

我们$\theta$想找到最小化$${\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f_w(x)]-{\mathbb{E}}_{z\sim p(z)}[f_w(g_{\theta }(z))]$$第一个组件是独立的$\theta$，所以我们最小化，$$-\frac{1}{m}\sum _{i=1}^m{f}_w(g_{\theta }(z^{(i)}))$$

118class GeneratorLoss(Module):

• f_fake 是$f_w(g_{\theta }(z))$

130    def forward(self, f_fake: torch.Tensor):