Milestone3_draft1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Key derivations for VAEs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1) Maximize Likelihood </br></br>\n",
    "\n",
    "\\begin{aligned}\n",
    "\\left[\n",
    "\\begin{array}\n",
    "pp_{\\theta}(y_n, z_n)\\\\\n",
    "p_{\\theta}(y_n|z_n)p_{\\theta}(z_n)\n",
    "\\end{array}\n",
    "\\right]\\\\\n",
    "\\end{aligned}\n",
    "\n",
    "2) Maximize observed $u$\n",
    "\n",
    "\\begin{aligned}\n",
    "\\theta_{\\text{MLE}} &= \\mathrm{argmax}_{\\theta}\\; \\mathrm{log} \\prod \\overbrace{\\int p_{\\theta}(y_n, z_n) dz_n}^{p_{\\theta}(y_n)}\\\\\n",
    "&= \\mathrm{argmax}_{\\theta}\\; \\nabla_{\\theta} \\sum_n \\mathrm{log}\\; \\mathbb{E}_{p_{\\theta}(z_n)} [p_{\\theta}(y_n, z_n)]\\\\\n",
    "\\end{aligned}\n",
    "However, this presents an issue when we consider the $\\log$ term and the $\\nabla_{\\theta}$ and $\\mathbb{E}_{p_{\\theta}}$, so we use EM steps:\n",
    "</br>  \n",
    "</br>  \n",
    "\\begin{aligned}\n",
    "&= \\mathrm{argmax}_{\\theta}\\; \\sum_n\\; \\mathrm{log} \\int p_{\\theta}(y_n|z_n)\\frac{p_{\\theta}(z_n)}{q(z_n)} q(z_n)d z_n\\\\\n",
    "&= \\mathrm{argmax}_{\\theta}\\; \\sum_n\\; \\mathrm{log}\\; \\mathbb{E}_{q(z_n)}\\; \\left[\\frac{p_{\\theta}(y_n|z_n)p_{\\theta}(z_n)}{q(z_n)}\\right]\\\\\n",
    "\\end{aligned}\n",
    "Apply Jensen's inequality, and we get:</br>  \n",
    "</br>  \n",
    "\\begin{aligned}\n",
    "&\\ge \\mathrm{argmax}_{\\theta}\\; \\underbrace{\\sum_n\\; \\mathbb{E}_{q(z_n)}\\; \\left[\\mathrm{log}\\left(\\frac{p_{\\theta}(y_n,z_n)}{q(z_n)}\\right)\\right]}_{\\text{ELBO}(q, \\theta)}\\\\\n",
    "\\end{aligned}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Maximize the ELBO**\n",
    "\\begin{aligned} \n",
    "\\theta_{\\text{MLE}} &= \\underset{\\theta, q}{\\text{argmax}}\\; \\text{ELBO}(\\theta, q)\\\\\n",
    "\\end{aligned}\n",
    "</br>  \n",
    "*M Step*:\n",
    "\\begin{aligned} \n",
    "\\underset{\\theta}{\\text{max}}\\; \\nabla_{\\theta}\\; \\text{ELBO}(\\theta, q^*)\\\\\n",
    "\\end{aligned}\n",
    "*E Step*:\n",
    "\\begin{aligned} \n",
    "\\underset{q}{\\text{max}}\\; \\text{ELBO}(\\theta^*, q)\\\\\n",
    "\\end{aligned}\n",
    "\n",
    "We use KL divergence to find:\n",
    "\\begin{aligned} \n",
    "q^* &= p_{theta^*}(z_n | y_n)\\\\\n",
    "\\end{aligned}\n",
    "\n",
    "But finding the closed form of this equation is hard -- so we rely on variational inference. We use the Gaussian family with mean field:\n",
    "\n",
    "\\begin{aligned} \n",
    "q(z_n) & = \\mathcal{N}(z_n, \\mu_n, \\Sigma_n)\\\\\n",
    "\\theta_{\\text{MLE}} &= \\underset{\\theta}{\\text{max}}\\; \\nabla_{\\theta}\\; \\text{ELBO}(\\theta, q^*)\\\\\n",
    "\\end{aligned}\n",
    "</br>  \n",
    "*A) Variational Inference*  \n",
    "Step E:\n",
    "\\begin{aligned}\n",
    "\\mu_n^*, \\Sigma_n^* &= \\underset{\\mu, \\sum}{\\text{arming}}\\; \\text{D}_{\\text{KL}} [q(z_n)||p_{\\theta^*}(z_n|y_n)]\\\\\n",
    "&\\equiv \\underbrace{\\underset{\\mu, \\sum}{\\text{argmax}}\\; \\mathbb{E}_{q_{\\mu, \\sum}(z_n)} \\left[\\mathrm{log}\\frac{p_{\\theta^*}(z_n, y_n)}{q_{\\mu_n, \\sum_n}(z_n)}\\right]}_{\\text{ELBO}(\\theta^*,q_{\\mu_n, \\sum_n} )}\\\\\n",
    "\\end{aligned}\n",
    "</br>\n",
    "As you can see, this is the ELBO of $(\\theta^*,q_{\\mu_n, \\sum_n})$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*B) Amortization*  \n",
    "\\begin{aligned}\n",
    "g_{\\phi^*} &= \\mu_{\\phi^*}(y_n), \\Sigma_{\\phi^*}(y_n)\\\\\n",
    "\\phi^* &= \\underset{\\phi}{\\text{argmin}}\\; \\sum_n\\; \\text{D}_{\\text{KL}}\\left[\\text{N}(\\mu_{\\phi^*}(y_n), \\Sigma_{\\phi^*}(y_n))||p_{\\theta^*}(z_n|y_n)\\right]\\\\\n",
    "\\end{aligned}\n",
    "</br>\n",
    "Where the term $\\phi^*$ comes from $g_{\\phi^*}(y_n)$.  \n",
    "</br>  \n",
    "\\begin{aligned}\n",
    "&= \\underset{\\phi}{\\text{argmax}} \\underbrace{ \\sum_n\\; \\mathbb{E}_{q_{\\phi}(z_n)}\\left[ \\mathrm{log} \\frac{p_{\\theta^*}(z_n, y_n)}{q_{\\phi}(z_n)} \\right]}_{\\text{ELBO}(\\theta^*,q_{\\phi} )}\\\\\n",
    "\\end{aligned}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*C) Joint Training -- do the E & M Steps together!*  \n",
    "</br>  \n",
    "\\begin{aligned} \n",
    "\\theta^*, \\phi^* &= \\underset{\\theta, \\phi}{\\text{argmax}}\\;\\text{ELBO}(\\theta, q_\\phi)\n",
    "\\end{aligned}\n",
    "Use gradient descent:  \n",
    "</br>  \n",
    "\\begin{aligned}\n",
    "\\nabla_{\\theta, \\phi} \\; \\text{ELBO}(\\theta, q_\\phi) &= \\sum_n\\; \\nabla_{\\theta, \\phi}\\; \\mathbb{E}_{q_\\phi(z_n)}\\;  \\left[ \\mathrm{log} \\frac{p_{\\theta^*}(z_n, y_n)}{q_\\phi(z_n)}\\right]\\\\\n",
    "\\end{aligned}\n",
    "<hr/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Not sure exactly where this belongs -- sorry!**\n",
    "\n",
    "\\begin{aligned}\n",
    "\\nabla_{\\theta, \\phi} \\; \\text{ELBO}(\\theta, q_\\phi) &= \\sum_n\\; \\nabla_{\\theta, \\phi}\\; \\mathbb{E}_{\\epsilon \\sim \\mathcal{N}(0, I)}\\;  \\left[ \\mathrm{log} \\frac{p_{\\theta}(\\mu_\\phi + \\epsilon\\Sigma_\\phi^{1/2}, y_n)}{q_\\phi(\\epsilon\\Sigma_\\phi^{1/2} + \\mu_\\phi)}\\right]\\\\\n",
    "\\end{aligned}\n",
    "<hr/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Added during Lab*\n",
    "</br>\n",
    "</br>\n",
    "\\begin{aligned}\n",
    "\\nabla_{\\theta, \\phi} \\; \\text{ELBO}(\\theta, q_\\phi) &= \\sum_n\\; \\underbrace{ \\nabla_{\\theta, \\phi}\\; \\mathbb{E}_{q_\\phi(z_n)}\\;   \\mathrm{log} p_{\\theta}(y_n|z_n)}_{\\text{(1)}}- \\underbrace{\\nabla_{\\theta, \\phi}\\text{D}_{\\text{KL}}\\; \\left[q_{\\phi}(z_n)||\\overbrace{p_{\\theta}(z_n)}^{\\text{Normal prior}\\; \\mathrm{N}(0,1)}\\right]}_{\\text{(2)}}\\\\\n",
    "\\\\\n",
    "\\end{aligned}\n",
    "</br>  \n",
    "Consider equation (1):\n",
    "</br>  \n",
    "\\begin{aligned}\n",
    "\\text{(1)}\\;  &= \\nabla_{\\theta, \\phi}\\; \\mathbb{E}\\; \\mathrm{log}\\; p_\\theta(y_n|\\mu_n+\\epsilon\\Sigma_n^{\\frac{1}{2}})\\\\\n",
    "\\end{aligned}\n",
    "</br>  \n",
    "Where we know that:\n",
    "\\begin{aligned}\n",
    "z_n &\\in \\mathbb{R}^J\\\\\n",
    "\\epsilon &\\sim \\mathcal{N}(0,1)\\\\\n",
    "z_n &\\sim \\mathcal{N}(\\mu_n, \\Sigma_n)\\\\\n",
    "y_n = \\theta_1 z_n &\\sim \\mathcal{N}(\\theta_1 z_n, \\Sigma_{data})\\\\\n",
    "\\end{aligned}\n",
    "So, we can say that:\n",
    "</br>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "p(y_n|z_n) &= \\nabla_{\\theta, \\phi} \\underset{\\epsilon\\sim\\mathcal{N}(0,1)}{\\mathbb{E}}\\; \\left( \\frac{-1}{2} \\log \\left( 2 \\pi |\\Sigma_{data}|\\right) - \\frac{1}{2}(y_n - \\hat{y_n})^T \\Sigma_{data}^{-1}(y_n - \\hat{y_n})\\right)\\\\\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "We use a Monte Carlo simulation to estimate this expectation.\n",
    "</br>  \n",
    "And equation (2):\n",
    "</br>  \n",
    "\\begin{aligned}\n",
    "\\text{(2)}\\; &=\\frac{1}{2}\\sum_{j=1}^{J}\\; (1+\\mathrm{log} (\\sigma_{j}^2 - \\mu_j^2 - \\sigma_j^2) \\\\\n",
    "\\end{aligned}\n",
    "</br>  \n"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}