tutorials/exponential_distribution_mle_gradient_descent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0ca8abd4",
   "metadata": {},
   "source": [
    "## Exponential Distribution and Gradient Descent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7507932",
   "metadata": {},
   "source": [
    "**Copyright (c) Meta Platforms, Inc. and affiliates.**\n",
    "\n",
    "This source code is licensed under the MIT license found in the\n",
    "LICENSE file in the root directory of this source tree."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a09ea8",
   "metadata": {},
   "source": [
    "The **exponential distribution** models the elapsed time between events, such as the time between each ad clicks or video views. For any given $ x $, the distribution gives the likelihood that $ x $ time passes between two consecutive events. \n",
    "\n",
    "We will fit the exponential distribution using the maximum likelihood estimation (MLE) and gradient descent, utilizing the power of DiffKt. \n",
    "\n",
    "Bring in DiffKt with the following command: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "240ca59f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@file:DependsOn(\"../kotlin/api/build/libs/api.jar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fc73476",
   "metadata": {},
   "source": [
    "The exponential distribution is modeled with the following function. Note that $ x $ has to be greater than 0. \n",
    "\n",
    "$$\n",
    "y = \\lambda e^{-\\lambda x}  \\{x > 0\\}\n",
    "$$\n",
    "\n",
    "To visualize how the exponential distribution morphs with different $ \\lambda $ values, here is an animation showing lambda values ranging from $ 0 $ to $ 5 $. \n",
    "\n",
    "![](./resources/MLPgNgVTsy.mp4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d784decc",
   "metadata": {},
   "source": [
    "Let's say we are running a video streaming website and tracking traffic to a particular video. We are trying to predict the length of time $ x $ between each video view. Logically, $ x $ cannot be negative because a negative length of time between events does not make sense. \n",
    "\n",
    "Lambda $ \\lambda $ is a parameter of rate, which in this case the number of views per minute. Therefore $ \\lambda = 2 $ means there are 2 viewers every minute on average. \n",
    "\n",
    "$$\n",
    "f(x) = 2 e^{-2 x} \\{x > 0\\}\n",
    "$$\n",
    "\n",
    "![](./resources/IUiwhJeriX.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "920fbcad",
   "metadata": {},
   "source": [
    "To interpret this, let's pose a question: what is the probability we would get a video view in 1 second or less? To do that we would find the area under the curve for that range. Since this is a probability distribution, the area under the entire curve is $ 1.0 $, but we are only interested in the area between $ 0 \\le x  \\le 1 $. \n",
    "\n",
    "Let's visualize the area of interest $ 0 \\le x  \\le 1 $ below, which will yield a probability of .8647. \n",
    "\n",
    "![](./resources/ZEKHazQuWs.mp4)\n",
    "\n",
    "That's how you interpret an exponential distribution. But how do you fit one when presented with a dataset? What is the correct $ \\lambda $ parameter? This is where we can use DiffKt with gradient descent. Below is a sample of data where we tracked the time between 15 video views as declared below in Kotlin code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "713cd5c5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "val xData = floatArrayOf(1.8929f, \n",
    "       6.3687f, \n",
    "       3.228f, \n",
    "       1.2192f, \n",
    "       0.2585f, \n",
    "       0.4404f, \n",
    "       3.0278f, \n",
    "       1.9918f, \n",
    "       3.4013f, \n",
    "       3.0343f, \n",
    "       1.0201f, \n",
    "       2.436f, \n",
    "       1.8981f, \n",
    "       2.9764f, \n",
    "       1.3621f\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "712d4669",
   "metadata": {},
   "source": [
    "Here is the data visualized on the x-axis below. \n",
    "\n",
    "![](./resources/ZhsZtuABdv.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71e40aaa",
   "metadata": {},
   "source": [
    "What is the best exponential distribution to fit to this data? What is the right lambda  $ \\lambda $ value? We do need to use maximum likelihood estimation, which finds the exponential function most likely to output the observed data. But rather than go through a lot of derivative calculations (as demonstrated in Josh Starmer's video [here](https://www.youtube.com/watch?v=p3T-_LMrvBc)) we can instead use DiffKT to do all the derivative work for us. \n",
    "\n",
    "We can take the derivative of the likelihood with respect to lambda $ \\lambda $ at a given $ x $ value. We can use gradient *ascent* to step up the gradient slope and find our way to the maximum likelihood. This means we find the parameter that is most likely to output the observed data, or their combined likelihoods. If we multiply the likelihoods together we will likely run into issues with floating point underflow. To combat this we can instead work in logspace, where we take the log of each likelihood and sum instead of multiply.\n",
    "\n",
    "Set the learning rate to `.01` and use `25` iterations, and that should be sufficient to find the exponential distribution producing the maximum likelihood of observing those 15 points. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b7c00080",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lambda=0.43408304"
     ]
    }
   ],
   "source": [
    "import org.diffkt.*\n",
    "\n",
    "val xTensor = tensorOf(*xData)\n",
    "\n",
    "// Start lambda at 1\n",
    "var lambda: DScalar = FloatScalar.ONE // intercept\n",
    "\n",
    "// Declare exponential function as scalar function \n",
    "fun f(lambda: DScalar) = lambda * exp(-lambda*xTensor)\n",
    "\n",
    "// Calculate total likelihood by taking natural log of each data point\n",
    "// and then sum them \n",
    "fun likelihood_est(lambda: DScalar): DScalar = ln(f(lambda)).sum()\n",
    "\n",
    "// The learning rate\n",
    "val L = .01F\n",
    "\n",
    "// The number of iterations to perform gradient descent\n",
    "val iterations = 25\n",
    "\n",
    "// Perform gradient descent\n",
    "for (i in 0..iterations) {\n",
    "\n",
    "    // get gradient for lambda \n",
    "    val lambdaGradient = forwardDerivative(lambda, ::likelihood_est)\n",
    "\n",
    "    // update m and b by adding the (learning rate) * (slope)\n",
    "    lambda += lambdaGradient * L\n",
    "}\n",
    "print(\"lambda=$lambda\") // "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53d389a9",
   "metadata": {},
   "source": [
    "Running the code above we get a lambda value of about $ 0.434 $. Here is a visualization of the gradient descent at work below. \n",
    "\n",
    "![](./resources/LPUPIctltW.mp4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20546cdf",
   "metadata": {},
   "source": [
    "This looks pretty good! As you can observe from the function above, we are much more likely to observe time between video videos closer to 0 seconds, but longer time lapses beyond 4 seconds are less likely. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "codemirror_mode": "text/x-kotlin",
   "file_extension": ".kt",
   "mimetype": "text/x-kotlin",
   "name": "kotlin",
   "nbconvert_exporter": "",
   "pygments_lexer": "kotlin",
   "version": "1.7.0-dev-3303"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}