In [1]:
import plotly.io as pio; pio.renderers.default = "notebook_connected"

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

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.

Notebook section Slide
Part 1: PyTorch (Lecture Demo 1) "Demo" after the PyTorch slides
Tensors "torch.tensor"
Automatic Differentiation "Computing Gradients using Autograd"
Datasets and Data Loaders "Big Concepts in PyTorch"
Models "torch.nn.functional and torch.nn", "torch.nn.Module"
Loss Functions "torch.nn.functional and torch.nn"
Training "Training Loop with an Optimizer"
Part 2: Parameter Initialization (Lecture Demo 2) "Parameter Initialization" through "Weight Initialization in PyTorch"
Part 3: Normalization (Lecture Demo 3) "Data Normalization" through "Batch-Norm and Layer-Norm in PyTorch"

Every cell runs in a few seconds on a laptop CPU.

In [2]:
import copy
import math

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from sklearn.datasets import make_moons, load_breast_cancer

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader, random_split

colors = px.colors.qualitative.Plotly
px.defaults.width = 800

Plotting Code

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.

In [3]:
# Berkeley blue for class 0 and California gold for class 1, as in Lecture 12.
CLASS_COLORS = {0: "#003262", 1: "#FDB515"}
PROB_COLORSCALE = [[0, "#B9CCE2"], [0.5, "#FFFFFF"], [1, "#FEE3A0"]]


def plot_data(X, t, title=None):
    """Scatter plot of a two-dimensional binary classification data set."""
    X, t = np.asarray(X), np.asarray(t).ravel()
    fig = go.Figure()
    for k, color in CLASS_COLORS.items():
        fig.add_scatter(x=X[t == k, 0], y=X[t == k, 1], mode="markers", name=f"class {k}",
                        marker=dict(color=color, size=7, line=dict(width=0.5, color="black")))
    fig.update_layout(title=title, xaxis_title="x1", yaxis_title="x2", width=600, height=500)
    fig.update_yaxes(scaleanchor="x", scaleratio=1)
    return fig


def boundary_traces(model, X, t, n=200, pad=0.5, showlegend=True):
    """Traces for the predicted probability p(t=1|x) = sigmoid(model(x)) on a grid, the
    decision boundary p = 0.5, and the data. The model is evaluated in eval mode."""
    X, t = np.asarray(X), np.asarray(t).ravel()
    lo, hi = X.min(0) - pad, X.max(0) + pad
    g1, g2 = np.linspace(lo[0], hi[0], n), np.linspace(lo[1], hi[1], n)
    G1, G2 = np.meshgrid(g1, g2)
    grid = torch.tensor(np.c_[G1.ravel(), G2.ravel()], dtype=torch.float32)
    was_training = model.training
    model.eval()
    with torch.no_grad():
        P = torch.sigmoid(model(grid)).reshape(n, n).numpy()
    model.train(was_training)
    traces = [go.Contour(x=g1, y=g2, z=P, zmin=0, zmax=1, colorscale=PROB_COLORSCALE,
                         contours=dict(start=0.5, end=0.5, size=1, coloring="heatmap"),
                         line=dict(width=2, color="black"), showscale=False, hoverinfo="skip")]
    for k, color in CLASS_COLORS.items():
        traces.append(go.Scatter(x=X[t == k, 0], y=X[t == k, 1], mode="markers",
                                 name=f"class {k}", legendgroup=f"class {k}", showlegend=showlegend,
                                 marker=dict(color=color, size=6, line=dict(width=0.5, color="black"))))
    return traces


def plot_decision_boundary(model, X, t, title=None):
    """Predicted probability, decision boundary, and data for one model."""
    fig = go.Figure(boundary_traces(model, X, t))
    fig.update_layout(title=title, xaxis_title="x1", yaxis_title="x2", width=600, height=500)
    return fig


def plot_decision_boundaries(models, titles, X, t):
    """Decision boundaries of several models side by side."""
    fig = make_subplots(rows=1, cols=len(models), subplot_titles=titles)
    for j, model in enumerate(models):
        fig.add_traces(boundary_traces(model, X, t, showlegend=(j == 0)), rows=1, cols=j + 1)
    fig.update_layout(width=330 * len(models) + 100, height=400)
    return fig


def plot_curves(curves, title=None, xaxis_title="epoch", yaxis_title="loss", log_y=False, dash=None):
    """Line plot of several curves given as a dict {name: list of values}."""
    fig = go.Figure()
    for i, (name, ys) in enumerate(curves.items()):
        line = dict(color=colors[i % len(colors)], width=3)
        if dash and dash.get(name):
            line["dash"] = dash[name]
        fig.add_scatter(x=np.arange(1, len(ys) + 1), y=ys, mode="lines", name=name, line=line)
    fig.update_layout(title=title, xaxis_title=xaxis_title, yaxis_title=yaxis_title,
                      width=800, height=450)
    if log_y:
        fig.update_yaxes(type="log", exponentformat="power")
    return fig

(Lecture Demo 1) Part 1: PyTorch

Lecture: run from the "Demo" slide after the PyTorch slides.

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.

Tensors

Lecture: slide "torch.tensor".

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.

In [4]:
x = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])
print(x)
print("shape:", x.shape, " dtype:", x.dtype)

tensor([[1., 2., 3.],
        [4., 5., 6.]])
shape: torch.Size([2, 3])  dtype: torch.float32

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.

In [5]:
print("default dtype:", torch.get_default_dtype())
n = torch.tensor([1, 2, 3])
print("from ints:", n.dtype)
try:
    n.mean()
except RuntimeError as err:
    print("RuntimeError:", err)
print("explicit dtype:", torch.tensor([1, 2, 3], dtype=torch.float32).mean())

default dtype: torch.float32
from ints: torch.int64
RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Long
explicit dtype: tensor(2.)

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().

In [6]:
a_np = np.random.default_rng(189).standard_normal((4, 3))
x64 = torch.tensor(a_np)
print("from NumPy:", x64.dtype)

layer = nn.Linear(3, 2)          # parameters are float32
try:
    layer(x64)
except RuntimeError as err:
    print("RuntimeError:", err)
print(layer(torch.tensor(a_np, dtype=torch.float32)).dtype)

from NumPy: torch.float64
RuntimeError: mat1 and mat2 must have the same dtype, but got Double and Float
torch.float32

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:

dtype bits (sign, exponent, mantissa) largest value relative precision $\epsilon$
torch.float32 1, 8, 23 $3.4 \times 10^{38}$ $1.2 \times 10^{-7}$
torch.float16 1, 5, 10 $65504$ $9.8 \times 10^{-4}$
torch.bfloat16 1, 8, 7 $3.4 \times 10^{38}$ $7.8 \times 10^{-3}$

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.

