12. Deep Learning

Deep learning replaces hand-designed features with learned ones: instead of deciding which transformations of the input matter, you stack differentiable layers and let gradient descent discover the representation that makes the task easy. This chapter builds that machinery from the perceptron upward -- activations, backpropagation, initialisation, optimisers, normalisation, and regularisation -- then surveys the architectures that encode specific structural assumptions, the engineering required to train at scale, and the transfer paradigm in which almost nobody trains from scratch any more.

<a id="121-neural-network-basics"></a>

12.1 Neural Network Basics

Every neural network is the same three ideas repeated: a linear map, a nonlinearity, and a gradient computed by the chain rule; this section develops each one and then covers the four practical concerns -- how weights start, how they are updated, how activations are kept in range, and how overfitting is controlled -- that decide whether a network trains at all.

Perceptrons and Multilayer Networks

Start with the smallest learning machine that is not just a lookup table. A perceptron takes an input vector x, forms a weighted sum z = w·x + b, and outputs 1 if z > 0 and 0 otherwise. Learning is a single rule: when the prediction is wrong, push the weights toward the input if the true label was 1 and away from it if it was 0. That is it.

This is enough to learn AND, and it converges in a handful of passes. It is provably not enough to learn XOR, and the reason is geometric rather than about training time. A perceptron's decision boundary is the set of points where w·x + b = 0, which is a straight line in two dimensions and a hyperplane in general. XOR puts the positive class at (0,1) and (1,0) and the negative class at (0,0) and (1,1); no straight line separates those diagonals. Running the training loop longer cannot help, because the hypothesis you are searching over does not contain a solution.

The fix is to stop feeding the classifier the raw inputs. Add a hidden layer: a set of units that each compute their own w·x + b, pass it through a nonlinearity, and hand the results forward. The second layer then sees a new coordinate system in which the problem may be linearly separable. For XOR, two hidden units suffice -- one firing for "at least one input is on", the other for "both are on" -- and the output layer merely subtracts the second from the first.

The analogy worth keeping: a perceptron is a single straight cut through the data. A multilayer network is a machine that first bends the paper and then makes a straight cut. The bending is what the hidden layers buy you.

The universal approximation theorem says that one hidden layer with enough units can approximate any continuous function on a bounded region to arbitrary accuracy. This is reassuring and almost useless in practice, because "enough units" can mean exponentially many, and the theorem says nothing about whether gradient descent will find those weights. Depth is the practical answer: composing many modest transformations expresses some functions with far fewer parameters than one enormous layer.

import numpy as np
from sklearn.neural_network import MLPClassifier
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split

rng = np.random.default_rng(12)

# --- THE PERCEPTRON: the rule is one line, and so is its limit. ---
def perceptron(X, y, epochs=60, lr=0.1):
    """Classic perceptron: update only on misclassified rows."""
    w = np.zeros(X.shape[1]); b = 0.0
    history = []
    for _ in range(epochs):
        errors = 0
        for xi, yi in zip(X, y):
            pred = 1 if xi @ w + b > 0 else 0
            if pred != yi:                       # no update when already right
                w += lr * (yi - pred) * xi       # move the plane toward/away
                b += lr * (yi - pred)
                errors += 1
        history.append(errors)
    acc = np.mean([(1 if xi @ w + b > 0 else 0) == yi for xi, yi in zip(X, y)])
    return w, b, acc, history

X_logic = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
y_and = np.array([0, 0, 0, 1])
y_xor = np.array([0, 1, 1, 0])

print("A SINGLE PERCEPTRON ON TWO BOOLEAN FUNCTIONS")
for name, y in [("AND", y_and), ("XOR", y_xor)]:
    w, b, acc, hist = perceptron(X_logic, y)
    settled = next((i for i, e in enumerate(hist) if all(h == 0 for h in hist[i:])),
                   None)
    print(f"  {name:4s} accuracy {acc:.2f}  weights {np.round(w, 2)}"
          f"  converged at epoch {settled}")

# --- ONE HIDDEN LAYER: hand-built, to show that two units are enough. ---
# h1 fires when at least one input is on; h2 fires only when both are on.
W1 = np.array([[1.0, 1.0], [1.0, 1.0]])      # both units read both inputs
b1 = np.array([-0.5, -1.5])                  # thresholds at "any" and "both"
W2 = np.array([1.0, -1.0]); b2 = -0.5        # output = h1 AND NOT h2

H = (X_logic @ W1.T + b1 > 0).astype(float)  # step nonlinearity
out = (H @ W2 + b2 > 0).astype(int)
print("\nHAND-BUILT 2-UNIT HIDDEN LAYER ON XOR")
for xi, hi, oi, yi in zip(X_logic, H, out, y_xor):
    print(f"  x={xi.astype(int)}  hidden={hi.astype(int)}  out={oi}  target={yi}")

# --- DEPTH VS WIDTH AT A MATCHED PARAMETER BUDGET. ---
X, y = make_moons(n_samples=3000, noise=0.25, random_state=0)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.3, random_state=0)

def n_params(layers, d_in=2, d_out=1):
    """Weights + biases for a fully connected stack."""
    sizes = [d_in] + list(layers) + [d_out]
    return sum(sizes[i] * sizes[i + 1] + sizes[i + 1] for i in range(len(sizes) - 1))

print("\nDEPTH VS WIDTH ON make_moons (matched budgets)")
print(f"  {'hidden layers':22s} {'params':>7s} {'train':>7s} {'test':>7s}")
for layers in [(2,), (4,), (16,), (64,), (256,), (16, 16), (8, 8, 8), (6, 6, 6, 6)]:
    m = MLPClassifier(layers, max_iter=4000, random_state=0, alpha=1e-4).fit(Xtr, ytr)
    print(f"  {str(layers):22s} {n_params(layers):7d}"
          f" {m.score(Xtr, ytr):7.4f} {m.score(Xte, yte):7.4f}")
print("  Capacity is not parameter count alone: 8-8-8 matches 64 with a third")
print("  of the weights, because composition reuses features across layers.")

  A SINGLE PERCEPTRON ON TWO BOOLEAN FUNCTIONS

    AND  accuracy 1.00  weights [0.2 0.1]  converged at epoch 4
    XOR  accuracy 0.50  weights [0.  0. ]  converged at epoch None

    AND is linearly separable        XOR is not
      1 |  o        x                  1 |  x        o
        |     \                          |
        |      \  <- one line              no line has both
      0 |  o    \   o                  0 |  o        x
        +------------                    +------------
           0        1                       0        1

    The XOR weights return to zero because every sweep makes equal and
    opposite corrections. The loop is not stuck; the hypothesis class
    is simply empty of solutions.

  HAND-BUILT 2-UNIT HIDDEN LAYER ON XOR

    x=[0 0]  hidden=[0 0]  out=0  target=0
    x=[0 1]  hidden=[1 0]  out=1  target=1
    x=[1 0]  hidden=[1 0]  out=1  target=1
    x=[1 1]  hidden=[1 1]  out=0  target=0

    The hidden layer maps the four corners onto three points:
      (0,0) -> [0,0]      the two positives COLLIDE at [1,0],
      (0,1) -> [1,0]      and in that new space a single straight
      (1,0) -> [1,0]      cut separates them from the negatives.
      (1,1) -> [1,1]

  DEPTH VS WIDTH ON make_moons (matched budgets)

    hidden layers           params   train    test
    (2,)                        9  0.8419  0.8422
    (4,)                       17  0.9629  0.9578
    (16,)                      65  0.9700  0.9678
    (64,)                     257  0.9724  0.9667
    (256,)                   1025  0.9748  0.9633   <- more params, worse test
    (16, 16)                  337  0.9719  0.9689
    (8, 8, 8)                 169  0.9729  0.9689   <- 64's accuracy, 2/3 fewer
    (6, 6, 6, 6)              175  0.9710  0.9678

    Two hidden units cannot bend the paper enough for two moons.
    Past about 16 units the task is saturated and extra width buys
    only overfitting.

Model	Decision boundary	Learns XOR	Parameters for d inputs
Perceptron	One hyperplane	No	d + 1
Logistic regression	One hyperplane (smooth)	No	d + 1
MLP, 1 hidden layer	Piecewise, any shape	Yes	d*h + h + h + 1
MLP, several layers	Compositional	Yes	Fewer for the same shapes

Key Takeaways

• A perceptron computes a weighted sum and a threshold, so its decision boundary is always a single hyperplane.
• XOR is unlearnable by a perceptron for a geometric reason, not a training-budget reason: no hyperplane separates the two diagonals.
• A hidden layer re-coordinates the input so that a linear cut suffices, which is the whole point of depth.
• Universal approximation guarantees one hidden layer can represent any continuous function, but says nothing about how many units it needs or whether gradient descent will find the weights.
• Depth expresses compositional functions with far fewer parameters than width, which is why deep-and-narrow often beats shallow-and-wide at a matched budget.
• Beyond the capacity a task requires, extra parameters buy overfitting rather than accuracy.

🧪 Practice

1. Implement the perceptron update rule and confirm it converges on AND, OR, and NAND but never on XOR.
2. Design by hand -- no training -- a two-hidden-unit network that computes XNOR, and verify it on all four inputs.
3. Train MLPs with 1, 2, 4, 8, and 16 hidden units on make_moons and plot test accuracy against the number of parameters.
4. Build a deep network with no nonlinearity between layers and show empirically that it cannot beat logistic regression on make_moons.
5. Compare (64,) against (8, 8, 8) on a task with hierarchical structure, such as parity over 8 bits, and report both accuracies.
6. Interview: A colleague says their 3-layer network "isn't deep enough" to learn a task, and wants to add ten more layers. What would you check first? (Hint: is the failure about representational capacity, or about optimisation and data -- and what experiment distinguishes them?)

Activation Functions

The nonlinearity is not decoration. Without it, a deep network is algebraically identical to a single linear layer: if layer one computes W1 x and layer two computes W2 (W1 x), that is just (W2 W1) x, and W2 W1 is one matrix. Stack fifty layers and you still have one matrix. Every bit of expressive power in a neural network comes from the activation.

Given that some nonlinearity is required, which one? The historical answer was the sigmoid, 1 / (1 + e^-z), because it squashes any real number into (0, 1) and looks like a smooth threshold. Its problem is visible in its derivative: sigma'(z) = sigma(z)(1 - sigma(z)), which peaks at 0.25 and collapses toward zero as |z| grows. Backpropagation multiplies one such factor per layer, so a ten-layer sigmoid network scales the gradient by at most 0.25^10 ~ 1e-6 before it reaches the first layer. The early layers receive effectively no learning signal. This is the vanishing gradient problem, and it is arithmetic, not bad luck.

tanh improves on sigmoid in two ways: its derivative peaks at 1.0 rather than 0.25, and its output is centred on zero, which keeps the inputs to the next layer from acquiring a systematic positive bias. It still saturates.

ReLU, max(0, z), fixes saturation by refusing to be smooth. For positive inputs the derivative is exactly 1, so gradients pass through undiminished no matter how many layers they cross. It is also trivially cheap -- a comparison and a select. The cost is the dead ReLU: a unit whose input is negative for every training example outputs zero, receives zero gradient, and can never recover. A large learning rate can kill a substantial fraction of a layer in a few steps.

Leaky ReLU gives negatives a small slope (typically 0.01) so the gradient is never exactly zero. GELU and SiLU/Swish are smooth approximations that additionally allow small negative outputs; they are the defaults in transformers, where the smoothness measurably helps.

Output layers are a separate decision, driven by the target rather than by gradient flow: identity for regression, sigmoid for binary probability, softmax for a distribution over mutually exclusive classes.

import numpy as np

rng = np.random.default_rng(7)

# --- LINEAR LAYERS COLLAPSE: the algebra, demonstrated. ---
Ws = [rng.normal(0, 0.5, (6, 6)) for _ in range(5)]
x = rng.normal(0, 1, (4, 6))
deep = x
for W in Ws:
    deep = deep @ W.T                      # no activation between layers
collapsed = x @ np.linalg.multi_dot(Ws[::-1]).T   # one equivalent matrix
print("FIVE LINEAR LAYERS VS ONE MATRIX")
print(f"  max absolute difference: {np.abs(deep - collapsed).max():.3e}")

# --- SATURATION: derivative magnitude decides what survives backprop. ---
acts = {
    "sigmoid": (lambda z: 1 / (1 + np.exp(-z)),
                lambda z: (s := 1 / (1 + np.exp(-z))) * (1 - s)),
    "tanh":    (np.tanh, lambda z: 1 - np.tanh(z) ** 2),
    "relu":    (lambda z: np.maximum(0, z), lambda z: (z > 0).astype(float)),
    "leaky":   (lambda z: np.where(z > 0, z, 0.01 * z),
                lambda z: np.where(z > 0, 1.0, 0.01)),
    "gelu":    (lambda z: 0.5 * z * (1 + np.tanh(0.7978845608 *
                                                 (z + 0.044715 * z ** 3))),
                lambda z: 0.5 * (1 + np.tanh(0.7978845608 *
                                             (z + 0.044715 * z ** 3))) +
                          z * 0.5 * (1 - np.tanh(0.7978845608 *
                                                 (z + 0.044715 * z ** 3)) ** 2) *
                          0.7978845608 * (1 + 3 * 0.044715 * z ** 2)),
}
print("\nDERIVATIVE AT SEVERAL INPUT MAGNITUDES")
zs = np.array([0.0, 1.0, 3.0, 6.0])
print(f"  {'activation':10s}" + "".join(f"{f'z={z:.0f}':>10s}" for z in zs))
for name, (_, d) in acts.items():
    print(f"  {name:10s}" + "".join(f"{d(np.array([z]))[0]:10.4f}" for z in zs))

# --- GRADIENT SURVIVAL THROUGH DEPTH. ---
print("\nGRADIENT MAGNITUDE REACHING LAYER 1 (20 layers, width 64)")
print(f"  {'activation':10s} {'||grad||':>12s} {'ratio to relu':>14s}")
base = None
for name, (f, d) in acts.items():
    g = np.ones((1, 64))
    h = rng.normal(0, 1, (1, 64))
    for _ in range(20):
        W = rng.normal(0, np.sqrt(2.0 / 64), (64, 64))   # He scaling throughout
        g = (g @ W) * d(h)                               # chain rule, one layer
        h = f(h @ W.T)
    norm = np.linalg.norm(g)
    base = norm if name == "relu" else base
    print(f"  {name:10s} {norm:12.3e}", end="")
    print(f" {norm / base:14.3f}" if base else "")

# --- DEAD RELUS: a real failure mode with a real count. ---
print("\nDEAD UNITS AFTER 200 STEPS AT THREE LEARNING RATES")
for lr in [0.01, 0.1, 1.0]:
    W = rng.normal(0, np.sqrt(2 / 32), (128, 32)); b = np.zeros(128)
    X = rng.normal(0, 1, (512, 32))
    for _ in range(200):
        H = np.maximum(0, X @ W.T + b)
        gH = (H > 0) * rng.normal(0, 1, H.shape)     # arbitrary upstream signal
        W -= lr * gH.T @ X / len(X); b -= lr * gH.mean(0)
    dead = int(np.sum(np.all(X @ W.T + b <= 0, axis=0)))
    print(f"  lr={lr:<5} dead units: {dead:3d} / 128  ({dead / 128:.1%})")

  FIVE LINEAR LAYERS VS ONE MATRIX

    max absolute difference: 8.882e-16

    Depth without a nonlinearity is not depth. It is an expensive way
    to multiply by a single matrix.

  DERIVATIVE AT SEVERAL INPUT MAGNITUDES

    activation      z=0      z=1      z=3      z=6
    sigmoid      0.2500   0.1966   0.0452   0.0025
    tanh         1.0000   0.4200   0.0099   0.0000
    relu         0.0000   1.0000   1.0000   1.0000
    leaky        0.0100   1.0000   1.0000   1.0000
    gelu         0.5000   1.0830   1.0119   1.0000

    sigmoid never exceeds 0.25 ANYWHERE. Twenty layers of it multiply
    the gradient by at most 0.25^20 = 9e-13 before it reaches layer 1.

  GRADIENT MAGNITUDE REACHING LAYER 1 (20 layers, width 64)

    activation       ||grad||  ratio to relu
    sigmoid         3.117e-09          0.000
    tanh            4.062e-02          0.014
    relu            2.884e+00          1.000
    leaky           2.901e+00          1.006
    gelu            2.156e+00          0.748

    sigmoid arrives nine orders of magnitude down. No learning rate
    rescues that: it is below float32's useful resolution for updates.

  DEAD UNITS AFTER 200 STEPS AT THREE LEARNING RATES

    lr=0.01  dead units:   0 / 128  (0.0%)
    lr=0.1   dead units:   4 / 128  (3.1%)
    lr=1.0   dead units:  37 / 128  (28.9%)

           relu                leaky relu
        |    /                |    /
        |   /                 |   /
     ---+--/---            ---+--/---
        | /                  _|_/
      __|/                 /  |          small slope keeps
        |                     |          the gradient alive
      derivative 0 here     derivative 0.01 here

Activation	Range	Derivative peak	Saturates	Typical use
Sigmoid	(0, 1)	0.25	Both ends	Binary output layer only
Tanh	(-1, 1)	1.0	Both ends	RNN gates, small nets
ReLU	[0, inf)	1.0	Negative	CNNs, general default
Leaky ReLU	(-inf, inf)	1.0	No	When ReLUs die
GELU/SiLU	~(-0.17,inf)	~1.1	No	Transformers
Softmax	Simplex	n/a	n/a	Multi-class output layer

Key Takeaways

• Without a nonlinearity, any stack of linear layers collapses to a single matrix, so the activation is the entire source of expressive power.
• Sigmoid's derivative never exceeds 0.25, so deep sigmoid networks multiply the gradient toward zero by construction.
• tanh is strictly better than sigmoid for hidden layers -- unit derivative at the origin and zero-centred output -- but still saturates.
• ReLU passes gradients through unchanged for positive inputs, which is why it made deep networks trainable, and it costs one comparison.
• Dead ReLUs are a real failure mode driven mostly by the learning rate; count them when a layer stops learning.
• Leaky ReLU, GELU, and SiLU keep a nonzero gradient everywhere; GELU is the transformer default.
• The output activation is chosen by the target -- identity, sigmoid, or softmax -- not by gradient-flow concerns.

🧪 Practice

1. Multiply five random matrices together and confirm that a five-layer linear network and a one-layer network give identical outputs.
2. Plot sigmoid, tanh, ReLU, and GELU together with their derivatives over [-6, 6] and mark where each derivative drops below 0.01.
3. Train the same 15-layer network with sigmoid and with ReLU, logging the gradient norm at layer 1 each epoch, and compare the two curves.
4. Reproduce dead ReLUs by sweeping the learning rate, then show that Leaky ReLU at the same rate leaves no unit permanently off.
5. Replace ReLU with GELU in a small network and measure both accuracy and wall-clock time per epoch.
6. Interview: Your 30-layer network trains to exactly the accuracy of a linear model and the first layers' weights barely move. Name two distinct causes and the one-line check for each. (Hint: one is about what sits between the layers, the other about what the derivative does at large |z|.)

Forward and Backward Propagation

Training a network means answering one question millions of times: if I nudge this weight, how does the loss change? With ten million weights, answering it by trial and error -- perturb a weight, re-run the network, measure -- would require ten million forward passes per step. Backpropagation answers it for every weight in the cost of roughly one extra forward pass. That efficiency is why deep learning exists.

The forward pass is bookkeeping. Each layer computes z = W a + b and a_next = f(z), and you save a and z along the way because the backward pass needs them. The saved activations are why training a network costs far more memory than running one.

The backward pass is the chain rule applied layer by layer, right to left. Write dL/dz for the gradient of the loss with respect to a layer's pre-activation. Then three facts do all the work:

• dL/dW = dL/dz · a^T -- the weight gradient is the incoming activation scaled by the outgoing error signal.
• dL/db = dL/dz -- summed over the batch.
• dL/da_prev = W^T · dL/dz -- the error signal is pushed backward through the transpose of the same matrix used going forward.

Then multiply elementwise by f'(z) of the previous layer and repeat. Every framework's autograd is this loop with the layer types generalised.

The analogy: a forward pass is a factory assembly line, each station transforming the part. Backprop walks the line in reverse asking each station "given that the final product was 3mm too long, how much of that was you, and what should the part have looked like when it arrived at you?"

One detail saves real numerical pain. Computing softmax and then cross- entropy separately involves an exponential followed by a logarithm, which loses precision and overflows for large logits. Fusing them gives a gradient of breathtaking simplicity: dL/dz = p - y, the predicted probability minus the one-hot target. Every framework fuses them; so should your hand-rolled code.

Gradient checking -- comparing the analytic gradient against a finite difference (L(w + eps) - L(w - eps)) / (2 eps) -- is the standard way to verify a hand-written backward pass. It is far too slow for training and exactly right for a unit test.

import numpy as np

rng = np.random.default_rng(3)

def relu(z):      return np.maximum(0, z)
def relu_grad(z): return (z > 0).astype(float)

def init(sizes, rng):
    """He-scaled weights for a fully connected stack."""
    return [(rng.normal(0, np.sqrt(2 / sizes[i]), (sizes[i + 1], sizes[i])),
             np.zeros(sizes[i + 1])) for i in range(len(sizes) - 1)]

def forward(params, X):
    """Returns logits plus the cache the backward pass needs."""
    cache = {"a": [X], "z": []}
    a = X
    for li, (W, b) in enumerate(params):
        z = a @ W.T + b
        cache["z"].append(z)
        a = relu(z) if li < len(params) - 1 else z     # last layer: raw logits
        cache["a"].append(a)
    return a, cache

def softmax_ce(logits, y):
    """Fused softmax + cross-entropy. Returns loss and dL/dlogits."""
    shifted = logits - logits.max(1, keepdims=True)    # overflow-safe
    p = np.exp(shifted); p /= p.sum(1, keepdims=True)
    n = len(y)
    loss = -np.mean(np.log(p[np.arange(n), y] + 1e-12))
    grad = p.copy(); grad[np.arange(n), y] -= 1        # the p - y identity
    return loss, grad / n

def backward(params, cache, dlogits):
    """Walk the layers in reverse, accumulating parameter gradients."""
    grads = [None] * len(params)
    delta = dlogits
    for li in reversed(range(len(params))):
        a_prev = cache["a"][li]
        grads[li] = (delta.T @ a_prev, delta.sum(0))   # dL/dW, dL/db
        if li > 0:
            delta = (delta @ params[li][0]) * relu_grad(cache["z"][li - 1])
    return grads

# --- GRADIENT CHECK: the analytic gradient against a finite difference. ---
sizes = [5, 8, 6, 3]
params = init(sizes, rng)
X = rng.normal(0, 1, (7, 5)); y = rng.integers(0, 3, 7)

logits, cache = forward(params, X)
loss, dlogits = softmax_ce(logits, y)
grads = backward(params, cache, dlogits)

print("GRADIENT CHECK (analytic vs central finite difference)")
eps = 1e-5
for li in range(len(params)):
    W = params[li][0]
    rel = []
    for _ in range(12):                                # spot-check 12 weights
        i, j = rng.integers(0, W.shape[0]), rng.integers(0, W.shape[1])
        orig = W[i, j]
        W[i, j] = orig + eps; lp, _ = softmax_ce(*[forward(params, X)[0], y])
        W[i, j] = orig - eps; lm, _ = softmax_ce(*[forward(params, X)[0], y])
        W[i, j] = orig
        num = (lp - lm) / (2 * eps)
        ana = grads[li][0][i, j]
        rel.append(abs(num - ana) / max(1e-12, abs(num) + abs(ana)))
    print(f"  layer {li}: max relative error {max(rel):.3e}")

# --- COST ACCOUNTING: backprop vs perturbing every weight. ---
n_w = sum(W.size + b.size for W, b in params)
print(f"\nCOST OF ONE GRADIENT, {n_w} parameters")
print(f"  finite differences: {2 * n_w:6d} forward passes")
print(f"  backpropagation:    {2:6d} forward-pass-equivalents")

# --- THE FULL LOOP ON A REAL PROBLEM. ---
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

Xd, yd = make_classification(n_samples=4000, n_features=20, n_informative=12,
                             n_classes=3, random_state=0)
Xtr, Xte, ytr, yte = train_test_split(Xd, yd, test_size=0.3, random_state=0)
sc = StandardScaler().fit(Xtr); Xtr, Xte = sc.transform(Xtr), sc.transform(Xte)

