Module 2: Neural networks and backpropagation

Module 01 ended with linear models on hand-made features, f(\mathbf{x}) = \mathbf{w}^\top\boldsymbol{\psi}(\mathbf{x}). (Module 01 wrote the feature map as \boldsymbol{\phi}; this module keeps \phi for the activation function.) Such a model works when someone knows \boldsymbol{\psi}: for a disc inside a ring, \boldsymbol{\psi}(\mathbf{x}) = (x_1^2, x_2^2, x_1x_2) makes the two classes linearly separable. For an image or a vibration spectrum nobody can write \boldsymbol{\psi} down. A neural network makes the feature map part of the model and learns it from the data:

with \mathbf{W} and \mathbf{b} trained together with \mathbf{w}, by the same gradient descent.

The multilayer perceptron

The multilayer perceptron (MLP) stacks layers of this kind. This module writes it as

with \mathbf{W}^{(l)} \in \mathbb{R}^{d_{l-1}\times d_l} (one row per input of the layer, one column per unit) and \mathbf{b}^{(l)} \in \mathbb{R}^{d_l}. This convention holds everywhere below. The activation function \phi is a fixed nonlinear function applied to each component separately (Section 5 compares the choices). The last layer has no \phi: its output \mathbf{z}^{(L)} is the prediction in regression and the vector of logits in classification, and the loss consumes it — squared error, or softmax cross-entropy (Module 01, Section 5). The vocabulary: \mathbf{x} is the input; layers 1 to L-1 are hidden layers, and each component of \mathbf{h}^{(l)} is a hidden unit; layer L is the output layer; d_l is the width of layer l; L, the number of weight layers, is the depth; and \mathbf{z}^{(l)} is a pre-activation. Figure 2.1 labels each of them on a small network.

Why the nonlinearity is necessary

Without \phi, depth buys nothing. Two layers give

using \mathbf{B}^\top\mathbf{A}^\top = (\mathbf{A}\mathbf{B})^\top: one affine map. By induction any stack of affine layers collapses to one, so a deep network without activations is a linear model with redundant parameters.

What the universal approximation theorem says

With \phi, one hidden layer is enough in principle. The universal approximation theorem: let \phi be continuous and not a polynomial. For every continuous f on a compact set K \subset \mathbb{R}^d and every \varepsilon > 0 there are a finite width N and weights such that the one-hidden-layer network g(\mathbf{x}) = \sum_{j=1}^{N} a_j\,\phi(\mathbf{w}_j^\top\mathbf{x} + b_j) + c satisfies

Cybenko (1989) proved it for sigmoids and Hornik (1991) for bounded, non-constant activations; Leshno et al. (1993) showed that “not a polynomial” is exactly the condition, which admits ReLU. The exception is easy to see: if \phi is a polynomial of degree p, every such g is a polynomial of degree at most p, a fixed family that cannot come arbitrarily close to \sin 3x.

The theorem says less than it seems to. It does not say how large N must be: the constructive proofs in effect tile K with a grid, and a grid of spacing h in d dimensions has about h^{-d} cells, so the unit count can grow exponentially with d. It does not say that gradient descent from a random start finds such weights. And it does not say that a network fitted to finitely many samples generalises. The rest of this module is about the second question; the evaluation discipline of Module 01, Section 10 is about the third.

The theorem made constructive in one dimension

In one dimension the weights can be written down. Take knots a = x_0 < x_1 < \dots < x_K = b and let p be the piecewise-linear interpolant of f, with slope s_k = \big(f(x_{k+1}) - f(x_k)\big)/(x_{k+1} - x_k) on segment k. On [a, b],

On the first segment only the first hinge is active, so p starts at f(x_0) with slope s_0; each later hinge switches on at its knot and changes the slope by exactly c_k. This is a one-hidden-layer ReLU network with K units: input weights 1, biases -x_k, output weights c_k and output bias f(x_0).

Its error follows from the remainder of linear interpolation. Fix x in a segment [x_k, x_{k+1}] of length h and choose the constant C so that e(u) = f(u) - p(u) - C(u - x_k)(u - x_{k+1}) vanishes at u = x. Then e has three zeros, x_k, x and x_{k+1}, so Rolle’s theorem applied twice gives a \xi with e''(\xi) = f''(\xi) - 2C = 0; hence f(x) - p(x) = \tfrac12 f''(\xi)(x - x_k)(x - x_{k+1}). The product is largest in size at the midpoint, where it is h^2/4, so with M = \max|f''|

Accuracy \varepsilon needs h \le \sqrt{8\varepsilon/M}: in one dimension, a number of units proportional to \varepsilon^{-1/2}.

Why depth helps

The construction spends one unit per linear piece. For some functions depth does far better. The tent map on [0, 1] uses two ReLU units:

It maps each half of [0, 1] onto all of [0, 1], rising on the first half and falling on the second, so applying it once more folds every linear piece in two: the k-fold composition t^k is a sawtooth with 2^k pieces (Telgarsky 2016). As a network, t^k is k layers of 2 units. The first layer has 2 weights and 2 biases; each further layer reads the previous pair (h_1, h_2) through 4 weights and 2 biases, since t = 2h_1 - 4h_2 is a linear function of the pair; the output 2h_1 - 4h_2 takes 2 weights and a bias. In all: 2k units and 4 + 6(k - 1) + 3 = 6k + 1 parameters.

A one-hidden-layer ReLU network on a scalar input, g(x) = \sum_j a_j\operatorname{ReLU}(w_jx + b_j) + c, changes slope only where a unit switches, at x = -b_j/w_j. With n units it has at most n breakpoints and so at most n + 1 pieces. Representing t^k therefore needs at least 2^k - 1 units and 3(2^k - 1) + 1 parameters (an input weight, a bias and an output weight per unit, plus c). Composition reuses the same two units on every piece; width pays for every piece separately.

What the argument shows is that functions exist that a deep network represents exponentially more cheaply than a shallow one; more generally, the number of linear regions a ReLU network can produce grows exponentially with depth (Montúfar et al. 2014). It does not show that the functions you care about are of this kind, nor that gradient descent finds the deep construction. Depth also has a price: the gradient must travel back through every layer, and Section 3 shows how it can vanish on the way, which Sections 5, 6 and 10 address. The informal case for depth is that real data look compositional — edges make textures, textures make parts (Module 03) — and a composition of simple maps matches that structure.

What the hidden units learn

A first-layer ReLU unit computes \operatorname{ReLU}(\mathbf{w}^\top\mathbf{x} + b): zero on one side of the line \mathbf{w}^\top\mathbf{x} + b = 0, rising linearly with distance on the other, and constant along the line. It is a ridge function, a ramp over a half-plane. The output layer adds the ramps with weights, so the logit is piecewise linear and the decision boundary, where the logit is zero, is a polygon. Three ramps can surround a disc with a triangle; two can only form a wedge, which encloses nothing.

Section 1 wrote one example as a column. Training processes a batch of B examples at once, stacked as the rows of a matrix. Most bugs in network code are shape bugs, so this section writes every shape down. It also counts what the forward pass costs in parameters, arithmetic and memory, because every cost estimate later in the series starts from these counts.

The batch form

With \mathbf{H}^{(0)} = \mathbf{X} \in \mathbb{R}^{B\times d_0}, one example per row,

where \mathbf{1} \in \mathbb{R}^{B} is a vector of ones, so \mathbf{1}\mathbf{b}^{(l)\top} copies the bias into every row. Row i of \mathbf{Z}^{(l)} is \mathbf{h}_i^{(l-1)\top}\mathbf{W}^{(l)} + \mathbf{b}^{(l)\top}, the transpose of Section 1’s \mathbf{W}^{(l)\top}\mathbf{h}_i^{(l-1)} + \mathbf{b}^{(l)}: the same computation, with examples as rows instead of columns. In NumPy the bias needs no copying: adding an array of shape (d_l,) to one of shape (B, d_l) broadcasts it over the rows. The output \mathbf{Z}^{(L)} \in \mathbb{R}^{B\times K} holds one row of K logits per example; the softmax happens inside the loss (Section 12), never as a layer in front of it.

import numpy as np

rng = np.random.default_rng(0)
sizes = [64, 128, 128, 10]                         # d0 ... d3: the digits network
params = [(rng.normal(0, np.sqrt(2 / m), (m, n)), np.zeros(n))    # W is (d_{l-1}, d_l)
          for m, n in zip(sizes[:-1], sizes[1:])]

X = rng.normal(size=(64, sizes[0]))                # a batch: B = 64 rows
H = X
for l, (W, b) in enumerate(params, start=1):
    Z = H @ W + b                                  # (B, d_{l-1}) @ (d_{l-1}, d_l) + (d_l,)
    assert Z.shape == (X.shape[0], W.shape[1])
    H = np.maximum(Z, 0) if l < len(params) else Z     # no activation on the logits
print(H.shape, sum(W.size + b.size for W, b in params))

(64, 10) 26122

Parameters, FLOPs and memory

Layer l has d_{l-1}d_l weights and d_l biases, so the network has \sum_{l=1}^{L}(d_{l-1}d_l + d_l) parameters. A product of a (B\times m) and an (m\times n) matrix computes Bn dot products of length m: Bmn multiplications and about as many additions, 2Bmn floating-point operations (FLOPs). Summed over the layers, the forward pass costs about 2B times the number of weights: each weight takes part in one multiplication and one addition per example. Bias additions and activations cost O(Bd_l) per layer, negligible next to Bd_{l-1}d_l for wide layers. This is where the rule of thumb comes from that a large model costs about 2N FLOPs per token for a forward pass.

The series fixes the convention for such counts once, in Module 06, Section 11: in a FLOP count N means the weights used in matrix products (an embedding lookup costs no FLOPs), attention adds a term that grows with the context, and a training step counts as three forward passes (Section 3 shows why the backward pass costs at most two). Memory, by contrast, counts every parameter.

