import plotly.io as pio; pio.renderers.default = "notebook_connected"

Lecture 15: Normalization and Regularization (CS 189, Fall 2026)

This notebook accompanies Lecture 15, Section 7.4, and Chapter 9 of Bishop and Bishop, Deep Learning: Foundations and Concepts. It begins with the normalization material carried over from Lecture 14: input standardization, batch normalization, layer normalization, and RMSNorm, implemented by hand and checked against PyTorch. A network with tens of thousands of parameters is trained on 100 noisy points, where it over-fits, and the notebook then applies the regularization methods of the lecture one at a time: weight decay, early stopping, model averaging (bagging), and dropout. It also compares a deep plain network with a residual network of the same depth, and ends with an interactive model explorer.

The PyTorch code uses the pieces introduced in Lecture 14: nn.Module, nn.BCEWithLogitsLoss, torch.optim, and the training loop with training and validation curves.

Notebook section	Slide
(Lecture Demo 1) Normalization	"Data Normalization" through "Batch-Norm and Layer-Norm in PyTorch"
Setup: A Flexible Model on 100 Points	"Learning Requires Inductive Biases"
(Lecture Demo 2) Weight Decay	"Weight Decay Regularization"
Learning Curves and Early Stopping	"Learning Curves and Early Stopping", "Early Stopping is a form of Weight Decay"
Ensembles and Bagging	"Ensembles of Models (Experts)", "Bagging: Bootstrap Aggregation"
(Lecture Demo 3) Dropout	"Dropout Regularization", "Drop-out and Model Averaging", "Implementing Drop-out in PyTorch"
(Lecture Demo 4) Residual Connections	"Vanishing Gradients in Deep Networks", "Residual Networks", "Expanding the Residual Network Graph"
(Lecture Demo 5) The Model Explorer	"Demo"

import copy
import time

import numpy as np
import torch
import torch.nn as nn
from sklearn.datasets import make_moons, make_circles

import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

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

# All models in this notebook are small, so everything runs on the CPU.
torch.set_num_threads(4)

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.

CLASS_COLORS = {0: '#1f77b4', 1: '#d62728'}   # class 0 blue, class 1 red
# Light blue where p(t=1|x) is near 0, white at 0.5, light red near 1.
PROB_COLORSCALE = [[0, '#9ecae1'], [0.5, '#ffffff'], [1, '#fc9272']]


def plot_grid(data, n=120, margin=0.5):
    """A square grid of inputs that covers the training and validation points."""
    x = torch.cat([data['x_train'], data['x_val']])
    lo, hi = x.min(0).values - margin, x.max(0).values + margin
    x1 = torch.linspace(lo[0], hi[0], n)
    x2 = torch.linspace(lo[1], hi[1], n)
    X1, X2 = torch.meshgrid(x1, x2, indexing='xy')
    return x1, x2, torch.stack([X1.ravel(), X2.ravel()], dim=1)


def boundary_traces(predict, data, showlegend=True, n=120):
    """Traces for the predicted probability p(t=1|x), the decision boundary p = 0.5,
    and the training points. `predict` is a model or a function that maps an (m, 2)
    tensor of inputs to m probabilities."""
    if isinstance(predict, nn.Module):
        model = predict
        predict = lambda x: predict_proba(model, x)
    x1, x2, grid = plot_grid(data, n)
    # The grid as start and step rather than coordinate arrays: the live FigureWidget in
    # VS Code can drop the arrays, which puts the contour at index coordinates 0..n-1.
    grid_xy = dict(x0=x1[0].item(), dx=(x1[1] - x1[0]).item(),
                   y0=x2[0].item(), dy=(x2[1] - x2[0]).item())
    p = predict(grid).reshape(n, n).numpy()
    traces = [
        go.Contour(**grid_xy, z=p, zmin=0, zmax=1, colorscale=PROB_COLORSCALE,
                   contours=dict(start=0, end=1, size=0.1), line=dict(width=0.3, color='gray'),
                   showscale=False, hoverinfo='skip'),
        go.Contour(**grid_xy, z=p, showscale=False, hoverinfo='skip',
                   contours=dict(start=0.5, end=0.5, size=1, coloring='lines'),
                   colorscale=[[0, 'black'], [1, 'black']], line=dict(width=2)),
    ]
    x, t = data['x_train'].numpy(), data['t_train'].numpy()
    for k in (0, 1):
        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=CLASS_COLORS[k], size=6,
                                             line=dict(width=0.5, color='black'))))
    return traces


def plot_boundaries(predictors, data, cols=None, height=380):
    """One decision-boundary panel per entry of the dictionary {title: model or function}."""
    cols = cols or len(predictors)
    rows = (len(predictors) + cols - 1) // cols
    fig = make_subplots(rows=rows, cols=cols, subplot_titles=list(predictors),
                        horizontal_spacing=0.04, vertical_spacing=0.08)
    for i, predict in enumerate(predictors.values()):
        for trace in boundary_traces(predict, data, showlegend=(i == 0)):
            fig.add_trace(trace, row=i // cols + 1, col=i % cols + 1)
    fig.update_xaxes(showticklabels=False)
    fig.update_yaxes(showticklabels=False)
    fig.update_layout(width=max(800, 260 * cols), height=height * rows,
                      margin=dict(l=10, r=10, t=40, b=10), legend=dict(orientation='h', y=-0.02))
    return fig


def plot_history(histories, title=None, metrics=('loss', 'acc'), mark=None, log_y=False):
    """Learning curves. histories = {name: history} where each history comes from fit().
    Training curves are dashed and validation curves solid, one color per run.
    mark = {name: epoch} adds a marker at that epoch on the validation curve."""
    labels = {'loss': 'cross-entropy loss', 'acc': 'accuracy'}
    fig = make_subplots(rows=1, cols=len(metrics), subplot_titles=[labels[m] for m in metrics])
    for i, (name, h) in enumerate(histories.items()):
        color = colors[i % len(colors)]
        epochs = np.arange(1, len(h['train_loss']) + 1)
        for j, m in enumerate(metrics):
            fig.add_scatter(x=epochs, y=h[f'train_{m}'], name=f'{name} (train)', legendgroup=name,
                            line=dict(color=color, dash='dash'), showlegend=(j == 0), row=1, col=j + 1)
            fig.add_scatter(x=epochs, y=h[f'val_{m}'], name=f'{name} (validation)', legendgroup=name,
                            line=dict(color=color), showlegend=(j == 0), row=1, col=j + 1)
            if mark and name in mark:
                e = mark[name]
                fig.add_scatter(x=[e + 1], y=[h[f'val_{m}'][e]], mode='markers', showlegend=False,
                                marker=dict(color='black', size=10, symbol='x'), row=1, col=j + 1)
        fig.update_xaxes(title_text='epoch')
    if log_y:
        fig.update_yaxes(type='log', row=1, col=1)
    fig.update_layout(title=title, width=1000, height=400)
    return fig


def explorer_figure(data, model, history):
    """The figure of the model explorer: decision boundary, loss, and accuracy."""
    fig = make_subplots(rows=1, cols=3, column_widths=[0.4, 0.3, 0.3],
                        subplot_titles=['p(t=1|x) and training points', 'cross-entropy loss', 'accuracy'])
    for trace in boundary_traces(model, data, n=80):
        fig.add_trace(trace, row=1, col=1)
    for j, m in enumerate(['loss', 'acc']):
        for split, dash in [('train', 'dash'), ('val', None)]:
            fig.add_scatter(x=[], y=[], name=f'{split} {m}', line=dict(dash=dash, color=colors[j]),
                            row=1, col=j + 2)
    fig.update_xaxes(title_text='epoch', col=2)
    fig.update_xaxes(title_text='epoch', col=3)
    fig.update_layout(width=1100, height=420, margin=dict(l=10, r=10, t=40, b=10),
                      legend=dict(orientation='h', y=-0.15))
    update_explorer_figure(fig, data, model, history)
    return fig


def update_explorer_figure(fig, data, model, history):
    traces = boundary_traces(model, data, n=80)
    epochs = np.arange(1, len(history['train_loss']) + 1)
    with fig.batch_update():
        fig.data[0].z = traces[0].z
        fig.data[1].z = traces[1].z
        for i, key in enumerate(['train_loss', 'val_loss', 'train_acc', 'val_acc']):
            fig.data[4 + i].x = epochs
            fig.data[4 + i].y = history[key]

# Plotting helpers from the Lecture 14 notebook, used in the Normalization section.
# Renamed where Lecture 15 has a function or constant of the same name.
# Berkeley blue for class 0 and California gold for class 1, as in Lecture 12.
L14_CLASS_COLORS = {0: "#003262", 1: "#FDB515"}
L14_PROB_COLORSCALE = [[0, "#B9CCE2"], [0.5, "#FFFFFF"], [1, "#FEE3A0"]]


def l14_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=L14_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 L14_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_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(l14_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) Normalization

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

This section continues Lecture 14, which ended before the normalization material. It uses a few definitions from the Lecture 14 notebook, and the next cell repeats them so that this notebook runs on its own: the moons data with its training and validation split, the LogisticRegression model, Lecture 14's MLP (renamed PlainMLP here, as this notebook defines its own MLP below), and the train_model function.

# Definitions repeated from the Lecture 14 notebook so that this notebook runs on its own.
from sklearn.datasets import load_breast_cancer
from torch.utils.data import TensorDataset, DataLoader, random_split

# The moons data of Lectures 12 to 14, with the same training and validation split.
X_moons, t_moons = make_moons(400, noise=0.2, random_state=189)
dataset = TensorDataset(torch.tensor(X_moons, dtype=torch.float32),
                        torch.tensor(t_moons, dtype=torch.float32).reshape(-1, 1))
train_data, val_data = random_split(dataset, [0.8, 0.2], generator=torch.Generator().manual_seed(189))


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)