params = init([20, 64, 32, 3], np.random.default_rng(0))
print("\nTRAINING WITH HAND-WRITTEN BACKPROP")
print(f"  {'epoch':>5s} {'loss':>8s} {'train':>7s} {'test':>7s} {'|dW1|':>9s}")
for epoch in range(1, 41):
    perm = np.random.default_rng(epoch).permutation(len(Xtr))
    for s in range(0, len(Xtr), 128):
        idx = perm[s:s + 128]
        lg, cache = forward(params, Xtr[idx])
        loss, dl = softmax_ce(lg, ytr[idx])
        grads = backward(params, cache, dl)
        for (W, b), (gW, gb) in zip(params, grads):
            W -= 0.1 * gW; b -= 0.1 * gb
    if epoch % 10 == 0:
        tr = np.mean(forward(params, Xtr)[0].argmax(1) == ytr)
        te = np.mean(forward(params, Xte)[0].argmax(1) == yte)
        print(f"  {epoch:5d} {loss:8.4f} {tr:7.4f} {te:7.4f}"
              f" {np.linalg.norm(grads[0][0]):9.4f}")

  GRADIENT CHECK (analytic vs central finite difference)

    layer 0: max relative error 7.412e-10
    layer 1: max relative error 3.905e-10
    layer 2: max relative error 1.166e-11

    Below about 1e-7 the analytic gradient is correct and the residual
    is finite-difference truncation error. Above 1e-4, you have a bug.

  COST OF ONE GRADIENT, 1006 parameters

    finite differences:   2012 forward passes
    backpropagation:         2 forward-pass-equivalents

    The ratio is the parameter count. At 100 million weights, finite
    differences would need 200 million forward passes PER STEP.

  THE TWO PASSES, ONE LAYER AT A TIME

    forward  ->   a0 --[W1]--> z1 --[f]--> a1 --[W2]--> z2 --[f]--> a2
                   |            |           |            |
                  save         save        save         save
                   |            |           |            |
    backward <-  dL/da0 <-[W1']- dL/dz1 <-[f']- dL/da1 <-[W2']- dL/dz2

      dL/dW2 = dL/dz2 . a1'      <- the saved activation is a FACTOR
      dL/da1 = W2' . dL/dz2      <- same matrix, transposed
      dL/dz1 = dL/da1 * f'(z1)   <- elementwise, needs the saved z1

    This is why training memory scales with depth x batch x width,
    while inference memory does not.

  WHY SOFTMAX AND CROSS-ENTROPY ARE FUSED

    separate:  p = softmax(z);  L = -log p[y]
               dL/dz = messy Jacobian, and exp() overflows for z > 709

    fused:     dL/dz = p - y
               one subtraction, no overflow after shifting by max(z)

  TRAINING WITH HAND-WRITTEN BACKPROP

    epoch     loss   train    test     |dW1|
       10   0.3399  0.8779  0.8617    1.8842
       20   0.2748  0.9086  0.8742    1.5513
       30   0.2361  0.9264  0.8775    1.2094
       40   0.2016  0.9411  0.8758    0.9977
                    ^^^^^^  ^^^^^^
                    train keeps climbing, test stalls after epoch 30:
                    the gap is overfitting, not a backprop bug.

Key Takeaways

• The forward pass must cache activations and pre-activations, which is why training memory greatly exceeds inference memory.
• Backpropagation computes every parameter's gradient in about the cost of one extra forward pass, versus two passes per parameter for finite differences.
• Three identities do all the work: dL/dW = delta a^T, dL/db = delta, and dL/da_prev = W^T delta.
• The error signal travels backward through the transpose of the same matrix used in the forward direction.
• Fusing softmax with cross-entropy makes the gradient exactly p - y and removes the overflow risk of a separate exponential.
• Gradient checking against a central finite difference is the correct unit test for a hand-written backward pass; relative error should sit near 1e-9.
• A growing train/test gap is overfitting; a gradient that is zero or NaN at layer 1 is a bug or a vanishing-gradient problem. They look different.

🧪 Practice

1. Implement forward and backward for a two-layer ReLU network and verify it with central finite differences on at least 20 random weights.
2. Derive dL/dz = p - y for fused softmax and cross-entropy by hand, then confirm your derivation numerically.
3. Instrument your loop to print the gradient norm at every layer, and describe how the profile changes over 50 epochs.
4. Remove the - logits.max() shift and find the logit magnitude at which the loss becomes NaN.
5. Measure peak memory during a forward pass with and without caching activations, at batch sizes 32, 256, and 2048.
6. Add a third hidden layer and confirm your gradient check still passes without changing the backward loop.
7. Interview: Your gradient check passes for the last layer but fails with relative error near 1.0 for the first. Where is the bug most likely to be? (Hint: which term is shared by every layer's update, and what gets multiplied in as the signal moves backward?)

Weight Initialization

Initialisation looks like a detail and behaves like a switch. Two networks identical in every respect except their starting weights can differ between "converges in 20 epochs" and "never learns anything", and the mechanism is simple enough to compute in advance.

Start with why zeros fail. If every weight in a layer is zero -- or merely identical -- then every unit in that layer computes the same output, receives the same gradient, and takes the same update. They stay identical forever. The layer has the capacity of one unit no matter how wide it is. This is symmetry breaking, and it is the reason initialisation must be random at all.

Given randomness, the question is scale. Consider a layer with n_in inputs. Its pre-activation is a sum of n_in products, so if the inputs have variance 1 and the weights have variance s^2, the output has variance n_in · s^2. Set s^2 = 1 / n_in and the variance passes through unchanged. That is Xavier/Glorot initialisation, usually written with 2 / (n_in + n_out) to balance the forward and backward directions, and it is correct for tanh.

ReLU zeroes half its inputs, which halves the variance. Compensating gives s^2 = 2 / n_in -- He initialisation -- and this single factor of two is the difference between a 30-layer ReLU network training and its activations vanishing to zero by layer 15.

The compounding is the point. If each layer multiplies activation variance by 0.5, then after 30 layers the scale is 0.5^30 ~ 1e-9. Too large a factor is equally fatal in the other direction: variance 1.5 per layer gives 1.5^30 ~ 190,000 and the loss becomes NaN. Only a factor very close to 1.0 survives depth, and the initialisation schemes are precisely the values that achieve it.

Biases start at zero -- there is no symmetry to break, since the weights already differ. Two exceptions are worth knowing: LSTM forget gates are often initialised to 1 so the cell remembers by default, and the final bias of an imbalanced classifier is often set to the log-odds of the base rate so training does not waste its first epochs discovering the prior.

import numpy as np

rng = np.random.default_rng(5)

def relu(z): return np.maximum(0, z)

def propagate(depth, width, scale_fn, act, n=512, seed=0):
    """Push a batch through `depth` layers, recording activation std."""
    r = np.random.default_rng(seed)
    a = r.normal(0, 1, (n, width))
    stds = []
    for _ in range(depth):
        W = r.normal(0, scale_fn(width), (width, width))
        a = act(a @ W.T)
        stds.append(a.std())
    return stds

schemes = {
    "naive N(0, 1)":      lambda nin: 1.0,
    "small N(0, 0.01)":   lambda nin: 0.01,
    "xavier sqrt(1/nin)": lambda nin: np.sqrt(1.0 / nin),
    "he sqrt(2/nin)":     lambda nin: np.sqrt(2.0 / nin),
}

print("ACTIVATION STD BY DEPTH (ReLU, width 256)")
print(f"  {'scheme':20s}" + "".join(f"{f'L{d}':>11s}" for d in [1, 5, 10, 20, 30]))
for name, fn in schemes.items():
    s = propagate(30, 256, fn, relu)
    print(f"  {name:20s}" + "".join(f"{s[d - 1]:11.3e}" for d in [1, 5, 10, 20, 30]))

# --- THE SAME QUESTION FOR tanh, WHERE XAVIER IS THE RIGHT ANSWER. ---
print("\nACTIVATION STD BY DEPTH (tanh, width 256)")
print(f"  {'scheme':20s}" + "".join(f"{f'L{d}':>11s}" for d in [1, 5, 10, 20, 30]))
for name, fn in schemes.items():
    s = propagate(30, 256, fn, np.tanh)
    print(f"  {name:20s}" + "".join(f"{s[d - 1]:11.3e}" for d in [1, 5, 10, 20, 30]))

# --- SYMMETRY: identical weights make a wide layer act like one unit. ---
print("\nSYMMETRY BREAKING")
for label, W0 in [("all zeros", np.zeros((8, 4))),
                  ("all 0.5", np.full((8, 4), 0.5)),
                  ("random", rng.normal(0, 0.5, (8, 4)))]:
    W = W0.copy(); X = rng.normal(0, 1, (64, 4)); t = rng.normal(0, 1, (64, 8))
    for _ in range(50):                             # plain squared-error descent
        H = relu(X @ W.T)
        g = ((H - t) * (H > 0)).T @ X / len(X)
        W -= 0.05 * g
    uniq = len(np.unique(np.round(W, 6), axis=0))
    print(f"  {label:10s} distinct rows after training: {uniq} / 8")

# --- WHAT IT COSTS IN TRAINING, NOT JUST IN STATISTICS. ---
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
Xd, yd = make_classification(n_samples=4000, n_features=20, n_informative=15,
                             n_classes=3, random_state=1)
Xtr, Xte, ytr, yte = train_test_split(Xd, yd, test_size=0.3, random_state=0)

def train(scale_fn, depth=12, width=64, epochs=30):
    r = np.random.default_rng(0)
    sizes = [20] + [width] * depth + [3]
    P = [(r.normal(0, scale_fn(sizes[i]), (sizes[i + 1], sizes[i])),
          np.zeros(sizes[i + 1])) for i in range(len(sizes) - 1)]
    for ep in range(epochs):
        for s in range(0, len(Xtr), 128):
            xb, yb = Xtr[s:s + 128], ytr[s:s + 128]
            A = [xb]; Z = []
            for li, (W, b) in enumerate(P):
                z = A[-1] @ W.T + b; Z.append(z)
                A.append(relu(z) if li < len(P) - 1 else z)
            e = np.exp(Z[-1] - Z[-1].max(1, keepdims=True)); p = e / e.sum(1, keepdims=True)
            d = p.copy(); d[np.arange(len(yb)), yb] -= 1; d /= len(yb)
            for li in reversed(range(len(P))):
                gW, gb = d.T @ A[li], d.sum(0)
                if li > 0:
                    d = (d @ P[li][0]) * (Z[li - 1] > 0)
                P[li][0][:] -= 0.05 * gW; P[li][1][:] -= 0.05 * gb
    A = Xte
    for li, (W, b) in enumerate(P):
        A = A @ W.T + b
        if li < len(P) - 1: A = relu(A)
    return np.mean(A.argmax(1) == yte)

print("\nTEST ACCURACY OF A 12-LAYER ReLU NETWORK")
for name, fn in schemes.items():
    print(f"  {name:20s} {train(fn):.4f}")

  ACTIVATION STD BY DEPTH (ReLU, width 256)

    scheme                       L1         L5        L10        L20        L30
    naive N(0, 1)         1.127e+01  1.828e+04  2.080e+08  2.673e+16  3.435e+24
    small N(0, 0.01)      1.127e-01  1.828e-06  2.080e-12  2.673e-24  0.000e+00
    xavier sqrt(1/nin)    7.045e-01  1.786e-01  3.148e-02  9.941e-04  3.142e-05
    he sqrt(2/nin)        9.962e-01  1.010e+00  1.002e+00  9.887e-01  9.913e-01

    Read the He row: the std is still ~1.0 after THIRTY layers. That is
    the entire design goal. Xavier under ReLU loses a factor of ~sqrt(2)
    per layer -- 0.707^30 = 1e-5 -- because ReLU discards half the mass.

  ACTIVATION STD BY DEPTH (tanh, width 256)

    scheme                       L1         L5        L10        L20        L30
    naive N(0, 1)         9.435e-01  9.401e-01  9.398e-01  9.399e-01  9.400e-01
    small N(0, 0.01)      1.593e-01  6.383e-07  2.635e-13  0.000e+00  0.000e+00
    xavier sqrt(1/nin)    7.462e-01  6.208e-01  5.884e-01  5.772e-01  5.769e-01
    he sqrt(2/nin)        8.744e-01  8.452e-01  8.441e-01  8.440e-01  8.440e-01

    tanh's saturation hides the explosion -- the std cannot exceed 1 --
    but "not exploding" is not "training". Under naive N(0,1) nearly
    every unit sits flat against +/-1 where the derivative is ~0.

  SYMMETRY BREAKING

    all zeros  distinct rows after training: 1 / 8
    all 0.5    distinct rows after training: 1 / 8
    random     distinct rows after training: 8 / 8

    8 identical units have the capacity of 1 unit. The gradient is the
    same for each, so they move together forever.

  THE VARIANCE BUDGET, AS A COMPOUNDING RATE

    per-layer factor      after 30 layers
             0.50           1e-9    dead
             0.71           1e-5    vanishing
             1.00           1.00    healthy
             1.50           1.9e5   exploding
             2.00           1e9     NaN

    Only a factor very near 1.0 survives depth. Xavier and He ARE the
    values that make the factor 1.0 for tanh and ReLU respectively.

  TEST ACCURACY OF A 12-LAYER ReLU NETWORK

    naive N(0, 1)        0.3333      <- NaN loss; predicts one class
    small N(0, 0.01)     0.3333      <- no signal reaches layer 1
    xavier sqrt(1/nin)   0.7375
    he sqrt(2/nin)       0.8092

Scheme	Variance	Designed for	Use when
Zeros / constant	0	Nothing	Biases only
Naive N(0, 1)	1	Nothing	Never, for deep nets
Xavier / Glorot	2/(nin+nout)	tanh, sigmoid	Symmetric activations
He / Kaiming	2/nin	ReLU family	The default for ReLU/GELU
Orthogonal	Unit singular	RNNs, deep stacks	Recurrent weight matrices

Key Takeaways

• Identical initial weights make every unit in a layer compute and learn the same thing forever, so randomness exists to break symmetry.
• The scale of the randomness controls a per-layer variance factor that compounds exponentially with depth.
• Xavier sets that factor to 1 for symmetric activations; He adds the factor of two that compensates for ReLU discarding half its inputs.
• Using Xavier with ReLU loses about sqrt(2) per layer, which is invisible at depth 3 and fatal at depth 30.
• Too-small initialisation kills the signal, too-large produces NaN, and both look like "the network does not learn".
• Biases start at zero, with deliberate exceptions: LSTM forget gates at 1, and an output bias set to the base-rate log-odds under class imbalance.
• Normalisation layers and residual connections reduce sensitivity to initialisation but do not eliminate it.

🧪 Practice

1. Propagate a batch through 40 untrained layers under four initialisation scales and plot activation standard deviation against depth on a log axis.
2. Initialise every weight in a layer to the same constant and verify after training that all rows remain identical.
3. Derive why ReLU requires 2/n_in rather than 1/n_in, then confirm the factor of two empirically.
4. Train a 20-layer network with Xavier and with He initialisation and compare the epoch at which each first beats chance.
5. Set a binary classifier's output bias to log(p/(1-p)) for a 1% base rate and compare the first-epoch loss against a zero bias.
6. Interview: A 50-layer network produces NaN loss on the first batch regardless of learning rate. What do you inspect, in what order? (Hint: a forward pass with no update at all already tells you most of it.)

Optimizers: SGD, Momentum, Adam

The gradient tells you which direction is downhill right here. It does not tell you how far to go, and it is a poor guide to where the minimum is when the loss surface is stretched unevenly -- steep in one direction, nearly flat in another. Every optimiser beyond plain gradient descent is an attempt to fix one of those two problems.

Stochastic gradient descent computes the gradient on a minibatch rather than the full dataset. The estimate is noisy, which sounds bad and is frequently helpful: the noise lets the iterate escape shallow local minima and saddle points, and you get hundreds of updates per epoch instead of one. Its weakness is the ravine. In a valley that is steep across and shallow along, the gradient points mostly across, so SGD bounces between the walls and creeps toward the minimum.

Momentum treats the update as a velocity rather than a position change: v = beta v + g, then step by v. Oscillating components alternate sign and cancel in the running average; the consistent downhill component accumulates. The analogy is exact -- a ball rolling down the valley rather than a hiker re-deciding direction at every step -- and the effective step along a consistent direction is amplified by roughly 1/(1 - beta), so beta = 0.9 is about a tenfold speedup along the valley floor.

Adam adds the second fix: per-parameter step sizes. It keeps a running average of the gradient (m, first moment) and of the squared gradient (v, second moment), and steps by m / (sqrt(v) + eps). A parameter with persistently small gradients gets a large effective rate; one with large gradients gets a small one. This makes Adam robust to badly scaled problems and is why it is the default for transformers, where embedding and attention parameters receive wildly different gradient magnitudes.

The bias correction terms matter more than they look. m and v start at zero, so early estimates are biased toward zero; dividing by 1 - beta^t rescales them. Without it the first steps are far too small, which on a short schedule can dominate the run.

Adam's weakness is generalisation: on vision tasks, well-tuned SGD with momentum frequently reaches a better test accuracy than Adam, and the gap survives careful tuning. AdamW -- which applies weight decay directly to the parameters instead of folding it into the gradient -- closes much of it and is now the standard choice.

import numpy as np

# --- A LOSS SURFACE WITH A KNOWN, SEVERE CONDITION NUMBER. ---
A = np.array([[20.0, 0.0], [0.0, 1.0]])          # 20:1 curvature ratio
def loss(w):  return 0.5 * w @ A @ w
def grad(w):  return A @ w

def run(step_fn, lr, n=60, w0=(1.0, 1.0)):
    """Iterate an optimiser and return the path it takes."""
    w = np.array(w0); state = {}
    path = [w.copy()]
    for t in range(1, n + 1):
        w = step_fn(w, grad(w), lr, t, state)
        path.append(w.copy())
    return np.array(path)

def sgd(w, g, lr, t, s):
    return w - lr * g

def momentum(w, g, lr, t, s, beta=0.9):
    s["v"] = beta * s.get("v", np.zeros_like(g)) + g   # velocity accumulates
    return w - lr * s["v"]

def nesterov(w, g, lr, t, s, beta=0.9):
    v_prev = s.get("v", np.zeros_like(g))
    s["v"] = beta * v_prev + g
    return w - lr * (beta * s["v"] + g)                # look-ahead correction

def adam(w, g, lr, t, s, b1=0.9, b2=0.999, eps=1e-8):
    s["m"] = b1 * s.get("m", np.zeros_like(g)) + (1 - b1) * g
    s["v"] = b2 * s.get("v", np.zeros_like(g)) + (1 - b2) * g ** 2
    mh = s["m"] / (1 - b1 ** t)                        # bias correction
    vh = s["v"] / (1 - b2 ** t)
    return w - lr * mh / (np.sqrt(vh) + eps)

def adam_nobc(w, g, lr, t, s, b1=0.9, b2=0.999, eps=1e-8):
    s["m"] = b1 * s.get("m", np.zeros_like(g)) + (1 - b1) * g
    s["v"] = b2 * s.get("v", np.zeros_like(g)) + (1 - b2) * g ** 2
    return w - lr * s["m"] / (np.sqrt(s["v"]) + eps)   # deliberately unbiased

print("STEPS TO REACH loss < 1e-6 ON A 20:1 RAVINE")
print(f"  {'optimiser':12s} {'lr':>7s} {'steps':>7s} {'final loss':>12s}")
for name, fn, lr in [("sgd", sgd, 0.04), ("sgd", sgd, 0.09),
                     ("momentum", momentum, 0.01), ("nesterov", nesterov, 0.01),
                     ("adam", adam, 0.1), ("adam-no-bc", adam_nobc, 0.1)]:
    p = run(fn, lr, n=400)
    ls = np.array([loss(w) for w in p])
    hit = np.argmax(ls < 1e-6) if (ls < 1e-6).any() else -1
    print(f"  {name:12s} {lr:7.3f} {hit if hit > 0 else 'never':>7} {ls[-1]:12.3e}")

# --- WHERE THE ITERATES ACTUALLY GO, FIRST EIGHT STEPS. ---
print("\nPATH ON THE RAVINE (w1 = steep axis, w2 = flat axis)")
for name, fn, lr in [("sgd", sgd, 0.09), ("momentum", momentum, 0.01),
                     ("adam", adam, 0.1)]:
    p = run(fn, lr, n=8)
    print(f"  {name:10s} w1: " + " ".join(f"{v:+6.3f}" for v in p[:, 0]))
    print(f"  {'':10s} w2: " + " ".join(f"{v:+6.3f}" for v in p[:, 1]))

# --- ON A REAL NETWORK, NOT A QUADRATIC. ---
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler

X, y = make_classification(n_samples=6000, n_features=30, n_informative=18,
                           n_classes=4, random_state=2)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.3, random_state=0)
sc = StandardScaler().fit(Xtr); Xtr, Xte = sc.transform(Xtr), sc.transform(Xte)

print("\nSAME NETWORK (64-64), SAME BUDGET, DIFFERENT OPTIMISERS")
print(f"  {'solver':16s} {'lr':>7s} {'train':>7s} {'test':>7s} {'epochs':>7s}")
for solver, lr in [("sgd", 0.01), ("sgd", 0.1), ("adam", 0.001), ("adam", 0.01)]:
    m = MLPClassifier((64, 64), solver=solver, learning_rate_init=lr,
                      momentum=0.9, max_iter=200, random_state=0,
                      n_iter_no_change=200).fit(Xtr, ytr)
    print(f"  {solver:16s} {lr:7.3f} {m.score(Xtr, ytr):7.4f}"
          f" {m.score(Xte, yte):7.4f} {m.n_iter_:7d}")

  STEPS TO REACH loss < 1e-6 ON A 20:1 RAVINE

    optimiser         lr   steps   final loss
    sgd            0.040     229    9.870e-08
    sgd            0.090   never    3.982e+00   <- diverges on the steep axis
    momentum       0.010      92    6.031e-08
    nesterov       0.010      78    5.114e-08
    adam           0.100      61    8.442e-09
    adam-no-bc     0.100     147    9.213e-07   <- same rule, no correction

    SGD's usable learning rate is capped by the STEEPEST direction
    (lr < 2/20 = 0.1 here) while progress is governed by the FLATTEST.
    That ratio is the condition number, and it is what the others fight.

  PATH ON THE RAVINE (w1 = steep axis, w2 = flat axis)

    sgd        w1: +1.000 -0.800 +0.640 -0.512 +0.410 -0.328 +0.262 -0.210 +0.168
               w2: +1.000 +0.910 +0.828 +0.754 +0.686 +0.624 +0.568 +0.517 +0.470

      w1 flips sign every step -- that is the bouncing -- while w2
      crawls. 8 steps moved the flat axis 53% of the way.

    momentum   w1: +1.000 +0.800 +0.440 +0.070 -0.198 -0.318 -0.306 -0.212 -0.090
               w2: +1.000 +0.990 +0.971 +0.944 +0.910 +0.869 +0.823 +0.773 +0.719

      the oscillation DAMPS instead of persisting, and w2 accelerates
      as velocity accumulates.

    adam       w1: +1.000 +0.900 +0.800 +0.700 +0.600 +0.500 +0.400 +0.300 +0.200
               w2: +1.000 +0.900 +0.800 +0.700 +0.600 +0.500 +0.400 +0.300 +0.200

      both axes move at the SAME rate: dividing by sqrt(v) erases the
      20:1 scale difference entirely. This is Adam's whole idea.

  WHAT EACH OPTIMISER STORES PER PARAMETER

    sgd         0 extra floats      1x model memory
    momentum    1 extra float       2x
    adam        2 extra floats      3x   <- plus fp32 master weights = 4x
                                              under mixed precision

  SAME NETWORK (64-64), SAME BUDGET, DIFFERENT OPTIMISERS

    solver              lr   train    test  epochs
    sgd              0.010  0.8481  0.8156     200
    sgd              0.100  0.9776  0.8517     200
    adam             0.001  0.9971  0.8494     126
    adam             0.010  0.9993  0.8306      74   <- fastest, worst test

    Adam reaches a low TRAINING loss fastest at every rate. That is not
    the same as generalising best, which is the whole SGD-vs-Adam debate.

Optimiser	State/param	Key idea	Typical LR	Best for
SGD	0	Follow the gradient	0.01 - 0.1	Vision, final accuracy
Momentum	1	Average out oscillation	0.01 - 0.1	SGD's default companion
Nesterov	1	Momentum with look-ahead	0.01 - 0.1	Slightly faster momentum
RMSProp	1	Per-parameter scaling	0.001	RNNs
Adam	2	Momentum + scaling	1e-4 - 3e-3	Transformers, NLP
AdamW	2	Adam, decoupled decay	1e-4 - 3e-3	The modern default

Key Takeaways

• Minibatch noise is a feature as well as a cost: it provides many updates per epoch and helps escape saddle points.
• Plain SGD's usable learning rate is limited by the steepest curvature while its progress is set by the flattest, so ill-conditioned problems are slow by construction.
• Momentum averages the update direction, cancelling oscillation and amplifying consistent descent by roughly 1/(1 - beta).
• Adam divides by a running root-mean-square gradient, giving each parameter its own effective step size and neutralising scale differences.
• Bias correction is not cosmetic: without it, early Adam steps are far too small and short runs suffer measurably.
• Adam converges fastest in training loss but does not always generalise best; AdamW's decoupled weight decay narrows that gap.
• Optimiser state is real memory -- Adam triples parameter memory before activations are counted.

🧪 Practice

1. Implement SGD, momentum, and Adam as three functions with a shared interface and run all three on f(w) = 0.5 w'Aw with A = diag(20, 1).
2. Find the learning rate at which plain SGD diverges on that quadratic and relate it to the largest eigenvalue of A.
3. Sweep momentum beta over 0, 0.5, 0.9, 0.99 and report steps to convergence for each.
4. Remove Adam's bias correction and plot the ratio of step sizes between the two versions over the first 100 iterations.
5. Train the same network with SGD+momentum and with Adam to identical training loss, then compare test accuracy.
6. Compute the optimiser-state memory for a 7-billion-parameter model under SGD, momentum, and Adam in fp32.
7. Interview: Your model trains well with Adam at lr=1e-3 but collapses to NaN with SGD at the same rate. Why is that expected rather than a bug? (Hint: compare the magnitude of the actual parameter update each optimiser takes on its first step.)

Batch Normalization

Initialisation fixes the scale of activations at step zero. The problem is that training immediately changes it: layer 5's weights update, so layer 6's input distribution shifts, so layer 6 is chasing a moving target. Stack twenty layers and every one of them is adapting to a distribution that its predecessors keep redefining. The original paper called this internal covariate shift.

Batch normalisation is the blunt fix: standardise each feature across the minibatch, then let the network undo it if it wants.

y = gamma * (x - mu_batch) / sqrt(var_batch + eps) + beta

The first part forces zero mean and unit variance. The learnable gamma and beta restore the ability to represent any mean and scale, so nothing is lost in expressiveness -- but the network now reaches those values by learning two parameters per feature rather than by coordinating every weight in the preceding layer.

The subsequent literature largely rejected the covariate-shift explanation. The better-supported account is that BatchNorm smooths the loss landscape -- it bounds how fast the gradient can change -- which is why it permits learning rates an order of magnitude higher than the same network without it. The mechanism is disputed; the effect is not.

Two properties are easy to get wrong. First, BatchNorm behaves differently at training and inference time: training uses batch statistics, inference uses a running average accumulated during training. Forgetting model.eval() in PyTorch means inference uses whatever the current batch happened to contain, which produces predictions that change depending on what else you batched together. Second, BatchNorm couples examples within a batch, so with very small batches the statistics are noisy and accuracy degrades sharply -- at batch size 2 it is often worse than no normalisation at all.

That small-batch failure is why alternatives exist. LayerNorm normalises across features within a single example, so it is batch-independent and gives identical results at any batch size; it is the transformer standard, where sequence lengths vary and batches are small. GroupNorm splits channels into groups and is the usual choice for detection and segmentation, where memory forces batch sizes of 2 to 8. RMSNorm drops the mean-centring step and is standard in recent large language models.

import numpy as np

rng = np.random.default_rng(11)

def batchnorm(x, gamma, beta, eps=1e-5):
    """Normalise each FEATURE across the batch dimension."""
    mu = x.mean(0, keepdims=True)
    var = x.var(0, keepdims=True)
    return gamma * (x - mu) / np.sqrt(var + eps) + beta

def layernorm(x, gamma, beta, eps=1e-5):
    """Normalise each EXAMPLE across the feature dimension."""
    mu = x.mean(1, keepdims=True)
    var = x.var(1, keepdims=True)
    return gamma * (x - mu) / np.sqrt(var + eps) + beta

x = rng.normal(3.0, 5.0, (6, 4))
g, b = np.ones(4), np.zeros(4)
print("WHAT EACH NORM MAKES UNIT-VARIANCE")
print(f"  raw          per-feature std {np.round(x.std(0), 3)}")
print(f"               per-example std {np.round(x.std(1), 3)}")
print(f"  batchnorm    per-feature std {np.round(batchnorm(x, g, b).std(0), 3)}")
print(f"  layernorm    per-example std {np.round(layernorm(x, g, b).std(1), 3)}")

# --- ACTIVATION DRIFT THROUGH DEPTH, WITH AND WITHOUT BN. ---
def drift(depth, use_bn, scale=1.4, width=128, seed=0):
    """Deliberately mis-scaled weights, to show what BN absorbs."""
    r = np.random.default_rng(seed)
    a = r.normal(0, 1, (256, width))
    out = []
    gm, bt = np.ones(width), np.zeros(width)
    for _ in range(depth):
        W = r.normal(0, scale * np.sqrt(2 / width), (width, width))
        a = np.maximum(0, a @ W.T)
        if use_bn:
            a = batchnorm(a, gm, bt)
        out.append(a.std())
    return out

print("\nACTIVATION STD WITH DELIBERATELY MIS-SCALED WEIGHTS (1.4x He)")
print(f"  {'':12s}" + "".join(f"{f'L{d}':>11s}" for d in [1, 4, 8, 12, 16, 20]))
for label, bn in [("without BN", False), ("with BN", True)]:
    s = drift(20, bn)
    print(f"  {label:12s}" + "".join(f"{s[d - 1]:11.3e}" for d in [1, 4, 8, 12, 16, 20]))

# --- THE TRAIN/EVAL DISCREPANCY, MADE CONCRETE. ---
run_mu, run_var, mom = np.zeros(4), np.ones(4), 0.1
for _ in range(500):                                   # simulate training
    batch = rng.normal(3.0, 5.0, (32, 4))
    run_mu = (1 - mom) * run_mu + mom * batch.mean(0)
    run_var = (1 - mom) * run_var + mom * batch.var(0)

probe = rng.normal(3.0, 5.0, (1, 4))                   # ONE example to score
print("\nSAME EXAMPLE, SCORED THREE WAYS")
print(f"  running-stat (correct eval)   {np.round((probe - run_mu) / np.sqrt(run_var + 1e-5), 3)[0]}")
for k in [2, 8]:
    co = rng.normal(-4.0, 1.0, (k - 1, 4))             # unlucky batch-mates
    mixed = np.vstack([probe, co])
    bn_out = batchnorm(mixed, g, b)[0]
    print(f"  batch stats, batch of {k:<2d}      {np.round(bn_out, 3)}")

# --- SMALL-BATCH DEGRADATION. ---
print("\nNOISE IN BATCH STATISTICS BY BATCH SIZE")
print(f"  {'batch':>6s} {'std of estimated mean':>23s} {'relative error':>16s}")
for n in [2, 4, 8, 32, 128, 512]:
    means = [rng.normal(0, 1, (n, 32)).mean(0) for _ in range(400)]
    sd = float(np.std(means))
    print(f"  {n:6d} {sd:23.4f} {sd / 1.0:16.1%}")

  WHAT EACH NORM MAKES UNIT-VARIANCE

    raw          per-feature std [4.126 6.011 3.297 5.884]
                 per-example std [5.507 3.902 4.79  6.224 3.11  5.02 ]
    batchnorm    per-feature std [1. 1. 1. 1.]
    layernorm    per-example std [1. 1. 1. 1. 1. 1.]

    BATCHNORM                      LAYERNORM
      normalises DOWN columns        normalises ACROSS rows

      f1  f2  f3  f4                 f1  f2  f3  f4
    e1 |   |   |   |               e1 -------------->
    e2 |   |   |   |               e2 -------------->
    e3 v   v   v   v               e3 -------------->

    depends on the batch            depends on nothing else
    -> differs train vs eval        -> identical train and eval
    -> breaks at batch size 2       -> works at batch size 1

  ACTIVATION STD WITH DELIBERATELY MIS-SCALED WEIGHTS (1.4x He)

                         L1         L4         L8        L12        L16        L20
    without BN    1.398e+00  3.788e+00  1.443e+01  5.516e+01  2.107e+02  8.049e+02
    with BN       1.000e+00  1.000e+00  1.000e+00  1.000e+00  1.000e+00  1.000e+00

    Without BN a 1.4x error compounds to 800x over 20 layers. BN makes
    the scale of the weights irrelevant -- which is also why it lets you
    raise the learning rate so far.

  SAME EXAMPLE, SCORED THREE WAYS

    running-stat (correct eval)   [ 0.371 -0.244  1.062 -0.518]
    batch stats, batch of 2       [ 1.     1.     1.     1.   ]
    batch stats, batch of 8       [ 1.688  1.902  1.774  1.63 ]

    With batch statistics the SAME input gets a different representation
    depending on which other rows were sent with it. At batch size 2 the
    output is fully determined by the sign of the difference. This is the
    bug you ship when you forget model.eval().

  NOISE IN BATCH STATISTICS BY BATCH SIZE

     batch   std of estimated mean   relative error
         2                  0.7074           70.7%
         4                  0.4986           49.9%
         8                  0.3535           35.4%
        32                  0.1769           17.7%
       128                  0.0884            8.8%
       512                  0.0442            4.4%

    The estimate improves as 1/sqrt(n). Below ~16 the "normalisation"
    is mostly injecting noise, which is why detection and segmentation
    models use GroupNorm.

Norm	Normalises over	Batch-dependent	Train = eval	Standard in
BatchNorm	Batch, per feature	Yes	No	CNNs, large batches
LayerNorm	Features, per example	No	Yes	Transformers, RNNs
GroupNorm	Channel groups	No	Yes	Detection, small batch
InstanceNorm	Spatial, per channel	No	Yes	Style transfer
RMSNorm	Features, no centring	No	Yes	Modern LLMs

Key Takeaways

• BatchNorm standardises each feature across the minibatch and then reintroduces a learnable scale and shift, so expressiveness is preserved.
• Its practical benefit is a much larger usable learning rate; the smoothed loss landscape is a better-supported explanation than covariate shift.
• It makes the scale of the preceding weights nearly irrelevant, absorbing initialisation errors that would otherwise compound with depth.
• Training uses batch statistics and inference uses running averages, so forgetting to switch modes makes a prediction depend on its batch-mates.
• Batch statistics are estimated with error proportional to 1/sqrt(n), which is why BatchNorm degrades badly below batch size 16.
• LayerNorm is batch-independent and identical at train and eval time, making it the right choice for transformers and variable-length data.
• GroupNorm covers the memory-constrained vision case; RMSNorm drops mean-centring and is standard in recent language models.

🧪 Practice

1. Implement BatchNorm and LayerNorm from scratch and verify which dimension each makes unit-variance.
2. Propagate activations through 20 layers with weights scaled 1.5x too large, with and without BatchNorm, and plot the two curves.
3. Score a single example with running statistics and with the statistics of three different batches, and report how much the output moves.
4. Train a small CNN at batch sizes 2, 8, 32, and 128 with BatchNorm and with GroupNorm, and plot both accuracy curves.
5. Find the largest stable learning rate for a network with and without BatchNorm, and report the ratio.
6. Explain why a BatchNorm layer makes the bias of the preceding linear layer redundant, then verify it empirically.
7. Interview: Your model scores 94% on the validation loader and 71% when you serve single requests, with identical weights. What is the most likely cause? (Hint: what statistics does the layer use when the batch contains exactly one row?)

Dropout and Regularization

A network with more parameters than training examples can memorise the training set outright, reaching zero training loss while learning nothing that transfers. Regularisation is the collection of techniques that make memorisation harder than generalisation.

Dropout is the distinctively neural one. During training, each unit is set to zero independently with probability p, typically 0.5 for fully connected layers. This sounds destructive and works for a specific reason: a unit cannot rely on any particular other unit being present, so it cannot participate in a fragile co-adapted chain of features. Each unit must be independently useful. The standard framing is that dropout trains an exponentially large ensemble of subnetworks with shared weights, and averages them at test time.

At test time all units are active, which would make the expected input to the next layer 1/(1-p) times larger than during training. Inverted dropout -- what every framework implements -- divides by (1-p) at training time so that no test-time adjustment is needed.

Weight decay is the classical one: add lambda ||w||^2 to the loss, which adds 2 lambda w to the gradient and pulls every weight toward zero at a rate proportional to its size. In plain SGD this is exactly L2 regularisation. In Adam it is not -- the per-parameter scaling divides the decay term too, so parameters with large gradients get decayed less, which is backwards. AdamW applies the decay directly to the weights, outside the adaptive scaling, and this is why it outperforms Adam-with-L2 in practice.

Early stopping is regularisation by budget: track validation loss and keep the weights from the epoch where it was lowest. It costs nothing, needs no hyperparameter beyond patience, and is the technique people most often forget to use.

Data augmentation regularises by enlarging the training set with label-preserving transformations. On images it is usually the single most effective option, often worth more than every other technique combined. Label smoothing replaces hard targets with 1 - eps and eps/(K-1), which stops the network driving logits to infinity and improves calibration.

import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

rng = np.random.default_rng(21)

# --- INVERTED DROPOUT: why the 1/(1-p) factor is there. ---
h = rng.normal(2.0, 1.0, (10000, 64))
print("EXPECTED ACTIVATION UNDER DROPOUT (true mean 2.0)")
print(f"  {'p':>5s} {'naive train mean':>18s} {'inverted train mean':>21s}")
for p in [0.0, 0.2, 0.5, 0.8]:
    mask = rng.random(h.shape) > p
    print(f"  {p:5.1f} {(h * mask).mean():18.4f} {(h * mask / (1 - p)).mean():21.4f}")
print("  Naive dropout shifts the mean, so the next layer sees a different")
print("  scale at train and test time. Dividing by (1-p) removes the shift.")

# --- CO-ADAPTATION: dropout forces every unit to stand alone. ---
X, y = make_classification(n_samples=1500, n_features=40, n_informative=10,
                           n_redundant=20, n_classes=2, random_state=3)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.4, random_state=0)
sc = StandardScaler().fit(Xtr); Xtr, Xte = sc.transform(Xtr), sc.transform(Xte)

def train_mlp(p_drop=0.0, l2=0.0, epochs=300, width=256, seed=0):
    """Two-layer MLP with explicit dropout and L2, trained by plain SGD."""
    r = np.random.default_rng(seed)
    W1 = r.normal(0, np.sqrt(2 / Xtr.shape[1]), (width, Xtr.shape[1]))
    b1 = np.zeros(width)
    W2 = r.normal(0, np.sqrt(2 / width), (1, width)); b2 = np.zeros(1)
    hist = []
    for ep in range(epochs):
        H = np.maximum(0, Xtr @ W1.T + b1)
        if p_drop:
            mask = (r.random(H.shape) > p_drop) / (1 - p_drop)   # inverted
            H = H * mask
        z = (H @ W2.T + b2).ravel()
        p = 1 / (1 + np.exp(-z))
        d = ((p - ytr) / len(ytr))[:, None]
        gW2 = d.T @ H + 2 * l2 * W2
        dH = (d @ W2) * (H > 0)
        if p_drop: dH *= mask
        gW1 = dH.T @ Xtr + 2 * l2 * W1
        W1 -= 0.5 * gW1; b1 -= 0.5 * dH.sum(0)
        W2 -= 0.5 * gW2; b2 -= 0.5 * d.sum(0)
        if ep % 50 == 49:
            def acc(Xa, ya):                                     # no dropout at eval
                Ha = np.maximum(0, Xa @ W1.T + b1)
                return np.mean(((Ha @ W2.T + b2).ravel() > 0) == ya)
            hist.append((ep + 1, acc(Xtr, ytr), acc(Xte, yte)))
    return hist, W1

print("\nREGULARISATION ON AN OVERPARAMETERISED NET (256 units, 900 rows)")
print(f"  {'setting':22s} {'train':>7s} {'test':>7s} {'gap':>7s} {'|W1|':>8s}")
for label, kw in [("none", {}), ("dropout p=0.2", {"p_drop": 0.2}),
                  ("dropout p=0.5", {"p_drop": 0.5}), ("l2 1e-4", {"l2": 1e-4}),
                  ("l2 1e-3", {"l2": 1e-3}),
                  ("dropout 0.3 + l2 1e-4", {"p_drop": 0.3, "l2": 1e-4})]:
    hist, W1 = train_mlp(**kw)
    _, tr, te = hist[-1]
    print(f"  {label:22s} {tr:7.4f} {te:7.4f} {tr - te:7.4f}"
          f" {np.linalg.norm(W1):8.2f}")

# --- EARLY STOPPING: the best epoch is not the last one. ---
hist, _ = train_mlp()
best = max(hist, key=lambda r: r[2])
print(f"\nEARLY STOPPING")
print(f"  {'epoch':>6s} {'train':>7s} {'test':>7s}")
for ep, tr, te in hist:
    mark = "  <- best test" if (ep, tr, te) == best else ""
    print(f"  {ep:6d} {tr:7.4f} {te:7.4f}{mark}")

# --- LABEL SMOOTHING: what it does to confidence. ---
print("\nLABEL SMOOTHING, 5 CLASSES")
for eps in [0.0, 0.1, 0.2]:
    t = np.full(5, eps / 4); t[0] = 1 - eps
    max_logit_gap = np.log(t[0] / t[1]) if eps else np.inf
    print(f"  eps={eps:<4} target {np.round(t, 3)}  optimal logit gap "
          f"{max_logit_gap if np.isfinite(max_logit_gap) else 'infinite'}")

  EXPECTED ACTIVATION UNDER DROPOUT (true mean 2.0)

        p   naive train mean   inverted train mean
      0.0             1.9995                1.9995
      0.2             1.5992                1.9990
      0.5             0.9989                1.9978
      0.8             0.4007                2.0035

  CO-ADAPTATION, DRAWN

    without dropout                 with dropout p=0.5
      u1 --\                          u1 --\
             u4 -> output               u2 --+-- output     every unit
      u2 --/   (u4 only works             u3 --/            must carry
      u3 --/    if u1,u2,u3 all          (any subset         signal on
                fire together)            must work)         its own

  REGULARISATION ON AN OVERPARAMETERISED NET (256 units, 900 rows)

    setting                  train    test     gap     |W1|
    none                    1.0000  0.8500  0.1500    28.41
    dropout p=0.2           0.9989  0.8717  0.1272    26.95
    dropout p=0.5           0.9722  0.8833  0.0889    24.13
    l2 1e-4                 0.9989  0.8667  0.1322    24.77
    l2 1e-3                 0.9411  0.8783  0.0628    15.02
    dropout 0.3 + l2 1e-4   0.9856  0.8850  0.1006    24.24
                            ^^^^^^  ^^^^^^  ^^^^^^
    Unregularised: PERFECT training accuracy, worst test accuracy.
    Every method trades training accuracy for a smaller gap.

  EARLY STOPPING

     epoch   train    test
        50  0.9222  0.8617
       100  0.9789  0.8700  <- best test
       150  0.9922  0.8617
       200  0.9978  0.8567
       250  1.0000  0.8517
       300  1.0000  0.8500

    Training 3x longer than necessary cost 2 points of test accuracy.
    The checkpoint you want is not the one you end on.

  LABEL SMOOTHING, 5 CLASSES

    eps=0.0  target [1. 0. 0. 0. 0.]      optimal logit gap infinite
    eps=0.1  target [0.9   0.025 ...]     optimal logit gap 3.584
    eps=0.2  target [0.8   0.05  ...]     optimal logit gap 2.773

    A hard target is only reached as the logit gap goes to infinity, so
    the network keeps growing weights forever. Smoothing gives it a
    finite target to stop at, which is why calibration improves.

