{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "33569260",
   "metadata": {},
   "outputs": [],
   "source": [
    "using PyPlot, LinearAlgebra\n",
    "\n",
    "# uncomment this, along with the @manipulate lines below,\n",
    "# to get an interactive version of these plots:\n",
    "# using Interact"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd782042",
   "metadata": {},
   "source": [
    "# Stochastic optimization\n",
    "\n",
    "Many practical optimization problems in machine learning, or other fields involving \"fitting to data\", have the property that their objective (\"loss\") functions are **additive**, i.e. you are minimizing an **average**:\n",
    "$$\n",
    "f(x) = \\frac{1}{m} \\sum_{i=1}^m f_i(x) \\longrightarrow_{m \\to \\infty} E[f_i(x)]\\, .\n",
    "$$\n",
    "where the different functions $f_i(x)$ correspond to **samples** from a **training set** of $m \\gg 1$ data points.   In the limit $m \\to \\infty$, this corresponds to minimizing the **expected value**.  (For example, if training examples are drawn randomly from some distribution (for [unsupervised learning](https://en.wikipedia.org/wiki/Unsupervised_learning)), you effectively have an infinite training set.)\n",
    "\n",
    "The most familiar example is probably **least-square fitting**, $f_i(x)$ is the squared difference between a model and a training data point $(\\text{input}_i, \\text{output}_i)$:\n",
    "$$\n",
    "f_i(x) = \\left\\Vert \\text{model}(x, \\text{input}_i) - \\text{output}_i \\right\\Vert_2^2 \\, .\n",
    "$$\n",
    "As another example, in [maximum-likelihood estimation](https://goodboychan.github.io/python/coursera/tensorflow_probability/icl/2021/08/19/01-Maximum-likelihood-estimation.html) one maximizes a \"likelihood\" probability, which can equivalently be formulated as minimizing the negative **log-likelihood** of the form:\n",
    "$$\n",
    "f_i(x) = -\\log L(\\text{output}_i | \\text{model}(x, \\text{input}_i))\n",
    "$$\n",
    "(where $L$ is a likelihood function in some statistical model, e.g. corresponding to [cross-entropy loss](https://en.wikipedia.org/wiki/Cross_entropy) in [classification problems](https://towardsdatascience.com/cross-entropy-for-classification-d98e7f974451)).  There are many other examples of this $f = \\sum_i f_i$ form.\n",
    "\n",
    "Gradient descent and related algorithms update the parameters $x$ using the gradient $\\nabla f$ evaluated at the current point (or possibly a linear combination of gradients at different points), but evaluating the gradient exactly might be very expensive if $m$ is large.  However, by linearity, what we are computing is\n",
    "$$\n",
    "\\nabla f = \\frac{1}{m} \\sum_{i=1}^m \\nabla f_i \\longrightarrow_{m \\to \\infty} \\boxed{ E[\\nabla f_i] }\\, ..\n",
    "$$\n",
    "which is effectively the **expected value of the gradient** on all the samples.\n",
    "\n",
    "In **stochastic optimization algorithms**, instead of computing the *exact* gradient $\\nabla f$, we instead **estimate the gradient** by **random sampling** of only a few points — maybe even just one! — and use this to update $x$.   A type of [Monte Carlo method](https://en.wikipedia.org/wiki/Monte_Carlo_method), this leads to a dramatic gain of comptuational efficiency, at the price of more **variance** in the convergence.\n",
    "\n",
    "In this notebook, we demonstrate a couple of basic algorithms for stochastic optimization on a simple linear-regression problem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baafb243",
   "metadata": {},
   "source": [
    "## A simple least-squares fitting problem\n",
    "\n",
    "Let's fit $m$ data points $(a_i, b_i)$ to a straight line $b(a) = x_1 + x_2 a$.  With a least-square fit, minimizing\n",
    "$$f(x) = \\Vert b - Ax \\Vert^2/m =\\frac{1}{m} \\sum_i \\underbrace{(b_i - (Ax)_i)^2}_{f_i(x)}\\, ,$$\n",
    "the solution is simply $\\hat{x} = A \\setminus b = (A^T A)^{-1} A^T b$ where\n",
    "\n",
    "$$\n",
    "A = \\begin{pmatrix} 1 & a_1 \\\\ 1 & a_2 \\\\ \\vdots & \\vdots \\\\ 1 & a_m \\end{pmatrix}\n",
    "$$\n",
    "\n",
    "Let's do that for 1000 data points given by a line $1 + 2a$ plus some noise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7413cb8b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element Vector{Float64}:\n",
       " 0.994280300046284\n",
       " 1.9584555442499423"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = 1000\n",
    "a = randn(m)\n",
    "b = @. 1 + 2a + randn()\n",
    "A = [one.(a) a]\n",
    "x̂ = A \\ b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f28865a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47444 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 640x480 with 1 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "PyObject <matplotlib.legend.Legend object at 0x2905f5370>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "title(\"ordinary least-square fit\")\n",
    "plot(a, b, \"b.\")\n",
    "plot(a, x̂[1] .+ x̂[2] .* a, \"r-\")\n",
    "legend([L\"data $1+2a+\\mathrm{noise}$\", L\"fit $%$(round(x̂[1],digits=3)) + %$(round(x̂[2],digits=3))a$\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c168087c",
   "metadata": {},
   "source": [
    "## Gradient descent\n",
    "\n",
    "We can instead solve this problem by gradient descent with a fixed learning rate $s$.  At each iteration\n",
    "$$\n",
    "x^{(k+1)} = x^{(k)} - s \\nabla f^{(k)}\n",
    "$$\n",
    "where for $f(x) = \\frac{1}{m}\\Vert b - Ax \\Vert^2 = \\frac{1}{m}(b-Ax)^T (b-Ax)$ we have $\\nabla f = \\frac{2}{m}A^T(Ax - b)$.   We'll start with an initial \"guess\" of $x^{(1)} = [0,0]$.\n",
    "\n",
    "This is sometimes called **\"full batch\" gradient descent\"**: we are looking at the **whole \"batch\"** of training data at once.\n",
    "\n",
    "We'll use the [Interact.jl package](https://github.com/JuliaGizmos/Interact.jl) to visualize the progress of this algorithm as we vary the \"learning rate\" $s$.\n",
    "\n",
    "For goodd convergence, we should typically have $s < 1/L$ where $L \\approx 2$ is the Lipschitz constant of $f$, corresponding to the largest eigenvalue of the Hessian $2A^T A / m$.  So, for convenience we'll specify $sL$ instead of $s$.\n",
    "\n",
    "For $s > 1/L$ it \"overshoots\" the optimum and starts to oscillate, but it still converges if it overshoots by less than a factor of 2 (so each step still ends up closer to the optimum).  If $s > 2/L$ then it overshoots by more than a factor of 2 (each step ends up *farther* from the optimum) and iterations diverge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4b186c51",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "L = (2 / m) * maximum(eigvals(A' * A)) = 2.01359649631084\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83312 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 1000x400 with 2 Axes>)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(x) = norm(b-A*x)^2 / m\n",
    "f(x1,x2) = f([x1,x2])\n",
    "x1 = range(-2,3, length=100)\n",
    "x2 = range(0,4,length=100)\n",
    "fvals = f.(x1',x2)\n",
    "\n",
    "@show L = (2/m) * maximum(eigvals(A'*A)) # Lipschitz constant\n",
    "\n",
    "f1 = figure(figsize=(10,4))\n",
    "#@manipulate for sL in slider(0.05:0.05:2.1, value=0.3, label=\"s × L\"),\n",
    "#                iterations in slider(10:100, value=20, label=\"# iterations\")\n",
    "let sL=0.3, iterations=20\n",
    "    withfig(f1) do\n",
    "        s = sL / L\n",
    "        x = [0.0,0.0]\n",
    "        X = [x]\n",
    "        ΔX = [norm(x - x̂)]\n",
    "        for i = 1:iterations\n",
    "            ∇f = (2/m)*(A'*(A*x - b))\n",
    "            x = x - s * ∇f\n",
    "            push!(X, x)\n",
    "            push!(ΔX, norm(x - x̂))\n",
    "        end\n",
    "        \n",
    "        subplot(1,2,1)\n",
    "        plt.subplots_adjust(wspace=0.4)\n",
    "        contour(x1', x2, fvals)\n",
    "        plot(x̂[1], x̂[2], \"r*\")\n",
    "        plot(getindex.(X,1), getindex.(X,2), \"ko-\")\n",
    "        xlabel(L\"x_1\")\n",
    "        ylabel(L\"x_2\")\n",
    "        title(\"gradient-descent iterates x\")\n",
    "        \n",
    "        subplot(1,2,2)\n",
    "        semilogy(ΔX, \"ko-\")\n",
    "        ylim(1e-10, 100)\n",
    "        xlabel(\"iteration\")\n",
    "        ylabel(L\"\\Vert x - \\hat{x} \\Vert\")\n",
    "        title(\"|error| in x\")\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b4d9fa9",
   "metadata": {},
   "source": [
    "## Stochastic gradient descent\n",
    "\n",
    "In the most basic form of stochastic gradient descent, we look at one data point at a time, chosen at random.  More generally, we can look at \"mini-batches\" of $M$ data points at a time.  That is, we will only **estimate**\n",
    "$$\n",
    "\\nabla f = E[\\nabla f_i] \\approx \\frac{1}{M} \\sum_{i \\in \\text{mini-batch}} \\nabla f_i\n",
    "$$\n",
    "The simplest case is $M=1$: we only look at one data point at a time!\n",
    "\n",
    "This is *much* cheaper than computing $\\nabla f$ when the size of the training is large.   (In fact, we can even handle *infinite* training sets this way, where the training data is generated randomly as needed.)\n",
    "\n",
    "Let's see how it works for the problem above, sampling data at random from an infinite ($m=\\infty$) training set generated as needed from the same model as above, again using the simplest possible algorithm with a fixed step size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cc7d34e5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88368 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 1000x400 with 2 Axes>)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x̂₀ = [1,2] # ground-truth model\n",
    "\n",
    "f2 = figure(figsize=(10,4))\n",
    "#@manipulate for sL in slider(0.01:0.01:0.4, value=0.2, label=\"s × L\"),\n",
    "#                M in slider(1:100, value=1, label=\"mini-batch M\"),\n",
    "#                iterations in slider(10:100, value=50, label=\"# iterations\")\n",
    "let sL = 0.1, M=1, iterations=50\n",
    "    withfig(f2) do\n",
    "        s = sL / L\n",
    "        x = [0.0,0.0]\n",
    "        X = [x]\n",
    "        ΔX = [norm(x - x̂₀)]\n",
    "        for i = 1:iterations\n",
    "            # generate a mini-batch of size M:\n",
    "            ã = randn(M)\n",
    "            b̃ = @. 1 + 2ã + randn()\n",
    "            Ã = [one.(ã) ã]\n",
    "            \n",
    "            # approximate gradient from this mini-batch\n",
    "            ∇f̃ = (2/M)*(Ã'*(Ã*x - b̃))\n",
    "            x = x - s * ∇f̃\n",
    "            push!(X, x)\n",
    "            push!(ΔX, norm(x - x̂₀))\n",
    "        end\n",
    "        \n",
    "        subplot(1,2,1)\n",
    "        plt.subplots_adjust(wspace=0.4)\n",
    "        contour(x1', x2, fvals)\n",
    "        plot(x̂₀[1], x̂₀[2], \"r*\")\n",
    "        plot(getindex.(X,1), getindex.(X,2), \"ko-\")\n",
    "        xlabel(L\"x_1\")\n",
    "        ylabel(L\"x_2\")\n",
    "        title(\"gradient-descent iterates x\")\n",
    "        \n",
    "        subplot(1,2,2)\n",
    "        semilogy(ΔX, \"ko-\")\n",
    "        ylim(1e-2, 2)\n",
    "        xlabel(\"iteration\")\n",
    "        ylabel(L\"\\Vert x - \\hat{x} \\Vert\")\n",
    "        title(\"|error| in x\")\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d572f97e",
   "metadata": {},
   "source": [
    "Some things to note:\n",
    "\n",
    "* Far from the solution, SGD converges quickly, but then starts to \"bounce around\" the minimum.\n",
    "* The process **converges \"on average\"**: $E[x^{(k)}]$ (the average after repeating SGD many times) converges similarly to exact/full-batch gradient descent.  The difference is that now we have **variance**.  (But the variance is not necessarily a bad thing, as discussed below, because it can suppress over-fitting!)\n",
    "\n",
    "One thing to notice is that increasing the mini-batch size $M$ has a very similar qualitative effect to decreasing the learning rate $s$.   The intuitive reason for this is that taking $M$ tiny steps with a single sample each (which hardly changes $x$, so the gradients of the different samples are *almost* at the same $x$) is effectively almost the same as taking a single larger step with $M$ samples (evaluated at the same $x$).\n",
    "\n",
    "Let's try this for 2000 iterations with $M=1$ and $s = 0.2/L$ vs. 100 iterations with $M=20$ and $s=0.2/20L = 0.01/L$.  (Notice that the total number of samples is the same in both cases!  In the former case we will only plot every 20th step so that the graphs are comparable.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b85c0555",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 118356 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 1000x400 with 2 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "PyObject Text(0.5, 1.0, '|error| in x')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sL = 0.2/20\n",
    "M = 1\n",
    "iterations = 20*100\n",
    "s = sL / L\n",
    "x = [0.0,0.0]\n",
    "X1 = [x]\n",
    "ΔX1 = [norm(x - x̂₀)]\n",
    "for i = 1:iterations\n",
    "    # generate a mini-batch of size M:\n",
    "    ã = randn(M)\n",
    "    b̃ = @. 1 + 2ã + randn()\n",
    "    Ã = [one.(ã) ã]\n",
    "\n",
    "    # approximate gradient from this mini-batch\n",
    "    ∇f̃ = (2/M)*(Ã'*(Ã*x - b̃))\n",
    "    x = x - s * ∇f̃\n",
    "    if i % 20 == 0\n",
    "        push!(X1, x)\n",
    "        push!(ΔX1, norm(x - x̂₀))\n",
    "    end\n",
    "end\n",
    "\n",
    "\n",
    "sL = 0.2\n",
    "M = 20\n",
    "iterations = 100\n",
    "s = sL / L\n",
    "x = [0.0,0.0]\n",
    "X20 = [x]\n",
    "ΔX20 = [norm(x - x̂₀)]\n",
    "for i = 1:iterations\n",
    "    # generate a mini-batch of size M:\n",
    "    ã = randn(M)\n",
    "    b̃ = @. 1 + 2ã + randn()\n",
    "    Ã = [one.(ã) ã]\n",
    "\n",
    "    # approximate gradient from this mini-batch\n",
    "    ∇f̃ = (2/M)*(Ã'*(Ã*x - b̃))\n",
    "    x = x - s * ∇f̃\n",
    "    push!(X20, x)\n",
    "    push!(ΔX20, norm(x - x̂₀))\n",
    "end\n",
    "\n",
    "\n",
    "figure(figsize=(10,4))\n",
    "suptitle(\"M = 1, s = 0.01/L vs. M = 20, s = 0.2/L\")\n",
    "\n",
    "subplot(1,2,1)\n",
    "plt.subplots_adjust(wspace=0.4)\n",
    "contour(x1', x2, fvals)\n",
    "plot(getindex.(X1,1), getindex.(X1,2), \"ko-\")\n",
    "plot(getindex.(X20,1), getindex.(X20,2), \"gs-\")\n",
    "legend([\"M=1, s=0.01/L\", \"M=20, s=0.2/L\"])\n",
    "xlabel(L\"x_1\")\n",
    "ylabel(L\"x_2\")\n",
    "title(\"gradient-descent iterates x\")\n",
    "\n",
    "subplot(1,2,2)\n",
    "semilogy(ΔX1, \"ko-\")\n",
    "semilogy(ΔX20, \"gs-\")\n",
    "legend([\"M=1, s=0.01/L\", \"M=20, s=0.2/L\"])\n",
    "ylim(1e-2, 2)\n",
    "xlabel(\"samples / 20\")\n",
    "ylabel(L\"\\Vert x - \\hat{x} \\Vert\")\n",
    "title(\"|error| in x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c6b619f",
   "metadata": {},
   "source": [
    "**Why use mini-batches at all, then?**  Why not always use $M=1$ and just shrink $s$ as needed?\n",
    "\n",
    "The answer is simple: **parallel computing**.   The big advantage of a mini-batch is that then you can evaluate $\\nabla f_i$ **in parallel** for each sample in the mini-batch, e.g. using a multi-core computer, GPUs, or distributed-memory clusters.\n",
    "\n",
    "At first glance, it may seem that we would want to increase $M$ as much as possible, within the limits of our compute power, since the \"ideal\" is surely to optimize over the whole training set at once, right?   Actually, no!   The danger of increasing $M$ by too much (or equivalently, making $s$ too small), is that it can easily lead to **over-fitting**.   The fact that stochastic optimization with small batches tends to \"bounce around\" the nominal optimum can be a benefit rather than a problem in machine-learning problems, because it tends to prevent you from over-fitting to the training data and failing to **generalize** well to other test data.\n",
    "\n",
    "* *In this simple linear-regression example, there are only 2 parameters x, so we are not in much danger of overfitting.  This is more of a concern with huge models like neural networks that have an enormous number of parameters.*\n",
    "\n",
    "There are many variations and refinements of mini-batching for stochastic gradient descent that are being explored in the machine-learning literature, and many open problems.  For example:\n",
    "\n",
    "* If you have a set of $m$ training data points, the statistical theory of mini-batching is clearest if you draw samples from this set *independently*, which means that you might occasionally draw the same sample more than once (i.e. you are \"sampling with replacement\").  However, for performance reasons nearly all machine learning is done by randomly shuffling the test data once and then streaming through it *in order* (*sequential* access = fast!), corresponding to \"sampling without replacement\".  Unfortunately, it's *much* harder to prove theoretical guarantees for without-replacement sampling; for recent progress, see e.g. [Shamir (2016)](https://arxiv.org/abs/1603.00570).\n",
    "* Not all training data is created equal!   Some data points, e.g. more \"unusual\" data points, might contribute more to $\\nabla f$ and hence be more \"informative\" about how to update the parameters $x$.   This leads to algorithms that try to sample training data non-uniformly, such as **importance sampling** (e.g. [Csiba and Richtárik (2016)](https://arxiv.org/abs/1602.02283)) and **active learning** (e.g. [Chakraborty et al. (2015)](https://doi.org/10.1109/TPAMI.2015.2389848) or [Lakshiminarayanan et al (2016)](https://arxiv.org/abs/1612.01474)).\n",
    "* Various strategies have been explored to compensate for the generalization problem of large batch sizes.  See e.g. [Keskar et al. (2016)](https://arxiv.org/abs/1609.04836) or [Hoffer et al. (2017)](https://arxiv.org/abs/1705.08741).  Of course, one alternative is to parallelize *within each* $f_i$ rather than over different samples, e.g. as described in [Huang et al. (2018)](https://arxiv.org/abs/1811.06965).\n",
    "* In training a neural network that has multiple \"layers\" (i.e. $f_i$ is a composition of several functions), it is common to use the mini-batch in order to insert additional rescaling transformations into each layer, an algorithm known as [batch normalization](https://en.wikipedia.org/wiki/Batch_normalization)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdb14cba",
   "metadata": {},
   "source": [
    "## The Adam algorithm for stochastic optimization\n",
    "\n",
    "A popular algorithm for stochastic optimization is [Adam (Kingma & Ba, 2015)](https://arxiv.org/pdf/1412.6980.pdf), which adds several features:\n",
    "\n",
    "* It has a \"momentum\" or \"memory\" term, where the gradient is estimated by a *weighted average* of the current mini-batch gradient and the previous gradient estimate.\n",
    "* It also estimates the *variance* of the gradient components in the same way, and essentially uses this as rough diagonal estimate of the Hessian: it takes a \"crude Newton\"-like step where the gradient components are weighed inversely by the square roots of their variances.\n",
    "\n",
    "It's quite simple to implement directly from the description in the paper linked above.  Let's try it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "77481ff3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 90240 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 1000x400 with 2 Axes>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x̂₀ = [1,2] # ground-truth model\n",
    "\n",
    "# typical Adam hyper-parameters\n",
    "ϵ = 1e-8\n",
    "β₁ = 0.9\n",
    "β₂ = 0.999\n",
    "\n",
    "f2 = figure(figsize=(10,4))\n",
    "#@manipulate for log₁₀α in slider(-4:0.2:0, value=-1, label=\"learning rate log₁₀α\"),\n",
    "#                M in slider(1:100, value=1, label=\"mini-batch M\"),\n",
    "#                iterations in slider(10:10:300, value=300, label=\"# iterations\")\n",
    "let log₁₀α=-1, M=1, iterations=300\n",
    "    withfig(f2) do\n",
    "        α = exp10(log₁₀α)\n",
    "        x = [0.0,0.0]\n",
    "        X = [x]\n",
    "        ΔX = [norm(x - x̂₀)]\n",
    "        mom = zero(x) # initial 1st-moment (mean gradient)\n",
    "        var = zero(x) # initial 2nd-moment (mean gradient.^2)\n",
    "        for t = 1:iterations\n",
    "            # generate a mini-batch of size M:\n",
    "            ã = randn(M)\n",
    "            b̃ = @. 1 + 2ã + randn()\n",
    "            Ã = [one.(ã) ã]\n",
    "            \n",
    "            # approximate gradient from this mini-batch\n",
    "            ∇f̃ = (2/M)*(Ã'*(Ã*x - b̃))\n",
    "            \n",
    "            # Adam update\n",
    "            mom = β₁*mom + (1-β₁)*∇f̃\n",
    "            var = β₂*var + (1-β₂)*(∇f̃.^2)\n",
    "            m̂ = mom / (1 - β₁^t) # normalize by total weights\n",
    "            v̂ = var / (1 - β₂^t) # normalize by total weights\n",
    "            x = @. x - α * m̂ / (√v̂ + ϵ)\n",
    "            \n",
    "            push!(X, x)\n",
    "            push!(ΔX, norm(x - x̂₀))\n",
    "        end\n",
    "        \n",
    "        subplot(1,2,1)\n",
    "        plt.subplots_adjust(wspace=0.4)\n",
    "        contour(x1', x2, fvals)\n",
    "        plot(x̂₀[1], x̂₀[2], \"r*\")\n",
    "        plot(getindex.(X,1), getindex.(X,2), \"ko-\")\n",
    "        xlabel(L\"x_1\")\n",
    "        ylabel(L\"x_2\")\n",
    "        title(\"gradient-descent iterates x\")\n",
    "        \n",
    "        subplot(1,2,2)\n",
    "        semilogy(ΔX, \"ko-\")\n",
    "        ylim(1e-2, 2)\n",
    "        xlabel(\"iteration\")\n",
    "        ylabel(L\"\\Vert x - \\hat{x} \\Vert\")\n",
    "        title(\"|error| in x\")\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5519b3b6",
   "metadata": {},
   "source": [
    "With correctly tuned hyperparameters, Adam converges a bit more smoothly — with *less variance* — than naive stochastic gradient descent, even for a mini-batch size $M=1$.   The \"memory\" effect of averaging the gradient over many iterations helps with this!\n",
    "\n",
    "I should also point out that this is a **very well-conditioned test problem** as evidenced by the nearly circular contours of $f(x)$ as seen in these examples.    In such cases naive gradient descent often works very well, and there is not as much benefit from either momentum terms or weighting the gradient inversely with the variance.   Many realistic learning problems are badly conditioned, however, and these enhancements of Adam (and related algorithms) may make a much bigger difference!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "aef8dcff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0393254191586012"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cond(A'A) # condition number of the Hessian"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "407ae852",
   "metadata": {},
   "source": [
    "(See, well-conditioned!)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3bfd367",
   "metadata": {},
   "source": [
    "## Other algorithms and generalizations\n",
    "\n",
    "There are *many, many* other stochastic optimization algorithms, some of which are linked in the Wikipedia article on [stochastic gradient descent](https://en.wikipedia.org/wiki/Stochastic_gradient_descent), with many variants on ways to choose parameters like learning rates and momentum coefficients.  This includes many algorithms similar in spirit to Adam and accelerated gradient descent, such as the [YOGI algorithm (Zaheer et al., 2018)](https://papers.nips.cc/paper_files/paper/2018/hash/90365351ccc7437a1309dc64e4db32a3-Abstract.html) which has better convergence guarantees than Adam; fancier quasi-Newton algorithms (see e.g. [Zhao et al. (2017)](https://arxiv.org/abs/1704.00116) and references therein); and also algorithms for more complicated optimization problems than simply minimizing $f(x)$: stochastic optimization with nonlinear constraints (see e.g. [Lan and Zhou (2016)](https://arxiv.org/abs/1604.03887), stochastic \"minimax\" problems (e.g. [Yang et al. (2022)](https://arxiv.org/abs/2202.07868)), and robust stochastic optimization (e.g. [Chen et al (2020)](https://doi.org/10.1287/mnsc.2020.3603)).\n",
    "\n",
    "Note that all of these algorithms can be trivially modified to accommodate objective functions that have a non-stochastic part, e.g.:\n",
    "$$\n",
    "f(x) = f_0(x) + E[f_i(x)] \\, ,\n",
    "$$\n",
    "since then $\\nabla f = \\nabla f_0 + E[\\nabla f_i]$ and you can estimate $\\nabla f$ by computing $\\nabla f_0$ *exactly* and then estimating $E[\\nabla f_i]$ from a mini-batch sample.   For example, this arises in [Tikhonov-regularized least-squares](https://en.wikipedia.org/wiki/Ridge_regression):\n",
    "$$\n",
    "f(x) = \\underbrace{\\lambda \\Vert x \\Vert_2^2}_{f_0(x)} + \\frac{1}{m} \\sum_i \\underbrace{\\left\\Vert \\text{model}(x, \\text{input}_i) - \\text{output}_i \\right\\Vert_2^2}_{f_i(x)}\\, ,\n",
    "$$\n",
    "where $\\nabla f_0 = 2\\lambda x$ is easy and cheap to compute since it is independent of $m$."
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.8.3",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
