Skip to content

Commit 4d60086

Browse files
pantelisclaude
andcommitted
feat: publish optimization SGD and optimizer-zoo notebooks
From-scratch gradient descent / SGD notebook and the momentum-to-Adam optimizer zoo, both executed in torch.dev.gpu, plus registry entries. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent de2f7f2 commit 4d60086

3 files changed

Lines changed: 636 additions & 0 deletions

File tree

notebooks/notebook-database.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,3 +533,17 @@ notebooks:
533533
description: SGD on the regularized ridge loss for degree-9 polynomial regression, lambda anchored to the closed-form optimum
534534
last_executed: 2026-05-30
535535
duration_seconds: 17.7
536+
- source: aiml-common/lectures/optimization/sgd/index.ipynb
537+
notebook: optimization/sgd/index.ipynb
538+
code_cells: 10
539+
environment: torch.dev.gpu
540+
description: Gradient descent, SGD and mini-batch from scratch in NumPy with a PyTorch coda
541+
last_executed: 2026-05-30
542+
duration_seconds: 5.4
543+
- source: aiml-common/lectures/optimization/optimizers/index.ipynb
544+
notebook: optimization/optimizers/index.ipynb
545+
code_cells: 8
546+
environment: torch.dev.gpu
547+
description: Momentum, Nesterov, RMSProp and Adam from scratch on ravine and saddle landscapes, with a PyTorch coda
548+
last_executed: 2026-05-31
549+
duration_seconds: 4.8
Lines changed: 297 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,297 @@
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

Comments
 (0)