The backward pass (Section 3) needs each layer’s input \mathbf{H}^{(l-1)} and its pre-activation \mathbf{Z}^{(l)} (or enough to evaluate \phi'), so a forward pass done for training stores O(B\sum_l d_l) numbers. Parameters are stored once; activations once per example in the batch, which is why activation memory grows with batch size times depth. Byte units follow the series convention: kB, MB and GB are decimal (10^3, 10^6 and 10^9 bytes); KiB, MiB and GiB are binary (2^{10}, 2^{20} and 2^{30} bytes).

Shapes in practice

PyTorch’s nn.Linear(d_in, d_out) stores its weight as a (d_out, d_in) tensor and computes \mathbf{X}\mathbf{W}^\top + \mathbf{b}. The mathematics is identical; only the stored layout differs, which matters when weights are copied between NumPy and PyTorch. Lab 1 copies W1.T into a PyTorch layer for exactly this reason.

import torch, torch.nn as nn

layer = nn.Linear(64, 128)                         # PyTorch stores (d_out, d_in)
print(tuple(layer.weight.shape), tuple(layer.bias.shape))

pred, t = torch.zeros(256, 1), torch.zeros(256)    # (B, 1) against (B,)
print(tuple((pred - t).shape))                     # broadcast, silently

(128, 64) (128,)
(256, 256)

Write the shape beside every array, in a comment or an assert. The second print shows the commonest silent shape bug in regression: a prediction of shape (B, 1) minus a target of shape (B,) broadcasts to (B, B), every prediction against every target. NumPy computes it without comment. PyTorch’s F.mse_loss warns, “Using a target size (torch.Size([256])) that is different to the input size (torch.Size([256, 1])). This will likely lead to incorrect results due to broadcasting”, and carries on. The model then learns the wrong thing. For prediction \hat y_i the mean over all targets is \frac1B\sum_j(\hat y_i - t_j)^2 = (\hat y_i - \bar t)^2 + \operatorname{Var}(t), so the loss is smallest when every prediction equals the mean target \bar t, whatever the input, and it then equals the variance of the targets.

Gradient descent needs \partial\mathcal{L}/\partial\mathbf{W}^{(l)} and \partial\mathcal{L}/\partial\mathbf{b}^{(l)} for every layer. Backpropagation is the chain rule organised so that each layer’s gradient is computed from the next layer’s in one backward sweep, at a cost comparable to the forward pass. This section derives it, so that every line of a backward pass can be written, checked and costed by hand.

Matrix calculus, as much as is needed

The gradient of a scalar with respect to a vector or matrix has the shape of that vector or matrix. For \mathbf{y} = f(\mathbf{x}) with \mathbf{y} \in \mathbb{R}^m and \mathbf{x} \in \mathbb{R}^n, the Jacobian \mathbf{J} \in \mathbb{R}^{m\times n} has entries J_{ij} = \partial y_i/\partial x_j. The chain rule composes Jacobians, \mathbf{J}_{g\circ f} = \mathbf{J}_g\mathbf{J}_f, and for a scalar loss it reads

Three facts follow, and this section uses nothing else.

1. For \mathbf{z} = \mathbf{W}^\top\mathbf{h} + \mathbf{b}, z_j = \sum_i W_{ij}h_i + b_j, so \partial z_j/\partial h_i = W_{ij}: the Jacobian is \mathbf{W}^\top, and \nabla_{\mathbf{h}}\mathcal{L} = \mathbf{W}\,\nabla_{\mathbf{z}}\mathcal{L}.
2. W_{ij} appears only in z_j, with coefficient h_i, so \partial\mathcal{L}/\partial W_{ij} = h_i\,\partial\mathcal{L}/\partial z_j: \nabla_{\mathbf{W}}\mathcal{L} = \mathbf{h}\,(\nabla_{\mathbf{z}}\mathcal{L})^\top, an outer product.
3. For an elementwise \mathbf{h} = \phi(\mathbf{z}), h_i depends only on z_i, so the Jacobian is \operatorname{diag}(\phi'(\mathbf{z})) and \nabla_{\mathbf{z}}\mathcal{L} = \nabla_{\mathbf{h}}\mathcal{L}\odot\phi'(\mathbf{z}), where \odot is the elementwise product.

When unsure where a transpose goes, check shapes: \nabla_{\mathbf{W}}\mathcal{L} must be d_{\text{in}}\times d_{\text{out}}, \mathbf{h} has length d_{\text{in}} and \nabla_{\mathbf{z}}\mathcal{L} length d_{\text{out}}, and only \mathbf{h}(\nabla_{\mathbf{z}}\mathcal{L})^\top has that shape.

The error signal, and the equation at the output

Define the error signal of layer l as the gradient with respect to its pre-activation, for one example: \boldsymbol{\delta}^{(l)} = \partial\mathcal{L}/\partial\mathbf{z}^{(l)} \in \mathbb{R}^{d_l}.

Equation 1, at the output. For softmax cross-entropy, \mathcal{L} = -\sum_k y_k\ln\hat p_k with \hat{\mathbf{p}} = \softmax(\mathbf{z}) and a one-hot \mathbf{y}. Module 01, Section 6 derived the gradient case by case; with the Jacobian it takes two lines. Differentiating \hat p_k = e^{z_k}/\sum_m e^{z_m} gives \partial\hat p_k/\partial z_j = \hat p_k([k = j] - \hat p_j), where [k = j] is 1 if k = j and 0 otherwise: the softmax Jacobian is \operatorname{diag}(\hat{\mathbf{p}}) - \hat{\mathbf{p}}\hat{\mathbf{p}}^\top. Then

because \sum_k y_k = 1. So \boldsymbol{\delta}^{(L)} = \hat{\mathbf{p}} - \mathbf{y}, the same \hat p - y as logistic regression, and its entries sum to zero. For squared error (\hat y - y)^2 on a linear output, \delta^{(L)} = 2(\hat y - y). When the loss is averaged over a batch, each example’s \boldsymbol{\delta} carries a factor 1/B.

Equation 2, between layers

Layer l + 1 computes \mathbf{z}^{(l+1)} = \mathbf{W}^{(l+1)\top}\phi(\mathbf{z}^{(l)}) + \mathbf{b}^{(l+1)}. Fact 1 carries the error from \mathbf{z}^{(l+1)} back to \mathbf{h}^{(l)}, and fact 3 carries it through \phi:

The error is pushed back through the same weights that carried the signal forward, and gated by the activation’s derivative: a unit with \phi'(z_i) = 0 passes nothing back.

Equations 3 and 4, the parameter gradients

Fact 2 applied to layer l, together with \partial\mathbf{z}^{(l)}/\partial\mathbf{b}^{(l)} = \mathbf{I}, gives

The weight gradient is the outer product of what came into the layer and the error that went out, and it has the shape of \mathbf{W}^{(l)}. Those four equations are the whole algorithm: a forward pass that stores what the equations will need, then one sweep backwards (Figure 2.5).

forward:   h(0) = x
           for l = 1 … L:   z(l) = W(l)ᵀ h(l−1) + b(l);   h(l) = φ(z(l))   (no φ at l = L)
           keep every h(l−1) and z(l)
backward:  δ(L) = ∂𝓛/∂z(L)                      p̂ − y, or 2(ŷ − y)            Equation 1
           for l = L … 1:
               ∂𝓛/∂W(l) = h(l−1) δ(l)ᵀ;   ∂𝓛/∂b(l) = δ(l)                  Equations 3, 4
               if l > 1:   δ(l−1) = (W(l) δ(l)) ⊙ φ′(z(l−1))                Equation 2

The batch form, and the code

For a batch, stack the error signals as rows, \boldsymbol{\Delta}^{(l)} \in \mathbb{R}^{B\times d_l} with row i equal to \boldsymbol{\delta}_i^{(l)\top}. Transposing the equations row by row and summing the parameter gradients over the batch:

Since (\mathbf{H}^\top\boldsymbol{\Delta})_{jk} = \sum_i H_{ij}\Delta_{ik}, the product \mathbf{H}^{(l-1)\top}\boldsymbol{\Delta}^{(l)} is the sum over the batch of the outer products \mathbf{h}_i\boldsymbol{\delta}_i^\top; the bias gradient is the column sums of \boldsymbol{\Delta}^{(l)}. Lab 1’s backward function, for a two-layer regression network with mean squared error, is these lines one for one. Run on the tiny network as a batch of one, it prints the numbers above:

import numpy as np

def forward(p, X):
    Z1 = X @ p["W1"] + p["b1"]; H1 = np.maximum(Z1, 0)        # (B, d1), ReLU
    Y = H1 @ p["W2"] + p["b2"]                                # (B, 1), no activation
    return Y, (X, Z1, H1)                                     # keep what backward needs

def backward(p, cache, Y, T):
    X, Z1, H1 = cache; B = X.shape[0]
    dY = 2 * (Y - T) / B                  # Delta(2): Equation 1, mean over the batch
    g = {"W2": H1.T @ dY, "b2": dY.sum(0)}    # Equations 3 and 4: H1^T Delta, column sums
    dH1 = dY @ p["W2"].T                  # Equation 2: the error passed down ...
    dZ1 = dH1 * (Z1 > 0)                  # ... gated by ReLU'(Z1)
    g["W1"] = X.T @ dZ1; g["b1"] = dZ1.sum(0)
    return g

p = {"W1": np.array([[0.5, -1.0], [0.25, 0.5]]), "b1": np.array([0.1, 0.0]),
     "W2": np.array([[0.8], [-0.6]]), "b2": np.array([0.2])}
X, T = np.array([[2.0, 1.0]]), np.array([[1.0]])             # one example: B = 1
Y, cache = forward(p, X)
g = backward(p, cache, Y, T)
for k in ("W1", "b1", "W2", "b2"):
    print(k, np.round(g[k], 4).tolist())

W1 [[0.896, 0.0], [0.448, 0.0]]
b1 [0.448, 0.0]
W2 [[0.756], [0.0]]
b2 [0.56]

What the backward pass costs

Layer l’s forward pass is one product, \mathbf{H}^{(l-1)}\mathbf{W}^{(l)}, of 2Bd_{l-1}d_l FLOPs. Its backward pass does two products of exactly that size: \boldsymbol{\Delta}^{(l)}\mathbf{W}^{(l)\top} to pass the error down, and \mathbf{H}^{(l-1)\top}\boldsymbol{\Delta}^{(l)} for the weight gradient. The gating and the bias sums cost O(Bd_l). So the backward pass costs at most twice the forward, and a training step at most about three forward passes. It costs less than twice when the first layer is a large share of the work, because the gradient with respect to the input \mathbf{X} is never needed and the first layer skips its error product. Module 06, Section 11 turns this argument into the training-FLOP count for transformers.

Two consequences

Two facts follow from the four equations and explain much of training practice.

Gradients vanish or explode with depth. Unrolling Equation 2 from the output,

a product of L - l matrices \mathbf{D}^{(k-1)}\mathbf{W}^{(k)}, each the transposed Jacobian of one layer. If their norms are below 1 the error shrinks geometrically on its way to the early layers, and the gradient vanishes; if above 1 it grows, and the gradient explodes. Deep sigmoid networks were notoriously hard to train in the 1990s largely for this reason, since \sigma' \le 1/4. ReLU, careful initialisation, normalisation and residual connections are the fixes, roughly in that historical order (Sections 5, 6 and 10).

Training memory grows with depth times batch. Equation 3 needs each layer’s input \mathbf{h}^{(l-1)}, so the forward pass must keep every layer’s input until the backward sweep reaches it: memory proportional to depth times batch, as Section 2 counted. Gradient checkpointing, which recomputes instead of storing, exists for this reason (Section 4).

Section 3 derived the backward pass of an MLP by hand. A new layer type, such as a convolution or an attention block, would need a new derivation, and every hand derivation is a chance for a bug. Frameworks avoid both by differentiating programs mechanically. Automatic differentiation (autodiff) applies the chain rule to the primitive operations a program actually executes, and backpropagation is one instance of its reverse mode. This section works both modes on a small example, shows why reverse mode is the right one for a scalar loss, and describes what PyTorch does when you call loss.backward().

Computational graphs

A computational graph has a node for each primitive operation (add, multiply, matrix product, exp, log, sin, tanh, ReLU) and edges that carry the intermediate values v_i, numbered in an evaluation order in which every node comes after its inputs (a topological order). Each primitive knows its local partial derivatives: \partial(uv)/\partial u = v, \partial\ln u/\partial u = 1/u. The running example, from Baydin et al. (2018), is

with v_1 = \ln x_1, v_2 = x_1x_2, v_3 = \sin x_2, v_4 = v_1 + v_2 and f = v_4 - v_3 (Figure 2.7). Each input feeds two nodes.

Forward mode

Forward mode carries a tangent \dot v_i, the derivative of v_i along a chosen direction in input space, alongside each value:

Seeding the inputs with a direction, \dot{\mathbf{x}} = \mathbf{u}, gives in one pass the directional derivative \mathbf{J}\mathbf{u}, a Jacobian-vector product (JVP). Dual numbers implement it: numbers a + b\epsilon with \epsilon^2 = 0. By Taylor’s theorem f(a + b\epsilon) = f(a) + f'(a)\,b\epsilon, every higher power of \epsilon being zero, so arithmetic on dual numbers carries the derivative along exactly. The product rule falls out of (a + b\epsilon)(c + d\epsilon) = ac + (ad + bc)\epsilon.

import math

class Dual:
    """a + b·ε with ε² = 0: the value a carries its derivative b along."""
    def __init__(self, a, b=0.0):
        self.a, self.b = a, b
    def __add__(self, o):
        return Dual(self.a + o.a, self.b + o.b)
    def __sub__(self, o):
        return Dual(self.a - o.a, self.b - o.b)
    def __mul__(self, o):                      # (a + bε)(c + dε) = ac + (ad + bc)ε
        return Dual(self.a * o.a, self.a * o.b + self.b * o.a)

def log(u): return Dual(math.log(u.a), u.b / u.a)
def sin(u): return Dual(math.sin(u.a), math.cos(u.a) * u.b)

def f(x1, x2):
    return log(x1) + x1 * x2 - sin(x2)

y = f(Dual(2.0, 1.0), Dual(5.0, 0.0))          # seed the tangent (1, 0)
print(f"f = {y.a:.4f}, df/dx1 = {y.b:.4f}")
y = f(Dual(2.0, 0.0), Dual(5.0, 1.0))          # seed the tangent (0, 1)
print(f"f = {y.a:.4f}, df/dx2 = {y.b:.4f}")

f = 11.6521, df/dx1 = 5.5000
f = 11.6521, df/dx2 = 1.7163

Reverse mode

Reverse mode first runs the program forward, keeping the values, then carries adjoints \bar v_i = \partial f/\partial v_i backwards from \bar f = 1:

The “+=” is the multivariable chain rule: a variable used in several places influences the output along several paths, and its adjoint is the sum of the contributions of all of them. Seeding the output with \mathbf{u} gives, in one pass, the vector-Jacobian product (VJP) \mathbf{u}^\top\mathbf{J}. For a scalar loss, u = 1 and that single row is the whole gradient.

Which mode, and what it costs

For f: \mathbb{R}^n \to \mathbb{R}^m the full Jacobian takes n forward-mode passes, a column each, or m reverse-mode passes, a row each (Figure 2.8), and every pass costs a small constant multiple of evaluating f; for reverse mode this is the cheap gradient principle (Griewank and Walther 2008). Training has m = 1 and n from about 10^4 to 10^{11}, so reverse mode wins by a factor of order n. Forward mode wins when n \ll m. The sensitivity of a simulated trajectory of 1,000 samples to 3 design parameters takes 3 forward passes against 1,000 reverse ones. Hessian-vector products are computed as forward mode applied to a reverse-mode gradient. And physics-informed networks (Module 05, Section 9) differentiate a network with respect to its few inputs, a natural use of forward mode.

The tape, and gradient checkpointing

Reverse mode’s price is memory. The backward sweep needs the forward values, so they are recorded (the tape) and held until the sweep has used them. Gradient checkpointing stores only some of them and recomputes the rest: keep the input of every k-th layer, and when the backward sweep reaches a segment, rerun that segment’s k layers forward from its checkpoint to rebuild their activations. The peak is about L/k checkpoints plus one segment of k layers, smallest near k = \sqrt L: activation memory falls from O(L) to O(\sqrt L) for about one extra forward pass (Chen et al. 2016).

Layers as VJP rules

A framework never forms a Jacobian. Each layer supplies a rule that maps the adjoint of its output to the adjoints of its inputs:

• linear, \mathbf{z} = \mathbf{W}^\top\mathbf{h} + \mathbf{b}: given \bar{\mathbf{z}}, \bar{\mathbf{h}} = \mathbf{W}\bar{\mathbf{z}}, \bar{\mathbf{W}} = \mathbf{h}\bar{\mathbf{z}}^\top and \bar{\mathbf{b}} = \bar{\mathbf{z}};
• activation, \mathbf{h} = \phi(\mathbf{z}): \bar{\mathbf{z}} = \bar{\mathbf{h}}\odot\phi'(\mathbf{z});
• softmax cross-entropy on logits: \bar{\mathbf{z}} = \hat{\mathbf{p}} - \mathbf{y}.

These are Section 3’s equations with \bar{\mathbf{z}}^{(l)} = \boldsymbol{\delta}^{(l)}: backpropagation is reverse mode with each layer treated as one primitive. Forming Jacobians instead would be hopeless. For a 4,096 → 4,096 layer applied to a batch of 512, the Jacobian of all outputs with respect to all inputs has (512\cdot 4{,}096)^2 \approx 4.4\times 10^{12} entries, almost all of them zero, while the VJP is one matrix product.

What PyTorch does when you call loss.backward()

Every operation on a tensor with requires_grad=True records a node, visible as the result’s .grad_fn, holding what its VJP rule will need. loss.backward() walks these nodes in reverse topological order and accumulates into each leaf tensor’s .grad with “+=”. The graph is built anew on every forward pass, as the code runs (define-by-run), and freed once backward() has used it. torch.no_grad() turns the recording off, for evaluation and for the optimiser’s own updates; .detach() returns a tensor cut out of the graph, which silently stops the gradient if done by mistake. On the tiny network of Section 2:

import torch

x = torch.tensor([2.0, 1.0])
W1 = torch.tensor([[0.5, -1.0], [0.25, 0.5]], requires_grad=True)   # leaves
b1 = torch.tensor([0.1, 0.0], requires_grad=True)
W2 = torch.tensor([[0.8], [-0.6]], requires_grad=True)
b2 = torch.tensor([0.2], requires_grad=True)

def loss_fn():                                # the network of s2, target t = 1
    h1 = torch.relu(W1.T @ x + b1)            # every operation records a node
    return ((W2.T @ h1 + b2 - 1.0) ** 2).sum()

loss = loss_fn()
node, chain = loss.grad_fn, []
while node is not None:                       # follow the first input of each node
    chain.append(node.name().split("::")[-1])
    node = node.next_functions[0][0] if node.next_functions else None
print(" <- ".join(chain))

loss.backward()                               # reverse sweep; the graph is then freed
print(W1.grad, b2.grad)
loss_fn().backward()                          # the next step's backward, not zeroed
print(W1.grad[:, 0], b2.grad)                 # added to the old values: doubled

for p in (W1, b1, W2, b2):
    p.grad = None                             # what optimizer.zero_grad() does
with torch.no_grad():
    print(loss_fn().requires_grad)            # nothing was recorded
print(W1.detach().requires_grad)              # a tensor cut out of the graph

SumBackward0 <- PowBackward0 <- SubBackward0 <- AddBackward0 <- MvBackward0 <- PermuteBackward0 <- AccumulateGrad
tensor([[0.8960, 0.0000],
        [0.4480, 0.0000]]) tensor([0.5600])
tensor([1.7920, 0.8960]) tensor([1.1200])
False
False

Following the first input of each node leads back from the loss through the sum, the square, the subtraction of t, the bias, the product \mathbf{W}^{(2)\top}\mathbf{h}^{(1)} and the transpose, to AccumulateGrad, the node that adds into W2.grad. The gradients are Section 3’s numbers. The second backward call adds a second copy: because .grad accumulates, every optimiser step must be preceded by optimizer.zero_grad(). The same “+=” that makes the chain rule work makes a forgotten zero_grad add all past gradients into every step (Lab 5, script B).

Lab 2 builds a scalar reverse-mode engine of about 100 lines that does exactly this, one number per node. It also shows why frameworks work on tensors: one training step of a 337-parameter network on 100 examples creates 66,440 scalar nodes, where PyTorch records about ten tensor operations.

Equation 2 of Section 3 multiplies the error signal by \phi'(\mathbf{z}) at every layer on its way down. Choosing an activation is therefore choosing, first of all, a derivative: the number by which every backpropagated error is scaled, once per layer. This section derives the derivatives and describes the two ways a unit stops passing gradient: saturation and death.

The choices

\Phi is the standard normal CDF and \phi_N the standard normal density. As of 2026 most open large language models use SiLU inside SwiGLU feed-forward blocks (Module 06, Section 9). PyTorch takes ReLU’s derivative at z = 0 to be 0. Figure 2.9 plots all six and their derivatives.

Deriving the derivatives

Write the sigmoid as (1+e^{-z})^{-1} and differentiate:

because e^{-z}/(1+e^{-z}) = 1 - 1/(1+e^{-z}). The product s(1-s) of a number in (0, 1) and its complement is largest at s = 1/2, that is at z = 0, where it equals 1/4. The backward pass computes it from the stored output.

For tanh, multiply the numerator and denominator of (e^z - e^{-z})/(e^z + e^{-z}) by e^{-z}:

tanh is a sigmoid stretched vertically to (-1, 1) and squeezed horizontally by 2. By the chain rule \tanh'(z) = 4\sigma'(2z), which at the origin is 4\cdot\tfrac14 = 1: the same shape, but centred on zero and with four times the slope. GELU and SiLU are products, so the product rule gives their derivatives directly: \Phi + z\phi_N and \sigma + z\sigma(1-\sigma).

Autograd reproduces these values, a cheap check on any activation you implement:

import torch
import torch.nn as nn

z = torch.tensor([-3.0, 0.0, 1.0, 2.0, 5.0], requires_grad=True)
acts = {"sigmoid": torch.sigmoid, "tanh": torch.tanh, "relu": torch.relu,
        "gelu": nn.GELU(), "silu": nn.SiLU()}
for name, phi in acts.items():
    (grad,) = torch.autograd.grad(phi(z).sum(), z)   # elementwise, so this is phi'(z)
    print(f"{name:8s}", " ".join(f"{g:8.4f}" for g in grad.tolist()))

sigmoid    0.0452   0.2500   0.1966   0.1050   0.0066
tanh       0.0099   1.0000   0.4200   0.0707   0.0002
relu       0.0000   0.0000   1.0000   1.0000   1.0000
gelu      -0.0119   0.5000   1.0833   1.0852   1.0000
silu      -0.0881   0.5000   0.9277   1.0908   1.0265

Saturation, and why ReLU won

Where \phi' \approx 0 a unit passes almost no gradient; it is saturated. The sigmoid’s derivative falls below 0.01 for |z| > 4.6 and tanh’s for |z| > 3.0. A saturated unit is not dead for ever: it still learns, but at a rate set by that tiny derivative.

ReLU won because of its derivative. On the active side it is exactly 1, so an error passes through many layers undiminished by the activations, where sigmoids multiply it by at most 1/4 per layer. It is also cheap: a comparison, with no exponential. GELU and SiLU keep the near-linear positive side and add smoothness.

GELU and SiLU are not monotonic. GELU has its minimum -0.170 at z = -0.752, SiLU its minimum -0.278 at z = -1.278, and below those points their derivatives are slightly negative. GELU is often computed with the approximation 0.5z\big(1 + \tanh(\sqrt{2/\pi}\,(z + 0.044715z^3))\big), which differs from z\Phi(z) by at most 4.7\times 10^{-4} (largest near |z| = 2.7). PyTorch’s nn.GELU() uses the exact error- function form; nn.GELU(approximate="tanh") uses the approximation.

Zero-centring

A sigmoid layer’s outputs are all positive. For one example the gradient of a weight into unit j is \partial\mathcal{L}/\partial W_{ij} = h_i\delta_j (Section 3), and with every h_i > 0 all these gradients share the sign of \delta_j. The incoming weight vector of unit j can then only move in directions whose components share a sign, so reaching a target takes a zig-zag. tanh avoids this by being zero-centred; so do zero-mean inputs and normalisation layers (Sections 6 and 10). The constraint holds for one example only; a mini-batch gradient sums over examples whose \delta_j differ in sign, which loosens it.

Dead ReLU units

If a ReLU unit’s pre-activation is negative for every training example, its derivative is zero everywhere: its incoming weights and bias receive zero gradient, and plain gradient descent can never revive it (Exercise 5). The usual causes are a large update (too high a learning rate) that pushes the bias far negative, and a poor initialisation. The remedies are a lower learning rate, He initialisation (Section 6), or an activation with a non-zero left side (leaky ReLU, GELU). Measure it: the dead fraction is the share of units whose output is 0 for every example of a batch, and Labs 4 and 5 log it.

The output layer is different

The output’s “activation” is fixed by the loss, not chosen for gradient flow: the identity for squared error, logits into a softmax or sigmoid inside a cross-entropy (Section 12).

Training starts from whatever the initial weights compute. If their scale is wrong, the forward signal and the backward error both grow or shrink geometrically with depth (the product of Jacobians in Section 3) before the optimiser has taken a single step. This section derives the scale from one requirement, that the signal have the same size at every layer, and measures what happens when the requirement is not met.

Symmetry

Set all weights of a layer to the same value, zero or any other constant, and every unit in the layer computes the same function of its input, receives the same error signal, and therefore the same gradient. After the update the units are still identical, and they stay identical for ever: the layer has the expressive power of one unit. With zero weights and ReLU it is worse, since \boldsymbol{\delta}^{(l)} = (\mathbf{W}^{(l+1)}\boldsymbol{\delta}^{(l+1)})\odot\phi' is zero and even the first layer’s gradient vanishes. Random weights break the symmetry; biases can start at zero, because the random weights already make the units differ. Rumelhart, Hinton and Williams (1986) started from small random weights for exactly this reason. What remains is the scale.

The forward derivation

Take one unit, z_j = \sum_{i=1}^{n} w_{ij}h_i with n = n_{\text{in}} inputs (biases are zero at initialisation). Assume the w_{ij} are independent, have zero mean and variance \sigma_w^2, and are independent of the inputs h_i. Then \mathbb{E}[z_j] = \sum_i \mathbb{E}[w_{ij}]\mathbb{E}[h_i] = 0. The cross terms of z_j^2 vanish too, since \mathbb{E}[w_{ij}h_iw_{kj}h_k] = \mathbb{E}[w_{ij}]\,\mathbb{E}[w_{kj}]\,\mathbb{E}[h_ih_k] = 0 for i \ne k, so

The input enters through its second moment \mathbb{E}[h^2], not its variance. The two agree only when h has zero mean, which ReLU outputs do not.

Now ask that \operatorname{Var}(z^{(l)}) = \operatorname{Var}(z^{(l-1)}). For tanh near the origin h \approx z, so \mathbb{E}[h^2] \approx \operatorname{Var}(z^{(l-1)}) and Equation 6.1 needs \sigma_w^2 = 1/n_{\text{in}} (LeCun et al. 1998). For ReLU, z^{(l-1)} is symmetric about zero, so \mathbb{E}[h^2] = \tfrac12\operatorname{Var}(z^{(l-1)}) (Section 5’s worked example) and preservation needs

He initialisation (He et al. 2015). A ReLU halves the second moment, and the factor 2 puts it back. It does not halve the variance, which falls to 0.341\operatorname{Var}(z).

The backward derivation

The error obeys \delta^{(l)}_i = \phi'(z^{(l)}_i)\sum_{j=1}^{n_{\text{out}}} w^{(l+1)}_{ij}\delta^{(l+1)}_j, a sum over the n_{\text{out}} units the unit feeds. The same argument, with the weights independent of the errors, gives

Near the origin tanh has \phi' \approx 1; ReLU has \mathbb{E}[\phi'^2] = P(z > 0) = 1/2. So the backward pass is preserved by \sigma_w^2 = 1/n_{\text{out}} for tanh and 2/n_{\text{out}} for ReLU. Glorot and Bengio (2010) split the difference with \sigma_w^2 = 2/(n_{\text{in}} + n_{\text{out}}), the Xavier or Glorot initialisation. In uniform form the weights are drawn from U(-a, a); since \operatorname{Var}(U(-a, a)) = a^2/3, the bound is a = \sqrt{6/(n_{\text{in}} + n_{\text{out}})}. He et al. noted that one direction suffices: with fan-in scaling the backward factor of each layer is n_{\text{out}}\cdot(2/n_{\text{in}})\cdot\tfrac12 = n_{\text{out}}/n_{\text{in}}, and the product over layers telescopes to the ratio of two widths, which does not grow with depth. Fan-in mode is the usual choice for ReLU.

Measured through ten layers

What wrong scales do

Too large, and tanh units saturate: the forward signal stays bounded, but \phi' \approx 0 at most units and the gradient vanishes; ReLU activations instead grow geometrically and overflow. Too small, and the signal shrinks geometrically, so the output equals the biases whatever the input. The error shrinks on the way back too, because it is multiplied by the same small weights, so the early layers receive tiny gradients and training starts on a plateau. Lab 1, step 7, shows the plateau on a shallow network initialised from \mathcal{N}(0, 10^{-4}).

Framework defaults are not He

PyTorch’s nn.Linear draws its weights and biases from U(-1/\sqrt{n_{\text{in}}}, 1/\sqrt{n_{\text{in}}}), a weight variance of 1/(3n_{\text{in}}). That is a sixth of He’s.

For deep ReLU stacks without normalisation, call the He initialiser explicitly. Shallow networks, and networks with normalisation layers, tolerate the default; Labs 3 to 5 use it at depth 3.

for m in model.modules():
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")  # std sqrt(2 / fan_in)
        nn.init.zeros_(m.bias)

Residual connections

Deep networks add one more device. A residual connection, \mathbf{h}^{(l)} = \mathbf{h}^{(l-1)} + f(\mathbf{h}^{(l-1)}), has the Jacobian \mathbf{I} + \partial f/\partial\mathbf{h}^{(l-1)}. The backward product of Section 3 then always contains an identity path, and the gradient cannot vanish along it. The forward variance now adds, \operatorname{Var}(\mathbf{h}^{(l)}) \approx \operatorname{Var}(\mathbf{h}^{(l-1)}) + \operatorname{Var}(f), and would double at every block if each branch preserved its input’s variance. Deep residual networks therefore scale down or zero-initialise the last layer of each branch, so that every block starts close to the identity. Module 03 introduces residual connections with ResNet, and Module 06 shows them in every transformer block.

The output layer

A small output scale makes the initial predictions near-uniform, so a K-class classifier starts at a loss near \ln K (\ln 10 = 2.303 for digits). Checking it is the first test of Section 14.

Plain stochastic gradient descent with a fixed learning rate is rarely used. The reason is visible on the simplest loss there is, a quadratic, and so is the first cure, momentum. This section analyses both on the quadratic and derives the largest learning rate each can take.

Gradient descent on a quadratic

Module 01, Section 3 analysed this case; here is the result in the form needed below. Take \mathcal{L}(\theta) = \tfrac12\theta^\top\mathbf{A}\theta with \mathbf{A} symmetric positive definite, eigenvalues \lambda_1 \le \dots \le \lambda_d and orthonormal eigenvectors \mathbf{q}_i. The gradient is \mathbf{A}\theta, so a step is \theta \leftarrow (\mathbf{I} - \eta\mathbf{A})\theta. In the eigen-coordinates u_i = \mathbf{q}_i^\top\theta the matrix is diagonal, and each coordinate evolves on its own:

Coordinate i shrinks if and only if |1 - \eta\lambda_i| < 1, so all of them shrink if and only if \eta < 2/\lambda_{\max}. The slowest coordinate sets the rate, \max_i|1 - \eta\lambda_i|, and only the two extreme eigenvalues compete. Balancing them, 1 - \eta\lambda_{\min} = -(1 - \eta\lambda_{\max}), gives \eta = 2/(\lambda_{\max} + \lambda_{\min}) and the best rate

For large \kappa, \ln\rho_{\text{GD}} \approx -2/\kappa: about \kappa/2 steps per factor e. The steep direction fixes the step size; the shallow direction sets the time.

Near a minimum \theta^*, where the gradient vanishes, the second-order Taylor expansion of a network’s loss is \mathcal{L}(\theta^*) + \tfrac12(\theta - \theta^*)^\top\mathbf{H}(\theta - \theta^*), with the Hessian \mathbf{H} in place of \mathbf{A}. The same limit \eta < 2/\lambda_{\max}(\mathbf{H}) applies locally. In full-batch training \lambda_{\max}(\mathbf{H}) tends to rise until it reaches about 2/\eta and then hovers there while the loss keeps falling, unevenly: the edge of stability (Cohen et al. 2021). The learning rate does not only have to respect the curvature; it shapes the curvature the network ends up in.

With mini-batches the gradient is the full gradient plus zero-mean noise whose covariance falls as 1/B (Module 01, Section 4). The noise makes the loss non-monotone and leaves a noise floor proportional to \eta, which Section 9 measures and schedules remove.

Heavy-ball momentum

Momentum keeps a velocity, the form of Polyak (1964) used by Rumelhart, Hinton and Williams (1986) and by PyTorch’s SGD:

Unrolled from \mathbf{v}_0 = \mathbf{0}, \mathbf{v}_t = \sum_{k=0}^{t-1}\mu^k\mathbf{g}_{t-k}, an exponentially weighted sum of past gradients. For a constant gradient it tends to \mathbf{g}/(1-\mu), so the step tends to \eta\mathbf{g}/(1-\mu): an effective learning rate \eta/(1-\mu), ten times \eta for \mu = 0.9. Eliminating \mathbf{v} gives the equivalent form \theta_{t+1} = \theta_t - \eta\mathbf{g}_t + \mu(\theta_t - \theta_{t-1}): a gradient step plus a fraction \mu of the previous step, the ball’s inertia. Figure 2.11 shows the difference on a ravine.

Why this helps is clearest as a filter. Feed the recursion v_t = \mu v_{t-1} + g_t a gradient component that oscillates with frequency \omega, g_t = e^{i\omega t}: \omega = 0 is a component that is the same every step, \omega = \pi one that flips sign every step. In the steady state v_t = H(\omega)e^{i\omega t}; substituting, H e^{i\omega t} = \mu H e^{i\omega(t-1)} + e^{i\omega t}, so

Momentum is a low-pass filter. Across a ravine the gradient flips sign at every step and is damped by 1/(1+\mu) = 0.53; along the ravine it is consistent and is amplified by 1/(1-\mu) = 10, a ratio of (1+\mu)/(1-\mu) = 19 (Figure 2.12).

On the quadratic the analysis is exact. Along an eigenvector with curvature \lambda the equivalent form reads u_{t+1} = (1 + \mu - \eta\lambda)u_t - \mu u_{t-1}, a linear recurrence whose solutions are u_t = r^t with

The two roots multiply to \mu, so when they are complex both have modulus \sqrt\mu: the error contracts by \sqrt\mu per step whatever \lambda is. Both roots lie inside the unit circle if and only if the polynomial is positive at r = 1 and r = -1 (with \mu < 1): at r = 1 it is \eta\lambda > 0, at r = -1 it is 2(1 + \mu) - \eta\lambda > 0. Heavy ball is therefore stable for \eta\lambda_{\max} < 2(1 + \mu), which is 3.8 for \mu = 0.9 against 2 for gradient descent. Choosing \eta = 4/(\sqrt{\lambda_{\max}} + \sqrt{\lambda_{\min}})^2 and \mu = \big((\sqrt\kappa - 1)/(\sqrt\kappa + 1)\big)^2 puts both extreme eigenvalues at the edge of the complex region (the algebra is in Polyak 1964) and gives the rate (\sqrt\kappa - 1)/(\sqrt\kappa + 1): about \sqrt\kappa/2 steps per factor e instead of \kappa/2.

Nesterov momentum

Nesterov momentum evaluates the gradient at the point the velocity is about to carry the parameters to: \mathbf{v} \leftarrow \mu\mathbf{v} + \nabla\mathcal{L}(\theta - \eta\mu\mathbf{v}), \theta \leftarrow \theta - \eta\mathbf{v}. PyTorch’s SGD(nesterov=True) uses an equivalent rewrite that evaluates the gradient at the stored parameters, \mathbf{v} \leftarrow \mu\mathbf{v} + \mathbf{g}, \theta \leftarrow \theta - \eta(\mathbf{g} + \mu\mathbf{v}). The look-ahead corrects the velocity before it overshoots, and at moderate \eta Nesterov is faster. It is not more stable. Its characteristic equation is r^2 - (1+\mu)(1-\eta\lambda)\,r + \mu(1-\eta\lambda) = 0, and the same test at r = -1 gives \eta\lambda_{\max} < 2(1+\mu)/(1+2\mu): 1.36 for \mu = 0.9, below gradient descent’s 2 and heavy ball’s 3.8.

In practice

Momentum acts mostly as a larger effective learning rate. In Lab 3, on the digits network, plain SGD at \eta = 0.05 never gets the training loss below 0.1 in 20 epochs, while SGD at \eta = 0.5 and SGD with \mu = 0.9 at \eta = 0.05 get there by epochs 3 and 4 and end within 0.3 points of each other in validation accuracy. Two conventions matter when moving between libraries or papers. PyTorch’s momentum buffer has no (1-\mu) factor (dampening=0), unlike Adam’s first moment (Section 8); switching between the two forms changes the effective learning rate by 1/(1-\mu). SGD with momentum 0.9 and a schedule remains strong for convolutional networks (Module 03); the adaptive methods of Section 8 are the default for transformers.

Momentum improves the direction of a step; it does nothing about scale. Gradient sizes differ by orders of magnitude across a network’s parameters: at the first step of Lab 3’s digits network the non-zero gradient entries range from 1.8\times 10^{-7} to 6.4\times 10^{-2}. A single \eta is then too large for some parameters and too small for others. The adaptive methods apply a diagonal preconditioner, \theta \leftarrow \theta - \eta\mathbf{D}^{-1}\mathbf{g}, with \mathbf{D} estimated from the gradients’ own history, so that each parameter gets its own step size.

AdaGrad and RMSProp

AdaGrad (Duchi, Hazan and Singer 2011) divides by the root of the accumulated squared gradients, elementwise:

A parameter that is rarely updated keeps a large step, which suits sparse features. But \mathbf{G} only grows: with gradients of roughly constant size, \sqrt{\mathbf{G}} grows like \sqrt t and the effective step decays like 1/\sqrt t, so long non-convex training stalls. RMSProp (Tieleman and Hinton 2012, from a lecture slide rather than a paper) replaces the sum by an exponential moving average, which forgets old gradients:

Adam

Adam (Kingma and Ba 2015) adds momentum to RMSProp, as a moving average of the gradient, and corrects both averages for their start at zero:

with \beta_1 = 0.9, \beta_2 = 0.999 and \epsilon = 10^{-8}. Large language models often use \beta_2 = 0.95, so that \mathbf{v} follows changes in the gradient’s scale faster.

Bias correction. With \mathbf{m}_0 = \mathbf{0}, unrolling the average gives m_t = (1-\beta_1)\sum_{s=1}^{t}\beta_1^{t-s}g_s. If the gradients have a constant mean,

by the geometric series. The average is biased towards zero by the factor 1-\beta_1^t, and dividing by it removes the bias; the same argument with \beta_2 and g^2 gives \hat{\mathbf{v}}. Without the correction the step m/\sqrt v equals the corrected step times (1-\beta_1^t)/\sqrt{1-\beta_2^t}. For \beta_2 = 0.999 that factor is 3.16 at t = 1, peaks at 6.57 at t = 12, and is still 1.26 at t = 1{,}000: early steps several times too large. For \beta_2 = 0.95 it is 0.45 at t = 1 and never exceeds 1.10. Which way the error goes depends on \beta_2 against \beta_1 (Figure 2.13).

What Adam’s step is. Parameter i moves by about \eta|\hat m_i|/\sqrt{\hat v_i}. Since \hat v_i estimates \mathbb{E}[g_i^2] = \mathbb{E}[g_i]^2 + \operatorname{Var}(g_i), the step is about \eta when the gradient is consistent and smaller when it is noisy. At t = 1, \hat m = g and \hat v = g^2, so the step is exactly \eta g/(|g| + \epsilon) \approx \eta\,\operatorname{sign}(g): every parameter with a non-negligible gradient moves by \eta.

The update is invariant to rescaling the loss: \mathbf{g} \to c\mathbf{g} gives \hat{\mathbf{m}} \to c\hat{\mathbf{m}} and \sqrt{\hat{\mathbf{v}}} \to |c|\sqrt{\hat{\mathbf{v}}}, and the ratio is unchanged (\epsilon aside). Multiply the loss by 1,000 and SGD’s step becomes 1,000 times larger, so it diverges at any \eta that was stable before; Adam’s step does not change. This is why Adam forgives unscaled losses. It is not invariant to rescaling the parameters, which is why unscaled inputs still hurt it (Exercise 15).

Adam is also per-coordinate: it equalises scales along the parameter axes but cannot undo curvature in rotated directions. On the optimiser-paths widget’s ravine (Section 7) it needs about 126 steps when the valley is aligned with the axes and about 214 when the valley is rotated by 45°, while gradient descent and heavy ball do not change.

Memory. Adam keeps two extra numbers per parameter. With fp32 weights and gradients that is 4 + 4 + 4 + 4 = 16 bytes per parameter; with bf16 weights and gradients plus an fp32 master copy it is 2 + 2 + 4 + 4 + 4 = 16 bytes as well. The digits network needs 26{,}122\cdot 16 = 417{,}952 bytes, 418 kB; a model of 7\times 10^9 parameters needs 112 GB before any activations. Module 08, Section 8 does the full accounting.

Weight decay is not L2 under Adam

Under SGD the two are the same. Adding (\lambda/2)\|\theta\|^2 to the loss adds \lambda\theta to the gradient, and

which shrinks the weights by the factor 1 - \eta\lambda every step: weight decay. PyTorch’s weight_decay=λ corresponds to the penalty (\lambda/2)\|\theta\|^2; Module 01 wrote the penalty as \lambda\|\theta\|^2, which doubles the coefficient.

Under Adam they differ (Loshchilov and Hutter 2019). With an L2 penalty, \mathbf{g} + \lambda\theta goes into \mathbf{m} and \mathbf{v}, and the decay term is divided by \sqrt{\hat{\mathbf{v}}} like everything else: parameter i shrinks by about \eta\lambda\theta_i/\sqrt{\hat v_i} per step. Parameters with a large gradient history are regularised less, those with small gradients more. AdamW applies the decay outside the normalisation,

shrinking every parameter by the same factor per step, and following the learning-rate schedule. No L2 coefficient reproduces it unless \hat{\mathbf{v}} is the same for every parameter.

In PyTorch, torch.optim.AdamW is the decoupled form (default weight_decay=0.01), and torch.optim.Adam(weight_decay=λ) is the coupled L2 form (default 0); recent versions also accept Adam(..., decoupled_weight_decay=True), which is AdamW. Decay the weight matrices, not biases or normalisation gains, which set offsets and scales that decay would only distort. Two parameter groups do it:

decay = [p for p in model.parameters() if p.ndim >= 2]      # weight matrices
no_decay = [p for p in model.parameters() if p.ndim < 2]    # biases, norm gains
opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.01},
                         {"params": no_decay, "weight_decay": 0.0}], lr=2e-3)
