Week 2. A General Framework for GANs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# f-GAN: Training Generative Neural Samplers using Variational Divergence Minimization\n",
    "\n",
    "In the [last note](Week%201.%20Generative%20Adversarial%20Networks.ipynb), we saw how to learn a distribution $p_g$ that approximated some unknown distribution $p_d$. We did this by introducing two neural networks. Namely, the generator and the discriminator, which were both neural networks. The discriminator and the generator played a game where the discriminator were to distinguish between real data sampled from $p_d$ and fake data sampled from the generator. The generator, on the other hand, were to fool the discriminator by transforming random noise $z_i \\sim \\mathcal{N}(0, 1)$ into something indistinguishable from the real data.\n",
    "\n",
    "In essence, fooling the discriminator corresponds to minimizing the \"distance\" between $p_d$ and $p_g$ for some notion of distance. We presented the [Kullbach-Liebler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence) as one such measure and argued that we were not able to calculate this measure for several reasons. The divergence is defined as\n",
    "\n",
    "$$\n",
    "D_{KL}(P || Q) = \\int_\\mathcal{X} p(x) \\log \\frac{p(x)}{q(x)} dx \\quad \\quad (1)\n",
    "$$\n",
    "\n",
    "over the domain $\\mathcal{X}$ for continuous probability density functions $p$ and $q$. \n",
    "\n",
    "Since we were unable to optimize this directly, we turned our attention towards Generative Adversarial Networks. They solved the problem by introducting a discriminator which is trained jointly with the generator by playing the following game: \n",
    "\n",
    "$$\n",
    "\\text{min }_g\\text{ max}_d\\; \\mathbb{E}_{x\\sim p_d}[\\text{log d(x)}] + \\mathbb{E}_{z \\sim p_z}[1-\\text{log(d(g(z))}]. \\quad \\quad (2)\n",
    "$$\n",
    "\n",
    "Given a fixed generator $G$ [[3]](#References) showed that the optimal\n",
    "discriminator $d^\\ast_g(x)$ is:\n",
    "\n",
    "$$\n",
    "d^\\ast_g(x) = \\frac{p_{d}(x)}{p_{d}(x) + p_g(x)}. \\quad \\quad (3)\n",
    "$$\n",
    "\n",
    "The generator $g$ that plays best against its $d^\\ast_g(x)$ is the generator so $p_g=p_d$. Furthermore, it turned out that playing this game approximately minimizes the [Jensen-Shannon divergence](https://en.wikipedia.org/wiki/Jensen%E2%80%93Shannon_divergence)\n",
    "\n",
    "$$\n",
    "D_{JS}(P || Q) = D_{KL}(P ||P + Q) + D_{KL}(Q || P + Q). \\quad \\quad (4)\n",
    "$$\n",
    "\n",
    "The goal of this note is to generalize Eqn. (4) to a broader class of divergences. Conceptually, this will allow us to \"change the cost function\".  <!--(which is essentially Eqn. (2)). It will also generalize Eqn. (3) such that the optimal discriminator will get different forms according to which divergence the GAN setup is minimizing. -->\n",
    "\n",
    "Figure 1 contains a simple motivation for why we may want to change the objective of the training. The only difference to the figure in the previous note, is that the real data is now generated from a biomodal distribution. The figure shows how the learned distribution $p_g$ (which is now underspecified since it can only represent \"simple\" uni-modal distribution) may learn to simply span the two modes of the data. It may be that we would rather want the generater to just focus on one of the modes of the true distribution. In that case we could hope that a change in traning objective would favour such a property.\n",
    "\n",
    "<div style=\"text-align: center;\">\n",
    "    <img src=\"gifs/dual-mode/gif.gif\">\n",
    "    <strong>Figure 1</strong> Animation of how $p_g$ converges towards $p_d$.\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generalizing GANs to f-divergences\n",
    "\n",
    "As stated above, the goal of this note is to generalize GANs to minimize a broader class of divergences. Specifically, we will show that we can generalize GANs to minimize any [f-divergence](https://en.wikipedia.org/wiki/F-divergence). This results was shown in [[1]](#References) and will be reproduced here. In order to do so, we need three ingredients. \n",
    "\n",
    "1. Not surprisingly, the first ingredient is the f-divergence. \n",
    "2. The second ingredient is the [convex conjugate](https://en.wikipedia.org/wiki/Convex_conjugate) of convex lower semi-continuous functions.\n",
    "3. [Jensens Inequality](https://en.wikipedia.org/wiki/Jensen%27s_inequality) is the third and final ingredient.\n",
    "\n",
    "The structure of the following sections is as follows. First, we describe what the three ingredients are and how we will use them. Second, we will use the ingredients to show a lower bound on the f-divergence, which is very close to the GAN objective in (2). Finally, we will discuss how to use this lower bound to minimize different f-divergences.  \n",
    "\n",
    "### f-divergences\n",
    "Until now, we have seen two divergences, the KL and the JS-divergence (see Eqn. (1) and (4) respectively). It turns out that these divergences are both part of a broughter class of divergences called the f-divergences. f-divergences are characterized by a convex and continuous function $f$ where $f(1) = 0$. The f-divergence is defined as follows\n",
    "\n",
    "$$\n",
    "D_f(P||Q) = \\int q(x) f\\left( \\frac{p(x)}{q(x)} \\right) \\quad \\quad (5)\n",
    "$$\n",
    "\n",
    "Let's verify that the KL-divergence is actually a part of this family. Let $f(u) = u \\log u$ and pug it into Eqn. (5):\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "D_f(P || Q) &= \\int q(x) \\left( \\frac{p(x)}{q(x)} \\log \\frac{p(x)}{q(x)} \\right)\\\\\n",
    "&= \\int p(x) \\log \\frac{p(x)}{q(x)} \\\\\n",
    "&= D_{KL}(P || Q)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Similarly, if we let $f(x) = u \\log u + (u + 1) \\log (u + 1)$, we recover the JS-divergence (up to a constant). For more divergences, please see Table 1 further down the note.\n",
    "\n",
    "### The Convex Conjugate\n",
    "The convex conjugate is defined over convex lower semi-continuous functions, but in the f-GAN setting, you can just think of it as convex and continuous functions. For any function $f$ in the set of all these convex and continuous functions, the _convex conjugate_ $f^\\ast$ is defined as\n",
    "\n",
    "$$\n",
    "f^\\ast(t) = \\sup_{u \\in \\text{dom}_f} \\{ut - f(u)\\}.\n",
    "$$\n",
    "\n",
    "In the definition of $f^\\ast$, we use the [supremum](https://en.wikipedia.org/wiki/Infimum_and_supremum). Intuitively, this can be thought of as the maximum. $f^\\ast$ has two nice properties. \n",
    "\n",
    "1. $f^\\ast$ it self is convex and continuous, so it has a convex conjugate $f^{\\ast\\ast}$  \n",
    "2. $f$ and $f^\\ast$ are each others dual, which means that $f = f^{\\ast\\ast}$\n",
    "\n",
    "Note how 2. allows us to reformulate $f$ in terms of it convex conjugate: \n",
    "\n",
    "$$\n",
    "f(x) = \\sup_{t \\in \\text{dom}_{f^\\ast}} \\left( xt - f^\\ast(t)\\right). \\quad \\quad (6)\n",
    "$$\n",
    "\n",
    "We will use these properties to derive our lower bound on all $f$-divergences. \n",
    "\n",
    "### Jensen's Inequality\n",
    "This inequality is simple but powerful. It says that for any convex function $f$, the following inequality holds:\n",
    "\n",
    "$$\n",
    "\\mathbb{E} [ f(x) ] \\geq f( \\mathbb{E} x )\n",
    "$$\n",
    "\n",
    "We will use this inequality to lower bound f-divergences by using the fact the $f^\\ast$ is convex.\n",
    "\n",
    "### Lower bounding f-divergences\n",
    "By using the above ingredients we construct a lowerbound for any $f$-divergences. Equation (6) allows us to rewrite any $f$-divergence in terms of its convex conjungate: \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "D_f (P||Q) &= \\int_\\mathcal{X} q(x) f\\left( \\frac{p(x)}{q(x)} \\right) dx\\\\\n",
    "&= \\int_\\mathcal{X} q(x) \\sup_{t \\in \\text{dom}_{f^\\ast}} \\left\\{ \\frac{p(x)}{q(x)} t - f^\\ast(t)\\right\\} dx\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "If we write this as an expectation over $x\\sim q$ and recall that $f^\\ast$ is convex we can lowerbound our expression with Jensens Inequality: \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathbb{E}_{x \\sim q} \\left[ \\sup_{t \\in \\text{dom}_{f\\ast}} \\left\\{ \\frac{p(x)}{q(x)} t - f^\\ast(t)\\right\\} \\right]\n",
    "\\geq \\sup_{t \\in \\text{dom}_{f^\\ast}} \\left\\{ \\mathbb{E}_{x \\sim q} \\left[ \\frac{p(x)}{q(x)} t - f^\\ast(t)\\right] \\right\\} \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Simplifying the equations we finally arrive at: \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\sup_{t \\in \\text{dom}_{f^\\ast}} \\left\\{ \\int_\\mathcal{X} q(x) \\frac{p(x)}{q(x)} t dx -\\mathbb{E}_{x \\sim q} \\left[ f^\\ast(t)\\right] \\right\\}\n",
    "&= \\sup_{t \\in {\\text{dom}}_{f^\\ast}} \\left\\{ \\mathbb{E}_{x \\sim p} \\left[t\\right] -\\mathbb{E}_{x \\sim q} \\left[ f^\\ast(t)\\right] \\right\\}\\\\\n",
    "&\\geq \\sup_{D \\in \\mathcal{D}} \\left\\{ \\mathbb{E}_{x \\sim p} \\left[D(x)\\right] -\\mathbb{E}_{x \\sim q} \\left[ f^\\ast(D(x))\\right] \\right\\} \\quad \\quad (7)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Where $\\mathcal{D}$ is any class of functions. <!--The first step utilizes the dual formulation of $f$ (Eqn. (6)). Second step uses the definition of expectations. Then, the inequality is introduced by Jensens Inequality. Afterwards, we realize that the $q(x)$s cancels and finally, we use the definition of expectations again.\n",
    "<span style=\"background: yellow;\">Please note that it is not clear from the article how to do the last step.(I don't know if I want to include this)</span>-->\n",
    "\n",
    "Everything above is bit on the technical side. The take-away is that we now have a lower bound on the f-divergences that looks a lot like the original GAN setup in Eqn. (2) if we let $D(x)$ be our discriminator and think of $q$ as our generator.\n",
    "\n",
    "As a side note, it turns out that under certain mild conditions [[2]](#References), the bound above is tight for:\n",
    "\n",
    "$$\n",
    "D^*(x) = f'\\left( \\frac{p(x)}{q(x)} \\right). \n",
    "$$\n",
    "\n",
    "This means that if $D(x) = f'\\left(\\frac{p(x)}{q(x)}\\right)$, then the f-divergence is exactly equal to the last line in (7).\n",
    "\n",
    "This result was originaly presented in [[2]](#References) along with a way to estimate the actual value of any f-divergence (called Variational Divergence Estimation). It was later modified by [[1]](#References) to minimize f-divergences inspired by adversarial trainign as used in GANs. We shall see how in the following section.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training $f$-GANs\n",
    "To minimize a f-divergence [1] use adversarial training similar to what you saw in the previous note. We let $D_\\omega(x)$ in Eqn. (7) be a neural network parameterized by $\\omega$ and let the distribution $q$ be defined by another neural network $G_\\theta$ applied to some random noise from some simple distribution $Z$. In that case, we get the following objective:\n",
    "\n",
    "$$\n",
    "F(\\omega, \\theta) = \\mathbb{E}_{x\\sim p_d}[{D_\\omega}(x)] - \\mathbb{E}_{x\\sim p_g} [f^\\ast({D_\\omega}(x))] \\quad = \\quad \\mathbb{E}_{x\\sim p_d}[{D_\\omega}(x)] - \\mathbb{E}_{z \\sim p_z} [f^\\ast({D_\\omega}({G_\\theta}(z)))]\n",
    "$$\n",
    "\n",
    "This objective allows us to train a GAN for any $f$-divergence. We simply do gradient descent on $F(\\omega, \\theta)$ alternating between updating $\\omega$ and $\\theta$. In [[1]](#References) they manage to show that under certain smoothness assumptions their algorithm converges to a saddelpoint.\n",
    "\n",
    "In summary, we can minimize any $f$-divergence by playing the the following game: \n",
    "\n",
    "$$\n",
    "\\min_{G_\\theta} \\max_{D_\\omega} F(\\omega, \\theta)\n",
    "$$\n",
    "\n",
    "\n",
    "\n",
    "Below, we show a couple of f-divergences along with their convex conjugates below.\n",
    "\n",
    "\n",
    "<table style=\"width: 85%; margin-left:auto; margin-right:auto; text-align: left;\">\n",
    "    <col width=\"15%\">\n",
    "    <col width=\"45%\">\n",
    "    <col width=\"25%\">\n",
    "    <col width=\"15%\">\n",
    "    <tr>\n",
    "        <th>Name</th>\n",
    "        <th>$D_f(P || Q)$</th>\n",
    "        <th>Generator $f(u)$</th>\n",
    "        <th>$D_{opt}(x)$</th>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>Kullback-Leibler</td>\n",
    "        <td>$\\int p(x) \\log \\frac{p(x)}{q(x)} dx$</td>\n",
    "        <td>$u \\log u$</td>\n",
    "        <td>$1 + \\log \\frac{p(x)}{q(x)}$</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>Reverse KL</td>\n",
    "        <td>$\\int q(x) \\log \\frac{q(x)}{p(x)} dx$</td>\n",
    "        <td align=\"left\">$- \\log u$</td>\n",
    "        <td>$- \\frac{q(x)}{p(x)}$</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>Jensen-Shannon</td>\n",
    "        <td>$\\frac{1}{2} \\int p(x) \\log \\frac{2p(x)}{p(x) + q(x)} + q(x) \\log \\frac{2q(x)}{p(x) + q(x)} dx$</td>\n",
    "        <td>$u \\log u - (u + 1) \\log \\frac{u + 1}{2}$</td>\n",
    "        <td>$\\log \\frac{2p(x)}{p(x) + q(x)}$</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>GAN</td>\n",
    "        <td>$\\int p(x) \\log \\frac{2p(x)}{p(x) + q(x)} + q(x) \\log \\frac{2q(x)}{p(x) + q(x)} dx - \\log (4)$</td>\n",
    "        <td>$u \\log u - (u + 1) \\log (u + 1) $</td>\n",
    "        <td>$\\log \\frac{p(x)}{p(x) + q(x)}$</td>\n",
    "    </tr>\n",
    "</table>\n",
    "<p style=\"text-align: center\"><b>Table 1</b> Overview of different f-divergences<p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code\n",
    "\n",
    "## Learning a multimodal distribution from an underspecified GAN\n",
    "\n",
    "In this section we use a gan to translate $N(0, 1)$ to a biomodal distribution composed of two normal distributions; $N(4, 1)$ and $N(8,1)$. The network that we use is identical to the one we used last week: <br><br>\n",
    "\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "<img src=\"figs/week1_architecture.png\" style=\"width: 80%;\">\n",
    "<b>Figure 2: </b>A schematic of the simple model for this example.  \n",
    "</div>\n",
    "\n",
    "The following animation show how during training the generated data from the fake data distribution $p_g$ approaches the real data distribution $p_d$. However, since the generator is underspecified, in the sence that it cannot model a biomodal distribution with only the two parameters $w$ and $b$, it cannot learn the real distribution. Notice also how the decision boundary of the discriminator ends up spiking right between the two modes, as there are more fake data in that region than real data. A final thing to notice in the animation is how much the generator and discriminator depends on each other, e.g., when the generator moves its data into a new region, the discriminator changes to reflect that, which makes the generator move again, etc. This is exactly what we would expect from the formulation of the game that they are playing.\n",
    "\n",
    "<br/>\n",
    "<br/>\n",
    "\n",
    "<div style=\"text-align: center\">\n",
    "<img src=\"gifs/dual-mode/gif.gif\">\n",
    "<b>Figure 3: </b>Animation of how $p_g$ converges towards $p_d$.\n",
    "</div>\n",
    "\n",
    "\n",
    "<b>Abstract Idea of Code: </b><br>\n",
    "The code is implemented using Keras. The code uses Keras to define a generator and discriminator neural network. For the generator we also define an optimizer so we can train it on real/fake images. The generator needs to use the discriminator as a cost function. We do this by freezing the discriminators weights $\\text{discriminator.trainable = False}$ so we can define an optimizer for the generator without having it change the discriminator. \n",
    "\n",
    "All plots are made using matplotlib. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from keras.models import Sequential\n",
    "from keras.layers import Dense, Activation\n",
    "from keras.optimizers import Adam\n",
    "\n",
    "n       = 1000\n",
    "fig, ax = plt.subplots(figsize=(5, 3))\n",
    "\n",
    "def plot(): \n",
    "    ax.cla()\n",
    "    ax.hist(generator.predict_on_batch(noise) , alpha=0.3, label=\"Fake p_g\")\n",
    "    ax.hist(real_data , alpha=0.3, label=\"Real p_d\")\n",
    "    xs = np.arange(-8, 16, 0.1)\n",
    "    pred = discriminator.predict_on_batch(xs)\n",
    "    ax.plot(xs, pred*250, label=\"D(x)\")\n",
    "    ax.legend(loc=3)\n",
    "    ax.set_ylim([0, 280])\n",
    "    ax.set_xlim([-9, 13])\n",
    "    \n",
    "    fig.suptitle(\"Iteration: [%i / %i]\"%(i, iterations))\n",
    "    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
    "    fig.canvas.draw()\n",
    "    plt.savefig(\"gifs/dual-mode/%i.png\"%i)\n",
    "    plt.pause(.01)\n",
    "\n",
    "# number of samples. \n",
    "n = 1000\n",
    "iterations = 200\n",
    "repeat = 10\n",
    "\n",
    "# define generator\n",
    "generator = Sequential()\n",
    "generator.add(Dense(1, input_dim=1)) # one neuron except bias, don't have the relu activation!\n",
    "\n",
    "# define discrimiantor\n",
    "discriminator = Sequential()\n",
    "discriminator.add(Dense(10, input_dim=1, activation=\"relu\")) # non linearity has some use here. \n",
    "discriminator.add(Dense(1,  activation=\"sigmoid\"))\n",
    "discriminator.compile(loss='binary_crossentropy', optimizer=Adam(0.005))\n",
    "\n",
    "# Combine models\n",
    "gan = Sequential()\n",
    "gan.add(generator)\n",
    "discriminator.trainable = False # from the gan model we freeze discriminator to use it as loss function\n",
    "gan.add(discriminator)\n",
    "gan.compile(loss='binary_crossentropy', optimizer=Adam(0.005))\n",
    "\n",
    "# Do several updates, plot some interval of iterations fake / real data. \n",
    "for i in range(iterations):\n",
    "    noise     = np.random.normal(0, 1, size=(n, 1))\n",
    "    fake_data = generator.predict_on_batch(noise)\n",
    "    \n",
    "    x1        = np.random.normal(4, 1, size=(n//2, 1))\n",
    "    x2        = np.random.normal(8, 1, size=(n//2, 1))\n",
    "    real_data = np.c_[x1, x2].reshape(-1, 1)\n",
    "    \n",
    "    disc_X = np.concatenate((real_data, fake_data), axis=0)\n",
    "    disc_y = np.concatenate((np.zeros(n), np.ones(n)), axis=0) # flip labels since we min instead of max. \n",
    "    \n",
    "    plot()\n",
    "    for j in range(repeat): discriminator.train_on_batch(x=disc_X, y=disc_y)\n",
    "    for j in range(repeat): gan.train_on_batch(x=noise, y=np.zeros(n))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# References\n",
    "[[1]](http://papers.nips.cc/paper/6066-f-gan-training-generative-neural-samplers-using-variational-divergence-minimization.pdf) Nowozin, S., Cseke, B. and Tomioka, R., 2016. f-gan: Training generative neural samplers using variational divergence minimization. In Advances in Neural Information Processing Systems (pp. 271-279).\n",
    "\n",
    "[[2]](https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=5605355) Nguyen, X., Wainwright, M.J. and Jordan, M.I., 2010. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11), pp.5847-5861.\n",
    "\n",
    "[[3]](https://papers.nips.cc/paper/5423-generative-adversarial-nets) Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A. and Bengio, Y., 2014. Generative adversarial nets. In Advances in neural information processing systems (pp. 2672-2680)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}