7.2.2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "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.7.6"
    },
    "colab": {
      "name": "7.2-EXE-variational-autoencoder.ipynb",
      "provenance": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UWBaZtwusUWW"
      },
      "source": [
        "# Analyzing the VAE\n",
        "\n",
        "## Mandatory Exercises\n",
        "\n",
        "### Exercise 1.\n",
        "\n",
        "1. Implement the class `ReparameterizedDiagonalGaussian` (`log_prob()` and `rsample()`).\n",
        "2. Import the class `Bernoulli`\n",
        "3. Implement the class `VariationalInference` (computation of the `elbo` and `beta_elbo`).\n",
        "\n",
        "### Exercise 2.\n",
        "\n",
        "**Trainnig and Evaluating a VAE model**\n",
        "\n",
        "1. Why do we use the reparameterization trick?\n",
        "2. What available metric can you use to estimate the marginal likelihood ($p_\\theta(\\mathbf{x})$) ?\n",
        "3. In the above plots, we display numerous model samples. If you had to pick one plot, which one would you pick to evaluate the quality of a VAE (i.e. using posterior samples $\\mathbf{z} \\sim q_\\phi(\\mathbf{z} | \\mathbf{x})$ or prior samples $\\mathbf{z} \\sim p(\\mathbf{z})$) ? Why?.\n",
        "4. How could you exploit the VAE model for classification?\n",
        "\n",
        "**Answers**:\n",
        "\n",
        "`[...]`\n",
        "\n",
        "### Exercise 3.\n",
        "\n",
        "**Experiment with the VAE model.**\n",
        "\n",
        "1. Experiment with the number of layers and activation functions in order to improve the reconstructions and latent representation. What solution did you find the best and why?\n",
        "2. Try to increase the number of digit classes in the training set and analyze the learning curves, latent space and reconstructions. For which classes and why does the VAE fail in reconstructing?  *HINT: Try the combination: `classes=[0, 1, 4, 9]`, to see how well VAE can separate these digits in the latent representation and reconstructions.*\n",
        "3. Increase the number of units in the latent layer. Does it increase the models representational power and how can you see and explain this? How does this affect the quality of the reconstructions?\n",
        "\n",
        "**Answers**:\n",
        "\n",
        "`[...]`\n",
        "\n",
        "### Exercise 4. \n",
        "\n",
        "**Analyze the purpose of the KL-term and the $\\beta$ parameter.**\n",
        "\n",
        "1. How does the KL-term, $\\mathcal{D}_{\\operatorname{KL}}\\left(q_\\phi(\\mathbf{z}|\\mathbf{x})\\ |\\ p(\\mathbf{z})\\right)$, work as a regulariser on the distributions over latent variables? *HINT*: When maximising the ELBO, the probability-distance measure is minimised $\\operatorname{KL} \\rightarrow 0$ so that $q(z|x) \\rightarrow p(z) = \\mathcal{N}(z|0,I)$ for all examples, x. At $\\operatorname{KL} = 0$ variations in x stops having an affect on the latent distribution and latent units are all described by the same distribution, $\\mathcal{N}(z|0,I)$, so they produce a noisy output without signal (i.e. $\\mathbf{z}=\\epsilon \\sim \\mathcal{N}(0,I)$) to the decoder.\n",
        "2. Try removing the KL-term (using the $\\beta$ parameter) and analyze what happens during training, in the learning curves, latent representation and reconstructions compared to before removing it.\n",
        "3. What does the loss reduces to? Explain how this can affect a VAE. *HINT*: Compare loss function for AE and VAE, and remember that we can use the pixel-wise binary crossentropy error as the loss in the AEs and for the reconstruction error, $\\log p_\\theta(\\mathbf{x}|\\mathbf{z}) = \\log \\mathcal{B}(\\mathbf{x} | g_\\theta(z))$, in VAEs.\n",
        "4. Experiment with different $\\beta$ values (e.g. [0.1, 10, 50]) and explain how this affects the reconstruction quality and the latent representations. *HINT* what is the tradeoff between reconstruction error ($\\log p_\\theta(\\mathbf{x} | \\mathbf{z})$) and the KL term?\n",
        "\n",
        "**Answers**:\n",
        "\n",
        "`[...]`\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "Dhx4VKMPsUWW"
      },
      "source": [
        "## Optional exercises\n",
        "\n",
        "- OPT: Use the original paper http://arxiv.org/pdf/1312.6114v10.pdf or [this blog](http://blog.shakirm.com/2015/10/machine-learning-trick-of-the-day-4-reparameterisation-tricks/) to explain what the reparameterization trick does.\n",
        "- OPT: Look through https://arxiv.org/abs/1611.00712 or https://arxiv.org/abs/1611.01144 and explain how one could instead introduce a categorical latent variable for $z$.\n",
        "- OPT: Implement the Gumbel softmax trick thereby letting $z$ take a categorical distribution.\n",
        "- OPT: The VAE is a probablistic model. We could model $p(x,z,y)$ where $y$ is the label information. Explain how this model could handle semi-supervised learning? You can look through the papers https://arxiv.org/pdf/1406.5298.pdf or  https://arxiv.org/pdf/1602.05473v4.pdf or again the two papers on Gumbel softmax.\n",
        "\n",
        "**Answers**:\n",
        "\n",
        "`[...]`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h4qmAYzQsUWX"
      },
      "source": [
        "# Credits \n",
        "\n",
        "- Original [Theano/Lasagne tutorial](https://github.com/DeepLearningDTU/nvidia_deep_learning_summercamp_2016/blob/master/lab1/lab1_FFN.ipynb) by \n",
        "Lars Maaløe ([larsmaaloee](https://github.com/larsmaaloee)),\n",
        "Søren Kaae Sønderby ([skaae](https://github.com/skaae)), and \n",
        "Casper Sønderby ([casperkaae](https://github.com/casperkaae)). \n",
        "- Converted to TensorFlow and updated by Maximillian F. Vording ([maximillian91](https://github.com/maximillian91)).\n",
        "- Converted to PyTorch and updated by Jesper Wohlert ([wohlert](https://github.com/wohlert)).\n",
        "- Major update in 2020: focus on the probabilistic interpretation, use of `torch.distributions` and improved visualizations by Valentin Lievin ([vlievin](http://vlievin.github.io))."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "nekjOIihsUWX"
      },
      "source": [
        "# Supplementary Material\n",
        "\n",
        "## Why a VAE learns a good approximate posterior $q_\\phi(\\mathbf{z} | \\mathbf{x}) \\approx p_\\theta(\\mathbf{z} | \\mathbf{x})$\n",
        "\n",
        "In order to compare the two distributions, let us focus on their KL-divergence (**important**: the divergence is expressed w.r.t $p_\\theta(\\mathbf{z}|\\mathbf{x})$ and not $p(\\mathbf{z})$, this is not the KL-divergence studied in the experiments):\n",
        "\n",
        "$$\\mathcal{D}_{\\operatorname{KL}}(q_\\phi(\\mathbf{z} | \\mathbf{x}) || p_\\theta(\\mathbf{z} | \\mathbf{x})) = \\int_z q_\\phi(\\mathbf{z} | \\mathbf{x}) \\log \\frac{q_\\phi(\\mathbf{z} | \\mathbf{x}) }{p_\\theta(\\mathbf{z} | \\mathbf{x})}dz = \\mathbb{E}_{q_\\phi(\\mathbf{z} | \\mathbf{x})} \\left[\\log \\frac{q_\\phi(\\mathbf{z} | \\mathbf{x}) }{p_\\theta(\\mathbf{z} | \\mathbf{x})}\\right]$$\n",
        "\n",
        "which is a non-negative distance measure between distributions, so by minimising it wrt. to the parameters in $q_\\phi(\\mathbf{z} | \\mathbf{x}) $, the distribution moves close to our unknown $p_\\theta(\\mathbf{z} | \\mathbf{x})$. But as $p_\\theta(\\mathbf{z} | \\mathbf{x})$ is unknown and would include some rather intractable integrals over neural networks, we can instead get rid of it by expressing it through Bayes rule $p(z|x) = p(x|z)p(z)/p(x)$ and thereby decompose the KL-divergence into our log-likelihood and lower bound:\n",
        "\n",
        "$$ \\mathcal{D}_{\\operatorname{KL}}(q_\\phi(\\mathbf{z} | \\mathbf{x})  || p_\\theta(\\mathbf{z} | \\mathbf{x})) \n",
        "= \\int_\\mathbf{z} q_\\phi(\\mathbf{z} | \\mathbf{x}) \\log \\frac{q_\\phi(\\mathbf{z} | \\mathbf{x}) p_\\theta(\\mathbf{x})}{p_\\theta(\\mathbf{x} | \\mathbf{z})\n",
        "p(z)}d\\mathbf{z}\n",
        "= \\int_\\mathbf{z} q_\\phi(\\mathbf{z} | \\mathbf{x}) \\log \\frac{q_\\phi(\\mathbf{z} | \\mathbf{x}) }{p_\\theta(\\mathbf{x} | \\mathbf{z})p_\\theta(\\mathbf{x} | \\mathbf{z})(\\mathbf{z})}d\\mathbf{z} + \\log p_\\theta(\\mathbf{x})$$ \n",
        "\n",
        "by seeing that the likelihood, $p_\\theta(\\mathbf{x})$, is independent of $\\mathbf{z}$ and pull it out of the integral. We can flip the sign and fraction in the integral term to recognise it as the negative lower bound\n",
        "\n",
        "$$\\mathcal{D}_{\\operatorname{KL}}(q_\\phi(\\mathbf{z} | \\mathbf{x})  || p_\\theta(\\mathbf{z} | \\mathbf{x})) = - \\int_z q_\\phi(\\mathbf{z} | \\mathbf{x}) \\log \\frac{p_\\theta(\\mathbf{x} | \\mathbf{z})p(\\mathbf{z})}{q_\\phi(\\mathbf{z} | \\mathbf{x}) }dz + \\log p_\\theta(\\mathbf{x}) =  -\\mathcal{L}(x) + \\log p_\\theta(\\mathbf{x})$$\n",
        "\n",
        "We then find the log-likelihood to consist of the two terms and hold the inequality\n",
        "\n",
        "$$\\log p_\\theta(\\mathbf{x}) =  \\mathcal{D}_{\\operatorname{KL}}(q_\\phi(\\mathbf{z} | \\mathbf{x})  || p_\\theta(\\mathbf{z} | \\mathbf{x})) + \\mathcal{L}(\\mathbf{x}) \\geq \\mathcal{L}(\\mathbf{x})$$\n",
        "\n",
        "where the KL-divergence is non-zero and the log-likelihood is $\\log p_\\theta(\\mathbf{x}) \\leq 0$. This means that maximising the lower bound from the negative domain towards $0$ will also maximise the log-likelihood, while pushing down the KL-divergence until $q_\\phi(\\mathbf{z} | \\mathbf{x}) $ cannot move closer to natures true distribution, $p_\\theta(\\mathbf{z} | \\mathbf{x})$. So how close the lower bound can get to the log-likelihood is dependent on the flexibility of the distribution we choose for $q_\\phi(\\mathbf{z} | \\mathbf{x}) $. \n",
        "\n",
        "\n",
        "\n",
        "## Importance Sampling in VAEs\n",
        "\n",
        "In VAEs, the true posterior distribution $p_\\theta(\\mathbf{z} | \\mathbf{x})$ is unknown. Instead we use an approximate posterior $q_\\theta(\\mathbf{z} | \\mathbf{x})$. $q$ is tractable and one can sample from it. This is an instance of importance-sampling: evaluating a quantity $\\mathbb{E}_{p(\\mathbf{z})} \\left[  f(\\mathbf{z})\\right]$ given samples from another distribtion $q$.\n",
        "\n",
        "### 1. Illustration\n",
        "\n",
        "Let us consider a small example to illustrate importance-sampling. In this example we define a distribution $p(\\mathbf{z})$ from which we (in practice) cannot sample and define a distribution $q(\\mathbf{z})$ from which we can sample. \n",
        "\n",
        "Then we can estimate $\\mathbb{E}_{p(\\mathbf{z})} \\left[  f(\\mathbf{z})\\right]$ using samples from $q$:\n",
        "\n",
        "$$\\mathbb{E}_{p(\\mathbf{z})} \\left[  f(\\mathbf{z})\\right] = \\mathbb{E}_{q(\\mathbf{z})} \\left[ w(\\mathbf{z}) f(\\mathbf{z})\\right], \\quad w(\\mathbf{z}) := \\frac{p(\\mathbf{z})}{q(\\mathbf{z})}$$\n",
        "\n",
        "In this example, we use $f(\\mathbf{z}) = \\mathbf{z}$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "jw4ntrsqsUWY"
      },
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from torch.distributions import Normal\n",
        "cuda = torch.cuda.is_available()\n",
        "    \n",
        "n_samples = 10000\n",
        "c1 = sns.color_palette()[0]\n",
        "c2 = sns.color_palette()[1]\n",
        "\n",
        "# define p(z) and q(z)\n",
        "p = Normal(loc=5 * torch.ones((1,)), scale=torch.ones((1,))) \n",
        "q = Normal(loc=torch.ones((1,)), scale=3*torch.ones((1,)))\n",
        "\n",
        "# sample z ~ q(z)\n",
        "q_samples = q.sample(sample_shape=torch.Size([n_samples]))\n",
        "\n",
        "# compute the densities p(z) and q(z)\n",
        "p_prob = p.log_prob(q_samples).exp()\n",
        "q_prob = q.log_prob(q_samples).exp()\n",
        "\n",
        "# define the importance-weights: w(z) = p(z) / q(z)\n",
        "w = p_prob / (q_prob + 1e-12)\n",
        "\n",
        "# compute the importance-weighted estimate \\hat{\\mu} = E_q[ w(z) f(z)] where f(z) = z\n",
        "estimate = (w * q_samples).mean()\n",
        "\n",
        "# compute the ground true\n",
        "p_samples = p.sample(sample_shape=torch.Size([n_samples]))\n",
        "empirical_mean = p_samples.mean()\n",
        "\n",
        "# display the results\n",
        "print(f\">> Empirical mean = {empirical_mean:.3f}, estimate = {estimate:.3f}\")\n",
        "plt.figure(figsize=(16,6))\n",
        "sns.distplot(p_samples, label=\"$p(z)$\", hist_kws={'alpha':0.2}, color=c1)\n",
        "sns.distplot(q_samples, label=\"$q(z)$\",hist_kws={'alpha':0.2}, color=c2)\n",
        "plt.axvline(empirical_mean, color=c1,linestyle=\"-\",label= \"Empirical Mean\")\n",
        "plt.axvline(estimate, color=c2,linestyle=\"-\", label= \"Importance-Weighted Estimate\")\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "I8hLuVx1sUWc"
      },
      "source": [
        "### 2. Application: Accurate Estimate of $\\log p_\\theta(\\mathbf{x})$ using Importance Sampling\n",
        "\n",
        "The ELBO is a lower bound to the marginal log likelihood, however it is not guaranteed to be accurate. [Importance-sampling can be leveraged in VAEs](https://arxiv.org/abs/1509.00519) to provide a more accurate estimate of the marginal log-likelihood $\\log p_\\theta(\\mathbf{x})$ using multiple samples from $q_\\phi(\\mathbf{z}_1,\\dots,\\mathbf{z}_K | \\mathbf{x})$. With $K+1$ samples $q_\\phi(\\mathbf{z} | \\mathbf{x})$, the following inequality holds:\n",
        "\n",
        "$$\\log p_\\theta(\\mathbf{x}) \\geq \\mathcal{L}_{K+1}(\\mathcal{x}) \\geq \\mathcal{L}_K(\\mathcal{x}) \\geq \\mathcal{L}(\\mathcal{x})$$\n",
        "\n",
        "where\n",
        "\n",
        "$$ \\mathcal{L}_{K}(\\mathcal{x}) := \\mathbb{E}_{q_\\phi(\\mathbf{z}_1,\\dots,\\mathbf{z}_K | \\mathbf{x})} \\left[ \n",
        "\\log \\frac{1}{K} \\sum_{k=1}^K \\frac{p_\\theta(\\mathbf{x}, \\mathbf{z}_k)}{q_\\phi(\\mathbf{z}_k | \\mathbf{x})} \n",
        "\\right]$$\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NyEI0xjAsUWc"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}