print(sum(p.numel() for p in decay), sum(p.numel() for p in no_decay))

25856 266

For the digits network, 25,856 weights are decayed and 266 biases are not. Adam’s original convergence proof was later shown to be flawed (Reddi, Kale and Kumar 2018). The method works in practice, and that is the evidence for it.

Three devices surround the optimiser: a schedule changes \eta during training, the range test finds its scale, and gradient clipping caps the occasional huge gradient.

Why decay

Module 01, Section 4 derived the reason. On one quadratic direction with curvature \lambda and mini-batch gradient noise of variance s^2/B, a constant step leaves the iterate hovering around the minimum with stationary variance

an excess loss \tfrac12\lambda V proportional to \eta/B. A constant learning rate therefore leaves the final model at a random point of a band whose width \eta sets; decaying \eta shrinks the band.

Schedules

Step decay multiplies \eta by 0.1 at fixed fractions of training, for example 50% and 75%. Cosine decay (Loshchilov and Hutter 2017) follows half a cosine from \eta_{\max} to \eta_{\min} over T steps:

Linear warmup raises \eta from near 0 to \eta_{\max} over the first T_w steps, typically 1–5% of training, a few hundred to a few thousand steps; the cosine then runs over the remaining T - T_w. Warmup plus cosine to a tenth of the peak or less is the language-model default; as of 2026, warmup–stable–decay schedules, which hold the peak and decay only near the end, are a common alternative (Module 08, Section 6). Figure 2.14 draws three schedules on one set of axes.

