{ "cells": [ { "cell_type": "markdown", "id": "be5ebaac", "metadata": {}, "source": [ "\n", "\n", "

Lecture 14: PyTorch, Initialization, and Normalization (CS 189, Fall 2026)

\n", "\n", "This notebook accompanies Lecture 14 and Sections 6.2.4, 7.2.5, and 7.4 of Bishop and Bishop, *Deep Learning: Foundations and Concepts*. It is organized in three parts. The first introduces PyTorch, the library we will use for the rest of the course to build and train neural networks. The second asks how the weights of a deep network should be initialized, and shows that the answer determines whether signals and gradients survive the passage through many layers. The third turns to normalization, implementing batch normalization, layer normalization, and RMSNorm by hand and checking each against its PyTorch counterpart.\n", "\n", "| Notebook section | Slide |\n", "|---|---|\n", "| Part 1: PyTorch (Lecture Demo 1) | \"Demo\" after the PyTorch slides |\n", "| Tensors | \"torch.tensor\" |\n", "| Automatic Differentiation | \"Computing Gradients using Autograd\" |\n", "| Datasets and Data Loaders | \"Big Concepts in PyTorch\" |\n", "| Models | \"torch.nn.functional and torch.nn\", \"torch.nn.Module\" |\n", "| Loss Functions | \"torch.nn.functional and torch.nn\" |\n", "| Training | \"Training Loop with an Optimizer\" |\n", "| Part 2: Parameter Initialization (Lecture Demo 2) | \"Parameter Initialization\" through \"Weight Initialization in PyTorch\" |\n", "| Part 3: Normalization (Lecture Demo 3) | \"Data Normalization\" through \"Batch-Norm and Layer-Norm in PyTorch\" |\n", "\n", "Every cell runs in a few seconds on a laptop CPU.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e2984942", "metadata": {}, "outputs": [], "source": [ "import copy\n", "import math\n", "\n", "import numpy as np\n", "import pandas as pd\n", "import plotly.express as px\n", "import plotly.graph_objects as go\n", "from plotly.subplots import make_subplots\n", "from sklearn.datasets import make_moons, load_breast_cancer\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torch.utils.data import TensorDataset, DataLoader, random_split\n", "\n", "colors = px.colors.qualitative.Plotly\n", "px.defaults.width = 800" ] }, { "cell_type": "markdown", "id": "970cfc16", "metadata": {}, "source": [ "\n", "\n", "

Plotting Code

\n", "\n", "The functions below produce the visualizations used throughout the notebook. They are not part of the course material, but they must be run before the rest of the notebook.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "cd6e2146", "metadata": {}, "outputs": [], "source": [ "# Berkeley blue for class 0 and California gold for class 1, as in Lecture 12.\n", "CLASS_COLORS = {0: \"#003262\", 1: \"#FDB515\"}\n", "PROB_COLORSCALE = [[0, \"#B9CCE2\"], [0.5, \"#FFFFFF\"], [1, \"#FEE3A0\"]]\n", "\n", "\n", "def plot_data(X, t, title=None):\n", " \"\"\"Scatter plot of a two-dimensional binary classification data set.\"\"\"\n", " X, t = np.asarray(X), np.asarray(t).ravel()\n", " fig = go.Figure()\n", " for k, color in CLASS_COLORS.items():\n", " fig.add_scatter(x=X[t == k, 0], y=X[t == k, 1], mode=\"markers\", name=f\"class {k}\",\n", " marker=dict(color=color, size=7, line=dict(width=0.5, color=\"black\")))\n", " fig.update_layout(title=title, xaxis_title=\"x1\", yaxis_title=\"x2\", width=600, height=500)\n", " fig.update_yaxes(scaleanchor=\"x\", scaleratio=1)\n", " return fig\n", "\n", "\n", "def boundary_traces(model, X, t, n=200, pad=0.5, showlegend=True):\n", " \"\"\"Traces for the predicted probability p(t=1|x) = sigmoid(model(x)) on a grid, the\n", " decision boundary p = 0.5, and the data. The model is evaluated in eval mode.\"\"\"\n", " X, t = np.asarray(X), np.asarray(t).ravel()\n", " lo, hi = X.min(0) - pad, X.max(0) + pad\n", " g1, g2 = np.linspace(lo[0], hi[0], n), np.linspace(lo[1], hi[1], n)\n", " G1, G2 = np.meshgrid(g1, g2)\n", " grid = torch.tensor(np.c_[G1.ravel(), G2.ravel()], dtype=torch.float32)\n", " was_training = model.training\n", " model.eval()\n", " with torch.no_grad():\n", " P = torch.sigmoid(model(grid)).reshape(n, n).numpy()\n", " model.train(was_training)\n", " traces = [go.Contour(x=g1, y=g2, z=P, zmin=0, zmax=1, colorscale=PROB_COLORSCALE,\n", " contours=dict(start=0.5, end=0.5, size=1, coloring=\"heatmap\"),\n", " line=dict(width=2, color=\"black\"), showscale=False, hoverinfo=\"skip\")]\n", " for k, color in CLASS_COLORS.items():\n", " traces.append(go.Scatter(x=X[t == k, 0], y=X[t == k, 1], mode=\"markers\",\n", " name=f\"class {k}\", legendgroup=f\"class {k}\", showlegend=showlegend,\n", " marker=dict(color=color, size=6, line=dict(width=0.5, color=\"black\"))))\n", " return traces\n", "\n", "\n", "def plot_decision_boundary(model, X, t, title=None):\n", " \"\"\"Predicted probability, decision boundary, and data for one model.\"\"\"\n", " fig = go.Figure(boundary_traces(model, X, t))\n", " fig.update_layout(title=title, xaxis_title=\"x1\", yaxis_title=\"x2\", width=600, height=500)\n", " return fig\n", "\n", "\n", "def plot_decision_boundaries(models, titles, X, t):\n", " \"\"\"Decision boundaries of several models side by side.\"\"\"\n", " fig = make_subplots(rows=1, cols=len(models), subplot_titles=titles)\n", " for j, model in enumerate(models):\n", " fig.add_traces(boundary_traces(model, X, t, showlegend=(j == 0)), rows=1, cols=j + 1)\n", " fig.update_layout(width=330 * len(models) + 100, height=400)\n", " return fig\n", "\n", "\n", "def plot_curves(curves, title=None, xaxis_title=\"epoch\", yaxis_title=\"loss\", log_y=False, dash=None):\n", " \"\"\"Line plot of several curves given as a dict {name: list of values}.\"\"\"\n", " fig = go.Figure()\n", " for i, (name, ys) in enumerate(curves.items()):\n", " line = dict(color=colors[i % len(colors)], width=3)\n", " if dash and dash.get(name):\n", " line[\"dash\"] = dash[name]\n", " fig.add_scatter(x=np.arange(1, len(ys) + 1), y=ys, mode=\"lines\", name=name, line=line)\n", " fig.update_layout(title=title, xaxis_title=xaxis_title, yaxis_title=yaxis_title,\n", " width=800, height=450)\n", " if log_y:\n", " fig.update_yaxes(type=\"log\", exponentformat=\"power\")\n", " return fig" ] }, { "cell_type": "markdown", "id": "db09eab6", "metadata": {}, "source": [ "\n", "\n", "

(Lecture Demo 1) Part 1: PyTorch

\n", "\n", "*Lecture:* run from the \"Demo\" slide after the PyTorch slides.\n", "\n", "PyTorch is the library we will use to build and train neural networks for the rest of the course, and four of its components appear in essentially every program we write. Tensors are multidimensional arrays, much like NumPy arrays, with the important addition that they can live on a GPU or another accelerator. Automatic differentiation computes gradients of tensor expressions with the same reverse-mode algorithm that we implemented in the `Scalar` class of Lecture 13. Modules, found in `torch.nn`, are the building blocks of models: layers, activation functions, and loss functions, each of which keeps track of its own parameters. Finally, the optimizers in `torch.optim` implement the update rules of Lectures 10 and 11, and the `Dataset` and `DataLoader` classes supply the mini-batches on which they operate.\n" ] }, { "cell_type": "markdown", "id": "9bc99eb5", "metadata": {}, "source": [ "\n", "\n", "

Tensors

\n", "\n", "*Lecture:* slide \"torch.tensor\".\n", "\n", "The tensor is PyTorch's basic data structure. A tensor can be created from a (nested) Python list, from a NumPy array, or with a constructor such as `torch.zeros`, and like a NumPy array it has a shape and a data type, or `dtype`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3af78e32", "metadata": {}, "outputs": [], "source": [ "x = torch.tensor([[1.0, 2.0, 3.0],\n", " [4.0, 5.0, 6.0]])\n", "print(x)\n", "print(\"shape:\", x.shape, \" dtype:\", x.dtype)" ] }, { "cell_type": "markdown", "id": "39272a74", "metadata": {}, "source": [ "When a tensor is created from Python numbers, PyTorch chooses the data type for us. Floating-point numbers become `torch.float32`, the library's default floating-point type, while integers become `torch.int64`. This matters more than it might first appear, because many operations, including the mean computed below, are defined only for floating-point tensors. For this reason it is good practice to state the `dtype` explicitly whenever a tensor is created.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0f5024f2", "metadata": {}, "outputs": [], "source": [ "print(\"default dtype:\", torch.get_default_dtype())\n", "n = torch.tensor([1, 2, 3])\n", "print(\"from ints:\", n.dtype)\n", "try:\n", " n.mean()\n", "except RuntimeError as err:\n", " print(\"RuntimeError:\", err)\n", "print(\"explicit dtype:\", torch.tensor([1, 2, 3], dtype=torch.float32).mean())" ] }, { "cell_type": "markdown", "id": "bbac8d21", "metadata": {}, "source": [ "A related pitfall arises when converting from NumPy. NumPy uses `float64` by default, and `torch.tensor` preserves the data type of the array it is given. Passing such a tensor to a layer causes an error, as the parameters of a PyTorch model are stored as `float32`. The remedy is to request the type explicitly with `dtype=torch.float32`, or to convert afterwards with `.float()`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5482a0c3", "metadata": {}, "outputs": [], "source": [ "a_np = np.random.default_rng(189).standard_normal((4, 3))\n", "x64 = torch.tensor(a_np)\n", "print(\"from NumPy:\", x64.dtype)\n", "\n", "layer = nn.Linear(3, 2) # parameters are float32\n", "try:\n", " layer(x64)\n", "except RuntimeError as err:\n", " print(\"RuntimeError:\", err)\n", "print(layer(torch.tensor(a_np, dtype=torch.float32)).dtype)" ] }, { "cell_type": "markdown", "id": "de24b794", "metadata": {}, "source": [ "**Low-precision data types.** *Lecture:* slide \"torch.tensor\" (key differences). In contrast to most NumPy code, deep learning code pays close attention to numerical precision. Large models are typically trained and run in 16-bit floating point, which halves the memory required by every tensor and runs considerably faster on the matrix units of modern GPUs. Two 16-bit formats are in common use, and they divide their bits differently between the exponent, which determines the range of representable values, and the mantissa, which determines their precision:\n", "\n", "| `dtype` | bits (sign, exponent, mantissa) | largest value | relative precision $\\epsilon$ |\n", "|---|---|---|---|\n", "| `torch.float32` | 1, 8, 23 | $3.4 \\times 10^{38}$ | $1.2 \\times 10^{-7}$ |\n", "| `torch.float16` | 1, 5, 10 | $65504$ | $9.8 \\times 10^{-4}$ |\n", "| `torch.bfloat16` | 1, 8, 7 | $3.4 \\times 10^{38}$ | $7.8 \\times 10^{-3}$ |\n", "\n", "The `float16` format retains more precision but covers a much smaller range, so large values overflow to `inf` and small gradients underflow to zero. The `bfloat16` (\"brain floating point\") format makes the opposite choice, keeping the full range of `float32` at the cost of precision. The cell below reads these limits from `torch.finfo` and demonstrates each kind of failure.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f87aa903", "metadata": {}, "outputs": [], "source": [ "for dt in (torch.float32, torch.float16, torch.bfloat16):\n", " info = torch.finfo(dt)\n", " print(f\"{str(dt):15s} bytes/entry {torch.tensor(0.0, dtype=dt).element_size()} \"\n", " f\"max {info.max:.3g} eps {info.eps:.2g}\")\n", "\n", "big, small = torch.tensor(70000.0), torch.tensor(1e-8)\n", "print(\"70000 ->\", big.to(torch.float16).item(), \"(float16),\", big.to(torch.bfloat16).item(), \"(bfloat16)\")\n", "print(\"1e-8 ->\", small.to(torch.float16).item(), \"(float16),\", small.to(torch.bfloat16).item(), \"(bfloat16)\")\n", "one = torch.tensor(1.0, dtype=torch.bfloat16)\n", "print(\"1 + 0.001 in bfloat16:\", (one + 0.001).item())\n", "\n", "n_params = 7e9 # a 7-billion-parameter language model\n", "for dt, nbytes in ((torch.float32, 4), (torch.bfloat16, 2)):\n", " print(f\"7B parameters in {dt}: {n_params * nbytes / 1e9:.0f} GB\")" ] }, { "cell_type": "markdown", "id": "50494992", "metadata": {}, "source": [ "These trade-offs motivate the common practice of mixed-precision training. The forward and backward passes run in `bfloat16`, for example inside a `with torch.autocast(device_type, dtype=torch.bfloat16):` block, while the parameters and the optimizer state are kept in `float32`. The reason for keeping the weights in higher precision is the last example above: a gradient step $\\eta \\nabla E$ is often much smaller than the weight it updates, and in `bfloat16` such an update would simply be rounded away. Individual tensors are converted with `x.to(torch.bfloat16)`, and an entire model with `model.to(torch.bfloat16)`.\n" ] }, { "cell_type": "markdown", "id": "b6692022", "metadata": {}, "source": [ "Random tensors are drawn from PyTorch's global random number generator, which `torch.manual_seed` seeds, or from an explicit `torch.Generator`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "7d412c55", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "print(torch.zeros(2, 3))\n", "print(torch.ones(2, 3))\n", "print(torch.arange(6))\n", "print(torch.randn(2, 3)) # N(0, 1) entries\n", "print(torch.rand(2, 3, generator=torch.Generator().manual_seed(189))) # U(0, 1) entries" ] }, { "cell_type": "markdown", "id": "b8d23067", "metadata": {}, "source": [ "**Shapes and indexing.** The shape of a tensor can be changed without changing its entries. The method `reshape` arranges the entries, in row-major order, into a new shape, `.T` transposes a matrix, and `unsqueeze` inserts a new axis of length one. Indexing and slicing work as they do in NumPy, including indexing with boolean masks.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2dede39a", "metadata": {}, "outputs": [], "source": [ "x = torch.arange(6, dtype=torch.float32)\n", "X = x.reshape(2, 3)\n", "print(X)\n", "print(\"X.T shape:\", X.T.shape)\n", "print(\"first row:\", X[0])\n", "print(\"second column:\", X[:, 1])\n", "print(\"entries greater than 2:\", X[X > 2])" ] }, { "cell_type": "markdown", "id": "6301efe4", "metadata": {}, "source": [ "Unsqueeze can be a bit more confusing so let's step through it." ] }, { "cell_type": "code", "execution_count": null, "id": "5bc3dc52", "metadata": {}, "outputs": [], "source": [ "print(X, X.shape, \"\\n\")\n", "print(X.unsqueeze(0), X.unsqueeze(0).shape, \"\\n\")\n", "print(X.unsqueeze(1), X.unsqueeze(1).shape, \"\\n\")\n", "print(X.unsqueeze(2), X.unsqueeze(2).shape, \"\\n\")" ] }, { "cell_type": "markdown", "id": "fbba4355", "metadata": {}, "source": [ "**Broadcasting.** When an elementwise operation combines tensors of different shapes, PyTorch aligns the shapes starting from the last axis and repeats any axis of length one, or any missing axis, as needed. A useful consequence is that subtracting the vector of column means from an $N \\times D$ matrix centers every column at once. This is exactly the first step of the data normalization that we will carry out in Part 3.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8ac0aa9c", "metadata": {}, "outputs": [], "source": [ "print(X + torch.tensor([10.0, 20.0, 30.0])) # (2, 3) + (3,): added to every row\n", "print(X - X.mean(dim=0)) # center each column" ] }, { "cell_type": "markdown", "id": "bcf77d62", "metadata": {}, "source": [ "**Matrix multiplication and reductions.** The operator `@` performs matrix multiplication. By convention a batch of $N$ inputs is stored as the rows of an $N \\times D$ matrix $\\mathbf{X}$, so a layer whose weight matrix $\\mathbf{W}$ has shape $(D_\\text{out}, D_\\text{in})$ computes $\\mathbf{X}\\mathbf{W}^\\top$, producing one row of outputs for each input. Reductions such as `sum` and `mean` take an argument `dim` that names the axis to be reduced. Reducing with `dim=0` collapses the rows and leaves one result per column, while `dim=1` collapses the columns and leaves one result per row.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8db355d1", "metadata": {}, "outputs": [], "source": [ "W = torch.tensor([[1.0, 0.0, -1.0],\n", " [0.5, 0.5, 0.5]]) # shape (D_out, D_in) = (2, 3)\n", "print(\"X @ W.T =\\n\", X @ W.T, \"\\n\") # shape (N, D_out) = (2, 2)\n", "\n", "print(\"sum of all entries:\", X.sum(), \"\\n\")\n", "\n", "print(\"column means (dim=0):\", X.mean(dim=0), \"\\n\")\n", "print(\"row means (dim=1): \", X.mean(dim=1), \"\\n\")\n", "print(\"row means, keepdim: \", X.mean(dim=1, keepdim=True), \n", " X.mean(dim=1, keepdim=True).shape, \"\\n\")" ] }, { "cell_type": "markdown", "id": "8cffcc6c", "metadata": {}, "source": [ "**NumPy interoperability.** Moving between NumPy and PyTorch is straightforward, with one subtlety. The functions `torch.from_numpy` and `.numpy()` share memory with the original array, so modifying one also modifies the other, whereas `torch.tensor` always makes a copy. A tensor that holds a single value can be converted to an ordinary Python number with `.item()`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ba82fa2f", "metadata": {}, "outputs": [], "source": [ "a = np.arange(3.0)\n", "shared = torch.from_numpy(a) # shares memory with a\n", "copied = torch.tensor(a) # copies a\n", "a[0] = 100.0\n", "print(\"from_numpy:\", shared)\n", "print(\"tensor: \", copied)\n", "print(\"back to NumPy:\", copied.numpy(), \" one entry as a float:\", copied[1].item())" ] }, { "cell_type": "markdown", "id": "004f2f67", "metadata": {}, "source": [ "**Devices.** Every tensor resides on a device, either the CPU or an accelerator such as a GPU (or `mps` on Apple silicon), and `torch.accelerator` reports whether an accelerator is available. Because an operation requires all of its operands to be on the same device, tensors and models are moved explicitly with `.to(device)`, and a tensor must be returned to the CPU before it can be converted to a NumPy array. The models in this notebook are small enough that the cost of moving data to an accelerator would outweigh any gain in speed, so the rest of the notebook runs on the CPU.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "b2f73f02", "metadata": {}, "outputs": [], "source": [ "if torch.accelerator.is_available():\n", " print(\"accelerator available:\", torch.accelerator.current_accelerator())\n", "else:\n", " print(\"no accelerator available\")\n", "\n", "device = torch.device(\"cpu\") # this notebook uses the CPU\n", "x = torch.randn(3, 2).to(device)\n", "print(x.device, x.cpu().numpy().shape)\n", "\n", "print(\"if there is a device this will move the tensor to it\")\n", "if torch.accelerator.is_available():\n", " device = torch.accelerator.current_accelerator()\n", " x = x.to(device)\n", " print(x.device,x)\n" ] }, { "cell_type": "markdown", "id": "24914483", "metadata": {}, "source": [ "\n", "\n", "

Automatic Differentiation

\n", "\n", "*Lecture:* slide \"Computing Gradients using Autograd\".\n", "\n", "Automatic differentiation in PyTorch works on tensors in the same way that the `Scalar` class of Lecture 13 worked on individual numbers. Setting `requires_grad=True` on a tensor instructs PyTorch to record every operation in which the tensor takes part, so that a computation graph is built as the expression is evaluated. Each result stores, in its attribute `grad_fn`, the node of the graph that produced it, and this node knows how to propagate adjoints back to its inputs. Calling `.backward()` on a scalar result then runs reverse-mode automatic differentiation through the graph, leaving the derivative $\\partial f / \\partial x$ in `x.grad` for every leaf tensor `x` that requires gradients.\n", "\n", "We return to the running example of Lecture 13, $f(x_1, x_2) = x_1 x_2 + e^{x_1 x_2} - \\sin(x_2)$, whose gradient is\n", "$$\n", "\\frac{\\partial f}{\\partial x_1} = x_2 + x_2 e^{x_1 x_2}, \\qquad\n", "\\frac{\\partial f}{\\partial x_2} = x_1 + x_1 e^{x_1 x_2} - \\cos(x_2),\n", "$$\n", "and which takes the value $(16.78, 8.81)$ at the point $(1, 2)$.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3988debb", "metadata": {}, "outputs": [], "source": [ "x1 = torch.tensor(1.0, requires_grad=True)\n", "x2 = torch.tensor(2.0, requires_grad=True)" ] }, { "cell_type": "code", "execution_count": null, "id": "e2c94386", "metadata": {}, "outputs": [], "source": [ "f = x1 * x2 + torch.exp(x1 * x2) - torch.sin(x2)\n", "\n", "print(f)\n", "print(\"f.grad_fn:\", f.grad_fn, \" inputs:\", f.grad_fn.next_functions)" ] }, { "cell_type": "code", "execution_count": null, "id": "163f6576", "metadata": {}, "outputs": [], "source": [ "f.backward()\n", "print(\"autograd:\", x1.grad.item(), x2.grad.item())\n", "print(\"exact: \", 2 + 2 * math.exp(2), 1 + math.exp(2) - math.cos(2))" ] }, { "cell_type": "markdown", "id": "f62fcead", "metadata": {}, "source": [ "**Gradients accumulate.** A call to `.backward()` adds the new gradient to whatever is already stored in `.grad` rather than overwriting it. This is the behavior required when a node feeds into several others, as we saw in Lecture 13, but it means that the stored gradient must be reset to zero before a new one is computed. PyTorch also frees the computation graph once the backward pass is complete, so calling `.backward()` a second time on the same `f` is an error, and we must first build $f$ again.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "552c686a", "metadata": {}, "outputs": [], "source": [ "f = x1 * x2 + torch.exp(x1 * x2) - torch.sin(x2)\n", "f.backward()\n", "print(\"after a second backward:\", x1.grad.item(), x2.grad.item()) # twice the gradient\n", "\n", "x1.grad.zero_()\n", "x2.grad.zero_()\n", "f = x1 * x2 + torch.exp(x1 * x2) - torch.sin(x2)\n", "f.backward()\n", "print(\"after zeroing first: \", x1.grad.item(), x2.grad.item())" ] }, { "cell_type": "markdown", "id": "5ca26280", "metadata": {}, "source": [ "**Turning the graph off.** Recording the graph is wasteful when we only want to evaluate a model, and it is wrong when we update the parameters themselves, since the update is not part of the model. Inside a `with torch.no_grad():` block no graph is recorded. A single tensor can also be separated from the graph with `.detach()`, which returns a tensor holding the same values but with no record of how they were computed.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ba181588", "metadata": {}, "outputs": [], "source": [ "with torch.no_grad():\n", " y = x1 * x2\n", "print(\"under no_grad:\", y.requires_grad)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f232504f", "metadata": {}, "outputs": [], "source": [ "z = (x1 * x2).detach()\n", "print(\"detached: \", z.requires_grad, \" with graph:\", (x1 * x2).requires_grad)" ] }, { "cell_type": "markdown", "id": "8d81e344", "metadata": {}, "source": [ "\n", "\n", "

Datasets and Data Loaders

\n", "\n", "*Lecture:* slide \"Big Concepts in PyTorch\".\n", "\n", "We use the moons data set from Lecture 12. A `TensorDataset` pairs each input with its target, so that indexing the data set as `dataset[n]` returns the pair $(\\mathbf{x}_n, t_n)$. The targets are stored as `float32` values with shape $(N, 1)$, which matches the shape of the output of a network with a single output unit.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "9689cb0a", "metadata": {}, "outputs": [], "source": [ "# Make numpy data\n", "X_moons, t_moons = make_moons(400, noise=0.2, random_state=189)" ] }, { "cell_type": "code", "execution_count": null, "id": "67d4aab2", "metadata": {}, "outputs": [], "source": [ "dataset = TensorDataset(torch.tensor(X_moons, dtype=torch.float32),\n", " torch.tensor(t_moons, dtype=torch.float32).reshape(-1, 1))\n", "print(\"N =\", len(dataset), \" first item:\", dataset[0])" ] }, { "cell_type": "code", "execution_count": null, "id": "bc52a69f", "metadata": {}, "outputs": [], "source": [ "plot_data(X_moons, t_moons, title=\"Moons data (N = 400)\").show()" ] }, { "cell_type": "markdown", "id": "a935f78c", "metadata": {}, "source": [ "So that the model can be evaluated on data it was not trained on, `random_split` divides the data set into a training set and a validation set. Passing a seeded `torch.Generator` makes the split the same from one run to the next.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0fbd9239", "metadata": {}, "outputs": [], "source": [ "train_data, val_data = random_split(dataset, [0.8, 0.2],\n", " generator=torch.Generator().manual_seed(189))\n", "print(\"training points:\", len(train_data), \" validation points:\", len(val_data))\n", "x_val, t_val = val_data[:] # all validation points as two tensors\n", "print(x_val.shape, t_val.shape)" ] }, { "cell_type": "markdown", "id": "1f9760a8", "metadata": {}, "source": [ "Mini-batches are produced by a `DataLoader`. With `shuffle=True` the loader draws a new random permutation of the training data at the start of every epoch, so a single pass over the loader visits each training point exactly once, in $\\lceil N/B \\rceil$ mini-batches of size $B$. This is the sampling scheme of mini-batch stochastic gradient descent in Algorithm 7.2 of Bishop. When $B$ does not divide $N$, the final mini-batch is simply smaller than the others.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "96af13ef", "metadata": {}, "outputs": [], "source": [ "for B in (32, 50):\n", " loader = DataLoader(train_data, batch_size=B, shuffle=True,\n", " generator=torch.Generator().manual_seed(189))\n", " sizes = [len(x) for x, t in loader]\n", " print(f\"B = {B}: {len(sizes)} batches = ceil({len(train_data)}/{B}) = \"\n", " f\"{math.ceil(len(train_data) / B)}, sizes {sizes}\")\n", "\n", "x, t = next(iter(loader))\n", "print(\"one batch:\", x.shape, t.shape)" ] }, { "cell_type": "markdown", "id": "cd2ffff7", "metadata": {}, "source": [ "\n", "\n", "

Models

\n", "\n", "*Lecture:* slides \"torch.nn.functional and torch.nn\" and \"torch.nn.Module\".\n", "\n", "PyTorch offers most operations in two forms. The module `torch.nn.functional`, conventionally imported as `F`, contains plain functions without parameters, such as activation functions. The module `torch.nn` contains the same operations packaged as modules, together with layers that own trainable parameters. The cell below confirms that three ways of computing a ReLU give the same result.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5d681d35", "metadata": {}, "outputs": [], "source": [ "a = torch.linspace(-2, 2, 5)\n", "print(F.relu(a), torch.relu(a), nn.ReLU()(a), sep=\"\\n\")" ] }, { "cell_type": "markdown", "id": "4be24cdc", "metadata": {}, "source": [ "A model is a subclass of `nn.Module`. Its `__init__` creates the parameters, and its `forward` method computes the output. Wrapping a tensor in `nn.Parameter` registers it with the module, so that it appears in `model.parameters()` and requires gradients.\n", "\n", "Logistic regression computes the logit $a = \\mathbf{w}^\\top \\mathbf{x} + b$ and the probability $p(t = 1 \\mid \\mathbf{x}) = \\sigma(a)$. The model returns the logit $a$; the sigmoid is applied inside the loss function (see Loss Functions below). " ] }, { "cell_type": "code", "execution_count": null, "id": "5a04bfc3", "metadata": {}, "outputs": [], "source": [ "class LogisticRegression(nn.Module):\n", " def __init__(self, d_in):\n", " super().__init__()\n", " self.w = nn.Parameter(torch.zeros(d_in, 1))\n", " self.b = nn.Parameter(torch.zeros(1))\n", "\n", " def forward(self, x):\n", " return x @ self.w + self.b # logits, shape (N, 1)\n", "\n", "\n", "model = LogisticRegression(d_in=2)\n", "for name, p in model.named_parameters():\n", " print(name, tuple(p.shape), p.requires_grad)" ] }, { "cell_type": "markdown", "id": "3effd51e", "metadata": {}, "source": [ "Although the computation is defined in a method named `forward`, a model is used by calling it directly, as `model(x)`. The call runs `forward(x)` together with any hooks that have been registered on the module, which is why models should always be invoked as `model(x)` rather than through `model.forward(x)`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5e1f6ca3", "metadata": {}, "outputs": [], "source": [ "x, t = train_data[:5]\n", "print(model(x).T)\n", "print(torch.equal(model(x), model.forward(x)))" ] }, { "cell_type": "markdown", "id": "2fd40000", "metadata": {}, "source": [ "The basic linear layer is implemented as `nn.Linear(d_in, d_out)`. Its weight matrix has shape $(D_\\text{out}, D_\\text{in})$, one row per output unit, and it computes $\\mathbf{x}\\mathbf{W}^\\top + \\mathbf{b}$ for a batch of row vectors. Unlike our class above, `nn.Linear` initializes its parameters randomly; Part 2 explains why and how.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "dac42975", "metadata": {}, "outputs": [], "source": [ "class LogisticRegressionLinear(nn.Module):\n", " def __init__(self, d_in):\n", " super().__init__()\n", " self.linear = nn.Linear(d_in, 1)\n", "\n", " def forward(self, x):\n", " return self.linear(x)\n", "\n", "\n", "torch.manual_seed(189)\n", "model = LogisticRegressionLinear(d_in=2)\n", "print(model)\n", "W, b = model.linear.weight, model.linear.bias\n", "print(\"weight shape:\", tuple(W.shape), \" bias shape:\", tuple(b.shape))\n", "print(\"same as x @ W.T + b:\", torch.allclose(model(x), x @ W.T + b))\n", "print(model.state_dict())" ] }, { "cell_type": "markdown", "id": "5681984e", "metadata": {}, "source": [ "`model.state_dict()` holds the values of all parameters (and buffers, see Part 3) by name. It is what `torch.save` writes to disk and `model.load_state_dict` restores." ] }, { "cell_type": "code", "execution_count": null, "id": "14848eac", "metadata": {}, "outputs": [], "source": [ "model.state_dict()" ] }, { "cell_type": "markdown", "id": "965e5201", "metadata": {}, "source": [ "A multilayer perceptron applies a sequence of linear layers with an activation function between consecutive layers. No activation follows the final layer, so the network returns logits. The layers are stored in an `nn.ModuleList` rather than an ordinary Python list, because only the former registers them with the model. With a plain list their parameters would not appear in `model.parameters()`, and the optimizer would never update them.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0c217570", "metadata": {}, "outputs": [], "source": [ "class MLP(nn.Module):\n", " def __init__(self, dims, activation=\"relu\"):\n", " super().__init__()\n", " self.layers = nn.ModuleList(\n", " [nn.Linear(dims[i], dims[i+1]) for i in range(len(dims)-1)])\n", " if activation == \"relu\": self.act = nn.functional.relu\n", " else: self.act = torch.tanh\n", "\n", " def forward(self, x):\n", " for i, layer in enumerate(self.layers):\n", " x = layer(x)\n", " if i < len(self.layers)-1:\n", " x = self.act(x)\n", " return x" ] }, { "cell_type": "code", "execution_count": null, "id": "38a7f0a0", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "model = MLP([2, 16, 16, 1], activation=\"tanh\")\n", "print(model)\n", "print(\"number of parameters:\", sum(p.numel() for p in model.parameters()))\n", "print(model(x).T)" ] }, { "cell_type": "markdown", "id": "c092d5be", "metadata": {}, "source": [ "\n", "\n", "

Loss Functions

\n", "\n", "*Lecture:* slide \"torch.nn.functional and torch.nn\".\n", "\n", "PyTorch's classification losses come in two versions. One takes logits, the outputs $a$ of the last linear layer before any sigmoid or softmax, and the other takes probabilities. In general the logit version should be used, as the model returns logits and the loss then applies the sigmoid or softmax internally in a way that is numerically stable.\n", "\n", "| Targets | Logit version (use) | Probability version |\n", "|---|---|---|\n", "| binary, $t_n \\in \\{0, 1\\}$ | `nn.BCEWithLogitsLoss()(a, t)` | `nn.BCELoss()(torch.sigmoid(a), t)` |\n", "| $K$ classes, $t_n \\in \\{0, \\dots, K-1\\}$ | `nn.CrossEntropyLoss()(a, t)` | `nn.NLLLoss()(torch.log(torch.softmax(a, dim=1)), t)` |\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ec205541", "metadata": {}, "outputs": [], "source": [ "loss_fn = nn.BCEWithLogitsLoss()\n", "x, t = train_data[:]\n", "a = model(x) # logits\n", "print(\"logit version: \", loss_fn(a, t).item())\n", "print(\"probability version:\", nn.BCELoss()(torch.sigmoid(a), t).item())" ] }, { "cell_type": "markdown", "id": "215b9721", "metadata": {}, "source": [ "On these data the two versions give the same value. They differ when the model is confident. In `float32` the sigmoid of 20 rounds to exactly 1, so for a confidently wrong prediction, with $a = 20$ and $t = 0$, the probability version must evaluate $\\ln(1 - 1)$. `BCELoss` guards against this by clipping the loss at 100, but the gradient is lost in the process. The logit version never forms the probability and computes $\\ln(1 + e^{20}) \\approx 20$ directly.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6d733482", "metadata": {}, "outputs": [], "source": [ "a_big, t_zero = torch.tensor([[20.0]]), torch.tensor([[0.0]])\n", "print(\"logit version: \", nn.BCEWithLogitsLoss()(a_big, t_zero).item())\n", "print(\"probability version:\", nn.BCELoss()(torch.sigmoid(a_big), t_zero).item())" ] }, { "cell_type": "markdown", "id": "b56ccb7f", "metadata": {}, "source": [ "For a problem with $K$ classes, the model produces $K$ logits for each data point, and `nn.CrossEntropyLoss` takes these logits together with the class labels $t_n$, supplied as an `int64` tensor of class indices.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f0ebc234", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "logits = torch.randn(5, 3) # N = 5 data points, K = 3 classes\n", "targets = torch.tensor([0, 2, 1, 1, 0]) # class indices\n", "print(\"CrossEntropyLoss:\", nn.CrossEntropyLoss()(logits, targets).item())" ] }, { "cell_type": "markdown", "id": "52870821", "metadata": {}, "source": [ "\n", "\n", "

Training

\n", "\n", "*Lecture:* slide \"Training Loop with an Optimizer\".\n", "\n", "Training a network repeats four steps: reset the gradients, compute the loss on a mini-batch (the forward pass), compute the gradient of the loss with respect to the parameters (the backward pass), and update the parameters. We begin by computing the gradient of the training loss for the MLP defined above.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "bc937277", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "model = MLP([2, 16, 16, 1], activation=\"tanh\")\n", "loss_fn = nn.BCEWithLogitsLoss()\n", "x, t = train_data[:]\n", "\n", "model.zero_grad()\n", "pred = model(x)\n", "loss = loss_fn(pred, t)\n", "loss.backward()\n", "print(model.layers[0].weight.grad)" ] }, { "cell_type": "markdown", "id": "d44e0ef3", "metadata": {}, "source": [ "The update itself is carried out by an optimizer. A call to `optimizer.step()` applies the optimizer's update rule to every parameter of the model. For `torch.optim.SGD` this is the gradient descent step $\\mathbf{w} \\leftarrow \\mathbf{w} - \\eta \\nabla E(\\mathbf{w})$ of Lecture 10, and for `torch.optim.Adam` it is the Adam update of Lecture 11. The complete mini-batch training loop combines a `DataLoader`, an optimizer, and the four steps above. Here we use Adam with learning rate $\\eta = 0.01$ and train for 100 epochs, each consisting of $\\lceil 320/32 \\rceil = 10$ mini-batches.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "657b1ac0", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "model = MLP([2, 16, 16, 1], activation=\"tanh\")\n", "loss_fn = nn.BCEWithLogitsLoss()\n", "eta, n_epochs = 0.01, 100\n", "\n", "loader = DataLoader(train_data, batch_size=32, shuffle=True)\n", "optimizer = torch.optim.Adam(model.parameters(), lr=eta)\n", "\n", "for epoch in range(n_epochs):\n", " for x, t in loader:\n", " optimizer.zero_grad()\n", " loss = loss_fn(model(x), t)\n", " loss.backward()\n", " optimizer.step()" ] }, { "cell_type": "code", "execution_count": null, "id": "0a41cc9d", "metadata": {}, "outputs": [], "source": [ "with torch.no_grad():\n", " accuracy = ((model(x_val) > 0) == (t_val > 0.5)).float().mean().item()\n", "print(f\"validation accuracy: {accuracy:.3f}\")\n", "plot_decision_boundary(model, X_moons, t_moons, title=\"MLP [2, 16, 16, 1], tanh, Adam\").show()" ] }, { "cell_type": "markdown", "id": "47ad262d", "metadata": {}, "source": [ "To monitor training, we record the training loss and the validation loss after every epoch. The training loss for an epoch is the average of its mini-batch losses, weighted by batch size, and so it reflects weights that changed over the course of the epoch. The validation loss, in contrast, is computed under `torch.no_grad()` once the epoch has finished, using the final weights.\n", "\n", "The methods `model.train()` and `model.eval()` switch a model between its training and evaluation behavior. They make no difference for the MLP, but they do matter for batch normalization (Part 3) and dropout (Lecture 15), so the loop calls them anyway. We use the function below for the rest of the notebook.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c659b88d", "metadata": {}, "outputs": [], "source": [ "def train_model(model, train_data, val_data, n_epochs=50, eta=1e-3,\n", " optimizer_class=torch.optim.Adam, batch_size=32, seed=189):\n", " \"\"\"Mini-batch training with BCEWithLogitsLoss. Returns the training and validation\n", " loss after every epoch.\"\"\"\n", " loss_fn = nn.BCEWithLogitsLoss()\n", " loader = DataLoader(train_data, batch_size=batch_size, shuffle=True,\n", " generator=torch.Generator().manual_seed(seed))\n", " optimizer = optimizer_class(model.parameters(), lr=eta)\n", " x_val, t_val = val_data[:]\n", " history = {\"train\": [], \"validation\": []}\n", " for epoch in range(n_epochs):\n", " model.train()\n", " total = 0.0\n", " for x, t in loader:\n", " optimizer.zero_grad()\n", " loss = loss_fn(model(x), t)\n", " loss.backward()\n", " optimizer.step()\n", " total += loss.item() * len(x)\n", " history[\"train\"].append(total / len(train_data))\n", " model.eval()\n", " with torch.no_grad():\n", " history[\"validation\"].append(loss_fn(model(x_val), t_val).item())\n", " return history" ] }, { "cell_type": "code", "execution_count": null, "id": "a1ef8fb3", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "model = MLP([2, 16, 16, 1], activation=\"relu\")\n", "history = train_model(model, train_data, val_data, n_epochs=100, eta=0.01)\n", "print(f\"final training loss {history['train'][-1]:.3f}, \"\n", " f\"validation loss {history['validation'][-1]:.3f}\")\n", "plot_curves(history, title=\"Learning curves: MLP [2, 16, 16, 1], ReLU, Adam (eta = 0.01)\").show()\n", "plot_decision_boundary(model, X_moons, t_moons, title=\"MLP [2, 16, 16, 1], ReLU\").show()" ] }, { "cell_type": "markdown", "id": "320ea9b2", "metadata": {}, "source": [ "Both losses fall quickly during the first 20 epochs and then level off. The validation loss remains close to the training loss, and here is even slightly below it, so the network is not overfitting these 320 points. Part of the reason the training loss is slightly higher is that it averages over the epoch, during which the weights were still improving. The decision boundary follows the two interleaved moons.\n" ] }, { "cell_type": "markdown", "id": "f5494fb7", "metadata": {}, "source": [ "\n", "\n", "

(Lecture Demo 2) Part 2: Parameter Initialization

\n", "\n", "*Lecture:* run from the slides \"Parameter Initialization\" through \"Weight Initialization in PyTorch\".\n", "\n", "Gradient descent requires a starting point $\\mathbf{w}^{(0)}$, and for a deep network this choice turns out to matter a great deal. Two considerations guide it (Bishop, Section 7.2.5). First, the initial weights must break the symmetry between the hidden units of a layer, since units that start out identical remain identical. Second, the weights must have the right scale, so that the pre-activations of every layer fall in the range where the activation functions have useful derivatives.\n" ] }, { "cell_type": "markdown", "id": "7132a135", "metadata": {}, "source": [ "\n", "\n", "

The Sweet Spot of an Activation Function

\n", "\n", "During backpropagation, the gradient that flows through a hidden unit is multiplied by the derivative of its activation function, $h'(a)$, evaluated at the unit's pre-activation (Lecture 13). How well a network trains therefore depends on where its pre-activations fall. The plot below shows the sigmoid, tanh, and ReLU functions together with their derivatives, which we compute using autograd. The shaded region marks the range of $a$ in which $h'(a)$ remains at least 10% of its maximum value. For the sigmoid and tanh this region is narrow. Once $|a|$ grows beyond it, these functions flatten out, or saturate, their derivatives approach zero, and very little gradient reaches the weights in earlier layers. ReLU behaves differently, as its derivative is exactly 1 for every positive pre-activation, however large, and exactly 0 for every negative one.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3cca688b", "metadata": {}, "outputs": [], "source": [ "a = torch.linspace(-6, 6, 1201, requires_grad=True)\n", "activations = {\"sigmoid\": torch.sigmoid, \"tanh\": torch.tanh, \"ReLU\": torch.relu}\n", "\n", "fig = make_subplots(rows=2, cols=3, subplot_titles=list(activations), shared_xaxes=True,\n", " vertical_spacing=0.08)\n", "for j, (name, h) in enumerate(activations.items()):\n", " a.grad = None\n", " z = h(a)\n", " z.sum().backward() # d/da_i of sum_j h(a_j) is h'(a_i)\n", " dz = a.grad\n", " aa, zz, dd = a.detach().numpy(), z.detach().numpy(), dz.numpy()\n", " fig.add_scatter(x=aa, y=zz, line=dict(color=colors[0], width=3), name=\"h(a)\",\n", " showlegend=(j == 0), row=1, col=j + 1)\n", " fig.add_scatter(x=aa, y=dd, line=dict(color=colors[1], width=3), name=\"h'(a)\",\n", " showlegend=(j == 0), row=2, col=j + 1)\n", " sweet = aa[dd >= 0.1 * dd.max()]\n", " for row in (1, 2):\n", " fig.add_vrect(x0=sweet.min(), x1=sweet.max(), fillcolor=\"LightGreen\", opacity=0.3,\n", " line_width=0, row=row, col=j + 1)\n", " print(f\"{name:8s} h'(a) >= 10% of its maximum for a in [{sweet.min():.2f}, {sweet.max():.2f}]\")\n", "fig.update_yaxes(range=[-1.2, 1.2], row=1)\n", "fig.update_yaxes(range=[-0.5, 6], row=1, col=3)\n", "fig.update_yaxes(range=[-0.05, 1.1], row=2)\n", "fig.update_xaxes(title_text=\"a\", row=2)\n", "fig.update_layout(width=950, height=550, title=\"Activations (top) and their derivatives (bottom)\")\n", "fig.show()" ] }, { "cell_type": "markdown", "id": "806c8ed7", "metadata": {}, "source": [ "\n", "\n", "

Symmetry

\n", "\n", "Suppose that two hidden units in the same layer receive the same inputs and have identical incoming and outgoing weights. They then compute the same function and receive the same gradient, so every gradient descent step changes them in the same way and they remain identical throughout training (Bishop, Sections 6.2.4 and 7.2.5). In effect, the layer contains only a single unit.\n", "\n", "**All weights zero.** The most extreme case is to initialize every weight and bias to zero. All hidden pre-activations are then $a = 0$, and since $h(0) = 0$ for both tanh and ReLU, the output layer receives $\\mathbf{z} = \\mathbf{0}$, which makes the gradient of its weights, $\\delta \\mathbf{z}^\\top$, zero as well. The error signal passed back to a hidden unit, $h'(a) \\sum_k w_{kj} \\delta_k$, also vanishes, because it is weighted by outgoing weights $w_{kj}$ that are themselves zero. The only parameter that receives a nonzero gradient is the output bias.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5e5e318a", "metadata": {}, "outputs": [], "source": [ "def constant_init(model, c):\n", " \"\"\"Set every weight to c and every bias to 0.\"\"\"\n", " with torch.no_grad():\n", " for layer in model.layers:\n", " layer.weight.fill_(c)\n", " layer.bias.zero_()\n", " return model\n", "\n", "\n", "x, t = train_data[:]\n", "model_zero = constant_init(MLP([2, 8, 8, 1], activation=\"tanh\"), 0.0)\n", "model_zero.zero_grad()\n", "nn.BCEWithLogitsLoss()(model_zero(x), t).backward()\n", "for name, p in model_zero.named_parameters():\n", " print(f\"{name:16s} largest |gradient| = {p.grad.abs().max().item():.4f}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "46b397e5", "metadata": {}, "outputs": [], "source": [ "history_zero = train_model(model_zero, train_data, val_data, n_epochs=50, eta=0.01)\n", "print(\"hidden weights after training, all zero:\",\n", " all((layer.weight == 0).all().item() for layer in model_zero.layers))\n", "print(\"output bias:\", model_zero.layers[-1].bias.item(), \" training loss:\", history_zero[\"train\"][-1])" ] }, { "cell_type": "markdown", "id": "ab75e9ed", "metadata": {}, "source": [ "The network can therefore learn nothing but its output bias, which converges to the logit of the fraction of class-1 training points. Since the classes are balanced this is close to zero, and the training loss remains at $\\ln 2 \\approx 0.693$.\n", "\n", "**All weights equal but nonzero.** If instead every weight is set to $0.5$, the gradients are no longer zero and the network does learn. The hidden units in each layer, however, begin as identical copies of one another and remain so throughout training. Only random initialization breaks this symmetry.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "68f2add2", "metadata": {}, "outputs": [], "source": [ "model_const = constant_init(MLP([2, 8, 8, 1], activation=\"tanh\"), 0.5)\n", "history_const = train_model(model_const, train_data, val_data, n_epochs=50, eta=0.01)\n", "torch.manual_seed(189)\n", "model_random = MLP([2, 8, 8, 1], activation=\"tanh\")\n", "history_random = train_model(model_random, train_data, val_data, n_epochs=50, eta=0.01)\n", "\n", "for name, m, h in [(\"constant 0.5\", model_const, history_const),\n", " (\"random\", model_random, history_random)]:\n", " W1 = m.layers[0].weight.detach()\n", " spread = (W1 - W1[0]).abs().max().item() # largest difference from the first row\n", " print(f\"{name:12s} first-layer rows differ by at most {spread:.2e}; \"\n", " f\"training loss {h['train'][-1]:.3f}\")\n", "print(\"first-layer weights after training from the constant 0.5:\\n\", model_const.layers[0].weight.data)" ] }, { "cell_type": "code", "execution_count": null, "id": "881db230", "metadata": {}, "outputs": [], "source": [ "plot_decision_boundaries([model_zero, model_const, model_random],\n", " [\"all zero\", \"all 0.5\", \"random (PyTorch default)\"],\n", " X_moons, t_moons).show()" ] }, { "cell_type": "markdown", "id": "fef1a2c5", "metadata": {}, "source": [ "The decision boundaries illustrate the three cases. The network initialized with zeros predicts the same class probability everywhere. The network started from equal weights has eight identical units in each layer, which gives it the capacity of a network with a single unit per layer, and its boundary is close to a straight line. Only the randomly initialized network fits the moons.\n" ] }, { "cell_type": "markdown", "id": "10a4a1f6", "metadata": {}, "source": [ "\n", "\n", "

PyTorch's Default Initialization

\n", "\n", "Unless told otherwise, `nn.Linear(fan_in, fan_out)` draws its weights and biases from the uniform distribution $\\mathcal{U}(-1/\\sqrt{M}, 1/\\sqrt{M})$, where $M$ = `fan_in` is the number of inputs to each unit. Internally, PyTorch implements this as `nn.init.kaiming_uniform_(weight, a=math.sqrt(5))`, whose bound $\\text{gain} \\cdot \\sqrt{3/M}$, with $\\text{gain} = \\sqrt{2/(1 + a^2)} = \\sqrt{1/3}$, reduces to exactly $1/\\sqrt{M}$. Since a uniform distribution on $[-b, b]$ has variance $b^2/3$, the default weights have variance $1/(3M)$, a value we will compare with the alternatives below.\n", "\n", "Why? Sadly, this is due to historically reasons. The early Pytorch initialization was more conservative and for backwards compatibility reasons the same initialization is still used.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "942dfb9c", "metadata": {}, "outputs": [], "source": [ "fan_in, fan_out = 256, 128\n", "torch.manual_seed(189)\n", "layer = nn.Linear(fan_in, fan_out)\n", "bound = 1 / math.sqrt(fan_in)\n", "print(f\"bound 1/sqrt(M) = {bound:.5f}; largest |weight| = {layer.weight.abs().max().item():.5f}, \"\n", " f\"largest |bias| = {layer.bias.abs().max().item():.5f}\")\n", "print(f\"weight std {layer.weight.std().item():.5f}; predicted 1/sqrt(3M) = {1 / math.sqrt(3 * fan_in):.5f}\")\n", "\n", "torch.manual_seed(189)\n", "w = torch.empty(fan_out, fan_in)\n", "nn.init.kaiming_uniform_(w, a=math.sqrt(5))\n", "print(\"same weights as kaiming_uniform_(a=sqrt(5)):\", torch.equal(w, layer.weight))" ] }, { "cell_type": "markdown", "id": "58b47fc1", "metadata": {}, "source": [ "\n", "\n", "

Variance of the Pre-activations

\n", "\n", "To understand how the scale of the weights affects a deep network, consider layer $l$ of a network with $M^{(l-1)}$ inputs per unit (Bishop, Eqs. 7.19 and 7.20):\n", "$$\n", "a_i^{(l)} = \\sum_{j=1}^{M^{(l-1)}} w_{ij} z_j^{(l-1)}, \\qquad z_i^{(l)} = h(a_i^{(l)}).\n", "$$\n", "Suppose that the weights are drawn independently with mean zero and variance $\\epsilon^2$, and independently of the inputs $z_j^{(l-1)}$. Each pre-activation then has mean zero, and because the terms of the sum are independent, its variance is the sum of the variances of the terms. For a single term, the identity $\\operatorname{var}[A] = \\mathbb{E}[A^2] - \\mathbb{E}[A]^2$ together with the independence of $w_{ij}$ and $z_j^{(l-1)}$ gives\n", "$$\n", "\\operatorname{var}[w_{ij} z_j^{(l-1)}] = \\mathbb{E}[w_{ij}^2]\\, \\mathbb{E}\\big[(z_j^{(l-1)})^2\\big] - \\mathbb{E}[w_{ij}]^2\\, \\mathbb{E}\\big[z_j^{(l-1)}\\big]^2 = \\epsilon^2\\, \\mathbb{E}\\big[(z_j^{(l-1)})^2\\big],\n", "$$\n", "as the zero mean of the weights removes the second term and makes $\\mathbb{E}[w_{ij}^2] = \\operatorname{var}[w_{ij}] = \\epsilon^2$. Summing the $M^{(l-1)}$ identical terms gives\n", "$$\n", "\\operatorname{var}[a_i^{(l)}] = M^{(l-1)} \\epsilon^2 \\, \\mathbb{E}\\big[(z_j^{(l-1)})^2\\big].\n", "$$\n", "The expectation $\\mathbb{E}[z^2]$ appears here, rather than the variance of $z$, because the outputs of a ReLU do not have zero mean.\n", "\n", "**He initialization (ReLU).** Suppose that $a = a_j^{(l-1)}$ has a density $p(a)$ that is symmetric about zero, $p(-a) = p(a)$, with variance $\\lambda^2$. Since $\\operatorname{ReLU}(a) = 0$ for $a < 0$, only positive values of $a$ contribute to $\\mathbb{E}[z^2]$:\n", "$$\n", "\\mathbb{E}[z^2] = \\int_{-\\infty}^{\\infty} \\operatorname{ReLU}(a)^2 \\, p(a)\\, da = \\int_{0}^{\\infty} a^2 \\, p(a)\\, da = \\frac{1}{2} \\int_{-\\infty}^{\\infty} a^2 \\, p(a)\\, da = \\frac{1}{2}\\lambda^2.\n", "$$\n", "The third step uses the symmetry of $p(a)$, as the integral of $a^2 p(a)$ over $a > 0$ equals the integral over $a < 0$ and is therefore half of the full integral. The full integral is $\\mathbb{E}[a^2]$, which equals $\\operatorname{var}[a] = \\lambda^2$ as $a$ has zero mean. (The variance of $z$ itself is smaller than $\\frac{1}{2}\\lambda^2$, as $z$ has a positive mean.) Substituting $\\mathbb{E}[z^2] = \\frac{1}{2}\\operatorname{var}[a^{(l-1)}]$ gives\n", "$$\n", "\\operatorname{var}[a^{(l)}] = \\tfrac{1}{2} M^{(l-1)} \\epsilon^2 \\operatorname{var}[a^{(l-1)}].\n", "$$\n", "The variance is thus preserved from one layer to the next when $\\epsilon^2 = 2/M^{(l-1)}$, a choice known as He initialization (He et al., 2015). If the weights are instead drawn from a uniform distribution on $[-b, b]$, whose variance is $b^2/3$, the same variance requires $b = \\sqrt{6/M^{(l-1)}}$. (Bishop's Eq. 7.22 states this result for the pre-activations, so the variance on its left-hand side should read $\\operatorname{var}[a^{(l)}]$.)\n", "\n", "**Xavier initialization (tanh).** For tanh units the argument changes slightly. Near zero $\\tanh(a) \\approx a$, so $\\mathbb{E}[z^2] \\approx \\operatorname{var}[a]$, and the forward pass preserves the variance when $\\epsilon^2 = 1/M^{(l-1)}$. The same reasoning applied to the backward pass shows that the variance of the gradients is preserved when $\\epsilon^2 = 1/M^{(l)}$, where $M^{(l)}$ is the number of outputs of the layer. Glorot and Bengio (2010) proposed the compromise $\\epsilon^2 = 2/(M^{(l-1)} + M^{(l)})$, which reduces to $1/M$ when a layer has as many outputs as inputs.\n", "\n", "**Experiment.** To see these effects directly, we pass $N = 1000$ standard normal inputs through a network of 30 layers, each of width 256, and record the standard deviation of the pre-activations $a^{(l)}$ in every layer. The biases are set to zero, except under PyTorch's default, which initializes them as well. We compare six weight distributions:\n", "- $\\mathcal{N}(0, 1)$, which is much too large;\n", "- $\\mathcal{N}(0, 0.01^2)$, a small fixed scale;\n", "- Xavier, $\\epsilon^2 = 2/(M^{(l-1)} + M^{(l)}) = 1/256$;\n", "- Xavier with the gain $5/3$ that PyTorch recommends for tanh, $\\epsilon = \\frac{5}{3}/16$;\n", "- He, $\\epsilon^2 = 2/M^{(l-1)} = 2/256$;\n", "- PyTorch's default, $\\mathcal{U}(\\pm 1/\\sqrt{M})$, so that $\\epsilon^2 = 1/(3M)$.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "16bf2789", "metadata": {}, "outputs": [], "source": [ "INITS = {\n", " \"N(0, 1)\": lambda w: nn.init.normal_(w, 0.0, 1.0),\n", " \"N(0, 0.01²)\": lambda w: nn.init.normal_(w, 0.0, 0.01),\n", " \"Xavier\": lambda w: nn.init.xavier_normal_(w),\n", " \"Xavier, gain 5/3\": lambda w: nn.init.xavier_normal_(w, gain=nn.init.calculate_gain(\"tanh\")),\n", " \"He\": lambda w: nn.init.kaiming_normal_(w, nonlinearity=\"relu\"),\n", " \"PyTorch default\": None,\n", "}\n", "\n", "\n", "def preactivation_stds(model, x):\n", " \"\"\"Standard deviation of the pre-activations a = layer(z) in every layer of an MLP.\"\"\"\n", " stds = []\n", " with torch.no_grad():\n", " for layer in model.layers:\n", " a = layer(x)\n", " stds.append(a.std().item())\n", " x = model.act(a)\n", " return stds\n", "\n", "\n", "depth, width = 30, 256\n", "x0 = torch.randn(1000, width, generator=torch.Generator().manual_seed(189))\n", "stds = {}\n", "for activation in (\"tanh\", \"relu\"):\n", " for name, init in INITS.items():\n", " torch.manual_seed(189)\n", " model = MLP([width] * (depth + 1), activation=activation)\n", " if init is not None:\n", " for layer in model.layers:\n", " init(layer.weight)\n", " nn.init.zeros_(layer.bias)\n", " stds[activation, name] = preactivation_stds(model, x0)\n", "\n", "table = pd.DataFrame({key: [s[l - 1] for l in (1, 2, 5, 10, 20, 30)] for key, s in stds.items()},\n", " index=[f\"layer {l}\" for l in (1, 2, 5, 10, 20, 30)]).T\n", "table.style.format(\"{:.3g}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "fb739ded", "metadata": {}, "outputs": [], "source": [ "fig = make_subplots(rows=1, cols=2, subplot_titles=[\"tanh\", \"ReLU\"], shared_yaxes=True)\n", "for j, activation in enumerate((\"tanh\", \"relu\")):\n", " for i, name in enumerate(INITS):\n", " fig.add_scatter(x=np.arange(1, depth + 1), y=stds[activation, name], mode=\"lines+markers\",\n", " name=name, legendgroup=name, showlegend=(j == 0),\n", " line=dict(color=colors[i], width=2), marker=dict(size=4),\n", " row=1, col=j + 1)\n", "fig.update_yaxes(type=\"log\", exponentformat=\"power\", range=[-30, 33])\n", "fig.update_yaxes(title_text=\"std of a^(l)\", col=1)\n", "fig.update_xaxes(title_text=\"layer l\")\n", "fig.update_layout(width=1000, height=500,\n", " title=\"Standard deviation of the pre-activations, 30 layers of width 256\")\n", "fig.show()" ] }, { "cell_type": "markdown", "id": "1b8c71ce", "metadata": {}, "source": [ "The poorly scaled initializations fail in opposite directions. Weights drawn from $\\mathcal{N}(0, 1)$ are far too large, and already in the first layer the pre-activations have standard deviation $\\sqrt{256} = 16$. With tanh almost every unit is then saturated at $\\pm 1$, so although the standard deviation stays near 16, the derivative $h'(a)$ is close to zero nearly everywhere and no gradient can flow. With ReLU nothing saturates, and the standard deviation instead grows by a factor of $\\sqrt{256/2} \\approx 11$ in every layer, reaching about $10^{31}$ by layer 30. The small fixed scale $\\mathcal{N}(0, 0.01^2)$ has the opposite effect. Each layer multiplies the standard deviation by $0.01 \\sqrt{256} = 0.16$ for tanh, or by $0.16/\\sqrt{2}$ for ReLU, so by layer 30 the signal has shrunk to about $10^{-24}$ and $10^{-29}$ respectively, and has effectively vanished.\n", "\n", "The scaled initializations behave well when they are matched to their activation function. With He initialization and ReLU, the standard deviation remains of order 1 in all 30 layers, between about 0.9 and 1.45. The variation from layer to layer is a random fluctuation caused by the finite width of 256 rather than a systematic trend, and the value of $\\sqrt{2} \\approx 1.41$ in the first layer arises because the network's inputs are standard normal rather than outputs of a ReLU. With Xavier initialization, tanh networks decay slowly, from 1 to about 0.13 at layer 30, because $\\tanh(a)$ is smaller in magnitude than $a$ whenever $a \\neq 0$, and the gain of $5/3$ compensates for this and holds the standard deviation near 1.1. Mismatched pairs do worse. Xavier with ReLU loses a factor of $\\sqrt{2}$ per layer and falls to about $3 \\times 10^{-5}$, while adding the gain of $5/3$ overshoots, and the standard deviation grows to about 130.\n", "\n", "PyTorch's default weight variance of $1/(3M)$ is six times smaller than He's, and the standard deviation of the pre-activations drops to about 0.04 within a few layers. It then stops falling, because the default also draws the biases from $\\mathcal{U}(\\pm 1/16)$, whose standard deviation is $1/\\sqrt{3 \\cdot 256} \\approx 0.036$. Beyond the first few layers the pre-activations consist mostly of these biases, and any dependence on the input has been lost.\n" ] }, { "cell_type": "markdown", "id": "91430923", "metadata": {}, "source": [ "\n", "\n", "

Weight Initialization in PyTorch

\n", "\n", "*Lecture:* slide \"Weight Initialization in PyTorch\".\n", "\n", "The functions in `nn.init` overwrite the values of a tensor in place. A convenient pattern is to apply them in a `reset_parameters` method that is called at the end of `__init__`, here using He initialization (`kaiming_normal_`) for ReLU networks, Xavier initialization (`xavier_normal_`) otherwise, and zero biases.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "be5a5970", "metadata": {}, "outputs": [], "source": [ "class MLPModel(nn.Module):\n", " def __init__(self, dims, activation=\"relu\"):\n", " super().__init__()\n", " self.activation = activation\n", " self.layers = nn.ModuleList(\n", " [nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)])\n", " self.act = torch.relu if activation == \"relu\" else torch.tanh\n", " self.reset_parameters()\n", "\n", " def reset_parameters(self):\n", " for layer in self.layers:\n", " if self.activation == \"relu\":\n", " nn.init.kaiming_normal_(layer.weight, nonlinearity=\"relu\") # He\n", " else:\n", " nn.init.xavier_normal_(layer.weight) # Xavier\n", " nn.init.zeros_(layer.bias)\n", "\n", " def forward(self, x):\n", " for i, layer in enumerate(self.layers):\n", " x = layer(x)\n", " if i < len(self.layers) - 1:\n", " x = self.act(x)\n", " return x" ] }, { "cell_type": "markdown", "id": "b099d8f7", "metadata": {}, "source": [ "By default `xavier_normal_` uses a gain of 1. For tanh networks PyTorch recommends the gain $5/3$ returned by `nn.init.calculate_gain(\"tanh\")`, passed as `nn.init.xavier_normal_(layer.weight, gain=5/3)`, and the experiment above shows why this correction helps. As a check, the sample standard deviations of the initialized weights agree with the formulas.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ddfd3343", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "for activation, predicted in [(\"relu\", math.sqrt(2 / 256)), (\"tanh\", math.sqrt(2 / (256 + 256)))]:\n", " model = MLPModel([256, 256, 256], activation=activation)\n", " print(f\"{activation}: weight std {model.layers[0].weight.std().item():.4f}, \"\n", " f\"predicted {predicted:.4f}\")" ] }, { "cell_type": "markdown", "id": "894a705a", "metadata": {}, "source": [ "\n", "\n", "

Initialization and Training

\n", "\n", "Whether a deep network trains at all depends on the scale of its initial weights. To demonstrate this, we train a tanh network with 20 hidden layers of width 64 on the moons data, using Adam with $\\eta = 10^{-3}$ for 60 epochs, from three different initializations: small weights drawn from $\\mathcal{N}(0, 0.01^2)$, PyTorch's default, and Xavier initialization through `MLPModel`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "aa5be227", "metadata": {}, "outputs": [], "source": [ "dims = [2] + [64] * 20 + [1]\n", "\n", "torch.manual_seed(189)\n", "model_small = MLPModel(dims, activation=\"tanh\")\n", "for layer in model_small.layers:\n", " nn.init.normal_(layer.weight, 0.0, 0.01)\n", "\n", "torch.manual_seed(189)\n", "model_default = MLP(dims, activation=\"tanh\")\n", "\n", "torch.manual_seed(189)\n", "model_xavier = MLPModel(dims, activation=\"tanh\")\n", "\n", "init_histories = {}\n", "for name, m in [(\"N(0, 0.01²)\", model_small), (\"PyTorch default\", model_default),\n", " (\"Xavier\", model_xavier)]:\n", " init_histories[name] = train_model(m, train_data, val_data, n_epochs=60, eta=1e-3)\n", " print(f\"{name:16s} training loss after 1 epoch {init_histories[name]['train'][0]:.3f}, \"\n", " f\"after 60 epochs {init_histories[name]['train'][-1]:.3f}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "9630cd09", "metadata": {}, "outputs": [], "source": [ "plot_curves({name: h[\"train\"] for name, h in init_histories.items()},\n", " title=\"Training loss, tanh network with 20 hidden layers of width 64\").show()\n", "plot_decision_boundaries([model_small, model_default, model_xavier],\n", " list(init_histories), X_moons, t_moons).show()" ] }, { "cell_type": "markdown", "id": "cc546138", "metadata": {}, "source": [ "With small initial weights, the signal vanishes long before it reaches the output layer. The last hidden pre-activations have a standard deviation of about $10^{-23}$, so the output of the network does not depend on its input, the gradients of all but the last few layers are vanishingly small, and the loss never leaves $\\ln 2 \\approx 0.693$. Under the default initialization the signal also decays, down to the floor set by the biases that we found above, and training is erratic, ending with a loss near 0.62 and an almost linear decision boundary. Xavier initialization keeps the signal alive through all 20 layers, and the network fits the moons.\n" ] }, { "cell_type": "markdown", "id": "6a8a794a", "metadata": {}, "source": [ "\n", "\n", "

(Lecture Demo 3) Part 3: Normalization

\n", "\n", "*Lecture:* run from the slides \"Data Normalization\" through \"Batch-Norm and Layer-Norm in PyTorch\".\n" ] }, { "cell_type": "markdown", "id": "4ecb7152", "metadata": {}, "source": [ "\n", "\n", "

Data Normalization

\n", "\n", "When the input variables are measured on very different scales, the error function has very different curvature along different directions in weight space, which slows gradient descent (Bishop, Section 7.4.1, and Lecture 10). To illustrate this, we take two features of the breast cancer data: the mean texture, which is around 19, and the mean area, which is around 650. Each input is standardized using the mean and standard deviation computed on the training set (Bishop, Eqs. 7.48 to 7.50):\n", "$$\n", "\\mu_i = \\frac{1}{N} \\sum_{n=1}^N x_{ni}, \\qquad\n", "\\sigma_i^2 = \\frac{1}{N} \\sum_{n=1}^N (x_{ni} - \\mu_i)^2, \\qquad\n", "\\tilde{x}_{ni} = \\frac{x_{ni} - \\mu_i}{\\sigma_i}.\n", "$$\n", "The same values of $\\mu_i$ and $\\sigma_i$ must then be used for the validation and test data, so that every input is transformed in the same way.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "34e89c48", "metadata": {}, "outputs": [], "source": [ "cancer = load_breast_cancer(as_frame=True).frame\n", "cancer_data = TensorDataset(\n", " torch.tensor(cancer[[\"mean texture\", \"mean area\"]].values, dtype=torch.float32),\n", " torch.tensor(cancer[\"target\"].values, dtype=torch.float32).reshape(-1, 1))\n", "cancer_train, cancer_val = random_split(cancer_data, [0.8, 0.2],\n", " generator=torch.Generator().manual_seed(189))\n", "x_train, t_train = cancer_train[:]\n", "x_val_c, t_val_c = cancer_val[:]\n", "\n", "mu = x_train.mean(dim=0)\n", "sigma = x_train.std(dim=0, correction=0) # 1/N, as in Eq. 7.49\n", "x_train_std = (x_train - mu) / sigma # broadcasting over the rows\n", "x_val_std = (x_val_c - mu) / sigma # training statistics for validation data\n", "\n", "print(\"training mean, std: \", mu.numpy(), sigma.numpy())\n", "print(\"standardized training mean, std: \", x_train_std.mean(dim=0).numpy().round(6),\n", " x_train_std.std(dim=0, correction=0).numpy().round(6))\n", "print(\"standardized validation mean, std:\", x_val_std.mean(dim=0).numpy().round(3),\n", " x_val_std.std(dim=0, correction=0).numpy().round(3))" ] }, { "cell_type": "markdown", "id": "e64e4b7b", "metadata": {}, "source": [ "After standardization the training data have mean 0 and standard deviation 1 exactly, and the validation data approximately so. To see the effect on the error function, recall that for logistic regression the Hessian is $\\frac{1}{N}\\sum_n \\sigma(a_n)(1 - \\sigma(a_n)) \\tilde{\\mathbf{x}}_n \\tilde{\\mathbf{x}}_n^\\top$, where $\\tilde{\\mathbf{x}}_n$ includes a constant 1 for the bias. Its shape is therefore governed by the matrix $\\frac{1}{N} \\sum_n \\tilde{\\mathbf{x}}_n \\tilde{\\mathbf{x}}_n^\\top$, whose condition number, the ratio of its largest to its smallest eigenvalue, falls from about $10^7$ for the raw inputs to about 2 after standardization.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e18524e6", "metadata": {}, "outputs": [], "source": [ "def condition_number(x):\n", " x1 = torch.cat([torch.ones(len(x), 1), x], dim=1).double() # constant 1 for the bias\n", " eigenvalues = torch.linalg.eigvalsh(x1.T @ x1 / len(x1))\n", " return (eigenvalues.max() / eigenvalues.min()).item()\n", "\n", "print(f\"condition number, raw inputs: {condition_number(x_train):.3g}\")\n", "print(f\"condition number, standardized inputs: {condition_number(x_train_std):.3g}\")" ] }, { "cell_type": "markdown", "id": "4c912ddd", "metadata": {}, "source": [ "Recall from Lecture 10 that gradient descent is stable only if $\\eta < 2/\\lambda_\\text{max}$, and that it converges along the direction of smallest curvature at the rate $1 - \\eta \\lambda_\\text{min}$. With the raw inputs $\\lambda_\\text{max}$ is so large that the learning rate must be about $10^{-5}$, and at that learning rate progress along the remaining directions is extremely slow. We train the `LogisticRegression` model from Part 1 with full-batch gradient descent for 200 steps on each version of the inputs, in each case with a learning rate close to the largest stable value.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "3297c6a3", "metadata": {}, "outputs": [], "source": [ "def gradient_descent(model, x, t, eta, n_steps=200):\n", " \"\"\"Full-batch gradient descent with BCEWithLogitsLoss; returns the loss at every step.\"\"\"\n", " loss_fn = nn.BCEWithLogitsLoss()\n", " optimizer = torch.optim.SGD(model.parameters(), lr=eta)\n", " losses = []\n", " for step in range(n_steps):\n", " optimizer.zero_grad()\n", " loss = loss_fn(model(x), t)\n", " loss.backward()\n", " optimizer.step()\n", " losses.append(loss.item())\n", " return losses\n", "\n", "losses_raw = gradient_descent(LogisticRegression(2), x_train, t_train, eta=1e-5)\n", "losses_std = gradient_descent(LogisticRegression(2), x_train_std, t_train, eta=1.0)\n", "print(f\"after 200 steps: raw inputs {losses_raw[-1]:.3f}, standardized inputs {losses_std[-1]:.3f}\")\n", "plot_curves({\"raw inputs (eta = 1e-5)\": losses_raw, \"standardized inputs (eta = 1)\": losses_std},\n", " title=\"Logistic regression on two breast cancer features\",\n", " xaxis_title=\"gradient descent step\").show()" ] }, { "cell_type": "markdown", "id": "1df6ba34", "metadata": {}, "source": [ "\n", "\n", "

Batch Normalization

\n", "\n", "Batch normalization applies the same idea to the pre-activations of a hidden layer (Bishop, Section 7.4.2). Unlike the inputs, however, the hidden pre-activations change whenever the weights are updated, so their statistics must be recomputed on every mini-batch. For a mini-batch of $K$ data points, hidden unit $i$ is normalized as\n", "$$\n", "\\mu_i = \\frac{1}{K} \\sum_{n=1}^K a_{ni}, \\qquad\n", "\\sigma_i^2 = \\frac{1}{K} \\sum_{n=1}^K (a_{ni} - \\mu_i)^2, \\qquad\n", "\\hat{a}_{ni} = \\frac{a_{ni} - \\mu_i}{\\sqrt{\\sigma_i^2 + \\delta}}, \\qquad\n", "\\tilde{a}_{ni} = \\gamma_i \\hat{a}_{ni} + \\beta_i\n", "$$\n", "(Bishop, Eqs. 7.52 to 7.55). The learnable parameters $\\gamma_i$ and $\\beta_i$ give the layer back the freedom to choose its own scale and shift, and the small constant $\\delta$, `eps = 1e-5` in PyTorch, prevents division by zero. The statistics are computed separately for each unit, that is, for each column, across the data points of the mini-batch, which are the rows; in code this corresponds to `dim=0`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "6a679df7", "metadata": {}, "outputs": [], "source": [ "def batch_norm(a, gamma, beta, delta=1e-5):\n", " \"\"\"Batch normalization of a mini-batch a with shape (K, M), Eqs. 7.52 to 7.55.\"\"\"\n", " mu = a.mean(dim=0) # (M,): one mean per unit\n", " var = ((a - mu) ** 2).mean(dim=0) # (M,): biased variance, 1/K\n", " a_hat = (a - mu) / torch.sqrt(var + delta)\n", " return gamma * a_hat + beta" ] }, { "cell_type": "markdown", "id": "5e0da55a", "metadata": {}, "source": [ "The layer `nn.BatchNorm1d(M)` stores $\\boldsymbol{\\gamma}$ as `weight`, initialized to 1, and $\\boldsymbol{\\beta}$ as `bias`, initialized to 0. So that the comparison also tests Eq. 7.55, we set them to arbitrary values before applying the layer, in training mode, to a mini-batch of $K = 8$ pre-activations of $M = 4$ units.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e773ec00", "metadata": {}, "outputs": [], "source": [ "torch.manual_seed(189)\n", "K, M_units = 8, 4\n", "a = 3 * torch.randn(K, M_units) + 5 # pre-activations with mean 5, std 3\n", "\n", "bn = nn.BatchNorm1d(M_units)\n", "with torch.no_grad():\n", " bn.weight.copy_(torch.tensor([1.0, 2.0, 0.5, -1.0])) # gamma\n", " bn.bias.copy_(torch.tensor([0.0, 1.0, -1.0, 0.5])) # beta\n", "\n", "bn.train()\n", "out = bn(a)\n", "print(\"output column means:\", out.mean(dim=0).detach().numpy().round(4), \"(= beta)\")\n", "print(\"same as by hand:\", torch.allclose(out, batch_norm(a, bn.weight, bn.bias), atol=1e-6))" ] }, { "cell_type": "markdown", "id": "538beb46", "metadata": {}, "source": [ "**Running statistics.** At test time we often need a prediction for a single data point, which has no mini-batch statistics. Batch normalization therefore uses moving averages of $\\mu_i$ and $\\sigma_i^2$ that are accumulated during training (Bishop, Eqs. 7.56 and 7.57). PyTorch updates these averages after every forward pass in training mode according to\n", "$$\n", "\\bar{\\mu}_i \\leftarrow (1 - m)\\,\\bar{\\mu}_i + m\\,\\mu_i, \\qquad\n", "\\bar{\\sigma}_i^2 \\leftarrow (1 - m)\\,\\bar{\\sigma}_i^2 + m\\,\\frac{K}{K - 1}\\,\\sigma_i^2,\n", "$$\n", "with `momentum` $m = 0.1$, which corresponds to Bishop's $\\alpha = 1 - m = 0.9$. PyTorch departs from the book in two details. It averages the variance rather than the standard deviation, and it uses the unbiased batch variance $\\frac{1}{K - 1}\\sum_n (a_{ni} - \\mu_i)^2$ in the running average while normalizing with the biased one. The running values start at $\\bar{\\mu} = 0$ and $\\bar{\\sigma}^2 = 1$, and they are stored as buffers, which are saved in the `state_dict` but are not trained.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e54020df", "metadata": {}, "outputs": [], "source": [ "def update_running_stats(running_mean, running_var, a, momentum=0.1):\n", " \"\"\"PyTorch's moving averages of the batch mean and the unbiased batch variance.\"\"\"\n", " mu = a.mean(dim=0)\n", " var_unbiased = a.var(dim=0, correction=1) # 1/(K - 1)\n", " return ((1 - momentum) * running_mean + momentum * mu,\n", " (1 - momentum) * running_var + momentum * var_unbiased)\n", "\n", "\n", "bn = nn.BatchNorm1d(M_units) # fresh layer: running mean 0, variance 1\n", "running_mean, running_var = torch.zeros(M_units), torch.ones(M_units)\n", "torch.manual_seed(189)\n", "for step in range(5):\n", " a_batch = 3 * torch.randn(K, M_units) + 5\n", " bn(a_batch) # training mode updates the buffers\n", " running_mean, running_var = update_running_stats(running_mean, running_var, a_batch)\n", "\n", "print(\"PyTorch running_mean:\", bn.running_mean.numpy().round(4))\n", "print(\"by hand: \", running_mean.numpy().round(4))\n", "print(\"PyTorch running_var: \", bn.running_var.numpy().round(4))\n", "print(\"by hand: \", running_var.numpy().round(4))\n", "print(\"batches tracked:\", bn.num_batches_tracked.item())" ] }, { "cell_type": "markdown", "id": "1c23cceb", "metadata": {}, "source": [ "After five mini-batches the running mean has moved only $1 - 0.9^5 \\approx 41\\%$ of the way from its initial value of 0 toward the true mean of 5, so the moving averages need many steps to settle. In evaluation mode, selected with `bn.eval()`, the layer normalizes with the running statistics instead of the batch statistics. The output for each data point then no longer depends on the other points in the batch, and a batch containing a single point is allowed. In training mode such a batch raises an error, because its variance is zero.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f2e8cbfe", "metadata": {}, "outputs": [], "source": [ "bn.eval()\n", "a_new = 3 * torch.randn(3, M_units) + 5\n", "expected = (a_new - bn.running_mean) / torch.sqrt(bn.running_var + bn.eps) * bn.weight + bn.bias\n", "print(\"eval mode uses the running statistics:\", torch.allclose(bn(a_new), expected, atol=1e-6))\n", "print(\"one data point in eval mode:\", bn(a_new[:1]).detach().numpy().round(3))\n", "\n", "bn.train()\n", "try:\n", " bn(a_new[:1])\n", "except ValueError as err:\n", " print(\"ValueError:\", err)" ] }, { "cell_type": "markdown", "id": "d0f34ef6", "metadata": {}, "source": [ "\n", "\n", "

Layer Normalization

\n", "\n", "Layer normalization takes the complementary approach (Bishop, Section 7.4.3). Rather than computing statistics for each unit across the mini-batch, it computes them for each data point $n$ across the $M$ hidden units of the layer:\n", "$$\n", "\\mu_n = \\frac{1}{M} \\sum_{i=1}^M a_{ni}, \\qquad\n", "\\sigma_n^2 = \\frac{1}{M} \\sum_{i=1}^M (a_{ni} - \\mu_n)^2, \\qquad\n", "\\hat{a}_{ni} = \\frac{a_{ni} - \\mu_n}{\\sqrt{\\sigma_n^2 + \\delta}}, \\qquad\n", "\\tilde{a}_{ni} = \\gamma_i \\hat{a}_{ni} + \\beta_i.\n", "$$\n", "(Bishop's Eq. 7.59 writes $\\mu_i$ inside the sum, but the mean $\\mu_n$ of data point $n$ is intended.) In code the statistics are taken across the columns of each row, `dim=1`, while the learnable parameters $\\gamma_i$ and $\\beta_i$ remain one per unit. Because the normalization of a data point no longer depends on any other data point, training and evaluation compute the same function, and no running statistics are needed.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "d5c6a4d2", "metadata": {}, "outputs": [], "source": [ "def layer_norm(a, gamma, beta, delta=1e-5):\n", " \"\"\"Layer normalization of a with shape (K, M), Eqs. 7.58 to 7.60 and 7.55.\"\"\"\n", " mu = a.mean(dim=1, keepdim=True) # (K, 1): one mean per data point\n", " var = ((a - mu) ** 2).mean(dim=1, keepdim=True) # (K, 1): biased variance, 1/M\n", " a_hat = (a - mu) / torch.sqrt(var + delta)\n", " return gamma * a_hat + beta\n", "\n", "\n", "ln = nn.LayerNorm(M_units)\n", "with torch.no_grad():\n", " ln.weight.copy_(torch.tensor([1.0, 2.0, 0.5, -1.0]))\n", " ln.bias.copy_(torch.tensor([0.0, 1.0, -1.0, 0.5]))\n", "for mode in (\"train\", \"eval\"):\n", " getattr(ln, mode)()\n", " print(f\"{mode} mode, same as by hand:\", torch.allclose(ln(a), layer_norm(a, ln.weight, ln.bias), atol=1e-6))\n", "\n", "ln_plain = nn.LayerNorm(M_units) # gamma = 1, beta = 0\n", "print(\"row means:\", ln_plain(a).mean(dim=1).detach().numpy().round(4))\n", "print(\"row stds: \", ln_plain(a).std(dim=1, correction=0).detach().numpy().round(4))" ] }, { "cell_type": "markdown", "id": "bb578ded", "metadata": {}, "source": [ "**Parameters and buffers.** Both layers hold the learnable parameters $\\boldsymbol{\\gamma}$ (`weight`) and $\\boldsymbol{\\beta}$ (`bias`), one of each per unit, and the optimizer updates them like any other parameters. Only batch normalization also needs buffers, to hold its running statistics.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0e81cc66", "metadata": {}, "outputs": [], "source": [ "for layer in (nn.BatchNorm1d(M_units), nn.LayerNorm(M_units)):\n", " print(layer)\n", " print(\" parameters:\", [(name, tuple(p.shape)) for name, p in layer.named_parameters()])\n", " print(\" buffers: \", [(name, tuple(b.shape)) for name, b in layer.named_buffers()])" ] }, { "cell_type": "markdown", "id": "702e684c", "metadata": {}, "source": [ "\n", "\n", "

RMSNorm

\n", "\n", "*Lecture:* slide \"RMS Normalization (RMSNorm)\".\n", "\n", "RMSNorm (Zhang and Sennrich, 2019) is a simplification of layer normalization that is not covered in Bishop but is now widely used. It retains the re-scaling step of layer normalization and drops the re-centering, dividing each data point by the root mean square of its pre-activations. It has a learnable scale $\\gamma_i$ for each unit but no shift $\\beta_i$:\n", "$$\n", "\\text{RMS}_n = \\sqrt{\\frac{1}{M} \\sum_{i=1}^M a_{ni}^2 + \\delta}, \\qquad\n", "\\tilde{a}_{ni} = \\gamma_i \\, \\frac{a_{ni}}{\\text{RMS}_n}.\n", "$$\n", "The slide writes the same quantities as $z_{km}^{(l)}$, with data point $k$ and unit $m$. When the mean $\\mu_n$ happens to be zero, $\\text{RMS}_n^2 = \\sigma_n^2 + \\delta$, and RMSNorm coincides with layer normalization without $\\beta$. PyTorch's `nn.RMSNorm` uses the machine epsilon of the data type as $\\delta$ unless `eps` is given.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "a6acc4da", "metadata": {}, "outputs": [], "source": [ "def rms_norm(a, gamma, delta=1e-5):\n", " # RMSNorm of a with shape (K, M): one statistic per data point, no shift\n", " rms = torch.sqrt((a ** 2).mean(dim=1, keepdim=True) + delta) # (K, 1)\n", " return gamma * a / rms\n", "\n", "\n", "rms = nn.RMSNorm(M_units, eps=1e-5)\n", "with torch.no_grad():\n", " rms.weight.copy_(torch.tensor([1.0, 2.0, 0.5, -1.0]))\n", "print(\"same as by hand:\", torch.allclose(rms(a), rms_norm(a, rms.weight), atol=1e-6))\n", "print(\"parameters:\", [(name, tuple(p.shape)) for name, p in rms.named_parameters()])\n", "\n", "rms_plain = nn.RMSNorm(M_units, eps=1e-5) # gamma = 1\n", "c = 10.0\n", "print(\"scale invariant, a -> 10 a: \", torch.allclose(rms_plain(c * a), rms_plain(a), atol=1e-4))\n", "print(\"shift invariant, a -> a + 10:\", torch.allclose(rms_plain(a + c), rms_plain(a), atol=1e-4))\n", "a0 = a - a.mean(dim=1, keepdim=True) # zero mean per data point\n", "ln0 = nn.LayerNorm(M_units, eps=1e-5, bias=False)\n", "print(\"equals layer norm when the mean is zero:\", torch.allclose(rms_plain(a0), ln0(a0), atol=1e-5))" ] }, { "cell_type": "markdown", "id": "d5123e71", "metadata": {}, "source": [ "The cell confirms the key difference between the two methods. Layer normalization is unaffected by either re-scaling or shifting the pre-activations of a data point, while RMSNorm is unaffected only by re-scaling. In exchange, RMSNorm needs a single statistic per data point instead of two and has no $\\boldsymbol{\\beta}$, which makes it somewhat cheaper to compute, and it is the normalization used in most recent large language models.\n" ] }, { "cell_type": "markdown", "id": "4c1250d4", "metadata": {}, "source": [ "\n", "\n", "

Batch-Norm and Layer-Norm in PyTorch

\n", "\n", "*Lecture:* slide \"Batch-Norm and Layer-Norm in PyTorch\".\n", "\n", "In a network, a normalization layer is inserted after each hidden linear layer, so that it normalizes the pre-activations before the activation function is applied, as in Bishop. With batch normalization the bias of the preceding linear layer becomes redundant, because subtracting the batch mean removes any constant offset and $\\beta_i$ takes over its role.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "759592f2", "metadata": {}, "outputs": [], "source": [ "class NormMLP(nn.Module):\n", " def __init__(self, dims, activation=\"relu\", norm_kind=\"batch\"):\n", " super().__init__()\n", " self.layers = nn.ModuleList(\n", " [nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)])\n", " self.act = torch.relu if activation == \"relu\" else torch.tanh\n", " if norm_kind == \"batch\":\n", " self.norms = nn.ModuleList([nn.BatchNorm1d(c) for c in dims[1:-1]])\n", " else: # norm_kind == \"layer\"\n", " self.norms = nn.ModuleList([nn.LayerNorm(c) for c in dims[1:-1]])\n", "\n", " def forward(self, x):\n", " for i, layer in enumerate(self.layers):\n", " x = layer(x)\n", " if i < len(self.layers) - 1:\n", " x = self.norms[i](x)\n", " x = self.act(x)\n", " return x\n", "\n", "\n", "print(NormMLP([2, 8, 8, 1], activation=\"tanh\", norm_kind=\"batch\"))" ] }, { "cell_type": "markdown", "id": "60beb030", "metadata": {}, "source": [ "To see the effect of normalization, we return to the deep tanh network with 20 hidden layers of width 64 and PyTorch's default initialization, which trained poorly in Part 2. We train it with plain mini-batch SGD ($\\eta = 0.02$, $B = 32$) for 60 epochs in three versions: without normalization, with batch normalization, and with layer normalization.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "83f5a7df", "metadata": {}, "outputs": [], "source": [ "dims = [2] + [64] * 20 + [1]\n", "norm_models = {}\n", "torch.manual_seed(189)\n", "norm_models[\"no normalization\"] = MLP(dims, activation=\"tanh\")\n", "torch.manual_seed(189)\n", "norm_models[\"batch norm\"] = NormMLP(dims, activation=\"tanh\", norm_kind=\"batch\")\n", "torch.manual_seed(189)\n", "norm_models[\"layer norm\"] = NormMLP(dims, activation=\"tanh\", norm_kind=\"layer\")\n", "\n", "norm_histories = {}\n", "for name, m in norm_models.items():\n", " norm_histories[name] = train_model(m, train_data, val_data, n_epochs=60, eta=0.02,\n", " optimizer_class=torch.optim.SGD)\n", " print(f\"{name:16s} training loss {norm_histories[name]['train'][-1]:.3f}, \"\n", " f\"validation loss {norm_histories[name]['validation'][-1]:.3f}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "b2af6c5d", "metadata": {}, "outputs": [], "source": [ "fig = make_subplots(rows=1, cols=2, subplot_titles=[\"training loss\", \"validation loss\"], shared_yaxes=True)\n", "for i, (name, h) in enumerate(norm_histories.items()):\n", " for j, key in enumerate((\"train\", \"validation\")):\n", " fig.add_scatter(x=np.arange(1, len(h[key]) + 1), y=h[key], mode=\"lines\", name=name,\n", " legendgroup=name, showlegend=(j == 0),\n", " line=dict(color=colors[i], width=3), row=1, col=j + 1)\n", "fig.update_xaxes(title_text=\"epoch\")\n", "fig.update_yaxes(title_text=\"loss\", col=1)\n", "fig.update_layout(width=1000, height=450,\n", " title=\"SGD on a tanh network with 20 hidden layers of width 64 (default initialization)\")\n", "fig.show()\n", "plot_decision_boundaries(list(norm_models.values()), list(norm_models), X_moons, t_moons).show()" ] }, { "cell_type": "markdown", "id": "9f9f1fc6", "metadata": {}, "source": [ "Without normalization the loss remains at $\\ln 2 \\approx 0.693$ for all 60 epochs, because the signal decays through the 20 layers and SGD receives almost no gradient. With either form of normalization, every layer receives pre-activations with zero mean and unit variance (before $\\gamma$ and $\\beta$ are applied), so the tanh units operate in their sweet spot and the loss falls within the first few epochs. The curves with normalization are noisy, since plain SGD with a fixed learning rate occasionally takes large steps in a 20-layer network. The validation loss of the batch-normalized network is computed in evaluation mode with the running statistics, which can differ from the statistics of the current mini-batch.\n" ] } ], "metadata": { "kernelspec": { "display_name": "py_3_11", "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.13" } }, "nbformat": 4, "nbformat_minor": 5 }