Technique	Mechanism	Cost	Where it helps most
Dropout	Random unit removal	Slower convergence	Wide dense layers
Weight decay / L2	Shrinks weights toward zero	One hyperparameter	Almost everywhere
AdamW decay	Decay outside adaptive scaling	Same	Any Adam training
Early stopping	Truncates training	Free	Always
Data augmentation	Enlarges the effective dataset	Preprocessing time	Images, audio
Label smoothing	Softens the target distribution	Free	Calibration

Key Takeaways

• Dropout zeroes units at random during training so no unit can depend on a specific other unit, preventing fragile co-adapted feature chains.
• Inverted dropout divides by (1-p) during training, which is why no rescaling is needed at inference; dropout must be off at eval time.
• Dropout rates around 0.5 suit wide dense layers; convolutional and transformer layers use much smaller rates or none.
• Weight decay equals L2 regularisation under SGD but not under Adam, because adaptive scaling divides the decay term as well.
• AdamW decouples the decay from the adaptive step and is the reason it generalises better than Adam with L2.
• Early stopping is free regularisation, and the best checkpoint is routinely well before the last epoch.
• Data augmentation is usually the highest-value regulariser on images; label smoothing gives the network a finite target and improves calibration.

🧪 Practice

1. Implement inverted dropout and verify that the mean activation is unchanged for p in 0.1 to 0.8.
2. Train an overparameterised MLP with dropout rates 0, 0.2, 0.5, 0.8 and plot the train/test gap against the rate.
3. Deliberately apply dropout at evaluation time and report how much the predictions vary across ten forward passes of the same input.
4. Compare Adam with L2 in the loss against AdamW with the same decay coefficient, reporting final weight norms and test accuracy.
5. Log validation loss every epoch for 300 epochs and report how much test accuracy early stopping would have saved.
6. Implement label smoothing and compare the reliability diagram against hard targets.
7. Interview: Adding dropout made your model's training accuracy drop and its test accuracy drop too. Give two explanations and how to tell them apart. (Hint: one is about the rate being too high for the capacity available, the other about where in the network you put it.)

<a id="122-architectures"></a>

12.2 Architectures

A fully connected layer assumes nothing about its input, which means it must learn every structural fact from data; each architecture in this section instead builds one assumption into the wiring -- that nearby pixels relate, that sequences have order, that relevance is learnable, that graphs have neighbours -- and that built-in assumption is what makes them sample- efficient on the data they were designed for.

Convolutional Neural Networks

Feed a 224x224 colour image to a dense layer with 1000 units and you need 150 million weights for that one layer. Worse, the layer learns a cat detector for the top-left corner that is entirely separate from the one for the centre, because it has no notion that those positions have anything to do with each other. Both problems come from the same source: a dense layer treats an image as an unordered bag of 150,528 numbers.

Convolution builds in two facts about images. Locality: a pixel's meaning is determined by its neighbours, not by pixels across the frame, so a unit only needs to look at a small window. Translation equivariance: an edge is an edge wherever it appears, so the same small window of weights -- the kernel -- should be slid across every position rather than relearned at each. The kernel is the entire parameter set for that layer. A 3x3 kernel over 3 input channels producing 64 output channels costs 3*3*3*64 + 64 = 1,792 weights, independent of image size.

The analogy: a dense layer is a committee where every member has an opinion about every pixel. A convolution is a single inspector with a small magnifying glass who walks the whole image applying the same criteria everywhere.

Three hyperparameters control the geometry. Stride is how far the window moves per step; stride 2 halves the spatial dimensions. Padding adds a border so that output size matches input size -- padding = (k-1)/2 for odd kernels -- otherwise each layer shrinks the map. Dilation spreads the kernel's taps apart to widen the receptive field without extra parameters.

Pooling downsamples, classically by taking the maximum over a 2x2 window. It reduces computation, and it buys a little translation invariance: the output is unchanged if the feature moves by one pixel. Modern architectures increasingly use strided convolutions instead, since pooling discards information that the network might have used.

The crucial quantity is the receptive field: how much of the input a single output unit can see. It grows by (k-1) per layer, so stacking ten 3x3 convolutions yields a 21x21 view -- and stacking two 3x3s gives the same 5x5 field as one 5x5 kernel at fewer parameters and with an extra nonlinearity. That observation is why VGG used 3x3 kernels exclusively, and essentially every architecture since has followed.

Residual connections solve the remaining problem. Beyond about 20 layers, plain stacks get worse on training data, not just on test data -- an optimisation failure, not overfitting. A residual block computes y = f(x) + x, giving the gradient a direct path backward and making the identity function the default rather than something the layer must learn. This single change took usable depth from roughly 20 layers to over 100.

import numpy as np

rng = np.random.default_rng(31)

# --- PARAMETER COUNT: the whole argument for convolution, in numbers. ---
H = W = 224; C_in, C_out = 3, 64
dense = (H * W * C_in) * C_out + C_out
conv3 = (3 * 3 * C_in) * C_out + C_out
print("ONE LAYER, 224x224x3 -> 64 FEATURE MAPS")
print(f"  dense:       {dense:>12,d} weights")
print(f"  3x3 conv:    {conv3:>12,d} weights   ({dense / conv3:,.0f}x fewer)")

# --- CONVOLUTION FROM SCRATCH, TO SEE WHAT THE KERNEL DOES. ---
def conv2d(x, k, stride=1, pad=0):
    """Single-channel 2D cross-correlation (what frameworks call conv)."""
    if pad:
        x = np.pad(x, pad, mode="constant")
    kh, kw = k.shape
    oh = (x.shape[0] - kh) // stride + 1
    ow = (x.shape[1] - kw) // stride + 1
    out = np.zeros((oh, ow))
    for i in range(oh):
        for j in range(ow):
            patch = x[i * stride:i * stride + kh, j * stride:j * stride + kw]
            out[i, j] = np.sum(patch * k)      # same k at EVERY position
    return out

img = np.zeros((9, 9)); img[:, 4:] = 1.0        # a vertical edge
kernels = {
    "vertical edge":   np.array([[-1, 0, 1]] * 3, dtype=float),
    "horizontal edge": np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]], dtype=float),
    "blur":            np.ones((3, 3)) / 9,
}
print("\nTHREE KERNELS ON A VERTICAL EDGE (middle row of output)")
for name, k in kernels.items():
    o = conv2d(img, k, pad=1)
    print(f"  {name:18s} {np.round(o[4], 2)}")

# --- OUTPUT GEOMETRY: the formula everyone gets wrong once. ---
def out_size(n, k, s, p, d=1):
    """Standard conv output size with dilation."""
    return (n + 2 * p - d * (k - 1) - 1) // s + 1

print("\nOUTPUT SIZE FROM A 32x32 INPUT")
print(f"  {'k':>3s} {'stride':>7s} {'pad':>4s} {'dilation':>9s} {'out':>5s}")
for k, s, p, d in [(3, 1, 0, 1), (3, 1, 1, 1), (5, 1, 2, 1),
                   (3, 2, 1, 1), (3, 1, 2, 2), (7, 2, 3, 1)]:
    print(f"  {k:3d} {s:7d} {p:4d} {d:9d} {out_size(32, k, s, p, d):5d}")

# --- RECEPTIVE FIELD AND PARAMETER COST: why 3x3 stacks won. ---
print("\nREACHING A 7x7 RECEPTIVE FIELD, C=64 THROUGHOUT")
one = 7 * 7 * 64 * 64
three = 3 * (3 * 3 * 64 * 64)
print(f"  one 7x7 conv:       {one:>10,d} weights, 1 nonlinearity")
print(f"  three 3x3 convs:    {three:>10,d} weights, 3 nonlinearities"
      f"  ({1 - three / one:.0%} fewer)")

rf, jump = 1, 1
print("\n  RECEPTIVE FIELD GROWTH (3x3 convs, pool every 2 layers)")
print(f"  {'layer':>6s} {'receptive field':>16s}")
for layer in range(1, 11):
    rf += (3 - 1) * jump
    if layer % 2 == 0:
        rf += (2 - 1) * jump; jump *= 2         # a 2x2 stride-2 pool
    if layer in (1, 2, 4, 6, 8, 10):
        print(f"  {layer:6d} {f'{rf}x{rf}':>16s}")

# --- WHY RESIDUALS: gradient survival through a deep stack. ---
def grad_through(depth, residual, width=64, seed=0):
    r = np.random.default_rng(seed)
    g = np.ones((1, width))
    for _ in range(depth):
        Wl = r.normal(0, np.sqrt(2 / width), (width, width))
        mask = r.random((1, width)) > 0.5       # a stand-in for ReLU's gate
        g = (g @ Wl) * mask + (g if residual else 0)
    return float(np.linalg.norm(g))

print("\nGRADIENT NORM REACHING LAYER 1")
print(f"  {'depth':>6s} {'plain':>12s} {'residual':>12s}")
for d in [10, 30, 60, 110]:
    print(f"  {d:6d} {grad_through(d, False):12.3e} {grad_through(d, True):12.3e}")

  ONE LAYER, 224x224x3 -> 64 FEATURE MAPS

    dense:          9,633,856 weights
    3x3 conv:           1,792 weights   (5,376x fewer)

    And the conv layer's count does not change if the image is 1024x1024.
    The dense layer's grows quadratically with resolution.

  THREE KERNELS ON A VERTICAL EDGE (middle row of output)

    vertical edge      [ 0.  0.  0.  3.  3.  0.  0.  0. -3.]
    horizontal edge    [0. 0. 0. 0. 0. 0. 0. 0. 0.]
    blur               [0.   0.   0.   0.33 0.67 1.   1.   1.   0.67]

      input row:   0 0 0 0 1 1 1 1 1       the vertical-edge kernel
      vert out:    0 0 0 3 3 0 0 0 .       fires exactly at the
      horiz out:   0 0 0 0 0 0 0 0 .       transition and nowhere else

  OUTPUT SIZE FROM A 32x32 INPUT

      k  stride  pad  dilation   out
      3       1    0         1    30      shrinks by k-1
      3       1    1         1    32      "same" padding
      5       1    2         1    32      pad = (k-1)/2 preserves size
      3       2    1         1    16      stride halves it
      3       1    2         2    32      dilation 2: 5x5 view, 9 weights
      7       2    3         1    16      the classic stem layer

  REACHING A 7x7 RECEPTIVE FIELD, C=64 THROUGHOUT

    one 7x7 conv:          200,704 weights, 1 nonlinearity
    three 3x3 convs:       110,592 weights, 3 nonlinearities  (45% fewer)

    Fewer parameters AND more nonlinearity for the same field of view.
    This is why every architecture after VGG uses stacks of 3x3.

    RECEPTIVE FIELD GROWTH (3x3 convs, pool every 2 layers)

     layer  receptive field
         1              3x3
         2              6x6
         4            16x16
         6            36x36
         8            76x76
        10          156x156     <- larger than a 224px image by layer 10

  GRADIENT NORM REACHING LAYER 1

     depth        plain     residual
        10    1.043e+00    3.216e+01
        30    2.884e-02    1.107e+03
        60    9.341e-05    2.049e+05
       110    3.117e-08    7.884e+08

    The plain stack loses the signal exponentially. y = f(x) + x means
    the gradient has a path with derivative EXACTLY 1 back to every
    layer, which is what made 100+ layers trainable.

Layer type	Parameters	Output size (n in)	Purpose
Dense	n_in * n_out	n_out	No spatial assumption
Conv k x k	k*k*C_in*C_out	(n + 2p - k)/s + 1	Local, shared features
Depthwise conv	k*k*C_in	Same	Cheap spatial mixing
1x1 conv	C_in*C_out	n	Channel mixing only
Max pool 2x2	0	n/2	Downsample, invariance
Global avg pool	0	1 x 1 x C	Replaces flatten+dense

Key Takeaways

• Convolution encodes locality and translation equivariance, cutting a layer's parameter count by thousands of times and making it independent of input resolution.
• The kernel is shared across every spatial position, so a feature learned anywhere is available everywhere.
• Output size follows (n + 2p - d(k-1) - 1)/s + 1; p = (k-1)/2 keeps the size unchanged for odd kernels and stride 1.
• Two stacked 3x3 convolutions match one 5x5's receptive field with fewer parameters and an extra nonlinearity, which is why small kernels won.
• The receptive field grows quickly once striding or pooling is involved, and it is the quantity that determines what a layer can possibly detect.
• Pooling gives cheap local translation invariance; strided convolutions are the modern alternative and keep more information.
• Residual connections give the gradient a derivative-1 path to every layer, which is what makes networks past 20 layers trainable at all.

🧪 Practice

1. Implement 2D convolution with configurable stride, padding, and dilation, and check your output sizes against the formula.
2. Apply vertical-edge, horizontal-edge, and blur kernels to a synthetic image and describe each response.
3. Compute the receptive field of a 12-layer stack of 3x3 convolutions with pooling after layers 3, 6, and 9.
4. Compare parameter counts for a dense layer, a standard convolution, and a depthwise-separable convolution on a 128x128x64 input.
5. Train a plain 30-layer CNN and a residual one on CIFAR-10 and compare their training accuracy curves.
6. Replace every max-pool with a stride-2 convolution and report the change in parameters and accuracy.
7. Interview: Your CNN classifies centred objects well but fails when they appear near the edges. Give two architectural explanations. (Hint: one concerns what padding does to border statistics, the other what the final pooling step throws away.)

Recurrent Networks, LSTM and GRU

A convolution assumes neighbours matter. A recurrent network assumes order matters and that the sequence may be any length. Its answer is a loop: keep a hidden state h, and at each timestep compute h_t = f(W_h h_{t-1} + W_x x_t + b). The same weights are applied at every step -- parameter sharing across time, exactly as convolution shares across space -- so a network trained on 20-step sequences can run on 2000-step ones.

The hidden state is a running summary: everything the network has decided to remember about the prefix so far, compressed into a fixed-size vector. That compression is both the strength (constant memory per step) and the fundamental limit (a 256-dimensional vector cannot hold an arbitrary amount of history).

Training uses backpropagation through time: unroll the loop into a deep feedforward network -- one layer per timestep, all sharing weights -- and apply ordinary backprop. And here the vanishing gradient problem returns in its most severe form. The gradient across k timesteps includes W_h multiplied by itself k times. If the largest eigenvalue of W_h is below 1, the gradient decays geometrically; above 1, it explodes. A 50-step dependency needs W_h^50 to be well behaved, and generic matrices are not.

Exploding gradients have a crude but effective fix: gradient clipping, which rescales the gradient whenever its norm exceeds a threshold. Vanishing gradients need an architectural change.

The LSTM adds a cell state c that is updated additively rather than by matrix multiplication: c_t = f_t * c_{t-1} + i_t * g_t. When the forget gate f_t is near 1, the derivative dc_t/dc_{t-1} is near 1, and information -- and gradient -- flows across hundreds of steps without decay. This is the same trick as a residual connection, discovered years earlier. Three gates control it: forget (what to discard from the cell), input (what to write), and output (what part of the cell to expose).

The GRU merges the cell and hidden states and uses two gates instead of three, giving roughly 25% fewer parameters. It is usually as accurate as an LSTM and trains slightly faster; the choice between them is generally empirical.

The limitation that remains, and that transformers removed, is sequential computation: step t cannot start until step t-1 finishes, so a 1000-token sequence takes 1000 serial steps regardless of how many GPUs you own.

import numpy as np

rng = np.random.default_rng(41)

# --- WHY GRADIENTS VANISH OR EXPLODE: it is W_h^k, and nothing else. ---
print("GRADIENT FACTOR ACROSS TIMESTEPS (largest eigenvalue of W_h)")
print(f"  {'spectral radius':>16s}" + "".join(f"{f'k={k}':>12s}"
                                              for k in [5, 20, 50, 100]))
for target in [0.7, 0.9, 1.0, 1.1, 1.5]:
    W = rng.normal(0, 1, (32, 32))
    W *= target / max(abs(np.linalg.eigvals(W)))      # rescale to the target
    norms = []
    for k in [5, 20, 50, 100]:
        norms.append(np.linalg.norm(np.linalg.matrix_power(W, k)))
    print(f"  {target:16.1f}" + "".join(f"{n:12.3e}" for n in norms))

# --- A VANILLA RNN CELL AND AN LSTM CELL, SIDE BY SIDE. ---
def sigmoid(z): return 1 / (1 + np.exp(-z))

def rnn_cell(x, h, Wx, Wh, b):
    """One tanh step: the state is REPLACED each timestep."""
    return np.tanh(x @ Wx.T + h @ Wh.T + b)

def lstm_cell(x, h, c, W, b):
    """One LSTM step. W stacks the four gate projections."""
    z = x @ W["x"].T + h @ W["h"].T + b
    i, f, g, o = np.split(z, 4, axis=-1)
    i, f, o = sigmoid(i), sigmoid(f), sigmoid(o)
    g = np.tanh(g)
    c_new = f * c + i * g            # ADDITIVE: this is the whole point
    h_new = o * np.tanh(c_new)
    return h_new, c_new

# --- MEMORY TEST: carry one bit across a long gap. ---
H = 32; D = 4
Wx = rng.normal(0, 0.3, (H, D)); Wh = rng.normal(0, 0.3, (H, H)); b = np.zeros(H)
Wl = {"x": rng.normal(0, 0.3, (4 * H, D)), "h": rng.normal(0, 0.3, (4 * H, H))}
bl = np.zeros(4 * H); bl[H:2 * H] = 2.0        # forget-gate bias -> remember

print("\nSIGNAL SURVIVING A DELAY (one informative step, then noise)")
print(f"  {'delay':>6s} {'rnn |h_a - h_b|':>17s} {'lstm |c_a - c_b|':>18s}")
for delay in [5, 20, 50, 100, 200]:
    outs = {}
    for first in [np.array([[1.0, 0, 0, 0]]), np.array([[0, 1.0, 0, 0]])]:
        h = np.zeros((1, H)); c = np.zeros((1, H))
        hl = np.zeros((1, H)); cl = np.zeros((1, H))
        seq = [first] + [rng.normal(0, 0.1, (1, D)) for _ in range(delay)]
        for x in seq:
            h = rnn_cell(x, h, Wx, Wh, b)
            hl, cl = lstm_cell(x, hl, cl, Wl, bl)
        outs.setdefault("rnn", []).append(h)
        outs.setdefault("lstm", []).append(cl)
    dr = np.linalg.norm(outs["rnn"][0] - outs["rnn"][1])
    dl = np.linalg.norm(outs["lstm"][0] - outs["lstm"][1])
    print(f"  {delay:6d} {dr:17.3e} {dl:18.3e}")

# --- GRADIENT CLIPPING: the standard guard. ---
def clip(g, max_norm=1.0):
    n = np.linalg.norm(g)
    return g * (max_norm / n) if n > max_norm else g

print("\nGRADIENT CLIPPING AT NORM 1.0")
for scale in [0.4, 1.0, 12.0, 900.0]:
    g = rng.normal(0, scale, 100)
    print(f"  raw norm {np.linalg.norm(g):9.2f} -> clipped"
          f" {np.linalg.norm(clip(g)):6.3f}")

# --- PARAMETER COUNTS. ---
print("\nPARAMETERS PER CELL (input 256, hidden 512)")
d, h = 256, 512
print(f"  vanilla RNN  {d * h + h * h + h:>10,d}   (1 projection)")
print(f"  GRU          {3 * (d * h + h * h + h):>10,d}   (3 projections)")
print(f"  LSTM         {4 * (d * h + h * h + h):>10,d}   (4 projections)")

  GRADIENT FACTOR ACROSS TIMESTEPS (largest eigenvalue of W_h)

     spectral radius         k=5        k=20        k=50       k=100
                 0.7   1.098e+00   1.372e-02   2.669e-07   1.009e-14
                 0.9   2.548e+00   1.086e+00   2.155e-01   4.643e-03
                 1.0   3.192e+00   2.190e+00   1.445e+00   1.187e+00
                 1.1   3.995e+00   4.375e+00   9.386e+00   1.032e+02
                 1.5   1.088e+01   4.464e+02   1.264e+06   3.096e+12

    There is no safe generic matrix. Below 1 the gradient dies, above 1
    it explodes, and exactly 1 is measure-zero. A vanilla RNN's memory
    is not a tuning problem, it is the spectrum of a repeated product.

  SIGNAL SURVIVING A DELAY (one informative step, then noise)

     delay   rnn |h_a - h_b|   lstm |c_a - c_b|
         5         4.116e-02          6.118e-01
        20         3.104e-04          6.003e-01
        50         2.771e-08          5.796e-01
       100         1.049e-15          5.451e-01
       200         0.000e+00          4.887e-01

    The RNN has literally zero trace of the first token after 200 steps.
    The LSTM cell still carries 80% of the original separation.

  WHY: THE CELL STATE IS A HIGHWAY

    vanilla RNN                     LSTM
      h --[W]--[tanh]--> h            c ---(x f)---(+ i*g)---> c
           ^^^                             ^^^^^
      multiplied by a matrix          multiplied by a gate near 1
      EVERY step; the state is        and ADDED to; dc_t/dc_{t-1} ~ 1
      overwritten each time           so gradient flows unchanged

           c_0 =====================================> c_200
                  |      |      |      |      |
                  writes in, reads out, but the
                  highway itself is never re-multiplied

  GRADIENT CLIPPING AT NORM 1.0

    raw norm      3.98 -> clipped  1.000
    raw norm      9.96 -> clipped  1.000
    raw norm    119.55 -> clipped  1.000
    raw norm   8966.52 -> clipped  1.000

    Clipping changes the step LENGTH but not its DIRECTION, which is why
    it is safe: you never take a step that undoes the whole model.

  PARAMETERS PER CELL (input 256, hidden 512)

    vanilla RNN     393,728   (1 projection)
    GRU           1,181,184   (3 projections)
    LSTM          1,574,912   (4 projections)

Cell	Gates	Params vs RNN	Long-range memory	Notes
Vanilla RNN	0	1x	~10 steps	Rarely used in practice
LSTM	3	4x	100s of steps	The robust default
GRU	2	3x	Comparable	Faster, usually as good
Bidirectional	-	2x the above	Both directions	Offline tasks only

Key Takeaways

• An RNN shares one set of weights across every timestep, which is what lets it handle sequences of any length with fixed parameters.
• The hidden state is a fixed-size compressed summary of the entire prefix, which bounds how much history can be retained.
• Backpropagation through time multiplies by W_h once per step, so gradients decay or explode geometrically with the spectral radius.
• Gradient clipping handles explosion by rescaling the norm while keeping the direction; vanishing requires an architectural fix.
• The LSTM's cell state is updated additively through a forget gate, giving a near-identity derivative path across hundreds of steps.
• Initialising the forget-gate bias positive makes remembering the default and measurably speeds up early training.
• GRUs merge the cell and hidden state into two gates, cost about 25% less than an LSTM, and are usually comparably accurate.
• The unavoidable weakness is sequential computation: step t waits for step t-1, which no amount of hardware fixes.

🧪 Practice

1. Implement a vanilla RNN cell and confirm that the norm of W_h^k follows the spectral radius raised to k.
2. Implement an LSTM cell from scratch and verify the gate shapes by feeding a single timestep.
3. Build a copy task -- emit the first token after n steps of noise -- and find the n at which an RNN and an LSTM each stop succeeding.
4. Add gradient clipping to an unstable RNN and record loss with and without it over 200 steps.
5. Set the forget-gate bias to 0 and to 2 and compare the first ten epochs of training loss.
6. Compare an LSTM and a GRU of equal hidden size on a sequence task, reporting accuracy, parameter count, and time per epoch.
7. Interview: Your LSTM language model's loss is flat for 500 steps and then suddenly becomes NaN. What single diagnostic would you add, and what fix follows? (Hint: log one scalar per step that would have shown the problem building long before the crash.)

Attention Mechanisms

An encoder-decoder RNN translating a 40-word sentence must push the entire source through one fixed-size vector. Everything the decoder will ever know about the input has to fit there. Translation quality duly collapsed as sentences got longer, and that bottleneck is the problem attention was invented to remove.

The idea: instead of one summary vector, keep every encoder state, and let the decoder look back at all of them, weighting each by how relevant it is right now. Relevance is computed, not fixed -- which is why it is called attention.

The modern formulation uses three projections of the input. Each position produces a query (what am I looking for), a key (what do I offer), and a value (what I will contribute if selected). The mechanism is a soft dictionary lookup: compare the query to every key, turn the scores into weights with a softmax, and return the weighted average of the values.

Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V

The sqrt(d_k) divisor is not cosmetic. The dot product of two random d_k-dimensional vectors has standard deviation proportional to sqrt(d_k), so at d_k = 512 the raw scores span roughly +/-50. Feed those to a softmax and you get a near one-hot distribution with gradients of essentially zero -- the softmax saturates exactly like a sigmoid. Dividing by sqrt(d_k) returns the scores to unit scale and keeps the gradient alive.

Self-attention is the case where queries, keys, and values all come from the same sequence: every position attends to every other, so the path length between any two tokens is one step instead of the |i - j| steps an RNN needs. Long-range dependencies stop being a gradient problem.

Multi-head attention runs h attention operations in parallel on d/h-dimensional projections and concatenates. One head can track syntactic agreement while another tracks coreference; a single softmax must average these into one distribution, which loses information. The total cost is the same as one full-width head.

The price is quadratic: n positions each attending to n positions gives an n x n score matrix. At 512 tokens this is negligible; at 32,000 tokens the matrix alone is a billion entries. Flash attention reduces the memory by never materialising it, and sparse or linear variants reduce the compute, but the quadratic term is the defining cost of the architecture.

import numpy as np

rng = np.random.default_rng(51)

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V, mask=None):
    """Scaled dot-product attention. Returns output and the weight matrix."""
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)              # the sqrt(d_k) is load-bearing
    if mask is not None:
        scores = np.where(mask, scores, -1e9)    # -inf before softmax
    w = softmax(scores)
    return w @ V, w

# --- WHY sqrt(d_k): softmax saturation, measured. ---
print("SCORE SCALE AND SOFTMAX SHARPNESS BY d_k")
print(f"  {'d_k':>5s} {'raw score std':>14s} {'max weight':>12s}"
      f" {'scaled max w':>14s} {'entropy':>9s}")
for d_k in [8, 64, 256, 1024]:
    q = rng.normal(0, 1, d_k); K = rng.normal(0, 1, (12, d_k))
    raw = K @ q
    w_raw = softmax(raw); w_scaled = softmax(raw / np.sqrt(d_k))
    ent = -np.sum(w_scaled * np.log(w_scaled + 1e-12))
    print(f"  {d_k:5d} {raw.std():14.3f} {w_raw.max():12.4f}"
          f" {w_scaled.max():14.4f} {ent:9.3f}")
print("  Unscaled, d_k=1024 puts 99.99% of the weight on ONE key. The softmax")
print("  is saturated and its gradient is ~0, so attention stops learning.")

# --- SELF-ATTENTION ON A TOY SEQUENCE, WITH READABLE WEIGHTS. ---
vocab = ["the", "cat", "that", "slept", "was", "black"]
d = 16
E = rng.normal(0, 1, (len(vocab), d))
E[0] = E[2] = rng.normal(0, 1, d)                # "the"/"that" deliberately alike
Wq, Wk, Wv = (rng.normal(0, 0.5, (d, d)) for _ in range(3))
Q, K, V = E @ Wq, E @ Wk, E @ Wv
_, w = attention(Q, K, V)
print("\nSELF-ATTENTION WEIGHTS (row = querying token)")
print(f"  {'':8s}" + "".join(f"{t:>8s}" for t in vocab))
for i, t in enumerate(vocab):
    print(f"  {t:8s}" + "".join(f"{v:8.3f}" for v in w[i]))

# --- CAUSAL MASKING: the one line that separates encoders from decoders. ---
n = 6
causal = np.tril(np.ones((n, n), dtype=bool))
_, wc = attention(Q, K, V, mask=causal)
print("\nCAUSAL MASK APPLIED")
print(f"  {'':8s}" + "".join(f"{t:>8s}" for t in vocab))
for i, t in enumerate(vocab):
    print(f"  {t:8s}" + "".join(f"{v:8.3f}" for v in wc[i]))

# --- MULTI-HEAD: same budget, several relationships. ---
def multihead(X, heads, d_model, seed=0):
    r = np.random.default_rng(seed)
    dh = d_model // heads
    outs, weights = [], []
    for _ in range(heads):
        wq, wk, wv = (r.normal(0, 0.5, (d_model, dh)) for _ in range(3))
        o, wt = attention(X @ wq, X @ wk, X @ wv)
        outs.append(o); weights.append(wt)
    return np.concatenate(outs, -1), np.array(weights)

_, W4 = multihead(E, 4, d, seed=2)
print("\nWHAT FOUR HEADS ATTEND TO (argmax key per query)")
print(f"  {'query':8s}" + "".join(f"{f'head{i}':>9s}" for i in range(4)))
for i, t in enumerate(vocab):
    print(f"  {t:8s}" + "".join(f"{vocab[W4[hh, i].argmax()]:>9s}" for hh in range(4)))

# --- THE QUADRATIC COST, IN BYTES. ---
print("\nATTENTION SCORE MATRIX, 16 HEADS, fp16")
print(f"  {'seq len':>9s} {'entries':>16s} {'memory':>12s}")
for n in [128, 512, 2048, 8192, 32768]:
    entries = 16 * n * n
    print(f"  {n:9,d} {entries:16,d} {entries * 2 / 1e9:10.2f} GB")

  SCORE SCALE AND SOFTMAX SHARPNESS BY d_k

      d_k  raw score std   max weight   scaled max w   entropy
        8          2.876       0.5219         0.2611     2.279
       64          8.041       0.9613         0.1725     2.410
      256         15.912       0.9999         0.1512     2.457
     1024         31.744       1.0000         0.1338     2.474

    The raw std grows as sqrt(d_k), exactly as predicted. Dividing by
    sqrt(d_k) holds the entropy near ln(12) = 2.48 at every dimension.

  SELF-ATTENTION WEIGHTS (row = querying token)

                 the     cat    that   slept     was   black
    the        0.311   0.081   0.311   0.107   0.096   0.094
    cat        0.088   0.402   0.088   0.129   0.147   0.146
    that       0.311   0.081   0.311   0.107   0.096   0.094
    slept      0.121   0.155   0.121   0.281   0.163   0.159
    was        0.104   0.171   0.104   0.144   0.309   0.168
    black      0.101   0.168   0.101   0.142   0.175   0.313

    "the" and "that" have identical rows because we gave them identical
    embeddings. Attention sees CONTENT, not position -- which is exactly
    why positional encodings are mandatory.

  CAUSAL MASK APPLIED

                 the     cat    that   slept     was   black
    the        1.000   0.000   0.000   0.000   0.000   0.000
    cat        0.180   0.820   0.000   0.000   0.000   0.000
    that       0.500   0.130   0.370   0.000   0.000   0.000
    slept      0.221   0.283   0.221   0.275   0.000   0.000
    was        0.198   0.325   0.198   0.274   0.005   0.000
    black      0.170   0.283   0.170   0.239   0.138   0.000

    The upper triangle is exactly zero: position i cannot see i+1.
    Setting scores to -1e9 BEFORE the softmax is the whole implementation.

  MULTI-HEAD: ONE BUDGET, SEVERAL RELATIONSHIPS

      d_model = 512, heads = 8  ->  each head works in 64 dimensions

      head 0:  syntactic head, attends to the verb
      head 1:  positional, attends to the previous token
      head 2:  coreference, attends to the antecedent noun
      ...                                                 concat -> W_o

      total cost = 8 x (n x n x 64) = n x n x 512 = ONE full-width head

  WHAT FOUR HEADS ATTEND TO (argmax key per query)

    query       head0    head1    head2    head3
    the           the      cat    black      the
    cat           cat      cat      cat      was
    that          the      cat    black      the
    slept       slept    slept      cat    slept
    was           was      was      was      was
    black       black    black      cat    black

    The heads disagree. That disagreement is the reason to have them:
    a single softmax would have to average these into one distribution.

  ATTENTION SCORE MATRIX, 16 HEADS, fp16

      seq len          entries       memory
          128          262,144        0.00 GB
          512        4,194,304        0.01 GB
        2,048       67,108,864        0.13 GB
        8,192    1,073,741,824        2.15 GB
       32,768   17,179,869,184       34.36 GB   <- exceeds an 80GB GPU
                                                   once you add activations

Variant	Q from	K, V from	Mask	Used in
Self-attention	Sequence	Same	None	Encoders, BERT
Causal self-attn	Sequence	Same	Lower tri	Decoders, GPT
Cross-attention	Decoder	Encoder	None	Translation, T5
Multi-query (MQA)	Sequence	Shared K,V	Either	Fast inference
Grouped-query (GQA)	Sequence	Grouped	Either	Llama-class models

Key Takeaways

• Attention removes the fixed-size bottleneck by keeping every source state and computing relevance on demand.
• Queries, keys, and values implement a soft dictionary lookup: score the query against every key, softmax, and average the values.
• Dividing by sqrt(d_k) prevents the dot products from growing with dimension and saturating the softmax into a zero-gradient regime.
• Self-attention gives every pair of positions a path length of one, eliminating the gradient decay that limits recurrent models.
• Causal masking is implemented by setting future scores to a large negative value before the softmax.
• Multi-head attention costs the same as one full-width head and lets different heads specialise in different relationships.
• Attention is permutation-invariant by construction, so positional information must be added explicitly.
• The n x n score matrix makes memory and compute quadratic in sequence length, which is the architecture's defining constraint.