Warmup has three reasons. At the start Adam’s \hat{\mathbf{v}} is estimated from a handful of gradients and is noisy even after bias correction, so the early steps are erratic (Liu et al. 2020). The curvature at initialisation can be high, so a step that is safe later is too large then. And large-batch SGD with a scaled-up learning rate needs a ramp: Goyal et al. (2017) scale \eta linearly with the batch size and warm up over 5 epochs. The large-batch regime belongs to Module 08.

The range test

The peak learning rate is the most important hyperparameter after the architecture. The learning-rate range test (Smith 2017) finds its scale in minutes: train for a few hundred steps while raising \eta geometrically from about 10^{-6} to 10, record the loss smoothed by an exponential moving average, and plot it against \log\eta. The curve is flat while \eta is too small, falls fastest over about a decade, reaches a minimum, and then rises steeply. Choose the peak about 3–10 times below the minimum, near the steepest fall: the full run is longer and noisier than the test.

In PyTorch, CosineAnnealingLR(opt, T_max) counts calls to scheduler.step(): epochs if stepped per epoch, batches if per batch. Warmup comes from LinearLR inside SequentialLR, or from a LambdaLR. Call scheduler.step() after optimizer.step(). The worked example’s schedule:

from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
total, warm = 10_000, 500
sched = SequentialLR(opt, milestones=[warm], schedulers=[
    LinearLR(opt, start_factor=1e-3, end_factor=1.0, total_iters=warm),
    CosineAnnealingLR(opt, T_max=total - warm, eta_min=3e-5)])   # T_max in batches