In [7]:
for dt in (torch.float32, torch.float16, torch.bfloat16):
    info = torch.finfo(dt)
    print(f"{str(dt):15s} bytes/entry {torch.tensor(0.0, dtype=dt).element_size()}   "
          f"max {info.max:.3g}   eps {info.eps:.2g}")

big, small = torch.tensor(70000.0), torch.tensor(1e-8)
print("70000 ->", big.to(torch.float16).item(), "(float16),", big.to(torch.bfloat16).item(), "(bfloat16)")
print("1e-8  ->", small.to(torch.float16).item(), "(float16),", small.to(torch.bfloat16).item(), "(bfloat16)")
one = torch.tensor(1.0, dtype=torch.bfloat16)
print("1 + 0.001 in bfloat16:", (one + 0.001).item())

n_params = 7e9                                    # a 7-billion-parameter language model
for dt, nbytes in ((torch.float32, 4), (torch.bfloat16, 2)):
    print(f"7B parameters in {dt}: {n_params * nbytes / 1e9:.0f} GB")

torch.float32   bytes/entry 4   max 3.4e+38   eps 1.2e-07
torch.float16   bytes/entry 2   max 6.55e+04   eps 0.00098
torch.bfloat16  bytes/entry 2   max 3.39e+38   eps 0.0078
70000 -> inf (float16), 70144.0 (bfloat16)
1e-8  -> 0.0 (float16), 1.0011717677116394e-08 (bfloat16)
1 + 0.001 in bfloat16: 1.0
7B parameters in torch.float32: 28 GB
7B parameters in torch.bfloat16: 14 GB

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).

Random tensors are drawn from PyTorch's global random number generator, which torch.manual_seed seeds, or from an explicit torch.Generator.

In [8]:
torch.manual_seed(189)
print(torch.zeros(2, 3))
print(torch.ones(2, 3))
print(torch.arange(6))
print(torch.randn(2, 3))                                            # N(0, 1) entries
print(torch.rand(2, 3, generator=torch.Generator().manual_seed(189)))  # U(0, 1) entries

tensor([[0., 0., 0.],
        [0., 0., 0.]])
tensor([[1., 1., 1.],
        [1., 1., 1.]])
tensor([0, 1, 2, 3, 4, 5])
tensor([[-0.6043, -2.2664,  1.7804],
        [-0.0055, -0.4880, -0.0740]])
tensor([[0.2136, 0.0424, 0.7420],
        [0.3987, 0.8749, 0.7304]])

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.

In [9]:
x = torch.arange(6, dtype=torch.float32)
X = x.reshape(2, 3)
print(X)
print("X.T shape:", X.T.shape)
print("first row:", X[0])
print("second column:", X[:, 1])
print("entries greater than 2:", X[X > 2])

tensor([[0., 1., 2.],
        [3., 4., 5.]])
X.T shape: torch.Size([3, 2])
first row: tensor([0., 1., 2.])
second column: tensor([1., 4.])
entries greater than 2: tensor([3., 4., 5.])

Unsqueeze can be a bit more confusing so let's step through it.

In [10]:
print(X, X.shape, "\n")
print(X.unsqueeze(0), X.unsqueeze(0).shape, "\n")
print(X.unsqueeze(1), X.unsqueeze(1).shape, "\n")
print(X.unsqueeze(2), X.unsqueeze(2).shape, "\n")

tensor([[0., 1., 2.],
        [3., 4., 5.]]) torch.Size([2, 3]) 

tensor([[[0., 1., 2.],
         [3., 4., 5.]]]) torch.Size([1, 2, 3]) 

tensor([[[0., 1., 2.]],

        [[3., 4., 5.]]]) torch.Size([2, 1, 3]) 

tensor([[[0.],
         [1.],
         [2.]],

        [[3.],
         [4.],
         [5.]]]) torch.Size([2, 3, 1])

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.

In [11]:
print(X + torch.tensor([10.0, 20.0, 30.0]))   # (2, 3) + (3,): added to every row
print(X - X.mean(dim=0))                      # center each column

tensor([[10., 21., 32.],
        [13., 24., 35.]])
tensor([[-1.5000, -1.5000, -1.5000],
        [ 1.5000,  1.5000,  1.5000]])

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.

In [12]:
W = torch.tensor([[1.0, 0.0, -1.0],
                  [0.5, 0.5, 0.5]])          # shape (D_out, D_in) = (2, 3)
print("X @ W.T =\n", X @ W.T, "\n")                 # shape (N, D_out) = (2, 2)

print("sum of all entries:", X.sum(), "\n")

print("column means (dim=0):", X.mean(dim=0), "\n")
print("row means (dim=1):   ", X.mean(dim=1), "\n")
print("row means, keepdim:  ", X.mean(dim=1, keepdim=True), 
      X.mean(dim=1, keepdim=True).shape, "\n")

X @ W.T =
 tensor([[-2.0000,  1.5000],
        [-2.0000,  6.0000]]) 

sum of all entries: tensor(15.) 

column means (dim=0): tensor([1.5000, 2.5000, 3.5000]) 

row means (dim=1):    tensor([1., 4.]) 

row means, keepdim:   tensor([[1.],
        [4.]]) torch.Size([2, 1])

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().

In [13]:
a = np.arange(3.0)
shared = torch.from_numpy(a)     # shares memory with a
copied = torch.tensor(a)         # copies a
a[0] = 100.0
print("from_numpy:", shared)
print("tensor:    ", copied)
print("back to NumPy:", copied.numpy(), " one entry as a float:", copied[1].item())

from_numpy: tensor([100.,   1.,   2.], dtype=torch.float64)
tensor:     tensor([0., 1., 2.], dtype=torch.float64)
back to NumPy: [0. 1. 2.]  one entry as a float: 1.0

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.

In [14]:
if torch.accelerator.is_available():
    print("accelerator available:", torch.accelerator.current_accelerator())
else:
    print("no accelerator available")

device = torch.device("cpu")    # this notebook uses the CPU
x = torch.randn(3, 2).to(device)
print(x.device, x.cpu().numpy().shape)

print("if there is a device this will move the tensor to it")
if torch.accelerator.is_available():
    device = torch.accelerator.current_accelerator()
    x = x.to(device)
    print(x.device,x)

accelerator available: mps
cpu (3, 2)
if there is a device this will move the tensor to it
mps:0 
tensor([[ 1.0426,  0.4177],
        [ 1.6498, -0.0516],
        [ 0.8123, -0.1253]], device='mps:0')

Automatic Differentiation

Lecture: slide "Computing Gradients using Autograd".

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.

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 $$ \frac{\partial f}{\partial x_1} = x_2 + x_2 e^{x_1 x_2}, \qquad \frac{\partial f}{\partial x_2} = x_1 + x_1 e^{x_1 x_2} - \cos(x_2), $$ and which takes the value $(16.78, 8.81)$ at the point $(1, 2)$.

