Lab_5_Optimization_PtII.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNY3IIxV7h3rwzUyCu6vPEP",
      "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_5_Optimization_PtII.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Ain't no mountain high enough, or: Optimization Pt II"
      ],
      "metadata": {
        "id": "D7LCHBjOhm0i"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qP4Zigt0zfMW"
      },
      "outputs": [],
      "source": [
        "!pip install lineax"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Newton's Method for Optimization\n",
        "Recall under [gradient descent](https://en.wikipedia.org/wiki/Gradient_descent) we can iteratively optimize a function $f(\\beta)$ by taking steps in the steepest direction,\n",
        "$$ \\hat{\\beta} = \\beta_t - \\rho_t \\nabla f(\\beta_t).$$\n",
        "\n",
        "But, can we do better, by considering higher-order information (ie geometry) of\n",
        "the function $f$?\n",
        "\n",
        "Let's consider a 2nd-order [Taylor-series approximation](https://en.wikipedia.org/wiki/Taylor_series) to $f$ around $\\beta_t$ as,\n",
        "\n",
        "$$f(\\beta) \\approx f(\\beta_t) + \\nabla f(\\beta_t)^T (\\beta - \\beta_t) + \\frac{1}{2} (\\beta - \\beta_t)^T H(\\beta_t)(\\beta - \\beta_t),$$ where $H(\\beta_t) = \\nabla^2 f(\\beta_t)$ (i.e. the [Hessian](https://en.wikipedia.org/wiki/Hessian_matrix) of $f$ at $\\beta_t$). If we minimize this _local_ approximation, we see\n",
        "\n",
        "$\\nabla_\\beta f(\\beta) \\approx \\nabla f(\\beta_t) + H(\\beta_t)(\\beta - \\beta_t) = \\nabla f(\\beta_t) + H(\\beta_t)\\beta - H(\\beta_t)\\beta_t ⇒$\n",
        "$$ H(\\beta_t)\\beta = H(\\beta_t)\\beta_t - \\nabla f(\\beta_t).$$\n",
        "\n",
        "We can recognize that this is a [system of linear equations](https://en.wikipedia.org/wiki/System_of_linear_equations) $A x = b$ where $A = H(\\beta_t)$, $x = \\beta$, and $b = H(\\beta_t)\\beta_t - \\nabla f(\\beta_t)$. The solution is given by, $\\hat{x} = A^{-1}b$, which in this case implies,\n",
        "$$ \\hat{\\beta} = H(\\beta_t)^{-1}\\left(H(\\beta_t)\\beta_t - \\nabla f(\\beta_t)\\right) = \\beta_t - H(\\beta_t)^{-1}\\nabla f(\\beta_t).$$\n",
        "\n",
        "\n",
        "\n",
        "[Newton's method](https://en.wikipedia.org/wiki/Newton%27s_method_in_optimization) is only guaranteed to converge _locally_, and can diverge even for _strongly_ [convex functions](https://en.wikipedia.org/wiki/Convex_function) (e.g., $f(\\beta) = \\sqrt{\\beta^2 + 1}$). To address this limitation, we can add a dampening parameter, $\\rho_t$, which gives us our final update form,\n",
        "$$ \\hat{\\beta} = H(\\beta_t)^{-1}(H(\\beta_t)\\beta_t - \\nabla f(\\beta_t)) = \\beta_t - \\rho_t H(\\beta_t)^{-1}\\nabla f(\\beta_t).$$\n",
        "\n",
        "## Quasi-Newton Methods for Optimization\n",
        "What if computing $H(\\beta_t)$ is prohibitive or too costly? Do we need _exact_ second order information to improve on gradient descent's convergence? Given an approximation of $H$, called $B$, i.e. $B(\\beta_t) \\approx H(\\beta_t)$, [_quasi_-Newton methods](https://en.wikipedia.org/wiki/Quasi-Newton_method) optimize for the form\n",
        "$$f(\\beta) \\approx f(\\beta_t) + \\nabla f(\\beta_t)^T (\\beta - \\beta_t) + \\frac{1}{2} (\\beta - \\beta_t)^T B(\\beta_t)(\\beta - \\beta_t),$$ where $B(\\beta_t) \\approx H(\\beta_t)$. Optimizing this statement gives us our update rule,\n",
        "$$ \\hat{\\beta} = \\beta_t - \\rho_t B(\\beta_t)^{-1}\\nabla f(\\beta_t).$$"
      ],
      "metadata": {
        "id": "_b5I2Q37z4_1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Poisson Regression\n",
        "\n",
        "$$y_i | x_i \\sim \\text{Poi}(\\lambda_i)$$ where $\\lambda_i := \\exp(x_i^T \\beta)$, and $\\text{Poi}(k | \\lambda) := \\frac{\\lambda^k \\exp(-\\lambda)}{k!}$ is the [PMF](https://en.wikipedia.org/wiki/Probability_mass_function) of the [Poisson distribution](https://en.wikipedia.org/wiki/Poisson_distribution). Given $\\{(y_i, x_i)\\}_{i=1}^n$, we would like to identify the [maximum likelihood parameter estimate](https://en.wikipedia.org/wiki/Maximum_likelihood_estimation) for $\\beta$. In other words, we would to find a value for $\\beta$ such that we maximize the log-likelihood given by,\n",
        "$$\\begin{align*}\n",
        "\\log \\ell(\\beta) &= \\sum_i \\log \\text{Poi}(y_i | \\exp(x_i^T \\beta)) \\\\\n",
        "&= \\sum_i \\log \\left[ \\frac{\\exp(y_i \\cdot x_i^T \\beta) \\exp(-\\exp(x_i^T \\beta))}{y_i!} \\right] \\\\\n",
        "&= \\sum_i \\log \\left[ \\frac{\\exp(y_i \\cdot x_i^T \\beta - \\exp(x_i^T \\beta))}{y_i!} \\right] \\\\\n",
        "&= \\sum_i \\log \\left[\\exp(y_i \\cdot x_i^T \\beta - \\exp(x_i^T \\beta))\\right] - \\log(y_i!) \\\\\n",
        "&= \\sum_i \\left[y_i \\cdot x_i^T \\beta - \\exp(x_i^T \\beta) - \\log(y_i!)\\right] \\\\\n",
        "&= y^T X\\beta - \\exp(X\\beta)^T 1_n - O(1) \\\\\n",
        "&= y^T X\\beta - \\lambda^T 1_n - O(1),\n",
        "\\end{align*}$$\n",
        "where $\\lambda = \\{\\lambda_1, \\dotsc, \\lambda_n\\}.$\n",
        "\n",
        "\n",
        "$$ \\begin{align*}\n",
        "\\nabla_\\beta \\ell &= \\nabla_\\beta \\left[ y^T X\\beta - \\lambda^T 1_n \\right] \\\\\n",
        "&= \\nabla_\\beta [ y^T X\\beta ] - \\nabla_\\beta [\\lambda^T 1_n] \\\\\n",
        "&= \\nabla_\\beta [ y^T X\\beta ] - \\nabla_\\beta [\\exp(X\\beta)^T 1_n] \\\\\n",
        "&= X^T y - \\exp(X\\beta)^T X  \\\\\n",
        "&= X^T y - \\lambda^T X  \\\\\n",
        "&= X^T(y - \\lambda) \\\\\n",
        "\\nabla^2_{\\beta \\beta} \\ell &= \\nabla_{\\beta} X^T(y - \\lambda) \\\\\n",
        "&= \\nabla_{\\beta} \\left[X^T y - X^T \\lambda \\right] \\\\\n",
        "&= - X^T \\nabla_{\\beta}  \\lambda \\\\\n",
        "&= -X^T \\nabla_{\\beta}  \\exp(X\\beta) \\\\\n",
        "&= -X^T \\Lambda X,\n",
        "\\end{align*}$$\n",
        "where $\\Lambda = \\text{diag}(\\lambda)$, i.e. $\\Lambda_{ii} = \\lambda_i$ and $\\Lambda_{ij} = 0$ for $i \\neq j$.\n",
        "\n",
        "To illustrate how $\\nabla_{\\beta}  \\exp(X\\beta) = \\Lambda X$ (i.e. last step in Hessian calculation), recall that the [Jacobian](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) of a function $f : \\mathbb{R}^n → \\mathbb{R}^m$ is the $m \\times n$ matrix $J$ such that $J_{ij} = \\frac{∂f_i}{∂j}$. In this case we are computing the Jacobian for $\\exp(X\\beta)$, which is $\\mathbb{R}^p → \\mathbb{R}^n$, so our final Jacobian for $\\exp(X\\beta)$ should have shape $n \\times p$. Notice that $J_{i,j} = \\frac{\\partial}{\\partial \\beta_j} \\exp(x_i^T \\beta) = x_{ij}\\exp(x_i^T \\beta)$, thus $J_{i, .} = \\exp(x_i^T \\beta) x_i^T$. Repeating this for each $i$ we have $$∇_\\beta \\exp(X \\beta) = J(\\exp(X \\beta)) = \\begin{bmatrix} J_{1,.} \\\\ ⋮ \\\\ J_{n,.} \\end{bmatrix} =\n",
        "\\begin{bmatrix} \\exp(x_1^T \\beta) x_1^T \\\\ ⋮ \\\\ \\exp(x_n^T \\beta) x_n^T \\end{bmatrix}  =\n",
        "\\begin{bmatrix} \\lambda_1 x_1^T \\\\ ⋮ \\\\ \\lambda_n x_n^T\\end{bmatrix} = \\Lambda X.$$\n",
        "\n",
        "We can fit using Newton's method. =>\n",
        "$$\\begin{align*}\n",
        "\\beta(t+1) &= \\beta(t) - H(\\beta(t))^{-1}\\nabla \\ell(\\beta_t) \\\\\n",
        "&= \\beta(t) + (X^T \\Lambda(t) X)^{-1} X^T (y - \\lambda) ⇒ \\\\\n",
        "&= (X^T \\Lambda(t) X)^{-1} X^T \\Lambda(t) (\\Lambda(t)^{-1}y + X\\beta(t) - 1)\n",
        "\\end{align*}$$\n",
        "where $\\Lambda(t) := \\text{diag}(\\lambda_1, \\dotsc, \\lambda_n)$."
      ],
      "metadata": {
        "id": "crQSiZATzm9D"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.random as rdm\n",
        "import jax.scipy.stats as stats\n",
        "\n",
        "import lineax as lx\n",
        "\n",
        "@jax.jit\n",
        "def loglikelihood(beta, y, X):\n",
        "  \"\"\"\n",
        "  Our loglikelihood function for $y_i | x_i ~ \\text{Poi}(\\exp(eta_i))$.\n",
        "\n",
        "  beta: beta\n",
        "  y: poisson-distributed observations\n",
        "  X: our design matrix\n",
        "\n",
        "  returns: sum of the logliklihoods of each sample\n",
        "  \"\"\"\n",
        "  mean_lambda = jnp.exp(X.mv(beta))\n",
        "  return jnp.sum(stats.poisson.logpmf(y, mean_lambda))\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def irwls_fit(beta, y, X, step_size):\n",
        "  \"\"\"\n",
        "  Perform MLE estimation for $\\beta$ under the model\n",
        "     $y_i | x_i ~ \\text{Poi}(\\exp(x_i^T \\beta))$.\n",
        "\n",
        "  beta: beta\n",
        "  y: poisson-distributed observations\n",
        "  X: our design matrix\n",
        "\n",
        "  returns: updated estimate of $\\beta$\n",
        "  \"\"\"\n",
        "  # compute lambda_i := exp(x_i @ beta)\n",
        "  eta = X.mv(beta)\n",
        "  d_i = jnp.exp(eta)\n",
        "  d_sqrt = jnp.sqrt(d_i)\n",
        "\n",
        "  # compute z_i := Lambda^{1/2}(Lambda^-1 y + X @beta - 1)\n",
        "  z = (y / d_i + eta - 1) * d_sqrt\n",
        "\n",
        "  # X* := Lambda^{1/2} X\n",
        "  # we use linear operators to postpone any computation\n",
        "  X_star = lx.DiagonalLinearOperator(d_sqrt) @ X\n",
        "\n",
        "  # lineax can solve normal equations iteratively as (t(X*) @ (X* @ guess)) - z\n",
        "  solution = lx.linear_solve(X_star, z, solver=lx.NormalCG(atol=1e-4, rtol=1e-3))\n",
        "  beta = solution.value\n",
        "  return beta\n",
        "\n",
        "\n",
        "def poiss_reg(y, X, fit_func, step_size = 1.0, max_iter=100, tol=1e-3):\n",
        "  \"\"\"\n",
        "  Perform MLE estimation for $\\beta$ under the model\n",
        "     $y_i | x_i ~ \\text{Poi}(\\exp(x_i^T \\beta))$.\n",
        "\n",
        "  y: poisson-distributed observations\n",
        "  X: our design matrix\n",
        "  max_iter: the maximum number of iterations to perform optimization\n",
        "  tol:\n",
        "\n",
        "  returns: updated estimate of $\\beta$\n",
        "  \"\"\"\n",
        "  # intialize eta := X @ beta\n",
        "  n, p = X.shape\n",
        "\n",
        "  # fake bookkeeping\n",
        "  loglike = -100000\n",
        "  delta = 10000\n",
        "\n",
        "  # convert to a linear operator for lineax\n",
        "  X = lx.MatrixLinearOperator(X)\n",
        "\n",
        "  # initialize using OLS estimate\n",
        "  sol = lx.linear_solve(X, (y - jnp.mean(y))/2, solver=lx.NormalCG(atol=1e-4, rtol=1e-3))\n",
        "  beta = sol.value\n",
        "  for epoch in range(max_iter):\n",
        "\n",
        "    # fit using our function\n",
        "    beta = fit_func(beta, y, X, step_size)\n",
        "\n",
        "    # evaluate log likelihood\n",
        "    newll = loglikelihood(beta, y, X)\n",
        "\n",
        "    # take delta and check if we can stop\n",
        "    delta = jnp.fabs(newll - loglike)\n",
        "    print(f\"Epoch[{epoch}] = {newll}\")\n",
        "    if delta < tol:\n",
        "      break\n",
        "\n",
        "    # replace old value\n",
        "    loglike = newll\n",
        "\n",
        "  return beta"
      ],
      "metadata": {
        "id": "DKK1oqZMztkU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Let's simulate a poisson regression model with N samples and P variables\n",
        "N = 1000\n",
        "P = 5\n",
        "\n",
        "# initialize PRNG env\n",
        "seed = 0\n",
        "key = rdm.PRNGKey(seed)\n",
        "\n",
        "# split key for each random call\n",
        "key, y_key, x_key, b_key = rdm.split(key, 4)\n",
        "X = rdm.normal(x_key, shape=(N, P))\n",
        "beta = rdm.normal(b_key, shape=(P,))\n",
        "\n",
        "# compute lambda_i\n",
        "mean_lambda = jnp.exp(X @ beta)\n",
        "\n",
        "# sample y from Poi(lambda_i)\n",
        "y = rdm.poisson(y_key, mean_lambda)\n",
        "\n",
        "# estimate beta using our irwls function\n",
        "# fit_func has signature (eta, y, X, step_size)\n",
        "beta_hat = poiss_reg(y, X, irwls_fit)\n",
        "print(f\"beta = {beta}\")\n",
        "print(f\"hat(beta) = {beta_hat}\")"
      ],
      "metadata": {
        "id": "nRai4qWiz1_R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# let's implement poisson regression using _only_ gradient informatino to perform inference\n",
        "# and measure how quickly it converges compared with the Newton method\n",
        "def grad_fit(beta, y, X, step_size):\n",
        "  eta = X.mv(beta)\n",
        "  mean_lambda = jnp.exp(eta)\n",
        "  grad = (X.transpose().mv(y - mean_lambda))\n",
        "  beta = beta + step_size * grad\n",
        "  return beta\n",
        "\n",
        "# NB: we can transpose a lx.MatrixLinearOperator (say X) as X.transpose()\n",
        "# NB: we compute matrix-vector produces using a lx.MatrixLinearOperator as X.mv(v)\n",
        "step_size = 1e-7\n",
        "beta_hat = poiss_reg(y, X, grad_fit, step_size, max_iter=1000)\n",
        "print(f\"beta = {beta}\")\n",
        "print(f\"hat(beta) = {beta_hat}\")"
      ],
      "metadata": {
        "id": "ZalxS2NOOfiV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Automatic differentiation\n",
        "Chain rules, okay! Notes TBD"
      ],
      "metadata": {
        "id": "WbJ1nlWE0R7T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# let's not worry and use autodiff\n",
        "auto_grad_ll = jax.grad(loglikelihood)\n",
        "\n",
        "def jax_grad_step(beta, y, X, step_size):\n",
        "  beta = beta + step_size * auto_grad_ll(beta, y, X)\n",
        "  return beta\n",
        "\n",
        "# NB: we can transpose a lx.MatrixLinearOperator (say X) as X.transpose()\n",
        "# NB: we compute matrix-vector produces using a lx.MatrixLinearOperator as X.mv(v)\n",
        "step_size = 1e-7\n",
        "beta_hat = poiss_reg(y, X, jax_grad_step, step_size, max_iter=1000)\n",
        "print(f\"beta = {beta}\")\n",
        "print(f\"hat(beta) = {beta_hat}\")"
      ],
      "metadata": {
        "id": "dSPl_smqUq9y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax.scipy.linalg as spla\n",
        "\n",
        "# Great! But can we use 2nd order information?\n",
        "auto_hess_ll = jax.hessian(loglikelihood)\n",
        "\n",
        "def jax_newton_step(beta, y, X, step_size):\n",
        "  H = auto_hess_ll(beta, y, X)\n",
        "  g = auto_grad_ll(beta, y, X)\n",
        "  # recall newton -step- is - inv(H)g; so maximization under\n",
        "  # beta = beta + (newton-step) = beta - inv(H)g\n",
        "  beta = beta - step_size * spla.solve(H, g, assume_a=\"sym\")\n",
        "  return beta\n",
        "\n",
        "# NB: we can transpose a lx.MatrixLinearOperator (say X) as X.transpose()\n",
        "# NB: we compute matrix-vector produces using a lx.MatrixLinearOperator as X.mv(v)\n",
        "step_size = 1.\n",
        "beta_hat = poiss_reg(y, X, jax_newton_step, step_size, max_iter=1000)\n",
        "print(f\"beta = {beta}\")\n",
        "print(f\"hat(beta) = {beta_hat}\")"
      ],
      "metadata": {
        "id": "qtbqZkAlWPuv"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}