Gradient clipping

Gradient clipping by global norm concatenates all parameters’ gradients into one vector and, if its norm exceeds a threshold c, rescales it: \mathbf{g} \leftarrow c\,\mathbf{g}/\|\mathbf{g}\|. The direction is kept. c = 1.0 is common for transformers and recurrent networks (Pascanu, Mikolov and Bengio 2013; Module 04, Section 4). Clipping each value separately (clip_grad_value_) changes the direction and is a cruder tool.

Clipping matters even under Adam, whose step is normalised, because a spike still enters \mathbf{m} and \mathbf{v}. In units of the typical gradient (g \approx 1, m \approx 1, v \approx 1), let one step bring a gradient of 100. Then v = 0.999 + 0.001\cdot 100^2 = 11.0 and \sqrt v = 3.3, so the following steps of that parameter shrink about 3.3 times. The excess decays as 10\cdot 0.999^t, with time constant 1/(1-\beta_2) = 1{,}000 steps: \sqrt v needs about 3,860 steps to return within 10% of normal. Meanwhile m = 0.9 + 0.1\cdot 100 = 10.9 points along the spike, and that step is 10.9/3.3 \approx 3 times the usual size. Clipping first prevents both. clip_grad_norm_ returns the norm before clipping: log it. If clipping fires on most steps, the threshold or the learning rate is wrong.

loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # pre-clip
opt.step()
sched.step()

Section 6 chose the initial weights so that pre-activations start with a sensible scale. That holds at step 0 only. Training changes the weights, and a pre-activation that drifts to a large mean or spread pushes its unit into saturation or into the dead region. A normalisation layer re-standardises activations inside the network at every step. Three variants matter; they differ in which numbers the mean and the spread are taken over (Figure 2.15).

Batch normalisation

Batch normalisation (Ioffe and Szegedy 2015) standardises each feature over the mini-batch. For a batch of pre-activations \mathbf{Z} \in \mathbb{R}^{B\times d} and one feature j,

Each feature of \hat{\mathbf{Z}} has mean 0 and variance just under 1 over the batch. The learned gain \gamma_j (initialised to 1) and shift \beta_j (initialised to 0) give back the freedom the standardisation took away: with \gamma_j = \sqrt{\sigma^2_{\mathcal{B}} + \epsilon} and \beta_j = \mu_{\mathcal{B}} the layer returns its input, so inserting it removes no function the network could represent. The constant \epsilon, typically 10^{-5}, keeps the division finite.

Because the mean is subtracted, a constant added to a feature disappears from \hat z. The bias of the layer in front of a batch norm is therefore redundant, and \beta takes over its job; that is why Module 03 writes bias=False before every batch norm.