In [15]:
x1 = torch.tensor(1.0, requires_grad=True)
x2 = torch.tensor(2.0, requires_grad=True)
In [16]:
f = x1 * x2 + torch.exp(x1 * x2) - torch.sin(x2)

print(f)
print("f.grad_fn:", f.grad_fn, " inputs:", f.grad_fn.next_functions)

tensor(8.4798, grad_fn=<SubBackward0>)
f.grad_fn: <SubBackward0 object at 0x103dc18a0>  inputs: ((<AddBackward0 object at 0x127a23af0>, 0), (<SinBackward0 object at 0x127a20d30>, 0))
In [17]:
f.backward()
print("autograd:", x1.grad.item(), x2.grad.item())
print("exact:   ", 2 + 2 * math.exp(2), 1 + math.exp(2) - math.cos(2))

autograd: 16.778112411499023 8.80520248413086
exact:    16.7781121978613 8.805202935477793

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.

In [18]:
f = x1 * x2 + torch.exp(x1 * x2) - torch.sin(x2)
f.backward()
print("after a second backward:", x1.grad.item(), x2.grad.item())   # twice the gradient

x1.grad.zero_()
x2.grad.zero_()
f = x1 * x2 + torch.exp(x1 * x2) - torch.sin(x2)
f.backward()
print("after zeroing first:    ", x1.grad.item(), x2.grad.item())

after a second backward: 33.55622482299805 17.61040496826172
after zeroing first:     16.778112411499023 8.80520248413086

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.

In [19]:
with torch.no_grad():
    y = x1 * x2
print("under no_grad:", y.requires_grad)

under no_grad: False
In [20]:
z = (x1 * x2).detach()
print("detached:     ", z.requires_grad, " with graph:", (x1 * x2).requires_grad)

detached:      False  with graph: True

Datasets and Data Loaders

Lecture: slide "Big Concepts in PyTorch".

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.

In [21]:
# Make numpy data
X_moons, t_moons = make_moons(400, noise=0.2, random_state=189)
In [22]:
dataset = TensorDataset(torch.tensor(X_moons, dtype=torch.float32),
                        torch.tensor(t_moons, dtype=torch.float32).reshape(-1, 1))
print("N =", len(dataset), " first item:", dataset[0])

N = 400  first item: (tensor([-0.1189,  0.2273]), tensor([1.]))
In [23]:
plot_data(X_moons, t_moons, title="Moons data (N = 400)").show()

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.

In [24]:
train_data, val_data = random_split(dataset, [0.8, 0.2],
                                    generator=torch.Generator().manual_seed(189))
print("training points:", len(train_data), " validation points:", len(val_data))
x_val, t_val = val_data[:]       # all validation points as two tensors
print(x_val.shape, t_val.shape)

training points: 320  validation points: 80
torch.Size([80, 2]) torch.Size([80, 1])

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.

In [25]:
for B in (32, 50):
    loader = DataLoader(train_data, batch_size=B, shuffle=True,
                        generator=torch.Generator().manual_seed(189))
    sizes = [len(x) for x, t in loader]
    print(f"B = {B}: {len(sizes)} batches = ceil({len(train_data)}/{B}) = "
          f"{math.ceil(len(train_data) / B)}, sizes {sizes}")

x, t = next(iter(loader))
print("one batch:", x.shape, t.shape)

B = 32: 10 batches = ceil(320/32) = 10, sizes [32, 32, 32, 32, 32, 32, 32, 32, 32, 32]
B = 50: 7 batches = ceil(320/50) = 7, sizes [50, 50, 50, 50, 50, 50, 20]
one batch: torch.Size([50, 2]) torch.Size([50, 1])

Models

Lecture: slides "torch.nn.functional and torch.nn" and "torch.nn.Module".

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.

In [26]:
a = torch.linspace(-2, 2, 5)
print(F.relu(a), torch.relu(a), nn.ReLU()(a), sep="\n")

tensor([0., 0., 0., 1., 2.])
tensor([0., 0., 0., 1., 2.])
tensor([0., 0., 0., 1., 2.])

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.

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).

In [27]:
class LogisticRegression(nn.Module):
    def __init__(self, d_in):
        super().__init__()
        self.w = nn.Parameter(torch.zeros(d_in, 1))
        self.b = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return x @ self.w + self.b      # logits, shape (N, 1)


model = LogisticRegression(d_in=2)
for name, p in model.named_parameters():
    print(name, tuple(p.shape), p.requires_grad)

w (2, 1) True
b (1,) True

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).

In [28]:
x, t = train_data[:5]
print(model(x).T)
print(torch.equal(model(x), model.forward(x)))

tensor([[0., 0., 0., 0., 0.]], grad_fn=<PermuteBackward0>)
True

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.

In [29]:
class LogisticRegressionLinear(nn.Module):
    def __init__(self, d_in):
        super().__init__()
        self.linear = nn.Linear(d_in, 1)

    def forward(self, x):
        return self.linear(x)


torch.manual_seed(189)
model = LogisticRegressionLinear(d_in=2)
print(model)
W, b = model.linear.weight, model.linear.bias
print("weight shape:", tuple(W.shape), " bias shape:", tuple(b.shape))
print("same as x @ W.T + b:", torch.allclose(model(x), x @ W.T + b))
print(model.state_dict())

LogisticRegressionLinear(
  (linear): Linear(in_features=2, out_features=1, bias=True)
)
weight shape: (1, 2)  bias shape: (1,)
same as x @ W.T + b: True
OrderedDict([('linear.weight', tensor([[-0.4051, -0.6472]])), ('linear.bias', tensor([0.3423]))])

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.

In [30]:
model.state_dict()
Out[30]:
OrderedDict([('linear.weight', tensor([[-0.4051, -0.6472]])),
             ('linear.bias', tensor([0.3423]))])

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.

In [31]:
class MLP(nn.Module):
    def __init__(self, dims, activation="relu"):
        super().__init__()
        self.layers = nn.ModuleList(
            [nn.Linear(dims[i], dims[i+1]) for i in range(len(dims)-1)])
        if activation == "relu": self.act = nn.functional.relu
        else: self.act = torch.tanh

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.layers)-1:
                x = self.act(x)
        return x
In [32]:
torch.manual_seed(189)
model = MLP([2, 16, 16, 1], activation="tanh")
print(model)
print("number of parameters:", sum(p.numel() for p in model.parameters()))
print(model(x).T)

MLP(
  (layers): ModuleList(
    (0): Linear(in_features=2, out_features=16, bias=True)
    (1): Linear(in_features=16, out_features=16, bias=True)
    (2): Linear(in_features=16, out_features=1, bias=True)
  )
)
number of parameters: 337
tensor([[ 0.2009,  0.1567,  0.0615, -0.0108,  0.1770]],
       grad_fn=<PermuteBackward0>)

Loss Functions