class PlainMLP(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


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

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.

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.

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 Lecture 14 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.

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.

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.

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.

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.

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.

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.

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.

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 layer's activation function, as in the batch normalization diagram on the slides, so that the next linear layer receives normalized inputs. Bishop (Section 7.4.2) and the original batch normalization paper instead normalize the pre-activations, before the activation function, and swapping the two lines in forward gives that version. Both orders are common in practice.

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.act(x)
                x = self.norms[i](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 Lecture 14. 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.

dims = [2] + [64] * 20 + [1]
norm_models = {}
torch.manual_seed(189)
norm_models["no normalization"] = PlainMLP(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.112, validation loss 0.072

layer norm       training loss 0.086, validation loss 0.077

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, the input to every linear layer has zero mean and unit variance (before $\gamma$ and $\beta$ are applied), so the pre-activations keep the same scale through all 20 layers, the tanh units stay 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.

Setup: A Flexible Model on 100 Points

Slides: "Learning Requires Inductive Biases".

The running example is binary classification on the two-moons data with a large amount of label noise near the boundary (noise=0.3). The training set has only $N = 100$ points, and a separate validation set of 1000 points from the same distribution estimates the generalization error. The model is a fully connected network with two hidden layers of 200 ReLU units, about 41,000 parameters, so it can fit the training set exactly, noise included. Without an inductive bias that prefers simple functions it over-fits, and each section below adds one such bias.

The functions in this section are course material. They use the PyTorch pieces from Lecture 14:

• make_datasets returns the training and validation inputs $\mathbf{x}_n$ and targets $t_n \in \{0, 1\}$ as tensors.
• MLP is a network with depth hidden layers of width units. Each hidden layer applies a linear map, an optional normalization (norm='batch' or 'layer', as in the Normalization section above), the activation, and dropout with drop probability dropout. The output is a single logit $a$, with $p(t = 1 \mid \mathbf{x}) = \sigma(a)$.
• fit trains with mini-batch Adam (or SGD) on the mean binary cross-entropy $E_D(\mathbf{w})$ and records the training and validation loss and accuracy after every epoch. Its argument weight_decay is the coefficient $\lambda$ of the next section.

def make_points(name, n, noise, seed):
    """n points of a two-class data set in the plane. The noise level is scaled per
    data set so that noise=0.3 gives a similar amount of class overlap."""
    rng = np.random.default_rng(seed)
    if name == 'moons':
        return make_moons(n, noise=noise, random_state=seed)
    if name == 'circles':
        return make_circles(n, noise=noise / 2, factor=0.5, random_state=seed)
    if name == 'xor':        # the label is the sign of x1 * x2
        x = rng.uniform(-1, 1, size=(n, 2))
        t = (x[:, 0] * x[:, 1] > 0).astype(int)
        return x + noise / 3 * rng.standard_normal((n, 2)), t
    if name == 'spiral':     # two interleaved spiral arms
        t = rng.integers(0, 2, n)
        r = rng.uniform(0.05, 1, n)
        theta = 3 * np.pi * r + np.pi * t
        x = np.stack([r * np.cos(theta), r * np.sin(theta)], axis=1)
        return x + noise / 6 * rng.standard_normal((n, 2)), t
    raise ValueError(f'unknown data set {name}')


def make_datasets(name='moons', n_train=100, n_val=1000, noise=0.3, seed=189):
    """Training and validation sets drawn independently from the same distribution."""
    x_train, t_train = make_points(name, n_train, noise, seed)
    x_val, t_val = make_points(name, n_val, noise, seed + 1)
    as_tensor = lambda a: torch.tensor(a, dtype=torch.float32)
    return {'x_train': as_tensor(x_train), 't_train': as_tensor(t_train),
            'x_val': as_tensor(x_val), 't_val': as_tensor(t_val)}


data = make_datasets()
print({k: tuple(v.shape) for k, v in data.items()})

{'x_train': (100, 2), 't_train': (100,), 'x_val': (1000, 2), 't_val': (1000,)}

ACTIVATIONS = {'relu': nn.ReLU, 'tanh': nn.Tanh, 'gelu': nn.GELU}


def make_norm(norm, width):
    """Normalization layer from Lecture 14, or the identity."""
    if norm == 'batch':
        return nn.BatchNorm1d(width)
    if norm == 'layer':
        return nn.LayerNorm(width)
    return nn.Identity()


class MLP(nn.Module):
    """Fully connected network with `depth` hidden layers of `width` units.
    Each hidden layer is Linear -> normalization -> activation -> Dropout(p=dropout).
    The output layer returns one logit a per input, with p(t=1|x) = sigmoid(a).
    The linear layers use PyTorch's default initialization."""

    def __init__(self, n_in=2, width=200, depth=2, activation='relu', norm='none', dropout=0.0):
        super().__init__()
        layers, d = [], n_in
        for _ in range(depth):
            layers += [nn.Linear(d, width), make_norm(norm, width),
                       ACTIVATIONS[activation](), nn.Dropout(p=dropout)]
            d = width
        self.hidden = nn.Sequential(*layers)
        self.head = nn.Linear(d, 1)

    def forward(self, x):
        return self.head(self.hidden(x)).squeeze(-1)


def count_parameters(model):
    return sum(p.numel() for p in model.parameters())


BIG_MLP = dict(width=200, depth=2)     # the over-parameterized model used in most sections
print('parameters:', count_parameters(MLP(**BIG_MLP)))

parameters: 41001

loss_fn = nn.BCEWithLogitsLoss()   # mean binary cross-entropy E_D, computed from the logits


def predict_proba(model, x):
    """p(t=1|x) for each row of x, with dropout off and batch norm using its running averages."""
    model.eval()
    with torch.no_grad():
        return torch.sigmoid(model(x))


def evaluate(model, x, t):
    """Mean cross-entropy and accuracy of the model on (x, t), in evaluation mode."""
    model.eval()
    with torch.no_grad():
        a = model(x)
        return loss_fn(a, t).item(), ((a > 0).float() == t).float().mean().item()


def fit(model, data, epochs=1000, lr=3e-3, weight_decay=0.0, batch_size=None,
        optimizer='adam', seed=189, callback=None):
    """Train the model on data['x_train'], data['t_train'] with mini-batches of size
    batch_size (None: the full training set in every step).

    weight_decay is the coefficient lambda of the penalty (lambda/2)||w||^2 on all
    parameters (PyTorch adds lambda * w to every gradient). After every epoch the function
    records the training and validation loss and accuracy and calls
    callback(epoch, model, history) if given; training stops early if the callback
    returns True. Returns the history: a dictionary of per-epoch lists.
    """
    torch.manual_seed(seed)                               # dropout masks
    shuffle = torch.Generator().manual_seed(seed)         # mini-batch order
    optimizers = {'adam': torch.optim.Adam, 'adamw': torch.optim.AdamW, 'sgd': torch.optim.SGD}
    opt = optimizers[optimizer](model.parameters(), lr=lr, weight_decay=weight_decay)
    x, t = data['x_train'], data['t_train']
    N = len(x)
    B = batch_size or N
    history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': []}
    for epoch in range(epochs):
        model.train()                                     # dropout on, batch norm uses batch statistics
        perm = torch.randperm(N, generator=shuffle)
        for i in range(0, N, B):
            idx = perm[i:i + B]
            opt.zero_grad()
            loss = loss_fn(model(x[idx]), t[idx])
            loss.backward()
            opt.step()
        train_loss, train_acc = evaluate(model, x, t)
        val_loss, val_acc = evaluate(model, data['x_val'], data['t_val'])
        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['train_acc'].append(train_acc)
        history['val_acc'].append(val_acc)
        if callback is not None and callback(epoch, model, history):
            break
    return history

The cell below plots the data, with the 100 training points drawn large and outlined and the 1000 validation points drawn small. The two classes overlap where the moons meet, so even the best possible classifier misclassifies some of the points, and a model that classifies every training point correctly has also fitted some of the noise.

fig = go.Figure()
x, t = data['x_val'].numpy(), data['t_val'].numpy()
for k in (0, 1):
    fig.add_scatter(x=x[t == k, 0], y=x[t == k, 1], mode='markers', name=f'validation, class {k}',
                    marker=dict(color=CLASS_COLORS[k], size=8, opacity=0.8))
x, t = data['x_train'].numpy(), data['t_train'].numpy()
for k in (0, 1):
    fig.add_scatter(x=x[t == k, 0], y=x[t == k, 1], mode='markers', name=f'training, class {k}',
                    marker=dict(color=CLASS_COLORS[k], size=12, line=dict(width=1, color='black')))
fig.update_layout(title='Two moons with noise 0.3: 100 training and 1000 validation points',
                  xaxis_title='x1', yaxis_title='x2', width=800, height=500)
fig

The over-fitting network

To see the problem that the rest of the notebook addresses, we train the large network on the 100 training points without any regularization (Adam, learning rate $3 \times 10^{-3}$, full-batch, 1000 epochs, about 2 seconds) and compare its loss and accuracy on the training and validation sets at the end of training with the smallest validation loss reached along the way.

torch.manual_seed(189)
overfit_model = MLP(**BIG_MLP)
overfit_history = fit(overfit_model, data, epochs=1000, lr=3e-3)

h = overfit_history
best_epoch = int(np.argmin(h['val_loss']))
print(f"after 1000 epochs:   training loss {h['train_loss'][-1]:.3f}, accuracy {h['train_acc'][-1]:.3f}")
print(f"                     validation loss {h['val_loss'][-1]:.3f}, accuracy {h['val_acc'][-1]:.3f}")
print(f"best epoch ({best_epoch + 1}):    validation loss {h['val_loss'][best_epoch]:.3f}, "
      f"accuracy {h['val_acc'][best_epoch]:.3f}")

after 1000 epochs:   training loss 0.001, accuracy 1.000
                     validation loss 0.886, accuracy 0.874
best epoch (74):    validation loss 0.213, accuracy 0.902

The plot below shows the trained network. In this plot and in the ones that follow, each panel draws the training points on top of the predicted probability $p(t = 1 \mid \mathbf{x})$ of a model, shaded from blue ($p$ near 0) through white to red ($p$ near 1), and the black curve is the decision boundary $p = 0.5$.

plot_boundaries({'no regularization (1000 epochs)': overfit_model}, data)

The network classifies every training point correctly, and its training loss is close to zero. To do so, its boundary bends around individual noisy points and encloses small islands of one class inside the region of the other. These islands fit the noise in the 100 training points and do not hold for new data, so the validation loss at the end of training is about four times its smallest value, which was reached within the first hundred epochs, and the validation accuracy is lower as well. The cross-entropy grows by much more than the error rate because the over-fitted network is confident on the validation points it gets wrong. Each of the following sections adds one inductive bias that prefers simpler functions and compares the result with this network.

(Lecture Demo 2) Weight Decay

Lecture: run from the slide "Weight Decay Regularization".

Weight decay adds a quadratic penalty on the parameters to the data error (Bishop, Eq. 9.1): $$ \widetilde{E}(\mathbf{w}) = E_D(\mathbf{w}) + \frac{\lambda}{2}\mathbf{w}^{\mathrm{T}}\mathbf{w}. $$ The factor $1/2$ cancels in the gradient (Bishop, Eq. 9.5): $$ \nabla \widetilde{E}(\mathbf{w}) = \nabla E_D(\mathbf{w}) + \lambda \mathbf{w}, $$ so a gradient descent step with learning rate $\eta$ is $$ \mathbf{w} \leftarrow \mathbf{w} - \eta\left(\nabla E_D(\mathbf{w}) + \lambda \mathbf{w}\right) = (1 - \eta\lambda)\,\mathbf{w} - \eta \nabla E_D(\mathbf{w}). $$ Every step first shrinks the weights by the factor $1 - \eta\lambda$, which gives the method its name. The argument weight_decay of the PyTorch optimizers is $\lambda$, and the optimizer adds $\lambda \mathbf{w}$ to the gradient of every parameter it updates.

One step of SGD, two ways

Below, model (a) takes one step of torch.optim.SGD(weight_decay=λ) on $E_D$, and model (b), an exact copy, takes one plain SGD step on $E_D + \frac{\lambda}{2}\lVert\mathbf{w}\rVert^2$, with the penalty added to the loss by hand.

lam, eta = 0.1, 0.5
x, t = data['x_train'], data['t_train']

torch.manual_seed(189)
model_a = MLP(width=20, depth=1)
model_b = copy.deepcopy(model_a)

# (a) the optimizer applies the weight decay
opt_a = torch.optim.SGD(model_a.parameters(), lr=eta, weight_decay=lam)
opt_a.zero_grad()
loss_fn(model_a(x), t).backward()
opt_a.step()

# (b) the penalty (lambda/2) ||w||^2 is part of the loss
opt_b = torch.optim.SGD(model_b.parameters(), lr=eta)
penalty = sum((p ** 2).sum() for p in model_b.parameters())
opt_b.zero_grad()
(loss_fn(model_b(x), t) + lam / 2 * penalty).backward()
opt_b.step()

def max_difference(m1, m2):
    return max((p1 - p2).abs().max().item() for p1, p2 in zip(m1.parameters(), m2.parameters()))

print(f"largest difference between the parameters after one step: {max_difference(model_a, model_b):.1e}")

largest difference between the parameters after one step: 7.5e-09

The two updates agree up to floating-point round-off. Two details of the PyTorch convention are worth knowing:

• weight_decay applies to every parameter given to the optimizer, biases included. Bishop (Section 9.2.1) notes that biases are usually left out of the penalty; this is done with two parameter groups, e.g. torch.optim.SGD([{'params': weights, 'weight_decay': lam}, {'params': biases, 'weight_decay': 0.0}], lr=eta).
• For Adam the two forms are still identical, because torch.optim.Adam(weight_decay=λ) also adds $\lambda \mathbf{w}$ to the gradient. The penalty gradient is then divided by Adam's running scale $\sqrt{\hat{v}} + \delta$ like the rest of the gradient, so weights with large gradients are decayed less. AdamW instead decouples the decay from the adaptive step and applies $\mathbf{w} \leftarrow \mathbf{w} - \eta\lambda\mathbf{w}$ directly, which is no longer the gradient of any fixed penalty added to the loss. The check below takes one step of each.

def one_step(opt_name, lam, add_penalty_to_loss=False, lr=0.01):
    torch.manual_seed(189)
    model = MLP(width=20, depth=1)
    opt = {'adam': torch.optim.Adam, 'adamw': torch.optim.AdamW}[opt_name](
        model.parameters(), lr=lr, weight_decay=0.0 if add_penalty_to_loss else lam)
    loss = loss_fn(model(x), t)
    if add_penalty_to_loss:
        loss = loss + lam / 2 * sum((p ** 2).sum() for p in model.parameters())
    opt.zero_grad()
    loss.backward()
    opt.step()
    return model

lam = 0.5
reference = one_step('adam', lam, add_penalty_to_loss=True)
print(f"Adam(weight_decay=λ)  vs  Adam on E_D + λ/2 ||w||^2: {max_difference(one_step('adam', lam), reference):.1e}")
print(f"AdamW(weight_decay=λ) vs  Adam on E_D + λ/2 ||w||^2: {max_difference(one_step('adamw', lam), reference):.1e}")

Adam(weight_decay=λ)  vs  Adam on E_D + λ/2 ||w||^2: 0.0e+00
AdamW(weight_decay=λ) vs  Adam on E_D + λ/2 ||w||^2: 2.0e-02

Weight decay on the over-fitting network

The cell below trains the large network on the 100 training points with four values of $\lambda$ (Adam, learning rate $3 \times 10^{-3}$, full-batch, 1000 epochs, about 2 seconds per model), so the run with $\lambda = 0$ repeats the over-fitting network of the Setup section. Because torch.optim.Adam(weight_decay=λ) adds the gradient of $\frac{\lambda}{2}\lVert\mathbf{w}\rVert^2$, each run minimizes $\widetilde{E}$ for its $\lambda$.

lambdas = [0.0, 1e-3, 3e-3, 3e-2]
wd_models, wd_histories = {}, {}
for lam in lambdas:
    torch.manual_seed(189)
    model = MLP(**BIG_MLP)
    wd_histories[f'λ = {lam:g}'] = fit(model, data, epochs=1000, lr=3e-3, weight_decay=lam)
    wd_models[f'λ = {lam:g}'] = model


def param_norm(model):
    """The Euclidean norm ||w|| of all parameters of the model."""
    return torch.sqrt(sum((p ** 2).sum() for p in model.parameters())).item()


print(f"{'':12s} {'train loss':>10s} {'val loss':>9s} {'train acc':>10s} {'val acc':>8s} {'||w||':>7s}")
for name, h in wd_histories.items():
    print(f"{name:12s} {h['train_loss'][-1]:10.3f} {h['val_loss'][-1]:9.3f} "
          f"{h['train_acc'][-1]:10.3f} {h['val_acc'][-1]:8.3f} {param_norm(wd_models[name]):7.1f}")

             train loss  val loss  train acc  val acc   ||w||
λ = 0             0.001     0.886      1.000    0.874    62.5
λ = 0.001         0.066     0.364      0.990    0.876    14.1
λ = 0.003         0.216     0.231      0.900    0.903     7.1
λ = 0.03          0.392     0.354      0.840    0.853     2.8

plot_boundaries(wd_models, data)

plot_history(wd_histories, title='Learning curves for four values of the weight decay λ')

Without weight decay ($\lambda = 0$) we recover the over-fitted network of the Setup section. Small values of $\lambda$ keep $\lVert\mathbf{w}\rVert$ small and the boundary smooth, with a lower validation loss and a higher validation accuracy. The largest value, $\lambda = 0.03$, under-fits, as the training loss stays high and the boundary is too simple. The coefficient $\lambda$ is a hyperparameter, chosen with the validation set.

Learning Curves and Early Stopping

Slides: "Learning Curves and Early Stopping" and "Early Stopping is a form of Weight Decay".

A learning curve plots the training and validation error against the iteration number (Bishop, Section 9.3). The training error decreases broadly monotonically, while the validation error typically decreases at first and then increases as the network starts to over-fit. Early stopping (Bishop, Section 9.3.1) returns the parameters from the epoch with the smallest validation error: $$ \tau^\star = \arg\min_{\tau} E_{\text{val}}\left(\mathbf{w}^{(\tau)}\right), \qquad \text{return } \mathbf{w}^{(\tau^\star)}. $$ In code this means keeping a copy of the best parameters (a checkpoint) during training. model.state_dict() returns references to the live parameter tensors, so the copy must be a deep copy, as otherwise the "checkpoint" would change with every later step.

torch.manual_seed(189)
model = MLP(**BIG_MLP)

best = {'val_loss': float('inf')}
norms = []

def checkpoint(epoch, model, history):
    norms.append(param_norm(model))
    if history['val_loss'][-1] < best['val_loss']:
        best['val_loss'] = history['val_loss'][-1]
        best['epoch'] = epoch
        best['state'] = copy.deepcopy(model.state_dict())

history = fit(model, data, epochs=1000, lr=3e-3, callback=checkpoint)
print(f"last epoch:  validation loss {history['val_loss'][-1]:.3f}, accuracy {history['val_acc'][-1]:.3f}")

model.load_state_dict(best['state'])     # restore the checkpoint
val_loss, val_acc = evaluate(model, data['x_val'], data['t_val'])
print(f"best epoch {best['epoch'] + 1}: validation loss {val_loss:.3f}, accuracy {val_acc:.3f} (restored model)")
early_stopped = model

last epoch:  validation loss 0.886, accuracy 0.874
best epoch 74: validation loss 0.213, accuracy 0.902 (restored model)

plot_history({'no regularization': history}, mark={'no regularization': best['epoch']},
             title='Learning curves of the large network (x: the early-stopping epoch)')

The validation loss reaches its minimum within the first hundred epochs and then rises steadily while the training loss continues toward zero. The validation accuracy changes much less than the validation loss, as most of the increase in the loss comes from confident mistakes. The restored model is the network at the marked epoch. Since the validation set was used to choose the epoch, its validation accuracy is a slightly optimistic estimate of the test accuracy, although with 1000 validation points the bias is small.

In practice training is stopped once the validation loss has not improved for a number of epochs called the patience. Here the callback returns True to stop training after 50 epochs without improvement:

def early_stopping(patience):
    best = {'val_loss': float('inf'), 'epoch': 0}
    def callback(epoch, model, history):
        if history['val_loss'][-1] < best['val_loss']:
            best.update(val_loss=history['val_loss'][-1], epoch=epoch,
                        state=copy.deepcopy(model.state_dict()))
        return epoch - best['epoch'] >= patience      # True stops the training loop
    return best, callback

torch.manual_seed(189)
model = MLP(**BIG_MLP)
best_p, callback = early_stopping(patience=50)
history_p = fit(model, data, epochs=1000, lr=3e-3, callback=callback)
model.load_state_dict(best_p['state'])
print(f"stopped after {len(history_p['val_loss'])} epochs; restored epoch {best_p['epoch'] + 1}, "
      f"validation accuracy {evaluate(model, data['x_val'], data['t_val'])[1]:.3f}")

stopped after 124 epochs; restored epoch 74, validation accuracy 0.902

Early stopping as weight decay

The network starts with small weights, and the norm $\lVert\mathbf{w}\rVert$ grows during training. Stopping early therefore returns a weight vector with a smaller norm, similar to the effect of weight decay (Bishop, Figure 9.8). For a quadratic error the product $\tau\eta$ of the number of steps and the learning rate plays the role of $1/\lambda$ (Bishop, Exercise 9.6). The plot shows $\lVert\mathbf{w}\rVert$ during the run above, next to the norms reached at the end of training with weight decay.

fig = make_subplots(rows=1, cols=2, subplot_titles=['validation loss', 'parameter norm ||w||'])
epochs = np.arange(1, len(norms) + 1)
fig.add_scatter(x=epochs, y=history['val_loss'], name='validation loss', row=1, col=1)
fig.add_scatter(x=epochs, y=norms, name='||w|| (λ = 0)', row=1, col=2)
for name in [n for n in wd_models if n != 'λ = 0']:
    fig.add_hline(y=param_norm(wd_models[name]), line_dash='dot', row=1, col=2,
                  annotation_text=f'{name} (end of training)', annotation_position='top right',
                  annotation_font_size=10)
for col in (1, 2):
    fig.add_vline(x=best['epoch'] + 1, line_dash='dash', line_color='black', row=1, col=col)
fig.update_xaxes(title_text='epoch')
fig.update_layout(width=1000, height=400, showlegend=False,
                  title=f"The early-stopping epoch {best['epoch'] + 1} (dashed) has ||w|| = {norms[best['epoch']]:.1f}")
fig

plot_boundaries({f"early stopping (epoch {best['epoch'] + 1})": early_stopped,
                 'λ = 0.003 (1000 epochs)': wd_models['λ = 0.003'],
                 'no regularization (1000 epochs)': wd_models['λ = 0']}, data)

Ensembles and Bagging

Slides: "Ensembles of Models (Experts)" and "Bagging: Bootstrap Aggregation".

An ensemble (committee) averages the predictive distributions of $L$ models (Bishop, Eq. 9.42): $$ p(t = 1 \mid \mathbf{x}) = \frac{1}{L}\sum_{l=1}^{L} p_l(t = 1 \mid \mathbf{x}). $$ Bagging (Breiman, 1996) trains model $l$ on a bootstrap data set: $N$ points drawn from the training set uniformly with replacement. A bootstrap set contains on average a fraction $1 - (1 - 1/N)^N \approx 1 - 1/e \approx 0.632$ of the distinct training points, some of them several times.

The cell below trains $L = 20$ copies of the large network, each on its own bootstrap sample, for 300 epochs (precomputed for the lecture: about 10 seconds).

def bootstrap_indices(N, generator):
    """N indices drawn uniformly with replacement from 0, ..., N-1."""
    return torch.randint(0, N, (N,), generator=generator)


L_max = 20
N = len(data['x_train'])
members, bootstrap_sets = [], []
sampler = torch.Generator().manual_seed(189)
start = time.time()
for l in range(L_max):
    idx = bootstrap_indices(N, sampler)
    bootstrap_sets.append(idx)
    data_l = dict(data, x_train=data['x_train'][idx], t_train=data['t_train'][idx])
    torch.manual_seed(189 + l)
    model = MLP(**BIG_MLP)
    fit(model, data_l, epochs=300, lr=3e-3, seed=189 + l)
    members.append(model)
print(f"trained {L_max} models in {time.time() - start:.1f} s")
print(f"average fraction of distinct training points per bootstrap set: "
      f"{np.mean([len(idx.unique()) / N for idx in bootstrap_sets]):.3f}")

trained 20 models in 7.5 s
average fraction of distinct training points per bootstrap set: 0.635

x_val, t_val = data['x_val'], data['t_val']
member_probs = torch.stack([predict_proba(m, x_val) for m in members])   # shape (L, number of points)


def cross_entropy(p, t, eps=1e-7):
    """Mean binary cross-entropy of predicted probabilities p. The clamp caps the loss of a
    confidently wrong prediction at -ln(1e-7) = 16.1, so it can differ slightly from the
    loss computed from logits in fit()."""
    p = p.clamp(eps, 1 - eps)
    return -(t * p.log() + (1 - t) * (1 - p).log()).mean().item()


def accuracy(p, t):
    return ((p > 0.5).float() == t).float().mean().item()


single = [cross_entropy(p, t_val) for p in member_probs]
print(f"single bagged model (mean over {L_max}): validation loss {np.mean(single):.3f}, "
      f"accuracy {np.mean([accuracy(p, t_val) for p in member_probs]):.3f}")
for L in [1, 2, 5, 10, 20]:
    p_avg = member_probs[:L].mean(dim=0)
    print(f"average of L = {L:2d} models: validation loss {cross_entropy(p_avg, t_val):.3f}, "
          f"accuracy {accuracy(p_avg, t_val):.3f}")

single bagged model (mean over 20): validation loss 0.743, accuracy 0.860
average of L =  1 models: validation loss 1.125, accuracy 0.867
average of L =  2 models: validation loss 0.457, accuracy 0.879
average of L =  5 models: validation loss 0.346, accuracy 0.875
average of L = 10 models: validation loss 0.311, accuracy 0.886
average of L = 20 models: validation loss 0.296, accuracy 0.891

ensemble = lambda x: torch.stack([predict_proba(m, x) for m in members]).mean(dim=0)
plot_boundaries({'model 1': members[0], 'model 2': members[1], 'model 3': members[2],
                 f'average of {L_max} models': ensemble}, data)

Each member over-fits its own bootstrap sample, and the boundaries of different members differ in where they bend around noisy points. Averaging cancels much of this variation, so the averaged boundary is smoother and the validation loss of the ensemble is far below that of a typical member.

How fast does averaging help?

Bishop (Section 9.6) analyzes the committee with the sum-of-squares error. Writing each model's prediction as the target plus an error, $y_l(\mathbf{x}) = h(\mathbf{x}) + \epsilon_l(\mathbf{x})$, the average error of the individual models is $E_{\text{AV}} = \frac{1}{L}\sum_l \mathbb{E}[\epsilon_l^2]$ and the error of the committee is $E_{\text{COM}} = \mathbb{E}\big[\big(\frac{1}{L}\sum_l \epsilon_l\big)^2\big]$. If the errors had zero mean and were uncorrelated, then $E_{\text{COM}} = E_{\text{AV}} / L$ (Bishop, Eq. 9.50). In general only $E_{\text{COM}} \le E_{\text{AV}}$ holds.

The plot measures both on the validation set with the squared error between the predicted probability and the label $t$ (the Brier score), for committees of the first $L$ members.

sq_err = (member_probs - t_val) ** 2                    # (L_max, points)
E_AV = [sq_err[:L].mean().item() for L in range(1, L_max + 1)]
E_COM = [((member_probs[:L].mean(dim=0) - t_val) ** 2).mean().item() for L in range(1, L_max + 1)]
Ls = np.arange(1, L_max + 1)

fig = go.Figure()
fig.add_scatter(x=Ls, y=E_COM, mode='lines+markers', name='E_COM: error of the average of L models')
fig.add_scatter(x=Ls, y=E_AV, mode='lines', name='E_AV: average error of the L models')
fig.add_scatter(x=Ls, y=E_AV[-1] / Ls, mode='lines', line_dash='dot', name='E_AV / L (uncorrelated errors)')
fig.update_layout(xaxis_title='number of models L', yaxis_title='mean squared error of p(t=1|x)',
                  title='Committee error versus committee size', width=800, height=450)
fig

The committee error drops quickly for the first few members and then levels off far above the $E_{\text{AV}}/L$ curve. The errors of the members are strongly correlated, as they are trained on overlapping bootstrap samples of the same 100 points, they share the architecture and the training procedure, and they all make the same mistakes on validation points that lie on the wrong side of the true boundary because of the label noise. Averaging removes only the part of the error that varies between members.

(Lecture Demo 3) Dropout

Lecture: run from the slide "Implementing Drop-out in PyTorch".

Dropout (Srivastava et al., 2014; Bishop, Section 9.6.1) multiplies each non-output activation $z_i$ by a random mask $R_i \in \{0, 1\}$, drawn anew for every data point and every training step.

Two conventions for the probability. Bishop writes $\rho$ for the probability of keeping a node ($R_i = 1$ with probability $\rho$; typical values $\rho = 0.5$ for hidden units and $\rho = 0.8$ for inputs). The slides and PyTorch's nn.Dropout(p) use $p$, the probability of dropping a node, so $p = 1 - \rho$. The code below uses $p$.

PyTorch implements inverted dropout. During training the kept activations are divided by $1 - p$, and in evaluation mode the layer is the identity, $$ \text{training: } \tilde{z}_i = \frac{R_i\, z_i}{1 - p}, \quad R_i \sim \mathrm{Bernoulli}(1 - p); \qquad \text{evaluation: } \tilde{z}_i = z_i. $$ The scaling makes the expected activation during training equal to the activation used at test time: $$ \mathbb{E}[\tilde{z}_i] = \frac{(1 - p)\, z_i}{1 - p} = z_i. $$ Bishop describes the equivalent original form, which does not scale during training and instead multiplies the outgoing weights of each node by $\rho = 1 - p$ at test time.

The code from the slide (seeded, so the mask is the same every time it runs):

torch.manual_seed(2026)
dropout = nn.Dropout(p=0.3)

dropout.train()
print(dropout(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])))

dropout.eval()
print(dropout(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])))