Training and evaluation. In training mode the layer uses the current batch’s statistics and also updates running averages,

In PyTorch the running variance uses the unbiased batch variance, as written, and momentum=0.1 is m, the weight of the new value, the opposite of an optimiser’s convention. In evaluation mode (model.eval()) the layer uses \mu_{\text{run}} and \sigma^2_{\text{run}}, so one example’s output no longer depends on its batch-mates. Forgetting model.eval() is a classic bug (Lab 5). So is judging a model early in training, when the running averages still lag fast-changing weights.

import torch, torch.nn as nn

bn = nn.BatchNorm1d(1)                           # gamma = 1, beta = 0, momentum = 0.1
z = torch.tensor([[2.0], [4.0], [6.0], [8.0]])   # one feature, a batch of four
y = bn(z)                                        # training mode: batch statistics
print("train out:", [f"{v:.3f}" for v in y.detach().ravel().tolist()])
print(f"running mean {bn.running_mean.item():.3f}, var {bn.running_var.item():.3f}")
bn.eval()                                        # evaluation mode: running statistics
print(f"eval out for 6: {bn(torch.tensor([[6.0]])).item():.3f}")

train out: ['-1.342', '-0.447', '0.447', '1.342']
running mean 0.500, var 1.567
eval out for 6: 4.394

Normalising over the batch has four consequences. An example’s output in training depends on its batch-mates, which adds noise (a mild regulariser) and invites the bugs above. Small batches give noisy statistics; below about 16 examples per batch, group norm or layer norm replace it (Module 03). Sequences of varying length fit badly, because padding pollutes the statistics. And a batch of one has zero variance in every feature; in training mode PyTorch refuses it with Expected more than 1 value per channel when training.

Why it helps

The original paper credited batch norm with reducing “internal covariate shift”, the drift of each layer’s input distribution. Santurkar et al. (2018) challenged that account: networks still trained well when noise re-created the drift after each batch norm, and the benefit they measured was mainly a smoother, better-conditioned loss landscape that tolerates larger learning rates.

One exact fact explains much of batch norm’s interaction with the optimiser. Let the weights \mathbf{W} feed a batch norm. Scaling them by c > 0 scales \mathbf{z}, \mu_{\mathcal{B}} and \sigma_{\mathcal{B}} by c, so \hat z is unchanged (ignoring \epsilon) and \mathcal{L}(c\mathbf{W}) = \mathcal{L}(\mathbf{W}). Differentiating both sides with respect to \mathbf{W},

Larger weights receive proportionally smaller gradients, so the relative change of a gradient step, \eta\|\nabla\mathcal{L}\|/\|\mathbf{W}\|, scales as \eta/\|\mathbf{W}\|^2: the effective learning rate depends on the weights’ norm. Differentiating \mathcal{L}(c\mathbf{W}) with respect to c at c = 1 also gives \langle\mathbf{W}, \nabla\mathcal{L}\rangle = 0: the gradient is orthogonal to \mathbf{W}, so plain gradient steps slowly grow \|\mathbf{W}\| and the effective rate falls. Weight decay on such weights cannot restrict the function, which does not depend on \|\mathbf{W}\|; it mostly keeps the norm small and the effective learning rate high.

Layer normalisation and RMSNorm

Layer normalisation (Ba, Kiros and Hinton 2016) takes the statistics over the d features of one example instead:

It is identical in training and evaluation and independent of the batch, so none of batch norm’s four consequences apply. Recurrent networks and transformers use it.

RMSNorm (Zhang and Sennrich 2019) drops the mean subtraction and the shift and divides by the root mean square:

It needs one reduction instead of two, works as well in practice, and is what most large language models use as of 2026; Module 06 places it in the block (pre-norm against post-norm). PyTorch provides nn.BatchNorm1d, nn.LayerNorm and, from version 2.4, nn.RMSNorm.

The backward pass of a normalisation

Normalisation changes what the gradient can do. Take layer norm with \boldsymbol{\gamma} = \mathbf{1}, \boldsymbol{\beta} = \mathbf{0} and \epsilon ignored, so \hat z_k = (z_k - \mu)/\sigma, and write \bar{\mathbf{g}} = \partial\mathcal{L}/\partial\mathbf{y}. The pieces:

• \partial\mu/\partial z_j = 1/d.
• \partial\sigma^2/\partial z_j = \frac{2}{d}\sum_k (z_k - \mu)\big([k = j] - \tfrac{1}{d}\big) = \frac{2}{d}(z_j - \mu), because the deviations sum to zero; so \partial\sigma/\partial z_j = (z_j - \mu)/(d\sigma) = \hat z_j/d.
• By the quotient rule, \dfrac{\partial\hat z_k}{\partial z_j} = \dfrac{[k = j] - 1/d}{\sigma} - \dfrac{z_k - \mu}{\sigma^2}\,\dfrac{\hat z_j}{d} = \dfrac{1}{\sigma}\Big([k = j] - \dfrac{1}{d} - \dfrac{\hat z_k\hat z_j}{d}\Big).
• Summing against \bar g_k:

Since \sum_j\hat z_j = 0 and \sum_j\hat z_j^2 = d, this gradient sums to zero and is orthogonal to \hat{\mathbf{z}}. Through a normalisation the layers below cannot change the mean or the scale of what they send up, only its pattern; \boldsymbol{\gamma} and \boldsymbol{\beta} set scale and offset. Batch norm’s backward pass is the same formula with the means taken over the batch instead of the features.

Module 01, Section 9 treated regularisation as trading variance for bias, with its strength tuned on the validation set. Networks stretch that picture. They are routinely over-parameterised (Lab 4 fits 26,122 parameters to 1,078 training images) and still generalise, partly through the implicit regularisation of SGD noise (Module 01, Section 4) and partly through the explicit methods below.

Weight decay

Weight decay is Module 01’s L_2 penalty, applied through AdamW so that every weight shrinks by the same fraction per step whatever its gradient history (Section 8). Typical \lambda runs from 10^{-4} to 10^{-1}, depending on the optimiser and its convention. Biases and normalisation gains are excluded, and on weights that feed a normalisation layer decay mostly raises the effective learning rate (Section 10).

Dropout

Dropout (Srivastava et al. 2014) corrupts the hidden units during training. Each unit is multiplied by an independent mask m_i \sim \text{Bernoulli}(1 - p), which is 0 with probability p, and the survivors are scaled by 1/(1 - p), the inverted form:

At evaluation nothing is dropped or scaled. The scaling is what makes the two modes agree. Taking the expectation over the mask,

so the next layer sees the same expected input in both modes. For the variance, \mathbb{E}[\tilde h_i^2] = (1 - p)\,h_i^2/(1 - p)^2 = h_i^2/(1 - p), so

multiplicative noise, proportional to the unit’s own size, with relative variance 1 at p = 0.5 and 0.11 at p = 0.1. The original formulation multiplied the weights by 1 - p at test time instead; frameworks implement the inverted form, so that evaluation needs no change at all.

Why it works: no unit can be relied on, because any may be missing. Equivalently, training samples from the 2^n thinned networks of an n-unit layer, which share weights, and evaluation approximates their ensemble: exactly, for a single layer feeding a softmax, as the renormalised geometric mean of their predictions (Hinton et al. 2012); approximately for deeper networks. Keeping dropout on at test time and averaging many passes gives a cheap uncertainty estimate, Monte Carlo dropout (Gal and Ghahramani 2016).

Typical p: up to 0.5 in older MLPs and in AlexNet’s fully connected layers; 0.1 in transformers. Many large language model pretraining runs, as of 2026, use none, because each token is seen about once (Module 08).

Early stopping

Early stopping evaluates the validation loss every epoch, keeps the checkpoint with the lowest value, and stops after a patience of several epochs without improvement; the best weights are then restored. Module 01 showed why it regularises: t steps of gradient descent from zero leave eigen-direction i of a quadratic at \big[1 - (1 - \eta\lambda_i)^t\big]\theta^*_i, close to ridge regression’s \lambda_i/(\lambda_i + \tau)\cdot\theta^*_i with \tau \approx 1/(\eta t), where \tau is the coefficient of the penalty (\tau/2)\|\theta\|^2 (Goodfellow et al. 2016, §7.8): approximately L_2 regularisation whose strength is set by the training time. What is new for networks is the procedure: checkpoints, patience, and restoring the best weights.

