View source for Wasserstein GAN

{{Short description|Generative adversarial network variant}}
{{Machine learning|Artificial neural network}}
The '''Wasserstein Generative Adversarial Network (WGAN)''' is a variant of [[generative adversarial network|generative adversarial network (GAN)]] proposed in 2017 that aims to "improve the stability of learning, get rid of problems like [[mode collapse]], and provide meaningful learning curves useful for debugging and hyperparameter searches".<ref name=":0">{{Cite journal |last1=Arjovsky |first1=Martin |last2=Chintala |first2=Soumith |last3=Bottou |first3=Léon |date=2017-07-17 |title=Wasserstein Generative Adversarial Networks |url=https://proceedings.mlr.press/v70/arjovsky17a.html |journal=International Conference on Machine Learning |language=en |publisher=PMLR |pages=214–223}}</ref><ref>{{cite arXiv |last=Weng |first=Lilian |date=2019-04-18 |title=From GAN to WGAN |class=cs.LG |eprint=1904.08994 }}</ref>

Compared with the original GAN discriminator, the Wasserstein GAN discriminator provides a better learning signal to the generator. This allows the training to be more stable when generator is learning distributions in very high dimensional spaces.

== Motivation ==

=== The GAN game ===
The original GAN method is based on the GAN game, a [[zero-sum game]] with 2 players: generator and discriminator. The game is defined over a [[probability space]] <math>(\Omega, \mathcal B, \mu_{ref})</math>, The generator's [[Strategy (game theory)|strategy set]] is the set of all probability measures <math>\mu_G</math> on <math>(\Omega, \mathcal B)</math>, and the discriminator's strategy set is the set of measurable functions <math>D: \Omega \to [0, 1]</math>.

The objective of the game is<math display="block">L(\mu_G, D) := \mathbb{E}_{x\sim \mu_{ref}}[\ln D(x)] + \mathbb{E}_{x\sim \mu_G}[\ln (1-D(x))].</math>
The generator aims to minimize it, and the discriminator aims to maximize it.

A [[Generative adversarial network#Main theorems for GAN game|basic theorem of the GAN game]] states that{{Math theorem
| math_statement = For any fixed generator strategy <math>\mu_G</math>, let the optimal reply be <math>D^* = \arg\max_{D} L(\mu_G, D)</math>, then

<math display="block">\begin{align}
D^*(x) &= \frac{d\mu_{ref}}{d(\mu_{ref} + \mu_G)}\\
L(\mu_G, D^*) &= 2D_{JS}(\mu_{ref}; \mu_G) - 2\ln 2,
\end{align}</math>

where the derivative is the [[Radon–Nikodym derivative]], and <math>D_{JS}</math> is the [[Jensen–Shannon divergence]].
| name = Theorem
| note = the optimal discriminator computes the Jensen–Shannon divergence
}}

Repeat the GAN game many times, each time with the generator moving first, and the discriminator moving second. Each time the generator <math>\mu_G</math> changes, the discriminator must adapt by approaching the ideal<math display="block">D^*(x) = \frac{d\mu_{ref}}{d(\mu_{ref} + \mu_G)}.</math>
Since we are really interested in <math>\mu_{ref}</math>, the discriminator function <math>D</math> is by itself rather uninteresting. It merely keeps track of the likelihood ratio between the generator distribution and the reference distribution. At equilibrium, the discriminator is just outputting <math>\frac 12</math> constantly, having given up trying to perceive any difference.{{NoteTag|note=In practice, the generator would never be able to reach perfect imitation, and so the discriminator would have motivation for perceiving the difference, which allows it to be used for other tasks, such as performing ImageNet classification [[unsupervised learning|without supervision]].|name=in practice}}

Concretely, in the GAN game, let us fix a generator <math>\mu_G</math>, and improve the discriminator step-by-step, with <math>\mu_{D, t}</math> being the discriminator at step <math>t</math>. Then we (ideally) have<math display="block">L(\mu_G, \mu_{D, 1}) \leq L(\mu_G, \mu_{D, 2}) \leq \cdots \leq \max_{\mu_D} L(\mu_G, \mu_D) = 2D_{JS}(\mu_{ref} \|  \mu_G) - 2\ln 2,</math>so we see that the discriminator is actually lower-bounding <math>D_{JS}(\mu_{ref} \| \mu_G)</math>.

=== Wasserstein distance ===
Thus, we see that the point of the discriminator is mainly as a critic to provide feedback for the generator, about "how far it is from perfection", where "far" is defined as Jensen–Shannon divergence.

Naturally, this brings the possibility of using a different criteria of farness. There are many possible [[Divergence (statistics)|divergences]] to choose from, such as the [[f-divergence]] family, which would give the f-GAN.<ref>{{Cite journal |last1=Nowozin |first1=Sebastian |last2=Cseke |first2=Botond |last3=Tomioka |first3=Ryota |date=2016 |title=f-GAN: Training Generative Neural Samplers using Variational Divergence Minimization |url=https://proceedings.neurips.cc/paper/2016/hash/cedebb6e872f539bef8c3f919874e9d7-Abstract.html |journal=Advances in Neural Information Processing Systems |publisher=Curran Associates, Inc. |volume=29|arxiv=1606.00709 }}</ref>

The Wasserstein GAN is obtained by using the [[Wasserstein metric]], which satisfies a "dual representation theorem" that renders it highly efficient to compute:

{{Math theorem
| math_statement = When the probability space <math>\Omega</math> is a metric space, then  
for any fixed <math>K > 0</math>, <math display="block">W_1(\mu, \nu) = \frac 1 K\sup_{\|f\|_L \leq K} \mathbb{E}_{x\sim \mu}[f(x)] -\mathbb E_{y\sim \nu}[f(y)]</math>
where <math>\|\cdot\|_L</math> is the [[Lipschitz norm]].
| name = Theorem
| note = Kantorovich-Rubenstein duality
}}

A proof can be found [[Wasserstein metric#Dual representation of W1|in the main page on Wasserstein metric]].

== Definition ==
By the Kantorovich-Rubenstein duality, the definition of Wasserstein GAN is clear:{{blockquote|A Wasserstein GAN game is defined by a [[probability space]] <math>(\Omega, \mathcal B, \mu_{ref})</math>, where <math>\Omega</math> is a [[metric space]], and a constant <math>K > 0</math>.

There are 2 players: generator and discriminator (also called "critic").

The generator's [[Strategy (game theory)|strategy set]] is the set of all probability measures <math>\mu_G</math> on <math>(\Omega, \mathcal B)</math>.

The discriminator's strategy set is the set of measurable functions of type <math>D: \Omega \to \R</math> with bounded Lipschitz-norm: <math>\|D\|_L \leq K</math>.

The Wasserstein GAN game is a [[zero-sum game]], with objective function<math display="block">L_{WGAN}(\mu_G, D) := \mathbb{E}_{x\sim \mu_G}[D(x)] -\mathbb E_{x\sim  \mu_{ref}}[D(x)].</math>
The generator goes first, and the discriminator goes second. The generator aims to minimize the objective, and the discriminator aims to maximize the objective:<math display="block">\min_{\mu_G} \max_{D} L_{WGAN}(\mu_G, D).</math>}}
By the Kantorovich-Rubenstein duality, for any generator strategy <math>\mu_G</math>, the optimal reply by the discriminator is <math>D^*</math>, such that <math display="block"> L_{WGAN}(\mu_G, D^*) = K \cdot W_1(\mu_G, \mu_{ref}).</math>Consequently, if the discriminator is good, the generator would be constantly pushed to minimize <math> W_1(\mu_G, \mu_{ref})</math>, and the optimal strategy for the generator is just <math> \mu_G = \mu_{ref}</math>, as it should.

== Comparison with GAN ==
In the Wasserstein GAN game, the discriminator provides a better gradient than in the GAN game.

Consider for example a game on the real line where both <math>\mu_G</math> and <math>\mu_{ref}</math> are Gaussian. Then the optimal Wasserstein critic <math>D_{WGAN}</math> and the optimal GAN discriminator <math>D</math> are plotted as below:
[[File:Wasserstein-GAN critic vs GAN discriminator.svg|thumb|The optimal Wasserstein critic <math>D_{WGAN}</math> and the optimal GAN discriminator <math>D</math> for a fixed reference distribution <math>\mu_{ref}</math> and generator distribution <math>\mu_G</math>. Both the Wasserstein critic <math>D_{WGAN}</math> and the GAN discriminator <math>D</math> are scaled down to fit the plot.]]
For fixed discriminator, the generator needs to minimize the following objectives:

* For GAN, <math>\mathbb E_{x\sim \mu_G} [\ln(1-D(x))]</math>.
* For Wasserstein GAN, <math>\mathbb E_{x\sim \mu_G} [D_{WGAN}(x)]</math>.

Let <math>\mu_G</math> be parametrized by <math>\theta</math>, then we can perform [[stochastic gradient descent]] by using two [[unbiased estimator]]s of the gradient:<math display="block">\nabla_{\theta} \mathbb E_{x\sim \mu_G} [\ln(1-D(x))] = \mathbb E_{x\sim \mu_G} [\ln(1-D(x))\cdot \nabla_{\theta} \ln\rho_{\mu_G}(x)]</math><math display="block">\nabla_{\theta} \mathbb E_{x\sim \mu_G} [D_{WGAN}(x)] = \mathbb E_{x\sim \mu_G} [D_{WGAN}(x)\cdot \nabla_{\theta} \ln\rho_{\mu_G}(x)]</math>where we used the [[reparameterization trick]].{{NoteTag|note=This is not how it is really done in practice, since 
<math>\nabla _{\theta }\ln \rho _{\mu _{G}}(x)</math>
 is in general intractable, but it is theoretically illuminating.|name=not really done in practice}}
[[File:Wasserstein-GAN critic vs log-minus GAN discriminator.svg|thumb|The same plot, but with the GAN discriminator <math>D</math> replaced by <math>\ln(1-D)</math> (and scaled down to fit the plot)]]
As shown, the generator in GAN is motivated to let its <math>\mu_G</math> "slide down the peak" of <math>\ln(1-D(x))</math>. Similarly for the generator in Wasserstein GAN.

For Wasserstein GAN, <math>D_{WGAN}</math> has gradient 1 almost everywhere, while for GAN, <math>\ln(1-D)</math> has flat gradient in the middle, and steep gradient elsewhere. As a result, the variance for the estimator in GAN is usually much larger than that in Wasserstein GAN. See also Figure 3 of.<ref name=":0" />

The problem with <math>D_{JS}</math> is much more severe in actual machine learning situations. Consider training a GAN to generate [[ImageNet]], a collection of photos of size 256-by-256. The space of all such photos is <math>\R^{256^2}</math>, and the distribution of ImageNet pictures, <math>\mu_{ref}</math>, concentrates on a manifold of much lower dimension in it. Consequently, any generator strategy <math>\mu_G</math> would almost surely be entirely disjoint from <math>\mu_{ref}</math>, making <math>D_{JS}(\mu_G \| \mu_{ref}) = +\infty</math>. Thus, a good discriminator can almost perfectly distinguish <math>\mu_{ref}</math> from <math>\mu_G</math>, as well as any <math>\mu_G'</math> close to <math>\mu_G</math>. Thus, the gradient <math>\nabla_{\mu_G} L(\mu_G, D) \approx 0</math>, creating no learning signal for the generator.

Detailed theorems can be found in.<ref>{{Cite journal |last1=Arjovsky |first1=Martin |last2=Bottou |first2=Léon |date=2017-01-01 |title=Towards Principled Methods for Training Generative Adversarial Networks |arxiv=1701.04862 |url=https://ui.adsabs.harvard.edu/abs/2017arXiv170104862A}}</ref>

== Training Wasserstein GANs ==
Training the generator in Wasserstein GAN is just [[gradient descent]], the same as in GAN (or most deep learning methods), but training the discriminator is different, as the discriminator is now restricted to have bounded Lipschitz norm. There are several methods for this.

=== Upper-bounding the Lipschitz norm ===

Let the discriminator function <math>D</math> to be implemented by a [[multilayer perceptron]]:<math display="block">D = D_n \circ D_{n-1} \circ \cdots \circ D_1</math>where <math>D_i(x) = h(W_i x)</math>, and <math>h:\R \to \R</math> is a fixed activation function with <math>\sup_x |h'(x)| \leq 1</math>. For example, the [[hyperbolic tangent function]] <math>h = \tanh</math> satisfies the requirement.

Then, for any <math>x</math>, let <math>x_i = (D_i \circ D_{i-1} \circ \cdots \circ D_1)(x)</math>, we have by the [[chain rule]]:<math display="block">d D(x) = diag(h'(W_n x_{n-1})) \cdot W_n \cdot diag(h'(W_{n-1} x_{n-2})) \cdot W_{n-1}  \cdots diag(h'(W_1 x)) \cdot W_1 \cdot dx</math>Thus, the Lipschitz norm of <math>D</math> is upper-bounded by<math display="block">\|D \|_L \leq \sup_{x}\| diag(h'(W_n x_{n-1})) \cdot W_n \cdot diag(h'(W_{n-1} x_{n-2})) \cdot W_{n-1}  \cdots diag(h'(W_1 x)) \cdot W_1\|_F</math>where <math>\|\cdot\|_s</math> is the [[operator norm]] of the matrix, that is, the largest [[singular value]] of the matrix, that is, the [[spectral radius]] of the matrix (these concepts are the same for matrices, but different for general [[linear operators]]).

Since <math>\sup_x |h'(x)| \leq 1</math>, we have <math>\|diag(h'(W_i x_{i-1}))\|_s = \max_j |h'(W_i x_{i-1, j})| \leq 1</math>, and consequently the upper bound:<math display="block">\|D \|_L \leq \prod_{i=1}^n \|W_i \|_s</math>Thus, if we can upper-bound operator norms <math>\|W_i\|_s</math> of each matrix, we can upper-bound the Lipschitz norm of <math>D</math>.

=== Weight clipping ===
Since for any <math>m\times l</math> matrix <math>W</math>, let <math>c = \max_{i, j} |W_{i, j}|</math>, we have<math display="block">\|W\|_s^2 = \sup_{\|x\|_2=1}\|W x\|_2^2 = \sup_{\|x\|_2=1}\sum_{i}\left(\sum_j W_{i, j} x_j\right)^2 = \sup_{\|x\|_2=1}\sum_{i, j, k}W_{ij}W_{ik}x_jx_k \leq c^2 ml^2</math>by clipping all entries of <math>W</math> to within some interval <math>[-c, c]</math>, we have can bound <math>\|W\|_s</math>.

This is the weight clipping method, proposed by the original paper.<ref name=":0" />

=== Spectral normalization ===

The spectral radius can be efficiently computed by the following algorithm:{{blockquote|

'''INPUT''' matrix <math>W</math> and initial guess <math>x</math>

Iterate <math>x \mapsto \frac{1}{\|Wx\|_2}Wx</math> to convergence <math>x^*</math>. This is the eigenvector of <math>W</math> with eigenvalue <math>\|W\|_s</math>.

'''RETURN''' <math>x^*, \|Wx^*\|_2</math>}}

By reassigning <math>W_i \leftarrow \frac{W_i}{\|W_i\|_s}</math> after each update of the discriminator, we can upper bound <math>\|W_i\|_s \leq 1</math>, and thus upper bound <math>\|D \|_L</math>.

The algorithm can be further accelerated by [[memoization]]: At step <math>t</math>, store <math>x^*_i(t)</math>. Then at step <math>t+1</math>, use <math>x^*_i(t)</math> as the initial guess for the algorithm. Since <math>W_i(t+1)</math> is very close to <math>W_i(t)</math>, so is <math>x^*_i(t)</math> close to <math>x^*_i(t+1)</math>, so this allows rapid convergence.

This is the spectral normalization method.<ref>{{cite arXiv |last1=Miyato |first1=Takeru |last2=Kataoka |first2=Toshiki |last3=Koyama |first3=Masanori |last4=Yoshida |first4=Yuichi |date=2018-02-16 |title=Spectral Normalization for Generative Adversarial Networks |class=cs.LG |eprint=1802.05957 }}</ref>

=== Gradient penalty ===

Instead of strictly bounding <math>\|D\|_L</math>, we can simply add a "gradient penalty" term for the discriminator, of form<math display="block">\mathbb{E}_{x\sim\hat\mu}[(\|\nabla D(x)\|_2 - a)^2]</math>where <math>\hat \mu</math> is a fixed distribution used to estimate how much the discriminator has violated the Lipschitz norm requirement.
The discriminator, in attempting to minimize the new loss function, would naturally bring <math>\nabla D(x)</math> close to <math>a</math> everywhere, thus making <math>\|D\|_L \approx a</math>.

This is the gradient penalty method.<ref>{{Cite journal |last1=Gulrajani |first1=Ishaan |last2=Ahmed |first2=Faruk |last3=Arjovsky |first3=Martin |last4=Dumoulin |first4=Vincent |last5=Courville |first5=Aaron C |date=2017 |title=Improved Training of Wasserstein GANs |url=https://proceedings.neurips.cc/paper/2017/hash/892c3b1c6dccd52936e27cbd0ff683d6-Abstract.html |journal=Advances in Neural Information Processing Systems |publisher=Curran Associates, Inc. |volume=30}}</ref>

== Further reading ==

* [https://lilianweng.github.io/posts/2017-08-20-gan/ From GAN to WGAN]
* [https://vincentherrmann.github.io/blog/wasserstein/ Wasserstein GAN and the Kantorovich-Rubinstein Duality]
* [https://www.depthfirstlearning.com/2019/WassersteinGAN Depth First Learning: Wasserstein GAN]

== See also ==

* [[Generative adversarial network]]
* [[Wasserstein metric]]
* [[Earth mover's distance]]
* [[Transportation theory (mathematics)|Transportation theory]]

== References ==
{{Reflist}}

== Notes ==
{{reflist|group=note}}

{{Differentiable computing}}

[[Category:Neural network architectures]]
[[Category:Cognitive science]]
[[Category:Unsupervised learning]]