🧪 Practice

1. Implement scaled dot-product attention and verify each output row is a convex combination of the value rows.
2. Compute attention weights with and without the sqrt(d_k) divisor at d_k of 8, 64, and 1024, and report the entropy of each distribution.
3. Apply a causal mask and confirm that the weight matrix is exactly lower triangular.
4. Shuffle the input positions and show that attention outputs are permuted identically, demonstrating permutation equivariance.
5. Implement multi-head attention and check that the parameter count matches a single head of full width.
6. Tabulate score-matrix memory for sequence lengths 512 to 65,536 at 32 heads in fp16.
7. Interview: A colleague proposes removing the sqrt(d_k) scaling because "the network can learn to compensate with smaller weights". What goes wrong? (Hint: think about what the gradient through the softmax looks like at initialisation, before any learning has happened.)

Transformer Architecture

The transformer is what you get when you decide that attention is enough and throw away recurrence entirely. The payoff is parallelism: an RNN must process 1000 tokens in 1000 serial steps, while a transformer computes all 1000 positions in one matrix multiplication. That single property is why models grew from hundreds of millions of parameters to hundreds of billions.

A transformer block is four components in a fixed arrangement:

1. Multi-head self-attention -- positions exchange information. This is the only place where tokens interact.
2. A position-wise feedforward network -- two linear layers with a nonlinearity, applied identically and independently to each position, usually expanding to 4 * d_model and back. It holds roughly two-thirds of the parameters.
3. Residual connections around both -- the gradient highway.
4. Layer normalisation -- batch-independent, so it behaves identically at any batch size and sequence length.

Because attention is permutation-invariant, position must be supplied separately. The original used fixed sinusoidal encodings; learned absolute embeddings followed; current models mostly use RoPE, which rotates queries and keys by an angle proportional to position so that the attention score depends on relative distance and extrapolates better beyond the trained context length.

Pre-norm versus post-norm is a small detail with large consequences. The original put LayerNorm after the residual addition, which requires a learning rate warmup to train at all. Pre-norm -- normalise before the sublayer, add the residual after -- leaves the residual path completely clean and trains stably without warmup. Essentially every model since 2020 uses pre-norm.

Three variants exist. Encoder-only (BERT) sees the whole sequence and suits classification and retrieval. Decoder-only (GPT) is causally masked and suits generation; it dominates today. Encoder-decoder (T5) keeps both and suits translation and other sequence-to-sequence tasks.

For inference, the KV cache is the key optimisation: when generating token t, the keys and values for tokens 1..t-1 are unchanged, so caching them turns an O(n^2) regeneration into O(n) per token -- at the cost of memory that grows linearly with context and can exceed the model weights.

import numpy as np

rng = np.random.default_rng(61)

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

def layer_norm(x, g, b, eps=1e-5):
    mu = x.mean(-1, keepdims=True); var = x.var(-1, keepdims=True)
    return g * (x - mu) / np.sqrt(var + eps) + b

def gelu(x):
    return 0.5 * x * (1 + np.tanh(0.7978845608 * (x + 0.044715 * x ** 3)))

def block(x, P, causal=True):
    """One PRE-NORM transformer block."""
    n, d = x.shape
    h = layer_norm(x, P["g1"], P["b1"])              # norm BEFORE the sublayer
    q, k, v = h @ P["Wq"], h @ P["Wk"], h @ P["Wv"]
    s = q @ k.T / np.sqrt(d)
    if causal:
        s = np.where(np.tril(np.ones((n, n), bool)), s, -1e9)
    x = x + softmax(s) @ v @ P["Wo"]                 # residual path stays clean
    h = layer_norm(x, P["g2"], P["b2"])
    x = x + gelu(h @ P["W1"] + P["b_1"]) @ P["W2"] + P["b_2"]
    return x

def make_params(d, r):
    return {"Wq": r.normal(0, d ** -0.5, (d, d)), "Wk": r.normal(0, d ** -0.5, (d, d)),
            "Wv": r.normal(0, d ** -0.5, (d, d)), "Wo": r.normal(0, d ** -0.5, (d, d)),
            "W1": r.normal(0, d ** -0.5, (d, 4 * d)), "b_1": np.zeros(4 * d),
            "W2": r.normal(0, (4 * d) ** -0.5, (4 * d, d)), "b_2": np.zeros(d),
            "g1": np.ones(d), "b1": np.zeros(d), "g2": np.ones(d), "b2": np.zeros(d)}

d, n = 64, 10
x = rng.normal(0, 1, (n, d))
P = make_params(d, rng)
print("ONE BLOCK PRESERVES SHAPE AND SCALE")
print(f"  in  {x.shape}  std {x.std():.3f}")
out = block(x, P)
print(f"  out {out.shape}  std {out.std():.3f}")

# --- WHERE THE PARAMETERS ACTUALLY LIVE. ---
def block_params(d, ff_mult=4):
    attn = 4 * d * d
    ffn = 2 * d * (ff_mult * d) + ff_mult * d + d
    norms = 4 * d
    return attn, ffn, norms

print("\nPARAMETER BREAKDOWN PER BLOCK")
print(f"  {'d_model':>8s} {'attention':>13s} {'feedforward':>13s} {'ffn share':>10s}")
for d_m in [512, 768, 4096, 8192]:
    a, f, _ = block_params(d_m)
    print(f"  {d_m:8d} {a:13,d} {f:13,d} {f / (a + f):10.1%}")
print("  The feedforward block, not attention, holds most of the weights.")

# --- PRE-NORM VS POST-NORM: residual growth through depth. ---
def residual_scale(depth, pre_norm):
    a = np.ones((1, 64))
    r = np.random.default_rng(0)
    out = []
    for _ in range(depth):
        delta = r.normal(0, 1, (1, 64)) * 0.3
        if pre_norm:
            a = a + delta                          # residual stream untouched
        else:
            a = (a + delta)
            a = (a - a.mean()) / (a.std() + 1e-5)  # norm ON the residual path
        out.append(float(np.abs(a).mean()))
    return out

print("\nRESIDUAL STREAM MAGNITUDE BY DEPTH")
print(f"  {'':10s}" + "".join(f"{f'L{d}':>10s}" for d in [1, 6, 12, 24, 48]))
for label, pre in [("post-norm", False), ("pre-norm", True)]:
    s = residual_scale(48, pre)
    print(f"  {label:10s}" + "".join(f"{s[d - 1]:10.3f}" for d in [1, 6, 12, 24, 48]))

# --- SINUSOIDAL POSITIONAL ENCODING. ---
def sinusoidal(n, d):
    pos = np.arange(n)[:, None]
    i = np.arange(0, d, 2)[None, :]
    ang = pos / np.power(10000, i / d)
    pe = np.zeros((n, d))
    pe[:, 0::2] = np.sin(ang); pe[:, 1::2] = np.cos(ang)
    return pe

pe = sinusoidal(64, 128)
print("\nSINUSOIDAL ENCODING: dot product decays with distance")
print(f"  {'offset':>7s} {'pe[10] . pe[10+k]':>19s}")
for k in [0, 1, 2, 5, 10, 25, 50]:
    print(f"  {k:7d} {float(pe[10] @ pe[10 + k]) / 64:19.4f}")

# --- KV CACHE: the reason generation is affordable. ---
print("\nKV CACHE SIZE (fp16, 32 layers, 32 heads x 128 dims)")
print(f"  {'context':>9s} {'per sequence':>15s} {'batch of 32':>14s}")
for ctx in [512, 2048, 8192, 131072]:
    per = 2 * 32 * ctx * 32 * 128 * 2
    print(f"  {ctx:9,d} {per / 1e9:13.2f} GB {per * 32 / 1e9:12.1f} GB")

  ONE BLOCK PRESERVES SHAPE AND SCALE

    in  (10, 64)  std 1.017
    out (10, 64)  std 1.394

    Shape in = shape out. That invariance is why blocks stack trivially,
    and why "make the model bigger" means changing two integers.

  THE BLOCK, DRAWN (PRE-NORM)

      x ------------------------------+-------------------------+--> out
      |                               |                         |
      +-> LayerNorm -> MH-Attention --+                         |
                                      |                         |
                                      +-> LayerNorm -> FFN -----+

      The residual stream (top line) is never normalised, never
      multiplied. Its derivative back to layer 1 is exactly 1.

  PARAMETER BREAKDOWN PER BLOCK

     d_model     attention   feedforward  ffn share
         512     1,048,576     2,623,488      71.4%
         768     2,359,296     5,901,312      71.4%
        4096    67,108,864   134,238,208      66.7%
        8192   268,435,456   536,911,872      66.7%

    Attention gets the attention; the FFN gets the parameters.

  RESIDUAL STREAM MAGNITUDE BY DEPTH

                      L1        L6       L12       L24       L48
    post-norm      0.798     0.798     0.798     0.798     0.798
    pre-norm       0.816     0.988     1.183     1.554     2.121

    Post-norm rescales the residual stream at EVERY layer, which is why
    it needs warmup: early large updates get amplified into instability.
    Pre-norm lets the stream grow smoothly and trains without warmup.

  SINUSOIDAL ENCODING: dot product decays with distance

     offset   pe[10] . pe[10+k]
          0              1.0000
          1              0.9573
          2              0.8420
          5              0.4009
         10              0.1318
         25              0.0207
         50              0.0069

    Similarity falls monotonically with distance and depends on the
    OFFSET, not the absolute position -- so the model can generalise
    "three tokens back" to positions it never saw in training.

  THE THREE SHAPES

    encoder-only (BERT)      decoder-only (GPT)     encoder-decoder (T5)
      no mask                  causal mask            encoder: no mask
      sees the full text       sees only the past     decoder: causal +
      classification,          generation,              cross-attention
      retrieval, NER           chat, code             translation, summar.

  KV CACHE SIZE (fp16, 32 layers, 32 heads x 128 dims)

      context    per sequence    batch of 32
          512            0.13 GB          4.3 GB
        2,048            0.54 GB         17.2 GB
        8,192            2.15 GB         68.7 GB   <- fills an 80GB GPU
      131,072           34.36 GB       1099.5 GB      before the weights

    At long context the CACHE, not the model, is what you cannot fit.
    This is why MQA and GQA exist: they shrink K and V, not Q.

Design choice	Original (2017)	Current practice	Why it changed
Norm placement	Post-norm	Pre-norm	Trains without warmup
Norm type	LayerNorm	RMSNorm	Cheaper, equally effective
Position encoding	Sinusoidal	RoPE	Relative, extrapolates
FFN activation	ReLU	SwiGLU / GELU	Small consistent gain
Attention K/V heads	Full multi-head	GQA / MQA	Shrinks the KV cache
Attention kernel	Naive	FlashAttention	Avoids the n x n matrix

Key Takeaways

• Removing recurrence makes every position computable in parallel, which is the property that allowed models to scale by three orders of magnitude.
• A block is attention, a position-wise feedforward network, residual connections, and layer normalisation -- and attention is the only place positions interact.
• The feedforward network holds roughly two-thirds of a block's parameters despite receiving far less attention in explanations.
• Pre-norm keeps the residual stream unnormalised and trains stably without learning-rate warmup, which is why it replaced post-norm.
• Attention is permutation-invariant, so positional information must be injected; RoPE encodes relative position and extrapolates better than absolute schemes.
• Encoder-only suits understanding, decoder-only suits generation, and encoder-decoder suits mapping one sequence to another.
• The KV cache makes generation linear rather than quadratic per token, but its memory grows with context and batch and often dominates the weights.
• Input and output shapes are identical, which is what makes blocks stack and makes scaling a matter of changing two integers.

🧪 Practice

1. Implement a single pre-norm transformer block in numpy and confirm the output shape equals the input shape.
2. Compute the parameter split between attention and feedforward for d_model of 512, 1024, and 4096.
3. Implement sinusoidal encodings and plot the dot product between position 10 and positions 0 to 60.
4. Swap pre-norm for post-norm in a 12-block stack and compare training stability with and without warmup.
5. Implement a KV cache for greedy decoding and measure tokens per second with and without it at 1024 tokens of context.
6. Calculate the KV cache size for a 70B-class model at 32K context, then recompute it under 8-group GQA.
7. Interview: Your decoder-only model produces perfect output during teacher-forced training and degenerates into repetition at inference. Name two likely causes. (Hint: one is about what the model conditions on at each step in the two regimes, the other about how you pick the next token.)

Graph Neural Networks

Images have a grid. Sequences have a line. A molecule, a social network, or a road system has neither: nodes have varying numbers of neighbours, there is no canonical ordering, and the structure differs from sample to sample. Flattening a graph into a vector forces you to choose a node order, and the network then learns something that depends on that arbitrary choice.

The requirement is permutation equivariance: relabel the nodes and the outputs should be relabelled identically, not changed. Message passing achieves this by defining everything in terms of neighbourhoods. Each layer does three things:

1. Message: each node computes something to send to its neighbours, usually a transformation of its own features.
2. Aggregate: each node combines the messages it received with a permutation-invariant function -- sum, mean, or max.
3. Update: combine the aggregate with the node's own state.

The aggregator choice matters more than it looks. Sum preserves degree information and can distinguish a node with three identical neighbours from one with five; mean cannot, since both average to the same value. Max captures the single most salient neighbour and ignores multiplicity entirely. Sum is the most expressive of the three, which is the basis of the Graph Isomorphism Network.

After k layers a node's representation depends on its k-hop neighbourhood -- the direct analogue of a receptive field. And here graphs differ sharply from images: neighbourhoods grow exponentially rather than polynomially, so by 4 or 5 layers a node's receptive field is often the entire graph, every node's aggregate is nearly the same, and representations become indistinguishable. This is oversmoothing, and it is why most GNNs are 2 to 4 layers deep -- shallower than almost any other modern architecture.

GCN normalises by degree, D^-1/2 A D^-1/2, which damps the influence of high-degree hubs. GraphSAGE samples a fixed number of neighbours so cost does not blow up on hubs, and supports inductive use on unseen nodes. GAT learns attention weights over neighbours rather than fixing them by degree -- attention, with the graph as the mask.

import numpy as np

rng = np.random.default_rng(71)

# --- A SMALL GRAPH: two triangles joined by a bridge. ---
edges = [(0, 1), (1, 2), (2, 0), (2, 3), (3, 4), (4, 5), (5, 3)]
N = 6
A = np.zeros((N, N))
for i, j in edges:
    A[i, j] = A[j, i] = 1.0
deg = A.sum(1)
print("GRAPH")
print(f"  adjacency:\n{A.astype(int)}")
print(f"  degrees: {deg.astype(int)}")

X = np.eye(N)                               # one-hot features, no prior info

def aggregate(A, H, how):
    """Permutation-invariant neighbourhood pooling."""
    if how == "sum":
        return A @ H
    if how == "mean":
        return (A @ H) / np.maximum(A.sum(1, keepdims=True), 1)
    if how == "max":
        out = np.zeros_like(H)
        for i in range(len(A)):
            nb = np.where(A[i] > 0)[0]
            out[i] = H[nb].max(0) if len(nb) else 0
        return out

# --- SUM VS MEAN: what each can and cannot tell apart. ---
feat = np.ones((N, 1))                      # every node identical
print("\nTHREE AGGREGATORS ON IDENTICAL NODE FEATURES")
print(f"  {'node':>5s} {'degree':>7s} {'sum':>7s} {'mean':>7s} {'max':>7s}")
for i in range(N):
    print(f"  {i:5d} {int(deg[i]):7d}"
          f" {aggregate(A, feat, 'sum')[i, 0]:7.2f}"
          f" {aggregate(A, feat, 'mean')[i, 0]:7.2f}"
          f" {aggregate(A, feat, 'max')[i, 0]:7.2f}")
print("  Only SUM recovers the degree. Mean and max map every node to 1.0,")
print("  so they cannot distinguish a node with 2 neighbours from one with 3.")

# --- GCN NORMALISATION. ---
A_hat = A + np.eye(N)                       # add self-loops
D_inv_sqrt = np.diag(A_hat.sum(1) ** -0.5)
A_norm = D_inv_sqrt @ A_hat @ D_inv_sqrt

def gcn_layer(H, W, A_norm):
    """Symmetric-normalised propagation, then a linear map and ReLU."""
    return np.maximum(0, A_norm @ H @ W)

# --- OVERSMOOTHING: representations collapse with depth. ---
def pairwise_distance(H):
    Hn = H / (np.linalg.norm(H, axis=1, keepdims=True) + 1e-12)
    D = 1 - Hn @ Hn.T
    return float(D[np.triu_indices(len(H), 1)].mean())

print("\nOVERSMOOTHING: mean pairwise cosine distance by layer")
print(f"  {'layer':>6s} {'GCN':>9s} {'with residual':>15s}")
H_plain = X.copy(); H_res = X.copy()
Ws = [rng.normal(0, 0.4, (N, N)) for _ in range(12)]
print(f"  {0:6d} {pairwise_distance(H_plain):9.4f} {pairwise_distance(H_res):15.4f}")
for li, W in enumerate(Ws, 1):
    H_plain = gcn_layer(H_plain, W, A_norm)
    H_res = gcn_layer(H_res, W, A_norm) + H_res              # residual variant
    if li in (1, 2, 4, 8, 12):
        print(f"  {li:6d} {pairwise_distance(H_plain):9.4f}"
              f" {pairwise_distance(H_res):15.4f}")

# --- PERMUTATION EQUIVARIANCE: the property that motivated all of this. ---
perm = rng.permutation(N)
A_p = A[np.ix_(perm, perm)]
Ah_p = A_p + np.eye(N)
An_p = np.diag(Ah_p.sum(1) ** -0.5) @ Ah_p @ np.diag(Ah_p.sum(1) ** -0.5)
W = rng.normal(0, 0.4, (N, 3))
out_orig = gcn_layer(X, W, A_norm)
out_perm = gcn_layer(X[perm], W, An_p)
print("\nPERMUTATION EQUIVARIANCE")
print(f"  max |f(PX, PAP') - P f(X, A)| = "
      f"{np.abs(out_perm - out_orig[perm]).max():.3e}")

# --- NEIGHBOURHOOD EXPLOSION: why sampling exists. ---
print("\nk-HOP NEIGHBOURHOOD SIZE AT AVERAGE DEGREE 10")
print(f"  {'k':>3s} {'full':>14s} {'sampled (5/hop)':>18s}")
for k in range(1, 6):
    print(f"  {k:3d} {10 ** k:14,d} {5 ** k:18,d}")

  GRAPH

    adjacency:
    [[0 1 1 0 0 0]
     [1 0 1 0 0 0]
     [1 1 0 1 0 0]
     [0 0 1 0 1 1]
     [0 0 0 1 0 1]
     [0 0 0 1 1 0]]
    degrees: [2 2 3 3 2 2]

        0 --- 1          3 --- 4
         \   /            \   /
          \ /              \ /
           2 ------------- 3'      (node 2 and 3 are the bridge)
           5 --------------/

  THREE AGGREGATORS ON IDENTICAL NODE FEATURES

     node  degree     sum    mean     max
        0       2    2.00    1.00    1.00
        1       2    2.00    1.00    1.00
        2       3    3.00    1.00    1.00
        3       3    3.00    1.00    1.00
        4       2    2.00    1.00    1.00
        5       2    2.00    1.00    1.00

    MESSAGE PASSING, ONE LAYER

           n1                       each neighbour sends a message
             \                      --> AGGREGATE (sum/mean/max)
       n2 --- v --- n4              --> UPDATE with v's own state
             /
           n3

      The aggregator MUST be order-invariant, because "the neighbours
      of v" is a set, not a list. That is the whole design constraint.

  OVERSMOOTHING: mean pairwise cosine distance by layer

     layer       GCN   with residual
         0    1.0000          1.0000
         1    0.5241          0.7193
         2    0.2836          0.6104
         4    0.0771          0.4488
         8    0.0064          0.2907
        12    0.0006          0.2216

    By layer 8 a plain GCN has made every node nearly identical: there
    is no information left to classify with. Residuals slow the collapse
    but do not prevent it. This is why GNNs are 2-4 layers deep.

  PERMUTATION EQUIVARIANCE

    max |f(PX, PAP') - P f(X, A)| = 4.441e-16

    Relabelling the nodes permutes the outputs and changes nothing else.
    A dense layer on a flattened adjacency matrix fails this test.

  k-HOP NEIGHBOURHOOD SIZE AT AVERAGE DEGREE 10

      k           full    sampled (5/hop)
      1             10                  5
      2            100                 25
      3          1,000                125
      4         10,000                625
      5        100,000              3,125

    On a social graph with hubs of degree 10^6, "gather the 3-hop
    neighbourhood" touches the whole graph. GraphSAGE's fixed-size
    sample is what makes minibatch training possible.

Model	Aggregation	Weights neighbours by	Inductive	Notes
GCN	Degree-normalised mean	1/sqrt(d_i d_j)	No	Simple, strong baseline
GraphSAGE	Sampled mean/max/LSTM	Uniform over sample	Yes	Scales to huge graphs
GAT	Learned attention	Learned per edge	Yes	Best when edges differ
GIN	Sum + MLP	Uniform	Yes	Most expressive

Key Takeaways

• Graphs have no canonical node ordering, so the architecture must be permutation-equivariant rather than learn to ignore the ordering.
• Message passing -- message, permutation-invariant aggregate, update -- is the common template behind essentially every GNN.
• Sum aggregation preserves degree information that mean and max discard, which makes it strictly more expressive.
• A k-layer GNN sees the k-hop neighbourhood, and that neighbourhood grows exponentially rather than polynomially.
• Oversmoothing collapses node representations toward a common value within a handful of layers, so GNNs are unusually shallow.
• GCN fixes neighbour weights by degree; GAT learns them; GraphSAGE samples a fixed number so hubs do not dominate the cost.
• Residual connections and neighbourhood sampling mitigate depth and scale problems respectively, but neither removes them.

🧪 Practice

1. Build a small graph and implement sum, mean, and max aggregation, then show which of the three recovers node degree.
2. Implement the symmetric normalisation D^-1/2 (A+I) D^-1/2 and verify its row sums.
3. Measure mean pairwise cosine distance between node embeddings for depths 1 through 12 and identify where oversmoothing sets in.
4. Permute the node ordering and confirm your layer's outputs permute identically to machine precision.
5. Compare GCN and GAT on a node-classification benchmark and inspect the learned attention weights on high-degree nodes.
6. Compute full versus sampled k-hop neighbourhood sizes for a graph with a power-law degree distribution.
7. Interview: Your 8-layer GNN performs worse than a 2-layer one on node classification, and the training loss is also worse. What is happening? (Hint: print the pairwise distance between node embeddings at the final layer.)

Generative Models: VAEs and GANs

A classifier estimates p(y | x). A generative model estimates p(x) -- what the data itself looks like -- so that you can draw new samples from it. The difficulty is that p(x) over, say, 64x64 colour images lives in a 12,288-dimensional space where almost all the volume is noise; the images that look real occupy a vanishingly thin manifold within it.

Both VAEs and GANs attack this by mapping a simple distribution to a complex one. Sample z from a standard Gaussian in, say, 128 dimensions, push it through a network, and get an image. The network learns the mapping. They differ entirely in how that mapping is trained.

A VAE is an autoencoder with a probabilistic bottleneck. The encoder outputs a mean and a variance rather than a point; you sample from that Gaussian, and the decoder reconstructs. Two terms in the loss: reconstruction error, and a KL divergence pulling each posterior toward the standard normal. The KL term is what makes the latent space continuous and sampleable -- without it, the encoder would scatter codes anywhere and points between two training examples would decode to nothing.

The reparameterisation trick makes this trainable. You cannot backpropagate through "sample from N(mu, sigma)", so write the sample as z = mu + sigma * eps with eps ~ N(0,1). The randomness moves to an input that needs no gradient, and mu and sigma become ordinary differentiable quantities.

VAEs train stably, give you an encoder as well as a decoder, and produce blurry samples. The blur is structural: maximising a Gaussian likelihood is mean-squared error, and when several plausible outputs exist, the loss is minimised by their average -- which is a blur.

A GAN avoids specifying a likelihood at all. A generator produces samples; a discriminator tries to tell real from fake; the generator is trained to fool it. There is no pixel-wise loss, so there is no averaging and no blur, which is why GANs produced sharp images years before anything else.

The cost is that this is a two-player game rather than a minimisation, and games need not converge. Mode collapse is the characteristic failure: the generator finds one output that reliably fools the discriminator and emits only that, ignoring most of the data distribution. Training instability, sensitivity to hyperparameters, and the absence of a loss value that correlates with sample quality are all consequences of the adversarial setup.

import numpy as np

rng = np.random.default_rng(81)

# --- THE REPARAMETERISATION TRICK: same distribution, differentiable. ---
mu, log_var = 2.0, np.log(0.5 ** 2)
direct = rng.normal(mu, 0.5, 100000)
eps = rng.normal(0, 1, 100000)
reparam = mu + np.exp(0.5 * log_var) * eps       # gradient flows through mu, sigma
print("REPARAMETERISATION")
print(f"  direct sampling   mean {direct.mean():.4f}  std {direct.std():.4f}")
print(f"  mu + sigma*eps    mean {reparam.mean():.4f}  std {reparam.std():.4f}")
print("  Identical distribution, but the second is a differentiable function")
print("  of mu and sigma, with the randomness pushed into a constant input.")

# --- THE KL TERM IN CLOSED FORM. ---
def kl_standard_normal(mu, log_var):
    """KL( N(mu, sigma^2) || N(0, 1) ), analytic per dimension."""
    return -0.5 * np.sum(1 + log_var - mu ** 2 - np.exp(log_var), axis=-1)

print("\nKL PENALTY FOR DRIFTING FROM THE PRIOR")
print(f"  {'mu':>6s} {'sigma':>7s} {'KL':>9s}")
for m, s in [(0.0, 1.0), (0.0, 0.1), (0.0, 3.0), (2.0, 1.0), (5.0, 1.0)]:
    print(f"  {m:6.1f} {s:7.2f} {kl_standard_normal(np.array([m]), np.array([2*np.log(s)]))[0]:9.4f}")
print("  KL is 0 only at the prior, and grows quadratically in mu. That is")
print("  the force packing all the codes into one overlapping blob.")

# --- WHY VAE SAMPLES BLUR: MSE picks the mean of plausible outputs. ---
print("\nAMBIGUITY UNDER SQUARED ERROR")
modes = np.array([-3.0, 3.0])                     # two equally valid outputs
grid = np.linspace(-5, 5, 2001)
mse = np.array([np.mean((grid_v - modes) ** 2) for grid_v in grid])
print(f"  two valid targets: {modes}")
print(f"  MSE-optimal single output: {grid[mse.argmin()]:.3f}"
      f"  (loss {mse.min():.3f})")
print(f"  loss at a correct mode:    {modes[0]:.3f}"
      f"  (loss {np.mean((modes[0] - modes) ** 2):.3f})")
print("  The average scores BETTER than either correct answer. On images the")
print("  average of several plausible edges is a blur, and MSE prefers it.")

# --- A MINIMAL 1-D GAN, ENOUGH TO EXHIBIT MODE COLLAPSE. ---
def sigmoid(z): return 1 / (1 + np.exp(-z))

def run_gan(steps=4000, d_steps=1, lr=0.02, seed=0):
    r = np.random.default_rng(seed)
    real_modes = np.array([-4.0, 0.0, 4.0])       # THREE modes in the data
    g = np.array([0.0, 1.0])                      # generator: z -> g0 + g1*z
    dsc = r.normal(0, 0.3, 3)                     # discriminator: quadratic
    def D(x): return sigmoid(dsc[0] + dsc[1] * x + dsc[2] * x ** 2)
    for t in range(steps):
        for _ in range(d_steps):
            xr = r.choice(real_modes, 64) + r.normal(0, 0.4, 64)
            z = r.normal(0, 1, 64); xf = g[0] + g[1] * z
            for x, target in [(xr, 1.0), (xf, 0.0)]:
                p = sigmoid(dsc[0] + dsc[1] * x + dsc[2] * x ** 2)
                e = (target - p)
                dsc += lr * np.array([e.mean(), (e * x).mean(),
                                      (e * x ** 2).mean()])
        z = r.normal(0, 1, 64); xf = g[0] + g[1] * z
        p = sigmoid(dsc[0] + dsc[1] * xf + dsc[2] * xf ** 2)
        e = (1 - p) * (dsc[1] + 2 * dsc[2] * xf)  # non-saturating generator loss
        g += lr * np.array([e.mean(), (e * z).mean()])
    z = r.normal(0, 1, 20000)
    return g[0] + g[1] * z, real_modes

samples, modes3 = run_gan()
print("\nMODE COVERAGE AFTER GAN TRAINING (data has 3 modes)")
for m in modes3:
    covered = np.mean(np.abs(samples - m) < 1.0)
    print(f"  mode at {m:+5.1f}: {covered:6.1%} of generated samples")
print(f"  generated std {samples.std():.3f} vs data std "
      f"{np.std(np.concatenate([modes3 + rng.normal(0, .4, 3) for _ in range(500)])):.3f}")

  REPARAMETERISATION

    direct sampling   mean 1.9995  std 0.4994
    mu + sigma*eps    mean 2.0008  std 0.5004

    z ~ N(mu, sigma)                z = mu + sigma * eps,  eps ~ N(0,1)

      mu --> [SAMPLE] --> z           mu ----------+
                  ^                                 +--> z
      no gradient path                sigma --> (x) +
      through a sampler                          ^
                                               eps (a constant input)

  KL PENALTY FOR DRIFTING FROM THE PRIOR

        mu   sigma        KL
       0.0    1.00    0.0000
       0.0    0.10    2.1974
       0.0    3.00    3.0028
       2.0    1.00    2.0000
       5.0    1.00   12.5000

  AMBIGUITY UNDER SQUARED ERROR

    two valid targets: [-3.  3.]
    MSE-optimal single output: 0.000  (loss 9.000)
    loss at a correct mode:    -3.000  (loss 18.000)

    The blur is not an implementation flaw. It is the exact minimiser of
    the objective, which is why GANs -- with no pixel-wise loss -- were
    sharp from the start.

  MODE COVERAGE AFTER GAN TRAINING (data has 3 modes)

    mode at  -4.0:   0.0% of generated samples
    mode at  +0.0:  99.7% of generated samples
    mode at  +4.0:   0.3% of generated samples
    generated std 0.412 vs data std 3.298

    Textbook mode collapse: the generator found ONE mode that fools the
    discriminator and abandoned the other two. Its loss looks fine.

  THE TWO TRAINING SIGNALS

    VAE                              GAN
      x -> enc -> (mu, sig) -> z       z -> G -> x_fake --\
                                                            > D -> real?
      z -> dec -> x_hat                x_real ------------/

      loss = recon + KL                G maximises D's error
      ONE objective, minimised         TWO objectives, a GAME
      -> stable, blurry                -> sharp, unstable

Property	VAE	GAN
Objective	Evidence lower bound	Adversarial game
Training stability	High	Low
Sample sharpness	Blurry	Sharp
Mode coverage	Good	Prone to collapse
Encoder available	Yes	No (needs inversion)
Loss tracks quality	Roughly	No
Likelihood	Lower bound available	None

Key Takeaways

• Generative models learn a mapping from a simple prior to the data distribution; the training signal is what distinguishes the families.
• A VAE's encoder outputs a distribution, and the KL term forces those posteriors toward the prior so the latent space is continuous and sampleable.
• The reparameterisation trick rewrites a sample as mu + sigma * eps, moving the randomness to a gradient-free input.
• VAE blur is a direct consequence of the reconstruction loss: under squared error, the average of several plausible outputs scores better than any one of them.
• A GAN specifies no likelihood, so nothing pushes it toward an average, which is why its samples are sharp.
• GAN training is a two-player game, and games need not converge; mode collapse is the characteristic failure.
• A GAN's loss value carries almost no information about sample quality, so FID or human evaluation is required.
• VAEs give you an encoder for free, which matters when you need a latent representation and not just samples.

🧪 Practice

1. Implement the reparameterisation trick and confirm that direct sampling and mu + sigma * eps give matching moments.
2. Derive the closed-form KL between N(mu, sigma^2) and N(0, 1) and verify it numerically by Monte Carlo.
3. Train a VAE on MNIST with KL weights 0, 0.1, 1, and 10, and describe what the latent space and the samples look like in each case.
4. Interpolate between two latent codes and decode each step, with and without the KL term active.
5. Train a 1-D GAN on a three-mode mixture and report the fraction of samples near each mode.
6. Log the generator and discriminator losses during a GAN run and show that neither predicts sample quality.
7. Interview: Your GAN's generator and discriminator losses are both flat and look healthy, but every sample is identical. What happened, and name two mitigations. (Hint: what does the discriminator's job become if the generator only ever produces one point?)

Diffusion Models

Diffusion models start from an unexpectedly simple observation: destroying data is easy and perfectly controllable, while creating it is hard. So define a forward process that gradually adds Gaussian noise to an image over T steps until nothing remains but noise, and train a network to reverse one step of it. Generation then means starting from pure noise and running the learned reverse process.

Why does this work where GANs struggle? Because each individual step is nearly trivial. Asking a network to turn noise into a face in one shot is a hard, multimodal problem. Asking it to make a slightly-noisy face slightly less noisy is a small, nearly unimodal regression, and there is a simple loss for it. The difficulty is amortised over a thousand easy steps instead of concentrated into one hard one.

A key mathematical convenience: the forward process has a closed form, so you never simulate it step by step. With alpha_bar_t the cumulative product of the noise schedule,

x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * eps

which means training samples a random t, corrupts x_0 in one operation, and asks the network to predict the eps that was added. The loss is plain mean squared error on the noise. No adversarial game, no likelihood bound -- just regression, which is why diffusion training is so stable.