Calibration and temperature scaling

Module 01, Section 7 measured calibration and deferred networks to this module. Networks trained long with cross-entropy tend to be overconfident: their top probability exceeds their accuracy. Label smoothing (below) acts at training time and can overshoot into underconfidence. Temperature scaling (Guo et al. 2017) divides the logits by one scalar T > 0, \hat{\mathbf{p}} = \softmax(\mathbf{z}/T), with T fitted on the validation set by minimising the negative log-likelihood. Dividing by a positive constant keeps the order of the logits, so the argmax and the accuracy are untouched; T > 1 softens an overconfident model and T < 1 sharpens an underconfident one.

Module 07 returns to calibration for language models.

Data augmentation and input noise

Data augmentation adds label-preserving transformations of the training data. For images it is the single most effective regulariser (Module 03). For sensor or tabular data the options are input noise, and time shifts or scalings to which the label really is invariant. Training with small Gaussian input noise is equivalent, to first order, to a Tikhonov penalty on the output’s derivatives with respect to the input (Bishop 1995).

Label smoothing

Label smoothing (Szegedy et al. 2016) trains on the softened target \mathbf{y}_{\text{LS}} = (1 - \alpha)\mathbf{y} + \alpha/K, which gives the true class 1 - \alpha + \alpha/K and every other class \alpha/K. The cross-entropy gradient with respect to the logits keeps the form of Section 3, \hat{\mathbf{p}} - \mathbf{y}_{\text{LS}}.

The reason to do this is the behaviour of hard labels. On separable training data, -\ln\hat p_y is positive for every finite logit margin and reaches 0 only as the margin goes to infinity, so the gradient never vanishes and the logits keep growing. With smoothed targets the cross-entropy -\sum_k y_{\text{LS},k}\ln\hat p_k is minimised at \hat{\mathbf{p}} = \mathbf{y}_{\text{LS}}, where the gradient is zero, and since z_{\text{true}} - z_{\text{other}} = \ln(\hat p_{\text{true}}/\hat p_{\text{other}}) for a softmax, the optimal margin is finite:

Label smoothing often improves accuracy and calibration, but can overshoot into underconfidence. It also erases the relative sizes of the wrong-class probabilities, which knowledge distillation would use (Müller et al. 2019). In PyTorch it is one argument, F.cross_entropy(logits, targets, label_smoothing=0.1).

Above all, more data

A network’s variance falls with N, and no regulariser substitutes for more data. Every regulariser also changes the best learning rate, because it changes the gradient: re-tune \eta after adding one.

A loss is built from exponentials and logarithms of numbers the network chooses. In exact arithmetic the formulas of Module 01, Section 5 are fine. In floating point e^z overflows for moderately large z, a small probability underflows to zero, \ln 0 = -\infty, and a single infinity turns into NaN throughout the network within one step. This section locates those limits and shows how losses are computed so that they are never reached.

Floating-point formats

A floating-point number has a sign bit, exponent bits that set its range, and mantissa bits that set its precision; the machine epsilon, the gap between 1 and the next number, is 2^{-\text{mantissa bits}}.

The values are those reported by torch.finfo and numpy.finfo. bf16 keeps fp32’s eight exponent bits, and so its range, but has only two to three significant digits; fp16 has three more mantissa bits but a range that ends at 65,504. Taking logarithms of the largest values, e^z overflows fp32 and bf16 above z \approx 88.7 and fp16 above z = \ln 65{,}504 = 11.09. In the other direction e^{-z} falls below the smallest subnormal at z \approx 103 in fp32 and z \approx 16.6 in fp16, and is rounded to exactly zero soon after (from about 104 and 17.3).

The log-sum-exp identity

Every loss over K classes needs \ln\sum_j e^{z_j}. For any constant m,

This is Module 01’s observation that adding a constant to every logit changes nothing, put to work. Choose m = \max_j z_j. Then every exponent z_j - m is at most 0, so nothing overflows, and one term equals e^0 = 1, so the sum is at least 1 and its logarithm is never \ln 0. The cross-entropy for true class y follows directly:

The fused function, F.cross_entropy applied to logits, computes the loss this way and its gradient \hat{\mathbf{p}} - \mathbf{y} from the same stable quantities. Module 01’s stable binary form is the case K = 2.

import torch, torch.nn.functional as F

z = torch.tensor([[1000.0, 999.0, 998.0]])
print(z.exp() / z.exp().sum())                 # by hand: inf / inf
print(f"{F.cross_entropy(z, torch.tensor([0])).item():.6f}")  # logsumexp(z) - z_0

z2 = torch.tensor([0.0, -120.0])
print(torch.log(torch.softmax(z2, dim=0)))     # softmax underflows to 0, then log 0
print(torch.log_softmax(z2, dim=0))            # z - logsumexp(z): finite

logit, target = torch.tensor([17.0]), torch.tensor([0.0])
print(F.binary_cross_entropy(torch.sigmoid(logit), target).item())  # clamped
print(F.binary_cross_entropy_with_logits(logit, target).item())     # correct

tensor([[nan, nan, nan]])
0.407606
tensor([0., -inf])
tensor([   0., -120.])
100.0
17.0

Three ways to get it wrong

Taking the log of a softmax in two steps. The softmax can underflow to exactly 0 for a very negative logit, and \ln 0 = -\infty gives an infinite loss and NaN gradients.

Applying a softmax before F.cross_entropy. The loss applies log-softmax itself, so the network’s probabilities are treated as logits confined to [0, 1]. At best the true class gets probability 1 and the others 0, and those “logits” give \hat p_y = e/(e + K - 1), so

The gradient must also pass through the extra softmax’s Jacobian, whose entries are at most 1/4 in size, so learning is very slow. Accuracy can still rise, which is what makes this a common, silent bug.

Applying a sigmoid before a binary loss. BCEWithLogitsLoss computes, from the logit, \ell = \max(z, 0) - zy + \ln(1 + e^{-|z|}), which never overflows. A sigmoid followed by BCELoss saturates: in fp32, 1 + e^{-z} rounds to exactly 1 once e^{-z} < 2^{-24}, that is from z = 24\ln 2 \approx 16.64, so \sigma(z) becomes exactly 1.0 and \ln(1 - 1) = -\infty. PyTorch clamps the logarithm at -100, so the loss is silently capped at 100 and its gradient is wrong.

Regression losses and reduction

Squared error is the Gaussian negative log-likelihood (Module 01) and punishes outliers quadratically. The Huber loss, \tfrac12 r^2 for |r| \le \delta_{\text{H}} and \delta_{\text{H}}(|r| - \delta_{\text{H}}/2) beyond, has its gradient clipped at \pm\delta_{\text{H}} and is the usual compromise: with \delta_{\text{H}} = 1 (PyTorch’s delta) a residual of 3 costs 2.5 instead of 4.5. Absolute error corresponds to Laplace noise.

The reduction matters too. A mean over the batch keeps the gradient’s scale independent of B; a sum multiplies the gradient, and so the effective learning rate, by B. PyTorch’s default is the mean; stay consistent when changing the batch size.

Mixed precision, and variances

Mixed precision computes matrix products in bf16 or fp16 but keeps the master weights, softmaxes, losses and normalisation statistics in fp32; torch.autocast chooses the precision per operation. fp16 also needs loss scaling: multiply the loss by S (for example 2^{16}) before backward and divide the gradients by S, so that small gradients do not underflow (Micikevicius et al. 2018). bf16 has fp32’s range and needs none, which is why the first advice for NaN losses is to use bf16 rather than fp16. Module 08 covers mixed precision at scale.

Compute variances as \operatorname{mean}\big((x - \mu)^2\big), never as \operatorname{mean}(x^2) - \mu^2, which subtracts two nearly equal large numbers. In fp32, for x = (10000, 10001, 10002) the first gives 0.6667 and the second exactly 0.

Every section so far justifies a line of a training loop. This one puts them together: a loop of about thirty lines trains a two-hidden-layer network to tell the inside of a circle from the outside of it in a square (a relative of the playground’s circles in Section 1), and the table after it says which section each line comes from, so that the loop can be read rather than copied.

import torch, torch.nn as nn, torch.nn.functional as F

torch.manual_seed(0)
X = torch.rand(2048, 2) * 2 - 1                       # points in the square
T = ((X ** 2).sum(1) < 0.5).long()                    # inside the circle of radius sqrt(0.5)
Xtr, Ttr, Xva, Tva = X[:1536], T[:1536], X[1536:], T[1536:]
mu, sd = Xtr.mean(0), Xtr.std(0)                      # standardise with TRAINING statistics
norm = lambda x: (x - mu) / sd

model = nn.Sequential(nn.Linear(2, 32), nn.GELU(), nn.Linear(32, 32), nn.GELU(),
                      nn.Linear(32, 2))
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-2)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=400)

for epoch in range(400):
    model.train()
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr), 64):
        idx = perm[i:i + 64]
        loss = F.cross_entropy(model(norm(Xtr[idx])), Ttr[idx])   # on logits
        opt.zero_grad(); loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    sched.step()
    if epoch % 100 == 99:
        model.eval()
        with torch.no_grad():
            acc = (model(norm(Xva)).argmax(1) == Tva).float().mean()
        print(epoch + 1, round(loss.item(), 4), round(acc.item(), 3))

100 0.0069 0.992
200 0.0205 0.996
300 0.0056 0.996
400 0.0107 0.996

Two reading notes. The printed loss is the last mini-batch’s loss, which is noisy: it jumps from 0.0069 to 0.0205 between epochs 100 and 200, while the epoch means fall steadily (0.0108, 0.0065, 0.0054, 0.0046). Log the epoch mean. And T_max counts scheduler.step() calls, here epochs.