|
| 1 | +{ |
| 2 | + "cells": [ |
| 3 | + { |
| 4 | + "cell_type": "markdown", |
| 5 | + "metadata": {}, |
| 6 | + "source": [ |
| 7 | + "Plain gradient descent, the subject of the [previous notebook](/aiml-common/lectures/optimization/sgd), struggles on two landscape features that are everywhere in deep learning. In a **ravine**, where the loss curves far more sharply in one direction than another, it zigzags across the steep walls while creeping along the shallow floor. Near a **saddle point**, where the gradient nearly vanishes, it stalls. This notebook builds the standard fixes from scratch, momentum, Nesterov, RMSProp, and Adam, compares their trajectories on both landscapes, and then shows the same optimizers through `torch.optim`. Throughout, $\\boldsymbol{\\theta}$ is the parameter vector and $\\mathbf{g} = \\nabla_{\\boldsymbol{\\theta}} L$ is the gradient." |
| 8 | + ] |
| 9 | + }, |
| 10 | + { |
| 11 | + "cell_type": "code", |
| 12 | + "metadata": { |
| 13 | + "tags": [ |
| 14 | + "remove-cell" |
| 15 | + ] |
| 16 | + }, |
| 17 | + "execution_count": null, |
| 18 | + "source": [ |
| 19 | + "import numpy as np\n", |
| 20 | + "import matplotlib.pyplot as plt\n", |
| 21 | + "import seaborn as sns\n", |
| 22 | + "sns.set_theme()" |
| 23 | + ], |
| 24 | + "outputs": [] |
| 25 | + }, |
| 26 | + { |
| 27 | + "cell_type": "markdown", |
| 28 | + "metadata": {}, |
| 29 | + "source": [ |
| 30 | + "## A ravine\n", |
| 31 | + "\n", |
| 32 | + "Take an anisotropic quadratic that is steep in one coordinate and shallow in the other,\n", |
| 33 | + "\n", |
| 34 | + "$$L(\\boldsymbol{\\theta}) = \\tfrac{1}{2}\\big(\\theta_0^2 + \\kappa\\,\\theta_1^2\\big), \\qquad \\nabla_{\\boldsymbol{\\theta}} L = (\\theta_0,\\; \\kappa\\,\\theta_1).$$\n", |
| 35 | + "\n", |
| 36 | + "The Hessian has eigenvalues $1$ and $\\kappa$, so the **condition number** is $\\kappa$. Gradient descent is stable only while $\\eta < 2/\\kappa$, set by the steep direction, but then the shallow direction contracts by just $(1 - \\eta)$ per step. With $\\kappa$ large the steep coordinate oscillates while the shallow one barely moves: the familiar zigzag down a narrow valley." |
| 37 | + ] |
| 38 | + }, |
| 39 | + { |
| 40 | + "cell_type": "code", |
| 41 | + "metadata": {}, |
| 42 | + "execution_count": null, |
| 43 | + "source": [ |
| 44 | + "KAPPA = 100.0\n", |
| 45 | + "def L(theta):\n", |
| 46 | + " return 0.5 * (theta[0]**2 + KAPPA * theta[1]**2)\n", |
| 47 | + "def grad(theta):\n", |
| 48 | + " return np.array([theta[0], KAPPA * theta[1]])\n", |
| 49 | + "\n", |
| 50 | + "start = np.array([-9.0, -1.0]) # common starting point; minimum is at the origin\n", |
| 51 | + "print(f\"condition number kappa = {KAPPA:.0f}\")" |
| 52 | + ], |
| 53 | + "outputs": [] |
| 54 | + }, |
| 55 | + { |
| 56 | + "cell_type": "markdown", |
| 57 | + "metadata": {}, |
| 58 | + "source": [ |
| 59 | + "## Momentum\n", |
| 60 | + "\n", |
| 61 | + "Momentum accumulates a running **velocity** $\\mathbf{v}$ that averages successive gradients. Oscillating components (the steep direction) cancel, while the consistent component (the shallow floor) builds up, so the iterate accelerates along the valley instead of bouncing across it:\n", |
| 62 | + "\n", |
| 63 | + "$$\\mathbf{v}_t = \\mu\\,\\mathbf{v}_{t-1} + \\mathbf{g}_t, \\qquad \\boldsymbol{\\theta}_t = \\boldsymbol{\\theta}_{t-1} - \\eta\\,\\mathbf{v}_t,$$\n", |
| 64 | + "\n", |
| 65 | + "with momentum coefficient $\\mu \\in [0,1)$. **Nesterov** momentum evaluates the gradient at a look-ahead point $\\boldsymbol{\\theta} + \\mu\\mathbf{v}$, which corrects overshoot and usually converges a little faster." |
| 66 | + ] |
| 67 | + }, |
| 68 | + { |
| 69 | + "cell_type": "markdown", |
| 70 | + "metadata": {}, |
| 71 | + "source": [ |
| 72 | + "## Adaptive methods: RMSProp and Adam\n", |
| 73 | + "\n", |
| 74 | + "A different fix rescales each coordinate by its own recent gradient magnitude, so steep directions take smaller steps and shallow directions larger ones automatically. **RMSProp** keeps an exponential average of squared gradients $\\mathbf{s}$ and divides by its root,\n", |
| 75 | + "\n", |
| 76 | + "$$\\mathbf{s}_t = \\rho\\,\\mathbf{s}_{t-1} + (1-\\rho)\\,\\mathbf{g}_t^2, \\qquad \\boldsymbol{\\theta}_t = \\boldsymbol{\\theta}_{t-1} - \\frac{\\eta}{\\sqrt{\\mathbf{s}_t} + \\epsilon}\\,\\mathbf{g}_t.$$\n", |
| 77 | + "\n", |
| 78 | + "**Adam** combines this with momentum, tracking averages of both the gradient ($\\mathbf{m}$) and its square ($\\mathbf{v}$), each bias-corrected,\n", |
| 79 | + "\n", |
| 80 | + "$$\\mathbf{m}_t = \\beta_1\\mathbf{m}_{t-1} + (1-\\beta_1)\\mathbf{g}_t, \\quad \\mathbf{v}_t = \\beta_2\\mathbf{v}_{t-1} + (1-\\beta_2)\\mathbf{g}_t^2, \\quad \\boldsymbol{\\theta}_t = \\boldsymbol{\\theta}_{t-1} - \\eta\\,\\frac{\\hat{\\mathbf{m}}_t}{\\sqrt{\\hat{\\mathbf{v}}_t} + \\epsilon}.$$\n", |
| 81 | + "\n", |
| 82 | + "Each optimizer below is a small step function that reads and updates its own state; a shared runner iterates it and records the path." |
| 83 | + ] |
| 84 | + }, |
| 85 | + { |
| 86 | + "cell_type": "code", |
| 87 | + "metadata": {}, |
| 88 | + "execution_count": null, |
| 89 | + "source": [ |
| 90 | + "def run(step, grad, theta0, n=120):\n", |
| 91 | + " theta, state, path = np.array(theta0, float), {}, [np.array(theta0, float)]\n", |
| 92 | + " for t in range(1, n + 1):\n", |
| 93 | + " theta = step(theta, grad, state, t)\n", |
| 94 | + " path.append(theta.copy())\n", |
| 95 | + " return np.array(path)\n", |
| 96 | + "\n", |
| 97 | + "def sgd(lr):\n", |
| 98 | + " def step(th, grad, s, t):\n", |
| 99 | + " return th - lr * grad(th)\n", |
| 100 | + " return step\n", |
| 101 | + "\n", |
| 102 | + "def momentum(lr, mu=0.9):\n", |
| 103 | + " def step(th, grad, s, t):\n", |
| 104 | + " s[\"v\"] = mu * s.get(\"v\", 0.0) + grad(th)\n", |
| 105 | + " return th - lr * s[\"v\"]\n", |
| 106 | + " return step\n", |
| 107 | + "\n", |
| 108 | + "def nesterov(lr, mu=0.9):\n", |
| 109 | + " def step(th, grad, s, t):\n", |
| 110 | + " v = s.get(\"v\", np.zeros_like(th))\n", |
| 111 | + " v = mu * v - lr * grad(th + mu * v)\n", |
| 112 | + " s[\"v\"] = v\n", |
| 113 | + " return th + v\n", |
| 114 | + " return step\n", |
| 115 | + "\n", |
| 116 | + "def rmsprop(lr, rho=0.99, eps=1e-8):\n", |
| 117 | + " def step(th, grad, s, t):\n", |
| 118 | + " g = grad(th)\n", |
| 119 | + " s[\"s\"] = rho * s.get(\"s\", 0.0) + (1 - rho) * g**2\n", |
| 120 | + " return th - lr * g / (np.sqrt(s[\"s\"]) + eps)\n", |
| 121 | + " return step\n", |
| 122 | + "\n", |
| 123 | + "def adam(lr, b1=0.9, b2=0.999, eps=1e-8):\n", |
| 124 | + " def step(th, grad, s, t):\n", |
| 125 | + " g = grad(th)\n", |
| 126 | + " s[\"m\"] = b1 * s.get(\"m\", 0.0) + (1 - b1) * g\n", |
| 127 | + " s[\"v\"] = b2 * s.get(\"v\", 0.0) + (1 - b2) * g**2\n", |
| 128 | + " mhat, vhat = s[\"m\"] / (1 - b1**t), s[\"v\"] / (1 - b2**t)\n", |
| 129 | + " return th - lr * mhat / (np.sqrt(vhat) + eps)\n", |
| 130 | + " return step" |
| 131 | + ], |
| 132 | + "outputs": [] |
| 133 | + }, |
| 134 | + { |
| 135 | + "cell_type": "code", |
| 136 | + "metadata": {}, |
| 137 | + "execution_count": null, |
| 138 | + "source": [ |
| 139 | + "paths = {\n", |
| 140 | + " \"SGD\": run(sgd(0.018), grad, start),\n", |
| 141 | + " \"Momentum\": run(momentum(0.012, 0.85), grad, start),\n", |
| 142 | + " \"Nesterov\": run(nesterov(0.008, 0.85), grad, start),\n", |
| 143 | + " \"RMSProp\": run(rmsprop(0.15), grad, start),\n", |
| 144 | + " \"Adam\": run(adam(0.5), grad, start),\n", |
| 145 | + "}\n", |
| 146 | + "for name, p in paths.items():\n", |
| 147 | + " print(f\"{name:9s} final L = {L(p[-1]):.3g}\")" |
| 148 | + ], |
| 149 | + "outputs": [] |
| 150 | + }, |
| 151 | + { |
| 152 | + "cell_type": "code", |
| 153 | + "metadata": { |
| 154 | + "tags": [ |
| 155 | + "hide-input" |
| 156 | + ] |
| 157 | + }, |
| 158 | + "execution_count": null, |
| 159 | + "source": [ |
| 160 | + "g0 = np.linspace(-10, 10, 240); g1 = np.linspace(-2.2, 2.2, 240)\n", |
| 161 | + "G0, G1 = np.meshgrid(g0, g1)\n", |
| 162 | + "Z = 0.5 * (G0**2 + KAPPA * G1**2)\n", |
| 163 | + "plt.figure(figsize=[11, 7])\n", |
| 164 | + "plt.contour(G0, G1, Z, levels=np.logspace(-0.5, 2.5, 22), cmap=\"Greys\", alpha=0.5)\n", |
| 165 | + "for name, p in paths.items():\n", |
| 166 | + " plt.plot(p[:, 0], p[:, 1], \"-\", lw=1.5, label=name)\n", |
| 167 | + "plt.scatter([0], [0], color=\"g\", zorder=5, label=\"minimum\")\n", |
| 168 | + "plt.xlabel(r\"$\\theta_0$\"); plt.ylabel(r\"$\\theta_1$\")\n", |
| 169 | + "plt.legend(); plt.title(r\"optimizer paths down a ravine ($\\kappa = 100$)\")\n", |
| 170 | + "plt.show()" |
| 171 | + ], |
| 172 | + "outputs": [] |
| 173 | + }, |
| 174 | + { |
| 175 | + "cell_type": "markdown", |
| 176 | + "metadata": {}, |
| 177 | + "source": [ |
| 178 | + "## Saddle points\n", |
| 179 | + "\n", |
| 180 | + "In high dimensions most critical points where the gradient vanishes are not minima but **saddles**, low along some directions and high along others. A clean two-dimensional model is\n", |
| 181 | + "\n", |
| 182 | + "$$L(\\boldsymbol{\\theta}) = \\tfrac{1}{2}\\big(\\theta_0^2 - \\theta_1^2\\big), \\qquad \\nabla_{\\boldsymbol{\\theta}} L = (\\theta_0,\\; -\\theta_1),$$\n", |
| 183 | + "\n", |
| 184 | + "with a saddle at the origin. Starting almost on the ridge ($\\theta_1 \\approx 0$) the gradient in the escape direction is tiny, so plain gradient descent and momentum dawdle near the origin, while the per-coordinate scaling in RMSProp and Adam amplifies the weak direction and breaks away sooner." |
| 185 | + ] |
| 186 | + }, |
| 187 | + { |
| 188 | + "cell_type": "code", |
| 189 | + "metadata": {}, |
| 190 | + "execution_count": null, |
| 191 | + "source": [ |
| 192 | + "def L_saddle(theta):\n", |
| 193 | + " return 0.5 * (theta[0]**2 - theta[1]**2)\n", |
| 194 | + "def grad_saddle(theta):\n", |
| 195 | + " return np.array([theta[0], -theta[1]])\n", |
| 196 | + "\n", |
| 197 | + "start_s = np.array([-1.8, 1e-2]) # almost on the ridge\n", |
| 198 | + "paths_s = {\n", |
| 199 | + " \"SGD\": run(sgd(0.08), grad_saddle, start_s, n=35),\n", |
| 200 | + " \"Momentum\": run(momentum(0.04), grad_saddle, start_s, n=35),\n", |
| 201 | + " \"RMSProp\": run(rmsprop(0.03), grad_saddle, start_s, n=35),\n", |
| 202 | + " \"Adam\": run(adam(0.05), grad_saddle, start_s, n=35),\n", |
| 203 | + "}\n", |
| 204 | + "for name, p in paths_s.items():\n", |
| 205 | + " print(f\"{name:9s} |theta_1| after 35 steps = {abs(p[-1, 1]):.3f}\")" |
| 206 | + ], |
| 207 | + "outputs": [] |
| 208 | + }, |
| 209 | + { |
| 210 | + "cell_type": "code", |
| 211 | + "metadata": { |
| 212 | + "tags": [ |
| 213 | + "hide-input" |
| 214 | + ] |
| 215 | + }, |
| 216 | + "execution_count": null, |
| 217 | + "source": [ |
| 218 | + "g0 = np.linspace(-2, 2, 200); g1 = np.linspace(-3, 3, 200)\n", |
| 219 | + "G0, G1 = np.meshgrid(g0, g1)\n", |
| 220 | + "Z = 0.5 * (G0**2 - G1**2)\n", |
| 221 | + "plt.figure(figsize=[9, 8])\n", |
| 222 | + "plt.contour(G0, G1, Z, levels=25, cmap=\"Greys\", alpha=0.5)\n", |
| 223 | + "for name, p in paths_s.items():\n", |
| 224 | + " plt.plot(p[:, 0], p[:, 1], \"-o\", ms=3, lw=1.2, label=name)\n", |
| 225 | + "plt.scatter([0], [0], color=\"r\", marker=\"x\", zorder=5, label=\"saddle\")\n", |
| 226 | + "plt.xlim(-2, 2); plt.ylim(-3, 3)\n", |
| 227 | + "plt.xlabel(r\"$\\theta_0$\"); plt.ylabel(r\"$\\theta_1$\")\n", |
| 228 | + "plt.legend(); plt.title(\"escaping a saddle point\"); plt.show()" |
| 229 | + ], |
| 230 | + "outputs": [] |
| 231 | + }, |
| 232 | + { |
| 233 | + "cell_type": "markdown", |
| 234 | + "metadata": {}, |
| 235 | + "source": [ |
| 236 | + "## The same optimizers in PyTorch\n", |
| 237 | + "\n", |
| 238 | + "`torch.optim` ships these as one-line choices. Vanilla SGD takes a `momentum` argument, and Adam is its own class. Optimizing the ravine through autograd reproduces what you built by hand." |
| 239 | + ] |
| 240 | + }, |
| 241 | + { |
| 242 | + "cell_type": "code", |
| 243 | + "metadata": {}, |
| 244 | + "execution_count": null, |
| 245 | + "source": [ |
| 246 | + "import torch\n", |
| 247 | + "\n", |
| 248 | + "def torch_descend(make_opt, n=120):\n", |
| 249 | + " theta = torch.tensor([-9.0, -1.0], requires_grad=True)\n", |
| 250 | + " opt = make_opt([theta])\n", |
| 251 | + " loss = None\n", |
| 252 | + " for _ in range(n):\n", |
| 253 | + " opt.zero_grad()\n", |
| 254 | + " loss = 0.5 * (theta[0]**2 + KAPPA * theta[1]**2)\n", |
| 255 | + " loss.backward()\n", |
| 256 | + " opt.step()\n", |
| 257 | + " return theta.detach().numpy(), loss.item()\n", |
| 258 | + "\n", |
| 259 | + "for name, make_opt in {\n", |
| 260 | + " \"SGD\": lambda p: torch.optim.SGD(p, lr=0.018),\n", |
| 261 | + " \"SGD+momentum\": lambda p: torch.optim.SGD(p, lr=0.012, momentum=0.85),\n", |
| 262 | + " \"Adam\": lambda p: torch.optim.Adam(p, lr=0.5),\n", |
| 263 | + "}.items():\n", |
| 264 | + " theta, loss = torch_descend(make_opt)\n", |
| 265 | + " print(f\"{name:13s} final L = {loss:.3g} theta = {np.round(theta, 3)}\")" |
| 266 | + ], |
| 267 | + "outputs": [] |
| 268 | + }, |
| 269 | + { |
| 270 | + "cell_type": "markdown", |
| 271 | + "metadata": {}, |
| 272 | + "source": [ |
| 273 | + "## Takeaways\n", |
| 274 | + "\n", |
| 275 | + "- Plain gradient descent is limited by the **steepest** direction, so on an ill-conditioned ravine it zigzags and the shallow direction crawls.\n", |
| 276 | + "- **Momentum** averages gradients into a velocity that cancels the oscillation and accelerates along the valley; **Nesterov** sharpens this with a look-ahead gradient.\n", |
| 277 | + "- **Adaptive** methods (RMSProp, Adam) rescale each coordinate by its own gradient history, which both fixes the conditioning and helps escape saddle points where one direction is nearly flat.\n", |
| 278 | + "- **Adam** is momentum plus per-coordinate scaling with bias correction, the common default; well-tuned SGD with momentum often matches or beats it on large problems.\n", |
| 279 | + "- `torch.optim` provides all of these; the update rules are exactly the ones implemented here by hand.\n", |
| 280 | + "\n", |
| 281 | + "**Key references**: [@Kingma2014-ua; @Ruder2016-overview; @Goodfellow2014-ub]" |
| 282 | + ] |
| 283 | + } |
| 284 | + ], |
| 285 | + "metadata": { |
| 286 | + "kernelspec": { |
| 287 | + "display_name": "Python 3", |
| 288 | + "language": "python", |
| 289 | + "name": "python3" |
| 290 | + }, |
| 291 | + "language_info": { |
| 292 | + "name": "python" |
| 293 | + } |
| 294 | + }, |
| 295 | + "nbformat": 4, |
| 296 | + "nbformat_minor": 5 |
| 297 | +} |
0 commit comments