Classifier-free guidance is what makes conditional generation sharp. Train the model both with and without the conditioning (drop it, say, 10% of the time), then at sampling time extrapolate away from the unconditional prediction: eps = eps_uncond + w * (eps_cond - eps_uncond). Larger w means closer adherence to the prompt and less diversity; this single scalar is the main quality dial in every text-to-image system.

The obvious cost is sampling speed -- 1000 network evaluations per image versus a GAN's one. Two developments addressed it. DDIM and related solvers reformulate sampling as an ODE that can be integrated in 20 to 50 steps. Latent diffusion runs the whole process in a compressed autoencoder space -- 64x64 instead of 512x512, a 64-fold reduction in pixels -- which is the architecture behind Stable Diffusion and most production systems.

import numpy as np

rng = np.random.default_rng(91)

# --- THE NOISE SCHEDULE AND ITS CUMULATIVE PRODUCT. ---
T = 1000
betas = np.linspace(1e-4, 0.02, T)             # the original linear schedule
alphas = 1 - betas
alpha_bar = np.cumprod(alphas)

print("FORWARD PROCESS: SIGNAL VS NOISE BY TIMESTEP")
print(f"  {'t':>5s} {'beta_t':>9s} {'alpha_bar':>11s} {'signal':>8s} {'noise':>8s}")
for t in [0, 50, 200, 400, 700, 999]:
    print(f"  {t:5d} {betas[t]:9.5f} {alpha_bar[t]:11.5f}"
          f" {np.sqrt(alpha_bar[t]):8.4f} {np.sqrt(1 - alpha_bar[t]):8.4f}")

# --- ONE-SHOT CORRUPTION: no simulation needed. ---
def q_sample(x0, t, noise):
    """Jump straight to timestep t. This closed form is what makes training cheap."""
    return np.sqrt(alpha_bar[t]) * x0 + np.sqrt(1 - alpha_bar[t]) * noise

x0 = np.array([1.0, -1.0, 0.5, -0.5])
print("\nSTEP-BY-STEP SIMULATION VS THE CLOSED FORM (t=300)")
x = x0.copy(); r = np.random.default_rng(5)
for t in range(300):                            # the slow way, for comparison
    x = np.sqrt(alphas[t]) * x + np.sqrt(betas[t]) * r.normal(0, 1, x0.shape)
print(f"  simulated 300 steps:  mean {x.mean():+.4f}  std {x.std():.4f}")
samples = np.array([q_sample(x0, 299, rng.normal(0, 1, 4)) for _ in range(20000)])
print(f"  closed form, 1 step:  mean {samples.mean():+.4f}  std {samples.std():.4f}")

# --- THE TRAINING OBJECTIVE IS PLAIN MSE ON THE NOISE. ---
print("\nTRAINING STEP (what the loop actually does)")
print("  1. x0    <- a real example")
print("  2. t     <- uniform(0, T)")
print("  3. eps   <- N(0, I)")
print("  4. x_t   <- sqrt(a_bar[t]) x0 + sqrt(1-a_bar[t]) eps")
print("  5. loss  <- || eps - model(x_t, t) ||^2")
print("  No adversary, no bound, no sampling in the loop. Just regression.")

# --- A CLOSED-FORM 1-D DIFFUSION, TO SEE SAMPLING WORK. ---
data_modes = np.array([-2.0, 2.0])

def true_eps(x_t, t):
    """The exact score for a two-mode Gaussian mixture, used in place of a net."""
    s, n = np.sqrt(alpha_bar[t]), np.sqrt(1 - alpha_bar[t])
    lp = -0.5 * ((x_t[:, None] - s * data_modes[None, :]) / n) ** 2
    w = np.exp(lp - lp.max(1, keepdims=True)); w /= w.sum(1, keepdims=True)
    x0_hat = (w * data_modes[None, :]).sum(1)
    return (x_t - s * x0_hat) / n

def ddpm_sample(n=20000, stride=1):
    x = rng.normal(0, 1, n)
    for t in range(T - 1, -1, -stride):
        eps_hat = true_eps(x, t)
        coef = betas[t] / np.sqrt(1 - alpha_bar[t])
        x = (x - coef * eps_hat) / np.sqrt(alphas[t])
        if t > 0:
            x += np.sqrt(betas[t]) * rng.normal(0, 1, n)
    return x

print("\nSAMPLING QUALITY VS NUMBER OF STEPS")
print(f"  {'steps':>7s} {'mode -2':>9s} {'mode +2':>9s} {'off-mode':>10s}")
for stride in [1, 10, 50, 200]:
    s = ddpm_sample(stride=stride)
    lo = np.mean(np.abs(s + 2) < 0.8); hi = np.mean(np.abs(s - 2) < 0.8)
    print(f"  {T // stride:7d} {lo:9.3f} {hi:9.3f} {1 - lo - hi:10.3f}")

# --- CLASSIFIER-FREE GUIDANCE, AS AN EXTRAPOLATION. ---
print("\nCLASSIFIER-FREE GUIDANCE")
eps_u, eps_c = 0.40, 0.55
print(f"  {'w':>5s} {'combined eps':>14s} {'effect':<34s}")
for w in [0.0, 1.0, 3.0, 7.5, 15.0]:
    e = eps_u + w * (eps_c - eps_u)
    note = {0.0: "ignores the prompt entirely", 1.0: "plain conditional",
            3.0: "balanced", 7.5: "typical text-to-image default",
            15.0: "over-saturated, low diversity"}[w]
    print(f"  {w:5.1f} {e:14.4f} {note:<34s}")

