mle_gaussian.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Maximum Likelihood Estimation for Gaussian Random Variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial demonstrates how to find the maximum likelihood estimates of the mean $\\mu$ and the variance $\\sigma^2$ both analytically and using gradient descent for a Gaussian random variable $X$. The probability density function of $X$ is:\n",
    "\n",
    "$$f_X(x;\\mu,\\sigma^2) = \\frac{1}{\\sqrt{2\\pi\\sigma^2}}\\exp\\left({-\\frac{(x-\\mu)^2}{2\\sigma^2}}\\right)$$\n",
    "\n",
    "Suppose that you observe 500 independent and identically distributed (i.i.d) samples of $X$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of samples: 500\n",
      "Variance: 9.338441848754883\n",
      "Mean: 2.0771899223327637\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "mu = 2\n",
    "sigma = 3\n",
    "num_samples = 500\n",
    "\n",
    "X = torch.randn((num_samples,),requires_grad = False)\n",
    "X = sigma*X + mu\n",
    "\n",
    "print('Number of samples: {}\\nVariance: {}\\nMean: {}'.format(X.shape[0],*torch.var_mean(X)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The independence assumption is valid if observing one of the samples does not give you information about the other samples, while the identical distribution assumption makes sense if these observations originated from the same underlying random experiment. Therefore, the dataset $\\mathcal{D}$ is:\n",
    "\n",
    "$$\\mathcal{D} = \\{X = x_1,X = x_2,...,X = x_{500}\\}$$\n",
    "\n",
    "The likelihood function is therefore:\n",
    "\n",
    "$$ p(\\mathcal{D};\\mu,\\sigma^2) = \\prod_{i=1}^{500}f_X(x_i;\\mu_i,\\sigma_{i}^2) = \\prod_{i=1}^{500} \\frac{1}{\\sqrt{2\\pi\\sigma_{i}^2}} \\exp\\left({-\\frac{(x_i-\\mu_i)^2}{2\\sigma_{i}^2}}\\right) $$\n",
    "\n",
    "Since the samples are identically distributed:\n",
    "\n",
    "$$\n",
    "\\mu_1 = \\mu_2 = ... = \\mu_{500} = \\mu \\\\\n",
    "\\sigma_{1}^2 = \\sigma_{2}^2 = ... = \\sigma_{500}^2 = \\sigma^2\n",
    "$$\n",
    "\n",
    "Then the log-likelihood function is:\n",
    "\n",
    "\\begin{align}\n",
    "\\ln(p(\\mathcal{D};\\mu,\\sigma^2)) &= \\sum_{i=1}^{500} \\ln\\left(\\frac{1}{\\sqrt{2\\pi\\sigma_{i}^2}}\\right) + \\ln\\left(\\exp\\left({-\\frac{(x_i-\\mu_i)^2}{2\\sigma_{i}^2}}\\right)\\right) \\\\\n",
    "&= -\\left(\\sum_{i=1}^{500} \\ln\\left(\\sqrt{2\\pi\\sigma_{i}^2}\\right) + \\sum_{i=1}^{500} \\frac{(x_i-\\mu_i)^2}{2\\sigma_{i}^2}\\right) \\\\\n",
    "&= -\\left(500 \\ln\\left(\\sqrt{2\\pi\\sigma^2}\\right) + \\frac{1}{2\\sigma^2}\\sum_{i=1}^{500} (x_i-\\mu)^2 \\right)\n",
    "\\end{align}\n",
    "\n",
    "The negative of the log-likelihood function can be implemented as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(4043.6475)\n"
     ]
    }
   ],
   "source": [
    "# Negative log-likelihood of X for a given mu and sigma\n",
    "\n",
    "from math import pi\n",
    "\n",
    "def normal_NLL(X,theta): # theta[0] = mu, theta[1] = sigma^2\n",
    "    first_term = X.shape[0]*torch.log(torch.sqrt(2*pi*theta[1]))\n",
    "    second_term = torch.div(1.,2*theta[1])*torch.sum(torch.pow(X-theta[0],2))\n",
    "    return (first_term + second_term)\n",
    "\n",
    "theta = torch.rand((2,))*10\n",
    "NLL = normal_NLL(X,theta)\n",
    "print(NLL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Therefore, the maximum likelihood estimates of $\\mu$ and $\\sigma^2$ are:\n",
    "\n",
    "\\begin{equation}\n",
    "\\DeclareMathOperator*{\\argmin}{\\arg\\!\\min}\n",
    "\\hat{\\mu} = \\argmin_{\\mu}{\\left(-\\ln\\left(p\\left(\\mathcal{D};\\mu,\\sigma^2\\right)\\right)\\right)} =  \\argmin_{\\mu}\\sum_{i=1}^{500} (x_i-\\mu)^2\\\\\n",
    "\\hat{\\sigma}^2 = \\argmin_{\\sigma^2}{\\left(-\\ln\\left(p\\left(\\mathcal{D};\\mu,\\sigma^2\\right)\\right)\\right)} = \\argmin_{\\sigma^2} \\left(500\\ln\\left(\\sqrt{2\\pi\\sigma^2}\\right) + \\frac{1}{2\\sigma^2}\\sum_{i=1}^{500}(x_i-\\mu)^2\\right)\n",
    "\\end{equation}\n",
    "\n",
    "$\\hat{\\mu}$ and $\\hat{\\sigma}^2$ can be computed analytically. To compute $\\hat{\\mu}$:\n",
    "\n",
    "\\begin{align}\n",
    "- \\frac{\\partial\\ln\\left(p\\left(\\mathcal{D};\\mu,\\sigma^2\\right)\\right)}{\\partial\\mu} &= -2 \\sum_{i=1}^{500} (x_i-\\mu) \\\\\n",
    "&= -2 \\left(\\sum_{i=1}^{500} x_i - \\sum_{i=1}^{500} \\mu\\right) \\\\\n",
    "&= -2 \\left(\\sum_{i=1}^{500} x_i - 500\\mu\\right) = 0\n",
    "\\end{align}\n",
    "\n",
    "Solving for $\\mu$:\n",
    "\n",
    "$$\\hat{\\mu} = \\frac{\\sum_{i=1}^{500} x_i}{500}$$\n",
    "\n",
    "Which is just the sample mean. Similarly, to compute the variance $\\hat{\\sigma}^2$:\n",
    "\n",
    "$$\n",
    "-\\frac{\\partial\\ln\\left(p\\left(\\mathcal{D};\\mu,\\sigma^2\\right)\\right)}{\\partial\\sigma^2} = \\frac{500}{2\\sigma^2} - \\frac{1}{2\\sigma^4} \\sum_{i=1}^{500} (x_i-\\mu)^2 = 0\n",
    "$$\n",
    "\n",
    "Solving for $\\sigma^2$:\n",
    "\n",
    "$$\n",
    "\\hat{\\sigma}^2 = \\frac{1}{500} \\sum_{i=1}^{500} (x_i-\\mu)^2\n",
    "$$\n",
    "\n",
    "Which is just the sample variance. However, this is a biased estimator of $\\sigma^2$. Instead, an unbiased estimate of $\\sigma^2$ can be computed as follows:\n",
    "\n",
    "$$\n",
    "\\hat{\\sigma}^2 = \\frac{1}{500-1} \\sum_{i=1}^{500} (x_i-\\mu)^2\n",
    "$$\n",
    "\n",
    "This modification is called [Bessel's correction](https://en.wikipedia.org/wiki/Bessel%27s_correction). The maximum likelihood estimates of $\\mu$ and $\\sigma^2$ can then be computed as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.0772)\n"
     ]
    }
   ],
   "source": [
    "# sample mean\n",
    "\n",
    "mu_hat = torch.sum(X)/X.shape[0]\n",
    "\n",
    "print(mu_hat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(9.3384)\n"
     ]
    }
   ],
   "source": [
    "# sample variance\n",
    "\n",
    "sigma_squared_hat = torch.sum(torch.pow(X-mu_hat,2))/(X.shape[0]-1)\n",
    "\n",
    "print(sigma_squared_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, the maximum likelihood estimates of $\\mu$ and $\\sigma^2$ can be computed using gradient descent. Gradient descent is an iterative algorithm based on the following general methodology:\n",
    "1. Guess an initial value for the parameter vector $\\mathbf{\\theta} = (\\mu,\\sigma^2)$.\n",
    "2. Compute the gradient of the negative log-likeihood function at these initial values of $\\mu$ and $\\sigma^2$. More precisely, if $\\mathbf{\\theta}_0 = (\\mu_0,\\sigma_0^2)$ is the initial parameter vector, then compute:\n",
    "\n",
    "$$\n",
    "-\\left.\\frac{\\partial\\ln\\left(p\\left(\\mathcal{D};\\mathbf{\\theta}\\right)\\right)}{\\partial\\mathbf{\\theta}} \\right\\rvert_{\\mathbf{\\theta} = \\mathbf{\\theta}_0}\n",
    "$$\n",
    "\n",
    "Note that the vector $\\mathbf{\\theta} = (\\mu,\\sigma^2)$ was used here because the individual partial derivatives of the negative log-likelihood function with respect to $\\mu$ and $\\sigma^2$ can only be used to perform minimization exclusively with respect to $\\mu$ or $\\sigma^2$ but not both at the same time.\n",
    "\n",
    "3. Use some function:\n",
    "\n",
    "$$\n",
    "\\mathbf{\\theta}_1 = g\\left(\\mathbf{\\theta}_0,-\\left.\\frac{\\partial\\ln\\left(p\\left(\\mathcal{D};\\mathbf{\\theta}\\right)\\right)}{\\partial\\mathbf{\\theta}} \\right\\rvert_{\\mathbf{\\theta} = \\mathbf{\\theta}_0}\\right)\n",
    "$$\n",
    "\n",
    "To choose a new parameter vector $\\mathbf{\\theta}_1$ that minimizes the negative log-likelihood function.\n",
    "\n",
    "4. Repeat steps 1,2, and 3 until convergence.\n",
    "\n",
    "The following explanation for how to choose $g(\\cdot)$ is adapted from [here](https://eli.thegreenplace.net/2016/understanding-gradient-descent/). Consider the simple function $f(x) = x^2$. The value of $f(x)$ decreases when x is negative and increasing, while its value increases when x is positive and increasing. Given an infinitesimally small change $dx$, such that:\n",
    "\n",
    "\\begin{align}\n",
    "f(x)+df(x) &= (x+dx)^2 \\\\\n",
    "f(x)+df(x) &= x^2 + 2x dx + {dx}^2 \\\\\n",
    "x^2 + df(x) &= x^2 + 2x dx \\\\\n",
    "df(x) &= 2x dx \\\\\n",
    "\\frac{df(x)}{dx} &= 2x\n",
    "\\end{align}\n",
    "\n",
    "Suppose that you do not know where the minimum of $f(x)$ is. Let us guess first that the minimum is at $x = -2$. The derivative of $f(x)$ at $x = -2$ is $-4$, which means that an infinitely small increase in $x$, denoted as $dx$, will result in a decrease in $f(x)$ by a factor of 4. Thus, it makes sense to increase $x$ in order to reach the minimum of $f(x)$. In fact, for any convex function $f(x)$, the negative of its derivative $\\frac{df(x)}{dx}$ will always \"point\" towards its minimum point. So, it makes sense to make a new guess $x_{new}$ such that:\n",
    "\n",
    "$$\n",
    "x_{new} = x_{old} - \\eta\\left.\\frac{df(x)}{dx}\\right\\rvert_{x = x_{old}}\n",
    "$$\n",
    "\n",
    "Where $\\eta$ is a small constant called the _learning rate_ that controls the magnitude of $\\frac{df(x)}{dx}$ and thus controls the magnitude of each change of $x$.\n",
    "\n",
    "For multivariate functions, such as $f(\\mathbf{\\theta})$ where $\\mathbf{\\theta} = (x,y,z)$, it can be shown that its gradient $\\nabla_{\\mathbf{\\theta}}f(\\mathbf{\\theta})$ points in the direction of steepest ascent, and so its negative points in the direction of steepest descent. Therefore, for multivariate functions, the update rule becomes:\n",
    "\n",
    "$$\n",
    "\\mathbf{\\theta}_{new} = \\mathbf{\\theta}_{old} - \\eta \\left.\\nabla_{\\mathbf{\\theta}}f(\\mathbf{\\theta})\\right\\rvert_{\\mathbf{\\theta} = \\mathbf{\\theta}_{old}}\n",
    "$$\n",
    "\n",
    "Therefore, in the case of $\\mu$ and $\\sigma^2$, the required gradient is:\n",
    "\n",
    "$$\n",
    "-\\left.\\frac{\\partial\\ln\\left(p\\left(\\mathcal{D};\\mathbf{\\theta}\\right)\\right)}{\\partial\\mathbf{\\theta}} \\right\\rvert_{\\mathbf{\\theta} = \\mathbf{\\theta}_{old}}\n",
    "$$\n",
    "\n",
    "Where $\\mathbf{\\theta}_{old} = (\\mu_{old},\\sigma^2_{old})$.\n",
    "\n",
    "Gradient descent can be implemented as follows. First, an initial theta is chosen at random, and its gradient is tracked:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([3.3409, 6.9727], requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "theta_old = torch.rand((2,),requires_grad = False)\n",
    "theta_old = theta_old*10 # sample uniformly from range [0,10]\n",
    "theta_old.requires_grad = True # theta_old is now a leaf tensor\n",
    "\n",
    "print(theta_old)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, a learning rate is chosen:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.001"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, the negative log-likelihood of $X$ with respect to the initial theta is computed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1336.3811, grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "NLL = normal_NLL(X,theta_old)\n",
    "\n",
    "print(NLL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The gradient of the negative log-likelihood function with respect to theta is then computed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "NLL.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The new value of theta is then computed using the gradient descent update rule:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([3.2503, 6.9930], grad_fn=<SubBackward0>)\n"
     ]
    }
   ],
   "source": [
    "theta_new = theta_old - lr*theta_old.grad\n",
    "theta_old.grad.zero_() # zero gradients accumulated in .grad\n",
    "\n",
    "print(theta_new)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, this process is repeated until convergence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "eps = 1e-10\n",
    "\n",
    "while torch.norm(theta_new - theta_old).item() > eps:\n",
    "    theta_old = theta_new.detach().requires_grad_() # copy \n",
    "    NLL = normal_NLL(X,theta_old)\n",
    "    NLL.backward()\n",
    "    theta_new = theta_old - lr*theta_old.grad\n",
    "    theta_old.grad.zero_() # zero gradients accumulated in .grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final estimates of $\\mu$ and $\\sigma^2$ using gradient descent are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Maximum likelihood estimate of mean: 2.0771918296813965\n",
      "Maximum likelihood estimate of variance: 9.319599151611328\n"
     ]
    }
   ],
   "source": [
    "print('Maximum likelihood estimate of mean: {}'.format(theta_new[0].item()))\n",
    "print('Maximum likelihood estimate of variance: {}'.format(theta_new[1].item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that the sample mean and variance that were analytically computed are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sample mean: 2.0771901607513428\n",
      "Sample variance: 9.338441848754883\n"
     ]
    }
   ],
   "source": [
    "print('Sample mean: {}\\nSample variance: {}'.format(mu_hat,sigma_squared_hat))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}