Lecture: slide "torch.nn.functional and torch.nn".

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.

Targets Logit version (use) Probability version
binary, $t_n \in \{0, 1\}$ nn.BCEWithLogitsLoss()(a, t) nn.BCELoss()(torch.sigmoid(a), t)
$K$ classes, $t_n \in \{0, \dots, K-1\}$ nn.CrossEntropyLoss()(a, t) nn.NLLLoss()(torch.log(torch.softmax(a, dim=1)), t)
In [33]:
loss_fn = nn.BCEWithLogitsLoss()
x, t = train_data[:]
a = model(x)                                    # logits
print("logit version:      ", loss_fn(a, t).item())
print("probability version:", nn.BCELoss()(torch.sigmoid(a), t).item())

logit version:       0.7091628313064575
probability version: 0.7091628313064575

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.

In [34]:
a_big, t_zero = torch.tensor([[20.0]]), torch.tensor([[0.0]])
print("logit version:      ", nn.BCEWithLogitsLoss()(a_big, t_zero).item())
print("probability version:", nn.BCELoss()(torch.sigmoid(a_big), t_zero).item())

logit version:       20.0
probability version: 100.0

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.

In [35]:
torch.manual_seed(189)
logits = torch.randn(5, 3)                      # N = 5 data points, K = 3 classes
targets = torch.tensor([0, 2, 1, 1, 0])         # class indices
print("CrossEntropyLoss:", nn.CrossEntropyLoss()(logits, targets).item())

CrossEntropyLoss: 1.2374517917633057

Training

Lecture: slide "Training Loop with an Optimizer".

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.

In [36]:
torch.manual_seed(189)
model = MLP([2, 16, 16, 1], activation="tanh")
loss_fn = nn.BCEWithLogitsLoss()
x, t = train_data[:]

model.zero_grad()
pred = model(x)
loss = loss_fn(pred, t)
loss.backward()
print(model.layers[0].weight.grad)

tensor([[-0.0144,  0.0148],
        [ 0.0205, -0.0210],
        [ 0.0072, -0.0081],
        [-0.0128,  0.0131],
        [ 0.0117, -0.0123],
        [-0.0050,  0.0081],
        [-0.0138,  0.0073],
        [ 0.0105, -0.0113],
        [-0.0039,  0.0066],
        [ 0.0061, -0.0090],
        [-0.0171,  0.0114],
        [ 0.0057, -0.0075],
        [-0.0162,  0.0140],
        [-0.0125,  0.0114],
        [ 0.0120, -0.0155],
        [ 0.0163, -0.0129]])

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.

In [37]:
torch.manual_seed(189)
model = MLP([2, 16, 16, 1], activation="tanh")
loss_fn = nn.BCEWithLogitsLoss()
eta, n_epochs = 0.01, 100

loader = DataLoader(train_data, batch_size=32, shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=eta)

for epoch in range(n_epochs):
    for x, t in loader:
        optimizer.zero_grad()
        loss = loss_fn(model(x), t)
        loss.backward()
        optimizer.step()
In [38]:
with torch.no_grad():
    accuracy = ((model(x_val) > 0) == (t_val > 0.5)).float().mean().item()
print(f"validation accuracy: {accuracy:.3f}")
plot_decision_boundary(model, X_moons, t_moons, title="MLP [2, 16, 16, 1], tanh, Adam").show()

validation accuracy: 0.975

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.

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.

In [39]:
def train_model(model, train_data, val_data, n_epochs=50, eta=1e-3,
                optimizer_class=torch.optim.Adam, batch_size=32, seed=189):
    """Mini-batch training with BCEWithLogitsLoss. Returns the training and validation
    loss after every epoch."""
    loss_fn = nn.BCEWithLogitsLoss()
    loader = DataLoader(train_data, batch_size=batch_size, shuffle=True,
                        generator=torch.Generator().manual_seed(seed))
    optimizer = optimizer_class(model.parameters(), lr=eta)
    x_val, t_val = val_data[:]
    history = {"train": [], "validation": []}
    for epoch in range(n_epochs):
        model.train()
        total = 0.0
        for x, t in loader:
            optimizer.zero_grad()
            loss = loss_fn(model(x), t)
            loss.backward()
            optimizer.step()
            total += loss.item() * len(x)
        history["train"].append(total / len(train_data))
        model.eval()
        with torch.no_grad():
            history["validation"].append(loss_fn(model(x_val), t_val).item())
    return history
In [40]:
torch.manual_seed(189)
model = MLP([2, 16, 16, 1], activation="relu")
history = train_model(model, train_data, val_data, n_epochs=100, eta=0.01)
print(f"final training loss {history['train'][-1]:.3f}, "
      f"validation loss {history['validation'][-1]:.3f}")
plot_curves(history, title="Learning curves: MLP [2, 16, 16, 1], ReLU, Adam (eta = 0.01)").show()
plot_decision_boundary(model, X_moons, t_moons, title="MLP [2, 16, 16, 1], ReLU").show()

final training loss 0.050, validation loss 0.041

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.

(Lecture Demo 2) Part 2: Parameter Initialization

Lecture: run from the slides "Parameter Initialization" through "Weight Initialization in PyTorch".

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.

The Sweet Spot of an Activation Function

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.

In [41]:
a = torch.linspace(-6, 6, 1201, requires_grad=True)
activations = {"sigmoid": torch.sigmoid, "tanh": torch.tanh, "ReLU": torch.relu}

fig = make_subplots(rows=2, cols=3, subplot_titles=list(activations), shared_xaxes=True,
                    vertical_spacing=0.08)
for j, (name, h) in enumerate(activations.items()):
    a.grad = None
    z = h(a)
    z.sum().backward()             # d/da_i of sum_j h(a_j) is h'(a_i)
    dz = a.grad
    aa, zz, dd = a.detach().numpy(), z.detach().numpy(), dz.numpy()
    fig.add_scatter(x=aa, y=zz, line=dict(color=colors[0], width=3), name="h(a)",
                    showlegend=(j == 0), row=1, col=j + 1)
    fig.add_scatter(x=aa, y=dd, line=dict(color=colors[1], width=3), name="h'(a)",
                    showlegend=(j == 0), row=2, col=j + 1)
    sweet = aa[dd >= 0.1 * dd.max()]
    for row in (1, 2):
        fig.add_vrect(x0=sweet.min(), x1=sweet.max(), fillcolor="LightGreen", opacity=0.3,
                      line_width=0, row=row, col=j + 1)
    print(f"{name:8s} h'(a) >= 10% of its maximum for a in [{sweet.min():.2f}, {sweet.max():.2f}]")
fig.update_yaxes(range=[-1.2, 1.2], row=1)
fig.update_yaxes(range=[-0.5, 6], row=1, col=3)
fig.update_yaxes(range=[-0.05, 1.1], row=2)
fig.update_xaxes(title_text="a", row=2)
fig.update_layout(width=950, height=550, title="Activations (top) and their derivatives (bottom)")
fig.show()

