ml_optimizers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "yNXcdUh_T7O_",
   "metadata": {
    "id": "yNXcdUh_T7O_"
   },
   "source": [
    "# Optimizers in ML using JAX\n",
    "\n",
    "> Optimization refers to the process of minimizing the loss function by systematically updating the network weights.\n",
    "\n",
    "In this notebook, I implement a few popular optimizers from scratch for a simple model i.e., Linear Regression on a dataset of 5 features. The goal of this notebook was to understand how these optimizers work under the hood and try to do a toy implementation myself. I also use a bit of JAX magic to perform the differentiation of the loss function w.r.t to the Weights and the Bias without explicitly writing their derivatives as a separate function. This can help to generalize this notebook for other types of loss functions as well.\n",
    "\n",
    "The optimizers I have implemented are - \n",
    "* Batch Gradient Descent\n",
    "* Batch Gradient Descent + Momentum\n",
    "* Nesterov Accelerated Momentum\n",
    "* Adagrad\n",
    "* RMSprop\n",
    "* Adam\n",
    "* Adamax\n",
    "* Nadam\n",
    "* Adabelief\n",
    "\n",
    "References -\n",
    "* https://ruder.io/optimizing-gradient-descent/\n",
    "* https://theaisummer.com/optimization/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1DRKSr9ULZU",
   "metadata": {
    "id": "b1DRKSr9ULZU"
   },
   "source": [
    "## Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ymTdjXeSUNsi",
   "metadata": {
    "id": "ymTdjXeSUNsi"
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as np\n",
    "from sklearn.datasets import make_regression\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3WNF0uapUOhM",
   "metadata": {
    "id": "3WNF0uapUOhM"
   },
   "source": [
    "## Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aeeafb3b",
   "metadata": {
    "id": "aeeafb3b"
   },
   "outputs": [],
   "source": [
    "def get_data():\n",
    "    \"\"\"Create dataset for linear regression using 5 features. Set fix bias = 2.0\n",
    "    \n",
    "    Returns:\n",
    "        X: training data\n",
    "        X_test: testing data\n",
    "        y: train labels\n",
    "        y_test: test labels\n",
    "        coef: true weight matrix (coefficients) for the dataset\n",
    "    \"\"\"\n",
    "    # create our dataset. Set fix bias of 2.0 and return weights (coef=True)\n",
    "    X, y, coef = make_regression(n_features=5, coef=True, bias=2.0)\n",
    "    X, X_test, y, y_test = train_test_split(X, y)\n",
    "    return (X, X_test, y, y_test, coef)\n",
    "\n",
    "def J(X, w, b, y):\n",
    "    \"\"\"Cost function for a linear regression. A forward pass of our model.\n",
    "\n",
    "    Args:\n",
    "        X: a features matrix.\n",
    "        w: weights (a column vector).\n",
    "        b: a bias.\n",
    "        y: a target vector.\n",
    "\n",
    "    Returns:\n",
    "        scalar: a cost of this solution.    \n",
    "    \"\"\"\n",
    "    y_hat = np.dot(X, w) + b # Predict values.\n",
    "    return ((y_hat - y)**2).mean() # Return cost."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "DeaSlrOHUZm0",
   "metadata": {
    "id": "DeaSlrOHUZm0"
   },
   "source": [
    "## Batch Gradient Descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e74873db",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "e74873db",
    "outputId": "7820eb77-45eb-4d27-bc06-3e6dfad884cc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [30.68836406  8.19425997 52.32139276 98.63300635 35.70556397]\n",
      "Calculated weights = [30.688354  8.194252 52.321373 98.632965 35.705536]\n",
      "True bias = 2.0\tCalculated bias = 2.000007\n",
      "Test loss: 0.000000006\n"
     ]
    }
   ],
   "source": [
    "from optimizers.batch_gradient_descent import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = batch_gradient_descent(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wJ4TINl7UcVV",
   "metadata": {
    "id": "wJ4TINl7UcVV"
   },
   "source": [
    "## Batch Gradient Descent + Momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "LkTZuoRRMW3U",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LkTZuoRRMW3U",
    "outputId": "046e4720-afa5-431c-a06d-a755b75739dc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [79.89865004 69.54844035  5.00857317 17.07146713 95.33941246]\n",
      "Calculated weights = [79.89865  69.54844   5.00857  17.071468 95.33941 ]\n",
      "True bias = 2.0\tCalculated bias = 1.999996\n",
      "Test loss: 0.000000000\n"
     ]
    }
   ],
   "source": [
    "from optimizers.momentum import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = momentum(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ceojilDUfH1",
   "metadata": {
    "id": "6ceojilDUfH1"
   },
   "source": [
    "## Nesterov accelerated momentum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5_niQkioN8Ll",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5_niQkioN8Ll",
    "outputId": "1d81ae85-b653-4c1d-c2c9-e57607c5b579"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [42.60473498 65.07277781 34.78011253 62.33332432 73.84550563]\n",
      "Calculated weights = [42.604736 65.07278  34.780113 62.333324 73.845505]\n",
      "True bias = 2.0\tCalculated bias = 2.000000\n",
      "Test loss: 0.000000000\n"
     ]
    }
   ],
   "source": [
    "from optimizers.nesterov_momentum import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = nesterov_momentum(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "x8cd4ImwUmvk",
   "metadata": {
    "id": "x8cd4ImwUmvk"
   },
   "source": [
    "## Adagrad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "VMfkP5uVN_9F",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VMfkP5uVN_9F",
    "outputId": "cdba9b42-9137-4f0b-ebc9-07c5f2edf16f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [14.77439555 65.65138433 13.16405989 56.48558691 90.74394822]\n",
      "Calculated weights = [14.773779 65.65003  13.163625 56.486057 90.73908 ]\n",
      "True bias = 2.0\tCalculated bias = 1.999916\n",
      "Test loss: 0.000044874\n"
     ]
    }
   ],
   "source": [
    "from optimizers.adagrad import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = adagrad(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "O7Yq1rylUoTk",
   "metadata": {
    "id": "O7Yq1rylUoTk"
   },
   "source": [
    "## RMSprop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "o2oZOLTtOHE0",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "o2oZOLTtOHE0",
    "outputId": "7dcef72b-13c6-44c0-b034-1fa359a54a92"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [41.97869789 49.60297967 28.45822788 46.6347428  48.8575031 ]\n",
      "Calculated weights = [41.9287   49.552982 28.40823  46.58474  48.807503]\n",
      "True bias = 2.0\tCalculated bias = 2.050000\n",
      "Test loss: 0.013983894\n"
     ]
    }
   ],
   "source": [
    "from optimizers.rmsprop import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = rmsprop(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ncONWqdHUqod",
   "metadata": {
    "id": "ncONWqdHUqod"
   },
   "source": [
    "## Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "PcclC8lzOKoN",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PcclC8lzOKoN",
    "outputId": "5a3d052a-6e54-43ff-8bf5-aa294c05cdba"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [91.82405194  5.98911475 51.21966463 76.09491592  3.97291692]\n",
      "Calculated weights = [91.82391    5.989142  51.21963   76.09506    3.9729562]\n",
      "True bias = 2.0\tCalculated bias = 2.000038\n",
      "Test loss: 0.000000036\n"
     ]
    }
   ],
   "source": [
    "from optimizers.adam import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = adam(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "HKklF0Q1Usk0",
   "metadata": {
    "id": "HKklF0Q1Usk0"
   },
   "source": [
    "## Adamax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2ZXobE06OQ01",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2ZXobE06OQ01",
    "outputId": "d7bce9b1-4b15-46f0-80cf-aacac0c0a8ad"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [70.74526421 40.12974786  4.90109456 56.99685724 58.66588867]\n",
      "Calculated weights = [70.745094  40.12968    4.9011374 56.99679   58.666    ]\n",
      "True bias = 2.0\tCalculated bias = 1.999996\n",
      "Test loss: 0.000000061\n"
     ]
    }
   ],
   "source": [
    "from optimizers.adamax import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = adamax(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1P8TTFz9UuAF",
   "metadata": {
    "id": "1P8TTFz9UuAF"
   },
   "source": [
    "## Nadam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "IbEwND3POXLU",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IbEwND3POXLU",
    "outputId": "f034239f-5e9c-4334-a086-7fe2743468d8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [ 1.58186482 19.09535888 16.30502928 81.38680601 73.46240111]\n",
      "Calculated weights = [ 1.581854 19.095356 16.305021 81.38672  73.46231 ]\n",
      "True bias = 2.0\tCalculated bias = 1.999992\n",
      "Test loss: 0.000000011\n"
     ]
    }
   ],
   "source": [
    "from optimizers.nadam import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = nadam(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "FtwPsK0fUwG2",
   "metadata": {
    "id": "FtwPsK0fUwG2"
   },
   "source": [
    "## Adabelief"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "kVU48zr3Oafk",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kVU48zr3Oafk",
    "outputId": "0bc70231-f81b-40cb-ab87-fe5f9771a7a6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True weights = [71.25219567 15.10311738  7.49360451 90.84284877 87.45135942]\n",
      "Calculated weights = [71.25219  15.10311   7.493615 90.84281  87.45141 ]\n",
      "True bias = 2.0\tCalculated bias = 1.999991\n",
      "Test loss: 0.000000003\n"
     ]
    }
   ],
   "source": [
    "from optimizers.adabelief import *\n",
    "X, X_test, y, y_test, coef = get_data()\n",
    "params = adabelief(J, X, y)\n",
    "\n",
    "print(\"True weights =\", coef)\n",
    "print(\"Calculated weights =\", params['w'])\n",
    "print(\"True bias = 2.0\\tCalculated bias = {:.6f}\".format(params['b']))\n",
    "print(\"Test loss: {:.9f}\".format(J(X_test, params['w'], params['b'], y_test)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "CsfOgDyEOhFV",
   "metadata": {
    "id": "CsfOgDyEOhFV"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ml-optimizers.ipynb",
   "provenance": []
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}