tensor([[ 1.4286,  0.0000,  4.2857],
        [ 5.7143,  7.1429,  8.5714],
        [ 0.0000, 11.4286, 12.8571]])
tensor([[1., 2., 3.],
        [4., 5., 6.],
        [7., 8., 9.]])

In training mode each entry is either $0$ or the input divided by $1 - p = 0.7$ (for example $1/0.7 = 1.4286$ and $9/0.7 = 12.8571$). In evaluation mode the input passes through unchanged. Averaging the training-mode output over many independent masks recovers the input, as the expectation above states:

z = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
dropout.train()
torch.manual_seed(189)
samples = torch.stack([dropout(z) for _ in range(100_000)])
kept = samples != 0
print("fraction of entries kept:", round(kept.float().mean().item(), 4), "(1 - p = 0.7)")
print("kept values equal z / (1 - p):", torch.allclose(samples[kept], (z / 0.7).expand_as(samples)[kept]))
print("average over 100,000 masks:")
print(samples.mean(dim=0))

fraction of entries kept: 0.6999 (1 - p = 0.7)
kept values equal z / (1 - p): True
average over 100,000 masks:
tensor([[1.0000, 2.0004, 2.9970],
        [4.0047, 4.9848, 5.9995],
        [7.0102, 7.9825, 9.0224]])