sigmoid  h'(a) >= 10% of its maximum for a in [-3.63, 3.63]
tanh     h'(a) >= 10% of its maximum for a in [-1.81, 1.81]
ReLU     h'(a) >= 10% of its maximum for a in [0.00, 6.00]

Symmetry

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.

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.

In [42]:
def constant_init(model, c):
    """Set every weight to c and every bias to 0."""
    with torch.no_grad():
        for layer in model.layers:
            layer.weight.fill_(c)
            layer.bias.zero_()
    return model


x, t = train_data[:]
model_zero = constant_init(MLP([2, 8, 8, 1], activation="tanh"), 0.0)
model_zero.zero_grad()
nn.BCEWithLogitsLoss()(model_zero(x), t).backward()
for name, p in model_zero.named_parameters():
    print(f"{name:16s} largest |gradient| = {p.grad.abs().max().item():.4f}")

layers.0.weight  largest |gradient| = 0.0000
layers.0.bias    largest |gradient| = 0.0000
layers.1.weight  largest |gradient| = 0.0000
layers.1.bias    largest |gradient| = 0.0000
layers.2.weight  largest |gradient| = 0.0000
layers.2.bias    largest |gradient| = 0.0031
In [43]:
history_zero = train_model(model_zero, train_data, val_data, n_epochs=50, eta=0.01)
print("hidden weights after training, all zero:",
      all((layer.weight == 0).all().item() for layer in model_zero.layers))
print("output bias:", model_zero.layers[-1].bias.item(), " training loss:", history_zero["train"][-1])

hidden weights after training, all zero: True
output bias: 0.01794794388115406  training loss: 0.6931521534919739

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$.

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.

In [44]:
model_const = constant_init(MLP([2, 8, 8, 1], activation="tanh"), 0.5)
history_const = train_model(model_const, train_data, val_data, n_epochs=50, eta=0.01)
torch.manual_seed(189)
model_random = MLP([2, 8, 8, 1], activation="tanh")
history_random = train_model(model_random, train_data, val_data, n_epochs=50, eta=0.01)

for name, m, h in [("constant 0.5", model_const, history_const),
                   ("random", model_random, history_random)]:
    W1 = m.layers[0].weight.detach()
    spread = (W1 - W1[0]).abs().max().item()    # largest difference from the first row
    print(f"{name:12s} first-layer rows differ by at most {spread:.2e}; "
          f"training loss {h['train'][-1]:.3f}")
print("first-layer weights after training from the constant 0.5:\n", model_const.layers[0].weight.data)

constant 0.5 first-layer rows differ by at most 0.00e+00; training loss 0.289
random       first-layer rows differ by at most 3.24e+00; training loss 0.083
first-layer weights after training from the constant 0.5:
 tensor([[ 0.1183, -0.5664],
        [ 0.1183, -0.5664],
        [ 0.1183, -0.5664],
        [ 0.1183, -0.5664],
        [ 0.1183, -0.5664],
        [ 0.1183, -0.5664],
        [ 0.1183, -0.5664],
        [ 0.1183, -0.5664]])
In [45]:
plot_decision_boundaries([model_zero, model_const, model_random],
                         ["all zero", "all 0.5", "random (PyTorch default)"],
                         X_moons, t_moons).show()

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.

PyTorch's Default Initialization

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.

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.

In [46]:
fan_in, fan_out = 256, 128
torch.manual_seed(189)
layer = nn.Linear(fan_in, fan_out)
bound = 1 / math.sqrt(fan_in)
print(f"bound 1/sqrt(M) = {bound:.5f}; largest |weight| = {layer.weight.abs().max().item():.5f}, "
      f"largest |bias| = {layer.bias.abs().max().item():.5f}")
print(f"weight std {layer.weight.std().item():.5f}; predicted 1/sqrt(3M) = {1 / math.sqrt(3 * fan_in):.5f}")

torch.manual_seed(189)
w = torch.empty(fan_out, fan_in)
nn.init.kaiming_uniform_(w, a=math.sqrt(5))
print("same weights as kaiming_uniform_(a=sqrt(5)):", torch.equal(w, layer.weight))

bound 1/sqrt(M) = 0.06250; largest |weight| = 0.06250, largest |bias| = 0.06197
weight std 0.03611; predicted 1/sqrt(3M) = 0.03608
same weights as kaiming_uniform_(a=sqrt(5)): True

Variance of the Pre-activations

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): $$ a_i^{(l)} = \sum_{j=1}^{M^{(l-1)}} w_{ij} z_j^{(l-1)}, \qquad z_i^{(l)} = h(a_i^{(l)}). $$ 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 $$ \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], $$ 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 $$ \operatorname{var}[a_i^{(l)}] = M^{(l-1)} \epsilon^2 \, \mathbb{E}\big[(z_j^{(l-1)})^2\big]. $$ 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.

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]$: $$ \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. $$ 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 $$ \operatorname{var}[a^{(l)}] = \tfrac{1}{2} M^{(l-1)} \epsilon^2 \operatorname{var}[a^{(l-1)}]. $$ 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)}]$.)

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.

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:

  • $\mathcal{N}(0, 1)$, which is much too large;
  • $\mathcal{N}(0, 0.01^2)$, a small fixed scale;
  • Xavier, $\epsilon^2 = 2/(M^{(l-1)} + M^{(l)}) = 1/256$;
  • Xavier with the gain $5/3$ that PyTorch recommends for tanh, $\epsilon = \frac{5}{3}/16$;
  • He, $\epsilon^2 = 2/M^{(l-1)} = 2/256$;
  • PyTorch's default, $\mathcal{U}(\pm 1/\sqrt{M})$, so that $\epsilon^2 = 1/(3M)$.
In [47]:
INITS = {
    "N(0, 1)": lambda w: nn.init.normal_(w, 0.0, 1.0),
    "N(0, 0.01²)": lambda w: nn.init.normal_(w, 0.0, 0.01),
    "Xavier": lambda w: nn.init.xavier_normal_(w),
    "Xavier, gain 5/3": lambda w: nn.init.xavier_normal_(w, gain=nn.init.calculate_gain("tanh")),
    "He": lambda w: nn.init.kaiming_normal_(w, nonlinearity="relu"),
    "PyTorch default": None,
}


def preactivation_stds(model, x):
    """Standard deviation of the pre-activations a = layer(z) in every layer of an MLP."""
    stds = []
    with torch.no_grad():
        for layer in model.layers:
            a = layer(x)
            stds.append(a.std().item())
            x = model.act(a)
    return stds


