lectures/lecture-13/lecture-13.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Lecture 13:  Generative models I"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Previous lecture: Adversarial attacks and training\n",
    "- Adversarial attacks\n",
    "- Adversarial training\n",
    "- Robustness of DL models (using randomized smoothing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Current lecture: Generative models I\n",
    "\n",
    "- Autoregressive models \n",
    "- Variational Autoencoders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Generative models\n",
    "\n",
    "So far, we talked about supervised learning and self-supervised learning, where the goal is to encode an object $x$\n",
    "\n",
    "into a meaningful feature representation $E(x)$ such that we can distinguish different object.\n",
    "\n",
    "The goal of generative models is different: \n",
    "\n",
    "Given samples $x_1, \\ldots, x_N$ from the unknown probability distribution $p(x)$, we want to learn the probability distribution itself.\n",
    "\n",
    "The minimum what we want to learn is a model that allows to **sample** from the distribution itself.\n",
    "\n",
    "In some cases, we can also learn the **density** $p(x)$, so we can evaluate the **log-likelihood** itself.\n",
    "\n",
    "We have encountered generative models when we talked about **transformer models**.\n",
    "\n",
    "GPT is a generative autoregressive model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Types of generative models\n",
    "\n",
    "Today, there are quite a lot generative models.\n",
    "\n",
    "Historically, one of the first family for generative models were **autoregressive models**, that focus on **predicting sequence values** in the future.\n",
    "\n",
    "Other important directions which we will cover in this course include:\n",
    "\n",
    "- Generative adversarial networks (implicit generative models)\n",
    "- Variational autoencoders\n",
    "- Diffusion models\n",
    "\n",
    "They are very different concepts, but they all solve the problem of learning (parametric form) of high-dimensional probability distribution from samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Probabilistic autoregressive model\n",
    "\n",
    "The autoregressive model is based on the representation of the probability density (for the continuous density) or probability (for discrete symbols as)\n",
    "\n",
    "\n",
    "$$p(x_1, \\ldots, x_T) = \\prod_{t=1}^T p(x_t \\vert x_{t-1}, \\ldots, x_1).$$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## NADe\n",
    "\n",
    "The first neural autoregressive distribution estimator has been proposed in \n",
    "[The Neural Autoregressive Distribution Estimator\n",
    "](https://proceedings.mlr.press/v15/larochelle11a.html) in 2011\n",
    "\n",
    "The architecture is not up-to-date, but quite instructive. It was motivated by fully visible sigmoid belief network, proposed as early as 1992! In FVSBN, the value is conditioned to the previous ones using 1-layer network.\n",
    "\n",
    "<img src='nade.png' width=80%>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## PixelCNN\n",
    "\n",
    "The idea of variational like-lihood based model for images have been proposed in two papers:\n",
    "\n",
    "[Conditional image generation with PixelCNN decoders](https://proceedings.neurips.cc/paper_files/paper/2016/hash/b1301141feffabac455e1f90a7de2054-Abstract.html)\n",
    "\n",
    "[Pixel Recurrent Neural Networks\n",
    "](https://proceedings.mlr.press/v48/oord16.html)\n",
    "\n",
    "Images formally is not a sequence, so we need to define a certain order."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## PixelCNN ordering\n",
    "\n",
    "In PixelCNN , the conditional distribution is modelled using convolutional neural network.\n",
    "\n",
    "There are several variants to model such conditional distributions.\n",
    "\n",
    "[PixelRNN paper](https://proceedings.mlr.press/v48/oord16.html) uses LSTM to build such conditional distribution.\n",
    "\n",
    "In [Conditional image generation with PixelCNN decoders](https://proceedings.neurips.cc/paper_files/paper/2016/hash/b1301141feffabac455e1f90a7de2054-Abstract.html)\n",
    " the ReLU units are replaced with **masked gated convolutions**:\n",
    "\n",
    "$$y = \\operatorname{tanh}W_{k, f} * x) \\odot \\sigma(W_{k, g} * x).$$\n",
    "\n",
    "It uses the **raster scan** with the following mask on the filter!\n",
    "\n",
    "In also uses two blocks to increase the **receptive field**: vertical and horizontal. One conditions only on the current row, the second conditions on all the previous rows.\n",
    "\n",
    "<img src='pixelcnn-mask.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Improving PixelCNN\n",
    "\n",
    "Certain improvents have been proposed in the paper\n",
    "[Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel. PixelSNAIL: An Improved Autoregressive Generative Model](http://proceedings.mlr.press/v80/chen18h/chen18h.pdf)\n",
    "\n",
    "The available options to model conditional distributions:\n",
    "1. Traditional RNNs, difficult to model long-range dependence for images\n",
    "2. Masked (gated) convolutions\n",
    "3. Of couse, we can use self-attention!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## PixelSNAIL architecture\n",
    "\n",
    "The whole decoder architecture has the following form:\n",
    "\n",
    "<img src='pixel-snail.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## PixelSNAIL-blocks\n",
    "\n",
    "The blocks mentioned in the previous slide have the following form.\n",
    "\n",
    "<img src='px-blocks.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Applications of autoregressive generative models\n",
    "\n",
    "The autoregressive generative models have been proposed for images, but the biggest usage and success have been achieved not for image, but for time-series (which is natural).\n",
    "\n",
    "For audio signals and tasks, such models held state-of-the art for quite a long time and still are being used.\n",
    "\n",
    "Namely, this is the WaveNet architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## WaveNet\n",
    "\n",
    "In [Generative model for raw audio](https://arxiv.org/abs/1609.03499) the WaveNet architecture has been proposed.\n",
    "\n",
    "Similarly to PixelCNNs (van den Oord et al., 2016a;b), the conditional probability distribution is\n",
    "modelled by a stack of convolutional layers. There are no pooling layers in the network, and the\n",
    "output of the model has the same time dimensionality as the input. The model outputs a categorical\n",
    "distribution over the next value xt with a softmax layer and it is optimized to maximize the loglikelihood of the data w.r.t. the parameters. Because log-likelihoods are tractable, we tune hyperparameters on a validation set and can easily measure if the model is overfitting or underfitting.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Causal convolutions\n",
    "\n",
    "The model can not depend on future timesteps, i.e. be **causal**.\n",
    "\n",
    "For images, the analogue of the causal convolution is **masked convolution**, where in the convolution the 'wrong' pixels are multiplied to zero.\n",
    "\n",
    "At training, such convolutions can be done simultaneously, whereas at inference we have to do prediction one at a time.\n",
    "\n",
    "Q: What is the difference to recurrent neural networks (RNN)?\n",
    "\n",
    "\n",
    "<img src='causal-conv.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Properties of causal convolutions\n",
    "\n",
    "- They are faster to train than recurrent neural networks, especially for longer sequences\n",
    "- For longer sequence we need **more hidden layers**: the number of hidden layers determine the **receptive field** of the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Dilated causal convolutions\n",
    "\n",
    "- A dilated convolution (also called a trous ` , or convolution with holes) is a convolution where the\n",
    "filter is applied over an area larger than its length by skipping input values with a certain step. It is\n",
    "equivalent to a convolution with a larger filter derived from the original filter by dilating it with zeros,\n",
    "but is significantly more efficient\n",
    "\n",
    "- **Stacked dilated convolutions** enable networks to have **very large receptive fields** with just a few layers, while preserving the input resolution throughout the network as well as computational efficiency.\n",
    "\n",
    "- The dilation is doubled every layer\n",
    "\n",
    "<img src='wavenet-animation.gif' width=100%>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Architecture details\n",
    "\n",
    "The model in the non-linearity uses **gated units** of the form\n",
    "\n",
    "$$\\mathbf{z}=\\tanh \\left(W_{f, k} * \\mathbf{x}\\right) \\odot \\sigma\\left(W_{g, k} * \\mathbf{x}\\right)$$\n",
    "\n",
    "<img src='wavenet-architecture.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Generating speech with Wavenet (conditioning)\n",
    "\n",
    "WaveNet architecture has been used for **Text-to-Speech generation.**\n",
    "\n",
    "In order to condition the generation on text, the **conditioning** has been introduced into the architecture.\n",
    "\n",
    "It is achieved by upsampling such local features to the level and using inside the **gating mechanism**\n",
    "\n",
    "Given $h_t$, possibly with a lower sampling frequency than the audio signal, e.g. linguistic features in a Text-to-Speech (TTS) model. \n",
    "\n",
    "This time series is first transformed  using a transposed convolutional network (learned upsampling) that maps it to a new time series $\\mathbf{y}=f(\\mathbf{h})$ with the same resolution as the audio signal, \n",
    "which is then used in the activation unit as follows:\n",
    "\n",
    "$$\n",
    "\\mathbf{z}=\\tanh \\left(W_{f, k} * \\mathbf{x}+V_{f, k} * \\mathbf{y}\\right) \\odot \\sigma\\left(W_{g, k} * \\mathbf{x}+V_{g, k} * \\mathbf{y}\\right)\n",
    "$$\n",
    "\n",
    "Note the gating mechanism again (actually, quite powerful idea)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Summary for simple autoregressive models\n",
    "\n",
    "- Autoregressive models follow the same ideology as modern transformers (modelling conditional distributions)\n",
    "- First attempts included RNN/GRU type model for the conditional distributions\n",
    "- PixelCNN has reached good results (at that time) for modelling image distribution at the pixel level. It used masked convolutions\n",
    "- WaveNet has got SOTA (at that time) for generation of speech from text, music and audio. It uses dilated causal convolutions.\n",
    "- Generation is **very slow** due to generation by one sample at a time.\n",
    "\n",
    "The autoregressive models represent the probability in a tractable manner. \n",
    "\n",
    "The variational autoencoders (VAE) represent probability in a way that likelihood can not be computed, but can be optimized using **variational inference** and samples.\n",
    "\n",
    "But first, we will discuss models based on the transformations of the random variable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Other models\n",
    "\n",
    "The autoregressive models can be considered as a special case of a more general construction:\n",
    "\n",
    "Learning an invertible transformation from a simple distribution with a tractable Jacobian.\n",
    "\n",
    "There are several examples of such approach: [Real NVP](https://arxiv.org/abs/1605.08803), [Glow](https://proceedings.neurips.cc/paper/2018/hash/d139db6a236200b21cc7f752979132d0-Abstract.html).\n",
    "\n",
    "Lets discuss the idea"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Change of variables formula\n",
    "\n",
    "The idea of real NVP is to parametrize the distribution as:\n",
    "\n",
    "$$z \\sim p(z), x = f^{-1}(z)  \\quad \\mbox{(generation)},$$\n",
    "\n",
    "where $f$ is a parametrized invertible transformation.\n",
    "\n",
    "The change of variables formula gives\n",
    "$$p_X(x)=p_Z(z)\\left|\\operatorname{det}\\left(\\frac{\\partial g(z)}{\\partial z^T}\\right)\\right|^{-1}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Parametrizing the mapping\n",
    "\n",
    "In [RealNVP model](https://arxiv.org/abs/1605.08803), $f$ is parametrized as a superposition of **affine coupling layers of the form**\n",
    "$$\n",
    "\\begin{aligned}\n",
    "y_{1: d} & =x_{1: d} \\\\\n",
    "y_{d+1: D} & =x_{d+1: D} \\odot \\exp \\left(s\\left(x_{1: d}\\right)\\right)+t\\left(x_{1: d}\\right)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "The Jacobian of such kind of mapping is triangular, and the determinant (and its logarithm) is easily computed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Glow\n",
    "\n",
    "In [Glow model](https://proceedings.neurips.cc/paper_files/paper/2018/file/d139db6a236200b21cc7f752979132d0-Paper.pdf)\n",
    "\n",
    "Invertible $1 \\times 1$ convolutions are added into the architecture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Later development\n",
    "\n",
    "- The formulas for change of variables are used in diffusion models\n",
    "- They are also used in FFJord approaches, when the forward model is parametrized by a **Neural Ordinary Differential Equation**.\n",
    "\n",
    "Now, lets get to the variational autoencoders!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Variational autoencoders\n",
    "\n",
    "Now we will discuss VAE, which is one of the most interesting and powerful class of generative models.\n",
    "\n",
    "\n",
    "VAE has been introduced in\n",
    "[Autoencoder Variational Bayes](https://arxiv.org/abs/1312.6114) and \n",
    "\n",
    "in the [paper](https://arxiv.org/abs/1401.4082) and has become one of the most influencial approaches in generative models.\n",
    "\n",
    "Several reviews are available, for example [this one](https://arxiv.org/abs/1906.02691)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Standard autoencoder\n",
    "\n",
    "<img src='vae-intuition.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VAE intuition \n",
    "\n",
    "We did not discuss too much classical **autoencoder**: we want to map an input $x$ to the latent space and then map back in order to reconstruct the initial image.\n",
    "\n",
    "Problems:\n",
    "- In the latent space we can just 'remember' the input\n",
    "- Such model can not generate new samples\n",
    "\n",
    "\n",
    "<img src='vae-int2.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VAE intuition (2)\n",
    "\n",
    "In VAE, we learn not the exact features, but some intervals (distributions of them).\n",
    "\n",
    "I.e., the encoder maps the input into a certain distribution (which can be, of course, close to a delta-function).\n",
    "\n",
    "Now, lets come up with a mathematical formalism.\n",
    "\n",
    "<img src='vae-int3.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## The model\n",
    "\n",
    "Suppose there exists a hidden variable $z$ that generates the output variable $x$.\n",
    "\n",
    "We know the distribution of $z$, $p(z)$, we know the distribution of $x$ conditioned to $z$, $p(x \\vert z)$.\n",
    "\n",
    "\n",
    "If we want to encode, we need to know the distribution $p(z \\vert x)$. This can be done by using Bayes formula:\n",
    "\n",
    "$$p(z \\vert x) = \\frac{p(x \\vert z) p(z)}{p(x)}, \\quad p(x) = \\int p(x \\vert z) p(z) dz. $$\n",
    "\n",
    "$p(x)$ is very difficult to evaluate in the general case!\n",
    "\n",
    "\n",
    "<img src='graph-model.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Approximating the distribution\n",
    "\n",
    "$$p(z \\vert x) = \\frac{p(x \\vert z) p(z)}{p(x)}, \\quad p(x) = \\int p(x \\vert z) p(z) dz. $$\n",
    "\n",
    "We can use **variational inference**!\n",
    "\n",
    "Lets approximate $p(z \\vert x)$ by another distribution $q(z \\vert x)$ \n",
    "\n",
    "Now we can minimize the distance $\\min \\mathrm{KL}\\left( {q\\left( {z|x} \\right)||p\\left( {z|x} \\right)} \\right)$, where $\\mathrm{KL}$ is the **Kullback–Leibler** divergence between to distributions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## KL-divergence\n",
    "\n",
    "KL-divergence between two continuous distributions is defined as\n",
    "\n",
    "$$\\mathrm{KL}(p || q) = E_p \\log \\frac{p}{q} = \\int p(x) \\log \\frac{p}{q}$$\n",
    "\n",
    "KL is a metric between two distributions, it has the following properties:\n",
    "\n",
    "1. $\\mathrm{KL}(p || q) \\geq 0$\n",
    "2. Is non-symmetric\n",
    "3. If it is zero, the distributions are the same."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Deriving VAEs formulas\n",
    "\n",
    "Recall, that we are minimizing $\\min \\mathrm{KL}\\left( {q\\left( {z|x} \\right)||p\\left( {z|x} \\right)} \\right)$ over $q$.\n",
    "\n",
    "\\begin{aligned}\n",
    "\\mathrm{KL}[\\mathrm{q}(\\mathrm{z}) \\| \\mathrm{p}(\\mathrm{z} \\mid \\mathrm{x})] & =\\int_z \\mathrm{q}(\\mathrm{z}) \\log \\frac{\\mathrm{q}(\\mathrm{z})}{\\mathrm{p}(\\mathrm{z} \\mid \\mathrm{x})} \\\\\n",
    "& =-\\int_z \\mathrm{q}(\\mathrm{z}) \\log \\frac{\\mathrm{p}(\\mathrm{z} \\mid \\mathrm{x})}{\\mathrm{q}(\\mathrm{z})} \\\\\n",
    "& =-\\left(\\int_z \\mathrm{q}(\\mathrm{z}) \\log \\frac{\\mathrm{p}(\\mathrm{x}, \\mathrm{z})}{\\mathrm{q}(\\mathrm{z})}-\\int_{\\mathrm{z}} \\mathrm{q}(\\mathrm{z}) \\log \\mathrm{p}(\\mathrm{x})\\right) \\\\\n",
    "& =-\\int_z \\mathrm{q}(\\mathrm{z}) \\log \\frac{\\mathrm{p}(\\mathrm{x}, \\mathrm{z})}{\\mathrm{q}(\\mathrm{z})}+\\log \\mathrm{p}(\\mathrm{x}) \\int_z \\mathrm{q}(\\mathrm{z}) \\\\\n",
    "& =-\\mathrm{L}+\\log \\mathrm{p}(\\mathrm{x})\n",
    "\\end{aligned}\n",
    "\n",
    "The value $L$ is called **variational lower bound**.\n",
    "\n",
    "We get\n",
    "\n",
    "$$L = \\log p(x) - \\mathrm{KL}(\\mathrm{q}(\\mathrm{z}) \\| \\mathrm{p}(\\mathrm{z} \\mid \\mathrm{x}))$$\n",
    "\n",
    "Since KL-divergence is non-negative, we get \n",
    "\n",
    "$$\\log p(x) \\geq L.$$\n",
    "\n",
    "Also, since for a fixed $x$ $\\log p(x)$ is a constant, minimizing over $q = q(z \\vert x)$ of the original KL divergence is equivalent to the maximization of $L$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## ELBO \n",
    "\n",
    "This bound is typically referred to as ELBO (Evidence Lower Bound).\n",
    "\n",
    "We have (recall that $p(x, z) = p(x \\vert z) p(z)$,\n",
    "\n",
    "$$L = \\int_z q(z) \\log \\frac{p(x, z)}{q(z)} = {E_{q\\left( {z|x} \\right)}}\\log p\\left( {x|z} \\right) - KL\\left( {q\\left( {z|x} \\right)||p\\left( z \\right)} \\right) $$\n",
    "\n",
    "These two terms have clear interpretation.\n",
    "\n",
    "The first term is a **reconstruction term** of the output from the latent vector.\n",
    "\n",
    "The second term is the **similarity** of the the distribution of $q(z \\vert x)$ to a prior distribution $p(z)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Using ELBO\n",
    "\n",
    "We can use $q(z \\vert x)$ as the **encoder** in the architecture to infer possible latent variables $z$ for a given $x$.\n",
    "\n",
    "We can parametrize $q(z \\vert x)$ by a neural network.\n",
    "\n",
    "But what actually this neural network will output? \n",
    "\n",
    "<img src='vae-4.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Parametrizing the latent distribution as Gaussians\n",
    "\n",
    "The classical VAE parametrizes the latent distribution $q(z \\vert x)$ as a Gaussian with mean value $\\mu$ and diagonal covariance matrix $\\sigma$.\n",
    "\n",
    "The **neural network** will predict $\\mu_j$ and $\\sigma_j$ for each latent variable.\n",
    "\n",
    "$\\sigma_j$ has to be non-negative, so typically $\\log \\sigma_j$ is predicted rather than $\\sigma_j$ itself.\n",
    "\n",
    "But, the computational graph is stochastic.\n",
    "\n",
    "How we differentiate through it?\n",
    "\n",
    "\n",
    "<img src='mean-sigma.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## But before...\n",
    "\n",
    "Before it is needed to mention that the KL-divergence between two Gaussian can be computed analytically. \n",
    "\n",
    "The decoder network $p(x \\vert z)$ will also produce mean value and variance for each pixel! \n",
    "\n",
    "\n",
    "So:\n",
    "\n",
    "1. We pass $x$ though a network to get $\\mu_{z \\vert x}, \\sigma_{z\\vert x}$\n",
    "2. Evaluate KL-divergence between $N(\\mu_z, \\sigma_z)$ and $p(z)$\n",
    "2. We sample from this distribution\n",
    "3. Using decoder network, predict $\\mu_{x \\vert z}$ and $\\sigma_{x \\vert z}$ for those samples\n",
    "4. Evaluate the log-likelihood of the original input $x$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## How we compute the gradients?\n",
    "\n",
    "ELBO has the following form:\n",
    "\n",
    "$$L = {E_{q\\left( {z|x} \\right)}}\\log p\\left( {x|z} \\right) - KL\\left( {q\\left( {z|x} \\right)||p\\left( z \\right)} \\right)$$\n",
    "\n",
    "We need to compute the gradient with respect to the parametres of $p(x \\vert z)$ and parameters of the distribution $q$ (also called variational parameters).\n",
    "\n",
    "\n",
    "The main challenge is the first terms, since for Gaussians the second term is an analytic formula.\n",
    "\n",
    "Gradient of $E_{q\\left( {z|x} \\right)}\\log p\\left( {x|z} \\right)$ with respect to $p$ can be computed as an expectation:\n",
    "\n",
    "$$ \\nabla_p E_{q\\left( {z|x} \\right)}\\log p\\left( {x|z} \\right) = E_{q\\left( {z|x} \\right)} \\nabla \\log p\\left( {x|z} \\right)$$\n",
    "\n",
    "The gradient with respect to the varitional parameters $q$ is much trickier. \n",
    "\n",
    "Fortunately, for basic distributions the expectation over parameter-dependent distribution can be estimated using **reparametrization trick**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Reparametrization trick\n",
    "\n",
    "For VAEs, a special technique called **reparametrization trick** has been proposed. \n",
    "\n",
    "For Gaussians, the trick is simple.\n",
    "\n",
    "1. Sample $\\varepsilon$ from $N(0, 1)$.\n",
    "2. Output $\\mu + \\sigma \\varepsilon$.\n",
    "\n",
    "Now we can differentiate with respect to $\\sigma$ and $\\varepsilon$. \n",
    "\n",
    "In general, when we have losses of the form \n",
    "\n",
    "$E_{z \\sim q} h(q)$ it is not so easy to compute their gradient efficiently.\n",
    "\n",
    "You can do quite "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Estimating the likelihood in VAEs\n",
    "\n",
    "Once the model is trained, we can:\n",
    "\n",
    "- Generate samples from it\n",
    "- Do encoding of the input to the latent space\n",
    "- We can also estimate the likelihood of the model by using **importance sampling** technique.\n",
    "\n",
    "$$\n",
    "\\log p_{\\boldsymbol{\\theta}}(\\mathbf{x})=\\log \\mathbb{E}_{q_\\phi(\\mathbf{z} \\mid \\mathbf{x})}\\left[p_{\\boldsymbol{\\theta}}(\\mathbf{x}, \\mathbf{z}) / q_{\\boldsymbol{\\phi}}(\\mathbf{z} \\mid \\mathbf{x})\\right]\n",
    "$$\n",
    "Taking random samples from $q_\\phi(\\mathbf{z} \\mid \\mathbf{x})$, a Monte Carlo estimator of this is:\n",
    "$$\n",
    "\\log p_{\\boldsymbol{\\theta}}(\\mathbf{x}) \\approx \\log \\frac{1}{L} \\sum_{l=1}^L p_{\\boldsymbol{\\theta}}\\left(\\mathbf{x}, \\mathbf{z}^{(l)}\\right) / q_\\phi\\left(\\mathbf{z}^{(l)} \\mid \\mathbf{x}\\right)\n",
    "$$\n",
    "\n",
    "Note that $L=1$ corresponds to the VAE objective (noisy). \n",
    "\n",
    "Better estimates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Neural discrete representation learning\n",
    "\n",
    "We have seen in several places that discrete latent codes could be more preferrable to Gaussian-type latent vectors.\n",
    "\n",
    "The original idea (VQ-VAE) has been proposed in the [paper](https://proceedings.neurips.cc/paper/2017/file/7a98af17e63a0ac09ce2e96d03992fbc-Paper.pdf)\n",
    "\n",
    "For example, for I-BERT we need visual tokens (i.e. discrete representation of each patch).\n",
    "\n",
    "Thus, we need to specify the architecture and the learning process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VQ-VAE\n",
    "\n",
    "In VQ-VAE (vector quantization with VAE) the latent space has defined as $\\mathbb{R}^{K \\times D}$, where\n",
    "\n",
    "$K$ is the size of the discrete latent space (i.e., $K$-way categorical) and $D$ is the dimensionality of each latent embedding vector $e_i$.\n",
    "\n",
    "The posterior categorical distribution $q(z \\mid x)$ probabilities are defined as one-hot as follows:\n",
    "$$\n",
    "q(z=k \\mid x)= \\begin{cases}1 & \\text { for } \\mathrm{k}=\\operatorname{argmin}_j\\left\\|z_e(x)-e_j\\right\\|_2, \\\\ 0 & \\text { otherwise }\\end{cases}\n",
    "$$\n",
    "where $z_e(x)$ is the output of the encoder network. This model is viewed as a VAE in which we can bound $\\log p(x)$ with the ELBO. Our proposal distribution $q(z=k \\mid x)$ is deterministic, and by defining a simple uniform prior over $z$ we obtain a KL divergence constant and equal to $\\log K$.\n",
    "\n",
    "The representation $z_e(x)$ is passed through the discretisation bottleneck followed by mapping onto the nearest element of embedding $e$ as given in equations 1 and 2 .\n",
    "$$\n",
    "z_q(x)=e_k, \\quad \\text { where } \\quad k=\\operatorname{argmin}_j\\left\\|z_e(x)-e_j\\right\\|_2\n",
    "$$\n",
    "\n",
    "**We will be training encoder, decoder and embeddings simulaneously**!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VQ-VAE: challenges and the loss\n",
    "\n",
    "The distribution over categories is assumed in the original paper to be uniform.\n",
    "\n",
    "This, KL-loss is constant an is not involved in optimization, so other losses are needed!\n",
    "\n",
    "Instead, the following loss is used:\n",
    "\n",
    "Thus, the total training objective becomes:\n",
    "$$\n",
    "L=\\log p\\left(x \\mid z_q(x)\\right)+\\left\\|\\operatorname{sg}\\left[z_e(x)\\right]-e\\right\\|_2^2+\\beta\\left\\|z_e(x)-\\operatorname{sg}[e]\\right\\|_2^2,\n",
    "$$\n",
    "\n",
    "where $\\operatorname{sg}$ is a stop-gradient operation (i.e. it is identity in the forward pass, and has zero partial derivatives in the backward pass).\n",
    "\n",
    "The last loss, the commitment loss, which only applies to the encoder weights, encourages\n",
    "the output of the encoder to stay close to the chosen codebook vector to prevent it from fluctuating\n",
    "too frequently from one code vector to another "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VQ-VAE: architecture\n",
    "\n",
    "<img src='vq-vae.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VQ-VAE version 2\n",
    "\n",
    "The improvement of VQ-VAE scheme has been proposed in\n",
    "[Generating Diverse High-Fidelity Images\n",
    "with VQ-VAE-2](https://proceedings.neurips.cc/paper/2019/file/5f8e2fa1718d1bbcadf1cd9c7a54fb8c-Paper.pdf)\n",
    "\n",
    "It uses the same loss, but tries to learn **hierarhical codes**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VQ-VAE version 2: stage 1\n",
    "\n",
    "In the first stage, hierarchical latent codes are used. \n",
    "\n",
    "For example, for 256 × 256 images, a two level latent hierarchy is used. The encoder\n",
    "network first transforms and downsamples the image by a factor of 4 to a 64 × 64 representation\n",
    "which is quantized to our bottom level latent map. Another stack of residual blocks then further\n",
    "scales down the representations by a factor of two, yielding a top-level 32 × 32 latent map after\n",
    "quantization. \n",
    "\n",
    "\n",
    "The decoder is similarly a feed-forward network that takes as input all levels of the\n",
    "quantized latent hierarchy. It consists of a few residual blocks followed by a number of strided\n",
    "transposed convolutions to upsample the representations back to the original image size.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VQ-VAE version 2: stage 2\n",
    "\n",
    "Once we have the encoder, we can model the distribution of the **discrete latent** \n",
    "\n",
    "by a PixelCNN autoregressive architecture (in fact, PixelSnail approach with self-attention).\n",
    "\n",
    "This is done independently.\n",
    "\n",
    "Of course, one can use other discrete distribution models over those latents."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Scheme of the approache\n",
    "<img src='vq-vae2.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Possible improvements over VQ-VAE\n",
    "\n",
    "VQ-VAE architecture:\n",
    "\n",
    "1. Uses uniform prior over discrete latent space \n",
    "2. Separate learning of decoder/endocoder\n",
    "\n",
    "One can easily use GPT to model the dependencies over the resulting latent variables, for example, as in [VideoGPT paper](https://arxiv.org/abs/2104.10157), where GPT is used to predict the discrete latent codes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## VQ-GAN\n",
    "\n",
    "Significant improvement of VQ-VAE has been achieved by combining the ideas of VQ-VAE with the ideas of generative networks (GANs), resulting in the VQ-GAN.\n",
    "\n",
    "In short (we will cover this in the lecture on GANs) is to add an additional loss on the quality of the generated images, and the loss is borrowed from GANs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Comparison with autoregressive models\n",
    "\n",
    "- **PixelCNN:** tractable probability density functions\n",
    "- **VAE:** intractable propability density functions, solved by variational inference.\n",
    "- **VQ-VAE**: strictly speaking, encoder-decoder architecture with codebook lookup."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "\n",
    "## Summary\n",
    "- Autoregressive models \n",
    "- Variational Autoencoders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Next lecture: Generative models II\n",
    "- Generative adversarial networks\n",
    "- DCGAN\n",
    "- WGAN\n",
    "- StyleGAN\n",
    "- Evaluating generative models"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "celltoolbar": "Слайд-шоу",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.0"
  },
  "nav_menu": {},
  "rise": {
   "scroll": true,
   "theme": "sky",
   "transition": "zoom"
  },
  "toc": {
   "navigate_menu": true,
   "number_sections": false,
   "sideBar": true,
   "threshold": 6,
   "toc_cell": false,
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}