{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-1 fw=all -->\n",
    "\n",
    "# Stochastic and Adaptive Methods\n",
    "\n",
    "The optimizer that trains essentially every modern network is AdamW with warmup\n",
    "and a decaying learning-rate schedule, not the plain gradient descent whose\n",
    "theory that section developed. This section\n",
    "closes the gap between that theory and that default. Two questions were left\n",
    "open there. First, the smooth nonconvex benchmark used to reason about\n",
    "deep networks, the $O(1/K)$ stationarity rate, was proved for the exact\n",
    "gradient, while training uses noisy minibatch estimates; what survives?\n",
    "Second, the whole $\\kappa$ story was about a *single* step size $\\eta$ forced\n",
    "to respect the stiffest curvature direction. The obvious escape, a *separate*\n",
    "step size per coordinate, is exactly what AdaGrad, RMSProp, and Adam\n",
    "implement, and it deserves the same treatment the global $\\eta$ received: a\n",
    "derivation, a rate, and an account of where it fails. Along the way we\n",
    "derive Adam's bias correction rather than assert it, state and verify the\n",
    "counterexample showing that Adam can converge to the *worst* point of a convex\n",
    "problem, work out why weight decay through a preconditioner is not weight\n",
    "decay, and give learning-rate schedules and warmup the mathematical reading\n",
    "they usually go without.\n",
    "\n",
    "We lean on that section throughout (the\n",
    "descent lemma, the condition number $\\kappa$, the minibatch variance\n",
    "$\\mathrm{tr}\\,\\Sigma/b$, and the noise ball are all consumed here) and on\n",
    "that section for the diagonal-metric view of\n",
    "preconditioning. The MAP reading of weight decay comes from\n",
    "that section, its constraint reading from\n",
    "that section, and the closing preconditioning\n",
    "ladder reaches back to the SVD and Newton--Schulz material of\n",
    "that section. Standard references are\n",
    "@Bottou.Curtis.Nocedal.2018 for the stochastic theory and the original\n",
    "papers cited as we go. As everywhere in this chapter, the code is plain NumPy:\n",
    "every optimizer below is under ten lines written by hand, which is the fastest\n",
    "way to see there is no magic in any of them."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "adaptive-stochastic-methods-imports"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "import numpy as np"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-2 fw=all -->\n",
    "\n",
    "## SGD Without Convexity\n",
    "\n",
    "### What the Deterministic Theorem Leaves Open\n",
    "\n",
    "Recall the two bookends from that section.\n",
    "For $L$-smooth $f$ bounded below, *deterministic* gradient descent satisfies\n",
    "$\\min_{k < K} \\|\\nabla f(\\mathbf{x}_k)\\|^2 \\le 2L(f(\\mathbf{x}_0) - f^\\star)/K$:\n",
    "stationarity at rate $O(1/K)$, no convexity needed. And for *stochastic*\n",
    "gradients on a strongly convex quadratic, a fixed step converges to a noise\n",
    "ball of squared radius $\\approx \\eta\\sigma^2/(2\\lambda)$ around the minimizer\n",
    "the equation. Training a\n",
    "network lives in the intersection of the two hypotheses, nonconvex *and*\n",
    "noisy, which neither result covers. The theorem that does is due to\n",
    "@Ghadimi.Lan.2013, and its proof is nothing more than the descent\n",
    "lemma with expectations taken at the right moment.\n",
    "\n",
    "Model one SGD step as $\\mathbf{x}_{k+1} = \\mathbf{x}_k - \\eta\\,\\mathbf{g}_k$,\n",
    "where the stochastic gradient is unbiased with bounded variance:\n",
    "\n",
    "$$\n",
    "\\mathbb{E}\\left[\\mathbf{g}_k \\mid \\mathbf{x}_k\\right] = \\nabla f(\\mathbf{x}_k),\n",
    "\\qquad\n",
    "\\mathbb{E}\\left[\\|\\mathbf{g}_k - \\nabla f(\\mathbf{x}_k)\\|^2 \\mid \\mathbf{x}_k\\right] \\;\\le\\; \\sigma^2.\n",
    "$$\n",
    "\n",
    "The minibatch proposition of that section\n",
    "supplies exactly these hypotheses, with $\\sigma^2 = \\mathrm{tr}\\,\\Sigma/b$.\n",
    "\n",
    "### The Ghadimi--Lan Rate\n",
    "\n",
    "**Proposition (SGD finds approximate stationary points).** *Let $f$ be\n",
    "$L$-smooth and bounded below by $f^\\star$, write\n",
    "$\\Delta = f(\\mathbf{x}_0) - f^\\star$, and run SGD with a constant step\n",
    "$\\eta \\le 1/L$. Then*\n",
    "\n",
    "$$\n",
    "\\frac{1}{K} \\sum_{k=0}^{K-1} \\mathbb{E}\\left[\\|\\nabla f(\\mathbf{x}_k)\\|^2\\right]\n",
    "\\;\\le\\; \\frac{2\\Delta}{\\eta K} + L\\eta\\sigma^2,\n",
    "$$\n",
    "\n",
    "*and choosing $\\eta = \\min\\bigl(1/L,\\ \\sqrt{2\\Delta/(L\\sigma^2 K)}\\bigr)$\n",
    "balances the two terms:*\n",
    "\n",
    "$$\n",
    "\\mathbb{E}\\left[\\|\\nabla f(\\mathbf{x}_R)\\|^2\\right]\n",
    "\\;\\le\\; \\frac{2L\\Delta}{K} + 2\\sigma\\,\\sqrt{\\frac{2L\\Delta}{K}},\n",
    "\\qquad R \\sim \\mathrm{Uniform}\\{0, \\ldots, K-1\\}.\n",
    "$$\n",
    "\n",
    "**Proof.** Apply the quadratic ceiling\n",
    "the equation to the step actually taken,\n",
    "$\\mathbf{x}_{k+1} - \\mathbf{x}_k = -\\eta\\,\\mathbf{g}_k$:\n",
    "\n",
    "$$\n",
    "f(\\mathbf{x}_{k+1}) \\;\\le\\; f(\\mathbf{x}_k) - \\eta\\, \\nabla f(\\mathbf{x}_k)^\\top \\mathbf{g}_k + \\tfrac{L\\eta^2}{2}\\, \\|\\mathbf{g}_k\\|^2.\n",
    "$$\n",
    "\n",
    "Take the expectation conditioned on $\\mathbf{x}_k$. Unbiasedness turns the\n",
    "middle term into $-\\eta\\,\\|\\nabla f(\\mathbf{x}_k)\\|^2$, and the variance\n",
    "decomposition\n",
    "$\\mathbb{E}\\|\\mathbf{g}_k\\|^2 = \\|\\nabla f(\\mathbf{x}_k)\\|^2 + \\mathbb{E}\\|\\mathbf{g}_k - \\nabla f(\\mathbf{x}_k)\\|^2$\n",
    "bounds the last by\n",
    "$\\tfrac{L\\eta^2}{2}(\\|\\nabla f(\\mathbf{x}_k)\\|^2 + \\sigma^2)$. Collecting,\n",
    "\n",
    "$$\n",
    "\\mathbb{E}\\left[f(\\mathbf{x}_{k+1}) \\mid \\mathbf{x}_k\\right]\n",
    "\\;\\le\\; f(\\mathbf{x}_k) - \\eta\\left(1 - \\tfrac{L\\eta}{2}\\right) \\|\\nabla f(\\mathbf{x}_k)\\|^2 + \\tfrac{L\\eta^2 \\sigma^2}{2},\n",
    "$$\n",
    "\n",
    "the descent lemma with one new term: a *noise tax* of\n",
    "$\\tfrac{L}{2}\\eta^2\\sigma^2$ per step, paid whether or not the gradient is\n",
    "large. For $\\eta \\le 1/L$ the bracket is at least $\\tfrac12$. Take total\n",
    "expectations, telescope over $k = 0, \\ldots, K-1$ (the left side collapses to\n",
    "$\\mathbb{E}[f(\\mathbf{x}_K)] - f(\\mathbf{x}_0) \\ge -\\Delta$), and divide by\n",
    "$\\eta K/2$ to get the equation. The first term of the bound\n",
    "falls in $\\eta$, the second grows; the stated $\\eta$ equalizes them, and the\n",
    "average over $k$ *is* the expectation at a uniformly random iterate $R$.\n",
    "$\\blacksquare$\n",
    "\n",
    "Read the optimized bound against its deterministic ancestor. Without noise\n",
    "($\\sigma = 0$) the second term vanishes and the $2L\\Delta/K$ term reproduces\n",
    "the $O(1/K)$ rate exactly. With noise, the second term dominates for large\n",
    "$K$, and it decays like $K^{-1/2}$: **noise turns the deterministic $1/K$\n",
    "into a square root**, the same square root that governed the minibatch\n",
    "variance ($1/b$ energy, $1/\\sqrt{b}$ amplitude). Be careful about what is\n",
    "bounded: the theorem controls the **expected squared gradient norm**. Running\n",
    "$100\\times$ longer reduces that bound by $10\\times$; interpreting it as a bound\n",
    "on the gradient norm itself gives only a factor $\\sqrt{10}$. A tenfold reduction\n",
    "of that norm would require $10{,}000\\times$ the budget under this worst-case\n",
    "rate. The prescribed step obeys the same\n",
    "logic as the noise-ball analysis: a longer budget affords a smaller $\\eta$,\n",
    "shrinking the floor the iterates can reach.\n",
    "\n",
    "One subtlety is a feature of the theorem rather than a defect. The\n",
    "deterministic result bounded $\\min_k \\|\\nabla f(\\mathbf{x}_k)\\|^2$: you may\n",
    "keep the best iterate, because you can *evaluate* which one it is. With\n",
    "stochastic gradients you never observe $\\|\\nabla f(\\mathbf{x}_k)\\|$, only\n",
    "noisy estimates, so \"return the best iterate\" is not an executable\n",
    "instruction. The theorem's output is therefore a *randomly selected* iterate,\n",
    "whose expected squared gradient norm is the average that\n",
    "the equation controls. The cell verifies precisely this\n",
    "distinction on a nonconvex toy (a one-neuron $\\tanh$ regression with noisy\n",
    "labels, single-example gradients), sweeping the budget $K$ with the\n",
    "theory-prescribed $\\eta \\propto 1/\\sqrt{K}$ and measuring both the\n",
    "random-iterate average and the (non-executable) trajectory minimum of\n",
    "$\\|\\nabla f\\|^2$, each averaged over $20$ seeds."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "adaptive-stochastic-methods-ghadimi-lan"
   },
   "source": [
    "n, d = 64, 5\n",
    "rng = np.random.default_rng(0)\n",
    "Xd = rng.normal(size=(n, d))\n",
    "yd = np.tanh(Xd @ rng.normal(size=d)) + 0.5 * rng.normal(size=n)\n",
    "\n",
    "def full_grad(w):                              # exact gradient, for measurement\n",
    "    t = np.tanh(Xd @ w)\n",
    "    return 2 * (Xd * ((t - yd) * (1 - t**2))[:, None]).mean(axis=0)\n",
    "\n",
    "def stoch_grad(w, i):                          # single-example gradient (b = 1)\n",
    "    t = np.tanh(Xd[i] @ w)\n",
    "    return 2 * (t - yd[i]) * (1 - t**2) * Xd[i]\n",
    "\n",
    "Ks, avg_res, min_res = [125, 500, 2000, 8000], [], []\n",
    "print('    K    eta_K     E|grad(x_R)|^2   min-so-far')\n",
    "for K in Ks:\n",
    "    eta = 0.5 / np.sqrt(K)                     # the theorem's eta ~ 1/sqrt(K)\n",
    "    avgs, mins = [], []\n",
    "    for seed in range(20):\n",
    "        rg, w, tot, m = np.random.default_rng(seed), np.zeros(d), 0.0, np.inf\n",
    "        for k in range(K):\n",
    "            g2 = (full_grad(w)**2).sum()\n",
    "            tot, m = tot + g2, min(m, g2)\n",
    "            w = w - eta * stoch_grad(w, rg.integers(0, n))\n",
    "        avgs.append(tot / K), mins.append(m)\n",
    "    avg_res.append(np.mean(avgs)), min_res.append(np.mean(mins))\n",
    "    print(f'{K:5d}   {eta:.4f}      {avg_res[-1]:.3e}      {min_res[-1]:.3e}')\n",
    "print(f'log-log slope, random iterate: '\n",
    "      f'{np.polyfit(np.log(Ks), np.log(avg_res), 1)[0]:.2f}  (theory: -1/2)')\n",
    "print(f'log-log slope, min-so-far:     '\n",
    "      f'{np.polyfit(np.log(Ks), np.log(min_res), 1)[0]:.2f}')"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "    K    eta_K     E|grad(x_R)|^2   min-so-far\n  125   0.0447      2.134e-01      9.385e-03\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  500   0.0224      8.485e-02      1.398e-03\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " 2000   0.0112      3.774e-02      2.973e-04\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      " 8000   0.0056      1.729e-02      5.428e-05\nlog-log slope, random iterate: -0.60  (theory: -1/2)\nlog-log slope, min-so-far:     -1.23\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-3 fw=all -->\n",
    "\n",
    "The random-iterate column scales as promised: a $64\\times$ increase in budget\n",
    "yields a $12\\times$ reduction, log-log slope $-0.60$, within reach of the\n",
    "theoretical $-1/2$. (The residual steepness comes from the average itself:\n",
    "the trajectory average still contains the early high-gradient iterates, whose\n",
    "contribution decays like $1/K$. The bound's transient $2\\Delta/(\\eta K)$ term\n",
    "is innocent here: with the cell's $\\eta \\propto 1/\\sqrt{K}$ it is of the same\n",
    "$K^{-1/2}$ order as the noise term, so it cannot steepen the slope.) The\n",
    "min-so-far column falls much faster,\n",
    "slope $-1.23$: on this benign toy the best iterate beats the guarantee by a\n",
    "wide margin, exactly why one would love to return it and exactly what the\n",
    "noise forbids you from *identifying*. Worst-case bounds are contracts: the\n",
    "guaranteed rate here is $K^{-1/2}$, and the schedule designs of this section\n",
    "build on it rather than on the lucky minimum.\n",
    "\n",
    "## Per-Coordinate Step Sizes\n",
    "\n",
    "### One Global Step Must Respect the Stiffest Mode\n",
    "\n",
    "In that section, stability chained the\n",
    "single step size to the stiffest mode ($\\eta < 2/\\lambda_{\\max}$), so the\n",
    "flattest mode advances by only $\\eta\\lambda_{\\min} \\approx 2/\\kappa$ per\n",
    "step: one scalar doing a matrix's job, at cost linear in $\\kappa$. But\n",
    "suppose the coordinate axes happen to *be* the eigendirections, as they are\n",
    "for\n",
    "$f(\\mathbf{x}) = \\tfrac12 \\mathbf{x}^\\top \\mathrm{diag}(\\boldsymbol{\\lambda})\\, \\mathbf{x}$,\n",
    "and we allow a *separate* step size $\\eta_i$ per coordinate. Then\n",
    "$\\eta_i = 1/\\lambda_i$ contracts every mode to zero in one step: a diagonal\n",
    "matrix of step sizes is a diagonal Newton's method, and the condition number\n",
    "simply disappears from the problem. This is the entire thesis of adaptive\n",
    "methods, drawn in the figure: gradient descent's\n",
    "zig-zag comes from using one number where $n$ are called for, and it\n",
    "disappears once each coordinate gets its own.\n",
    "\n",
    "![The section's thesis, drawn. Both panels show the same elongated quadratic ($\\kappa = 20$) from the same start, and both paths are actual runs of the stated update. Left: gradient descent with one global step size near the stability ceiling zig-zags across the stiff axis while crawling along the slow one. Right: Adam's per-coordinate normalization (with $\\beta_1 = 0$, to isolate the rescaling from momentum) moves both coordinates at comparable speed and curves smoothly into the minimum: the valley is effectively rounded, at the cost of estimating the rescaling from gradients alone.](../img/mdl-opt-per-coordinate.svg)\n",
    "\n",
    "Two obstacles stand between this observation and an algorithm. The\n",
    "eigenvalues are unknown, and off the quadratic model they vary from place to\n",
    "place; and with stochastic gradients we never see clean curvature, only noisy\n",
    "first-order information. Everything in the adaptive family is a scheme for\n",
    "estimating a useful diagonal rescaling *from the gradients themselves*.\n",
    "\n",
    "### AdaGrad: Calibrating Steps by Accumulated Evidence\n",
    "\n",
    "The first such scheme, **AdaGrad** [@Duchi.Hazan.Singer.2011] (shown in\n",
    "action in that section), keeps a\n",
    "running sum of squared gradients per coordinate and divides by its square\n",
    "root:\n",
    "\n",
    "$$\n",
    "\\mathbf{s}_t = \\mathbf{s}_{t-1} + \\mathbf{g}_t^2,\n",
    "\\qquad\n",
    "\\mathbf{x}_{t+1} = \\mathbf{x}_t - \\frac{\\eta}{\\sqrt{\\mathbf{s}_t} + \\epsilon}\\; \\mathbf{g}_t,\n",
    "$$\n",
    "\n",
    "with all operations elementwise and $\\epsilon$ a small constant that keeps the\n",
    "division finite. Two independent derivations land on this rule.\n",
    "\n",
    "*The metric view.* Exercise 2 of\n",
    "that section showed that steepest descent is\n",
    "relative to a choice of norm: measuring step length by\n",
    "$\\|\\mathbf{d}\\|_A = \\sqrt{\\mathbf{d}^\\top A\\, \\mathbf{d}}$ makes the steepest\n",
    "direction $-A^{-1}\\nabla f$, a *preconditioned* gradient. AdaGrad is steepest\n",
    "descent in the metric $A_t = \\mathrm{diag}(\\sqrt{\\mathbf{s}_t})$: it declares\n",
    "a coordinate \"expensive to move\" in proportion to the evidence\n",
    "$\\sqrt{\\sum_{s \\le t} g_{s,i}^2}$ that its gradient has been large. Newton\n",
    "would use the true curvature $\\mathrm{diag}(\\lambda_i)$; AdaGrad substitutes\n",
    "the cheapest available proxy, the gradient's own history (Exercise 1 derives\n",
    "the update from this view).\n",
    "\n",
    "*The regret view, and why the square root.* AdaGrad was invented for\n",
    "**online convex optimization**: convex losses $f_1, f_2, \\ldots$ arrive one\n",
    "at a time, the algorithm must play $\\mathbf{x}_t$ *before* seeing $f_t$, and\n",
    "performance is measured by **regret**, the gap\n",
    "$\\sum_{t=1}^{T} f_t(\\mathbf{x}_t) - \\min_{\\mathbf{x}} \\sum_{t=1}^{T} f_t(\\mathbf{x})$\n",
    "between the accumulated loss and that of the best *fixed* point chosen in\n",
    "hindsight; an algorithm learns, in this sense, when its average regret\n",
    "(regret divided by $T$) vanishes as $T \\to \\infty$. The problems AdaGrad was\n",
    "built for have **sparse features**: think of a bag-of-words\n",
    "model where the coordinate for a rare word receives a nonzero gradient once\n",
    "in ten thousand steps. A global $\\eta_t \\propto 1/\\sqrt{t}$ decays that rare\n",
    "coordinate's step size along with everyone else's, so by the time the rare\n",
    "evidence arrives the step is too small to use it. Per-coordinate\n",
    "accumulation fixes this: coordinate $i$'s effective step\n",
    "$\\eta/\\sqrt{s_{t,i}}$ decays with *its own* activity, not with wall-clock\n",
    "time, so rare-but-informative coordinates keep large steps. The square root\n",
    "is what the regret analysis selects: with the per-coordinate rates chosen\n",
    "this way, the accumulated-gradient terms in the regret bound sum to exactly\n",
    "$\\sum_i \\sqrt{\\sum_t g_{t,i}^2}$, a quantity that can beat the dimension-free\n",
    "$G\\sqrt{T}$ regret of plain SGD by a factor up to $\\sqrt{d}$ when gradients\n",
    "are sparse [@Duchi.Hazan.Singer.2011].\n",
    "\n",
    "The same accumulation is AdaGrad's defect outside the convex setting. Since\n",
    "$\\mathbf{s}_t$ never forgets, the effective step decays like\n",
    "$\\eta/(\\sigma\\sqrt{t})$ under persistent gradient noise: a\n",
    "Robbins--Monro-style schedule *hard-wired into the optimizer*. On a convex\n",
    "problem that is precisely what convergence demands; on a nonconvex landscape\n",
    "the optimizer may need to cross a plateau or a saddle region long after\n",
    "$\\mathbf{s}_t$ has grown large, and it arrives with steps ground down to\n",
    "nothing. The method stalls because it remembers too much.\n",
    "\n",
    "### RMSProp: Forgetting on Purpose\n",
    "\n",
    "The repair is to replace the all-time sum with an **exponential moving\n",
    "average**, the same fix that turns a cumulative mean into\n",
    "something that can track a moving target. **RMSProp**\n",
    "[@Tieleman.Hinton.2012] (shown in action in that section) keeps\n",
    "\n",
    "$$\n",
    "\\mathbf{v}_t = \\beta_2\\, \\mathbf{v}_{t-1} + (1 - \\beta_2)\\, \\mathbf{g}_t^2,\n",
    "\\qquad\n",
    "\\mathbf{x}_{t+1} = \\mathbf{x}_t - \\frac{\\eta}{\\sqrt{\\mathbf{v}_t} + \\epsilon}\\; \\mathbf{g}_t,\n",
    "$$\n",
    "\n",
    "so $\\mathbf{v}_t$ is a weighted average of recent squared gradients with an\n",
    "effective memory of about $1/(1 - \\beta_2)$ steps ($10$ at RMSProp's standard\n",
    "decay $\\beta_2 = 0.9$). The preconditioner now estimates the *current* gradient\n",
    "scale per coordinate rather than the lifetime total: steps no longer decay to\n",
    "zero by construction, and the method can keep moving on a nonstationary\n",
    "landscape. The tradeoff, invisible here but fatal below, is that forgetting\n",
    "discards exactly the kind of rare evidence AdaGrad was built to retain.\n",
    "\n",
    "### Adam: Momentum, Second Moments, and Bias Correction\n",
    "\n",
    "**Adam** [@Kingma.Ba.2014] (shown in action in that section)\n",
    "completes the recipe by applying the same exponential averaging to the\n",
    "gradient itself and correcting both averages for their startup bias. The\n",
    "gradient average $\\mathbf{m}_t$ is momentum\n",
    "(that section) in averaged form: where the\n",
    "heavy-ball buffer of that section *accumulates* gradients, the\n",
    "$(1 - \\beta_1)$ factor here makes $\\mathbf{m}_t$ a weighted average of them,\n",
    "the same method up to a rescaling of the step size:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathbf{m}_t &= \\beta_1\\, \\mathbf{m}_{t-1} + (1 - \\beta_1)\\, \\mathbf{g}_t,\n",
    "\\qquad &\n",
    "\\hat{\\mathbf{m}}_t &= \\mathbf{m}_t / (1 - \\beta_1^t), \\\\\n",
    "\\mathbf{v}_t &= \\beta_2\\, \\mathbf{v}_{t-1} + (1 - \\beta_2)\\, \\mathbf{g}_t^2,\n",
    "\\qquad &\n",
    "\\hat{\\mathbf{v}}_t &= \\mathbf{v}_t / (1 - \\beta_2^t), \\\\\n",
    "\\mathbf{x}_{t+1} &= \\mathbf{x}_t - \\eta\\; \\frac{\\hat{\\mathbf{m}}_t}{\\sqrt{\\hat{\\mathbf{v}}_t} + \\epsilon},\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "with defaults $\\beta_1 = 0.9$, $\\beta_2 = 0.999$, $\\epsilon = 10^{-8}$, and\n",
    "$\\mathbf{m}_0 = \\mathbf{v}_0 = \\mathbf{0}$. The default $\\beta_2 = 0.999$\n",
    "gives the second-moment estimate an effective memory of\n",
    "about $1/(1-\\beta_2) = 1000$ steps, a hundred times longer than RMSProp's\n",
    "customary ten. The bias correction is a two-line computation. Unroll the\n",
    "recursion:\n",
    "\n",
    "$$\n",
    "\\mathbf{v}_t = (1 - \\beta_2) \\sum_{s=1}^{t} \\beta_2^{\\,t-s}\\, \\mathbf{g}_s^2 .\n",
    "$$\n",
    "\n",
    "If the squared-gradient scale is (locally) stationary,\n",
    "$\\mathbb{E}[\\mathbf{g}_s^2] = \\bar{\\mathbf{g}^2}$ for all $s$, then by the\n",
    "geometric series $\\sum_{s=1}^{t} \\beta_2^{\\,t-s} = (1 - \\beta_2^t)/(1 - \\beta_2)$,\n",
    "\n",
    "$$\n",
    "\\mathbb{E}[\\mathbf{v}_t]\n",
    "= (1 - \\beta_2)\\, \\bar{\\mathbf{g}^2} \\sum_{s=1}^{t} \\beta_2^{\\,t-s}\n",
    "= \\left(1 - \\beta_2^t\\right) \\bar{\\mathbf{g}^2},\n",
    "$$\n",
    "\n",
    "so dividing by $1 - \\beta_2^t$ makes $\\hat{\\mathbf{v}}_t$ an *exactly*\n",
    "unbiased estimate of $\\bar{\\mathbf{g}^2}$ at every $t$: the correction\n",
    "factors cancel the zero initialization identically, at every step, rather\n",
    "than merely fading with it. The same computation with $\\beta_1$ handles\n",
    "$\\hat{\\mathbf{m}}_t$ (Exercise 2). The transient being corrected is\n",
    "large: at the defaults, $\\mathbf{v}_{10}$ carries only about $1\\%$ of its\n",
    "stationary value, and the raw ratio\n",
    "$\\mathbf{m}_t/(\\sqrt{\\mathbf{v}_t} + \\epsilon)$ mis-scales early steps by\n",
    "the factor $(1-\\beta_1^t)/\\sqrt{1-\\beta_2^t}$: already $3.16$ at\n",
    "$t = 1$, peaking above $6\\times$ near $t = 12$, because the momentum\n",
    "average saturates in about ten steps while the second-moment average needs\n",
    "about a thousand. the figure plots the transient\n",
    "and its cancellation.\n",
    "\n",
    "![Why Adam corrects for bias. Left: with zero initialization, the running average $\\mathbf{v}_t$ carries only the fraction $1-\\beta_2^t$ of the true squared-gradient scale; at $\\beta_2=0.999$ it has recovered just $63\\%$ after $1000$ steps, and dividing by $1-\\beta_2^t$ cancels the deficit exactly, at every $t$. Right: the resulting mis-scaling of the raw update $\\mathbf{m}_t/\\sqrt{\\mathbf{v}_t}$, the factor $(1-\\beta_1^t)/\\sqrt{1-\\beta_2^t}$: it is $3.16$ at $t=1$, peaks above $6\\times$ near $t=12$ (the numerator saturates in about ten steps, the denominator needs about a thousand), and decays back to $1$ only on the $1/(1-\\beta_2)$ timescale. Uncorrected Adam takes its *largest* steps precisely when its preconditioner is estimated from the fewest samples.](../img/mdl-opt-bias-correction.svg)\n",
    "\n",
    "Assembled, Adam is three ideas stacked: RMSProp's per-coordinate scale\n",
    "estimate, momentum's variance-averaged direction, and an exact startup\n",
    "correction for both. Near a minimum of a quadratic with diagonal Hessian, the\n",
    "scale estimate obeys\n",
    "$\\sqrt{\\hat{v}_i} \\approx |g_i| = \\lambda_i |x_i|$, so the update on\n",
    "coordinate $i$ is\n",
    "$\\eta\\, g_i / \\sqrt{\\hat{v}_i} \\approx \\eta\\,\\mathrm{sign}(x_i)$:\n",
    "constant-magnitude *sign descent*, where a diagonal Newton step would move by\n",
    "$x_i$ itself. What survives of the opening argument is the *ratio* between\n",
    "coordinates: the effective step $\\eta/(\\lambda_i |x_i|)$ scales like\n",
    "$1/\\lambda_i$ wherever the $|x_i|$ are comparable (as at the equal-coordinate\n",
    "start of the race below), so stiff coordinates take small steps and flat ones\n",
    "take large steps, estimated from first-order information alone.\n",
    "\n",
    "### When Adam Fails to Converge\n",
    "\n",
    "Adam's forgetting has a provable cost, found by\n",
    "@Reddi.Kale.Kumar.2019 after the method had already trained thousands\n",
    "of models: on some convex problems, Adam converges to the *worst* point.\n",
    "\n",
    "**Proposition (Adam need not converge).** *Consider online optimization over\n",
    "the interval $x \\in [-1, 1]$ with the periodic sequence of convex losses*\n",
    "\n",
    "$$\n",
    "f_t(x) =\n",
    "\\begin{cases}\n",
    "C\\,x, & t \\bmod 3 = 1, \\\\\n",
    "-x, & \\textrm{otherwise},\n",
    "\\end{cases}\n",
    "\\qquad C > 2,\n",
    "$$\n",
    "\n",
    "*whose gradients over one period sum to $C - 2 > 0$, so the best fixed point\n",
    "is $x^\\star = -1$. Adam with $\\beta_1 = 0$, $\\beta_2 = 1/(1 + C^2)$, and any\n",
    "step sequence $\\eta_t = \\eta/\\sqrt{t}$ makes* positive *net progress each\n",
    "period: its iterates converge to $x = +1$, where every third loss is\n",
    "maximal, and its average regret does not vanish.*\n",
    "[@Reddi.Kale.Kumar.2019] *carry the computation to general\n",
    "$\\beta_1 < \\sqrt{\\beta_2}$ (every practical setting): for any fixed\n",
    "exponential decay there is a $C$ that defeats it.*\n",
    "\n",
    "The mechanism fits in one sentence: **the effective step on the informative\n",
    "gradient shrinks faster than its information accrues.** When the rare large\n",
    "gradient $C$ arrives, it lands in $\\mathbf{v}$ *quadratically*: $v_t$\n",
    "jumps to nearly $C^2$, so the very step that should exploit the information\n",
    "is normalized by nearly $C$ and moves the iterate by only $O(\\eta_t)$. Then\n",
    "$\\beta_2$ forgets: over the next two steps $v_t$ collapses back toward $1$,\n",
    "and the two small $-1$ gradients each get full-sized steps in the *wrong*\n",
    "direction. Signal is throttled by its own magnitude, noise is not, and the\n",
    "mismatch compounds forever. This is a different failure from the noise ball\n",
    "of that section (which shrinks with the step\n",
    "size): the drift has a *sign*, and the iterate travels to the\n",
    "worst boundary and stays. The one-line fix, **AMSGrad**\n",
    "[@Reddi.Kale.Kumar.2019], replaces $\\hat{\\mathbf{v}}_t$ by the running\n",
    "*maximum* $\\tilde{\\mathbf{v}}_t = \\max(\\tilde{\\mathbf{v}}_{t-1}, \\mathbf{v}_t)$,\n",
    "which makes the per-coordinate step nonincreasing (the large gradient is\n",
    "never forgotten) and restores the convergence guarantee. The cell runs the\n",
    "construction with $C = 4$ against AMSGrad and plain projected SGD with the\n",
    "same $1/\\sqrt{t}$ decay:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "adaptive-stochastic-methods-reddi"
   },
   "source": [
    "def online(method, C=4.0, T=15000, alpha=0.5):\n",
    "    beta2 = 1 / (1 + C * C)                    # the proposition's beta_2\n",
    "    x, v, vmax, xs = 0.0, 0.0, 0.0, []\n",
    "    for t in range(1, T + 1):\n",
    "        g = C if t % 3 == 1 else -1.0          # rare +C, frequent -1\n",
    "        v = beta2 * v + (1 - beta2) * g * g\n",
    "        vmax = max(vmax, v)\n",
    "        if method == 'sgd':\n",
    "            step = alpha / np.sqrt(t) * g / C\n",
    "        else:                                  # adam (beta_1 = 0) or amsgrad\n",
    "            vv = vmax if method == 'amsgrad' else v\n",
    "            step = alpha / np.sqrt(t) * g / (np.sqrt(vv) + 1e-8)\n",
    "        x = min(1.0, max(-1.0, x - step))      # project back onto [-1, 1]\n",
    "        xs.append(x)\n",
    "    return xs\n",
    "\n",
    "print('        x_t at t =    30      300     3000    15000     (x* = -1)')\n",
    "for method in ['adam', 'amsgrad', 'sgd']:\n",
    "    xs = online(method)\n",
    "    print(f'{method:8s}' + ''.join(f'{xs[t - 1]:9.3f}'\n",
    "                                   for t in (30, 300, 3000, 15000)))"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "        x_t at t =    30      300     3000    15000     (x* = -1)\nadam        0.745    1.000    1.000    1.000\namsgrad    -0.953   -0.985   -0.995   -0.998\nsgd        -0.954   -0.986   -0.995   -0.998\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-4 fw=all -->\n",
    "\n",
    "Adam marches to $+1$, the maximizer of the periodic loss, and is pinned\n",
    "there by $t = 300$; AMSGrad and SGD settle onto the optimum $-1$. Nothing\n",
    "about the problem is pathological from SGD's point of view (it is convex,\n",
    "one-dimensional, with bounded gradients); the pathology is entirely in the\n",
    "interaction between exponential forgetting and rare informative gradients. In\n",
    "practice the construction's fingerprint, occasional large gradients that\n",
    "matter, is not exotic: rare tokens, rare classes, and loss spikes all\n",
    "qualify. Practice mostly retains plain Adam anyway (the drift needs the rare\n",
    "gradients to be *consistently* informative in the same direction, which real\n",
    "noise usually breaks up), but the theorem establishes the narrower conclusion that vanilla Adam lacks a\n",
    "general convergence guarantee without additional assumptions or modification,\n",
    "even for convex problems.\n",
    "\n",
    "### The Valley, Revisited\n",
    "\n",
    "The counterpart of that failure is what Adam does to the condition number,\n",
    "and here the news is good. The cell returns to the ill-conditioned quadratic\n",
    "of that section at $\\kappa = 10^3$ and races\n",
    "tuned gradient descent (the optimal single step\n",
    "$\\eta^\\star = 2/(\\lambda_{\\min} + \\lambda_{\\max})$) against hand-rolled\n",
    "Adam, printing Adam's effective per-coordinate steps\n",
    "$\\eta/(\\sqrt{\\hat{\\mathbf{v}}_t} + \\epsilon)$ as it runs:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "adaptive-stochastic-methods-adam-vs-gd"
   },
   "source": [
    "lam = np.array([1.0, 1000.0])                  # kappa = 1000\n",
    "f = lambda x: 0.5 * (lam * x * x).sum()\n",
    "\n",
    "def race(optimizer, eta, K=20000, tol=1e-8):\n",
    "    x, m, v, hit = np.array([1.0, 1.0]), np.zeros(2), np.zeros(2), None\n",
    "    for k in range(1, K + 1):\n",
    "        g = lam * x\n",
    "        if optimizer == 'gd':\n",
    "            x = x - eta * g\n",
    "        else:                                  # adam, the update of the text\n",
    "            m = 0.9 * m + 0.1 * g\n",
    "            v = 0.999 * v + 0.001 * g * g\n",
    "            step = eta / (np.sqrt(v / (1 - 0.999**k)) + 1e-8)\n",
    "            if k in (1, 100, 1000):\n",
    "                print(f'   k = {k:4d}: effective steps {step.round(6)}')\n",
    "            x = x - step * (m / (1 - 0.9**k))\n",
    "        if hit is None and f(x) < tol:\n",
    "            hit = k\n",
    "    return hit\n",
    "\n",
    "hit = race('gd', 2 / lam.sum())\n",
    "print(f'GD, optimal single eta = {2 / lam.sum():.1e}: f < 1e-8 at k = {hit}')\n",
    "print('Adam, eta = 0.01:')\n",
    "print(f'Adam reaches f < 1e-8 at k = {race(\"adam\", 0.01)}')"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "GD, optimal single eta = 2.0e-03: f < 1e-8 at k = 6160\nAdam, eta = 0.01:\n   k =    1: effective steps [1.e-02 1.e-05]\n   k =  100: effective steps [1.6534e-02 1.7000e-05]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "   k = 1000: effective steps [6.5755e-02 6.6000e-05]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Adam reaches f < 1e-8 at k = 344\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-5 fw=all -->\n",
    "\n",
    "Gradient descent with the *optimal* single step needs $6160$ iterations, and\n",
    "the $(\\kappa-1)/(\\kappa+1)$ contraction law of\n",
    "that section predicts that count exactly:\n",
    "with $\\eta^\\star$ both coordinates contract by $\\rho = 999/1001$ per step, so\n",
    "$f_k = 500.5\\,\\rho^{2k}$ crosses $10^{-8}$ at $k = 6159.07$, making $6160$\n",
    "the first iterate under tolerance. Adam gets there in $344$. The printed\n",
    "effective steps say why: from the\n",
    "first iteration Adam's step on the stiff coordinate is $10^{-5}$ and on the\n",
    "flat coordinate $10^{-2}$, the $1000\\times$ ratio of the eigenvalues,\n",
    "reconstructed from gradient magnitudes alone, with no eigendecomposition and\n",
    "no Hessian. Both coordinates then travel toward zero at comparable speed\n",
    "(this is the right panel of the figure): the valley\n",
    "has been rescaled into a bowl.\n",
    "\n",
    "Now the counterpoint. On *this* problem, exact preconditioning is\n",
    "available: Newton's method solves it in one step, and even a fixed\n",
    "diagonal $\\eta_i = 1/\\lambda_i$ needs no adaptivity at all. What Adam offers\n",
    "is an approximation to that oracle built *from noisy first-order data,\n",
    "continuously, at $O(d)$ cost*. And the approximation is crude:\n",
    "$\\sqrt{\\hat{\\mathbf{v}}}$\n",
    "conflates curvature with gradient noise (both inflate squared gradients), the\n",
    "diagonal ignores every off-axis correlation (rotate this quadratic by\n",
    "$45°$ and per-coordinate rescaling loses most of its effect), and the Reddi\n",
    "construction shows the estimate can be steered adversarially. Adaptive\n",
    "methods occupy a well-chosen point partway along a cost--fidelity curve; the\n",
    "preconditioning ladder at the end of this section climbs further up it.\n",
    "\n",
    "## Decoupled Weight Decay\n",
    "\n",
    "### The Penalty Gradient Goes Through the Preconditioner\n",
    "\n",
    "Weight decay looks too simple to interact with any of this: add\n",
    "$\\tfrac{\\lambda}{2}\\|\\mathbf{w}\\|^2$ to the loss, get $\\lambda \\mathbf{w}$\n",
    "added to the gradient. Under SGD the two implementations, \"penalize the\n",
    "loss\" and \"shrink the weights\", are *identical*:\n",
    "\n",
    "$$\n",
    "\\mathbf{w}_{t+1} = \\mathbf{w}_t - \\eta\\,(\\mathbf{g}_t + \\lambda \\mathbf{w}_t)\n",
    "= (1 - \\eta\\lambda)\\, \\mathbf{w}_t - \\eta\\, \\mathbf{g}_t .\n",
    "$$\n",
    "\n",
    "Under Adam they are not, and the two-line derivation is the whole story.\n",
    "Feed the penalized gradient through the equation and look at the\n",
    "decay term alone (take $\\beta_1 = 0$ for clarity, so the update is\n",
    "$-\\eta\\,(\\mathbf{g}_t + \\lambda\\mathbf{w}_t)/(\\sqrt{\\hat{\\mathbf{v}}_t} + \\epsilon)$):\n",
    "the shrinkage applied to coordinate $i$ is\n",
    "\n",
    "$$\n",
    "\\textrm{$\\ell_2$ through Adam:}\\quad\n",
    "\\frac{\\eta\\,\\lambda}{\\sqrt{\\hat{v}_{t,i}} + \\epsilon}\\; w_{t,i}\n",
    "\\qquad \\textrm{versus} \\qquad\n",
    "\\textrm{decoupled:}\\quad \\eta\\,\\lambda\\, w_{t,i} .\n",
    "$$\n",
    "\n",
    "The penalty's gradient is divided by the same $\\sqrt{\\hat{\\mathbf{v}}}$ as\n",
    "everything else, so the *regularization strength is no longer a constant of\n",
    "the problem*: it varies per coordinate and over time, inversely with the\n",
    "gradient scale. A weight whose loss gradients are large or noisy (large\n",
    "$\\hat{v}_i$) is barely regularized at all; a weight whose gradients have gone\n",
    "quiet is regularized hard. Whatever $\\tfrac{\\lambda}{2}\\|\\mathbf{w}\\|^2$ was\n",
    "supposed to mean (the Gaussian prior of the MAP estimate in\n",
    "that section, or the norm constraint\n",
    "$\\|\\mathbf{w}\\|^2 \\le r^2$ whose multiplier it is in\n",
    "that section), that meaning assumed one\n",
    "$\\lambda$ for all coordinates, and Adam's preconditioner silently discards it.\n",
    "\n",
    "**AdamW** [@Loshchilov.Hutter.2019] restores the intended semantics by\n",
    "*decoupling*: the loss gradient goes through the preconditioner, the decay\n",
    "does not,\n",
    "\n",
    "$$\n",
    "\\mathbf{w}_{t+1} = (1 - \\eta\\lambda)\\, \\mathbf{w}_t - \\eta\\; \\frac{\\hat{\\mathbf{m}}_t}{\\sqrt{\\hat{\\mathbf{v}}_t} + \\epsilon},\n",
    "$$\n",
    "\n",
    "which is verbatim the SGD shrinkage $(1 - \\eta\\lambda)$, applied uniformly.\n",
    "The empirical finding that motivated the fix ($\\ell_2$-regularized Adam\n",
    "generalizing worse than SGD with weight decay, with the gap closed by\n",
    "decoupling) is downstream of exactly the per-coordinate distortion in\n",
    "the equation. The cell isolates that distortion: two weights\n",
    "whose loss gradients are pure noise (so decay is the only systematic force)\n",
    "at very different noise scales, $\\sigma = (10, 0.1)$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "adaptive-stochastic-methods-adamw"
   },
   "source": [
    "sig, alpha, lam_wd, T = np.array([10.0, 0.1]), 1e-3, 0.1, 4000\n",
    "\n",
    "def decay_race(decoupled, seed=0):\n",
    "    rg = np.random.default_rng(seed)\n",
    "    w, m, v = np.array([1.0, 1.0]), np.zeros(2), np.zeros(2)\n",
    "    for t in range(1, T + 1):\n",
    "        g = sig * rg.standard_normal(2)        # loss gradient: pure noise\n",
    "        if not decoupled:\n",
    "            g = g + lam_wd * w                 # l2 penalty inside the gradient\n",
    "        m, v = 0.9 * m + 0.1 * g, 0.999 * v + 0.001 * g * g\n",
    "        w = w - alpha * (m / (1 - 0.9**t)) / (np.sqrt(v / (1 - 0.999**t)) + 1e-8)\n",
    "        if decoupled:\n",
    "            w = w - alpha * lam_wd * w         # decay outside the preconditioner\n",
    "    return w\n",
    "\n",
    "print('per-step decay rate, l2 through Adam (alpha*lam/sigma_i):',\n",
    "      (alpha * lam_wd / sig).round(6))\n",
    "print('per-step decay rate, AdamW (uniform):                    ',\n",
    "      alpha * lam_wd)\n",
    "print(f'|w| after {T} steps,  Adam + l2: {np.abs(decay_race(False)).round(4)}')\n",
    "print(f'|w| after {T} steps,  AdamW    : {np.abs(decay_race(True)).round(4)}')\n",
    "print(f'AdamW prediction (1 - alpha*lam)^T = {(1 - alpha * lam_wd)**T:.4f}')"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "per-step decay rate, l2 through Adam (alpha*lam/sigma_i): [1.e-05 1.e-03]\nper-step decay rate, AdamW (uniform):                     0.0001\n|w| after 4000 steps,  Adam + l2: [0.9677 0.0235]\n|w| after 4000 steps,  AdamW    : [0.6681 0.6598]\nAdamW prediction (1 - alpha*lam)^T = 0.6703\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-6 fw=all -->\n",
    "\n",
    "Under $\\ell_2$-through-Adam the noisy coordinate has decayed from $1$ to only\n",
    "$0.968$ after $4000$ steps while the quiet coordinate has collapsed to\n",
    "$0.024$: effective decay rates of $10^{-5}$ versus $10^{-3}$, the\n",
    "predicted $\\eta\\lambda/\\sigma_i$, a $100\\times$ disparity that nobody chose.\n",
    "AdamW shrinks both to $\\approx 0.66$, within $2\\%$ of the pure-decay\n",
    "prediction $(1 - \\eta\\lambda)^T = 0.6703$: one $\\lambda$ with one meaning.\n",
    "(A subtlety about fixed points: at a *deterministic* stationary point,\n",
    "coupled Adam is stationary exactly where SGD with $\\ell_2$ is, at\n",
    "$\\nabla L + \\lambda\\mathbf{w} = \\mathbf{0}$, while *AdamW's* stationary\n",
    "points sit elsewhere; Exercise 4 derives both conditions and explains why the\n",
    "noisy, never-stationary training regime, where paths matter more than fixed\n",
    "points, is where decoupling matters.) In the major libraries, `AdamW` and\n",
    "`Adam` with a `weight_decay` flag implement two different regularizers: the\n",
    "former is the decoupled update exactly, the latter the coupled\n",
    "the equation.\n",
    "\n",
    "## Schedules and Warmup\n",
    "\n",
    "### What Decay Does, and Which Shape\n",
    "\n",
    "A learning-rate schedule is a sequence $\\eta_t$ (the main book surveys the\n",
    "common ones in that section), and the mathematics it must\n",
    "negotiate was already proved in\n",
    "that section: a constant step parks SGD on a\n",
    "noise floor proportional to $\\eta$ the equation, decay must\n",
    "respect Robbins--Monro the equation to reach the optimum,\n",
    "and the textbook $\\eta_k = c/k$ carries the trap that undersized $c$ silently\n",
    "degrades the rate class.\n",
    "\n",
    "Add to this the Ghadimi--Lan prescription of this section: for a fixed\n",
    "nonconvex budget $K$, a *constant* $\\eta \\propto 1/\\sqrt{K}$ is already\n",
    "**minimax-optimal**, meaning no method that sees only the same\n",
    "unbiased, bounded-variance gradient oracle can guarantee a better worst-case\n",
    "rate than the $K^{-1/2}$ it achieves; the matching lower bound is due to\n",
    "@Arjevani.Carmon.Duchi.ea.2023 . The theory therefore assigns decay a\n",
    "specific job: it lowers the noise floor, at a cost in transient progress, and\n",
    "beyond convex problems no theorem ranks one decay *shape* against another.\n",
    "Which of the statements below are proved and which are read off practice is\n",
    "marked as we go. The three shapes to know, drawn in\n",
    "the figure, are these:\n",
    "\n",
    "![The schedule zoo, at equal budget. A constant step never leaves its noise floor (the floor is $\\propto \\eta$, per the noise-ball analysis); $c/k$ decay reaches the optimum but must get $c$ right; cosine spends a long tail at small steps; warmup--stable--decay (WSD) holds the constant plateau as long as possible and drops the noise floor in a final decay phase. All four begin with the linear warmup that adaptive preconditioners need while $\\hat{\\mathbf{v}}$ is still estimated from a handful of gradients.](../img/mdl-opt-schedule-zoo.svg)\n",
    "\n",
    "**Cosine decay** [@Loshchilov.Hutter.2016] sets\n",
    "$\\eta_t = \\tfrac12 \\eta_0 (1 + \\cos(\\pi t / K))$: a smooth descent from\n",
    "$\\eta_0$ to $0$ with most of its time spent at moderate steps and a long,\n",
    "gentle tail. It satisfies no optimality theorem; it is popular because it has\n",
    "one parameter ($\\eta_0$), no kinks to tune, and behaves gracefully across\n",
    "budgets. **Warmup--stable--decay (WSD)** [@Hu.Tu.Han.ea.2024] holds\n",
    "$\\eta_0$ constant for most of the run and decays only in a final fraction\n",
    "(often the last $10$--$20\\%$). Its practical appeal is operational: the\n",
    "long constant plateau means checkpoints from one run can be re-decayed to any\n",
    "budget, rather than committing to $K$ at launch. Its loss-landscape\n",
    "reading is the river-valley picture: during the plateau the iterate travels\n",
    "fast *along* the valley while rattling in its noise ball across it, and the\n",
    "final decay shrinks the ball, dropping the loss sharply as the iterate\n",
    "settles to the valley floor. That reading is this chapter's noise-ball\n",
    "result used as a design principle, and it is the level\n",
    "of theory available: the plateau's benefit on real losses (more transient\n",
    "progress per unit budget) is an empirical regularity, not a theorem. The cell\n",
    "races the shapes on the noisy quadratic, the miniature where the noise\n",
    "floor is exact, at equal budget, printing the loss at $80\\%$ of the budget\n",
    "and at the end:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "adaptive-stochastic-methods-schedules"
   },
   "source": [
    "lam_s, sigma, K, eta0 = 1.0, 1.0, 2000, 0.3\n",
    "\n",
    "def run_schedule(eta_fn, seeds=200):\n",
    "    at80, final = [], []\n",
    "    for s in range(seeds):\n",
    "        rg, x = np.random.default_rng(s), 3.0\n",
    "        for k in range(K):\n",
    "            x = x - eta_fn(k) * (lam_s * x + sigma * rg.standard_normal())\n",
    "            if k == int(0.8 * K) - 1:\n",
    "                at80.append(0.5 * lam_s * x * x)\n",
    "        final.append(0.5 * lam_s * x * x)\n",
    "    return np.mean(at80), np.mean(final)\n",
    "\n",
    "schedules = [\n",
    "    ('constant', lambda k: eta0),\n",
    "    ('c/k decay', lambda k: eta0 / (1 + k / 50)),\n",
    "    ('cosine', lambda k: 0.5 * eta0 * (1 + np.cos(np.pi * k / K))),\n",
    "    ('WSD 80/20', lambda k: eta0 if k < 0.8 * K else eta0 * (K - k) / (0.2 * K)),\n",
    "]\n",
    "print(f'         E f at 80% budget   E f at end     (floor ~ '\n",
    "      f'{eta0 * sigma**2 / (2 * (2 - eta0 * lam_s)):.3f})')\n",
    "for name, fn in schedules:\n",
    "    m80, mK = run_schedule(fn)\n",
    "    print(f'{name:10s}    {m80:9.2e}      {mK:9.2e}')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "         E f at 80% budget   E f at end     (floor ~ 0.088)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "constant       9.53e-02       9.19e-02\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "c/k decay      2.82e-03       1.94e-03\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "cosine         8.71e-03       1.58e-03\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "WSD 80/20      9.53e-02       5.48e-03\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-7 fw=all -->\n",
    "\n",
    "The constant run sits on its floor at both checkpoints ($9.5$ and\n",
    "$9.2 \\times 10^{-2}$, against the printed prediction\n",
    "$\\eta\\sigma^2/(2(2 - \\eta\\lambda)) \\approx 0.088$), while every decaying\n",
    "schedule ends one to two orders lower. The instructive row is WSD: at $80\\%$\n",
    "of the budget it is *still on the constant's floor* ($9.5 \\times 10^{-2}$,\n",
    "equal to the constant row to every printed digit, since the two runs are the\n",
    "same trajectory until the decay begins), then the final $400$ steps drop it\n",
    "by a factor of $17$. All of WSD's visible gain\n",
    "arrives in the decay phase: the loss-curve cliff that WSD training runs\n",
    "are known for, reproduced here by a one-line quadratic. On this pure\n",
    "quadratic, cosine's longer tail of small steps wins the endgame\n",
    "($1.6 \\times 10^{-3}$); the reasons to prefer WSD (re-decayable\n",
    "checkpoints, more time at large steps on losses whose transient is the\n",
    "bottleneck) live off this toy, in exactly the gap between the noise-ball\n",
    "theorem and practice.\n",
    "\n",
    "### Warmup: Do Not Trust an Estimated Preconditioner Cold\n",
    "\n",
    "Every schedule above begins with **warmup**: $\\eta_t$ ramped linearly from\n",
    "$\\approx 0$ over the first fraction of training\n",
    "[@Goyal.Dollar.Girshick.ea.2017]. For SGD, warmup is mild prudence\n",
    "(early gradients are large and the landscape unexplored). For adaptive\n",
    "methods it is close to structural, and this section has already assembled\n",
    "both reasons.\n",
    "\n",
    "First, the preconditioner starts out *statistically worthless*. At $t = 1$,\n",
    "$\\hat{\\mathbf{v}}_1 = \\mathbf{g}_1^2$: the per-coordinate rescaling is\n",
    "estimated from a single sample, and bias correction makes it *unbiased*, not\n",
    "*accurate*. The step direction\n",
    "$\\hat{\\mathbf{m}}_1/(\\sqrt{\\hat{\\mathbf{v}}_1} + \\epsilon) = \\mathrm{sign}(\\mathbf{g}_1)$\n",
    "commits to full-magnitude movement on every coordinate, including ones whose\n",
    "lone observed gradient was pure noise. The estimate needs on the order of\n",
    "$1/(1 - \\beta_2) = 1000$ steps to average over a meaningful window, and\n",
    "until then Adam is applying its most aggressive, least justified rescaling.\n",
    "Ramping $\\eta$ over that same window is the obvious hedge: keep the steps\n",
    "small until the preconditioner has data to justify them. (This diagnosis can\n",
    "also be read off the figure: the correction fixes\n",
    "the *mean* of the early estimate; nothing fixes its variance but data.)\n",
    "\n",
    "Second, the target step size is only safe *after* the landscape has adapted\n",
    "to it. The edge-of-stability measurement in\n",
    "that section showed sharpness equilibrating\n",
    "onto $2/\\eta$, but that equilibration takes time, and a full-size $\\eta$\n",
    "imposed at initialization, before progressive sharpening has anywhere to\n",
    "settle, simply violates the local stability ceiling and diverges or spikes.\n",
    "Warmup gives the curvature time to meet the step size it will be asked to\n",
    "tolerate. Neither argument is a theorem about deep networks; both are this\n",
    "chapter's instruments (estimator variance and the stability ceiling)\n",
    "pointed at the first thousand steps, and together they predict the practice\n",
    "that every large training run has converged on independently.\n",
    "\n",
    "## Variance Reduction for Finite Sums\n",
    "\n",
    "Minibatching reduces variance by averaging more examples, but a fixed-size\n",
    "minibatch still has noise at the optimum. A finite training set offers another\n",
    "possibility: reuse information from earlier gradients as a **control variate**.\n",
    "Write the empirical objective as\n",
    "\n",
    "$$\n",
    "F(\\mathbf x)=\\frac1n\\sum_{i=1}^n f_i(\\mathbf x).\n",
    "$$\n",
    "\n",
    "At a reference point $\\widetilde{\\mathbf x}$, compute the full gradient\n",
    "$\\widetilde{\\boldsymbol\\mu}=\\nabla F(\\widetilde{\\mathbf x})$. The **SVRG**\n",
    "estimator at the current point is\n",
    "\n",
    "$$\n",
    "\\mathbf v_i(\\mathbf x)\n",
    "=\\nabla f_i(\\mathbf x)-\\nabla f_i(\\widetilde{\\mathbf x})\n",
    " +\\widetilde{\\boldsymbol\\mu}.\n",
    "$$\n",
    "\n",
    "For a uniformly sampled index,\n",
    "\n",
    "$$\n",
    "\\mathbb E_i[\\mathbf v_i(\\mathbf x)]\n",
    "=\\nabla F(\\mathbf x)-\\nabla F(\\widetilde{\\mathbf x})\n",
    " +\\nabla F(\\widetilde{\\mathbf x})\n",
    "=\\nabla F(\\mathbf x),\n",
    "$$\n",
    "\n",
    "so the estimator is unbiased. Its advantage is not the mean but the variance:\n",
    "when $\\mathbf x$ is close to the snapshot, the two component gradients are\n",
    "highly correlated and their difference is much less variable than either one.\n",
    "After an inner loop moves too far from the snapshot, refresh the full gradient\n",
    "and repeat. Under smooth strong convexity, this mechanism yields geometric\n",
    "convergence with a fixed step size, rather than SGD's fixed-step noise floor\n",
    "[@Bottou.Curtis.Nocedal.2018].\n",
    "\n",
    "The cell compares SVRG with SGD on ridge regression. One SVRG epoch pays $n$\n",
    "component-gradient equivalents for the full gradient and two per inner step;\n",
    "SGD receives the same total budget of $3n$ component gradients per reported\n",
    "epoch. Both use a fixed step size, so the comparison isolates the control\n",
    "variate rather than a schedule."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "mdl-adaptive-svrg"
   },
   "source": [
    "rng = np.random.default_rng(4)\n",
    "n_vr, d_vr = 300, 12\n",
    "X_vr = rng.standard_normal((n_vr, d_vr))\n",
    "w_true_vr = rng.standard_normal(d_vr)\n",
    "y_vr = X_vr @ w_true_vr + 0.5 * rng.standard_normal(n_vr)\n",
    "lam_vr = 0.1\n",
    "H_vr = X_vr.T @ X_vr / n_vr + lam_vr * np.eye(d_vr)\n",
    "w_star_vr = np.linalg.solve(H_vr, X_vr.T @ y_vr / n_vr)\n",
    "\n",
    "def component_grad(w, i):\n",
    "    return X_vr[i] * (X_vr[i] @ w - y_vr[i]) + lam_vr * w\n",
    "\n",
    "def full_grad(w):\n",
    "    return X_vr.T @ (X_vr @ w - y_vr) / n_vr + lam_vr * w\n",
    "\n",
    "def objective_vr(w):\n",
    "    return (0.5 * np.mean((X_vr @ w - y_vr)**2)\n",
    "            + 0.5 * lam_vr * (w @ w))\n",
    "\n",
    "eta_vr, epochs_vr = 0.05, 20\n",
    "snapshot, gaps_svrg = np.zeros(d_vr), []\n",
    "for _ in range(epochs_vr):\n",
    "    mean_grad = full_grad(snapshot)\n",
    "    w = snapshot.copy()\n",
    "    for _ in range(n_vr):\n",
    "        i = rng.integers(n_vr)\n",
    "        v = component_grad(w, i) - component_grad(snapshot, i) + mean_grad\n",
    "        w -= eta_vr * v\n",
    "    snapshot = w\n",
    "    gaps_svrg.append(objective_vr(snapshot) - objective_vr(w_star_vr))\n",
    "\n",
    "rng_sgd = np.random.default_rng(5)\n",
    "w_sgd, gaps_sgd = np.zeros(d_vr), []\n",
    "for _ in range(epochs_vr):\n",
    "    for _ in range(3 * n_vr):               # same component-gradient budget\n",
    "        i = rng_sgd.integers(n_vr)\n",
    "        w_sgd -= eta_vr * component_grad(w_sgd, i)\n",
    "    gaps_sgd.append(objective_vr(w_sgd) - objective_vr(w_star_vr))\n",
    "\n",
    "d2l.plot(np.arange(1, epochs_vr + 1), [gaps_sgd, gaps_svrg],\n",
    "         'equal-budget epoch', 'optimality gap',\n",
    "         legend=['SGD, fixed step', 'SVRG, fixed step'], yscale='log')"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"253.878125pt\" height=\"183.35625pt\" viewBox=\"0 0 253.878125 183.35625\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T13:54:40.053397</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 183.35625 \nL 253.878125 183.35625 \nL 253.878125 0 \nL 0 0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 51.378125 145.8 \nL 246.678125 145.8 \nL 246.678125 7.2 \nL 51.378125 7.2 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <path d=\"M 97.633388 145.8 \nL 97.633388 7.2 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_2\">\n      <defs>\n       <path id=\"m94aeeb4d20\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m94aeeb4d20\" x=\"97.633388\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 5 -->\n      <g transform=\"translate(94.452138 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-35\" d=\"M 691 4666 \nL 3169 4666 \nL 3169 4134 \nL 1269 4134 \nL 1269 2991 \nQ 1406 3038 1543 3061 \nQ 1681 3084 1819 3084 \nQ 2600 3084 3056 2656 \nQ 3513 2228 3513 1497 \nQ 3513 744 3044 326 \nQ 2575 -91 1722 -91 \nQ 1428 -91 1123 -41 \nQ 819 9 494 109 \nL 494 744 \nQ 775 591 1075 516 \nQ 1375 441 1709 441 \nQ 2250 441 2565 725 \nQ 2881 1009 2881 1497 \nQ 2881 1984 2565 2268 \nQ 2250 2553 1709 2553 \nQ 1456 2553 1204 2497 \nQ 953 2441 691 2322 \nL 691 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-35\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_3\">\n      <path d=\"M 144.355876 145.8 \nL 144.355876 7.2 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#m94aeeb4d20\" x=\"144.355876\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 10 -->\n      <g transform=\"translate(137.993376 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_5\">\n      <path d=\"M 191.078364 145.8 \nL 191.078364 7.2 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m94aeeb4d20\" x=\"191.078364\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 15 -->\n      <g transform=\"translate(184.715864 160.398438) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_7\">\n      <path d=\"M 237.800852 145.8 \nL 237.800852 7.2 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#m94aeeb4d20\" x=\"237.800852\" y=\"145.8\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 20 -->\n      <g transform=\"translate(231.438352 160.398438) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-32\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_5\">\n     <!-- equal-budget epoch -->\n     <g transform=\"translate(98.80625 174.076563) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-71\" d=\"M 947 1747 \nQ 947 1113 1208 752 \nQ 1469 391 1925 391 \nQ 2381 391 2643 752 \nQ 2906 1113 2906 1747 \nQ 2906 2381 2643 2742 \nQ 2381 3103 1925 3103 \nQ 1469 3103 1208 2742 \nQ 947 2381 947 1747 \nz\nM 2906 525 \nQ 2725 213 2448 61 \nQ 2172 -91 1784 -91 \nQ 1150 -91 751 415 \nQ 353 922 353 1747 \nQ 353 2572 751 3078 \nQ 1150 3584 1784 3584 \nQ 2172 3584 2448 3432 \nQ 2725 3281 2906 2969 \nL 2906 3500 \nL 3481 3500 \nL 3481 -1331 \nL 2906 -1331 \nL 2906 525 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-75\" d=\"M 544 1381 \nL 544 3500 \nL 1119 3500 \nL 1119 1403 \nQ 1119 906 1312 657 \nQ 1506 409 1894 409 \nQ 2359 409 2629 706 \nQ 2900 1003 2900 1516 \nL 2900 3500 \nL 3475 3500 \nL 3475 0 \nL 2900 0 \nL 2900 538 \nQ 2691 219 2414 64 \nQ 2138 -91 1772 -91 \nQ 1169 -91 856 284 \nQ 544 659 544 1381 \nz\nM 1991 3584 \nL 1991 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-61\" d=\"M 2194 1759 \nQ 1497 1759 1228 1600 \nQ 959 1441 959 1056 \nQ 959 750 1161 570 \nQ 1363 391 1709 391 \nQ 2188 391 2477 730 \nQ 2766 1069 2766 1631 \nL 2766 1759 \nL 2194 1759 \nz\nM 3341 1997 \nL 3341 0 \nL 2766 0 \nL 2766 531 \nQ 2569 213 2275 61 \nQ 1981 -91 1556 -91 \nQ 1019 -91 701 211 \nQ 384 513 384 1019 \nQ 384 1609 779 1909 \nQ 1175 2209 1959 2209 \nL 2766 2209 \nL 2766 2266 \nQ 2766 2663 2505 2880 \nQ 2244 3097 1772 3097 \nQ 1472 3097 1187 3025 \nQ 903 2953 641 2809 \nL 641 3341 \nQ 956 3463 1253 3523 \nQ 1550 3584 1831 3584 \nQ 2591 3584 2966 3190 \nQ 3341 2797 3341 1997 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-2d\" d=\"M 313 2009 \nL 1997 2009 \nL 1997 1497 \nL 313 1497 \nL 313 2009 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-62\" d=\"M 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\nM 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nL 1159 0 \nL 581 0 \nL 581 4863 \nL 1159 4863 \nL 1159 2969 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-64\" d=\"M 2906 2969 \nL 2906 4863 \nL 3481 4863 \nL 3481 0 \nL 2906 0 \nL 2906 525 \nQ 2725 213 2448 61 \nQ 2172 -91 1784 -91 \nQ 1150 -91 751 415 \nQ 353 922 353 1747 \nQ 353 2572 751 3078 \nQ 1150 3584 1784 3584 \nQ 2172 3584 2448 3432 \nQ 2725 3281 2906 2969 \nz\nM 947 1747 \nQ 947 1113 1208 752 \nQ 1469 391 1925 391 \nQ 2381 391 2643 752 \nQ 2906 1113 2906 1747 \nQ 2906 2381 2643 2742 \nQ 2381 3103 1925 3103 \nQ 1469 3103 1208 2742 \nQ 947 2381 947 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-67\" d=\"M 2906 1791 \nQ 2906 2416 2648 2759 \nQ 2391 3103 1925 3103 \nQ 1463 3103 1205 2759 \nQ 947 2416 947 1791 \nQ 947 1169 1205 825 \nQ 1463 481 1925 481 \nQ 2391 481 2648 825 \nQ 2906 1169 2906 1791 \nz\nM 3481 434 \nQ 3481 -459 3084 -895 \nQ 2688 -1331 1869 -1331 \nQ 1566 -1331 1297 -1286 \nQ 1028 -1241 775 -1147 \nL 775 -588 \nQ 1028 -725 1275 -790 \nQ 1522 -856 1778 -856 \nQ 2344 -856 2625 -561 \nQ 2906 -266 2906 331 \nL 2906 616 \nQ 2728 306 2450 153 \nQ 2172 0 1784 0 \nQ 1141 0 747 490 \nQ 353 981 353 1791 \nQ 353 2603 747 3093 \nQ 1141 3584 1784 3584 \nQ 2172 3584 2450 3431 \nQ 2728 3278 2906 2969 \nL 2906 3500 \nL 3481 3500 \nL 3481 434 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-74\" d=\"M 1172 4494 \nL 1172 3500 \nL 2356 3500 \nL 2356 3053 \nL 1172 3053 \nL 1172 1153 \nQ 1172 725 1289 603 \nQ 1406 481 1766 481 \nL 2356 481 \nL 2356 0 \nL 1766 0 \nQ 1100 0 847 248 \nQ 594 497 594 1153 \nL 594 3053 \nL 172 3053 \nL 172 3500 \nL 594 3500 \nL 594 4494 \nL 1172 4494 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-20\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-70\" d=\"M 1159 525 \nL 1159 -1331 \nL 581 -1331 \nL 581 3500 \nL 1159 3500 \nL 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nz\nM 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6f\" d=\"M 1959 3097 \nQ 1497 3097 1228 2736 \nQ 959 2375 959 1747 \nQ 959 1119 1226 758 \nQ 1494 397 1959 397 \nQ 2419 397 2687 759 \nQ 2956 1122 2956 1747 \nQ 2956 2369 2687 2733 \nQ 2419 3097 1959 3097 \nz\nM 1959 3584 \nQ 2709 3584 3137 3096 \nQ 3566 2609 3566 1747 \nQ 3566 888 3137 398 \nQ 2709 -91 1959 -91 \nQ 1206 -91 779 398 \nQ 353 888 353 1747 \nQ 353 2609 779 3096 \nQ 1206 3584 1959 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-63\" d=\"M 3122 3366 \nL 3122 2828 \nQ 2878 2963 2633 3030 \nQ 2388 3097 2138 3097 \nQ 1578 3097 1268 2742 \nQ 959 2388 959 1747 \nQ 959 1106 1268 751 \nQ 1578 397 2138 397 \nQ 2388 397 2633 464 \nQ 2878 531 3122 666 \nL 3122 134 \nQ 2881 22 2623 -34 \nQ 2366 -91 2075 -91 \nQ 1284 -91 818 406 \nQ 353 903 353 1747 \nQ 353 2603 823 3093 \nQ 1294 3584 2113 3584 \nQ 2378 3584 2631 3529 \nQ 2884 3475 3122 3366 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-68\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 4863 \nL 1159 4863 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-65\"/>\n      <use xlink:href=\"#DejaVuSans-71\" transform=\"translate(61.523438 0)\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(188.378906 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(249.658203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-2d\" transform=\"translate(277.441406 0)\"/>\n      <use xlink:href=\"#DejaVuSans-62\" transform=\"translate(313.525391 0)\"/>\n      <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(377.001953 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(440.380859 0)\"/>\n      <use xlink:href=\"#DejaVuSans-67\" transform=\"translate(503.857422 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(567.333984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(628.857422 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(668.066406 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(699.853516 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(761.376953 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(824.853516 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(886.035156 0)\"/>\n      <use xlink:href=\"#DejaVuSans-68\" transform=\"translate(941.015625 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_9\">\n      <path d=\"M 51.378125 117.397285 \nL 246.678125 117.397285 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_10\">\n      <defs>\n       <path id=\"m5ea4ea6052\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m5ea4ea6052\" x=\"51.378125\" y=\"117.397285\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- $\\mathdefault{10^{-6}}$ -->\n      <g transform=\"translate(20.878125 121.196504) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-2212\" d=\"M 678 2272 \nL 4684 2272 \nL 4684 1741 \nL 678 1741 \nL 678 2272 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-36\" d=\"M 2113 2584 \nQ 1688 2584 1439 2293 \nQ 1191 2003 1191 1497 \nQ 1191 994 1439 701 \nQ 1688 409 2113 409 \nQ 2538 409 2786 701 \nQ 3034 994 3034 1497 \nQ 3034 2003 2786 2293 \nQ 2538 2584 2113 2584 \nz\nM 3366 4563 \nL 3366 3988 \nQ 3128 4100 2886 4159 \nQ 2644 4219 2406 4219 \nQ 1781 4219 1451 3797 \nQ 1122 3375 1075 2522 \nQ 1259 2794 1537 2939 \nQ 1816 3084 2150 3084 \nQ 2853 3084 3261 2657 \nQ 3669 2231 3669 1497 \nQ 3669 778 3244 343 \nQ 2819 -91 2113 -91 \nQ 1303 -91 875 529 \nQ 447 1150 447 2328 \nQ 447 3434 972 4092 \nQ 1497 4750 2381 4750 \nQ 2619 4750 2861 4703 \nQ 3103 4656 3366 4563 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-36\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_11\">\n      <path d=\"M 51.378125 85.479241 \nL 246.678125 85.479241 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#m5ea4ea6052\" x=\"51.378125\" y=\"85.479241\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- $\\mathdefault{10^{-4}}$ -->\n      <g transform=\"translate(20.878125 89.27846) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \nL 825 1625 \nL 2419 1625 \nL 2419 4116 \nz\nM 2253 4666 \nL 3047 4666 \nL 3047 1625 \nL 3713 1625 \nL 3713 1100 \nL 3047 1100 \nL 3047 0 \nL 2419 0 \nL 2419 1100 \nL 313 1100 \nL 313 1709 \nL 2253 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.684375)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 38.965625) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(186.855469 38.965625) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_13\">\n      <path d=\"M 51.378125 53.561197 \nL 246.678125 53.561197 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#m5ea4ea6052\" x=\"51.378125\" y=\"53.561197\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- $\\mathdefault{10^{-2}}$ -->\n      <g transform=\"translate(20.878125 57.360416) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-2212\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(186.855469 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_15\">\n      <path d=\"M 51.378125 21.643153 \nL 246.678125 21.643153 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #b0b0b0; stroke-width: 0.8; stroke-linecap: square\"/>\n     </g>\n     <g id=\"line2d_16\">\n      <g>\n       <use xlink:href=\"#m5ea4ea6052\" x=\"51.378125\" y=\"21.643153\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- $\\mathdefault{10^{0}}$ -->\n      <g transform=\"translate(26.778125 25.442372) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(0 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(63.623047 0.765625)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(128.203125 39.046875) scale(0.7)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_10\">\n     <!-- optimality gap -->\n     <g transform=\"translate(14.798438 112.716406) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6d\" d=\"M 3328 2828 \nQ 3544 3216 3844 3400 \nQ 4144 3584 4550 3584 \nQ 5097 3584 5394 3201 \nQ 5691 2819 5691 2113 \nL 5691 0 \nL 5113 0 \nL 5113 2094 \nQ 5113 2597 4934 2840 \nQ 4756 3084 4391 3084 \nQ 3944 3084 3684 2787 \nQ 3425 2491 3425 1978 \nL 3425 0 \nL 2847 0 \nL 2847 2094 \nQ 2847 2600 2669 2842 \nQ 2491 3084 2119 3084 \nQ 1678 3084 1418 2786 \nQ 1159 2488 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1356 3278 1631 3431 \nQ 1906 3584 2284 3584 \nQ 2666 3584 2933 3390 \nQ 3200 3197 3328 2828 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-79\" d=\"M 2059 -325 \nQ 1816 -950 1584 -1140 \nQ 1353 -1331 966 -1331 \nL 506 -1331 \nL 506 -850 \nL 844 -850 \nQ 1081 -850 1212 -737 \nQ 1344 -625 1503 -206 \nL 1606 56 \nL 191 3500 \nL 800 3500 \nL 1894 763 \nL 2988 3500 \nL 3597 3500 \nL 2059 -325 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-6f\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(61.181641 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(124.658203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(163.867188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6d\" transform=\"translate(191.650391 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(289.0625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(350.341797 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(378.125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(405.908203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(445.117188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(504.296875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-67\" transform=\"translate(536.083984 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(599.560547 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(660.839844 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_17\">\n    <path d=\"M 60.255398 41.244154 \nL 69.599895 43.333111 \nL 78.944393 36.663201 \nL 88.288891 40.287314 \nL 97.633388 33.630372 \nL 106.977886 35.72657 \nL 116.322383 39.521804 \nL 125.666881 42.863955 \nL 135.011379 40.841411 \nL 144.355876 31.736413 \nL 153.700374 41.286513 \nL 163.044871 36.124942 \nL 172.389369 39.796283 \nL 181.733867 43.875488 \nL 191.078364 39.576758 \nL 200.422862 44.636606 \nL 209.767359 44.981211 \nL 219.111857 38.283564 \nL 228.456355 42.768892 \nL 237.800852 40.169691 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n   </g>\n   <g id=\"line2d_18\">\n    <path d=\"M 60.255398 13.5 \nL 69.599895 21.257937 \nL 78.944393 27.009605 \nL 88.288891 35.161382 \nL 97.633388 43.026391 \nL 106.977886 53.246308 \nL 116.322383 57.534646 \nL 125.666881 58.075453 \nL 135.011379 70.535571 \nL 144.355876 80.372502 \nL 153.700374 87.973748 \nL 163.044871 97.278031 \nL 172.389369 104.29171 \nL 181.733867 105.306967 \nL 191.078364 112.525766 \nL 200.422862 119.139357 \nL 209.767359 126.90599 \nL 219.111857 131.689576 \nL 228.456355 136.554309 \nL 237.800852 139.5 \n\" clip-path=\"url(#pb3d7d49b1a)\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 51.378125 145.8 \nL 51.378125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 246.678125 145.8 \nL 246.678125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 51.378125 145.8 \nL 246.678125 145.8 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 51.378125 7.2 \nL 246.678125 7.2 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_7\">\n     <path d=\"M 58.378125 140.8 \nL 173.835938 140.8 \nQ 175.835938 140.8 175.835938 138.8 \nL 175.835938 110.44375 \nQ 175.835938 108.44375 173.835938 108.44375 \nL 58.378125 108.44375 \nQ 56.378125 108.44375 56.378125 110.44375 \nL 56.378125 138.8 \nQ 56.378125 140.8 58.378125 140.8 \nz\n\" style=\"fill: #ffffff; opacity: 0.8; stroke: #cccccc; stroke-linejoin: miter\"/>\n    </g>\n    <g id=\"line2d_19\">\n     <path d=\"M 60.378125 116.542188 \nL 70.378125 116.542188 \nL 80.378125 116.542188 \n\" style=\"fill: none; stroke: #1f77b4; stroke-width: 1.5; stroke-linecap: square\"/>\n    </g>\n    <g id=\"text_11\">\n     <!-- SGD, fixed step -->\n     <g transform=\"translate(88.378125 120.042188) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-53\" d=\"M 3425 4513 \nL 3425 3897 \nQ 3066 4069 2747 4153 \nQ 2428 4238 2131 4238 \nQ 1616 4238 1336 4038 \nQ 1056 3838 1056 3469 \nQ 1056 3159 1242 3001 \nQ 1428 2844 1947 2747 \nL 2328 2669 \nQ 3034 2534 3370 2195 \nQ 3706 1856 3706 1288 \nQ 3706 609 3251 259 \nQ 2797 -91 1919 -91 \nQ 1588 -91 1214 -16 \nQ 841 59 441 206 \nL 441 856 \nQ 825 641 1194 531 \nQ 1563 422 1919 422 \nQ 2459 422 2753 634 \nQ 3047 847 3047 1241 \nQ 3047 1584 2836 1778 \nQ 2625 1972 2144 2069 \nL 1759 2144 \nQ 1053 2284 737 2584 \nQ 422 2884 422 3419 \nQ 422 4038 858 4394 \nQ 1294 4750 2059 4750 \nQ 2388 4750 2728 4690 \nQ 3069 4631 3425 4513 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-47\" d=\"M 3809 666 \nL 3809 1919 \nL 2778 1919 \nL 2778 2438 \nL 4434 2438 \nL 4434 434 \nQ 4069 175 3628 42 \nQ 3188 -91 2688 -91 \nQ 1594 -91 976 548 \nQ 359 1188 359 2328 \nQ 359 3472 976 4111 \nQ 1594 4750 2688 4750 \nQ 3144 4750 3555 4637 \nQ 3966 4525 4313 4306 \nL 4313 3634 \nQ 3963 3931 3569 4081 \nQ 3175 4231 2741 4231 \nQ 1884 4231 1454 3753 \nQ 1025 3275 1025 2328 \nQ 1025 1384 1454 906 \nQ 1884 428 2741 428 \nQ 3075 428 3337 486 \nQ 3600 544 3809 666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-44\" d=\"M 1259 4147 \nL 1259 519 \nL 2022 519 \nQ 2988 519 3436 956 \nQ 3884 1394 3884 2338 \nQ 3884 3275 3436 3711 \nQ 2988 4147 2022 4147 \nL 1259 4147 \nz\nM 628 4666 \nL 1925 4666 \nQ 3281 4666 3915 4102 \nQ 4550 3538 4550 2338 \nQ 4550 1131 3912 565 \nQ 3275 0 1925 0 \nL 628 0 \nL 628 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-2c\" d=\"M 750 794 \nL 1409 794 \nL 1409 256 \nL 897 -744 \nL 494 -744 \nL 750 256 \nL 750 794 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-66\" d=\"M 2375 4863 \nL 2375 4384 \nL 1825 4384 \nQ 1516 4384 1395 4259 \nQ 1275 4134 1275 3809 \nL 1275 3500 \nL 2222 3500 \nL 2222 3053 \nL 1275 3053 \nL 1275 0 \nL 697 0 \nL 697 3053 \nL 147 3053 \nL 147 3500 \nL 697 3500 \nL 697 3744 \nQ 697 4328 969 4595 \nQ 1241 4863 1831 4863 \nL 2375 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-78\" d=\"M 3513 3500 \nL 2247 1797 \nL 3578 0 \nL 2900 0 \nL 1881 1375 \nL 863 0 \nL 184 0 \nL 1544 1831 \nL 300 3500 \nL 978 3500 \nL 1906 2253 \nL 2834 3500 \nL 3513 3500 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-53\"/>\n      <use xlink:href=\"#DejaVuSans-47\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-44\" transform=\"translate(140.966797 0)\"/>\n      <use xlink:href=\"#DejaVuSans-2c\" transform=\"translate(217.96875 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(249.755859 0)\"/>\n      <use xlink:href=\"#DejaVuSans-66\" transform=\"translate(281.542969 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(316.748047 0)\"/>\n      <use xlink:href=\"#DejaVuSans-78\" transform=\"translate(344.53125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(400.585938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(462.109375 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(525.585938 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(557.373047 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(609.472656 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(648.681641 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(710.205078 0)\"/>\n     </g>\n    </g>\n    <g id=\"line2d_20\">\n     <path d=\"M 60.378125 131.220313 \nL 70.378125 131.220313 \nL 80.378125 131.220313 \n\" style=\"fill: none; stroke-dasharray: 5.55,2.4; stroke-dashoffset: 0; stroke: #bf00bf; stroke-width: 1.5\"/>\n    </g>\n    <g id=\"text_12\">\n     <!-- SVRG, fixed step -->\n     <g transform=\"translate(88.378125 134.720313) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-56\" d=\"M 1831 0 \nL 50 4666 \nL 709 4666 \nL 2188 738 \nL 3669 4666 \nL 4325 4666 \nL 2547 0 \nL 1831 0 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-52\" d=\"M 2841 2188 \nQ 3044 2119 3236 1894 \nQ 3428 1669 3622 1275 \nL 4263 0 \nL 3584 0 \nL 2988 1197 \nQ 2756 1666 2539 1819 \nQ 2322 1972 1947 1972 \nL 1259 1972 \nL 1259 0 \nL 628 0 \nL 628 4666 \nL 2053 4666 \nQ 2853 4666 3247 4331 \nQ 3641 3997 3641 3322 \nQ 3641 2881 3436 2590 \nQ 3231 2300 2841 2188 \nz\nM 1259 4147 \nL 1259 2491 \nL 2053 2491 \nQ 2509 2491 2742 2702 \nQ 2975 2913 2975 3322 \nQ 2975 3731 2742 3939 \nQ 2509 4147 2053 4147 \nL 1259 4147 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-53\"/>\n      <use xlink:href=\"#DejaVuSans-56\" transform=\"translate(63.476562 0)\"/>\n      <use xlink:href=\"#DejaVuSans-52\" transform=\"translate(131.884766 0)\"/>\n      <use xlink:href=\"#DejaVuSans-47\" transform=\"translate(201.367188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-2c\" transform=\"translate(278.857422 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(310.644531 0)\"/>\n      <use xlink:href=\"#DejaVuSans-66\" transform=\"translate(342.431641 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(377.636719 0)\"/>\n      <use xlink:href=\"#DejaVuSans-78\" transform=\"translate(405.419922 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(461.474609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-64\" transform=\"translate(522.998047 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(586.474609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(618.261719 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(670.361328 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(709.570312 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(771.09375 0)\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"pb3d7d49b1a\">\n   <rect x=\"51.378125\" y=\"7.2\" width=\"195.3\" height=\"138.6\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=mdl-adaptive-stochastic-methods-md-8 fw=all -->\n",
    "\n",
    "SGD makes cheaper early progress, but its constant-step iterates continue to\n",
    "wander around the optimum. SVRG pays for its snapshots and then drives the gap\n",
    "down geometrically on this finite, strongly convex problem. The result does not\n",
    "say SVRG universally beats SGD: full passes are expensive on enormous or\n",
    "streaming datasets, the linear-rate theorem uses finite-sum smoothness and\n",
    "convexity, and stale control variates are less effective when the objective\n",
    "changes during training.\n",
    "\n",
    "**SAGA** removes the periodic full-gradient pass by storing the most recently\n",
    "seen gradient for every example and using their table average as the control\n",
    "variate. It spends one new component gradient per step but stores $O(nd)$\n",
    "numbers in the literal form (structure can reduce this for generalized linear\n",
    "models). SARAH-style estimators update the control variate recursively. The\n",
    "shared principle is more important than the acronyms: correlate a noisy new\n",
    "estimate with a stored estimate whose mean is known, subtract the shared noise,\n",
    "and add the known mean back.\n",
    "\n",
    "Variance reduction is distinct from momentum. Momentum averages gradients\n",
    "across recent iterates and can reduce high-frequency noise, but its moving\n",
    "average is generally biased for the gradient at the current point. SVRG and\n",
    "SAGA engineer an estimator that is unbiased at the current point and whose\n",
    "variance shrinks as the optimization converges.\n",
    "\n",
    "## Beyond Diagonals: the Preconditioning Ladder\n",
    "\n",
    "Step back and every method in this section is one answer to a single design\n",
    "question: *which matrix $B_t$ should multiply the gradient*,\n",
    "$\\mathbf{x}_{t+1} = \\mathbf{x}_t - \\eta\\, B_t\\, \\mathbf{g}_t$? Gradient\n",
    "descent answers $B = I$; Newton answers $B = (\\nabla^2 f)^{-1}$ at $O(d^3)$;\n",
    "Adam answers with a diagonal estimated from gradient history at $O(d)$.\n",
    "Between the diagonal and the full inverse Hessian runs a ladder of structured\n",
    "compromises, and 2020s practice has climbed it rung by rung. What makes the\n",
    "climb affordable is one structural fact the \"one long vector\n",
    "$\\mathbf{x} \\in \\mathbb{R}^d$\" notation hides: a network's parameters come in\n",
    "*layers of matrices*, and curvature between layers can be ignored while\n",
    "curvature within a layer is captured cheaply.\n",
    "\n",
    "**Kronecker-factored curvature (K-FAC)** [@Martens.Grosse.2015] treats\n",
    "each layer's weight matrix $W \\in \\mathbb{R}^{m \\times n}$ separately and\n",
    "approximates its block of the **Fisher information matrix**, the\n",
    "outer-product curvature\n",
    "$\\mathbb{E}[\\nabla \\log p\\, \\nabla \\log p^\\top]$ that\n",
    "that section introduced. The word *curvature* is\n",
    "precise here: for log-likelihood losses the Fisher equals the Gauss--Newton matrix,\n",
    "the PSD part of the Hessian obtained by dropping the second-derivative terms\n",
    "of the network [@Martens.Grosse.2015]. K-FAC's approximation is a\n",
    "Kronecker\n",
    "product $A \\otimes G$ (the block matrix whose $(i,j)$ block is $a_{ij}\\, G$),\n",
    "where $A = \\mathbb{E}[\\mathbf{a}\\mathbf{a}^\\top]$ is\n",
    "the second moment of the layer's *inputs* and\n",
    "$G = \\mathbb{E}[\\boldsymbol{\\delta}\\boldsymbol{\\delta}^\\top]$ that of the\n",
    "gradients at its *outputs*. The factorization is exact if the layer's inputs\n",
    "and output-gradients are statistically independent under the model's\n",
    "distribution; that independence, which holds only approximately in trained\n",
    "networks, is the approximation. The payoff is the\n",
    "Kronecker inverse identity: writing $\\mathrm{vec}(V)$ for the vector obtained\n",
    "by stacking the columns of the layer's gradient matrix $V$,\n",
    "\n",
    "$$\n",
    "(A \\otimes G)^{-1}\\, \\mathrm{vec}(V) \\;=\\; \\mathrm{vec}\\left(G^{-1}\\, V\\, A^{-1}\\right),\n",
    "$$\n",
    "\n",
    "so preconditioning by a curvature matrix with $(mn)^2$ entries costs two\n",
    "small inverses, of sizes $n \\times n$ and $m \\times m$: the whole trick in\n",
    "one display. Preconditioning by the Fisher has its own name and pedigree:\n",
    "it is the **natural gradient** [@Amari.1998], steepest descent when\n",
    "distance is measured\n",
    "between the *distributions* the parameters define rather than between\n",
    "parameter vectors, the metric view of AdaGrad again, with the Fisher as\n",
    "the metric. **Shampoo** [@Gupta.Koren.Singer.2018] keeps the same\n",
    "two-sided structure but builds the factors AdaGrad-style from accumulated\n",
    "gradient statistics, preconditioning $V \\mapsto L^{-1/4} V R^{-1/4}$ with\n",
    "$L = \\sum_t V_t V_t^\\top$ and $R = \\sum_t V_t^\\top V_t$, a two-sided matrix\n",
    "cousin of the equation. It is a cousin rather than a strict\n",
    "generalization: restricting $L, R$ to their diagonals normalizes entry\n",
    "$(i,j)$ by the factored $(l_i r_j)^{1/4}$ of its row and column statistics,\n",
    "not by AdaGrad's per-entry root, and the precise relation is an inequality\n",
    "(the Kronecker root $L^{1/4} \\otimes R^{1/4}$ upper-bounds the full-matrix\n",
    "AdaGrad preconditioner [@Gupta.Koren.Singer.2018]). **Muon**\n",
    "[@Jordan.Jin.Boza.ea.2024] takes the limiting view: rather than\n",
    "estimate second moments at all, it replaces each layer's momentum matrix by\n",
    "its nearest orthogonal matrix, the polar factor $UV^\\top$ of the SVD,\n",
    "computed by the Newton--Schulz iteration demonstrated in\n",
    "that section. This whitens the update's spectrum\n",
    "directly: every singular direction of the update moves at the same rate,\n",
    "per-*direction* step equalization where Adam manages only per-coordinate.\n",
    "\n",
    "The ladder, assembled: diagonal (Adam, $O(d)$) $\\to$ per-layer Kronecker\n",
    "factors (K-FAC) $\\to$ per-layer full-matrix roots (Shampoo) $\\to$ per-layer\n",
    "spectral normalization (Muon) $\\to$ full second order (Newton, unreachable).\n",
    "Each rung captures geometry the previous one ignored and costs more per-step\n",
    "arithmetic; where training runs land on the ladder is set by hardware\n",
    "economics as much as mathematics, which is why the frontier keeps moving.\n",
    "The same matching of step size to geometry also extends *across model\n",
    "sizes*: the maximal-update parametrization (muP) [@Yang.Hu.2021] derives\n",
    "how per-layer learning rates must scale with width so that one tuned\n",
    "$\\eta$ transfers from a small model to a large one.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* The **Ghadimi--Lan theorem** extends the no-convexity guarantee to\n",
    "  SGD: with $\\eta \\le 1/L$ the descent lemma survives in expectation with a\n",
    "  noise tax $\\tfrac{L}{2}\\eta^2\\sigma^2$ per step, and at the balanced\n",
    "  $\\eta \\propto 1/\\sqrt{K}$ a random iterate satisfies\n",
    "  $\\mathbb{E}\\|\\nabla f\\|^2 = O(\\sigma\\sqrt{L\\Delta/K})$. Noise turns the\n",
    "  deterministic $1/K$ rate into $K^{-1/2}$, and the\n",
    "  guarantee is for a *random* iterate (the best one cannot be identified\n",
    "  from noisy evaluations).\n",
    "* **Per-coordinate step sizes** attack $\\kappa$ directly: on a diagonal\n",
    "  quadratic, $\\eta_i = 1/\\lambda_i$ removes the condition number entirely.\n",
    "  **AdaGrad** is steepest descent in the metric\n",
    "  $\\mathrm{diag}(\\sqrt{\\sum_t \\mathbf{g}_t^2})$, built for sparse\n",
    "  gradients, but its never-forgetting denominator can drive steps toward zero\n",
    "  and make progress impractically slow on persistent noisy or nonconvex\n",
    "  problems; **RMSProp** substitutes an EMA; **Adam** adds momentum\n",
    "  and *exact* bias correction: $\\mathbb{E}[\\mathbf{v}_t] =\n",
    "  (1 - \\beta_2^t)\\,\\bar{\\mathbf{g}^2}$ under stationarity, so dividing by\n",
    "  $1 - \\beta_2^t$ cancels the zero-initialization deficit identically.\n",
    "* Vanilla Adam is **not guaranteed to converge in general**, even on convex\n",
    "  problems: in the Reddi--Kale--Kumar construction the effective step on a rare informative gradient shrinks\n",
    "  (by $\\sqrt{\\mathbf{v}}\\approx C$) faster than its information accrues, and\n",
    "  Adam converges to the *worst* point while AMSGrad's running-max fix (a\n",
    "  nonincreasing per-coordinate step) and plain SGD find the optimum. What\n",
    "  adaptivity wins back is the valley: on the $\\kappa = 10^3$ quadratic Adam\n",
    "  reconstructs the $1000\\times$ eigenvalue ratio from gradient magnitudes\n",
    "  alone and beats optimally tuned GD by $18\\times$.\n",
    "* **Weight decay through Adam's preconditioner is not weight decay**: the\n",
    "  shrinkage becomes $\\eta\\lambda/(\\sqrt{\\hat{v}_i} + \\epsilon)$ per\n",
    "  coordinate, strong where gradients are quiet, negligible where they are\n",
    "  large or noisy. **AdamW** decouples the decay from the preconditioner,\n",
    "  restoring the uniform $(1 - \\eta\\lambda)$ that the MAP-prior and\n",
    "  norm-constraint readings of $\\lambda$ presuppose.\n",
    "* **Schedules** negotiate the noise floor ($\\propto \\eta$) against transient\n",
    "  progress; beyond convexity no theorem ranks decay shapes. Cosine spends a\n",
    "  long small-step tail; **WSD** stays on the constant-step floor until a\n",
    "  final decay phase drops it (the loss cliff, reproduced on a noisy\n",
    "  quadratic). **Warmup** is the adaptive family's structural need: at small\n",
    "  $t$ the preconditioner is estimated from a handful of gradients (unbiased\n",
    "  $\\ne$ accurate), and the stability ceiling has not yet adapted to the\n",
    "  target step.\n",
    "* On finite-sum objectives, **SVRG** subtracts a component gradient at a\n",
    "  snapshot and adds the snapshot's full gradient, producing an unbiased\n",
    "  control-variate estimator whose variance shrinks near the snapshot. **SAGA**\n",
    "  replaces full-gradient refreshes by a table of stored component gradients.\n",
    "  These methods can reach linear fixed-step convergence under smooth strong\n",
    "  convexity, but require repeated access to a finite dataset.\n",
    "* Between Adam's diagonal and Newton's full inverse runs the\n",
    "  **preconditioning ladder**: K-FAC's Kronecker-factored Fisher (natural\n",
    "  gradient at two small inverses via\n",
    "  $(A \\otimes G)^{-1}\\mathrm{vec}(V) = \\mathrm{vec}(G^{-1}VA^{-1})$),\n",
    "  Shampoo's two-sided full-matrix roots, Muon's Newton--Schulz\n",
    "  orthogonalization of momentum, each rung exploiting the fact that\n",
    "  parameters come in matrices.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. **AdaGrad from the metric view.** Exercise 2 of\n",
    "   that section showed that the steepest\n",
    "   descent direction under the norm\n",
    "   $\\|\\mathbf{d}\\|_A = \\sqrt{\\mathbf{d}^\\top A \\mathbf{d}}$ is proportional\n",
    "   to $-A^{-1}\\nabla f$. Set\n",
    "   $A_t = \\mathrm{diag}\\bigl(\\sqrt{\\mathbf{s}_t}\\bigr) + \\epsilon I$ with\n",
    "   $\\mathbf{s}_t = \\sum_{s \\le t} \\mathbf{g}_s^2$ and recover\n",
    "   the equation. Why is $A_t$ built from\n",
    "   $\\sqrt{\\mathbf{s}_t}$ rather than $\\mathbf{s}_t$? (Two answers: match the\n",
    "   units of a curvature, since squared gradients have the wrong dimensions,\n",
    "   and check which choice makes the effective step on a coordinate with\n",
    "   i.i.d. noise decay like the Robbins--Monro-compatible $1/\\sqrt{t}$\n",
    "   rather than the too-fast $1/t$.)\n",
    "2. **Bias correction for the first moment.** Repeat the computation of\n",
    "   the equation for\n",
    "   $\\mathbf{m}_t = \\beta_1 \\mathbf{m}_{t-1} + (1-\\beta_1)\\mathbf{g}_t$ under\n",
    "   $\\mathbb{E}[\\mathbf{g}_s] = \\bar{\\mathbf{g}}$, obtaining\n",
    "   $\\mathbb{E}[\\mathbf{m}_t] = (1 - \\beta_1^t)\\,\\bar{\\mathbf{g}}$. Then\n",
    "   compute the scale of *uncorrected* Adam's first step: with\n",
    "   $\\mathbf{m}_1 = (1-\\beta_1)\\mathbf{g}_1$ and\n",
    "   $\\mathbf{v}_1 = (1-\\beta_2)\\mathbf{g}_1^2$, show the raw ratio\n",
    "   $\\mathbf{m}_1/\\sqrt{\\mathbf{v}_1}$ has magnitude\n",
    "   $(1-\\beta_1)/\\sqrt{1-\\beta_2} \\approx 3.16$ at the defaults. Finally,\n",
    "   find (numerically or by calculus) where the inflation factor\n",
    "   $(1-\\beta_1^t)/\\sqrt{1-\\beta_2^t}$ *peaks* and how large it gets,\n",
    "   and explain, from the two averaging timescales $1/(1-\\beta_1)$ and\n",
    "   $1/(1-\\beta_2)$, why the peak sits where it does.\n",
    "3. **The Reddi example by hand.** Take $C = 4$, $\\beta_1 = 0$,\n",
    "   $\\beta_2 = 1/17$, a fixed step $\\eta_t = \\eta$, and an iterate $x$ in the\n",
    "   interior of $[-1, 1]$ with $v$ at its period-limit value. Compute the\n",
    "   three Adam updates of one period of the equation and show\n",
    "   the net displacement over the period is positive (toward the wrong\n",
    "   endpoint) for small $\\eta$. Where exactly does the factor\n",
    "   $\\sqrt{v} \\approx C$ throttle the informative step?\n",
    "4. **Coupled versus decoupled fixed points.** For\n",
    "   $L(\\mathbf{w}) = \\tfrac12 (\\mathbf{w} - \\mathbf{w}_0)^\\top\n",
    "   \\mathrm{diag}(h_1, h_2) (\\mathbf{w} - \\mathbf{w}_0)$ with\n",
    "   $h = (100, 1)$ and deterministic gradients, write the stationarity\n",
    "   conditions of (a) Adam on $L + \\tfrac{\\lambda}{2}\\|\\mathbf{w}\\|^2$ and\n",
    "   (b) AdamW on $L$. Show that (a) demands\n",
    "   $\\nabla L + \\lambda\\mathbf{w} = \\mathbf{0}$ (the *same* point as SGD\n",
    "   with $\\ell_2$), while (b) demands\n",
    "   $\\nabla L_i = -\\lambda w_i (\\sqrt{\\hat{v}_i} + \\epsilon)$, a\n",
    "   *different* point. Reconcile this with the section's cell: in the noisy,\n",
    "   never-stationary regime, which of the two updates applies the uniform\n",
    "   per-step shrinkage, and why is that the property that matters during\n",
    "   training?\n",
    "5. **Where $\\epsilon$ matters.** On the $\\kappa = 10^3$ quadratic of the\n",
    "   `#adaptive-stochastic-methods-adam-vs-gd` cell, sweep\n",
    "   $\\epsilon \\in \\{10^{-12}, 10^{-8}, 10^{-4}, 10^{-2}, 10^{-1}\\}$ and\n",
    "   record iterations to $f < 10^{-8}$. Explain the two regimes you find:\n",
    "   $\\epsilon$ is invisible while $\\sqrt{\\hat{v}_i} \\gg \\epsilon$, and it\n",
    "   converts Adam into plain (momentum) SGD with step $\\eta/\\epsilon$ once\n",
    "   gradients shrink below it. Which coordinate's gradients cross that\n",
    "   threshold first, and what does $\\epsilon$ therefore control near\n",
    "   convergence?\n",
    "6. **The optimal constant step, exactly.** For the noisy scalar quadratic of\n",
    "   the `#adaptive-stochastic-methods-schedules` cell, iterate the noise-ball\n",
    "   recursion of that section to get the\n",
    "   exact finite-horizon loss\n",
    "   $\\mathbb{E}[x_K^2] = (1-\\eta\\lambda)^{2K} x_0^2 +\n",
    "   \\bigl(1 - (1-\\eta\\lambda)^{2K}\\bigr)\\, \\eta\\sigma^2/(\\lambda(2-\\eta\\lambda))$.\n",
    "   For a fixed budget $K$, show the minimizing constant step scales like\n",
    "   $\\eta^\\star = \\Theta(\\log K / K)$ (balance the exponential transient\n",
    "   against the linear-in-$\\eta$ floor), *not* like the $1/\\sqrt{K}$ of the\n",
    "   Ghadimi--Lan prescription, and explain why there is no contradiction\n",
    "   (strong convexity versus the nonconvex worst case). Verify your\n",
    "   $\\eta^\\star$ numerically against a grid search at $K = 2000$.\n",
    "7. **AMSGrad's monotone steps.** Show that with the running maximum\n",
    "   $\\tilde{\\mathbf{v}}_t = \\max(\\tilde{\\mathbf{v}}_{t-1}, \\mathbf{v}_t)$ the\n",
    "   effective per-coordinate step $\\eta_t/\\sqrt{\\tilde{v}_{t,i}}$ is\n",
    "   nonincreasing in $t$ whenever $\\eta_t$ is, and explain in one sentence\n",
    "   why this property blocks the drift mechanism of\n",
    "   the equation. Then modify the\n",
    "   `#adaptive-stochastic-methods-reddi` cell to *interpolate*: replace the\n",
    "   hard max by $\\tilde{v}_t = \\max(\\gamma\\tilde{v}_{t-1}, v_t)$ and find the\n",
    "   largest forgetting factor $\\gamma < 1$ at which the drift to $+1$\n",
    "   reappears.\n",
    "\n",
    "8. **SVRG unbiasedness and cost.** Prove the conditional unbiasedness of\n",
    "   the equation. Count component-gradient evaluations in one epoch\n",
    "   of the cell, then vary the inner-loop length from $n/2$ to $4n$. Explain the\n",
    "   tradeoff between amortizing the snapshot gradient and letting the control\n",
    "   variate become stale.\n",
    "\n",
    "## Discussions\n",
    "\n",
    "This section is the bridge between this chapter's theory and the optimizer\n",
    "actually running in your training loop. The main book shows the methods in\n",
    "action (that section and that section for stochastic\n",
    "descent, that section for velocity, that section for the\n",
    "adaptive family, that section for schedules), while the\n",
    "Ghadimi--Lan rate, the diagonal-metric derivation, the bias-correction\n",
    "identity, the Reddi counterexample, and the decoupling algebra proved here\n",
    "are the mathematics those recipes stand on. Within this part,\n",
    "that section supplied every instrument\n",
    "(descent lemma, $\\kappa$, minibatch variance, noise ball);\n",
    "that section explains which guarantees return when convexity\n",
    "does; and the stationary-distribution view of constant-step SGD, the noise\n",
    "ball as an invariant measure, is developed properly with the SDE machinery\n",
    "of that section.\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/adaptive-stochastic-methods)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