depth, width = 30, 256
x0 = torch.randn(1000, width, generator=torch.Generator().manual_seed(189))
stds = {}
for activation in ("tanh", "relu"):
    for name, init in INITS.items():
        torch.manual_seed(189)
        model = MLP([width] * (depth + 1), activation=activation)
        if init is not None:
            for layer in model.layers:
                init(layer.weight)
                nn.init.zeros_(layer.bias)
        stds[activation, name] = preactivation_stds(model, x0)

table = pd.DataFrame({key: [s[l - 1] for l in (1, 2, 5, 10, 20, 30)] for key, s in stds.items()},
                     index=[f"layer {l}" for l in (1, 2, 5, 10, 20, 30)]).T
table.style.format("{:.3g}")
Out[47]:
    layer 1 layer 2 layer 5 layer 10 layer 20 layer 30
tanh N(0, 1) 16 15.6 15.6 15.6 15.5 15.6
N(0, 0.01²) 0.16 0.025 0.000103 1.09e-08 1.2e-16 1.32e-24
Xavier 1 0.628 0.359 0.241 0.166 0.132
Xavier, gain 5/3 1.67 1.26 1.1 1.08 1.09 1.08
He 1.41 1.02 0.82 0.787 0.792 0.776
PyTorch default 0.588 0.27 0.0649 0.0406 0.0449 0.0472
relu N(0, 1) 16 183 2.52e+05 4.77e+10 1.51e+21 3.74e+31
N(0, 0.01²) 0.16 0.0183 2.52e-05 4.77e-10 1.51e-19 3.74e-29
Xavier 1 0.713 0.241 0.0434 0.00125 2.81e-05
Xavier, gain 5/3 1.67 1.98 3.1 7.18 34.1 127
He 1.41 1.43 1.36 1.39 1.28 0.921
PyTorch default 0.588 0.239 0.0441 0.0369 0.0392 0.0422
In [48]:
fig = make_subplots(rows=1, cols=2, subplot_titles=["tanh", "ReLU"], shared_yaxes=True)
for j, activation in enumerate(("tanh", "relu")):
    for i, name in enumerate(INITS):
        fig.add_scatter(x=np.arange(1, depth + 1), y=stds[activation, name], mode="lines+markers",
                        name=name, legendgroup=name, showlegend=(j == 0),
                        line=dict(color=colors[i], width=2), marker=dict(size=4),
                        row=1, col=j + 1)
fig.update_yaxes(type="log", exponentformat="power", range=[-30, 33])
fig.update_yaxes(title_text="std of a^(l)", col=1)
fig.update_xaxes(title_text="layer l")
fig.update_layout(width=1000, height=500,
                  title="Standard deviation of the pre-activations, 30 layers of width 256")
fig.show()

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.

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.

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.

Weight Initialization in PyTorch

Lecture: slide "Weight Initialization in PyTorch".

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.

In [49]:
class MLPModel(nn.Module):
    def __init__(self, dims, activation="relu"):
        super().__init__()
        self.activation = activation
        self.layers = nn.ModuleList(
            [nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)])
        self.act = torch.relu if activation == "relu" else torch.tanh
        self.reset_parameters()

    def reset_parameters(self):
        for layer in self.layers:
            if self.activation == "relu":
                nn.init.kaiming_normal_(layer.weight, nonlinearity="relu")  # He
            else:
                nn.init.xavier_normal_(layer.weight)                         # Xavier
            nn.init.zeros_(layer.bias)

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.layers) - 1:
                x = self.act(x)
        return x

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.

In [50]:
torch.manual_seed(189)
for activation, predicted in [("relu", math.sqrt(2 / 256)), ("tanh", math.sqrt(2 / (256 + 256)))]:
    model = MLPModel([256, 256, 256], activation=activation)
    print(f"{activation}: weight std {model.layers[0].weight.std().item():.4f}, "
          f"predicted {predicted:.4f}")

relu: weight std 0.0888, predicted 0.0884
tanh: weight std 0.0625, predicted 0.0625

Initialization and Training

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.

In [51]:
dims = [2] + [64] * 20 + [1]

torch.manual_seed(189)
model_small = MLPModel(dims, activation="tanh")
for layer in model_small.layers:
    nn.init.normal_(layer.weight, 0.0, 0.01)

torch.manual_seed(189)
model_default = MLP(dims, activation="tanh")

torch.manual_seed(189)
model_xavier = MLPModel(dims, activation="tanh")

init_histories = {}
for name, m in [("N(0, 0.01²)", model_small), ("PyTorch default", model_default),
                ("Xavier", model_xavier)]:
    init_histories[name] = train_model(m, train_data, val_data, n_epochs=60, eta=1e-3)
    print(f"{name:16s} training loss after 1 epoch {init_histories[name]['train'][0]:.3f}, "
          f"after 60 epochs {init_histories[name]['train'][-1]:.3f}")

N(0, 0.01²)      training loss after 1 epoch 0.693, after 60 epochs 0.693

PyTorch default  training loss after 1 epoch 0.695, after 60 epochs 0.618

Xavier           training loss after 1 epoch 0.438, after 60 epochs 0.080
In [52]:
plot_curves({name: h["train"] for name, h in init_histories.items()},
            title="Training loss, tanh network with 20 hidden layers of width 64").show()
plot_decision_boundaries([model_small, model_default, model_xavier],
                         list(init_histories), X_moons, t_moons).show()

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.

(Lecture Demo 3) Part 3: Normalization

Lecture: run from the slides "Data Normalization" through "Batch-Norm and Layer-Norm in PyTorch".

Data Normalization

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): $$ \mu_i = \frac{1}{N} \sum_{n=1}^N x_{ni}, \qquad \sigma_i^2 = \frac{1}{N} \sum_{n=1}^N (x_{ni} - \mu_i)^2, \qquad \tilde{x}_{ni} = \frac{x_{ni} - \mu_i}{\sigma_i}. $$ 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.

In [53]:
cancer = load_breast_cancer(as_frame=True).frame
cancer_data = TensorDataset(
    torch.tensor(cancer[["mean texture", "mean area"]].values, dtype=torch.float32),
    torch.tensor(cancer["target"].values, dtype=torch.float32).reshape(-1, 1))
cancer_train, cancer_val = random_split(cancer_data, [0.8, 0.2],
                                        generator=torch.Generator().manual_seed(189))
x_train, t_train = cancer_train[:]
x_val_c, t_val_c = cancer_val[:]

mu = x_train.mean(dim=0)
sigma = x_train.std(dim=0, correction=0)     # 1/N, as in Eq. 7.49
x_train_std = (x_train - mu) / sigma         # broadcasting over the rows
x_val_std = (x_val_c - mu) / sigma           # training statistics for validation data

print("training mean, std:  ", mu.numpy(), sigma.numpy())
print("standardized training mean, std:  ", x_train_std.mean(dim=0).numpy().round(6),
      x_train_std.std(dim=0, correction=0).numpy().round(6))
