stochastic_variational_inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyOdrA7FbetffxZ+EOLpm7dO",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Antony-gitau/probabilistic_AI_playgraound/blob/main/stochastic_variational_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this notebook, we want to implement a simple example that will help us understand the concept of SVI while building a probabilistic model.\n",
        "\n",
        "recap:\n",
        "\n",
        "- probabilistic models are statistical models that include one or more probability distribution functions in the model to account for uncertainities.\n",
        "\n",
        "- pyro is one of the programming languages that helps us code probabilistic models. Other languages include pymc\n",
        "\n",
        "- stochastic variational inference is a way for calculating the variational inference. [Variational inference](https://en.wikipedia.org/wiki/Variational_Bayesian_methods) is one of the technique of approximating intractable integrals of a bayesian inference. \n",
        "\n",
        "\n",
        "So in this simple example, we will try to determine whether the coin is fair or not. \n",
        "\n",
        "we will first start off by genearing some datapoints with 6 observed heads and 4 observed tails."
      ],
      "metadata": {
        "id": "y-qAL4ylF3nK"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "F6NEskAaFbYW"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install -q --upgrade pyro-ppl torch\n",
        "import torch\n",
        "import torch.distributions.constraints as constraints\n",
        "import pyro\n",
        "import pyro.distributions as dist"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#create an empty list of data\n",
        "data = []\n",
        "\n",
        "#populate the data variable with tensors. 6 1s (heads) and 4 0s (tails)\n",
        "for _ in range(6):\n",
        "  data.append(torch.tensor(1.0))\n",
        "\n",
        "for _ in range(4):\n",
        "  data.append(torch.tensor(0.0))\n"
      ],
      "metadata": {
        "id": "1hKEUD3XVkRC"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7hMlqbV4XDI4",
        "outputId": "20241bc2-8025-443b-d9e8-8d1b64b4aca5"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[tensor(1.),\n",
              " tensor(1.),\n",
              " tensor(1.),\n",
              " tensor(1.),\n",
              " tensor(1.),\n",
              " tensor(1.),\n",
              " tensor(0.),\n",
              " tensor(0.),\n",
              " tensor(0.),\n",
              " tensor(0.)]"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's now create the probabilistic model. \n",
        "\n",
        "We know that;\n",
        "- heads are encoded as 1 and tails as 0\n",
        "- the fairness of a coin is 0.5 withot bias. we can encode the 0.5 into a value f, where f=0.5 and $f 𝞊 [0.0, 1.0]$\n",
        "- we can know that coin might be unbiased. this is our prior information. but we dont know how unfair this coin is. that is why we will be quantifying.\n",
        "- we know that beta distribution can be used to model probabilities of binary events. We only need to choose the α and β paramaters, and in this case, we can experiment with a probability distribution of just 10 samples.\n",
        "\n"
      ],
      "metadata": {
        "id": "lEAcUWiSXL5V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# define the model function taking in the data\n",
        "\n",
        "def model(data):\n",
        "  #we define the prior infomation. we chose a beta distribution \n",
        "  alpha = torch.tensor(10.0)\n",
        "  beta = torch.tensor(10.0)\n",
        "\n",
        "# we generate sample f(fairness) as a beta distrbution with hyperparameters alpha and beta\n",
        "  f = pyro.sample(\"latent_fairness\", dist.Beta(alpha, beta))\n",
        " \n",
        " # loop over the observed data\n",
        "  for i in range(len(data)):\n",
        "    # we model the likelihood of abserving heads or tail (fairness of the coin) in a \n",
        "    #bernoulli function\n",
        "    pyro.sample(\"obs_{}\".format(i), dist.Bernoulli(f), obs = data[i])\n"
      ],
      "metadata": {
        "id": "mPyhJEm4XDrU"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notes on the thought process of choosing likelihood function:\n",
        "\n",
        "(so in the code above, we choose bernoulli disrtibution as the likelihood function. Its is a good choice for binary outcomes)\n",
        "\n",
        "- type of observed data. continuos discrete and binary data require different likelihood functions\n",
        "\n",
        "- the assumptions of the model. that is the parameters of the model.\n",
        "- distribution should also be tractable\n",
        "\n",
        "---\n",
        "next up, we define a guide. Guide is the fancy name used by pyro in place of variational distribution, which are the parametirized distribution $q_𝜙(z)$. \n",
        "\n",
        "- the guide function will use the same names for random variables as in the model function\n",
        "\n",
        "- the difference between alpha and beta in the model function and alpha_q and beta_q in the guide function is that the latter require gradient, hence are learnable parameters. unlike the former which are samples that define the shape of prior distribution. \n",
        "\n",
        "- also, we initiate the parameters at 15.0 from which they will be optimized by svi to get closer to the true posterior as possible."
      ],
      "metadata": {
        "id": "kcr44qdKk32R"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# guide takes in same arguement as model\n",
        "def guide(data):\n",
        "\n",
        "  #these are the parametirized distributions over latent varibles with\n",
        "  # learnable parameters that are adjusted during optimization\n",
        "  alpha_q = pyro.param(\"alpha_q\", torch.tensor(15.0),constraint = constraints.positive)\n",
        "  beta_q = pyro.param(\"beta_q\", torch.tensor(15.0), constraint =constraints.positive)\n",
        "\n",
        "  # sample the latent fairness using beta(alpha_q and beta_q)\n",
        "  pyro.sample(\"latent_fairness\", dist.Beta(alpha_q, beta_q))"
      ],
      "metadata": {
        "id": "lF5hd7fFl8Bi"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we optmize stochatically.\n",
        "\n",
        "notes of learning rate.\n",
        "A higher rate means larger steps and faster convergence but may results in overshooting the optimal values. vice versa for lower rates. \n",
        "\n",
        "a note on betas. this parameter controls the weight given to past gradient when computing the update direction for the model parameters. That is controls the moving average. a large value of beta means we are averaging over more gradients.\n",
        " \n",
        "\n",
        " a higher beta1 means more weight is given to recent gradients, and a higher beta2 values means more weight is given to the magnitude if the gradient and lower beta2 means more weight is given to direction of the gradient."
      ],
      "metadata": {
        "id": "rgvILCnFr6rQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# set up an adam optimizer\n",
        "from pyro.optim import Adam\n",
        "adam_params = {\"lr\": 0.005, \"betas\": (0.90, 0.999)}\n",
        "optimizer = Adam(adam_params)"
      ],
      "metadata": {
        "id": "MSMjHrRwr9qp"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we can use svi as the inference algorithm"
      ],
      "metadata": {
        "id": "oCAzlRo1sm3n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pyro.infer import SVI, Trace_ELBO\n",
        "svi = SVI(model, guide, optimizer, loss=Trace_ELBO())\n",
        "\n",
        "n_steps = 5000\n",
        "for step in range(n_steps):\n",
        "  svi.step(data)"
      ],
      "metadata": {
        "id": "Rzz0GF3bsqsi"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# we can assign a varible to the variational parameters\n",
        "alpha_q = pyro.param(\"alpha_q\").item()\n",
        "alpha_q"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fXGVaXfIu0wh",
        "outputId": "0b8b1bc4-fb75-4b02-849a-46523fe8c996"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "15.508711814880371"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "beta_q = pyro.param(\"beta_q\").item()\n",
        "beta_q"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qep-3OQ6vRi4",
        "outputId": "a8d348b4-a763-42f7-bc2d-e5c3cdfdf354"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "13.823347091674805"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can now compute the infered mean and standard deviation of the coin's fairness. "
      ],
      "metadata": {
        "id": "a0EdSdUIvrZW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "# compute the inferred mean of the coin's fairness\n",
        "inferred_mean = alpha_q / (alpha_q + beta_q)\n",
        "print(\"Inferred mean is \" + str(inferred_mean))\n",
        "\n",
        "# compute inferred standard deviation\n",
        "factor = beta_q / (alpha_q * (1.0 + alpha_q + beta_q))\n",
        "print(\"This is the factor \" + str(factor))\n",
        "\n",
        "inferred_std = inferred_mean * math.sqrt(factor)\n",
        "print(\"Inferred std \" + str(inferred_std))\n",
        "\n",
        "print(\"\\nBased on the data and our prior belief, the fairness \" +\n",
        "      \"of the coin is %.3f +- %.3f\" % (inferred_mean, inferred_std))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aohGsjIcu2jK",
        "outputId": "6aca982c-cb12-424f-8812-6e12b6e1153b"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Inferred mean is 0.5287290559550342\n",
            "This is the factor 0.029385669935780522\n",
            "Inferred std 0.09063605108814138\n",
            "\n",
            "Based on the data and our prior belief, the fairness of the coin is 0.529 +- 0.091\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Recall that our prior had 6 heads and 4 tails. Meaning, the fairness of the frequencies was $6\\over10$ = 0.60. \n",
        "\n",
        "To compute the exact posterior mean, we need to intergrate the product of likelihod and prior\n",
        "\n",
        "In this case, we used bernoulli likelihood and beta prior. It is known that the exact posterior is also a beta distribution.\n",
        "\n",
        "exact mean = $alpha\\over(alpha + beta)$ = $10 \\over (10+10)$ = 0.5\n"
      ],
      "metadata": {
        "id": "nj2BMtkbzhAW"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "xUzMP1r128re"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}