{ "cells": [ { "cell_type": "markdown", "id": "a29fbc49", "metadata": {}, "source": [ "# Lecture 4 · Computation graphs, backpropagation, and autograd\n", "\n", "This is the **single complete companion** to the lecture. It follows the same\n", "numbers and notation as the slides:\n", "\n", "1. local backward rules;\n", "2. one complete scalar graph;\n", "3. symbolic differentiation and finite differences as independent checks;\n", "4. the same graph in PyTorch autograd;\n", "5. gradient accumulation and `zero_grad()`;\n", "6. branch accumulation and a vector neuron;\n", "7. one dense-layer vector-Jacobian product;\n", "8. a three-example mean batch and equivalent microbatches;\n", "9. one concrete optimizer step and a tiny MLP.\n", "\n", "**Evidence contract.** Every input is a declared teaching construction. Every\n", "displayed result is freshly computed by the cells below. There is no dataset\n", "download, randomness, or hidden model training.\n" ] }, { "cell_type": "markdown", "id": "4c60af69", "metadata": {}, "source": [ "## Checkpoint 1 · Setup and a compact display helper\n", "\n", "The notebook uses only PyTorch and standard Colab/Jupyter display utilities.\n", "Double precision keeps the small arithmetic easy to compare with the slides.\n" ] }, { "cell_type": "code", "execution_count": 1, "id": "ccafe4e4", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:21.513521Z", "iopub.status.busy": "2026-08-20T12:58:21.513444Z", "iopub.status.idle": "2026-08-20T12:58:22.006289Z", "shell.execute_reply": "2026-08-20T12:58:22.005581Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "torch 2.13.0\n", "default dtype torch.float64\n" ] } ], "source": [ "import html\n", "import math\n", "import torch\n", "from IPython.display import HTML, display\n", "\n", "torch.set_default_dtype(torch.float64)\n", "\n", "def fmt(value):\n", " if isinstance(value, torch.Tensor):\n", " value = value.detach().cpu().tolist()\n", " return html.escape(str(value))\n", "\n", "def show_table(headers, rows, caption=None):\n", " head = \"\".join(f\"{html.escape(str(h))}\" for h in headers)\n", " body = \"\".join(\n", " \"\" + \"\".join(f\"{fmt(v)}\" for v in row) + \"\"\n", " for row in rows\n", " )\n", " cap = f\"{html.escape(caption)}\" if caption else \"\"\n", " display(HTML(\n", " \"\"\n", " f\"
{cap}{head}\"\n", " f\"{body}
\"\n", " ))\n", "\n", "print(\"torch\", torch.__version__)\n", "print(\"default dtype\", torch.get_default_dtype())\n" ] }, { "cell_type": "markdown", "id": "02e8021f", "metadata": {}, "source": [ "## Checkpoint 2 · Three local backward rules\n", "\n", "If a later part of the graph sends an upstream gradient $g_v=\\partial L/\\partial v$,\n", "each primitive returns $g_v$ times its local derivative.\n", "\n", "- square: $v=u^2 \\Rightarrow g_u=g_v(2u)$;\n", "- addition: $v=a+b \\Rightarrow (g_a,g_b)=(g_v,g_v)$;\n", "- multiplication: $v=ab \\Rightarrow (g_a,g_b)=(g_vb,g_va)$.\n", "\n", "The square example explicitly **assumes the rest of the graph sends $g_v=7$**.\n", "That number is not produced by the square; it arrives from downstream.\n" ] }, { "cell_type": "code", "execution_count": 2, "id": "aceb12a1", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.007627Z", "iopub.status.busy": "2026-08-20T12:58:22.007493Z", "iopub.status.idle": "2026-08-20T12:58:22.011295Z", "shell.execute_reply": "2026-08-20T12:58:22.010870Z" } }, "outputs": [ { "data": { "text/html": [ "
Local rules: upstream × local
operationforwardarriving gradientreturned gradient(s)
squarev = 3^2 = 9g_v = 7 (given)g_u = 7(2·3) = 42
additionv = 2 + 5 = 7g_v = 3g_a = 3, g_b = 3
multiplicationv = 2·5 = 10g_v = 3g_a = 15, g_b = 6
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# Square: the arriving upstream gradient is deliberately supplied as 7.\n", "u = 3.0\n", "g_v_square = 7.0\n", "v_square = u**2\n", "g_u_square = g_v_square * (2*u)\n", "\n", "# Addition and multiplication use the same arriving gradient 3.\n", "a_local, b_local, g_v_local = 2.0, 5.0, 3.0\n", "g_a_add, g_b_add = g_v_local, g_v_local\n", "g_a_mul = g_v_local * b_local\n", "g_b_mul = g_v_local * a_local\n", "\n", "assert (v_square, g_u_square) == (9.0, 42.0)\n", "assert (g_a_add, g_b_add) == (3.0, 3.0)\n", "assert (g_a_mul, g_b_mul) == (15.0, 6.0)\n", "\n", "show_table(\n", " [\"operation\", \"forward\", \"arriving gradient\", \"returned gradient(s)\"],\n", " [\n", " [\"square\", \"v = 3^2 = 9\", \"g_v = 7 (given)\", \"g_u = 7(2·3) = 42\"],\n", " [\"addition\", \"v = 2 + 5 = 7\", \"g_v = 3\", \"g_a = 3, g_b = 3\"],\n", " [\"multiplication\", \"v = 2·5 = 10\", \"g_v = 3\", \"g_a = 15, g_b = 6\"],\n", " ],\n", " \"Local rules: upstream × local\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "4fabc7bc", "metadata": {}, "source": [ "## Checkpoint 3 · Forward through the complete scalar graph\n", "\n", "We decompose\n", "\n", "$$L=(wx+b-y)^2$$\n", "\n", "into primitive values $m=wx$, $a=m+b$, $e=a-y$, and $L=e^2$.\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "6b970b02", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.012316Z", "iopub.status.busy": "2026-08-20T12:58:22.012249Z", "iopub.status.idle": "2026-08-20T12:58:22.014606Z", "shell.execute_reply": "2026-08-20T12:58:22.014228Z" } }, "outputs": [ { "data": { "text/html": [ "
Forward pass
stored valuecalculationvalue
mw·x6.0
am+b7.0
ea-y-3.0
Le^29.0
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "x, w, b, y = 3.0, 2.0, 1.0, 10.0\n", "m = w*x\n", "a = m+b\n", "e = a-y\n", "L = e**2\n", "\n", "assert (m, a, e, L) == (6.0, 7.0, -3.0, 9.0)\n", "show_table(\n", " [\"stored value\", \"calculation\", \"value\"],\n", " [\n", " [\"m\", \"w·x\", m],\n", " [\"a\", \"m+b\", a],\n", " [\"e\", \"a-y\", e],\n", " [\"L\", \"e^2\", L],\n", " ],\n", " \"Forward pass\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "5ae554cc", "metadata": {}, "source": [ "## Checkpoint 4 · Reverse sweep through the same stored values\n", "\n", "Seed $g_L=\\partial L/\\partial L=1$. Then repeatedly apply\n", "**upstream × local**. When a value fans out, returned contributions add.\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "f76ef53f", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.015615Z", "iopub.status.busy": "2026-08-20T12:58:22.015556Z", "iopub.status.idle": "2026-08-20T12:58:22.018410Z", "shell.execute_reply": "2026-08-20T12:58:22.018053Z" } }, "outputs": [ { "data": { "text/html": [ "
All stored values and their gradients
stored value qqg_q = ∂L/∂q
L9.01.0
e-3.0-6.0
a7.0-6.0
m6.0-6.0
w2.0-18.0
x3.0-12.0
b1.0-6.0
y10.06.0
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "gL = 1.0\n", "ge = gL * (2*e) # L=e^2\n", "ga = ge * 1.0 # e=a-y\n", "gy = ge * (-1.0)\n", "gm = ga * 1.0 # a=m+b\n", "gb = ga * 1.0\n", "gw = gm * x # m=w*x\n", "gx = gm * w\n", "\n", "expected = {\"L\": 1.0, \"e\": -6.0, \"a\": -6.0, \"m\": -6.0,\n", " \"w\": -18.0, \"x\": -12.0, \"b\": -6.0, \"y\": 6.0}\n", "actual = {\"L\": gL, \"e\": ge, \"a\": ga, \"m\": gm,\n", " \"w\": gw, \"x\": gx, \"b\": gb, \"y\": gy}\n", "assert actual == expected\n", "\n", "show_table(\n", " [\"stored value q\", \"q\", \"g_q = ∂L/∂q\"],\n", " [[name, {\"L\":L,\"e\":e,\"a\":a,\"m\":m,\"w\":w,\"x\":x,\"b\":b,\"y\":y}[name], grad]\n", " for name, grad in actual.items()],\n", " \"All stored values and their gradients\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "6694d4f9", "metadata": {}, "source": [ "## Checkpoint 5 · Two independent checks\n", "\n", "**Symbolic differentiation** transforms a formula into a derivative formula;\n", "it can be done on paper or by a computer-algebra system. **Finite differences**\n", "instead probe nearby loss values and approximate one derivative number.\n", "Neither is how PyTorch performs reverse-mode autograd.\n" ] }, { "cell_type": "code", "execution_count": 5, "id": "c8cd0f79", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.019492Z", "iopub.status.busy": "2026-08-20T12:58:22.019423Z", "iopub.status.idle": "2026-08-20T12:58:22.023393Z", "shell.execute_reply": "2026-08-20T12:58:22.023058Z" } }, "outputs": [ { "data": { "text/html": [ "
Three routes agree
qmanual reversesymbolic formula at this pointcentral difference
w-18.0-18.0-18.00000000
x-12.0-12.0-12.00000000
b-6.0-6.0-6.00000000
y6.06.06.00000000
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "residual = w*x+b-y\n", "symbolic = {\n", " \"w\": 2*residual*x,\n", " \"x\": 2*residual*w,\n", " \"b\": 2*residual,\n", " \"y\": -2*residual,\n", "}\n", "\n", "def loss_value(w_, x_, b_, y_):\n", " return (w_*x_+b_-y_)**2\n", "\n", "eps = 1e-5\n", "base = {\"w\": w, \"x\": x, \"b\": b, \"y\": y}\n", "finite = {}\n", "order = [\"w\", \"x\", \"b\", \"y\"]\n", "for name in order:\n", " plus = base.copy(); minus = base.copy()\n", " plus[name] += eps; minus[name] -= eps\n", " finite[name] = (\n", " loss_value(plus[\"w\"], plus[\"x\"], plus[\"b\"], plus[\"y\"])\n", " - loss_value(minus[\"w\"], minus[\"x\"], minus[\"b\"], minus[\"y\"])\n", " ) / (2*eps)\n", "\n", "manual = {\"w\": gw, \"x\": gx, \"b\": gb, \"y\": gy}\n", "for name in order:\n", " assert symbolic[name] == manual[name]\n", " assert math.isclose(finite[name], manual[name], rel_tol=1e-9, abs_tol=1e-8)\n", "\n", "show_table(\n", " [\"q\", \"manual reverse\", \"symbolic formula at this point\", \"central difference\"],\n", " [[name, manual[name], symbolic[name], f\"{finite[name]:.8f}\"] for name in order],\n", " \"Three routes agree\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "b37c1a8b", "metadata": {}, "source": [ "## Checkpoint 6 · One gradient step lowers this loss\n", "\n", "Backprop computes gradients; an optimizer later uses them. For this tiny\n", "example, one gradient-descent step with $\\eta=0.01$ is easy to audit.\n" ] }, { "cell_type": "code", "execution_count": 6, "id": "5fad3c1d", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.024503Z", "iopub.status.busy": "2026-08-20T12:58:22.024440Z", "iopub.status.idle": "2026-08-20T12:58:22.026767Z", "shell.execute_reply": "2026-08-20T12:58:22.026413Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "w: 2.00 -> 2.18\n", "b: 1.00 -> 1.06\n", "loss: 9.00 -> 5.76\n" ] } ], "source": [ "eta = 0.01\n", "w_new = w - eta*gw\n", "b_new = b - eta*gb\n", "prediction_new = w_new*x+b_new\n", "loss_new = (prediction_new-y)**2\n", "\n", "assert math.isclose(w_new, 2.18)\n", "assert math.isclose(b_new, 1.06)\n", "assert math.isclose(prediction_new, 7.6)\n", "assert math.isclose(loss_new, 5.76)\n", "assert loss_new < L\n", "\n", "print(f\"w: {w:.2f} -> {w_new:.2f}\")\n", "print(f\"b: {b:.2f} -> {b_new:.2f}\")\n", "print(f\"loss: {L:.2f} -> {loss_new:.2f}\")\n" ] }, { "cell_type": "markdown", "id": "752a0279", "metadata": {}, "source": [ "## Checkpoint 7 · PyTorch records and replays the same scalar graph\n", "\n", "`retain_grad()` lets us inspect intermediate non-leaf gradients for teaching.\n", "Ordinary training normally keeps only parameter gradients.\n" ] }, { "cell_type": "code", "execution_count": 7, "id": "e0839c11", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.027898Z", "iopub.status.busy": "2026-08-20T12:58:22.027834Z", "iopub.status.idle": "2026-08-20T12:58:22.275831Z", "shell.execute_reply": "2026-08-20T12:58:22.275352Z" } }, "outputs": [ { "data": { "text/html": [ "
PyTorch reproduces the complete ledger
nodestored value.grad after backward()
L9.01.0
e-3.0-6.0
a7.0-6.0
m6.0-6.0
w2.0-18.0
x3.0-12.0
b1.0-6.0
y10.06.0
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "x_t = torch.tensor(3.0, requires_grad=True)\n", "w_t = torch.tensor(2.0, requires_grad=True)\n", "b_t = torch.tensor(1.0, requires_grad=True)\n", "y_t = torch.tensor(10.0, requires_grad=True)\n", "\n", "m_t = w_t*x_t\n", "a_t = m_t+b_t\n", "e_t = a_t-y_t\n", "L_t = e_t**2\n", "for node in (m_t, a_t, e_t, L_t):\n", " node.retain_grad()\n", "L_t.backward()\n", "\n", "torch_expected = {\n", " \"L\": (L_t, L_t.grad, 9.0, 1.0),\n", " \"e\": (e_t, e_t.grad, -3.0, -6.0),\n", " \"a\": (a_t, a_t.grad, 7.0, -6.0),\n", " \"m\": (m_t, m_t.grad, 6.0, -6.0),\n", " \"w\": (w_t, w_t.grad, 2.0, -18.0),\n", " \"x\": (x_t, x_t.grad, 3.0, -12.0),\n", " \"b\": (b_t, b_t.grad, 1.0, -6.0),\n", " \"y\": (y_t, y_t.grad, 10.0, 6.0),\n", "}\n", "for _, (node, grad, value_expected, grad_expected) in torch_expected.items():\n", " torch.testing.assert_close(node.detach(), torch.tensor(value_expected))\n", " torch.testing.assert_close(grad, torch.tensor(grad_expected))\n", "\n", "show_table(\n", " [\"node\", \"stored value\", \".grad after backward()\"],\n", " [[name, float(node.detach()), float(grad)]\n", " for name, (node, grad, _, _) in torch_expected.items()],\n", " \"PyTorch reproduces the complete ledger\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "d652b57f", "metadata": {}, "source": [ "## Checkpoint 8 · Gradients accumulate until we clear them\n", "\n", "Each `backward()` adds into a leaf's `.grad`. A training loop therefore starts\n", "each intended update window with `zero_grad()` (or sets gradients to `None`).\n" ] }, { "cell_type": "code", "execution_count": 8, "id": "b31c4f7a", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.277076Z", "iopub.status.busy": "2026-08-20T12:58:22.276975Z", "iopub.status.idle": "2026-08-20T12:58:22.280482Z", "shell.execute_reply": "2026-08-20T12:58:22.280163Z" } }, "outputs": [ { "data": { "text/html": [ "
Accumulation is addition, not replacement
backward callw.grad
1-18.0
2-36.0
3-54.0
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "after zeroing: w.grad = 0.0\n" ] } ], "source": [ "w_acc = torch.tensor(2.0, requires_grad=True)\n", "history = []\n", "for backward_index in range(1, 4):\n", " loss_i = (w_acc*3.0+1.0-10.0)**2 # fresh graph, same leaf w_acc\n", " loss_i.backward()\n", " history.append(float(w_acc.grad))\n", "\n", "assert history == [-18.0, -36.0, -54.0]\n", "w_acc.grad.zero_()\n", "assert float(w_acc.grad) == 0.0\n", "\n", "show_table(\n", " [\"backward call\", \"w.grad\"],\n", " [[i, value] for i, value in enumerate(history, 1)],\n", " \"Accumulation is addition, not replacement\",\n", ")\n", "print(\"after zeroing: w.grad =\", float(w_acc.grad))\n" ] }, { "cell_type": "markdown", "id": "e75b214b", "metadata": {}, "source": [ "## Checkpoint 9 · At a branch, gradient contributions add\n", "\n", "The value $x=4$ is used along two paths:\n", "\n", "$$u=x^2,\\qquad v=3x,\\qquad L=u+v.$$\n", "\n", "The two returned contributions are $2x=8$ and $3$, so autograd must add them\n", "at the shared leaf: $x.\\text{grad}=8+3=11$.\n" ] }, { "cell_type": "code", "execution_count": 9, "id": "5f9adb04", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.281548Z", "iopub.status.busy": "2026-08-20T12:58:22.281488Z", "iopub.status.idle": "2026-08-20T12:58:22.284474Z", "shell.execute_reply": "2026-08-20T12:58:22.284212Z" } }, "outputs": [ { "data": { "text/html": [ "
Autograd accumulates at a shared value
pathlocal returncontribution to x.grad
u=x^21·2x8.0
v=3x1·33.0
shared xadd both paths11.0
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "x_branch = torch.tensor(4.0, requires_grad=True)\n", "u_branch = x_branch**2\n", "v_branch = 3*x_branch\n", "L_branch = u_branch+v_branch\n", "L_branch.backward()\n", "\n", "assert tuple(float(t.detach()) for t in (u_branch, v_branch, L_branch)) == (16.0, 12.0, 28.0)\n", "torch.testing.assert_close(x_branch.grad, torch.tensor(11.0))\n", "show_table(\n", " [\"path\", \"local return\", \"contribution to x.grad\"],\n", " [\n", " [\"u=x^2\", \"1·2x\", 8.0],\n", " [\"v=3x\", \"1·3\", 3.0],\n", " [\"shared x\", \"add both paths\", float(x_branch.grad)],\n", " ],\n", " \"Autograd accumulates at a shared value\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "550f13cd", "metadata": {}, "source": [ "## Checkpoint 10 · One vector neuron is still one affine operation\n", "\n", "For $z=w^Tx+b$ and $L=z^2$, the arriving scalar is $g_z=2z$. The affine\n", "operation returns $g_w=g_zx$, $g_x=g_zw$, and $g_b=g_z$.\n" ] }, { "cell_type": "code", "execution_count": 10, "id": "ba1105d6", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.285761Z", "iopub.status.busy": "2026-08-20T12:58:22.285701Z", "iopub.status.idle": "2026-08-20T12:58:22.289182Z", "shell.execute_reply": "2026-08-20T12:58:22.288885Z" } }, "outputs": [ { "data": { "text/html": [ "
Vector dot-product example from the slides
quantityvaluegradient
w[1.0, 3.0][-4.0, 2.0]
x[2.0, -1.0][-2.0, -6.0]
b0.0-2.0
z-1.0g_z = 2z = -2
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "x_vec = torch.tensor([2.0, -1.0], requires_grad=True)\n", "w_vec = torch.tensor([1.0, 3.0], requires_grad=True)\n", "b_vec = torch.tensor(0.0, requires_grad=True)\n", "z_vec = w_vec@x_vec+b_vec\n", "L_vec = z_vec**2\n", "L_vec.backward()\n", "\n", "torch.testing.assert_close(z_vec, torch.tensor(-1.0))\n", "torch.testing.assert_close(L_vec, torch.tensor(1.0))\n", "torch.testing.assert_close(w_vec.grad, torch.tensor([-4.0, 2.0]))\n", "torch.testing.assert_close(x_vec.grad, torch.tensor([-2.0, -6.0]))\n", "torch.testing.assert_close(b_vec.grad, torch.tensor(-2.0))\n", "show_table(\n", " [\"quantity\", \"value\", \"gradient\"],\n", " [\n", " [\"w\", w_vec.detach().tolist(), w_vec.grad.tolist()],\n", " [\"x\", x_vec.detach().tolist(), x_vec.grad.tolist()],\n", " [\"b\", float(b_vec.detach()), float(b_vec.grad)],\n", " [\"z\", float(z_vec.detach()), \"g_z = 2z = -2\"],\n", " ],\n", " \"Vector dot-product example from the slides\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "df2608ef", "metadata": {}, "source": [ "## Checkpoint 11 · Dense backward is the scalar affine rule repeated by row\n", "\n", "For column vectors, $z=Wx+b$. A later graph sends\n", "$g_z=(4,-2)^T$. Coordinate calculus gives\n", "\n", "$$g_x=W^Tg_z,\\qquad g_W=g_zx^T,\\qquad g_b=g_z.$$\n", "\n", "The transpose is not a memorized decoration: the same input feeds both output\n", "rows, so their returned contributions add at $x$.\n" ] }, { "cell_type": "code", "execution_count": 11, "id": "342e9b00", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.290191Z", "iopub.status.busy": "2026-08-20T12:58:22.290130Z", "iopub.status.idle": "2026-08-20T12:58:22.294542Z", "shell.execute_reply": "2026-08-20T12:58:22.294090Z" } }, "outputs": [ { "data": { "text/html": [ "
Dense VJP: derive by coordinate, then verify shapes
recipientrow-level returnstacked resultshape
x[[4.0, 12.0], [4.0, -2.0]][8.0, 10.0](2,)
W[[8.0, -4.0], [-4.0, 2.0]][[8.0, -4.0], [-4.0, 2.0]](2,2)
b[4.0, -2.0][4.0, -2.0](2,)
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "x_dense = torch.tensor([2.0, -1.0])\n", "W_dense = torch.tensor([[1.0, 3.0], [-2.0, 1.0]])\n", "b_dense = torch.tensor([0.0, 1.0])\n", "g_z = torch.tensor([4.0, -2.0])\n", "z_dense = W_dense@x_dense+b_dense\n", "\n", "gx_rows = g_z[:, None]*W_dense\n", "gW_rows = g_z[:, None]*x_dense[None, :]\n", "g_x_dense = gx_rows.sum(dim=0)\n", "g_W_dense = gW_rows\n", "g_b_dense = g_z\n", "\n", "torch.testing.assert_close(z_dense, torch.tensor([-1.0, -4.0]))\n", "torch.testing.assert_close(gx_rows, torch.tensor([[4.0, 12.0], [4.0, -2.0]]))\n", "torch.testing.assert_close(g_x_dense, torch.tensor([8.0, 10.0]))\n", "torch.testing.assert_close(g_W_dense, torch.tensor([[8.0, -4.0], [-4.0, 2.0]]))\n", "torch.testing.assert_close(g_b_dense, torch.tensor([4.0, -2.0]))\n", "\n", "# Verify the same vector-Jacobian product with autograd.\n", "x_leaf = x_dense.clone().requires_grad_()\n", "W_leaf = W_dense.clone().requires_grad_()\n", "b_leaf = b_dense.clone().requires_grad_()\n", "z_leaf = W_leaf@x_leaf+b_leaf\n", "z_leaf.backward(g_z)\n", "torch.testing.assert_close(x_leaf.grad, g_x_dense)\n", "torch.testing.assert_close(W_leaf.grad, g_W_dense)\n", "torch.testing.assert_close(b_leaf.grad, g_b_dense)\n", "\n", "show_table(\n", " [\"recipient\", \"row-level return\", \"stacked result\", \"shape\"],\n", " [\n", " [\"x\", gx_rows.tolist(), g_x_dense.tolist(), \"(2,)\"],\n", " [\"W\", gW_rows.tolist(), g_W_dense.tolist(), \"(2,2)\"],\n", " [\"b\", g_z.tolist(), g_b_dense.tolist(), \"(2,)\"],\n", " ],\n", " \"Dense VJP: derive by coordinate, then verify shapes\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "c49dfcea", "metadata": {}, "source": [ "## Checkpoint 12 · A batch adds an example axis, not new parameters\n", "\n", "PyTorch stores examples as rows $X\\in\\mathbb R^{B\\times d}$, so\n", "$Z=XW^T+b$. With half-squared error per example,\n", "\n", "$$\\ell_n=\\tfrac12\\|z^{(n)}-y^{(n)}\\|^2,\\qquad\n", "L=\\tfrac1B\\sum_n\\ell_n,$$\n", "\n", "the upstream batch gradient is $G_Z=(Z-Y)/B$.\n" ] }, { "cell_type": "code", "execution_count": 12, "id": "51f25a18", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.295565Z", "iopub.status.busy": "2026-08-20T12:58:22.295500Z", "iopub.status.idle": "2026-08-20T12:58:22.300275Z", "shell.execute_reply": "2026-08-20T12:58:22.299964Z" } }, "outputs": [ { "data": { "text/html": [ "
Three examples reuse the same W and b
nx^(n)z^(n)y^(n)r^(n)ell_n
1[2.0, -1.0][-1.0, -4.0][-5.0, -2.0][4.0, -2.0]10.0
2[-1.0, 2.0][5.0, 5.0][7.0, 1.0][-2.0, 4.0]10.0
3[2.0, 2.0][8.0, -1.0][7.0, -2.0][1.0, 1.0]1.0
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "mean loss L = 7.0\n", "Z.grad should be R/B =\n", " tensor([[ 1.3333, -0.6667],\n", " [-0.6667, 1.3333],\n", " [ 0.3333, 0.3333]])\n" ] } ], "source": [ "X = torch.tensor([[2.0, -1.0], [-1.0, 2.0], [2.0, 2.0]])\n", "Y = torch.tensor([[-5.0, -2.0], [7.0, 1.0], [7.0, -2.0]])\n", "B = X.shape[0]\n", "Z = X@W_dense.T+b_dense\n", "R = Z-Y\n", "loss_each = 0.5*(R**2).sum(dim=1)\n", "loss_mean = loss_each.mean()\n", "G_Z = R/B\n", "\n", "torch.testing.assert_close(Z, torch.tensor([[-1.0,-4.0],[5.0,5.0],[8.0,-1.0]]))\n", "torch.testing.assert_close(R, torch.tensor([[4.0,-2.0],[-2.0,4.0],[1.0,1.0]]))\n", "torch.testing.assert_close(loss_each, torch.tensor([10.0,10.0,1.0]))\n", "torch.testing.assert_close(loss_mean, torch.tensor(7.0))\n", "torch.testing.assert_close(G_Z, R/3)\n", "\n", "show_table(\n", " [\"n\", \"x^(n)\", \"z^(n)\", \"y^(n)\", \"r^(n)\", \"ell_n\"],\n", " [[n+1, X[n].tolist(), Z[n].tolist(), Y[n].tolist(), R[n].tolist(), float(loss_each[n])]\n", " for n in range(B)],\n", " \"Three examples reuse the same W and b\",\n", ")\n", "print(\"mean loss L =\", float(loss_mean))\n", "print(\"Z.grad should be R/B =\\n\", G_Z)\n" ] }, { "cell_type": "markdown", "id": "6a65529b", "metadata": {}, "source": [ "## Checkpoint 13 · Shared parameter paths add; the mean then scales\n", "\n", "Each example proposes one outer product $r^{(n)}(x^{(n)})^T$ to the\n", "same $W$. The batch gradient adds those proposals and divides by $B$.\n", "The three input rows are distinct, so their gradients remain separate.\n" ] }, { "cell_type": "code", "execution_count": 13, "id": "869a9e24", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.301370Z", "iopub.status.busy": "2026-08-20T12:58:22.301294Z", "iopub.status.idle": "2026-08-20T12:58:22.306216Z", "shell.execute_reply": "2026-08-20T12:58:22.305783Z" } }, "outputs": [ { "data": { "text/html": [ "
Manual batch arithmetic equals PyTorch autograd
quantitymanual batch ruleautograd resultshape
Z.grad[[1.3333333333333333, -0.6666666666666666], [-0.6666666666666666, 1.3333333333333333], [0.3333333333333333, 0.3333333333333333]][[1.3333333333333333, -0.6666666666666666], [-0.6666666666666666, 1.3333333333333333], [0.3333333333333333, 0.3333333333333333]](3, 2)
W.grad[[4.0, -2.0], [-2.0, 4.0]][[3.9999999999999996, -2.0], [-2.0, 3.9999999999999996]](2, 2)
b.grad[1.0, 1.0][1.0, 1.0](2,)
X.grad[[2.6666666666666665, 3.3333333333333335], [-3.333333333333333, -0.6666666666666667], [-0.3333333333333333, 1.3333333333333333]][[2.6666666666666665, 3.3333333333333335], [-3.333333333333333, -0.6666666666666667], [-0.3333333333333333, 1.3333333333333333]](3, 2)
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "gW_each_unscaled = torch.stack([torch.outer(R[n], X[n]) for n in range(B)])\n", "gW_batch = gW_each_unscaled.mean(dim=0)\n", "gb_batch = R.mean(dim=0)\n", "gX_batch = G_Z@W_dense\n", "\n", "expected_gW_each = torch.tensor([\n", " [[8.0,-4.0],[-4.0,2.0]],\n", " [[2.0,-4.0],[-4.0,8.0]],\n", " [[2.0,2.0],[2.0,2.0]],\n", "])\n", "torch.testing.assert_close(gW_each_unscaled, expected_gW_each)\n", "torch.testing.assert_close(gW_batch, torch.tensor([[4.0,-2.0],[-2.0,4.0]]))\n", "torch.testing.assert_close(gb_batch, torch.tensor([1.0,1.0]))\n", "torch.testing.assert_close(gX_batch, torch.tensor([\n", " [8/3,10/3],[-10/3,-2/3],[-1/3,4/3]\n", "]))\n", "\n", "# Full-batch autograd verification, including the non-leaf Z gradient.\n", "X_full = X.clone().requires_grad_()\n", "W_full = W_dense.clone().requires_grad_()\n", "b_full = b_dense.clone().requires_grad_()\n", "Z_full = X_full@W_full.T+b_full\n", "Z_full.retain_grad()\n", "L_full = 0.5*((Z_full-Y)**2).sum(dim=1).mean()\n", "L_full.backward()\n", "\n", "torch.testing.assert_close(L_full, loss_mean)\n", "torch.testing.assert_close(Z_full.grad, G_Z)\n", "torch.testing.assert_close(W_full.grad, gW_batch)\n", "torch.testing.assert_close(b_full.grad, gb_batch)\n", "torch.testing.assert_close(X_full.grad, gX_batch)\n", "\n", "show_table(\n", " [\"quantity\", \"manual batch rule\", \"autograd result\", \"shape\"],\n", " [\n", " [\"Z.grad\", G_Z.tolist(), Z_full.grad.tolist(), tuple(Z_full.grad.shape)],\n", " [\"W.grad\", gW_batch.tolist(), W_full.grad.tolist(), tuple(W_full.grad.shape)],\n", " [\"b.grad\", gb_batch.tolist(), b_full.grad.tolist(), tuple(b_full.grad.shape)],\n", " [\"X.grad\", gX_batch.tolist(), X_full.grad.tolist(), tuple(X_full.grad.shape)],\n", " ],\n", " \"Manual batch arithmetic equals PyTorch autograd\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "08365504", "metadata": {}, "source": [ "## Checkpoint 14 · Microbatches reproduce the same mean gradient\n", "\n", "Clear once, backpropagate each $\\ell_n/B$, and update once. This equals one\n", "full-batch backward when parameters stay fixed and the reduction is scaled\n", "consistently. It need not remain equivalent with batch-coupled operations such\n", "as BatchNorm.\n" ] }, { "cell_type": "code", "execution_count": 14, "id": "ed4b36fa", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.307189Z", "iopub.status.busy": "2026-08-20T12:58:22.307122Z", "iopub.status.idle": "2026-08-20T12:58:22.310852Z", "shell.execute_reply": "2026-08-20T12:58:22.310497Z" } }, "outputs": [ { "data": { "text/html": [ "
The running sum reaches the full mean-batch gradient
after microbatchaccumulated W.grad
1[[2.6666666666666665, -1.3333333333333333], [-1.3333333333333333, 0.6666666666666666]]
2[[3.333333333333333, -2.6666666666666665], [-2.6666666666666665, 3.333333333333333]]
3[[3.9999999999999996, -2.0], [-2.0, 3.9999999999999996]]
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "full-batch W.grad =\n", " tensor([[ 4.0000, -2.0000],\n", " [-2.0000, 4.0000]])\n" ] } ], "source": [ "W_micro = W_dense.clone().requires_grad_()\n", "b_micro = b_dense.clone().requires_grad_()\n", "snapshots = []\n", "\n", "for n in range(B):\n", " z_n = W_micro@X[n]+b_micro\n", " scaled_loss_n = 0.5*((z_n-Y[n])**2).sum()/B\n", " scaled_loss_n.backward()\n", " snapshots.append(W_micro.grad.detach().clone())\n", "\n", "torch.testing.assert_close(W_micro.grad, W_full.grad)\n", "torch.testing.assert_close(b_micro.grad, b_full.grad)\n", "\n", "show_table(\n", " [\"after microbatch\", \"accumulated W.grad\"],\n", " [[n+1, snapshots[n].tolist()] for n in range(B)],\n", " \"The running sum reaches the full mean-batch gradient\",\n", ")\n", "print(\"full-batch W.grad =\\n\", W_full.grad)\n" ] }, { "cell_type": "markdown", "id": "bc94a9ca", "metadata": {}, "source": [ "## Checkpoint 15 · A concrete training step uses the batch gradient\n", "\n", "Now the familiar loop is executable: clear, forward, reduce to one scalar,\n", "backward, and update. With learning rate $0.01$, the same fixed batch moves\n", "from loss $7$ to $6.5865$.\n" ] }, { "cell_type": "code", "execution_count": 15, "id": "be86ca89", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.311950Z", "iopub.status.busy": "2026-08-20T12:58:22.311871Z", "iopub.status.idle": "2026-08-20T12:58:22.475840Z", "shell.execute_reply": "2026-08-20T12:58:22.475428Z" } }, "outputs": [ { "data": { "text/html": [ "
clear → forward → scalar loss → backward → update
quantitybeforeafter one SGD step
W[[1.0, 3.0], [-2.0, 1.0]][[0.96, 3.02], [-1.98, 0.96]]
b[0.0, 1.0][-0.01, 0.99]
mean loss7.06.586499999999998
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "W_train = W_dense.clone().requires_grad_()\n", "b_train = b_dense.clone().requires_grad_()\n", "optimizer = torch.optim.SGD([W_train, b_train], lr=0.01)\n", "\n", "optimizer.zero_grad()\n", "Z_train = X@W_train.T+b_train\n", "loss_train = 0.5*((Z_train-Y)**2).sum(dim=1).mean()\n", "loss_train.backward()\n", "optimizer.step()\n", "\n", "with torch.no_grad():\n", " loss_after = 0.5*((X@W_train.T+b_train-Y)**2).sum(dim=1).mean()\n", "\n", "torch.testing.assert_close(W_train, torch.tensor([[0.96,3.02],[-1.98,0.96]]))\n", "torch.testing.assert_close(b_train, torch.tensor([-0.01,0.99]))\n", "torch.testing.assert_close(loss_after, torch.tensor(6.5865))\n", "assert loss_after < loss_train\n", "\n", "show_table(\n", " [\"quantity\", \"before\", \"after one SGD step\"],\n", " [\n", " [\"W\", W_dense.tolist(), W_train.detach().tolist()],\n", " [\"b\", b_dense.tolist(), b_train.detach().tolist()],\n", " [\"mean loss\", float(loss_train.detach()), float(loss_after)],\n", " ],\n", " \"clear → forward → scalar loss → backward → update\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "ff9a507d", "metadata": {}, "source": [ "## Checkpoint 16 · A tiny MLP is the same three-block pattern\n", "\n", "The slides finish with dense $\\rightarrow$ ReLU $\\rightarrow$ dense. We set\n", "deterministic weights so every forward value is auditable, then let PyTorch\n", "return one gradient for every parameter. The inactive hidden unit returns zero.\n" ] }, { "cell_type": "code", "execution_count": 16, "id": "7cc649f1", "metadata": { "execution": { "iopub.execute_input": "2026-08-20T12:58:22.477093Z", "iopub.status.busy": "2026-08-20T12:58:22.477001Z", "iopub.status.idle": "2026-08-20T12:58:22.483293Z", "shell.execute_reply": "2026-08-20T12:58:22.482902Z" } }, "outputs": [ { "data": { "text/html": [ "
Dense → ReLU → dense
stagevalue
dense 1 preactivation[2.5, -0.5, 3.0]
ReLU[2.5, 0.0, 3.0]
dense 2 prediction[4.0, 1.75]
squared-error loss12.0625
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
Autograd follows the recorded blocks in reverse
parametershapegradient shape
layer1.weight(3, 2)(3, 2)
layer1.bias(3,)(3,)
layer2.weight(2, 3)(2, 3)
layer2.bias(2,)(2,)
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from torch import nn\n", "\n", "layer1 = nn.Linear(2, 3)\n", "layer2 = nn.Linear(3, 2)\n", "with torch.no_grad():\n", " layer1.weight.copy_(torch.tensor([[1.0,-0.5],[0.0,0.0],[1.0,-1.0]]))\n", " layer1.bias.copy_(torch.tensor([0.0,-0.5,0.0]))\n", " layer2.weight.copy_(torch.tensor([[0.4,0.3,1.0],[0.7,-0.2,0.0]]))\n", " layer2.bias.zero_()\n", "\n", "x_mlp = torch.tensor([2.0,-1.0])\n", "y_mlp = torch.tensor([1.0,0.0])\n", "preactivation = layer1(x_mlp)\n", "preactivation.retain_grad()\n", "hidden = torch.relu(preactivation)\n", "prediction = layer2(hidden)\n", "loss_mlp = ((prediction-y_mlp)**2).sum()\n", "loss_mlp.backward()\n", "\n", "torch.testing.assert_close(preactivation.detach(), torch.tensor([2.5,-0.5,3.0]))\n", "torch.testing.assert_close(hidden.detach(), torch.tensor([2.5,0.0,3.0]))\n", "torch.testing.assert_close(prediction.detach(), torch.tensor([4.0,1.75]))\n", "torch.testing.assert_close(loss_mlp.detach(), torch.tensor(12.0625))\n", "torch.testing.assert_close(preactivation.grad[1], torch.tensor(0.0))\n", "torch.testing.assert_close(layer1.weight.grad[1], torch.zeros(2))\n", "torch.testing.assert_close(layer1.bias.grad[1], torch.tensor(0.0))\n", "\n", "for name, parameter in [\n", " (\"layer1.weight\", layer1.weight), (\"layer1.bias\", layer1.bias),\n", " (\"layer2.weight\", layer2.weight), (\"layer2.bias\", layer2.bias),\n", "]:\n", " assert parameter.grad is not None\n", " assert parameter.grad.shape == parameter.shape\n", " assert torch.isfinite(parameter.grad).all()\n", "\n", "show_table(\n", " [\"stage\", \"value\"],\n", " [\n", " [\"dense 1 preactivation\", preactivation.detach().tolist()],\n", " [\"ReLU\", hidden.detach().tolist()],\n", " [\"dense 2 prediction\", prediction.detach().tolist()],\n", " [\"squared-error loss\", float(loss_mlp.detach())],\n", " ],\n", " \"Dense → ReLU → dense\",\n", ")\n", "show_table(\n", " [\"parameter\", \"shape\", \"gradient shape\"],\n", " [[name, tuple(parameter.shape), tuple(parameter.grad.shape)] for name, parameter in [\n", " (\"layer1.weight\", layer1.weight), (\"layer1.bias\", layer1.bias),\n", " (\"layer2.weight\", layer2.weight), (\"layer2.bias\", layer2.bias),\n", " ]],\n", " \"Autograd follows the recorded blocks in reverse\",\n", ")\n" ] }, { "cell_type": "markdown", "id": "886dfa81", "metadata": {}, "source": [ "## Takeaway\n", "\n", "- Reverse-mode backprop repeatedly applies **upstream × local** and adds at\n", " shared values.\n", "- Symbolic differentiation derives a formula; finite differences approximate\n", " a point derivative; autograd records the executed graph and runs the reverse\n", " sweep.\n", "- A dense layer repeats the scalar affine rule row by row.\n", "- A batch adds an example axis. Shared parameter gradients add across examples;\n", " the loss reduction determines whether the final result is a sum or a mean.\n", "- Multiple backward calls accumulate into `.grad`; clear once per intended\n", " update window and step once after all correctly scaled contributions arrive.\n", "- A deeper MLP is not a new differentiation idea: each dense or activation\n", " block receives one arriving gradient and returns the next one.\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.11" } }, "nbformat": 4, "nbformat_minor": 5 }