print("standardized validation mean, std:", x_val_std.mean(dim=0).numpy().round(3),
      x_val_std.std(dim=0, correction=0).numpy().round(3))

training mean, std:   [ 19.217762 648.6888  ] [  4.3579426 347.44595  ]
standardized training mean, std:   [0. 0.] [1. 1.]
standardized validation mean, std: [0.083 0.09 ] [0.925 1.056]

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.

In [54]:
def condition_number(x):
    x1 = torch.cat([torch.ones(len(x), 1), x], dim=1).double()   # constant 1 for the bias
    eigenvalues = torch.linalg.eigvalsh(x1.T @ x1 / len(x1))
    return (eigenvalues.max() / eigenvalues.min()).item()

print(f"condition number, raw inputs:          {condition_number(x_train):.3g}")
print(f"condition number, standardized inputs: {condition_number(x_train_std):.3g}")

condition number, raw inputs:          1.12e+07
condition number, standardized inputs: 1.96

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.

In [55]:
def gradient_descent(model, x, t, eta, n_steps=200):
    """Full-batch gradient descent with BCEWithLogitsLoss; returns the loss at every step."""
    loss_fn = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=eta)
    losses = []
    for step in range(n_steps):
        optimizer.zero_grad()
        loss = loss_fn(model(x), t)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    return losses

losses_raw = gradient_descent(LogisticRegression(2), x_train, t_train, eta=1e-5)
losses_std = gradient_descent(LogisticRegression(2), x_train_std, t_train, eta=1.0)
print(f"after 200 steps: raw inputs {losses_raw[-1]:.3f}, standardized inputs {losses_std[-1]:.3f}")
plot_curves({"raw inputs (eta = 1e-5)": losses_raw, "standardized inputs (eta = 1)": losses_std},
            title="Logistic regression on two breast cancer features",
            xaxis_title="gradient descent step").show()

after 200 steps: raw inputs 0.676, standardized inputs 0.268

Batch Normalization

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 $$ \mu_i = \frac{1}{K} \sum_{n=1}^K a_{ni}, \qquad \sigma_i^2 = \frac{1}{K} \sum_{n=1}^K (a_{ni} - \mu_i)^2, \qquad \hat{a}_{ni} = \frac{a_{ni} - \mu_i}{\sqrt{\sigma_i^2 + \delta}}, \qquad \tilde{a}_{ni} = \gamma_i \hat{a}_{ni} + \beta_i $$ (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.

In [56]:
def batch_norm(a, gamma, beta, delta=1e-5):
    """Batch normalization of a mini-batch a with shape (K, M), Eqs. 7.52 to 7.55."""
    mu = a.mean(dim=0)                          # (M,): one mean per unit
    var = ((a - mu) ** 2).mean(dim=0)           # (M,): biased variance, 1/K
    a_hat = (a - mu) / torch.sqrt(var + delta)
    return gamma * a_hat + beta

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.

In [57]:
torch.manual_seed(189)
K, M_units = 8, 4
a = 3 * torch.randn(K, M_units) + 5          # pre-activations with mean 5, std 3

bn = nn.BatchNorm1d(M_units)
with torch.no_grad():
    bn.weight.copy_(torch.tensor([1.0, 2.0, 0.5, -1.0]))     # gamma
    bn.bias.copy_(torch.tensor([0.0, 1.0, -1.0, 0.5]))       # beta

bn.train()
out = bn(a)
print("output column means:", out.mean(dim=0).detach().numpy().round(4), "(= beta)")
print("same as by hand:", torch.allclose(out, batch_norm(a, bn.weight, bn.bias), atol=1e-6))

output column means: [-0.   1.  -1.   0.5] (= beta)
same as by hand: True

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 $$ \bar{\mu}_i \leftarrow (1 - m)\,\bar{\mu}_i + m\,\mu_i, \qquad \bar{\sigma}_i^2 \leftarrow (1 - m)\,\bar{\sigma}_i^2 + m\,\frac{K}{K - 1}\,\sigma_i^2, $$ 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.

In [58]:
def update_running_stats(running_mean, running_var, a, momentum=0.1):
    """PyTorch's moving averages of the batch mean and the unbiased batch variance."""
    mu = a.mean(dim=0)
    var_unbiased = a.var(dim=0, correction=1)   # 1/(K - 1)
    return ((1 - momentum) * running_mean + momentum * mu,
            (1 - momentum) * running_var + momentum * var_unbiased)


bn = nn.BatchNorm1d(M_units)                    # fresh layer: running mean 0, variance 1
running_mean, running_var = torch.zeros(M_units), torch.ones(M_units)
torch.manual_seed(189)
for step in range(5):
    a_batch = 3 * torch.randn(K, M_units) + 5
    bn(a_batch)                                 # training mode updates the buffers
    running_mean, running_var = update_running_stats(running_mean, running_var, a_batch)

print("PyTorch running_mean:", bn.running_mean.numpy().round(4))
print("by hand:             ", running_mean.numpy().round(4))
print("PyTorch running_var: ", bn.running_var.numpy().round(4))
print("by hand:             ", running_var.numpy().round(4))
print("batches tracked:", bn.num_batches_tracked.item())

PyTorch running_mean: [2.2007 2.1236 2.3704 1.9995]
by hand:              [2.2007 2.1236 2.3704 1.9995]
PyTorch running_var:  [5.2393 5.5549 4.9969 3.4115]
by hand:              [5.2393 5.5549 4.9969 3.4115]
batches tracked: 5

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.

In [59]:
bn.eval()
a_new = 3 * torch.randn(3, M_units) + 5
expected = (a_new - bn.running_mean) / torch.sqrt(bn.running_var + bn.eps) * bn.weight + bn.bias
print("eval mode uses the running statistics:", torch.allclose(bn(a_new), expected, atol=1e-6))
print("one data point in eval mode:", bn(a_new[:1]).detach().numpy().round(3))

bn.train()
try:
    bn(a_new[:1])
except ValueError as err:
    print("ValueError:", err)

eval mode uses the running statistics: True
one data point in eval mode: [[-1.173  1.022  2.665 -2.416]]
ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 4])

