02_variational_inference_intuition.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cae4e954-9d88-4946-977a-9fcbec485b77",
   "metadata": {},
   "source": [
    "# Introduction to variational inference\n",
    "\n",
    "Given a probabilistic model where the joint distribution of latent states $x$ and observation $y$ is given by,\n",
    "$ p(y,x) = p(y|x)p(x) $, we would like to infer the posterior distribution,\n",
    "$$ p(x|y) = \\frac{p(x,y)}{p(y)} \\propto p(y|x)p(x) $$\n",
    "\n",
    "Let's use a 1-D grid to evaluate these distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2ea739a-1c07-4fd7-8755-313652068c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "\n",
    "xr = np.linspace(-10., 10., 201)\n",
    "dx = xr[1] - xr[0]\n",
    "\n",
    "likelihood = np.exp(-(xr - 2)**2/6)\n",
    "likelihood /= sum(likelihood) * dx\n",
    "\n",
    "#prior = np.array(xr < 0).astype(float) + np.array(xr >= 0).astype(float) * 0.1\n",
    "prior = np.exp(-np.abs(xr)/3)\n",
    "prior /= sum(prior) * dx\n",
    "\n",
    "# Bayes Rule - note that this approximate inference on a grid doesn't scale to high-dim\n",
    "posterior = prior * likelihood\n",
    "posterior /= sum(posterior) * dx\n",
    "\n",
    "plt.plot(xr, likelihood, label=\"p(y|x) likelihood\")\n",
    "plt.plot(xr, prior, label=\"p(x) prior dist.\")\n",
    "plt.plot(xr, posterior, '--', label=\"p(x|y) posterior dist.\")\n",
    "plt.xlabel(\"x\"); plt.ylabel(\"probability density\"); plt.legend(); plt.grid();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52986eb3-108a-4e1e-a8ce-09e3891a3ee7",
   "metadata": {},
   "source": [
    "This problem is intractible for most likelihood $p(y|x)$ and prior $p(x)$ unless they are *conjugate*, a rare mathematically convienient coincidence.\n",
    "Therefore, we seek an approximate inference method.\n",
    "\n",
    "We assume that we can evaluate $p(x)$ and $p(y|x)$ and also sample from both distributions, computationally efficiently.\n",
    "We will use PyTorch to represent these distributions, so that we can autodifferentiate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2676018-225a-4114-ab4d-f353fb0def09",
   "metadata": {},
   "outputs": [],
   "source": [
    "lik = torch.distributions.normal.Normal(torch.tensor([2.0]), torch.tensor([np.sqrt(3)]))\n",
    "pri = torch.distributions.laplace.Laplace(torch.tensor([0.0]), torch.tensor([3.0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c35ac74-3881-4b87-bf00-840147d574cd",
   "metadata": {},
   "source": [
    "### Step 1: choose a parametric family of distributions\n",
    "\n",
    "Let's choose a parametric family of distributions to approximate the desired posterior distribution $p(x|y)$.\n",
    "We denote the approximate distribution with\n",
    "$$ q(x;\\phi) \\approx p(x|y) $$\n",
    "where the goal of the inference is to find the \"best\" parameter $\\phi \\in \\Phi$.\n",
    "We assume that it is easy to evaluate and sample from $q$, entropy of $q$ is differentiable with respect to $\\phi$ and is easy to compute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2934fd80-1867-4360-846e-f43958e5a698",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = torch.tensor([0.0],requires_grad=True); sigma = torch.tensor([5.0],requires_grad=True)\n",
    "q = torch.distributions.normal.Normal(mu, sigma)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ab94f9c-c59e-47c3-8b0d-4f4e0f24f1c7",
   "metadata": {},
   "source": [
    "## Step 2: define the loss function\n",
    "Variational inference turns the inference problem into an optimization problem by defining the loss function that measures the quality of approximation based on a divergence measure.\n",
    "A divergence measure $d(p,q)$ is non-negative, and returns 0 if and only if the two distributions are identical.\n",
    "A typical choice for variational inference is the Kullback-Leibler (KL) divergence:\n",
    "$$ d_{\\text{KL}}(p || q) = \\int \\log\\left(\\frac{dp}{dq}\\right) dp $$\n",
    "KL is a central quantity in Shannon's information theory in measuring the amount of bits wasted in compression by using $q$ where the true distribution is $p$.\n",
    "\n",
    "KL works particularly nicely with exponential family distributions and hence plays a key role in Amari's information geometry.\n",
    "\n",
    "KL is not symmetric, and variational inference uses the following loss function:\n",
    "$$\n",
    "ELBO(\\phi) = -d_{\\text{KL}}(q(x;\\phi) || p(x|y)) \\\\\n",
    "= -E_q[ \\log(p(x|y)) ] + H(q(x;\\phi)) \\\\\n",
    "= -E_q[ \\log(p(y|x)) ] - E_q[ \\log(p(x)) ] + E_q[ \\log(p(y)) ]+ H(q(x;\\phi))\n",
    "$$\n",
    "where $H(\\cdot)$ denotes the entropy.\n",
    "\n",
    "Since $E_q[ \\log(p(y)) ]$ is constant, we can drop it for the optimization.\n",
    "$$\n",
    "ELBO'(\\phi) \\approx \\frac{1}{n} \\sum_i [ - \\log(p( x_i | y )) - \\log(p(x_i))] + H(q(x;\\phi))\n",
    "$$\n",
    "where $x_i$ are $n$ independent samples from $q(x;\\phi)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89437094-7506-4f03-9976-0430923f1418",
   "metadata": {},
   "outputs": [],
   "source": [
    "nMC = 100\n",
    "X = q.sample(torch.Size([nMC])) # <-- this doesn't propagate gradients; each realization is a constant, not a function of phi\n",
    "ELBO = -torch.mean(lik.log_prob(X) + pri.log_prob(X)) + q.entropy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cc184bd-ad47-4361-8eb1-92672e4bde82",
   "metadata": {},
   "source": [
    "## Step 3: use the reparameterization trick\n",
    "\n",
    "Once a realization is drawn from a distribution, they no longer depend on the parameters.\n",
    "Therefore, the above Monte Carlo approximation of the expectation doesn't work.\n",
    "However, if the parameters are location and/or scale parameters, we could simply shift and scale the samples, and making them differentiable.\n",
    "This is the reparameterization trick. Fortunately, PyTorch has implemented samplers for many common distributions, so that differentiable samples may be drawn from a `torch.Distribution` object using `.rsample()`\n",
    "\n",
    "There are a number of variational distributions that allow the reparameterization trick.\n",
    "\n",
    "- Rezende, D. J., Mohamed, S., & Wierstra, D. (2014, May 30). Stochastic Backpropagation and Approximate Inference in Deep Generative Models. International Conference on Machine Learning. http://jmlr.org/proceedings/papers/v32/rezende14.html\n",
    "- Kingma, D. P., & Welling, M. (2014, May 1). Auto-Encoding Variational Bayes. International Conference on Learning Representation. http://arxiv.org/abs/1312.6114"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7cef74b-dfdd-496c-ad47-61bd36fbd558",
   "metadata": {},
   "outputs": [],
   "source": [
    "sn = torch.distributions.normal.Normal(torch.tensor([0.0]), torch.tensor([1.0]))\n",
    "X = (sn.sample(torch.Size([nMC])) + mu) * sigma # <-- reparametrization trick (only works for certain distributions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a87de282-4830-4087-968a-1cd900d00a2a",
   "metadata": {},
   "source": [
    "## Step 4: maximize the ELBO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ebf66e3-1748-485e-802e-606f8ee6f1b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.SGD([mu, sigma], lr=1e-3) # you could try Adam, if you want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c9a234-7767-48af-baf5-09ff776183af",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import trange, tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f76e6fe-055d-4d41-976a-2359e4d8c2db",
   "metadata": {},
   "outputs": [],
   "source": [
    "ELBO_trace = []\n",
    "for k in trange(10000):\n",
    "    X = (sn.sample(torch.Size([nMC])) + mu) * sigma # <-- reparametrization trick (only works for certain distributions)\n",
    "    nELBO = -torch.mean(lik.log_prob(X) + pri.log_prob(X)) - q.entropy() # negative of the ELBO to be minimized\n",
    "    ELBO_trace.append(-nELBO.item())\n",
    "    optimizer.zero_grad()\n",
    "    nELBO.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35f1837c-446c-4cdf-9f3e-92f995595af6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(ELBO_trace);\n",
    "plt.title(\"convergence\"); plt.ylabel(\"ELBO\"); plt.xlabel(\"gradient steps\"); plt.grid();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00996425-7b16-4aea-91f8-019e17357b9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "q_plot = torch.distributions.normal.Normal(mu.detach(), sigma.detach())\n",
    "xrt = torch.tensor(xr)\n",
    "plt.plot(xr, np.exp(lik.log_prob(xrt).numpy()), label=\"likelihood\")\n",
    "plt.plot(xr, np.exp(pri.log_prob(xrt).numpy()), label=\"prior\")\n",
    "plt.plot(xr, posterior, '--', label=\"true posterior\")\n",
    "plt.plot(xr, np.exp(q_plot.log_prob(xrt).numpy()), label=\"variational posterior\")\n",
    "plt.legend(); plt.grid();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5a5911b-aa53-4792-9947-5508e2f98f93",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}