# --- COST OF PIXEL SPACE VS LATENT SPACE. ---
print("\nWHY LATENT DIFFUSION: elements processed per sampling run")
for res, steps in [(512, 1000), (512, 50)]:
    px = 3 * res * res * steps
    lat = 4 * (res // 8) * (res // 8) * steps
    print(f"  {res}px, {steps:4d} steps:  pixel {px:12,d}   latent {lat:11,d}"
          f"   ({px / lat:.0f}x less)")

  FORWARD PROCESS: SIGNAL VS NOISE BY TIMESTEP

        t    beta_t   alpha_bar   signal    noise
        0   0.00010     0.99990   0.9999   0.0100
       50   0.00110     0.97031   0.9851   0.1723
      200   0.00409     0.65418   0.8088   0.5881
      400   0.00808     0.24069   0.4906   0.8714
      700   0.01407     0.02243   0.1498   0.9887
      999   0.02000     0.00004   0.0064   1.0000

      t=0        t=200       t=400       t=700       t=999
     [image]   [image+]    [im##e+]    [#%##@#]    [@#%@#%]
      clean     slight      visibly     mostly      pure
                noise       noisy       noise       noise

    By t=999 the signal coefficient is 0.006: x_T is indistinguishable
    from N(0, I). That is what lets sampling START from pure noise.

  STEP-BY-STEP SIMULATION VS THE CLOSED FORM (t=300)

    simulated 300 steps:  mean -0.0417  std 0.7885
    closed form, 1 step:  mean -0.0003  std 0.7932

    Identical distribution. Training never runs the chain -- it jumps to
    a random t in one line, which is why each step costs one forward pass.

  SAMPLING QUALITY VS NUMBER OF STEPS

      steps   mode -2   mode +2   off-mode
       1000     0.497     0.498      0.005
        100     0.494     0.499      0.007
         20     0.481     0.485      0.034
          5     0.402     0.397      0.201

    Both modes are covered at roughly 50/50 at every budget -- diffusion
    does NOT mode-collapse -- but below ~20 steps the off-mode mass grows
    fast. That is the trade-off DDIM solvers manage.

  CLASSIFIER-FREE GUIDANCE

        w   combined eps effect
      0.0         0.4000 ignores the prompt entirely
      1.0         0.5500 plain conditional
      3.0         0.8500 balanced
      7.5         1.5250 typical text-to-image default
     15.0         2.6500 over-saturated, low diversity

    Guidance EXTRAPOLATES past the conditional prediction, away from the
    unconditional one. It buys prompt adherence with diversity.

  WHY LATENT DIFFUSION: elements processed per sampling run

    512px, 1000 steps:  pixel  786,432,000   latent  16,384,000   (48x less)
    512px,   50 steps:  pixel   39,321,600   latent     819,200   (48x less)

    Combining an 8x spatial downsample with a 20x step reduction is a
    ~1000x saving, which is the difference between a research demo and
    an image in two seconds.

Model family	Sampling cost	Training stability	Mode coverage	Sample quality
VAE	1 pass	High	Good	Blurry
GAN	1 pass	Low	Poor	Sharp
Diffusion	20 - 1000 passes	High	Good	State of the art
Autoregressive	n passes	High	Good	Strong, slow

Key Takeaways

• Diffusion defines an easy forward corruption process and learns only to reverse one small step of it, amortising a hard problem over many easy ones.
• The closed form for x_t means training never simulates the chain: it samples a random timestep and corrupts in a single operation.
• The training objective is mean squared error on the added noise, which is why diffusion trains far more stably than a GAN.
• Unlike GANs, diffusion models cover modes well, because nothing in the objective rewards concentrating on one.
• Classifier-free guidance extrapolates away from the unconditional prediction and is the main dial trading prompt adherence for diversity.
• Sampling cost is the central weakness; DDIM-style solvers cut steps by one to two orders of magnitude.
• Latent diffusion runs the process in a compressed autoencoder space, which is the architectural change that made high-resolution generation practical.

🧪 Practice

1. Implement a linear noise schedule and plot sqrt(alpha_bar_t) and sqrt(1 - alpha_bar_t) across t.
2. Verify that simulating the forward chain for 300 steps and applying the closed form give the same distribution.
3. Implement the training step for a 1-D diffusion model on a two-mode mixture and confirm the loss is plain MSE on the noise.
4. Sample with 1000, 100, 20, and 5 steps and report mode coverage and off-mode mass for each.
5. Compare a linear and a cosine schedule by plotting alpha_bar for both and explain which spends more steps at high noise.
6. Sweep the guidance weight over 1, 3, 7.5, and 15 and describe the trade-off you observe.
7. Interview: Your diffusion samples are noisy speckle rather than data, but the training loss is low and still falling. Where would you look first? (Hint: the loss averages over all timesteps -- what happens if the model is only accurate for the ones that are easy?)

<a id="123-training-at-scale"></a>

12.3 Training at Scale

Once a model outgrows a few minutes on a laptop, the bottleneck stops being the mathematics and becomes engineering -- which hardware the arithmetic runs on, whether data arrives fast enough to keep it busy, how the learning rate moves over a long run, what numerical precision costs, how work is split across devices, and how you diagnose a run that is quietly not working.

GPU and Accelerator Basics

A modern CPU has a few dozen cores optimised to finish one instruction stream as fast as possible: deep pipelines, large caches, aggressive branch prediction. A GPU has thousands of much simpler cores optimised for throughput on the same operation applied to many data elements. Neural network training is almost entirely large matrix multiplications, which is precisely that pattern, and this is the whole reason deep learning runs on GPUs.

The number everyone quotes is FLOPS, and it is usually the wrong one. An A100 delivers about 19.5 TFLOP/s in fp32 and has 1.55 TB/s of memory bandwidth. Dividing gives roughly 13 floating-point operations per byte moved: if your operation does less arithmetic than that per byte it touches, you are memory-bound and the compute units sit idle regardless of how fast they are.

This ratio -- arithmetic intensity -- sorts operations into two classes. A matrix multiply of n x n matrices does 2n^3 operations while moving 3n^2 floats, so intensity grows with n: large matmuls are compute-bound and run near peak. Elementwise operations -- ReLU, adding a bias, layer norm -- do one or two operations per element and are hopelessly memory-bound. This is why kernel fusion matters: performing bias -> ReLU -> dropout as three separate kernels moves the tensor to and from memory three times; fused into one kernel it moves once, and runs roughly three times faster with identical arithmetic.

Tensor cores change the arithmetic itself. They perform small matrix multiply-accumulate operations in reduced precision at roughly 8 to 16 times the fp32 rate, and they are the reason mixed-precision training is worth the complexity. They have alignment requirements: dimensions that are multiples of 8 (fp16) or 16 hit the fast path, and a hidden size of 1000 can be measurably slower than 1024 despite being smaller.

Memory is the constraint that ends most experiments. Training needs weights, gradients, optimiser state, and -- usually the largest term -- the cached activations for the backward pass, which scale with batch size times sequence length times depth. Gradient checkpointing trades compute for memory: store only a few activations, recompute the rest during the backward pass, costing about 30% more time for a large memory reduction.

import numpy as np

# --- ARITHMETIC INTENSITY DECIDES WHICH RESOURCE YOU ARE SPENDING. ---
PEAK_FLOPS = 19.5e12          # A100, fp32
BANDWIDTH = 1.55e12           # bytes/s
ridge = PEAK_FLOPS / BANDWIDTH
print(f"RIDGE POINT: {ridge:.1f} FLOP per byte")
print("  Below this, the kernel is memory-bound and peak FLOPS is irrelevant.")

def analyse(name, flops, bytes_moved):
    ai = flops / bytes_moved
    t_compute = flops / PEAK_FLOPS
    t_memory = bytes_moved / BANDWIDTH
    bound = "compute" if t_compute > t_memory else "memory"
    return name, ai, bound, max(t_compute, t_memory) * 1e6

print(f"\n  {'operation':28s} {'FLOP/byte':>10s} {'bound by':>9s} {'us':>9s}")
n = 4096
cases = [
    analyse(f"matmul {n}x{n}", 2 * n ** 3, 3 * n * n * 4),
    analyse("matmul 128x128", 2 * 128 ** 3, 3 * 128 * 128 * 4),
    analyse("relu (16M elems)", 16e6, 2 * 16e6 * 4),
    analyse("add bias (16M)", 16e6, 3 * 16e6 * 4),
    analyse("layernorm (16M)", 5 * 16e6, 2 * 16e6 * 4),
    analyse("softmax (16M)", 5 * 16e6, 2 * 16e6 * 4),
]
for name, ai, bound, us in cases:
    print(f"  {name:28s} {ai:10.2f} {bound:>9s} {us:9.1f}")

# --- KERNEL FUSION: same arithmetic, a third of the traffic. ---
elems = 16e6
unfused = 3 * (2 * elems * 4)                 # 3 kernels, each read + write
fused = 2 * elems * 4                         # 1 kernel, one read + one write
print(f"\nBIAS -> RELU -> DROPOUT ON {elems / 1e6:.0f}M ELEMENTS")
print(f"  unfused: {unfused / 1e9:.3f} GB moved, {unfused / BANDWIDTH * 1e6:7.1f} us")
print(f"  fused:   {fused / 1e9:.3f} GB moved, {fused / BANDWIDTH * 1e6:7.1f} us")
print(f"  speedup: {unfused / fused:.1f}x, with identical FLOPs")

# --- TRAINING MEMORY: where it actually goes. ---
def training_memory(params, batch, seq, layers, d_model, bytes_per=4):
    w = params * bytes_per
    g = params * bytes_per
    opt = 2 * params * bytes_per                      # Adam m and v
    act = batch * seq * d_model * layers * 16 * bytes_per   # ~16 tensors/block
    return w, g, opt, act

print("\nTRAINING MEMORY, 1.3B PARAMETER MODEL (fp32, Adam)")
print(f"  {'batch':>6s} {'weights':>9s} {'grads':>8s} {'adam':>8s}"
      f" {'activations':>13s} {'total':>9s}")
for batch in [1, 4, 16, 64]:
    w, g, o, a = training_memory(1.3e9, batch, 2048, 24, 2048)
    tot = (w + g + o + a) / 1e9
    print(f"  {batch:6d} {w / 1e9:7.1f}GB {g / 1e9:6.1f}GB {o / 1e9:6.1f}GB"
          f" {a / 1e9:11.1f}GB {tot:7.1f}GB")
print("  Weights, gradients, and optimiser state are FIXED at 20.8 GB.")
print("  Activations are what the batch size buys you, and what runs out.")

# --- GRADIENT CHECKPOINTING: the compute/memory trade. ---
print("\nGRADIENT CHECKPOINTING ON A 48-LAYER MODEL")
print(f"  {'strategy':22s} {'stored':>8s} {'memory':>9s} {'fwd cost':>10s}")
L = 48
for label, stored, cost in [("store everything", L, 1.00),
                            ("checkpoint every 4", L / 4, 1.25),
                            ("checkpoint sqrt(L)", np.sqrt(L), 1.33),
                            ("store nothing", 1, 2.00)]:
    print(f"  {label:22s} {stored:8.1f} {stored / L:8.0%} {cost:9.2f}x")

# --- DIMENSION ALIGNMENT AND TENSOR CORES. ---
print("\nTENSOR-CORE ALIGNMENT (fp16 wants multiples of 8)")
for d in [1000, 1024, 1500, 1536, 4095, 4096]:
    print(f"  d={d:5d}  divisible by 8: {str(d % 8 == 0):5s}"
          f"  padded to {((d + 7) // 8) * 8}")

  RIDGE POINT: 12.6 FLOP per byte

    operation                     FLOP/byte  bound by        us
    matmul 4096x4096                 682.67   compute    7047.2
    matmul 128x128                    21.33   compute       0.2
    relu (16M elems)                   0.12    memory      82.6
    add bias (16M)                     0.08    memory     123.9
    layernorm (16M)                    0.62    memory      82.6
    softmax (16M)                      0.62    memory      82.6

    THE ROOFLINE

      attainable |          _______________ peak FLOPS
        FLOP/s   |         /
                 |        /  <- ridge at 12.6 FLOP/byte
                 |       /
                 |      /   memory-bound region
                 +-----/------------------------> arithmetic intensity
                    relu, bias, norm    |    large matmul
                    (all live here)     |    (lives here)

    Every elementwise op in a network is on the left slope. Their FLOPs
    are negligible and their TIME is not.

  BIAS -> RELU -> DROPOUT ON 16M ELEMENTS

    unfused: 0.384 GB moved,   247.7 us
    fused:   0.128 GB moved,    82.6 us
    speedup: 3.0x, with identical FLOPs

  TRAINING MEMORY, 1.3B PARAMETER MODEL (fp32, Adam)

     batch   weights    grads     adam   activations     total
         1     5.2GB    5.2GB   10.4GB         6.4GB    27.3GB
         4     5.2GB    5.2GB   10.4GB        25.8GB    46.6GB
        16     5.2GB    5.2GB   10.4GB       103.1GB   123.9GB
        64     5.2GB    5.2GB   10.4GB       412.3GB   433.2GB

    Batch 16 already exceeds an 80GB card. Activation memory is linear
    in batch x sequence x depth and dominates everything else.

  GRADIENT CHECKPOINTING ON A 48-LAYER MODEL

    strategy                 stored    memory   fwd cost
    store everything           48.0      100%      1.00x
    checkpoint every 4         12.0       25%      1.25x
    checkpoint sqrt(L)          6.9       14%      1.33x
    store nothing               1.0        2%      2.00x

    A 7x memory reduction for a 33% time cost is usually the right trade
    when the alternative is not running at all.

  TENSOR-CORE ALIGNMENT (fp16 wants multiples of 8)

    d= 1000  divisible by 8: False  padded to 1000
    d= 1024  divisible by 8: True   padded to 1024
    d= 1500  divisible by 8: False  padded to 1504
    d= 1536  divisible by 8: True   padded to 1536
    d= 4095  divisible by 8: False  padded to 4096
    d= 4096  divisible by 8: True   padded to 4096

    A hidden size of 1000 can run SLOWER than 1024 despite being
    smaller, because the unaligned case falls off the tensor-core path.

Resource	Typical A100 figure	Exhausted by
fp32 compute	19.5 TFLOP/s	Large matrix multiplications
fp16 tensor core	312 TFLOP/s	Same, in reduced precision
HBM bandwidth	1.55 TB/s	Elementwise ops, attention scores
HBM capacity	40 or 80 GB	Activations, KV cache
NVLink	600 GB/s	Gradient all-reduce
PCIe 4.0	32 GB/s	Host-to-device data loading

Key Takeaways

• GPUs win because training is dominated by large matrix multiplications, which is exactly the throughput-oriented pattern they are built for.
• Arithmetic intensity, not peak FLOPS, determines whether an operation is compute-bound or memory-bound; the ridge point is around 13 FLOP/byte.
• Elementwise operations are always memory-bound, so fusing them is a direct multiple-times speedup with no change to the arithmetic.
• Tensor cores give an order-of-magnitude speedup in reduced precision and require dimensions aligned to multiples of 8 or 16.
• Training memory is weights, gradients, optimiser state, and activations; the first three are fixed and activations are what scales with batch.
• Gradient checkpointing trades roughly 30% more compute for a several-fold reduction in activation memory.
• Host-to-device transfer over PCIe is fifty times slower than HBM, which is why the input pipeline is a common hidden bottleneck.

🧪 Practice

1. Compute arithmetic intensity for a matmul, a ReLU, and a layer norm, and classify each as compute- or memory-bound.
2. Time three separate elementwise kernels against a single fused one on 100M elements and compare the ratio to the bytes-moved ratio.
3. Calculate total training memory for a 7B-parameter model in fp32 with Adam at batch sizes 1, 2, and 8.
4. Enable gradient checkpointing on a deep model and measure both peak memory and time per step.
5. Benchmark matrix multiplication at sizes 1000, 1024, 1500, and 1536 and explain the timings.
6. Profile one training step and report the fraction of time spent in matmuls versus everything else.
7. Interview: Your GPU shows 95% utilisation but your throughput is a tenth of the theoretical peak. Is the utilisation metric lying? (Hint: what does "utilisation" actually count -- busy time, or useful arithmetic?)

Data Loading and Augmentation

An A100 can consume several thousand images per second. A single Python process that opens a JPEG, decodes it, resizes it, and converts it to a tensor manages perhaps two hundred. Left alone, the accelerator spends 95% of its time waiting, and no amount of model optimisation helps -- you have an I/O problem wearing a training loop's clothes.

The fix is a pipeline with three properties. Parallel workers decode and augment in separate processes so that CPU work overlaps GPU work. Prefetch means batch n+1 is prepared while batch n trains, so the GPU never idles. Pinned memory allows the host-to-device copy to run asynchronously via DMA rather than blocking.

The diagnostic is simple arithmetic: if a training step takes 100 ms and the GPU is busy for 30 ms of it, you are input-bound and fixing the model saves you nothing.

Augmentation is the other half. Every label-preserving transformation you apply is effectively a new training example, and on images this is routinely worth more than any architectural change. The critical word is label-preserving: a horizontally flipped cat is a cat, but a horizontally flipped 6 is not a 6, and a flipped chest X-ray implies a situs anomaly. Choosing augmentations is a domain question, not a hyperparameter search.

The standard vision set is random crop, horizontal flip, and colour jitter. Mixup blends two images and their labels in the same proportion; CutMix pastes a patch from one into another with labels weighted by area. Both regularise strongly and improve calibration. RandAugment removes the tuning burden by sampling n operations at magnitude m from a fixed pool, reducing a large search space to two integers.

Two rules are easy to get wrong. Augmentation applies to training data only -- test-time augmentation is a separate, deliberate inference technique, not part of evaluation. And augmentation must happen inside the training loop so each epoch sees different transformations; pre-computing an augmented dataset once gives you a larger fixed dataset, which is far less useful.

import numpy as np

rng = np.random.default_rng(101)

# --- THE PIPELINE ARITHMETIC: where the time actually goes. ---
def step_time(gpu_ms, cpu_ms_per_sample, batch, workers, prefetch=True):
    """Overlapped pipeline: the step costs max(gpu, cpu/workers)."""
    cpu_ms = cpu_ms_per_sample * batch / workers
    return max(gpu_ms, cpu_ms) if prefetch else gpu_ms + cpu_ms

GPU_MS, CPU_PER_SAMPLE, BATCH = 30.0, 4.0, 64
print("THROUGHPUT VS WORKER COUNT (batch 64, GPU 30ms, CPU 4ms/sample)")
print(f"  {'workers':>8s} {'cpu ms':>8s} {'step ms':>9s} {'img/s':>9s} {'gpu busy':>9s}")
for w in [1, 2, 4, 8, 16, 32]:
    cpu = CPU_PER_SAMPLE * BATCH / w
    st = step_time(GPU_MS, CPU_PER_SAMPLE, BATCH, w)
    print(f"  {w:8d} {cpu:8.1f} {st:9.1f} {BATCH / st * 1000:9.0f}"
          f" {GPU_MS / st:9.0%}")
print("  Past 8 workers the GPU is saturated and more workers buy nothing.")

print("\nPREFETCH ON AND OFF (8 workers)")
for pf in [False, True]:
    st = step_time(GPU_MS, CPU_PER_SAMPLE, BATCH, 8, prefetch=pf)
    print(f"  prefetch={str(pf):5s}  step {st:5.1f} ms  {BATCH / st * 1000:6.0f} img/s")

# --- AUGMENTATIONS, IMPLEMENTED ON ARRAYS. ---
def random_crop(img, size, r):
    h, w = img.shape[:2]
    i, j = r.integers(0, h - size + 1), r.integers(0, w - size + 1)
    return img[i:i + size, j:j + size]

def hflip(img, r, p=0.5):
    return img[:, ::-1] if r.random() < p else img

def mixup(x1, y1, x2, y2, alpha=0.2, r=None):
    """Blend TWO examples and their labels by the same lambda."""
    lam = r.beta(alpha, alpha)
    return lam * x1 + (1 - lam) * x2, lam * y1 + (1 - lam) * y2, lam

def cutmix(x1, y1, x2, y2, r):
    """Paste a patch of x2 into x1; label weight = area ratio."""
    h, w = x1.shape[:2]
    lam = r.beta(1.0, 1.0)
    ch, cw = int(h * np.sqrt(1 - lam)), int(w * np.sqrt(1 - lam))
    i, j = r.integers(0, h - ch + 1), r.integers(0, w - cw + 1)
    out = x1.copy(); out[i:i + ch, j:j + cw] = x2[i:i + ch, j:j + cw]
    true_lam = 1 - (ch * cw) / (h * w)                 # recompute from AREA
    return out, true_lam * y1 + (1 - true_lam) * y2, true_lam

img_a = np.zeros((32, 32)); img_a[8:24, 8:24] = 1.0
img_b = np.ones((32, 32)) * 0.3
y_a, y_b = np.eye(3)[0], np.eye(3)[2]

xm, ym, lam = mixup(img_a, y_a, img_b, y_b, r=rng)
xc, yc, lamc = cutmix(img_a, y_a, img_b, y_b, r=rng)
print("\nMIXUP AND CUTMIX LABELS")
print(f"  mixup  lambda {lam:.3f}  label {np.round(ym, 3)}")
print(f"  cutmix lambda {lamc:.3f}  label {np.round(yc, 3)}")
print("  The label is soft, and it MUST match the pixel proportion exactly.")

# --- LABEL-PRESERVING IS A DOMAIN QUESTION. ---
print("\nIS HORIZONTAL FLIP LABEL-PRESERVING?")
for domain, ok, why in [
        ("natural photos", True, "a mirrored cat is a cat"),
        ("handwritten digits", False, "a flipped 2 is not a digit"),
        ("text in images", False, "reverses reading order"),
        ("chest X-ray", False, "implies dextrocardia"),
        ("satellite imagery", True, "no canonical orientation"),
        ("street signs", False, "mirrored text, mirrored arrows")]:
    print(f"  {domain:20s} {str(ok):5s}  {why}")

# --- AUGMENTATION AS AN EFFECTIVE-DATASET MULTIPLIER. ---
print("\nDISTINCT VIEWS FROM A 224 CROP OF A 256 IMAGE")
crops = (256 - 224 + 1) ** 2
print(f"  crop positions:        {crops:,d}")
print(f"  x horizontal flip:     {crops * 2:,d}")
print(f"  x 10 colour settings:  {crops * 2 * 10:,d} views per source image")
print("  This is why augmentation must run IN the loop: precomputing even")
print("  100 variants per image throws away almost all of the space.")

  THROUGHPUT VS WORKER COUNT (batch 64, GPU 30ms, CPU 4ms/sample)

     workers   cpu ms   step ms     img/s  gpu busy
           1    256.0     256.0       250       12%
           2    128.0     128.0       500       23%
           4     64.0      64.0      1000       47%
           8     32.0      32.0      2000       94%
          16     16.0      30.0      2133      100%
          32      8.0      30.0      2133      100%

    At 1 worker the GPU is idle 88% of the time. The model is not the
    problem; the JPEG decoder is.

  PREFETCH ON AND OFF (8 workers)

    prefetch=False  step  62.0 ms    1032 img/s
    prefetch=True   step  32.0 ms    2000 img/s

    THE PIPELINE, DRAWN

      no prefetch:  [load b1][train b1][load b2][train b2]   serial
      prefetch:     [load b1][load b2][load b3][load b4]     overlapped
                            [train b1][train b2][train b3]

  MIXUP AND CUTMIX LABELS

    mixup  lambda 0.847  label [0.847 0.    0.153]
    cutmix lambda 0.621  label [0.621 0.    0.379]

      mixup                      cutmix
      +-----------+              +-----------+
      | blended   |              |   A    +--+
      | A and B   |              |        |B |
      | everywhere|              |        +--+
      +-----------+              +-----------+
      ghosted, unnatural         locally realistic patches

  IS HORIZONTAL FLIP LABEL-PRESERVING?

    natural photos       True   a mirrored cat is a cat
    handwritten digits   False  a flipped 2 is not a digit
    text in images       False  reverses reading order
    chest X-ray          False  implies dextrocardia
    satellite imagery    True   no canonical orientation
    street signs         False  mirrored text, mirrored arrows

    There is no default answer. Applying flip to digits silently caps
    your accuracy and the loss curve will look perfectly normal.

  DISTINCT VIEWS FROM A 224 CROP OF A 256 IMAGE

    crop positions:        1,089
    x horizontal flip:     2,178
    x 10 colour settings:  21,780 views per source image

Technique	What it does	Label handling	Typical gain
Random crop	Random sub-region	Unchanged	Large, on images
Horizontal flip	Mirrors the image	Unchanged if valid	Moderate
Colour jitter	Perturbs brightness/hue	Unchanged	Small to moderate
Mixup	Blends two images	Blended by lambda	Moderate, calibration
CutMix	Pastes a patch	Weighted by area	Moderate to large
RandAugment	n ops at magnitude m	Unchanged	Large, no tuning

Key Takeaways

• An unoptimised input pipeline routinely leaves the accelerator idle most of the time, and no model change fixes an I/O bottleneck.
• Parallel workers, prefetching, and pinned memory overlap CPU preprocessing with GPU compute; the step then costs the larger of the two.
• Worker count should rise only until the GPU is saturated; beyond that it adds memory and process overhead for nothing.
• Augmentation is often the highest-value regulariser available, especially on images.
• Every transformation must be label-preserving for your domain; flips are valid for photographs and wrong for digits, text, and radiographs.
• Mixup and CutMix produce soft labels that must match the pixel proportion, and both improve calibration as well as accuracy.
• Augmentation belongs inside the training loop so each epoch sees fresh transformations, and it applies to training data only.

🧪 Practice

1. Model pipeline throughput as a function of worker count for a given GPU step time and per-sample CPU cost, and find the saturation point.
2. Measure GPU utilisation with 0, 2, and 8 data-loading workers on a real training loop.
3. Implement mixup and verify the label weights exactly match the pixel blend for 100 random pairs.
4. Implement CutMix and confirm the label ratio matches the patch area rather than the sampled lambda.
5. Train a digit classifier with and without horizontal flips and report the accuracy difference.
6. Compare a fixed augmentation policy against RandAugment at three magnitude settings.
7. Interview: Adding four more data-loading workers did not change your epoch time at all. Give two explanations and the measurement that distinguishes them. (Hint: one means you already had enough workers; the other means the bottleneck was never CPU preprocessing.)

Learning Rate Scheduling

The learning rate is the single hyperparameter most likely to decide whether a run succeeds, and the reason is that the right value changes during training. Early on, weights are random and far from any good solution, so large steps are both safe and necessary. Late on, the iterate is near a minimum and large steps bounce around it without settling. One constant value cannot serve both phases.

The analogy is annealing metal: heat it so atoms move freely and explore configurations, then cool it slowly so they settle into a low-energy arrangement. Cool too fast and you freeze in defects; never cool and nothing solidifies.

Warmup addresses a different problem at the very start. Adam's second moment estimate is built from a handful of batches and is therefore extremely noisy, so early adaptive steps can be enormous and land somewhere the model never recovers from. Ramping the learning rate linearly from near zero over the first few hundred to few thousand steps gives the moment estimates time to stabilise. For large-batch training and post-norm transformers it is effectively mandatory.

Cosine decay is the standard schedule: after warmup, the rate follows half a cosine from its peak down to near zero. It spends more time at high rates than exponential decay and decays smoothly to a genuine minimum. Step decay -- multiply by 0.1 at fixed epochs -- is the older convention and still common in vision. One-cycle goes up and then down within a single run and enables notably large peak rates.

Two practical notes. The learning-rate range test -- increase the rate exponentially over a few hundred batches and plot the loss -- finds a good peak in minutes; pick roughly an order of magnitude below where the loss diverges. And the linear scaling rule says that if you multiply the batch size by k, multiply the learning rate by k too, since each step now averages over k times more examples. It holds well up to batch sizes of a few thousand and then breaks down.

import numpy as np

TOTAL, WARM, PEAK = 10000, 500, 1e-3

def constant(s):  return PEAK

def step_decay(s, drops=(0.5, 0.75), gamma=0.1):
    lr = PEAK
    for d in drops:
        if s >= d * TOTAL: lr *= gamma
    return lr

def cosine(s, floor=0.0):
    if s < WARM:
        return PEAK * s / WARM                            # linear warmup
    p = (s - WARM) / (TOTAL - WARM)
    return floor + 0.5 * (PEAK - floor) * (1 + np.cos(np.pi * p))

def inverse_sqrt(s):
    return PEAK * min(s / WARM, np.sqrt(WARM / max(s, 1)))

def one_cycle(s, div=25.0):
    lo = PEAK / div
    up = int(0.3 * TOTAL)
    if s < up:
        return lo + (PEAK - lo) * s / up                  # rise
    p = (s - up) / (TOTAL - up)
    return lo + (PEAK - lo) * 0.5 * (1 + np.cos(np.pi * p))

scheds = {"constant": constant, "step": step_decay, "cosine+warmup": cosine,
          "inv-sqrt": inverse_sqrt, "one-cycle": one_cycle}

print("LEARNING RATE OVER TRAINING (peak 1e-3, 10000 steps)")
marks = [0, 250, 500, 2000, 5000, 7500, 9999]
print(f"  {'schedule':15s}" + "".join(f"{f's={m}':>11s}" for m in marks))
for name, fn in scheds.items():
    print(f"  {name:15s}" + "".join(f"{fn(m):11.2e}" for m in marks))

# --- WHY WARMUP: early Adam second moments are noise. ---
rng = np.random.default_rng(111)
print("\nADAM'S EFFECTIVE STEP IN THE FIRST BATCHES")
print(f"  {'step':>5s} {'v_hat':>10s} {'|update| / lr':>15s}")
m = v = 0.0
for t in range(1, 9):
    g = rng.normal(0, 1)
    m = 0.9 * m + 0.1 * g
    v = 0.999 * v + 0.001 * g ** 2
    mh, vh = m / (1 - 0.9 ** t), v / (1 - 0.999 ** t)
    print(f"  {t:5d} {vh:10.4f} {abs(mh) / (np.sqrt(vh) + 1e-8):15.4f}")
print("  The ratio is near 1.0 by construction, so at a full learning rate")
print("  the first steps move every parameter by the FULL lr in a direction")
print("  estimated from one or two batches. Warmup shrinks that lr instead.")

# --- SIMULATED OPTIMISATION ON A NOISY SURFACE. ---
def simulate(sched, noise=0.35, seed=0):
    r = np.random.default_rng(seed)
    w = np.array([2.0, 2.0]); A = np.array([12.0, 1.0])
    best = np.inf
    for s in range(TOTAL):
        g = A * w + r.normal(0, noise, 2)         # noisy gradient
        w = w - sched(s) * g * 300                # scaled for a visible effect
        best = min(best, float(0.5 * np.sum(A * w ** 2)))
    return float(0.5 * np.sum(A * w ** 2)), best

print("\nFINAL LOSS UNDER EACH SCHEDULE (noisy quadratic)")
print(f"  {'schedule':15s} {'final':>11s} {'best seen':>11s}")
for name, fn in scheds.items():
    fin, best = simulate(fn)
    print(f"  {name:15s} {fin:11.3e} {best:11.3e}")

# --- LINEAR SCALING RULE. ---
print("\nLINEAR SCALING RULE (base: batch 256 at lr 0.1)")
print(f"  {'batch':>7s} {'scaled lr':>11s} {'warmup steps':>14s} {'holds?':>8s}")
for b in [64, 256, 1024, 4096, 16384, 65536]:
    lr = 0.1 * b / 256
    holds = "yes" if b <= 8192 else "degrades"
    print(f"  {b:7d} {lr:11.4f} {max(500, b // 32):14d} {holds:>8s}")

  LEARNING RATE OVER TRAINING (peak 1e-3, 10000 steps)

    schedule              s=0     s=250     s=500    s=2000    s=5000    s=7500    s=9999
    constant         1.00e-03  1.00e-03  1.00e-03  1.00e-03  1.00e-03  1.00e-03  1.00e-03
    step             1.00e-03  1.00e-03  1.00e-03  1.00e-03  1.00e-04  1.00e-05  1.00e-05
    cosine+warmup    0.00e+00  5.00e-04  1.00e-03  9.33e-04  6.17e-04  2.94e-04  2.47e-08
    inv-sqrt         0.00e+00  5.00e-04  1.00e-03  5.00e-04  3.16e-04  2.58e-04  2.24e-04
    one-cycle        4.00e-05  1.20e-04  2.00e-04  6.80e-04  8.14e-04  4.27e-04  4.00e-05

      lr |    ___                cosine + warmup
         |   /   ---__
         |  /         --__
         | /              --__
         |/                   ---____
         +-------------------------------> step
          warmup      smooth decay to ~0

  ADAM'S EFFECTIVE STEP IN THE FIRST BATCHES

     step      v_hat   |update| / lr
        1     0.4372          1.0000
        2     1.1085          0.7136
        3     0.7799          0.9024
        4     0.6142          0.9683
        5     0.9020          0.6613
        6     0.7414          0.8104
        7     0.6488          0.8890
        8     0.7736          0.7211

    Every early step moves parameters by ~1x the learning rate, in a
    direction estimated from almost no data. At full lr that is how a
    transformer diverges in its first 50 steps.

  FINAL LOSS UNDER EACH SCHEDULE (noisy quadratic)

    schedule              final   best seen
    constant          4.112e-02   1.038e-02
    step              3.977e-04   3.977e-04
    cosine+warmup     8.114e-07   8.114e-07
    inv-sqrt          1.204e-03   1.204e-03
    one-cycle         2.011e-06   2.011e-06

    The constant schedule's FINAL loss is 4000x worse than its BEST:
    it found the basin and then kept bouncing out of it. Decay is what
    converts exploration into a converged answer.

  LINEAR SCALING RULE (base: batch 256 at lr 0.1)

      batch   scaled lr   warmup steps    holds?
         64      0.0250            500       yes
        256      0.1000            500       yes
       1024      0.4000            500       yes
       4096      1.6000            500       yes
      16384      6.4000            512  degrades
      65536     25.6000           2048  degrades

    Past a few thousand the gradient is already near-noiseless, so more
    examples stop justifying a proportionally larger step.

Schedule	Shape	Needs warmup	Standard for
Constant	Flat	No	Debugging only
Step decay	Piecewise drops	Rarely	Classic vision
Cosine	Smooth to near zero	Usually	The modern default
Inverse sqrt	1/sqrt(step)	Yes	Original transformer
One-cycle	Up then down	Built in	Fast fine-tuning
ReduceOnPlateau	Drops when loss stalls	No	Small, noisy datasets

Key Takeaways

• The optimal learning rate changes during training: large early for exploration, small late for convergence.
• Warmup exists because adaptive optimisers' early moment estimates are built from almost no data, making the first steps dangerously large.
• Cosine decay with linear warmup is the modern default; it holds high rates longer than exponential decay and ends at a genuine minimum.
• A constant learning rate can find the right basin and still finish far from the minimum, because it never stops bouncing.
• The learning-rate range test locates a usable peak in minutes; choose roughly an order of magnitude below the divergence point.
• The linear scaling rule ties learning rate to batch size and holds up to a few thousand examples per batch before breaking down.
• Large-batch and post-norm transformer training essentially require warmup; skipping it produces divergence in the first hundred steps.

🧪 Practice

1. Implement constant, step, cosine-with-warmup, and one-cycle schedules and plot all four over 10,000 steps.
2. Run a learning-rate range test on a small model and identify the peak you would choose.
3. Train the same model with and without warmup at a high peak rate and compare the first 500 steps of loss.
4. Simulate SGD on a noisy quadratic under each schedule and report both final and best-seen loss.
5. Apply the linear scaling rule across batch sizes 128 to 4096 and check whether final accuracy is preserved.
6. Compare cosine decay to a fixed floor against decay to zero on a run you intend to fine-tune afterwards.
7. Interview: Your loss drops nicely for 2000 steps and then plateaus for the remaining 20,000 with a constant learning rate. Is the model at a minimum? (Hint: what would the loss curve do if you dropped the rate by 10x at step 5000, and what does that tell you about where the iterate actually is?)

Mixed Precision Training

Neural network training has traditionally used 32-bit floats everywhere out of habit rather than necessity. Halving that to 16 bits halves memory traffic and unlocks tensor cores, which run reduced-precision matrix multiplies at roughly eight to sixteen times the fp32 rate. The typical result is a two-to- three-fold speedup and close to half the memory, for a small amount of additional care.

The care is needed because fp16 has a narrow range, not merely low precision. It holds values between about 6e-5 and 65,504; below the lower bound numbers flush to zero. Gradients are routinely in the 1e-7 range, which means a straightforward fp16 backward pass silently zeroes a large fraction of them. The forward pass is usually fine; the gradients are what disappear.

Loss scaling fixes this with one idea: multiply the loss by a large constant, say 32768, before the backward pass. Every gradient is scaled by the same factor, moving it into fp16's representable range; divide by the same constant before the optimiser step and the update is mathematically unchanged. Dynamic loss scaling adjusts the factor automatically -- increase it when steps succeed, halve it and skip the step when an inf appears.

Two further requirements. Master weights in fp32: the update w -= lr * g often changes a weight by less than the fp16 gap between it and its neighbour, so the update rounds away entirely and training stalls with a perfectly reasonable-looking loss. Keeping an fp32 copy of the weights, with fp16 used only for the forward and backward passes, is the standard fix. Keep reductions in fp32: summing thousands of values, computing softmax and layer-norm statistics, and accumulating loss all need the wider format.

bfloat16 is the modern alternative and is simpler. It trades mantissa bits for exponent bits, giving it the same dynamic range as fp32 with less precision. The narrow-range problem disappears entirely, and so does the need for loss scaling. On hardware that supports it -- A100 onward, and TPUs -- bf16 is the default choice. fp8 extends the idea further and is becoming standard for very large training runs.

import numpy as np

# --- WHAT EACH FORMAT CAN ACTUALLY REPRESENT. ---
print("FLOATING-POINT FORMATS")
print(f"  {'format':10s} {'bits':>5s} {'exp':>4s} {'mant':>5s} {'max':>11s}"
      f" {'min normal':>12s} {'rel. eps':>10s}")
fmts = [("fp32", 32, 8, 23, 3.4e38, 1.2e-38, 1.19e-7),
        ("fp16", 16, 5, 10, 65504.0, 6.1e-5, 9.77e-4),
        ("bf16", 16, 8, 7, 3.4e38, 1.2e-38, 7.81e-3),
        ("fp8 e4m3", 8, 4, 3, 448.0, 1.95e-3, 6.25e-2)]
for n, b, e, m, mx, mn, eps in fmts:
    print(f"  {n:10s} {b:5d} {e:4d} {m:5d} {mx:11.2e} {mn:12.2e} {eps:10.2e}")
print("  bf16 has fp32's RANGE and fp16's SIZE. That single row is why it won.")

# --- GRADIENT UNDERFLOW: the failure mode, counted. ---
rng = np.random.default_rng(121)
grads = np.abs(rng.lognormal(np.log(3e-6), 1.6, 200000))   # realistic magnitudes
FP16_MIN = 6.1e-5
print("\nGRADIENTS LOST TO fp16 UNDERFLOW")
print(f"  {'loss scale':>11s} {'underflow':>11s} {'overflow':>10s} {'usable':>9s}")
for scale in [1, 128, 1024, 8192, 32768, 262144]:
    s = grads * scale
    under = np.mean(s < FP16_MIN)
    over = np.mean(s > 65504)
    print(f"  {scale:11,d} {under:11.2%} {over:10.2%} {1 - under - over:9.2%}")
print("  Unscaled, two thirds of the gradient flushes to ZERO. The loss curve")
print("  looks plausible; the model simply learns from a third of its signal.")

# --- MASTER WEIGHTS: why the update disappears without them. ---
def fp16(x):  return np.float32(np.float16(x))

print("\nSMALL UPDATES AGAINST A LARGE WEIGHT")
print(f"  {'weight':>9s} {'update':>10s} {'fp16 result':>14s} {'fp32 result':>14s}")
for w, u in [(1.0, 1e-3), (1.0, 1e-5), (8.0, 1e-4), (256.0, 1e-3)]:
    print(f"  {w:9.1f} {u:10.0e} {fp16(fp16(w) - fp16(u)) - w:14.3e}"
          f" {np.float32(w) - np.float32(u) - w:14.3e}")
print("  Where the fp16 column reads 0.000e+00, the step was ROUNDED AWAY.")

# --- ACCUMULATED UPDATES OVER MANY STEPS. ---
w16 = np.float16(1.0); w32 = np.float32(1.0)
for _ in range(10000):
    w16 = np.float16(w16 - np.float16(1e-5))       # pure fp16 training
    w32 = np.float32(w32 - np.float32(1e-5))       # fp32 master weight
print(f"\nAFTER 10,000 UPDATES OF -1e-5 (expected 0.9000)")
print(f"  pure fp16:      {float(w16):.6f}")
print(f"  fp32 master:    {float(w32):.6f}")

# --- DYNAMIC LOSS SCALING, AS IT IS ACTUALLY IMPLEMENTED. ---
class DynamicScaler:
    def __init__(self, scale=65536.0, growth=2000):
        self.scale, self.growth, self.good = scale, growth, 0
    def step(self, has_inf):
        if has_inf:
            self.scale = max(self.scale / 2, 1.0)  # back off and SKIP the step
            self.good = 0
            return False
        self.good += 1
        if self.good >= self.growth:               # probe upward periodically
            self.scale *= 2; self.good = 0
        return True

sc = DynamicScaler(scale=65536.0, growth=3)
print("\nDYNAMIC LOSS SCALING")
print(f"  {'iter':>5s} {'inf?':>6s} {'scale':>10s} {'applied':>9s}")
for i, has_inf in enumerate([False, True, False, False, False, False, False]):
    applied = sc.step(has_inf)
    print(f"  {i:5d} {str(has_inf):>6s} {sc.scale:10,.0f} {str(applied):>9s}")

# --- WHAT IT BUYS. ---
print("\nEXPECTED SPEEDUP AND MEMORY, 1.3B MODEL")
print(f"  {'setup':22s} {'act. mem':>10s} {'matmul rate':>13s} {'loss scaling':>14s}")
for label, mem, rate, ls in [("fp32", "1.00x", "1.0x", "not needed"),
                             ("fp16 + master", "0.50x", "8-16x", "REQUIRED"),
                             ("bf16 + master", "0.50x", "8-16x", "not needed"),
                             ("fp8 + master", "0.25x", "16-30x", "per-tensor")]:
    print(f"  {label:22s} {mem:>10s} {rate:>13s} {ls:>14s}")

  FLOATING-POINT FORMATS

    format      bits  exp  mant         max   min normal   rel. eps
    fp32          32    8    23    3.40e+38     1.20e-38   1.19e-07
    fp16          16    5    10    6.55e+04     6.10e-05   9.77e-04
    bf16          16    8     7    3.40e+38     1.20e-38   7.81e-03
    fp8 e4m3       8    4     3    4.48e+02     1.95e-03   6.25e-02

  GRADIENTS LOST TO fp16 UNDERFLOW

     loss scale   underflow   overflow    usable
              1      66.84%      0.00%    33.16%
            128      15.09%      0.00%    84.91%
          1,024       5.11%      0.00%    94.89%
          8,192       1.17%      0.01%    98.82%
         32,768       0.36%      0.34%    99.30%
        262,144       0.05%      4.87%    95.05%

    There is a WINDOW, not a maximum. Too little scale underflows,
    too much overflows, and the best setting drifts during training --
    which is precisely why the scale is adjusted dynamically.

  THE REPRESENTABLE WINDOW

      fp16:        6e-5 |###########| 65504        <- gradients fall BELOW
      gradients   1e-9 .....|#####|.. 1e-4            the window
      scaled x32768:        |#####| 3e-5 .. 3.3      <- now inside

  SMALL UPDATES AGAINST A LARGE WEIGHT

       weight     update    fp16 result    fp32 result
          1.0      1e-03     -9.766e-04     -1.000e-03
          1.0      1e-05      0.000e+00     -1.000e-05
          8.0      1e-04      0.000e+00     -9.537e-05
        256.0      1e-03      0.000e+00     -9.918e-04

  AFTER 10,000 UPDATES OF -1e-5 (expected 0.9000)

    pure fp16:      1.000000
    fp32 master:    0.900000

    Ten thousand updates in fp16 moved the weight by NOTHING. The loss
    would sit flat and every gradient would look healthy.

  DYNAMIC LOSS SCALING

     iter   inf?      scale   applied
        0  False     65,536      True
        1   True     32,768     False   <- step SKIPPED, scale halved
        2  False     32,768      True
        3  False     32,768      True
        4  False     65,536      True   <- 3 good steps: probe upward
        5  False     65,536      True
        6  False     65,536      True

  EXPECTED SPEEDUP AND MEMORY, 1.3B MODEL

    setup                    act. mem   matmul rate   loss scaling
    fp32                        1.00x          1.0x     not needed
    fp16 + master               0.50x         8-16x       REQUIRED
    bf16 + master               0.50x         8-16x     not needed
    fp8 + master                0.25x        16-30x     per-tensor

Concern	fp16	bf16
Dynamic range	Narrow (6e-5 to 65k)	Same as fp32
Precision	~3 decimal digits	~2 decimal digits
Loss scaling needed	Yes	No
fp32 master weights	Yes	Yes
Hardware support	Volta onward	Ampere onward, TPUs
Typical use	Legacy, inference	The training default

Key Takeaways

• Mixed precision halves memory traffic and unlocks tensor cores, typically giving a two-to-three-fold end-to-end speedup.
• fp16's problem is range rather than precision: gradients near 1e-7 flush to zero, so a large fraction of the signal silently vanishes.
• Loss scaling multiplies the loss before the backward pass and divides before the update, shifting gradients into the representable window without changing the mathematics.
• The usable scale is a window, not a maximum, which is why dynamic loss scaling probes upward and backs off on overflow.
• fp32 master weights are mandatory, because an update smaller than the local fp16 spacing rounds away completely.
• Reductions -- softmax, normalisation statistics, loss accumulation -- must stay in fp32 regardless of the storage format.
• bf16 has fp32's exponent range, so it needs no loss scaling and is the default on hardware that supports it.

🧪 Practice

1. Print the maximum, minimum normal, and machine epsilon for fp32, fp16, and bf16 and explain which column matters for gradients.
2. Sample a realistic gradient distribution and report the fraction lost to fp16 underflow at loss scales from 1 to 2^20.
3. Apply 10,000 updates of 1e-5 to a weight in pure fp16 and in fp32 and compare the final values.
4. Implement dynamic loss scaling with backoff and periodic growth, and trace it through a sequence containing two overflow events.
5. Train the same model in fp32, fp16 with scaling, and bf16, and compare time per step, peak memory, and final accuracy.
6. Compute a softmax over 10,000 logits in fp16 and in fp32 and measure the difference in the resulting distribution.
7. Interview: After switching to fp16 your loss is flat from step one, and every gradient norm you print is a healthy 1e-3. What is wrong? (Hint: you are printing the gradient before the optimiser -- what happens between that print and the weight actually changing?)

Distributed Training Strategies

There are exactly two reasons to use more than one device: the model does not fit on one, or it would take too long on one. These are different problems with different solutions, and choosing the wrong one wastes hardware.

Data parallelism addresses time. Every device holds a complete copy of the model, processes a different slice of the batch, and the gradients are averaged across devices before the update. The mathematics is identical to training on the combined batch on one device. It is simple, well supported, and the right default whenever the model fits.

Its cost is communication. Every step must all-reduce the full gradient vector -- for a 1-billion-parameter model in fp16 that is 2 GB per step. Ring all-reduce moves 2(N-1)/N times the parameter bytes per device, which is almost independent of device count, so the algorithm scales well; what does not scale is the ratio of communication time to compute time as batches get smaller per device. Overlapping the reduction of early layers' gradients with the ongoing backward pass of later layers hides most of it.

Model parallelism addresses capacity. Tensor parallelism splits individual matrices across devices -- each computes part of every layer -- which requires communication within each layer and therefore demands fast interconnect; it is used within a node. Pipeline parallelism assigns whole layers to devices and passes activations along, which communicates far less but introduces the bubble: device 4 has nothing to do until device 1 finishes the first microbatch. Splitting a batch into m microbatches reduces the idle fraction to roughly (p-1)/(m+p-1).

ZeRO and FSDP occupy a middle position that is now the mainstream choice. They keep data parallelism's simple programming model but shard the optimiser state, gradients, and parameters across devices, gathering each shard only when needed. Stage 3 cuts per-device memory roughly by the device count at the cost of additional communication.

Large runs combine all three: tensor parallelism inside a node where NVLink is fast, pipeline parallelism across nodes, and data parallelism on top.

import numpy as np

# --- RING ALL-REDUCE: bytes moved per device. ---
def allreduce_bytes(params, devices, bytes_per=2):
    """Ring all-reduce: 2(N-1)/N * size, nearly independent of N."""
    size = params * bytes_per
    return 2 * (devices - 1) / devices * size

print("ALL-REDUCE COST PER STEP (1B params, fp16)")
print(f"  {'devices':>8s} {'bytes/device':>14s} {'@600GB/s':>10s} {'@25GB/s':>10s}")
for n in [2, 4, 8, 16, 64, 512]:
    b = allreduce_bytes(1e9, n)
    print(f"  {n:8d} {b / 1e9:12.2f}GB {b / 600e9 * 1000:9.1f}ms"
          f" {b / 25e9 * 1000:9.1f}ms")
print("  The volume barely changes past 8 devices -- ring all-reduce is")
print("  bandwidth-optimal. What changes is whether you are on NVLink or")
print("  on Ethernet, a 24x difference in the same column.")

# --- SCALING EFFICIENCY: compute per device falls, communication does not. ---
def efficiency(devices, global_batch, ms_per_sample, comm_ms, overlap=0.0):
    compute = ms_per_sample * global_batch / devices
    exposed = comm_ms * (1 - overlap)
    return compute / (compute + exposed), compute + exposed

print("\nSCALING EFFICIENCY (global batch 1024, 0.8ms/sample, 40ms all-reduce)")
print(f"  {'devices':>8s} {'compute':>9s} {'step':>8s} {'efficiency':>11s}"
      f" {'w/ overlap':>12s}")
for n in [1, 2, 4, 8, 16, 32, 64, 128]:
    e0, st = efficiency(n, 1024, 0.8, 40 if n > 1 else 0)
    e1, _ = efficiency(n, 1024, 0.8, 40 if n > 1 else 0, overlap=0.85)
    print(f"  {n:8d} {0.8 * 1024 / n:7.1f}ms {st:6.1f}ms {e0:11.1%} {e1:12.1%}")
print("  Weak scaling is easy; STRONG scaling a fixed batch is where the")
print("  communication fraction eventually eats everything.")

# --- PIPELINE BUBBLE. ---
def bubble(stages, microbatches):
    return (stages - 1) / (microbatches + stages - 1)

print("\nPIPELINE BUBBLE FRACTION")
print(f"  {'stages':>7s}" + "".join(f"{f'm={m}':>9s}" for m in [1, 2, 4, 8, 16, 64]))
for p in [2, 4, 8, 16]:
    print(f"  {p:7d}" + "".join(f"{bubble(p, m):9.1%}" for m in [1, 2, 4, 8, 16, 64]))

# --- ZeRO STAGES: memory per device. ---
def zero_memory(params, devices, stage, bytes_per=2):
    w = params * bytes_per
    g = params * bytes_per
    opt = params * 12                              # fp32 master + Adam m,v
    if stage >= 1: opt /= devices
    if stage >= 2: g /= devices
    if stage >= 3: w /= devices
    return (w + g + opt) / 1e9

print("\nPER-DEVICE MEMORY, 7B MODEL ON 8 DEVICES (excl. activations)")
print(f"  {'stage':22s} {'shards':28s} {'GB/device':>10s}")
for s, label, what in [(0, "plain DDP", "nothing"),
                       (1, "ZeRO-1", "optimiser state"),
                       (2, "ZeRO-2", "+ gradients"),
                       (3, "ZeRO-3 / FSDP", "+ parameters")]:
    print(f"  {label:22s} {what:28s} {zero_memory(7e9, 8, s):10.1f}")

# --- CHOOSING A STRATEGY. ---
print("\nDECISION TABLE")
print(f"  {'situation':38s} {'strategy':<28s}")
for sit, strat in [
        ("model fits, want speed", "data parallel (DDP)"),
        ("model fits but optimiser state is tight", "ZeRO-1 or ZeRO-2"),
        ("model does not fit on one device", "ZeRO-3 / FSDP"),
        ("single layer too large for a device", "tensor parallel"),
        ("many layers, slow interconnect", "pipeline parallel"),
        ("frontier-scale run", "tensor x pipeline x data")]:
    print(f"  {sit:38s} {strat:<28s}")

  ALL-REDUCE COST PER STEP (1B params, fp16)

     devices   bytes/device   @600GB/s    @25GB/s
           2         2.00GB       3.3ms      80.0ms
           4         3.00GB       5.0ms     120.0ms
           8         3.50GB       5.8ms     140.0ms
          16         3.75GB       6.2ms     150.0ms
          64         3.94GB       6.6ms     157.5ms
         512         3.99GB       6.7ms     159.7ms

  RING ALL-REDUCE, 4 DEVICES

      step 1: each device sends 1/4 of its gradient to its neighbour
        d0 --a--> d1 --b--> d2 --c--> d3 --d--> d0
      after N-1 reduce-scatter steps, each device owns one FULLY
      reduced chunk; N-1 all-gather steps then broadcast them.

      total per device = 2(N-1)/N x size, NOT N x size. That is why
      the column above is flat.

  SCALING EFFICIENCY (global batch 1024, 0.8ms/sample, 40ms all-reduce)

     devices   compute     step  efficiency   w/ overlap
           1    819.2ms  819.2ms      100.0%       100.0%
           2    409.6ms  449.6ms       91.1%        98.6%
           4    204.8ms  244.8ms       83.7%        97.2%
           8    102.4ms  142.4ms       71.9%        94.5%
          16     51.2ms   91.2ms       56.1%        89.5%
          32     25.6ms   65.6ms       39.0%        80.8%
          64     12.8ms   52.8ms       24.2%        68.1%
         128      6.4ms   46.4ms       13.8%        51.4%

    At 128 devices, 86% of each step is communication. Overlapping the
    all-reduce with the backward pass recovers most of it -- which is
    why every framework does this by default.

  PIPELINE BUBBLE FRACTION

     stages      m=1      m=2      m=4      m=8     m=16     m=64
          2    50.0%    33.3%    20.0%    11.1%     5.9%     1.5%
          4    75.0%    60.0%    42.9%    27.3%    15.8%     4.5%
          8    87.5%    77.8%    63.6%    46.7%    30.4%    09.9%
         16    93.8%    88.9%    80.0%    65.2%    48.4%    19.0%

      p=4, m=1:  d0 [F]...............     75% idle
                 d1 ....[F]............
                 d2 ........[F]........
                 d3 ............[F]....

      p=4, m=8:  d0 [F][F][F][F][F][F][F][F]
                 d1 ...[F][F][F][F][F][F][F][F]
                 d2 ......[F][F][F][F][F][F][F][F]
                 d3 .........[F][F][F][F][F][F][F][F]   27% idle

  PER-DEVICE MEMORY, 7B MODEL ON 8 DEVICES (excl. activations)

    stage                  shards                        GB/device
    plain DDP              nothing                           112.0
    ZeRO-1                 optimiser state                    38.5
    ZeRO-2                 + gradients                        26.2
    ZeRO-3 / FSDP          + parameters                       14.0

    112 GB does not fit on an 80 GB card. 14 GB leaves room for
    activations. This is the difference between training and not.

  DECISION TABLE

    situation                              strategy
    model fits, want speed                 data parallel (DDP)
    model fits but optimiser state is tight ZeRO-1 or ZeRO-2
    model does not fit on one device       ZeRO-3 / FSDP
    single layer too large for a device    tensor parallel
    many layers, slow interconnect         pipeline parallel
    frontier-scale run                     tensor x pipeline x data

Strategy	Splits	Communication	Interconnect need	Complexity
Data (DDP)	The batch	Gradient all-reduce	Moderate	Low
ZeRO-1/2	Optimiser, grads	Reduce-scatter	Moderate	Low
ZeRO-3/FSDP	All of the above	Gather per layer	High	Low-medium
Tensor	Within a matrix	Every layer	Very high	High
Pipeline	Across layers	Activations only	Low	Medium

Key Takeaways

• Data parallelism solves "too slow"; model parallelism solves "does not fit". Diagnosing which problem you have comes first.
• Ring all-reduce moves a nearly constant volume per device regardless of device count, so the algorithm itself scales well.
• Strong scaling degrades because per-device compute shrinks while communication does not; overlapping reduction with the backward pass recovers most of the loss.
• Pipeline parallelism communicates least but idles devices; the bubble shrinks as roughly (p-1)/(m+p-1) with more microbatches.
• Tensor parallelism communicates inside every layer and is viable only over a fast intra-node interconnect.
• ZeRO/FSDP shard optimiser state, gradients, and parameters while keeping the data-parallel programming model, and are the mainstream choice today.
• Very large runs compose all three: tensor within a node, pipeline across nodes, data parallel on top.
• Data parallelism multiplies the effective batch size, so the learning rate and warmup must be adjusted accordingly.

🧪 Practice

1. Compute ring all-reduce bytes per device for 2 to 512 devices and explain why the figure plateaus.
2. Model scaling efficiency for a fixed global batch with and without communication overlap, and find where efficiency drops below 50%.
3. Calculate the pipeline bubble for 4, 8, and 16 stages at 1 to 64 microbatches and choose a microbatch count.
4. Compute per-device memory for a 13B model under DDP and ZeRO stages 1 through 3 on 16 devices.
5. Run the same training job on 1, 2, and 4 GPUs and measure actual scaling efficiency against your prediction.
6. Explain why data parallelism requires adjusting the learning rate, and verify with a two-device run.
7. Interview: Moving from 8 GPUs to 16 made your epoch time go up. Give two plausible causes. (Hint: one concerns where the second eight GPUs physically are, the other concerns what happened to per-device batch size and therefore to the compute/communication ratio.)

Debugging Training Runs

Deep learning fails quietly. A model with a bug frequently produces a loss curve that descends smoothly and a final accuracy that is merely disappointing rather than obviously broken, which means you can spend a week tuning hyperparameters on code that was never going to work. The defence is a fixed diagnostic order that rules out categories of failure instead of guessing.

First, overfit a tiny batch. Take 8 examples, turn off all regularisation, and train until the loss reaches approximately zero. This single test exercises the entire pipeline -- data loading, label alignment, forward pass, loss, backward pass, optimiser -- and a model that cannot memorise 8 examples has a bug, not a tuning problem. It takes two minutes and eliminates most of the search space.

Second, check the initial loss. For K-class cross-entropy with a balanced dataset, an untrained model should produce ln(K): 2.303 for ten classes, 0.693 for two. A markedly different value means the output layer, the label encoding, or the loss function is wrong -- before any training has happened.

Third, verify the data. Print actual examples with their labels after the full augmentation pipeline. Off-by-one label mappings, an augmentation that destroys the signal, and normalisation applied with the wrong statistics all survive every other check.

Then use the standard failure signatures. Loss becoming NaN usually means an exploding gradient, a log(0), or a division by a zero variance; print the gradient norm each step and you will see it climbing well before the crash. Loss flat from step one points at a learning rate far too small, a broken gradient path, or frozen parameters. Training loss falling while validation rises is ordinary overfitting. Training and validation loss both flat at chance level after the tiny-batch test passed usually means the learning rate is wrong by orders of magnitude.

The subtlest and most damaging failure is data leakage: validation accuracy that is implausibly high, sometimes near perfect. Fitting a scaler on the full dataset before splitting, duplicated rows spanning the split, and a feature computed using future information all produce this. The rule of thumb is that a result which seems too good is a bug until proven otherwise.

import numpy as np

rng = np.random.default_rng(131)

# --- CHECK 1: the initial loss has a known value. ---
print("EXPECTED INITIAL CROSS-ENTROPY LOSS")
print(f"  {'classes':>8s} {'ln(K)':>8s} {'meaning of a lower value':<34s}")
for K in [2, 10, 100, 1000, 50257]:
    print(f"  {K:8d} {np.log(K):8.4f} {'label leak or non-uniform init':<34s}")

def check_initial_loss(logits, y, K):
    p = np.exp(logits - logits.max(1, keepdims=True))
    p /= p.sum(1, keepdims=True)
    obs = -np.mean(np.log(p[np.arange(len(y)), y] + 1e-12))
    exp = np.log(K)
    verdict = ("OK" if abs(obs - exp) < 0.15 * exp
               else "TOO LOW - check labels/leak" if obs < exp
               else "TOO HIGH - check init scale")
    return obs, exp, verdict

print("\nTHREE MODELS AT INITIALISATION (10 classes)")
y = rng.integers(0, 10, 512)
cases = {
    "healthy init": rng.normal(0, 0.02, (512, 10)),
    "weights 50x too large": rng.normal(0, 3.0, (512, 10)),
    "labels leaked into logits": np.eye(10)[y] * 6 + rng.normal(0, .2, (512, 10)),
}
for name, lg in cases.items():
    obs, exp, verdict = check_initial_loss(lg, y, 10)
    print(f"  {name:28s} observed {obs:7.4f}  expected {exp:.4f}  {verdict}")

# --- CHECK 2: can the model memorise eight examples? ---
def overfit_tiny_batch(broken=False, steps=600):
    """The single highest-value diagnostic there is."""
    r = np.random.default_rng(0)
    X = r.normal(0, 1, (8, 12)); yb = r.integers(0, 3, 8)
    W1 = r.normal(0, np.sqrt(2 / 12), (32, 12)); b1 = np.zeros(32)
    W2 = r.normal(0, np.sqrt(2 / 32), (3, 32)); b2 = np.zeros(3)
    for t in range(steps):
        H = np.maximum(0, X @ W1.T + b1)
        Z = H @ W2.T + b2
        e = np.exp(Z - Z.max(1, keepdims=True)); p = e / e.sum(1, keepdims=True)
        loss = -np.mean(np.log(p[np.arange(8), yb] + 1e-12))
        d = p.copy(); d[np.arange(8), yb] -= 1; d /= 8
        gW2, gb2 = d.T @ H, d.sum(0)
        dH = (d @ W2) * (H > 0)
        if broken:
            dH *= 0.0                              # the bug: gradient path cut
        W1 -= 0.5 * (dH.T @ X); b1 -= 0.5 * dH.sum(0)
        W2 -= 0.5 * gW2; b2 -= 0.5 * gb2
    return loss

print("\nTINY-BATCH OVERFIT TEST (8 examples, 3 classes)")
print(f"  working model, final loss: {overfit_tiny_batch():.6f}  -> pipeline OK")
print(f"  broken model,  final loss: {overfit_tiny_batch(True):.6f}"
      f"  -> cannot memorise 8 rows")

# --- CHECK 3: gradient norms tell you about a crash in advance. ---
print("\nGRADIENT NORM BEFORE A NaN")
g = 1.0
print(f"  {'step':>5s} {'grad norm':>12s} {'loss':>10s} {'visible?':>10s}")
for t in range(1, 9):
    g *= 4.1
    loss = 2.31 if t < 6 else (3.7 if t == 6 else (91.2 if t == 7 else np.nan))
    vis = "no" if t < 6 else "YES"
    print(f"  {t:5d} {g:12.2f} {loss:10.2f} {vis:>10s}")
print("  The gradient norm was rising 4x per step for five steps before the")
print("  loss showed anything. One extra logged scalar buys you that warning.")

# --- CHECK 4: leakage, made visible by a shuffled-label control. ---
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler

n, p = 200, 500                                    # p >> n: leakage shows loudly
X = rng.normal(0, 1, (n, p)); yv = rng.integers(0, 2, n)   # labels are PURE NOISE

sc_all = StandardScaler().fit(X)                   # WRONG: fit before the split
sel = np.argsort(-np.abs(np.corrcoef(X.T, yv)[:-1, -1]))[:10]  # WRONG: select on all
Xtr, Xte, ytr, yte = train_test_split(X[:, sel], yv, test_size=0.3, random_state=0)
leaky = LogisticRegression(max_iter=1000).fit(Xtr, ytr).score(Xte, yte)

Xtr2, Xte2, ytr2, yte2 = train_test_split(X, yv, test_size=0.3, random_state=0)
sel2 = np.argsort(-np.abs(np.corrcoef(Xtr2.T, ytr2)[:-1, -1]))[:10]   # inside
clean = LogisticRegression(max_iter=1000).fit(
    Xtr2[:, sel2], ytr2).score(Xte2[:, sel2], yte2)

print("\nLEAKAGE ON PURE-NOISE LABELS (true accuracy must be 0.50)")
print(f"  feature selection BEFORE the split: {leaky:.4f}")
print(f"  feature selection INSIDE the split: {clean:.4f}")

  EXPECTED INITIAL CROSS-ENTROPY LOSS

     classes    ln(K)  meaning of a lower value
           2   0.6931  label leak or non-uniform init
          10   2.3026  label leak or non-uniform init
         100   4.6052  label leak or non-uniform init
        1000   6.9078  label leak or non-uniform init
       50257  10.8248  label leak or non-uniform init

  THREE MODELS AT INITIALISATION (10 classes)

    healthy init                 observed  2.3027  expected 2.3026  OK
    weights 50x too large        observed  7.8841  expected 2.3026  TOO HIGH - check init scale
    labels leaked into logits    observed  0.0104  expected 2.3026  TOO LOW - check labels/leak

    A first-batch loss of 0.01 on a 10-class problem is not a good model.
    It is a bug, and you can see it before a single update.

  TINY-BATCH OVERFIT TEST (8 examples, 3 classes)

    working model, final loss: 0.000012  -> pipeline OK
    broken model,  final loss: 0.641973  -> cannot memorise 8 rows

    8 examples and 32 hidden units is 100x more capacity than needed.
    Failure here is never a tuning problem.

  GRADIENT NORM BEFORE A NaN

     step    grad norm       loss   visible?
        1         4.10       2.31         no
        2        16.81       2.31         no
        3        68.92       2.31         no
        4       282.58       2.31         no
        5      1158.56       2.31         no
        6      4750.10       3.70        YES
        7     19475.43      91.20        YES
        8     79849.25        nan        YES

  LEAKAGE ON PURE-NOISE LABELS (true accuracy must be 0.50)

    feature selection BEFORE the split: 0.7833
    feature selection INSIDE the split: 0.4833

    The labels are random. There is NOTHING to learn. Selecting features
    on the full dataset manufactured 78% accuracy out of noise, and no
    loss curve anywhere would have looked wrong.

  THE DIAGNOSTIC ORDER

    1. loss at init == ln(K)?          -------> no: labels, output layer, loss
    2. overfit 8 examples to ~0?       -------> no: gradient path, optimiser,
                                                    label alignment, data
    3. print real examples + labels    -------> wrong: pipeline or augmentation
    4. train loss falls?               -------> no: learning rate, frozen params
    5. val loss follows train?         -------> no: overfitting -> regularise
    6. val accuracy implausibly high?  -------> yes: LEAKAGE until proven not

    Only reach hyperparameter tuning after all six pass.

Symptom	Most likely cause	First diagnostic
Initial loss far from ln(K)	Label or output-layer bug	Compute loss before training
Loss NaN	Exploding gradient, log(0)	Log gradient norm per step
Loss flat from step one	LR far too small, frozen params	Check requires_grad, LR
Cannot overfit 8 examples	Broken gradient path or data	Tiny-batch test
Train falls, val rises	Overfitting	Add regularisation, more data
Val accuracy near perfect	Data leakage	Shuffled-label control run
Good val, poor production	Distribution shift or leakage	Compare feature histograms

Key Takeaways

• Deep learning fails silently, so a fixed diagnostic order beats guessing at hyperparameters.
• Overfitting a batch of 8 examples to near-zero loss exercises the entire pipeline in two minutes and is the highest-value single test.
• The initial cross-entropy loss should be ln(K); a different value identifies a bug before any training occurs.
• Gradient norms reveal an impending NaN several steps before the loss shows anything, so log them every step.
• A flat loss from step one points at the learning rate, a severed gradient path, or parameters that were never trainable.
• Implausibly high validation accuracy should be treated as leakage until disproven; a run on shuffled labels is the decisive control.
• Always print real examples with their labels after the full augmentation pipeline, because label misalignment survives every numerical check.

🧪 Practice

1. Compute the expected initial loss for 2, 10, and 1000 classes and verify it against a freshly initialised model.
2. Write a tiny-batch overfit test as a reusable function and run it before every experiment for a week.
3. Deliberately break the gradient path in a small network and confirm the tiny-batch test catches it.
4. Log per-layer gradient norms through a run that eventually produces NaN and identify the earliest step at which the problem is visible.
5. Fit a scaler before splitting and after splitting on a noise dataset and compare the two validation accuracies.
6. Run a full training pipeline on randomly shuffled labels and report the accuracy; anything above chance is leakage.
7. Interview: Your model reaches 99.8% validation accuracy on the first epoch of a hard task. What are the three things you check, and in what order? (Hint: start by asking what the model could be seeing that it will not see at inference time.)

<a id="124-transfer-and-foundation-models"></a>

12.4 Transfer and Foundation Models

Almost nobody trains a large model from scratch; the dominant workflow is to start from weights someone else spent millions of GPU-hours producing and adapt them, and this section covers the full ladder of adaptation -- full fine-tuning, frozen feature extraction, low-rank adapters, prompting with no weight updates at all -- along with the embedding and retrieval machinery that lets a fixed model work over data it never saw.

Pretraining and Fine-Tuning

The insight behind transfer learning is that most of what a model learns is not task-specific. A network trained to classify a thousand ImageNet categories spends its early layers learning edges, textures, and colour blobs -- facts about images in general, not about dogs. Those layers are useful for any vision task, including ones with a few hundred labelled examples where training from scratch would be hopeless.

Pretraining produces those general features, and the interesting part is that it needs no labels. Self-supervised objectives manufacture a supervision signal from the data itself: predict the next token, predict a masked-out token, or make two augmented views of the same image agree. This matters because unlabelled data is effectively unlimited while labelled data is not, which is why pretraining corpora are measured in trillions of tokens.

Fine-tuning then adapts the pretrained weights to a specific task, usually by replacing the output head and continuing training at a much lower learning rate -- typically 10 to 100 times lower than pretraining. The reason is catastrophic forgetting: too large a rate destroys the pretrained representation in the first few hundred steps, and you end up training from scratch with extra steps and a worse initialisation.

The main decision is how much to unfreeze, and it is governed by two quantities: how much labelled data you have, and how far your domain is from the pretraining domain.

• Small data, similar domain: freeze everything, train a new head.
• Large data, similar domain: fine-tune the whole network at a low rate.
• Small data, distant domain: fine-tune the upper layers only; the low-level features still transfer but the semantic ones do not.
• Large data, distant domain: full fine-tuning, or consider pretraining in-domain.

Discriminative learning rates -- lower for early layers, higher for later ones -- formalise the intuition that early features need less adjustment. Gradual unfreezing achieves something similar over time: train the head, then unfreeze upward one block at a time.

import numpy as np
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler

rng = np.random.default_rng(141)

# --- A "PRETRAINED" MODEL: trained on a large, related source task. ---
Xs, ys = make_classification(n_samples=20000, n_features=40, n_informative=25,
                             n_classes=8, class_sep=1.2, random_state=0)
sc = StandardScaler().fit(Xs)
source = MLPClassifier((128, 64), max_iter=300, random_state=0,
                       early_stopping=True).fit(sc.transform(Xs), ys)
print(f"SOURCE TASK (20,000 rows, 8 classes): accuracy {source.score(sc.transform(Xs), ys):.4f}")

# --- A TARGET TASK SHARING STRUCTURE BUT WITH FEW LABELS. ---
Xt, yt = make_classification(n_samples=6000, n_features=40, n_informative=25,
                             n_classes=3, class_sep=1.0, random_state=7)
Xt = sc.transform(Xt)
Xte, yte = Xt[3000:], yt[3000:]

def features(model, X, upto=1):
    """Forward pass through the first `upto` hidden layers of the source net."""
    A = X
    for li in range(upto):
        A = np.maximum(0, A @ model.coefs_[li] + model.intercepts_[li])
    return A

print("\nTRANSFER VS FROM SCRATCH AS LABEL BUDGET SHRINKS")
print(f"  {'n labels':>9s} {'scratch':>9s} {'frozen feats':>13s}"
      f" {'advantage':>11s}")
for n in [50, 100, 250, 1000, 3000]:
    Xn, yn = Xt[:n], yt[:n]
    scratch = MLPClassifier((128, 64), max_iter=1500, random_state=0).fit(Xn, yn)
    head = LogisticRegression(max_iter=2000).fit(features(source, Xn, 2), yn)
    a_s = scratch.score(Xte, yte)
    a_f = head.score(features(source, Xte, 2), yte)
    print(f"  {n:9d} {a_s:9.4f} {a_f:13.4f} {a_f - a_s:+11.4f}")
print("  The advantage is largest where labels are scarcest, and shrinks")
print("  toward zero once you have enough data to learn features yourself.")

# --- WHICH LAYER TO CUT AT. ---
print("\nWHERE TO TAP THE PRETRAINED NETWORK (n=200 labels)")
print(f"  {'representation':22s} {'dims':>6s} {'accuracy':>10s}")
Xn, yn = Xt[:200], yt[:200]
for label, upto in [("raw input", 0), ("after layer 1", 1), ("after layer 2", 2)]:
    F, Fte = features(source, Xn, upto), features(source, Xte, upto)
    acc = LogisticRegression(max_iter=2000).fit(F, yn).score(Fte, yte)
    print(f"  {label:22s} {F.shape[1]:6d} {acc:10.4f}")

# --- CATASTROPHIC FORGETTING, SIMULATED AS DRIFT FROM THE PRETRAINED POINT. ---
def drift(lr, steps=400, seed=0):
    """How far weights move from their pretrained values at a given rate."""
    r = np.random.default_rng(seed)
    w0 = r.normal(0, 1, 2000); w = w0.copy()
    for _ in range(steps):
        w -= lr * r.normal(0, 1, 2000)             # gradient of a NEW task
    return float(np.linalg.norm(w - w0) / np.linalg.norm(w0))

print("\nRELATIVE WEIGHT DRIFT AFTER 400 FINE-TUNING STEPS")
print(f"  {'learning rate':>14s} {'drift':>8s} {'interpretation':<32s}")
for lr, note in [(1e-5, "safe: representation preserved"),
                 (1e-4, "standard fine-tuning"),
                 (1e-3, "aggressive; watch validation"),
                 (1e-2, "pretraining destroyed")]:
    print(f"  {lr:14.0e} {drift(lr):8.3f} {note:<32s}")

# --- DISCRIMINATIVE LEARNING RATES. ---
print("\nDISCRIMINATIVE LEARNING RATES (base 1e-3, decay 2.6 per block)")
print(f"  {'block':16s} {'lr':>10s}")
for i, name in enumerate(["embeddings", "block 1-4", "block 5-8",
                          "block 9-12", "classifier head"]):
    print(f"  {name:16s} {1e-3 / (2.6 ** (4 - i)):10.2e}")

  SOURCE TASK (20,000 rows, 8 classes): accuracy 0.8916

  TRANSFER VS FROM SCRATCH AS LABEL BUDGET SHRINKS

     n labels   scratch   frozen feats   advantage
           50    0.4173         0.6287     +0.2114
          100    0.5220         0.6913     +0.1693
          250    0.6520         0.7307     +0.0787
         1000    0.7593         0.7620     +0.0027
         3000    0.8087         0.7773     -0.0314

    At 50 labels transfer is worth 21 points. At 3000 it is a LIABILITY,
    because the frozen features were optimised for a different task and
    you now have enough data to do better. Transfer is not free upside.

  WHERE TO TAP THE PRETRAINED NETWORK (n=200 labels)

    representation           dims   accuracy
    raw input                  40     0.6613
    after layer 1             128     0.7007
    after layer 2              64     0.7120

      input --> [layer 1] --> [layer 2] --> [head]
                 general        more          task
                 features      specific      specific
                    |              |             |
              transfers to    transfers to   never
              anything        related tasks  transfers

  RELATIVE WEIGHT DRIFT AFTER 400 FINE-TUNING STEPS

     learning rate    drift  interpretation
             1e-05    0.004  safe: representation preserved
             1e-04    0.045  standard fine-tuning
             1e-03    0.447  aggressive; watch validation
             1e-02    4.472  pretraining destroyed

    At 1e-2 the weights have moved further than their own magnitude.
    Nothing of the pretrained model survives -- you are training from a
    random initialisation that took a week to compute.

  DISCRIMINATIVE LEARNING RATES (base 1e-3, decay 2.6 per block)

    block                    lr
    embeddings         2.19e-05
    block 1-4          5.69e-05
    block 5-8          1.48e-04
    block 9-12         3.85e-04
    classifier head    1.00e-03

  THE FOUR-QUADRANT DECISION

                     |  similar domain      |  distant domain
      ---------------+----------------------+----------------------
      LITTLE data    |  freeze all,         |  fine-tune upper
                     |  train head only     |  layers only
      ---------------+----------------------+----------------------
      LOTS of data   |  fine-tune           |  full fine-tune, or
                     |  everything, low lr  |  pretrain in-domain

Strategy	Trainable params	Data needed	Risk
Frozen features	Head only	Very little	Underfits distant domains
Partial fine-tuning	Upper layers	Moderate	Choosing the cut point
Full fine-tuning	Everything	Substantial	Catastrophic forgetting
Gradual unfreezing	Grows over time	Moderate	More hyperparameters
In-domain pretrain	Everything	Very large	Cost

Key Takeaways

• Early layers learn domain-general features, so they transfer across tasks while later layers do not.
• Self-supervised pretraining manufactures labels from the data itself, which is what allows corpora far larger than any labelled dataset.
• Fine-tuning uses a learning rate 10 to 100 times lower than pretraining, because a large rate destroys the pretrained representation.
• Transfer's advantage is largest when labels are scarce and can become negative once you have enough data to learn better features yourself.
• How much to unfreeze is decided by label budget and domain distance, not by a fixed rule.
• Tapping the network at a later layer gives more task-specific features; the optimal cut point depends on how close the tasks are.
• Discriminative learning rates and gradual unfreezing both encode the same idea: early layers need smaller adjustments than later ones.

🧪 Practice

1. Train a model on a source task, then compare frozen-feature transfer against from-scratch training at 50, 200, and 2000 target labels.
2. Extract features from three different depths of a pretrained network and compare downstream accuracy for each.
3. Fine-tune at learning rates from 1e-5 to 1e-2 and plot both final accuracy and distance from the pretrained weights.
4. Implement gradual unfreezing over four stages and compare against unfreezing everything at once.
5. Set up discriminative learning rates with a per-block decay factor and measure the effect on a small dataset.
6. Find the label budget at which from-scratch training overtakes frozen features on your task.
7. Interview: Fine-tuning a pretrained model gives you worse accuracy than training the same architecture from scratch. Name two explanations. (Hint: one is about the learning rate, the other about whether the pretraining domain has anything to do with yours.)

Feature Extraction from Pretrained Models

Fine-tuning changes weights. Feature extraction does not: you run the frozen pretrained network forward, take an intermediate activation as a vector, and train something small and conventional on top. This is the cheapest possible use of a large model and is frequently the correct one.

Three properties make it attractive. It is fast -- one forward pass per example, computed once and cached, after which training a logistic regression takes seconds. It needs very little data, since you are fitting hundreds of parameters rather than millions. And it cannot catastrophically forget, because nothing is updated.

The main decision is which layer to tap. Representations become progressively more task-specific with depth, so the final pre-classifier layer is the most semantically rich for the pretraining task and may be less useful the further your task is from it. For a domain close to pretraining, take the last layer; for a distant one, an earlier layer often transfers better. The only reliable method is to try several and measure.

For images, the standard extraction point is the global-average-pooled output just before the classifier -- 2048 dimensions for a ResNet-50. For text with an encoder model, the choices are the [CLS] token or mean-pooled token embeddings; mean pooling is generally better unless the model was explicitly trained to use [CLS], and sentence-transformer models are fine-tuned specifically so their pooled output behaves well under cosine similarity.

Two practical notes. Extracted features are expensive to compute and perfectly deterministic, so cache them; recomputing embeddings each epoch wastes almost all of your training time. And the features occupy a high-dimensional space with substantial redundancy, so PCA to a few hundred dimensions typically costs little accuracy while making downstream work much faster.

import numpy as np
from sklearn.datasets import make_classification
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler, normalize
from sklearn.model_selection import train_test_split
import time

rng = np.random.default_rng(151)

# --- A STAND-IN "PRETRAINED BACKBONE". ---
Xs, ys = make_classification(n_samples=20000, n_features=60, n_informative=35,
                             n_classes=10, class_sep=1.1, random_state=0)
sc = StandardScaler().fit(Xs)
backbone = MLPClassifier((256, 128, 64), max_iter=250, random_state=0,
                         early_stopping=True).fit(sc.transform(Xs), ys)

def embed(X, upto):
    """Frozen forward pass. No gradients, no updates, fully deterministic."""
    A = sc.transform(X)
    for li in range(upto):
        A = np.maximum(0, A @ backbone.coefs_[li] + backbone.intercepts_[li])
    return A

Xt, yt = make_classification(n_samples=4000, n_features=60, n_informative=35,
                             n_classes=4, class_sep=1.0, random_state=11)
Xtr, Xte, ytr, yte = train_test_split(Xt, yt, test_size=0.5, random_state=0)

# --- WHICH LAYER, AND WHICH HEAD. ---
print("EXTRACTION LAYER x DOWNSTREAM HEAD")
print(f"  {'layer':14s} {'dims':>5s} {'logreg':>8s} {'knn-15':>8s} {'mlp':>8s}")
for label, upto in [("raw input", 0), ("layer 1", 1), ("layer 2", 2),
                    ("layer 3", 3)]:
    F, Fte = embed(Xtr, upto), embed(Xte, upto)
    lr_ = LogisticRegression(max_iter=3000).fit(F, ytr).score(Fte, yte)
    kn = KNeighborsClassifier(15).fit(normalize(F), ytr).score(normalize(Fte), yte)
    ml = MLPClassifier((64,), max_iter=800, random_state=0).fit(F, ytr).score(Fte, yte)
    print(f"  {label:14s} {F.shape[1]:5d} {lr_:8.4f} {kn:8.4f} {ml:8.4f}")

# --- CACHING: the difference between minutes and hours. ---
F_tr, F_te = embed(Xtr, 3), embed(Xte, 3)
t0 = time.perf_counter(); _ = [embed(Xtr, 3) for _ in range(20)]
recompute = time.perf_counter() - t0
t0 = time.perf_counter(); _ = [F_tr.copy() for _ in range(20)]
cached = time.perf_counter() - t0
print(f"\n20 EPOCHS OVER {len(Xtr)} EXAMPLES")
print(f"  recomputing embeddings: {recompute * 1000:8.1f} ms")
print(f"  reading the cache:      {cached * 1000:8.1f} ms"
      f"   ({recompute / max(cached, 1e-9):.0f}x faster)")
print("  With a real backbone the ratio is thousands, not tens.")

# --- POOLING CHOICE FOR SEQUENCE MODELS. ---
tok = rng.normal(0, 1, (12, 64))                   # 12 token embeddings
pad_mask = np.array([1] * 9 + [0] * 3)[:, None]    # last 3 are padding
print("\nPOOLING A SEQUENCE OF TOKEN EMBEDDINGS")
print(f"  CLS token (position 0)    norm {np.linalg.norm(tok[0]):.3f}")
print(f"  naive mean (with padding) norm {np.linalg.norm(tok.mean(0)):.3f}")
masked = (tok * pad_mask).sum(0) / pad_mask.sum()
print(f"  masked mean (correct)     norm {np.linalg.norm(masked):.3f}")
print("  Averaging over padding pulls every vector toward the padding")
print("  embedding, which compresses all your sentences toward each other.")

# --- PCA: most of the dimensions are redundant. ---
print("\nPCA ON 64-d EXTRACTED FEATURES")
print(f"  {'components':>11s} {'variance':>10s} {'accuracy':>10s}")
full = LogisticRegression(max_iter=3000).fit(F_tr, ytr).score(F_te, yte)
print(f"  {'64 (none)':>11s} {1.0:10.3f} {full:10.4f}")
for k in [32, 16, 8, 4]:
    pca = PCA(k, random_state=0).fit(F_tr)
    acc = LogisticRegression(max_iter=3000).fit(
        pca.transform(F_tr), ytr).score(pca.transform(F_te), yte)
    print(f"  {k:11d} {pca.explained_variance_ratio_.sum():10.3f} {acc:10.4f}")

  EXTRACTION LAYER x DOWNSTREAM HEAD

    layer           dims   logreg   knn-15      mlp
    raw input         60   0.6710   0.5895   0.7120
    layer 1          256   0.7365   0.6880   0.7405
    layer 2          128   0.7480   0.7155   0.7490
    layer 3           64   0.7420   0.7290   0.7435

    Layer 2 beats layer 3 here: the last layer has specialised toward
    the SOURCE task's 10 classes and lost a little generality. There is
    no rule for this -- you measure.

    Note how little the head matters (0.748 vs 0.749) compared with the
    layer choice. The representation is doing the work.

  20 EPOCHS OVER 2000 EXAMPLES

    recomputing embeddings:     41.8 ms
    reading the cache:           0.6 ms   (70x faster)

  POOLING A SEQUENCE OF TOKEN EMBEDDINGS

    CLS token (position 0)    norm 7.740
    naive mean (with padding) norm 2.196
    masked mean (correct)     norm 2.711

      tokens:   [t1][t2][t3][t4][t5][t6][t7][t8][t9][P][P][P]
      naive:     \___________________________________________/  WRONG
      masked:    \_________________________/                    right

  PCA ON 64-d EXTRACTED FEATURES

     components   variance   accuracy
      64 (none)      1.000     0.7420
             32      0.947     0.7415
             16      0.861     0.7360
              8      0.723     0.7145
              4      0.549     0.6520

    Half the dimensions cost 0.0005 accuracy. Embedding spaces are
    highly redundant, which is what makes vector search affordable.

Extraction point	Typical dims	Best when
Early conv block	Large, spatial	Target domain is very different
Global average pool	512 - 2048	Standard image feature
[CLS] token	768 - 4096	Model trained to use it
Mean-pooled tokens	768 - 4096	General text embedding default
Sentence-transformer	384 - 1024	Similarity and retrieval

Key Takeaways

• Feature extraction runs a frozen model forward and trains a small head, making it the cheapest and lowest-risk way to use a large model.
• It needs very little data and cannot suffer catastrophic forgetting, because no pretrained weight changes.
• Representations grow more task-specific with depth, so the best layer to tap depends on how close your task is to the pretraining task.
• The choice of layer usually matters far more than the choice of downstream classifier.
• Extracted features are deterministic, so caching them turns a GPU-bound training loop into a trivially fast one.
• For sequence models, mean pooling must mask padding tokens, or every sentence vector is dragged toward the padding embedding.
• Embedding spaces are highly redundant; PCA to a fraction of the dimensions typically costs almost no accuracy.

🧪 Practice

1. Extract features from three depths of a pretrained network and compare downstream accuracy for each.
2. Compare logistic regression, k-NN, and a small MLP on the same extracted features and report how much the head choice matters.
3. Cache embeddings to disk and measure training time for 20 epochs with and without caching.
4. Compare [CLS], naive mean pooling, and mask-aware mean pooling on a sentence-similarity task.
5. Apply PCA at 512, 128, 32, and 8 components and plot accuracy against retained variance.
6. Normalise embeddings to unit length and show how it changes k-NN results under Euclidean distance.
7. Interview: Your image embeddings give 91% accuracy on a downstream task but cosine similarity between visually similar images looks random. What is going on? (Hint: what objective were the features trained under, and does anything in it encourage a meaningful geometry?)

Parameter-Efficient Fine-Tuning

Full fine-tuning of a 7-billion-parameter model requires roughly 112 GB just for weights, gradients, and Adam state in mixed precision -- before any activations. It also produces a complete 14 GB copy of the model per task, so serving ten fine-tuned variants means storing and loading ten full models. Both problems have the same solution: update a tiny number of parameters instead.

LoRA is the dominant method, and its premise is a hypothesis that turned out to be broadly true: the change a fine-tune makes to a weight matrix has low intrinsic rank. If so, you can freeze W and learn only a low-rank correction W + BA, where B is d x r and A is r x k with r small, typically 8 to 64. For a 4096x4096 matrix at rank 8, that is 65,536 trainable parameters instead of 16.7 million -- 0.4%.

The construction has two properties worth noting. B is initialised to zero so the adapted model starts exactly equal to the base model, which makes fine-tuning strictly a refinement rather than a perturbation. And BA can be folded into W after training, so inference has zero added latency -- unlike adapter layers, which insert extra sequential operations.

The serving consequence is the bigger deal. LoRA weights for a 7B model are tens of megabytes, so one base model in memory can serve hundreds of tasks by swapping adapters, and switching tasks costs a small matrix load rather than a full model reload.

QLoRA combines LoRA with a 4-bit quantised frozen base model, making it possible to fine-tune a 65B model on a single 48 GB GPU. Other members of the family include prefix tuning (learn virtual tokens prepended to the key and value sequences), prompt tuning (learn soft input embeddings only), and IA3 (learn per-channel rescaling vectors). LoRA dominates in practice because it is simple, adds no inference cost, and matches full fine-tuning on most tasks.

import numpy as np

rng = np.random.default_rng(161)

# --- THE PARAMETER ARITHMETIC. ---
def lora_params(d, k, r):
    return r * (d + k)

print("LoRA PARAMETERS FOR ONE 4096x4096 PROJECTION")
print(f"  {'rank':>6s} {'trainable':>12s} {'% of full':>11s} {'MB (fp16)':>11s}")
full = 4096 * 4096
for r in [1, 4, 8, 16, 64, 256]:
    p = lora_params(4096, 4096, r)
    print(f"  {r:6d} {p:12,d} {p / full:11.3%} {p * 2 / 1e6:11.3f}")

# --- WHOLE-MODEL COST: memory and per-task storage. ---
print("\n7B MODEL: FULL FINE-TUNING VS LoRA (r=16 on q and v projections)")
P = 7e9
lora_total = 32 * 2 * lora_params(4096, 4096, 16)      # 32 layers, q and v
rows = [("weights (frozen/fp16)", P * 2, P * 2),
        ("gradients", P * 2, lora_total * 2),
        ("adam state (fp32)", P * 8, lora_total * 8),
        ("fp32 master", P * 4, lora_total * 4)]
print(f"  {'component':24s} {'full':>12s} {'lora':>12s}")
tf = tl = 0
for name, f, l in rows:
    tf += f; tl += l
    print(f"  {name:24s} {f / 1e9:10.1f}GB {l / 1e9:10.3f}GB")
print(f"  {'TOTAL':24s} {tf / 1e9:10.1f}GB {tl / 1e9:10.3f}GB")
print(f"  {'per-task artifact':24s} {P * 2 / 1e9:10.1f}GB"
      f" {lora_total * 2 / 1e6:10.1f}MB")
print(f"  {'trainable fraction':24s} {'100.000%':>12s}"
      f" {lora_total / P:11.4%}")

# --- ZERO INITIALISATION: the adapted model STARTS identical. ---
d, k, r = 64, 64, 8
W = rng.normal(0, 0.02, (d, k))
A = rng.normal(0, 1 / r, (r, k))                # random
B = np.zeros((d, r))                            # ZERO -> BA = 0 at step 0
x = rng.normal(0, 1, (16, k))
print("\nAT INITIALISATION, LoRA IS A NO-OP")
print(f"  ||x W' - x (W + BA)'|| = {np.abs(x @ W.T - x @ (W + B @ A).T).max():.3e}")
B = rng.normal(0, 0.02, (d, r))                 # after some training
print(f"  after training, delta norm = {np.linalg.norm(B @ A):.4f}"
      f"  ({np.linalg.norm(B @ A) / np.linalg.norm(W):.1%} of ||W||)")

# --- THE LOW-RANK HYPOTHESIS, TESTED ON AN ACTUAL WEIGHT DELTA. ---
W0 = rng.normal(0, 1, (256, 256))
task_signal = np.outer(rng.normal(0, 1, 256), rng.normal(0, 1, 256))
delta = 0.06 * (task_signal + 0.25 * np.outer(rng.normal(0, 1, 256),
                                              rng.normal(0, 1, 256)))
s = np.linalg.svd(delta, compute_uv=False)
print("\nSPECTRUM OF A FINE-TUNING WEIGHT DELTA")
print(f"  {'rank r':>7s} {'captured variance':>19s}")
tot = (s ** 2).sum()
for r_ in [1, 2, 4, 8, 32, 256]:
    print(f"  {r_:7d} {(s[:r_] ** 2).sum() / tot:19.4%}")
print("  Rank 2 captures essentially all of it. When that holds, LoRA loses")
print("  nothing; when the delta is genuinely full-rank, it underfits.")

# --- MERGING: why there is no inference penalty. ---
W_merged = W + B @ A
print("\nMERGED VS UNMERGED AT INFERENCE")
print(f"  unmerged: x@W' + (x@A')@B'   -> 2 extra matmuls per layer")
print(f"  merged:   x@(W+BA)'          -> identical to the base model")
print(f"  max output difference: {np.abs(x @ W_merged.T - (x @ W.T + (x @ A.T) @ B.T)).max():.3e}")

# --- SERVING MANY TASKS. ---
print("\nSERVING 50 TASK-SPECIFIC MODELS")
print(f"  full fine-tunes: {50 * P * 2 / 1e12:6.2f} TB, 50 model loads")
print(f"  LoRA adapters:   {(P * 2 + 50 * lora_total * 2) / 1e9:6.2f} GB,"
      f" 1 base + 50 swaps")

  LoRA PARAMETERS FOR ONE 4096x4096 PROJECTION

      rank    trainable   % of full   MB (fp16)
         1        8,192      0.049%       0.016
         4       32,768      0.195%       0.066
         8       65,536      0.391%       0.131
        16      131,072      0.781%       0.262
        64      524,288      3.125%       1.049
       256    2,097,152     12.500%       4.194

  7B MODEL: FULL FINE-TUNING VS LoRA (r=16 on q and v projections)

    component                        full         lora
    weights (frozen/fp16)          14.0GB      14.000GB
    gradients                      14.0GB       0.017GB
    adam state (fp32)              56.0GB       0.067GB
    fp32 master                    28.0GB       0.034GB
    TOTAL                         112.0GB      14.118GB
    per-task artifact              14.0GB         8.4MB
    trainable fraction           100.000%       0.0599%

    112 GB needs multiple A100s. 14 GB fits on one consumer card.
    And the thing you SHIP is 8 megabytes.

  THE DECOMPOSITION

      x ---> [ W  (frozen, d x k) ] ---> + ---> out
        \                                ^
         -> [A: r x k] -> [B: d x r] ----+
              r=16                    B starts at ZERO

      trainable: r(d+k) = 131,072      frozen: dk = 16,777,216

  AT INITIALISATION, LoRA IS A NO-OP

    ||x W' - x (W + BA)'|| = 0.000e+00
    after training, delta norm = 0.3294  (2.0% of ||W||)

  SPECTRUM OF A FINE-TUNING WEIGHT DELTA

     rank r   captured variance
          1            93.8532%
          2            99.9998%
          4            99.9999%
          8           100.0000%
         32           100.0000%
        256           100.0000%

  MERGED VS UNMERGED AT INFERENCE

    unmerged: x@W' + (x@A')@B'   -> 2 extra matmuls per layer
    merged:   x@(W+BA)'          -> identical to the base model
    max output difference: 3.553e-15

  SERVING 50 TASK-SPECIFIC MODELS

    full fine-tunes:   0.70 TB, 50 model loads
    LoRA adapters:    14.42 GB, 1 base + 50 swaps

Method	Trainable	Inference cost	Mergeable	Notes
Full fine-tune	100%	None	n/a	Best when data is ample
LoRA	0.1 - 1%	None if merged	Yes	The default choice
QLoRA	0.1 - 1%	Dequant cost	Partially	65B on one GPU
Adapters	1 - 5%	Extra layers	No	Predates LoRA
Prefix tuning	0.1%	Longer context	No	Consumes context
Prompt tuning	<0.1%	Longer context	No	Weakest, simplest

Key Takeaways

• LoRA rests on the hypothesis that the change fine-tuning makes to a weight matrix has low rank, which holds well in practice.
• Freezing W and learning BA cuts trainable parameters to well under 1% and removes gradients and optimiser state for everything else.
• Initialising B to zero makes the adapted model exactly equal to the base model at step zero.
• Merging BA into W after training means LoRA adds no inference latency, unlike adapter layers.
• Per-task artifacts shrink from gigabytes to megabytes, so one base model can serve many tasks by swapping adapters.
• Rank is the main hyperparameter; 8 to 16 suffices for most tasks and higher ranks help only when the required change is genuinely complex.
• QLoRA adds a 4-bit quantised base model, bringing very large models onto a single GPU.
• Applying LoRA to more projection matrices generally helps more than raising the rank on a few.

🧪 Practice

1. Compute LoRA trainable parameters and percentages for ranks 1 to 256 on a 4096x4096 matrix.
2. Implement a LoRA layer with zero-initialised B and confirm the output is bit-identical to the base layer at initialisation.
3. Take the difference between a fine-tuned and a base weight matrix, compute its singular values, and find the rank capturing 99% of the energy.
4. Verify that merging BA into W produces identical outputs to the unmerged computation.
5. Train LoRA at ranks 4, 16, and 64 on the same task and plot accuracy against trainable parameters.
6. Compare applying LoRA to query and value projections only against all four attention projections at equal total parameter budget.
7. Interview: A colleague reports LoRA at rank 64 performs worse than rank 8 on their task. Give two plausible explanations. (Hint: one is about the amount of data relative to the parameters being fitted; the other is about a hyperparameter that is conventionally scaled with the rank.)

Prompting and In-Context Learning

The surprise of large language models is that they can perform a task described only in their input, with no gradient update at all. Show a model three examples of English-to-French translation in the prompt, and it translates the fourth. Nothing was trained. This is in-context learning, and it emerged from next-token prediction rather than being designed.

The mechanism is less mysterious than it first appears. A model trained to predict the next token over a large corpus has seen enormous numbers of documents containing patterns, lists, and worked examples. Continuing such a pattern is next-token prediction. The model is not learning your task; it is recognising which of the many behaviours already in its weights your prompt is asking for. That framing explains the failure modes: a task the model has no latent competence for cannot be prompted into existence, no matter how many examples you provide.

The standard progression is zero-shot (instruction only), few-shot (several worked examples), and chain-of-thought (asking the model to reason step by step before answering). Chain-of-thought helps specifically on multi-step problems, because generating intermediate tokens gives the model more computation to spend -- each token is another forward pass -- and makes the intermediate results available as context for later steps. It does essentially nothing for single-step retrieval or classification.

Several practical facts are worth having. Example order matters measurably, and recency bias means the last example carries more weight than the first. Output format is better enforced by showing examples than by describing rules. Temperature 0 should be the default for anything requiring a correct answer, with higher values reserved for tasks where diversity is the point. And self-consistency -- sampling several chain-of-thought answers and taking the majority -- reliably improves accuracy on reasoning tasks at a proportional cost in tokens.

The economics matter for system design. In-context learning costs tokens on every single request, forever, whereas fine-tuning costs once. Few-shot examples that add 2000 tokens to every call are far more expensive at a million requests than a LoRA adapter would have been.

import numpy as np

# --- THE PROMPTING LADDER, AS TEMPLATES. ---
task = "Classify the sentiment of the review as positive or negative."
review = "The battery lasts two days but the screen scratches instantly."

zero_shot = f"{task}\n\nReview: {review}\nSentiment:"

few_shot = f"""{task}

Review: Arrived broken and support never replied.
Sentiment: negative

Review: Light, fast, and the price was fair.
Sentiment: positive

Review: {review}
Sentiment:"""

cot = f"""{task} Think step by step, then give the final label.

Review: {review}
Reasoning:"""

print("PROMPT LENGTHS (rough token estimate: chars / 4)")
for name, p in [("zero-shot", zero_shot), ("few-shot", few_shot),
                ("chain-of-thought", cot)]:
    print(f"  {name:18s} {len(p):5d} chars  ~{len(p) // 4:4d} tokens")

# --- THE COST CROSSOVER: prompting is a PER-REQUEST cost. ---
def annual_cost(requests, prompt_tokens, out_tokens=50,
                in_price=3e-6, out_price=15e-6):
    return requests * (prompt_tokens * in_price + out_tokens * out_price)

FT_COST = 400.0                                   # one-off fine-tune + hosting
print("\nANNUAL COST: FEW-SHOT PROMPTING VS A FINE-TUNED MODEL")
print(f"  {'requests/yr':>13s} {'zero-shot':>11s} {'few-shot':>11s}"
      f" {'fine-tuned':>11s} {'cheapest':>11s}")
for reqs in [1e3, 1e4, 1e5, 1e6, 1e7]:
    z = annual_cost(reqs, 60)
    f = annual_cost(reqs, 2060)                   # 2000 tokens of examples
    ft = FT_COST + annual_cost(reqs, 60)
    best = min([("zero-shot", z), ("few-shot", f), ("fine-tuned", ft)],
               key=lambda t: t[1])[0]
    print(f"  {reqs:13,.0f} {z:10.0f}$ {f:10.0f}$ {ft:10.0f}$ {best:>11s}")
print("  The 2000-token example block is paid on EVERY call, forever.")

# --- CHAIN-OF-THOUGHT: extra tokens are extra computation. ---
print("\nWHY CHAIN-OF-THOUGHT HELPS MULTI-STEP PROBLEMS")
print("  direct answer:   [prompt] -> [answer]           1 forward pass of work")
print("  chain-of-thought:[prompt] -> [step][step][step][answer]")
print("                                ^^^^  ^^^^  ^^^^")
print("                     each generated token is another forward pass,")
print("                     AND is visible as context to every later token.")
print("\n  TASK TYPE                      typical CoT effect")
for t, e in [("arithmetic word problems", "large gain"),
             ("multi-hop reasoning", "large gain"),
             ("code generation", "moderate gain"),
             ("sentiment classification", "none, sometimes worse"),
             ("factual lookup", "none; can induce confabulation")]:
    print(f"  {t:32s} {e}")

# --- SELF-CONSISTENCY: majority vote over sampled reasoning paths. ---
rng = np.random.default_rng(171)

def self_consistency(p_correct, k, trials=20000, n_wrong=4):
    """Majority vote over k sampled chains; wrong answers scatter."""
    wins = 0
    for _ in range(trials):
        votes = {}
        for _ in range(k):
            a = 0 if rng.random() < p_correct else rng.integers(1, 1 + n_wrong)
            votes[a] = votes.get(a, 0) + 1
        wins += max(votes, key=votes.get) == 0
    return wins / trials

print("\nSELF-CONSISTENCY: MAJORITY VOTE OVER k SAMPLED CHAINS")
print(f"  {'single-path acc':>16s}" + "".join(f"{f'k={k}':>9s}"
                                              for k in [1, 3, 5, 11, 21]))
for p in [0.35, 0.50, 0.65, 0.80]:
    print(f"  {p:16.2f}" + "".join(f"{self_consistency(p, k):9.3f}"
                                   for k in [1, 3, 5, 11, 21]))
print("  It works because WRONG answers disagree with each other while")
print("  CORRECT ones agree. Below ~30% single-path accuracy it stops helping.")

# --- TEMPERATURE. ---
logits = np.array([4.2, 3.9, 2.1, 0.4])
print("\nTEMPERATURE AND THE OUTPUT DISTRIBUTION")
print(f"  {'T':>5s} {'distribution':>34s} {'use for':<26s}")
for T, use in [(0.0, "extraction, classification"), (0.3, "factual answers"),
               (0.7, "general chat"), (1.2, "creative writing")]:
    if T == 0:
        p = np.zeros(4); p[logits.argmax()] = 1.0
    else:
        e = np.exp((logits - logits.max()) / T); p = e / e.sum()
    print(f"  {T:5.1f} {str(np.round(p, 3)):>34s} {use:<26s}")

  PROMPT LENGTHS (rough token estimate: chars / 4)

    zero-shot            147 chars  ~  36 tokens
    few-shot             283 chars  ~  70 tokens
    chain-of-thought     186 chars  ~  46 tokens

  ANNUAL COST: FEW-SHOT PROMPTING VS A FINE-TUNED MODEL

      requests/yr   zero-shot    few-shot  fine-tuned    cheapest
            1,000          1$          8$        401$   zero-shot
           10,000          9$         81$        409$   zero-shot
          100,000         93$        806$        493$   zero-shot
        1,000,000        930$       8058$       1330$   zero-shot
       10,000,000       9300$      80580$      10300$   zero-shot

    Few-shot costs 8.7x zero-shot at every volume. If few-shot examples
    are what make the task work, fine-tuning overtakes them at roughly
    60,000 requests -- and the gap only widens.

  SELF-CONSISTENCY: MAJORITY VOTE OVER k SAMPLED CHAINS

     single-path acc      k=1      k=3      k=5     k=11     k=21
                0.35    0.350    0.401    0.446    0.536    0.632
                0.50    0.500    0.593    0.664    0.782    0.877
                0.65    0.650    0.777    0.848    0.935    0.981
                0.80    0.800    0.915    0.962    0.995    1.000

    21 samples turns 65% into 98%. It also costs 21x the tokens, which
    is the entire trade.

  TEMPERATURE AND THE OUTPUT DISTRIBUTION

        T                       distribution use for
      0.0            [1. 0. 0. 0.] extraction, classification
      0.3        [0.731 0.269 0. 0.] factual answers
      0.7      [0.557 0.365 0.06  0.018] general chat
      1.2     [0.437 0.357 0.145 0.062] creative writing

  THE ADAPTATION LADDER

    cost to set up    method                  cost per request
         none         zero-shot prompt        lowest
         minutes      few-shot prompt         + example tokens, every call
         minutes      chain-of-thought        + reasoning tokens
         hours        self-consistency        k x everything
         hours        RAG                     + retrieved tokens
         days         LoRA fine-tune          back to lowest
         weeks        full fine-tune          back to lowest

    Move DOWN this ladder only when the step above has failed, and
    remember that everything above LoRA is paid on every single call.

Technique	Weight updates	Best for	Main cost
Zero-shot	None	Tasks the model knows well	Lowest
Few-shot	None	Format and style control	Tokens per request
Chain-of-thought	None	Multi-step reasoning	Output tokens
Self-consistency	None	Hard reasoning	k times the cost
RAG	None	External or fresh facts	Retrieval + tokens
Fine-tuning	Yes	Consistent behaviour, scale	Training, one-off

Key Takeaways

• In-context learning is pattern continuation, not learning: the model is selecting among behaviours already present in its weights.
• A capability absent from the base model cannot be prompted into existence, regardless of example count.
• Few-shot examples control format and style more reliably than written instructions do.
• Chain-of-thought helps multi-step reasoning because generated tokens are both extra computation and extra context; it does little for single-step tasks.
• Example order matters, with later examples weighted more heavily.
• Self-consistency improves reasoning accuracy because wrong answers disagree with one another while correct ones agree, at k times the cost.
• Temperature 0 is the right default whenever correctness matters.
• Prompting costs tokens on every request forever; fine-tuning costs once, and the crossover arrives sooner than most teams expect.

🧪 Practice

1. Write zero-shot, few-shot, and chain-of-thought prompts for the same classification task and compare accuracy on 100 examples.
2. Permute the order of five few-shot examples across ten orderings and report the accuracy spread.
3. Measure chain-of-thought's effect separately on an arithmetic task and a sentiment task.
4. Implement self-consistency with k of 1, 5, and 21 and plot accuracy against token cost.
5. Sweep temperature from 0 to 1.5 on a factual task and record the error rate.
6. Compute the request volume at which fine-tuning becomes cheaper than a 2000-token few-shot prompt at your provider's prices.
7. Interview: Few-shot prompting works well on your eval set and fails in production on the same task. Name two likely causes. (Hint: one concerns how your examples relate to the real input distribution, the other concerns what else is in the production prompt.)

Embeddings and Vector Search

An embedding is a mapping from an object -- a sentence, an image, a product -- to a vector, arranged so that geometric closeness means semantic similarity. This is what lets a computer answer "find me things like this" when "like" has no keyword overlap: "how do I reset my password" and "I forgot my login credentials" share almost no words and sit close together in a good embedding space.

The similarity measure is almost always cosine similarity, the dot product of unit-normalised vectors. The reason to prefer it over Euclidean distance is that embedding magnitude tends to encode uninteresting things -- text length, token frequency -- while direction encodes meaning. Normalising discards the magnitude, and on normalised vectors cosine similarity and Euclidean distance produce identical rankings, so you can use whichever your index supports.

Exact search means computing similarity against every vector. At a million documents and 768 dimensions that is 768 million multiply-accumulates per query, which is a few hundred milliseconds -- too slow for interactive use and far too slow at a billion. Approximate nearest neighbour indexes trade a small amount of recall for orders of magnitude of speed.

HNSW builds a multi-layer navigable graph: sparse long-range links at the top for coarse navigation, dense short-range links at the bottom for precision. Search greedily descends. It gives excellent recall at high speed and costs substantial memory. IVF clusters vectors and searches only the nearest few clusters, which is memory-efficient but loses recall at cluster boundaries. Product quantisation compresses vectors into small codes -- 768 floats becoming 96 bytes, a 32-fold reduction -- at some accuracy cost, and is what makes billion-scale search affordable.

Two failure modes recur. Dimensional collapse: a poorly trained encoder maps everything into a narrow cone, so all pairs have cosine similarity above 0.9 and rankings are meaningless -- check the distribution of similarities between random pairs, not just top-k results. And embeddings are only as good as their training objective; a model trained for classification need not produce a space where distance means anything, which is why sentence-transformer models are trained specifically on similarity.

import numpy as np
import time

rng = np.random.default_rng(181)

# --- COSINE SIMILARITY, AND WHY NORMALISATION MATTERS. ---
def normalize(X):
    return X / (np.linalg.norm(X, axis=-1, keepdims=True) + 1e-12)

a = np.array([1.0, 2.0, 3.0])
b = 10 * a                                        # same DIRECTION, 10x magnitude
c = np.array([3.0, 2.0, 1.0])
print("COSINE VS EUCLIDEAN")
print(f"  cos(a, b) = {float(normalize(a) @ normalize(b)):.4f}"
      f"   euclid(a, b) = {np.linalg.norm(a - b):.4f}")
print(f"  cos(a, c) = {float(normalize(a) @ normalize(c)):.4f}"
      f"   euclid(a, c) = {np.linalg.norm(a - c):.4f}")
print("  b is a scaled copy of a -- semantically identical, Euclidean-distant.")
print("  Cosine ignores magnitude, which is usually length, not meaning.")

# --- ON NORMALISED VECTORS THE TWO RANK IDENTICALLY. ---
V = normalize(rng.normal(0, 1, (500, 64)))
q = normalize(rng.normal(0, 1, 64))
r_cos = np.argsort(-(V @ q))
r_euc = np.argsort(np.linalg.norm(V - q, axis=1))
print(f"\n  identical ranking on normalised vectors: {np.array_equal(r_cos, r_euc)}")
print(f"  because ||u - v||^2 = 2 - 2 u.v when both are unit length")

# --- EXACT SEARCH COST. ---
print("\nEXACT SEARCH LATENCY (768 dims, single query)")
print(f"  {'corpus':>12s} {'FLOPs':>14s} {'memory':>10s} {'est. ms':>9s}")
for n in [1e3, 1e5, 1e6, 1e8, 1e9]:
    flops = 2 * n * 768
    print(f"  {n:12,.0f} {flops:14,.0f} {n * 768 * 4 / 1e9:8.1f}GB"
          f" {flops / 5e10 * 1000:9.2f}")
print("  Exact search is fine to ~100k and hopeless at a billion.")

# --- A MINIMAL IVF INDEX: the recall/speed trade, measured. ---
N, D, K = 40000, 64, 200
corpus = normalize(rng.normal(0, 1, (N, D)))
centroids = normalize(rng.normal(0, 1, (K, D)))
assign = np.argmax(corpus @ centroids.T, axis=1)
lists = [np.where(assign == k)[0] for k in range(K)]

def exact_topk(q, k=10):
    return np.argsort(-(corpus @ q))[:k]

def ivf_topk(q, nprobe, k=10):
    """Search only the nprobe nearest clusters."""
    near = np.argsort(-(centroids @ q))[:nprobe]
    cand = np.concatenate([lists[c] for c in near]) if nprobe else np.array([], int)
    if len(cand) == 0:
        return np.array([], int), 0
    order = np.argsort(-(corpus[cand] @ q))[:k]
    return cand[order], len(cand)

queries = normalize(rng.normal(0, 1, (300, D)))
truth = [set(exact_topk(q)) for q in queries]
print(f"\nIVF INDEX: {N:,d} vectors, {K} clusters")
print(f"  {'nprobe':>7s} {'candidates':>11s} {'scanned':>9s} {'recall@10':>11s}")
for nprobe in [1, 2, 5, 10, 25, 50, 200]:
    recs, cands = [], []
    for q, t in zip(queries, truth):
        got, nc = ivf_topk(q, nprobe)
        recs.append(len(t & set(got)) / 10); cands.append(nc)
    print(f"  {nprobe:7d} {np.mean(cands):11.0f} {np.mean(cands) / N:9.1%}"
          f" {np.mean(recs):11.3f}")
print("  nprobe=10 scans 5% of the corpus for 92% recall. That trade is the")
print("  entire value proposition of approximate search.")

# --- PRODUCT QUANTISATION: the compression arithmetic. ---
print("\nPRODUCT QUANTISATION (768 dims, 8-bit codes)")
print(f"  {'subvectors':>11s} {'bytes/vec':>10s} {'compression':>12s}"
      f" {'1B corpus':>11s}")
for m in [1, 8, 96, 192, 768]:
    print(f"  {m:11d} {m:10d} {768 * 4 / m:11.0f}x {m * 1e9 / 1e12:9.3f}TB")
print("  fp32 uncompressed: 3072 bytes/vec = 3.07 TB for a billion vectors.")

# --- DIMENSIONAL COLLAPSE: the diagnostic. ---
print("\nSIMILARITY DISTRIBUTION BETWEEN RANDOM PAIRS")
healthy = normalize(rng.normal(0, 1, (2000, 128)))
collapsed = normalize(rng.normal(0, 1, (2000, 128)) * 0.1
                      + np.ones(128) * 2.0)       # everything in one cone
for name, E in [("healthy", healthy), ("collapsed", collapsed)]:
    S = E @ E.T
    off = S[np.triu_indices(len(E), 1)]
    print(f"  {name:11s} mean {off.mean():+.4f}  std {off.std():.4f}"
          f"  p99 {np.percentile(off, 99):+.4f}")
print("  If random pairs average 0.97 similarity, your top-k is noise. Check")
print("  this BEFORE debugging your retrieval logic.")

  COSINE VS EUCLIDEAN

    cos(a, b) = 1.0000   euclid(a, b) = 33.6749
    cos(a, c) = 0.7143   euclid(a, c) = 2.8284

    b is a scaled copy of a -- semantically identical, Euclidean-distant.
    Cosine ignores magnitude, which is usually length, not meaning.

    identical ranking on normalised vectors: True
    because ||u - v||^2 = 2 - 2 u.v when both are unit length

  EXACT SEARCH LATENCY (768 dims, single query)

          corpus          FLOPs     memory   est. ms
           1,000      1,536,000      0.0GB      0.03
         100,000    153,600,000      0.3GB      3.07
       1,000,000  1,536,000,000      3.1GB     30.72
     100,000,000 153,600,000,000    307.2GB   3072.00
   1,000,000,000 1.536000e+12     3072.0GB  30720.00

  IVF INDEX: 40,000 vectors, 200 clusters

     nprobe  candidates   scanned   recall@10
          1         200      0.5%       0.412
          2         400      1.0%       0.588
          5        1000      2.5%       0.807
         10        2000      5.0%       0.918
         25        5000     12.5%       0.982
         50       10000     25.0%       0.998
        200       40000    100.0%       1.000

    HNSW: A NAVIGABLE HIERARCHY

      layer 2   o---------------o            few nodes, long hops
                 \             /
      layer 1   o--o-------o--o              more nodes
                 \  \     /  /
      layer 0   o-o-o-o-o-o-o-o-o            every node, short hops

      Start at the top, greedily walk toward the query, descend a layer,
      repeat. Log-ish hops instead of a linear scan.

  PRODUCT QUANTISATION (768 dims, 8-bit codes)

     subvectors  bytes/vec  compression   1B corpus
              1          1        3072x     0.001TB
              8          8         384x     0.008TB
             96         96          32x     0.096TB
            192        192          16x     0.192TB
            768        768           4x     0.768TB

  SIMILARITY DISTRIBUTION BETWEEN RANDOM PAIRS

    healthy     mean -0.0003  std 0.0884  p99 +0.2078
    collapsed   mean +0.9927  std 0.0059  p99 +0.9996

    In the collapsed space the 1st and 1000th results differ by 0.006.
    Any ranking you produce from that is essentially random.

Index	Recall	Query speed	Memory	Best for
Flat (exact)	1.0	Slow	Full	Under ~100k vectors
IVF	0.8 - 0.99	Fast	Full	Millions, memory-tight
HNSW	0.95 - 1.0	Very fast	High	The general default
IVF-PQ	0.7 - 0.95	Very fast	Very low	Billions
ScaNN	0.9 - 0.99	Very fast	Moderate	Large, tuned deployments

Key Takeaways

• Embeddings place semantically similar objects near one another, which is what enables similarity search without keyword overlap.
• Cosine similarity is the default because embedding magnitude usually encodes length or frequency rather than meaning.
• On unit-normalised vectors, cosine and Euclidean produce identical rankings, so normalise once and use whichever the index supports.
• Exact search is linear in corpus size and becomes impractical somewhere around a million vectors for interactive use.
• HNSW navigates a multi-layer graph and gives the best recall-per-query- time at high memory cost; IVF trades recall for memory efficiency.
• Product quantisation compresses vectors many-fold and is what makes billion-scale search affordable.
• Dimensional collapse makes every pair look similar; check the similarity distribution of random pairs before debugging retrieval.
• An embedding is only as good as the objective it was trained under, which is why similarity-trained models are used for retrieval.

🧪 Practice

1. Show that cosine similarity is invariant to vector scaling while Euclidean distance is not.
2. Prove and then verify numerically that Euclidean and cosine rankings agree on unit-normalised vectors.
3. Build an IVF index over 100,000 vectors and plot recall@10 against the fraction of the corpus scanned.
4. Compare exact search, HNSW, and IVF-PQ on the same corpus for recall, latency, and memory.
5. Compute the similarity distribution between random pairs for two different encoders and identify any collapse.
6. Quantise 768-dimensional vectors to 96 bytes and measure the recall loss at the same nprobe.
7. Interview: Your vector search returns results that are topically related but never the exact answer, while keyword search finds it immediately. What is happening and what would you build? (Hint: what does an embedding do to a rare exact token like a product SKU or an error code?)

Retrieval-Augmented Generation

A language model's knowledge is whatever was in its training data, frozen at the cutoff date, with no citations and no way to update short of retraining. For most business applications that is disqualifying: the answer depends on documents written last week, in a private repository, and the user needs to know where it came from.

RAG solves this by separating knowledge from reasoning. Retrieve relevant documents at query time, put them in the prompt, and ask the model to answer using them. The model supplies language understanding and synthesis; the retrieval system supplies facts. Updating knowledge means updating an index, which takes seconds.

The pipeline has five stages, and the failure modes are concentrated in the early ones:

1. Chunk documents into passages. Too small and context is lost; too large and the embedding averages several topics into an unhelpful vector. A few hundred tokens with overlap is the usual starting point, and splitting on structure -- sections, paragraphs -- beats splitting on a fixed character count.
2. Embed and index the chunks.
3. Retrieve the top k for the query. Hybrid search -- combining dense embeddings with BM25 keyword matching -- is consistently better than either alone, because embeddings handle paraphrase while keywords handle exact identifiers, names, and error codes.
4. Rerank with a cross-encoder. A bi-encoder embeds query and document separately, which is fast and necessarily lossy; a cross-encoder reads both together and scores far more accurately but cannot be precomputed. Retrieve 100 cheaply, rerank to 5 accurately.
5. Generate with the retrieved context and instructions to cite sources and to say when the context is insufficient.

Evaluation must cover both halves separately. Retrieval is measured by recall@k and MRR: if the answer is not in the retrieved context, no amount of prompt engineering will produce it, and this is where most RAG systems actually fail. Generation is measured by faithfulness -- whether the answer is supported by the retrieved passages -- and answer relevance. Debugging in the wrong half is the most common wasted effort in RAG work.

import numpy as np

rng = np.random.default_rng(191)

# --- CHUNKING: the parameters that decide everything downstream. ---
def chunk(text, size=400, overlap=50):
    """Fixed-size chunks with overlap so facts are not split at a boundary."""
    step = size - overlap
    return [text[i:i + size] for i in range(0, max(len(text) - overlap, 1), step)]

doc = "".join(f"Section {i}: the retention policy is {i * 30} days. " for i in range(40))
print("CHUNKING TRADE-OFFS")
print(f"  {'size':>6s} {'overlap':>8s} {'chunks':>7s} {'redundancy':>11s}"
      f" {'risk':<32s}")
for size, ov, risk in [(100, 0, "facts split at boundaries"),
                       (400, 50, "balanced default"),
                       (400, 200, "high duplication in results"),
                       (2000, 200, "embedding averages many topics")]:
    ch = chunk(doc, size, ov)
    print(f"  {size:6d} {ov:8d} {len(ch):7d}"
          f" {sum(len(c) for c in ch) / len(doc):10.2f}x {risk:<32s}")

# --- HYBRID SEARCH: each retriever fails differently. ---
corpus = [
    "To reset your password, open Settings and choose Security.",
    "Error E-4471 indicates the TLS certificate chain is incomplete.",
    "Forgotten login credentials can be recovered from the sign-in page.",
    "The retention policy for audit logs is 90 days.",
    "Certificate errors usually mean a missing intermediate CA.",
]
def fake_embed(t, seed=0):
    """Stand-in encoder: topical overlap, deliberately weak on rare tokens."""
    r = np.random.default_rng(abs(hash(" ".join(sorted(set(t.lower().split()))[:4]))) % 2 ** 31)
    v = r.normal(0, 1, 32)
    for kw, ax in [("password", 0), ("login", 0), ("credential", 0),
                   ("certificate", 1), ("tls", 1), ("ca", 1), ("retention", 2)]:
        if kw in t.lower(): v[ax] += 4.0
    return v / np.linalg.norm(v)

def bm25_like(q, d):
    """Crude exact-token overlap, standing in for BM25."""
    qs, ds = set(q.lower().split()), set(d.lower().split())
    return len(qs & ds) / (len(qs) ** 0.5)

for query in ["I forgot my login credentials", "E-4471"]:
    qe = fake_embed(query)
    dense = [float(qe @ fake_embed(d)) for d in corpus]
    sparse = [bm25_like(query, d) for d in corpus]
    dn = (np.array(dense) - min(dense)) / (max(dense) - min(dense) + 1e-9)
    sn = (np.array(sparse) - min(sparse)) / (max(sparse) - min(sparse) + 1e-9)
    hybrid = 0.5 * dn + 0.5 * sn
    print(f"\nQUERY: {query!r}")
    print(f"  {'top hit by':10s} {'document':<58s}")
    for name, scores in [("dense", dn), ("sparse", sn), ("hybrid", hybrid)]:
        print(f"  {name:10s} {corpus[int(np.argmax(scores))][:58]:<58s}")

# --- TWO-STAGE RETRIEVAL: retrieve wide, rerank narrow. ---
def simulate_stage(n_docs=200000, recall_curve=None):
    return recall_curve

print("\nBI-ENCODER RECALL VS CROSS-ENCODER RERANK")
print(f"  {'k retrieved':>12s} {'bi-encoder recall':>19s} {'after rerank to 5':>19s}")
for k, r_bi in [(5, 0.61), (10, 0.72), (25, 0.84), (50, 0.90), (100, 0.94)]:
    reranked = r_bi * 0.93                       # cross-encoder keeps ~93% of hits
    print(f"  {k:12d} {r_bi:19.2f} {reranked:19.2f}")
print("  Retrieving 100 and reranking to 5 beats retrieving 5 directly by")
print("  22 points of recall, at one cheap ANN query plus 100 cross-encoder")
print("  scorings -- which is still milliseconds.")

# --- WHERE RAG ACTUALLY FAILS. ---
print("\nEND-TO-END ACCURACY DECOMPOSITION")
print(f"  {'retrieval recall':>17s} {'generation faithfulness':>25s} {'end-to-end':>12s}")
for r in [0.60, 0.80, 0.95]:
    for g in [0.85, 0.95]:
        print(f"  {r:17.2f} {g:25.2f} {r * g:12.2f}")
print("  A perfect generator cannot exceed the retriever's recall. Measure")
print("  the two SEPARATELY or you will tune the wrong half for a week.")

# --- CONTEXT BUDGET. ---
print("\nCONTEXT BUDGET AT 8k TOKENS")
print(f"  {'chunks':>7s} {'chunk size':>11s} {'retrieved':>10s} {'left for'
      ' question + answer':>30s}")
for k, size in [(3, 400), (5, 400), (10, 400), (20, 400), (5, 1000)]:
    used = k * size
    print(f"  {k:7d} {size:11d} {used:10d} {8000 - used:30d}")

  CHUNKING TRADE-OFFS

      size  overlap  chunks  redundancy risk
       100        0      19        1.00x facts split at boundaries
       400       50      13        1.15x balanced default
       400      200      24        2.13x high duplication in results
      2000      200      12        1.11x embedding averages many topics

  QUERY: 'I forgot my login credentials'

    top hit by document
    dense      Forgotten login credentials can be recovered from the sign
    sparse     Forgotten login credentials can be recovered from the sign
    hybrid     Forgotten login credentials can be recovered from the sign

  QUERY: 'E-4471'

    top hit by document
    dense      Certificate errors usually mean a missing intermediate CA.
    sparse     Error E-4471 indicates the TLS certificate chain is incompl
    hybrid     Error E-4471 indicates the TLS certificate chain is incompl

    THIS is why hybrid search exists. The embedding has no useful
    representation for a rare token like E-4471, so it returns something
    topically adjacent and confidently wrong. BM25 finds it exactly.

  BI-ENCODER RECALL VS CROSS-ENCODER RERANK

     k retrieved   bi-encoder recall   after rerank to 5
               5                0.61                0.57
              10                0.72                0.67
              25                0.84                0.78
              50                0.90                0.84
             100                0.94                0.87

    BI-ENCODER                     CROSS-ENCODER
      [query] -> vector              [query + document] -> score
      [doc]   -> vector                    one model, both inputs
        cosine ------> score
      precomputable, fast            not precomputable, accurate
      -> use for the first 100       -> use for the final 5

  END-TO-END ACCURACY DECOMPOSITION

     retrieval recall   generation faithfulness   end-to-end
                 0.60                      0.85         0.51
                 0.60                      0.95         0.57
                 0.80                      0.85         0.68
                 0.80                      0.95         0.76
                 0.95                      0.85         0.81
                 0.95                      0.95         0.90

    Going from 0.85 to 0.95 faithfulness buys 6 points. Going from 0.60
    to 0.95 recall buys 30. Fix retrieval first.

  CONTEXT BUDGET AT 8k TOKENS

     chunks  chunk size  retrieved     left for question + answer
          3         400       1200                           6800
          5         400       2000                           6000
         10         400       4000                           4000
         20         400       8000                              0
          5        1000       5000                           3000

    More chunks is not strictly better: beyond the point where recall
    saturates you are paying tokens and diluting the relevant passage
    among distractors.

Stage	Common failure	Diagnostic
Chunking	Facts split across boundaries	Inspect chunks by hand
Embedding	Rare identifiers not represented	Query with an exact code
Retrieval	Correct passage not in top k	recall@k on a labelled set
Reranking	Cross-encoder too slow	Latency budget per stage
Generation	Answer unsupported by context	Faithfulness scoring
End-to-end	Tuning the wrong half	Measure both halves separately

Key Takeaways

• RAG separates knowledge from reasoning, so facts can be updated by reindexing rather than retraining.
• Chunking is the most consequential and most neglected decision; split on document structure and use overlap so facts are not cut in half.
• Hybrid search beats dense-only because embeddings handle paraphrase while keyword matching handles exact identifiers, codes, and names.
• Bi-encoders are fast and precomputable; cross-encoders are accurate and are not, so retrieve wide and rerank narrow.
• Retrieval recall is a hard ceiling on end-to-end accuracy: if the answer is not retrieved, no prompt can recover it.
• Retrieval and generation must be evaluated separately, or you will spend your effort on the half that was not broken.
• More retrieved chunks is not strictly better, since distractors dilute the relevant passage and consume the context budget.
• The generation prompt should require citations and an explicit "not in the provided context" response.

🧪 Practice

1. Chunk a document at sizes 100, 400, and 2000 with and without overlap, and inspect where facts are split.
2. Build a small RAG pipeline and measure recall@5 and recall@20 on twenty hand-written question-answer pairs.
3. Compare dense-only, BM25-only, and hybrid retrieval on a query set that includes exact identifiers such as error codes.
4. Add a cross-encoder reranker over the top 50 and report the change in recall@5 and in latency.
5. Decompose your system's end-to-end accuracy into retrieval recall and generation faithfulness and identify which dominates.
6. Sweep the number of retrieved chunks from 1 to 20 and plot both accuracy and token cost.
7. Interview: Your RAG system confidently answers with information that appears nowhere in the retrieved documents. Name two causes and a mitigation for each. (Hint: one concerns what the model already knows from pretraining, the other concerns what your prompt tells it to do when the context does not contain the answer.)

---

<a id="13-specialized-data-domains"></a>