Layer Normalization

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: $$ \mu_n = \frac{1}{M} \sum_{i=1}^M a_{ni}, \qquad \sigma_n^2 = \frac{1}{M} \sum_{i=1}^M (a_{ni} - \mu_n)^2, \qquad \hat{a}_{ni} = \frac{a_{ni} - \mu_n}{\sqrt{\sigma_n^2 + \delta}}, \qquad \tilde{a}_{ni} = \gamma_i \hat{a}_{ni} + \beta_i. $$ (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.

In [60]:
def layer_norm(a, gamma, beta, delta=1e-5):
    """Layer normalization of a with shape (K, M), Eqs. 7.58 to 7.60 and 7.55."""
    mu = a.mean(dim=1, keepdim=True)                     # (K, 1): one mean per data point
    var = ((a - mu) ** 2).mean(dim=1, keepdim=True)      # (K, 1): biased variance, 1/M
    a_hat = (a - mu) / torch.sqrt(var + delta)
    return gamma * a_hat + beta


ln = nn.LayerNorm(M_units)
with torch.no_grad():
    ln.weight.copy_(torch.tensor([1.0, 2.0, 0.5, -1.0]))
    ln.bias.copy_(torch.tensor([0.0, 1.0, -1.0, 0.5]))
for mode in ("train", "eval"):
    getattr(ln, mode)()
    print(f"{mode} mode, same as by hand:", torch.allclose(ln(a), layer_norm(a, ln.weight, ln.bias), atol=1e-6))

ln_plain = nn.LayerNorm(M_units)         # gamma = 1, beta = 0
print("row means:", ln_plain(a).mean(dim=1).detach().numpy().round(4))
print("row stds: ", ln_plain(a).std(dim=1, correction=0).detach().numpy().round(4))

train mode, same as by hand: True
eval mode, same as by hand: True
row means: [ 0.  0. -0.  0. -0.  0.  0. -0.]
row stds:  [1. 1. 1. 1. 1. 1. 1. 1.]

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.

In [61]:
for layer in (nn.BatchNorm1d(M_units), nn.LayerNorm(M_units)):
    print(layer)
    print("  parameters:", [(name, tuple(p.shape)) for name, p in layer.named_parameters()])
    print("  buffers:   ", [(name, tuple(b.shape)) for name, b in layer.named_buffers()])

BatchNorm1d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  parameters: [('weight', (4,)), ('bias', (4,))]
  buffers:    [('running_mean', (4,)), ('running_var', (4,)), ('num_batches_tracked', ())]
LayerNorm((4,), eps=1e-05, elementwise_affine=True)
  parameters: [('weight', (4,)), ('bias', (4,))]
  buffers:    []

RMSNorm

Lecture: slide "RMS Normalization (RMSNorm)".

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$: $$ \text{RMS}_n = \sqrt{\frac{1}{M} \sum_{i=1}^M a_{ni}^2 + \delta}, \qquad \tilde{a}_{ni} = \gamma_i \, \frac{a_{ni}}{\text{RMS}_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.

In [62]:
def rms_norm(a, gamma, delta=1e-5):
    # RMSNorm of a with shape (K, M): one statistic per data point, no shift
    rms = torch.sqrt((a ** 2).mean(dim=1, keepdim=True) + delta)   # (K, 1)
    return gamma * a / rms


rms = nn.RMSNorm(M_units, eps=1e-5)
with torch.no_grad():
    rms.weight.copy_(torch.tensor([1.0, 2.0, 0.5, -1.0]))
print("same as by hand:", torch.allclose(rms(a), rms_norm(a, rms.weight), atol=1e-6))
print("parameters:", [(name, tuple(p.shape)) for name, p in rms.named_parameters()])

rms_plain = nn.RMSNorm(M_units, eps=1e-5)          # gamma = 1
c = 10.0
print("scale invariant, a -> 10 a:  ", torch.allclose(rms_plain(c * a), rms_plain(a), atol=1e-4))
print("shift invariant, a -> a + 10:", torch.allclose(rms_plain(a + c), rms_plain(a), atol=1e-4))
a0 = a - a.mean(dim=1, keepdim=True)               # zero mean per data point
ln0 = nn.LayerNorm(M_units, eps=1e-5, bias=False)
print("equals layer norm when the mean is zero:", torch.allclose(rms_plain(a0), ln0(a0), atol=1e-5))

same as by hand: True
parameters: [('weight', (4,))]
scale invariant, a -> 10 a:   True
shift invariant, a -> a + 10: False
equals layer norm when the mean is zero: True

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.

Batch-Norm and Layer-Norm in PyTorch

Lecture: slide "Batch-Norm and Layer-Norm in PyTorch".

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.

In [63]:
class NormMLP(nn.Module):
    def __init__(self, dims, activation="relu", norm_kind="batch"):
        super().__init__()
        self.layers = nn.ModuleList(
            [nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)])
        self.act = torch.relu if activation == "relu" else torch.tanh
        if norm_kind == "batch":
            self.norms = nn.ModuleList([nn.BatchNorm1d(c) for c in dims[1:-1]])
        else:  # norm_kind == "layer"
            self.norms = nn.ModuleList([nn.LayerNorm(c) for c in dims[1:-1]])

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.layers) - 1:
                x = self.norms[i](x)
                x = self.act(x)
        return x


print(NormMLP([2, 8, 8, 1], activation="tanh", norm_kind="batch"))

NormMLP(
  (layers): ModuleList(
    (0): Linear(in_features=2, out_features=8, bias=True)
    (1): Linear(in_features=8, out_features=8, bias=True)
    (2): Linear(in_features=8, out_features=1, bias=True)
  )
  (norms): ModuleList(
    (0-1): 2 x BatchNorm1d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  )
)

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.

In [64]:
dims = [2] + [64] * 20 + [1]
norm_models = {}
torch.manual_seed(189)
norm_models["no normalization"] = MLP(dims, activation="tanh")
torch.manual_seed(189)
norm_models["batch norm"] = NormMLP(dims, activation="tanh", norm_kind="batch")
torch.manual_seed(189)
norm_models["layer norm"] = NormMLP(dims, activation="tanh", norm_kind="layer")

norm_histories = {}
for name, m in norm_models.items():
    norm_histories[name] = train_model(m, train_data, val_data, n_epochs=60, eta=0.02,
                                       optimizer_class=torch.optim.SGD)
    print(f"{name:16s} training loss {norm_histories[name]['train'][-1]:.3f}, "
          f"validation loss {norm_histories[name]['validation'][-1]:.3f}")

no normalization training loss 0.693, validation loss 0.693

batch norm       training loss 0.179, validation loss 0.104

layer norm       training loss 0.107, validation loss 0.246
In [65]:
fig = make_subplots(rows=1, cols=2, subplot_titles=["training loss", "validation loss"], shared_yaxes=True)
for i, (name, h) in enumerate(norm_histories.items()):
    for j, key in enumerate(("train", "validation")):
        fig.add_scatter(x=np.arange(1, len(h[key]) + 1), y=h[key], mode="lines", name=name,
                        legendgroup=name, showlegend=(j == 0),
                        line=dict(color=colors[i], width=3), row=1, col=j + 1)
fig.update_xaxes(title_text="epoch")
fig.update_yaxes(title_text="loss", col=1)
fig.update_layout(width=1000, height=450,
                  title="SGD on a tanh network with 20 hidden layers of width 64 (default initialization)")
fig.show()
plot_decision_boundaries(list(norm_models.values()), list(norm_models), X_moons, t_moons).show()

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.