Lab_11_Variational_Inference_PtIV.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOWz5qtHi7RX8m3Ca/1vkxF",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/USCbiostats/PM520/blob/main/Lab_11_Variational_Inference_PtIV.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Get Lucky, or: Non-Conjugate Variational Inference Pt II\n",
        "\n",
        "Recall, that  CAVI derivations assume that our surrogate models are the result of conditional conjugacy between the expected log likelihood and the prior.\n",
        "\n",
        "Last week, we _assumed_ that $Q_j(\\theta_j)$ is in the _same_ exponential family as its corresponding prior $\\Pr(\\theta_j)$. We then derived the analytic expectations required for the ELBO and performed gradient ascent.\n",
        "\n",
        "> _What if there are no closed form/analytic solutions for the expectations (i.e. ELBO)_?\n",
        "\n",
        "One solution to this problem, is to leverage _stochastic_ gradient descent by performing Monte Carlo sampling of the necessary gradients."
      ],
      "metadata": {
        "id": "jrD68Sr315c_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Why is this a problem? Lemme differentiate under the integral!\n",
        "$$\\begin{align*}\n",
        "\\text{ELBO}(\\theta) &:= \\mathbb{E}_Q\\left[\\log \\Pr(\\mathbf{X} | \\mathbf{z})\\right] - \\mathsf{D}_{KL}(Q(\\mathbf{z})||P(\\mathbf{z}))\\\\\n",
        "\\nabla_{\\theta}\\text{ELBO}(\\theta) &=\n",
        "  \\nabla_{\\theta}\\mathbb{E}_Q\\left[\\log \\Pr(\\mathbf{X} | \\mathbf{z})\\right]\n",
        "  -\\underbrace{\\nabla_{\\theta}\\mathsf{D}_{KL}(Q(\\mathbf{z})||P(\\mathbf{z}))}_{\\text{typically analytically tractable!}}\\\\\n",
        "\\nabla_{\\theta}\\mathbb{E}_Q\\left[\\log \\Pr(\\mathbf{X} | \\mathbf{z})\\right] &= \\nabla_{\\theta}\\int Q_{\\theta}(\\mathbf{z})\\log \\Pr(\\mathbf{X} | \\mathbf{z}) d\\mathbf{z}\\\\\n",
        "  &= \\int \\nabla_{\\theta}Q_{\\theta}(\\mathbf{z})\\log \\Pr(\\mathbf{X}|\\mathbf{z})d \\mathbf{z} \\\\\n",
        "  &\\neq \\mathbb{E}_{Q}\\left[\\nabla_{\\theta} \\log \\Pr(\\mathbf{X} | \\mathbf{z})\\right]\n",
        "\\end{align*}$$\n",
        "\n",
        "\n",
        "Our expectation depends on the parameters $\\theta$, which complicates our expression, and typically doesn't result in a known closed form solution. Additionally, by the straightforward derivation above, we see that we can't sample gradients directly as it isn't the same expression!\n",
        "\n",
        "> _What gives?_\n",
        "\n",
        "## Reparameterization Trick\n",
        "If we can define $\\mathbf{z}$ as a _deterministic_ function $g_{\\theta}(ɛ) \\mapsto \\mathbf{z}$, then we may be able to circumvent this issue. For _location-scale_ families $f$, this is trivial! Namely, $\\mathbf{z} = \\mu + \\sigma \\circ ɛ$, where $ɛ_j \\sim f(0, 1)$, $\\theta = \\{\\mu, \\sigma\\}$. Now we have gone from $\\mathbb{E}_{Q}\\left[ \\log \\Pr(\\mathbf{X} | \\mathbf{z}) \\right]$ to $\\mathbb{E}_{ɛ \\sim f(0, 1)}\\left[ \\log \\Pr(\\mathbf{X} | g_{\\theta}(ɛ))\\right]$.\n",
        "\n",
        "> _How does this help us?_\n",
        "\n",
        "We can use Monte-Carlo estimates of the gradient under this reparameterization to approximate the exact gradient. This can be\n",
        "shown by,\n",
        "$$\\begin{align*}\n",
        "\\nabla_{\\theta}\\mathbb{E}_Q\\left[\\log \\Pr(\\mathbf{X} | \\mathbf{z})\\right] &=\n",
        "  \\nabla_{\\theta}\\mathbb{E}_{ɛ \\sim f(0, 1)}\\left[ \\log \\Pr(\\mathbf{X} | g_{\\theta}(ɛ))\\right] \\\\\n",
        "&= \\nabla_{\\theta} \\int f(ɛ) \\log \\Pr(\\mathbf{X} | g_{\\theta}(ɛ)) dɛ\\\\\n",
        "&= \\int f(ɛ) \\nabla_{\\theta}\\log \\Pr(\\mathbf{X} | g_{\\theta}(ɛ)) dɛ\\\\\n",
        "&= \\mathbb{E}_{ɛ \\sim f(0, 1)}\\left[\\nabla_{\\theta}\\log \\Pr(\\mathbf{X} | g_{\\theta}(ɛ)) \\right] \\\\\n",
        "&\\approx \\dfrac{1}{T} \\sum_{t=1}^T \\nabla_{\\theta}\\log \\Pr(\\mathbf{X} | g_{\\theta}(ɛ^t)),\n",
        "\\end{align*}$$\n",
        "where $ɛ^t \\sim f(0, 1)$.\n"
      ],
      "metadata": {
        "id": "pMTNVVCP8Cnu"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Lab: Logistic Regression with Normal priors on effects\n",
        "We would like to perform variational inference under the following model:\n",
        "$$\\begin{align*}\n",
        "\\mathbf{y}_i | \\beta &\\sim \\text{Bernoulli}(\\text{sigmoid}(\\mathbf{x}_i^T \\beta))\\\\\n",
        "\\beta_j &\\sim N(0, \\sigma^2_b).\n",
        "\\end{align*}$$\n",
        "\n",
        "\n",
        "Let's assume $Q(\\beta) = \\prod_{j=1}^p Q_j(\\beta_j) = N(\\beta_j | \\mu_j, \\sigma^2_j)$. This model doesn't exhibit (any known) closed-form expectations, so we must rely on stochastic optimization. Let's code the reparameterization trick to optimize the ELBO for this model."
      ],
      "metadata": {
        "id": "lu-DHtx0Wku7"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "BiGZ0RfieKCk"
      },
      "outputs": [],
      "source": [
        "import operator as op\n",
        "\n",
        "from typing import NamedTuple\n",
        "\n",
        "import jax\n",
        "import jax.random as rdm\n",
        "import jax.nn as nn\n",
        "import jax.numpy as jnp\n",
        "import jax.scipy as jsp\n",
        "\n",
        "from jax import Array\n",
        "from jax.tree_util import tree_map, tree_reduce\n",
        "from jax.typing import ArrayLike\n",
        "\n",
        "from jax import config\n",
        "config.update(\"jax_enable_x64\", True)\n",
        "\n",
        "\n",
        "class PriorParams(NamedTuple):\n",
        "  mean_b: Array\n",
        "  var_b: Array\n",
        "\n",
        "\n",
        "class PosteriorParams(NamedTuple):\n",
        "  mean_b: Array\n",
        "  log_var_b: Array\n",
        "\n",
        "\n",
        "def _kl_divergence(post: PosteriorParams, prior: PriorParams) -> Array:\n",
        "  var_b = jnp.exp(post.log_var_b)\n",
        "\n",
        "  term1 = ((post.mean_b - prior.mean_b) ** 2) / prior.var_b\n",
        "  term2 = var_b / prior.var_b\n",
        "  term3 = post.log_var_b - jnp.log(prior.var_b)\n",
        "\n",
        "  return 0.5 * (term1 + term2 - term3 - 1)\n",
        "\n",
        "\n",
        "def kl_divergence(post: PosteriorParams, prior: PriorParams) -> float:\n",
        "  \"\"\"KL divergence for scalar normals\n",
        "  \"\"\"\n",
        "  return jnp.sum(_kl_divergence(post, prior))\n",
        "\n",
        "\n",
        "def complete_log_likelihood(y: ArrayLike, X: ArrayLike, beta: Array) -> float:\n",
        "  \"\"\" log-likelihood of binary outcomes y, given X, beta\n",
        "  \"\"\"\n",
        "  eps = jnp.finfo(beta.dtype).eps\n",
        "  probs = jnp.clip(nn.sigmoid(X @ beta),\n",
        "                   eps,\n",
        "                   1. - eps,\n",
        "  )\n",
        "  return jnp.sum(jsp.stats.bernoulli.logpmf(y, probs))\n",
        "\n",
        "\n",
        "def sample_log_likelihood(post: PosteriorParams, y: ArrayLike, X: ArrayLike, key: ArrayLike) -> float:\n",
        "  n, p = X.shape\n",
        "  # reparameterization trick\n",
        "  std_dev_b = jnp.exp(0.5 * post.log_var_b)\n",
        "  beta = post.mean_b + std_dev_b * rdm.normal(key, shape=(p,))\n",
        "  return complete_log_likelihood(y, X, beta)\n",
        "\n",
        "\n",
        "def elbo(post: PosteriorParams, prior: PriorParams, y: ArrayLike, X: ArrayLike, key: ArrayLike) -> float:\n",
        "  e_ll = sample_log_likelihood(post, y, X, key)\n",
        "  kl = kl_divergence(post, prior)\n",
        "  return e_ll - kl\n",
        "\n",
        "\n",
        "def fit(y: ArrayLike, X: ArrayLike, prior_var_b: float = 1e-3, num_samples: int = 5, step_size = 1e-3, seed = 0, max_iter=100) -> PosteriorParams:\n",
        "  #initialize our random keye\n",
        "  n, p = X.shape\n",
        "  key = rdm.PRNGKey(seed)\n",
        "\n",
        "  # split to initalize our variational parameters\n",
        "  key, init_mean_key = rdm.split(key, 2)\n",
        "  post = PosteriorParams(\n",
        "      mean_b = jnp.sqrt(prior_var_b) * rdm.normal(init_mean_key, (p,)),\n",
        "      log_var_b = jnp.log(prior_var_b * jnp.ones((p,))),\n",
        "  )\n",
        "  prior = PriorParams(\n",
        "      mean_b = jnp.zeros((p,)),\n",
        "      var_b = prior_var_b * jnp.ones((p,)),\n",
        "  )\n",
        "\n",
        "  # sample gradient using reparam trick for expected log like\n",
        "  def _step(post, key):\n",
        "    elboval, elboggrad = jax.value_and_grad(elbo)(post, prior, y, X, key)\n",
        "\n",
        "    return elboval, elboggrad\n",
        "\n",
        "  for epoch in range(max_iter):\n",
        "    key, *skey = rdm.split(key, num_samples + 1)\n",
        "    skey = jnp.array(skey)\n",
        "    evals, grads = jax.vmap(_step, (None, 0,))(post, skey)\n",
        "    elboval = jnp.mean(evals)\n",
        "    grad = PosteriorParams(mean_b = jnp.mean(grads.mean_b, axis=0),\n",
        "                           log_var_b = jnp.mean(grads.log_var_b, axis=0))\n",
        "\n",
        "    print(f\"ELBO[{epoch}] ≈ {elboval}\")\n",
        "    post = tree_map(lambda _post, _grad: _post + step_size * _grad, post, grad)\n",
        "    #print(f\"params = {post}\")\n",
        "    # sample to -compute/evaluate- the ELBO\n",
        "\n",
        "  return post"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# simulate binary outcome\n",
        "N, P = 250, 25\n",
        "prior_var_b = 1e-2\n",
        "\n",
        "seed = 0\n",
        "key = rdm.PRNGKey(seed)\n",
        "key, beta_key, x_key, y_key = rdm.split(key, 4)\n",
        "\n",
        "beta = jnp.sqrt(prior_var_b) * rdm.normal(beta_key, shape=(P,))\n",
        "X = rdm.normal(x_key, (N, P))\n",
        "\n",
        "pred = X @ beta\n",
        "prob = nn.sigmoid(pred)\n",
        "y = rdm.bernoulli(y_key, prob)\n",
        "\n",
        "params = fit(y, X, step_size=0.01, prior_var_b=prior_var_b, num_samples=10)"
      ],
      "metadata": {
        "id": "QbCzlIWRP6dU"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}