Dropout on the over-fitting network

The same network and training run as in the weight decay section, with dropout after each hidden layer ($p = 0.5$, Bishop's $\rho = 0.5$) and no weight decay. Each epoch now trains a different random sub-network on every data point, and the evaluation uses the full network (model.eval() inside evaluate).

dropout_models, dropout_histories = {}, {}
for p in [0.0, 0.5]:
    torch.manual_seed(189)
    model = MLP(**BIG_MLP, dropout=p)
    dropout_histories[f'p = {p}'] = fit(model, data, epochs=1000, lr=3e-3)
    dropout_models[f'p = {p}'] = model
    h = dropout_histories[f'p = {p}']
    print(f"dropout p = {p}: final train loss {h['train_loss'][-1]:.3f}, val loss {h['val_loss'][-1]:.3f}, "
          f"val accuracy {h['val_acc'][-1]:.3f}; smallest val loss {min(h['val_loss']):.3f}")
plot_history(dropout_histories, title='Learning curves with and without dropout')

dropout p = 0.0: final train loss 0.001, val loss 0.886, val accuracy 0.874; smallest val loss 0.213

dropout p = 0.5: final train loss 0.132, val loss 0.302, val accuracy 0.895; smallest val loss 0.214

plot_boundaries({'no dropout': dropout_models['p = 0.0'], 'dropout p = 0.5': dropout_models['p = 0.5']}, data)

Dropout slows over-fitting: after 1000 epochs the validation loss with dropout is a fraction of the loss without it, and the boundary is smoother. It does not remove over-fitting entirely, since the validation loss still creeps up late in training, and the smallest validation loss along the way is about the same in both runs. In this small problem dropout is therefore best combined with early stopping or weight decay, which is also common practice (Bishop notes that weight decay is often used alongside dropout).

Monte Carlo dropout

The full network in evaluation mode approximates the average over the exponentially many sub-networks. The average can also be estimated directly by sampling masks at test time (Bishop, Eq. 9.51): $$ p(t \mid \mathbf{x}) = \sum_{\mathbf{R}} p(\mathbf{R})\, p(t \mid \mathbf{x}, \mathbf{R}) \approx \frac{1}{S}\sum_{s=1}^{S} p(t \mid \mathbf{x}, \mathbf{R}^{(s)}), $$ with $S = 20$ masks below. Calling model.train() turns the dropout layers back on (this network has no batch normalization, which train() would also switch to batch statistics).

def mc_dropout_proba(model, x, S=20, seed=189):
    """Average of p(t=1|x) over S random dropout masks (Bishop, Eq. 9.51)."""
    torch.manual_seed(seed)
    model.train()                        # dropout on
    with torch.no_grad():
        p = torch.stack([torch.sigmoid(model(x)) for _ in range(S)]).mean(dim=0)
    model.eval()
    return p

model = dropout_models['p = 0.5']
p_eval = predict_proba(model, x_val)
p_mc = mc_dropout_proba(model, x_val, S=20)
print(f"evaluation mode (rescaled full network): val loss {cross_entropy(p_eval, t_val):.3f}, accuracy {accuracy(p_eval, t_val):.3f}")
print(f"Monte Carlo dropout, S = 20 masks:       val loss {cross_entropy(p_mc, t_val):.3f}, accuracy {accuracy(p_mc, t_val):.3f}")
plot_boundaries({'evaluation mode': model, 'Monte Carlo dropout (S = 20)': lambda x: mc_dropout_proba(model, x)}, data)

evaluation mode (rescaled full network): val loss 0.286, accuracy 0.895
Monte Carlo dropout, S = 20 masks:       val loss 0.268, accuracy 0.896

The two predictions have nearly the same decision boundary and accuracy, and Monte Carlo dropout gives a slightly lower validation loss. The speckle in the right panel is Monte Carlo noise, as every grid point averages its own 20 random masks. The cost is $S$ forward passes per prediction.

(Lecture Demo 4) Residual Connections

Lecture: run from the slide "Expanding the Residual Network Graph" (the gradient plot also goes with "Vanishing Gradients in Deep Networks").

A plain network computes $\mathbf{z}^{(l)} = h\big(\mathbf{W}^{(l)}\mathbf{z}^{(l-1)} + \mathbf{b}^{(l)}\big)$. A residual block adds its input back onto the output of a learned function $F_l$ (Bishop, Eqs. 9.35 to 9.37): $$ \mathbf{z}^{(l)} = \mathbf{z}^{(l-1)} + F_l\big(\mathbf{z}^{(l-1)}\big). $$ Following the slide we use $F_l(\mathbf{z}) = \mathbf{W}^{(l)}\,\mathrm{ReLU}(\mathbf{z}) + \mathbf{b}^{(l)}$, with the activation before the linear map (Bishop, Figure 9.16(b)), so that $F_l$ can output negative as well as positive values. A first linear layer maps the input $\mathbf{x} \in \mathbb{R}^2$ to the width of the blocks, playing the role of the matrix $\mathbf{W}$ in Bishop, Eq. 9.41.

By the chain rule the Jacobian of a block is the identity plus the Jacobian of $F_l$, $$ \frac{\partial \mathbf{z}^{(l)}}{\partial \mathbf{z}^{(l-1)}} = \mathbf{I} + \frac{\partial F_l}{\partial \mathbf{z}^{(l-1)}}, $$ so the error signal reaching layer $l - 1$ always contains the term passed through the identity unchanged. In a plain network the signal is multiplied by $\mathrm{diag}(h')\,\mathbf{W}^{(l)}$ at every layer and can shrink geometrically with depth.

class ResidualBlock(nn.Module):
    """z -> z + F(z) with F(z) = W Dropout(ReLU(Norm(z))) + b."""

    def __init__(self, width, activation='relu', norm='none', dropout=0.0):
        super().__init__()
        self.F = nn.Sequential(make_norm(norm, width), ACTIVATIONS[activation](),
                               nn.Dropout(p=dropout), nn.Linear(width, width))

    def forward(self, z):
        return z + self.F(z)


class ResidualNetwork(nn.Module):
    """Linear input layer, `depth` residual blocks of `width` units, and a linear output layer.
    The arguments match MLP, so the two can be swapped."""

    def __init__(self, n_in=2, width=200, depth=2, activation='relu', norm='none', dropout=0.0):
        super().__init__()
        self.input = nn.Linear(n_in, width)
        self.blocks = nn.Sequential(*[ResidualBlock(width, activation, norm, dropout) for _ in range(depth)])
        self.output = nn.Sequential(make_norm(norm, width), ACTIVATIONS[activation]())
        self.head = nn.Linear(width, 1)

    def forward(self, x):
        return self.head(self.output(self.blocks(self.input(x)))).squeeze(-1)

A check of the block equation, and of the expanded form of three blocks (Bishop, Eq. 9.40, the slide "Expanding the Residual Network Graph"), in which the output is a sum of paths through subsets of the blocks: $$ \mathbf{z}^{(3)} = F_3\big(F_2(F_1(\mathbf{x}) + \mathbf{x}) + F_1(\mathbf{x}) + \mathbf{x}\big) + F_2\big(F_1(\mathbf{x}) + \mathbf{x}\big) + F_1(\mathbf{x}) + \mathbf{x}. $$

torch.manual_seed(189)
F1, F2, F3 = [ResidualBlock(width=4).F for _ in range(3)]
blocks = nn.Sequential(*[ResidualBlock(width=4) for _ in range(3)])
for block, F in zip(blocks, (F1, F2, F3)):
    block.F = F
x0 = torch.randn(5, 4)
with torch.no_grad():
    z3_blocks = blocks(x0)
    z3_expanded = F3(F2(F1(x0) + x0) + F1(x0) + x0) + F2(F1(x0) + x0) + F1(x0) + x0
    block = blocks[0]
    print("block(z) == z + W ReLU(z) + b:", torch.allclose(block(x0), x0 + torch.relu(x0) @ block.F[-1].weight.T + block.F[-1].bias))
print("three blocks == expanded form (Bishop, Eq. 9.40):", torch.allclose(z3_blocks, z3_expanded))

block(z) == z + W ReLU(z) + b: True
three blocks == expanded form (Bishop, Eq. 9.40): True

Gradients at initialization

Take a plain network and a residual network with 30 hidden layers of 32 ReLU units each, both with PyTorch's default initialization, and compute the norm of the gradient of $E_D$ with respect to each weight matrix on the training set. The default initialization draws weights from $\mathcal{U}(-1/\sqrt{d}, 1/\sqrt{d})$ for fan-in $d$, a variance of $1/(3d)$, which is smaller than the He initialization $2/d$ for ReLU from Lecture 14. The plot also shows the plain network with He initialization.

def gradient_norms(model, data):
    """||dE_D/dW|| for the weight matrix of every linear layer, from input to output."""
    model.train()
    model.zero_grad()
    loss_fn(model(data['x_train']), data['t_train']).backward()
    return [m.weight.grad.norm().item() for m in model.modules() if isinstance(m, nn.Linear)]


def he_init(model):
    for m in model.modules():
        if isinstance(m, nn.Linear):
            nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
            nn.init.zeros_(m.bias)
    return model


DEEP = dict(width=32, depth=30)
def make_deep(kind):
    torch.manual_seed(189)
    if kind == 'plain':
        return MLP(**DEEP)
    if kind == 'plain, He init':
        return he_init(MLP(**DEEP))
    return ResidualNetwork(**DEEP)

kinds = ['plain', 'plain, He init', 'residual']
fig = go.Figure()
for kind in kinds:
    g = gradient_norms(make_deep(kind), data)
    fig.add_scatter(x=np.arange(1, len(g) + 1), y=g, mode='lines+markers', name=kind)
    print(f"{kind:15s} first layer {g[0]:.1e}   last hidden layer {g[-2]:.1e}   output layer {g[-1]:.1e}")
fig.update_yaxes(type='log', title='gradient norm ||dE/dW||', exponentformat='power')
fig.update_xaxes(title='linear layer (1 = input side)')
fig.update_layout(title='Gradient norms at initialization, 30 hidden layers of 32 units', width=900, height=450)
fig

plain           first layer 6.9e-13   last hidden layer 5.3e-04   output layer 1.0e-03
plain, He init  first layer 6.3e-01   last hidden layer 3.7e+00   output layer 2.9e+00
residual        first layer 1.3e+00   last hidden layer 3.8e+00   output layer 9.0e+00

In the plain network with the default initialization the gradient shrinks by roughly a constant factor per layer, and the first layer's gradient is about ten orders of magnitude smaller than the last one's, so the early layers do not learn. He initialization keeps the gradient scale roughly constant at initialization, as derived in Lecture 14. The residual network has gradients of similar size in all layers even with the default initialization, because every block passes the error signal through its identity path.

Training the deep networks

Each network trains for 300 epochs with Adam on the 100 training points (about 5 seconds in total).

deep_models, deep_histories = {}, {}
for kind in kinds:
    model = make_deep(kind)
    deep_histories[kind] = fit(model, data, epochs=300, lr=1e-3)
    deep_models[kind] = model
    h = deep_histories[kind]
    print(f"{kind:15s} final train loss {h['train_loss'][-1]:.3f}, val accuracy {h['val_acc'][-1]:.3f}")
plot_history(deep_histories, title='Training 30-layer networks')

plain           final train loss 0.693, val accuracy 0.500

plain, He init  final train loss 0.077, val accuracy 0.845

residual        final train loss 0.001, val accuracy 0.854

plot_boundaries(deep_models, data)

The plain network with the default initialization does not train at all, as its loss stays at $\ln 2 \approx 0.693$ and it predicts $p \approx 0.5$ everywhere. With He initialization the plain network trains at this depth, and so does the residual network with the default initialization. Both then over-fit the 100 training points (the validation loss rises after about 50 epochs), and the spikes in their curves are single steps that temporarily increase the loss, which deep networks trained with a fixed learning rate show occasionally. On ImageNet, He et al. (2016) found that a 34-layer plain network with batch normalization reached a higher training error than an 18-layer one, while residual networks improved with depth up to 152 layers. The residual connection makes the result much less sensitive to such details, which is why it is part of nearly every deep architecture.

(Lecture Demo 5) The Model Explorer

Lecture: run from the final "Demo" slide.

The explorer trains a model with the chosen settings and shows the decision boundary together with the training and validation loss and accuracy, updated live every 10 epochs. Press Run Interact to train. Every setting combines the pieces of this and the previous lecture: the data set and its size, MLP or residual network, depth and width, activation, normalization (Demo 1), dropout, learning rate, weight decay, batch size, and epochs. Training uses Adam with weight_decay as the L2 penalty, and the seeds are fixed, so the same settings give the same result. Suggested experiments:

• Moons with 100 points: increase the width and the epochs until the model over-fits, then add weight decay or dropout.
• Spiral: compare a shallow and a deep MLP, then a deep MLP with a ResNet of the same depth.
• Deep networks (15 or more layers) with tanh, with and without batch or layer normalization.
• Batch size 8 versus the full training set (100) at the same number of epochs.

def run_explorer(dataset='moons', n_train=100, model_type='MLP', depth=3, width=64,
                 activation='relu', norm='none', dropout=0.0, lr=3e-3, weight_decay=0.0,
                 batch_size=32, epochs=200, live=True, update_every=10):
    """Train one model with the given settings and plot the result.
    With live=True the figure is a FigureWidget that is updated during training."""
    data = make_datasets(dataset, n_train=n_train)
    torch.manual_seed(189)
    Model = MLP if model_type == 'MLP' else ResidualNetwork
    model = Model(width=width, depth=depth, activation=activation, norm=norm, dropout=dropout)
    empty = {k: [] for k in ['train_loss', 'val_loss', 'train_acc', 'val_acc']}
    fig = explorer_figure(data, model, empty)
    if live:
        fig = go.FigureWidget(fig)
        display(fig)

    def show_progress(epoch, model, history):
        if live and (epoch + 1) % update_every == 0:
            update_explorer_figure(fig, data, model, history)

    start = time.time()
    history = fit(model, data, epochs=epochs, lr=lr, weight_decay=weight_decay,
                  batch_size=batch_size, callback=show_progress)
    update_explorer_figure(fig, data, model, history)
    if not live:
        fig.show()
    print(f"{count_parameters(model):,} parameters, trained in {time.time() - start:.1f} s.  "
          f"train loss {history['train_loss'][-1]:.3f}, accuracy {history['train_acc'][-1]:.3f} | "
          f"validation loss {history['val_loss'][-1]:.3f}, accuracy {history['val_acc'][-1]:.3f}")
    return model, history

from ipywidgets import interact_manual, fixed, Dropdown, IntSlider, FloatSlider, SelectionSlider

style = {'description_width': '110px'}
interact_manual(
    run_explorer,
    dataset=Dropdown(options=['moons', 'circles', 'xor', 'spiral'], value='moons', description='data set', style=style),
    n_train=SelectionSlider(options=[50, 100, 200, 400], value=100, description='training points', style=style),
    model_type=Dropdown(options=['MLP', 'ResNet'], value='MLP', description='model', style=style),
    depth=IntSlider(min=1, max=30, value=3, description='hidden layers', style=style),
    width=SelectionSlider(options=[4, 8, 16, 32, 64, 128, 256], value=64, description='units per layer', style=style),
    activation=Dropdown(options=['relu', 'tanh', 'gelu'], value='relu', description='activation', style=style),
    norm=Dropdown(options=['none', 'batch', 'layer'], value='none', description='normalization', style=style),
    dropout=FloatSlider(min=0, max=0.8, step=0.1, value=0.0, description='dropout p', style=style),
    lr=SelectionSlider(options=[1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1], value=3e-3, description='learning rate', style=style),
    weight_decay=SelectionSlider(options=[0.0, 1e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1], value=0.0, description='weight decay λ', style=style),
    batch_size=SelectionSlider(options=[8, 16, 32, 64, 400], value=32, description='batch size', style=style),
    epochs=SelectionSlider(options=[50, 100, 200, 300, 500, 1000], value=200, description='epochs', style=style),
    live=fixed(True), update_every=fixed(10));

A batch size of 400 means the full training set for every training-set size offered. For readers of the exported HTML version, where the widget does not run, the cell below trains four fixed settings on the spiral data with 400 training points (ReLU, learning rate $3 \times 10^{-3}$, batch size 32, 200 epochs; about 15 seconds in total). The one-layer network is too simple for the spiral, four layers fit it, and at 20 layers the plain network with the default initialization fails to train while the residual network of the same depth fits the spiral.

settings = {
    'MLP, 1 layer of 16':              dict(model_type='MLP', depth=1, width=16),
    'MLP, 4 layers of 64':             dict(model_type='MLP', depth=4, width=64),
    'MLP, 20 layers of 64':            dict(model_type='MLP', depth=20, width=64),
    'ResNet, 20 blocks of 64':         dict(model_type='ResNet', depth=20, width=64),
}
spiral = make_datasets('spiral', n_train=400)
fixed_models, fixed_histories = {}, {}
for name, s in settings.items():
    torch.manual_seed(189)
    Model = MLP if s['model_type'] == 'MLP' else ResidualNetwork
    model = Model(width=s['width'], depth=s['depth'])
    fixed_histories[name] = fit(model, spiral, epochs=200, lr=3e-3, batch_size=32)
    fixed_models[name] = model
    print(f"{name:32s} validation accuracy {fixed_histories[name]['val_acc'][-1]:.3f}")
plot_boundaries(fixed_models, spiral, cols=4, height=330)

MLP, 1 layer of 16               validation accuracy 0.700

MLP, 4 layers of 64              validation accuracy 0.992

MLP, 20 layers of 64             validation accuracy 0.480

ResNet, 20 blocks of 64          validation accuracy 0.992