Math for Deep Learning — Handbook

A derivation-first tour of the mathematics behind modern deep learning: autodiff and backprop, attention and transformers, convolution, generative and sequence models, RL, and the frontier topics — each idea explained plainly, derived, illustrated, and coded.

98 topics · 24 chapters · companion to the Math for DL Concept Lab

Each topic is laid out the same way: The idea (the plain-language concept), Picture it (what the interactive in the lab shows), The math (the key equations), Derivation (where they come from), Code (a short runnable illustration), and In practice (how it is used). Use the search box on the left to jump to any topic.

Calculus & autodiff

The derivative as slope

the slope of the tangent

The idea

The derivative of a function at a point is the slope of its tangent there — the limit of the secant through two points as they collapse together. Drag the point along the curve and shrink h to watch the secant swing into the tangent.

Picture it

In the companion lab you can drag the point along the curve; shrink h to make the secant approach the tangent.

The math

f′(x) = lim_{h→0} (f(x+h)−f(x))/h. For f(x)=½x²: f′(x) = x (power rule).

Derivation

(½(x+h)² − ½x²)/h = (½(2xh+h²))/h = x + h/2 → x as h→0. The secant slope limits to the tangent slope.

Code

python

import numpy as np
# Numerical derivative (finite difference) vs the analytic one
f  = lambda x: x**3
df = lambda x: 3*x**2
x, h = 2.0, 1e-5
num = (f(x+h) - f(x-h)) / (2*h)        # central difference
print(num, df(x))                       # ~12.0, 12.0

In practice

Every gradient is a derivative — the entire training signal. Finite differences (this secant) are how you gradient-check a hand-written backward pass.

Gradient descent

follow the slope downhill

The idea

Training minimizes a loss by repeatedly stepping downhill: subtract a small multiple of the gradient from the current point. The learning rate sets the step size — too small crawls, too large overshoots or diverges. Drop a ball anywhere and step it down.

Picture it

In the companion lab you can click the curve to place the ball, then step. Adjust the learning rate.

The math

Update: x ← x − η·f′(x). ∇f points uphill, so −∇f points downhill; η is the step size.

Derivation

First-order Taylor: f(x−ηf′) ≈ f(x) − η·f′(x)². For small η the change is ≤ 0, so loss decreases — until η is too big and the approximation breaks.

Code

python

import numpy as np
# Gradient = vector of partial derivatives; points uphill
f = lambda w: w[0]**2 + 3*w[1]**2
def grad(w, h=1e-5):
    g = np.zeros_like(w)
    for i in range(len(w)):
        e = np.zeros_like(w); e[i] = h
        g[i] = (f(w+e) - f(w-e)) / (2*h)
    return g
print(grad(np.array([1.0, 1.0])))       # [2, 6] -> analytic [2w0, 6w1]

In practice

This is SGD, the workhorse of training. Too-large η diverges, too-small crawls; momentum and Adam adapt the step. Saddles and poor conditioning slow it — hence normalization and adaptive optimizers.

The chain rule

slopes multiply through a composition

The idea

Composing functions multiplies their slopes. If y = f(g(x)), then dy/dx = f′(g(x)) · g′(x): the local sensitivity of the outer function times that of the inner. This single rule is what makes backpropagation possible. Slide x and watch the factors multiply.

Picture it

In the companion lab you can slide x through y = sin(x)² = f(g(x)) with g=sin, f=square.

The math

If y = f(g(x)) then dy/dx = f′(g(x)) · g′(x). Local slopes multiply along the composition.

Derivation

Δy/Δx = (Δy/Δu)(Δu/Δx) with u=g(x); taking limits gives the product of derivatives. For sin²x: 2·sin x·cos x.

Code

python

import numpy as np
# Chain rule: dz/dx = dz/dy * dy/dx, composed through a pipeline
x = 1.5
y = np.sin(x);      dy_dx = np.cos(x)
z = y**2;           dz_dy = 2*y
print("dz/dx =", dz_dy * dy_dx)          # = 2 sin(x) cos(x)

In practice

Backpropagation is the chain rule applied right-to-left through the network. Each layer multiplies by its local Jacobian; many factors <1 give vanishing, >1 exploding gradients.

Backprop on a tiny graph

forward values, backward gradients

The idea

A neuron computes pred = w·x + b, and squared-error loss compares it to a target. The forward pass fills in values left to right; the backward pass sends gradients right to left by the chain rule. Move w and b, watch the loss and gradients, and take a step.

Picture it

In the companion lab you can set the weight and bias; step to descend the loss.

The math

pred = w·x + b, L = (pred − y)². dL/dw = 2(pred−y)·x, dL/db = 2(pred−y).

Derivation

Let e = pred−y. dL/dpred = 2e (square). Then chain back: dpred/dw = x, dpred/db = 1 ⇒ dL/dw = 2e·x, dL/db = 2e.

Code

python

import numpy as np
# Backprop through a tiny net: forward, then reverse-mode gradients
x, w1, w2 = 1.0, 0.8, -0.5
h = np.tanh(w1*x)                        # forward
y = w2*h
L = 0.5*(y - 1.0)**2
dy  = (y - 1.0)                          # backward
dw2 = dy*h
dh  = dy*w2
dw1 = dh*(1 - h**2)*x                    # tanh' = 1 - tanh^2
print(round(dw1,4), round(dw2,4))

In practice

Reverse-mode autodiff: one forward pass caches values, one backward pass produces every parameter’s gradient by reusing these local derivatives — the engine inside PyTorch/JAX.

Activations & their derivatives

shape and slope of a nonlinearity

The idea

Nonlinear activations are what let networks bend. Each has a shape and a derivative — and that derivative is what backprop multiplies by. Flat regions (where the derivative is near zero) starve learning: the vanishing-gradient problem. Switch activations and slide the input.

Picture it

In the companion lab you can pick an activation and slide the input marker.

The math

sigmoid σ(x)=1/(1+e⁻ˣ); tanh; ReLU=max(0,x). σ′ = σ(1−σ); tanh′ = 1−tanh²; ReLU′ = 1[x>0].

Derivation

σ′: differentiate (1+e⁻ˣ)⁻¹ → e⁻ˣ/(1+e⁻ˣ)² = σ(1−σ), peaking at ¼. Saturated regions (|x| large) have derivative ≈ 0.

Code

python

import numpy as np
# Activations and their derivatives (the gradient signal that flows back)
relu   = lambda z: np.maximum(0, z);       d_relu = lambda z: (z > 0).astype(float)
sig    = lambda z: 1/(1+np.exp(-z));       d_sig  = lambda z: sig(z)*(1-sig(z))
z = np.array([-2., 0., 2.])
print("relu'", d_relu(z), "  sigmoid'", np.round(d_sig(z), 3))

In practice

Nonlinearity is what makes depth expressive. Saturating derivatives cause vanishing gradients; ReLU keeps a derivative of 1 on the positive side (but can "die" at 0) — GELU/Swish smooth this.

Loss & probability

Softmax & cross-entropy

logits → probabilities → loss

The idea

A classifier outputs raw scores (logits); softmax turns them into a probability distribution, and cross-entropy scores how surprised the model is by the true label. Their famous pairing exists because the gradient collapses to something beautifully simple. Tune the logits and the true class.

Picture it

In the companion lab you can set the logits; pick the true class.

The math

pᵢ = e^{zᵢ} / Σⱼ e^{zⱼ} (softmax) CE = −log pₜ (t = true class) ∂CE/∂zᵢ = pᵢ − yᵢ (y = one-hot)

Derivation

∂pⱼ/∂zᵢ = pᵢ(δᵢⱼ − pⱼ). With CE = −log pₜ, ∂CE/∂zᵢ = −(δᵢₜ − pᵢ) = pᵢ − yᵢ. The exp in softmax cancels the log in CE — hence the clean result.

Code

python

import numpy as np
# Softmax: logits -> probability distribution (stable version)
def softmax(z):
    z = z - z.max()                      # subtract max for stability
    e = np.exp(z)
    return e / e.sum()
print(np.round(softmax(np.array([2.0, 1.0, 0.1])), 3))   # sums to 1

In practice

This is the standard classification head. Because the gradient is just predicted − target, backprop starts from an error signal with no awkward factors — numerically stable (with the log-sum-exp trick) and cheap.

Maximum likelihood → your loss

why MSE and cross-entropy are losses

The idea

Where do loss functions come from? Maximum likelihood: pick the parameters that make the observed data most probable. Maximizing the product of probabilities is the same as minimizing the summed negative log — and for Gaussian noise that is exactly squared error. Slide μ to find the peak.

Picture it

In the companion lab you can slide the model mean μ to fit the data points (σ fixed). The negative log-likelihood is minimized at the sample mean.

The math

L(μ) = ∏ᵢ N(xᵢ | μ, σ²) ℓ(μ) = log L = −(1/2σ²)Σ(xᵢ−μ)² + const argmaxᵥ ℓ = argminᵥ Σ(xᵢ−μ)² (= MSE)

Derivation

Set dℓ/dμ = (1/σ²)Σ(xᵢ−μ) = 0 ⇒ μ* = (1/n)Σ xᵢ = sample mean. Gaussian likelihood → squared-error loss falls right out.

Code

python

import numpy as np
# Max likelihood = minimize negative log-likelihood; for a Gaussian -> MSE
data = np.random.normal(3.0, 1.0, 1000)
# NLL of mu given unit variance is sum (x-mu)^2 / 2  -> minimized at the mean
mu_hat = data.mean()
print("MLE mean:", round(mu_hat, 3))

In practice

"Loss = negative log-likelihood." Gaussian noise → MSE; a categorical/Bernoulli model → cross-entropy. Adding a parameter prior turns MLE into MAP — i.e. weight regularization (L2 = Gaussian prior).

Logistic regression & BCE

a probability and a boundary

The idea

For a yes/no label, squash a linear score through the sigmoid to get a probability, then score it with binary cross-entropy. Remarkably the gradient is again predicted − target. Move the weight and bias to slide and steepen the decision boundary between the two classes.

Picture it

In the companion lab you can set weight and bias; the boundary sits where σ = 0.5.

The math

p = σ(z) = 1/(1+e^{−z}), z = wx + b BCE = −[y·log p + (1−y)·log(1−p)] ∂BCE/∂z = p − y

Derivation

σ′(z) = σ(z)(1−σ(z)). Differentiating BCE and using this identity, the (1−p) and p factors cancel, leaving ∂BCE/∂z = p − y — the same clean signal as softmax.

Code

python

import numpy as np
# Sigmoid maps a logit to P(y=1); its inverse is the log-odds (logit)
sigmoid = lambda z: 1/(1+np.exp(-z))
p = sigmoid(np.array([-2., 0., 2.]))
logit = np.log(p/(1-p))                  # recovers the input
print(np.round(p,3), np.round(logit,3))

In practice

The sigmoid output unit for binary tasks and multi-label heads. The boundary z = 0 is linear in the features; deeper nets first learn nonlinear features, then apply this same logistic layer on top.

Entropy, cross-entropy & KL

CE = entropy + KL gap

The idea

Three quantities, one identity. Entropy is the irreducible uncertainty in the true distribution p; cross-entropy is the cost of coding p’s outcomes with model q; and their gap is the KL divergence. Since CE = H(p) + KL(p∥q), minimizing cross-entropy is minimizing KL to the data. Drag q toward p.

Picture it

In the companion lab you can adjust the model q over 3 outcomes (true p is fixed). KL hits 0 when q = p.

The math

H(p) = −Σ pᵢ log pᵢ CE(p,q) = −Σ pᵢ log qᵢ KL(p∥q) = Σ pᵢ log(pᵢ/qᵢ) = CE − H

Derivation

CE − H = −Σ p log q + Σ p log p = Σ p log(p/q) = KL ≥ 0 (Gibbs), with equality iff q = p. So CE = H(p) + KL(p∥q), and H(p) is fixed w.r.t. q.

Code

python

import numpy as np
# Cross-entropy = entropy(p) + KL(p||q); training minimizes it
def cross_entropy(p, q): return -(p*np.log(q)).sum()
def kl(p, q):           return (p*np.log(p/q)).sum()
p = np.array([1.,0.,0.])                 # one-hot label
q = np.array([0.7,0.2,0.1])              # prediction
print("CE =", round(cross_entropy(p, np.clip(q,1e-9,1)), 3))

In practice

Training a classifier minimizes CE = (constant) + KL to the data distribution. KL also drives knowledge distillation (match a teacher), the VAE ELBO, and the KL penalty that keeps RLHF policies near the reference model.

MSE, MAE, Huber — gradients

the loss is its gradient

The idea

The loss you pick decides how the model reacts to errors, because training only ever sees the gradient. Squared error’s gradient grows with the error (outliers dominate); absolute error’s is constant (robust but kinked); Huber blends both. Drag the error and compare the slopes.

Picture it

In the companion lab you can drag the residual e = prediction − target; set Huber’s δ.

The math

MSE = e², grad = 2e MAE = |e|, grad = sign(e) Huberδ: ½e² if |e|≤δ, else δ(|e|−½δ); grad = e if |e|≤δ else δ·sign(e)

Derivation

MSE’s gradient 2e is unbounded — a single outlier sends a huge update. MAE’s is bounded at ±1 — robust, but nonsmooth at 0. Huber is quadratic (smooth) near 0 and linear (bounded) far out: best of both.

Code

python

import numpy as np
# Picking a loss = picking a noise model. MSE<->Gaussian, CE<->categorical
y, yhat = np.array([1.,0.,1.]), np.array([0.9,0.2,0.6])
mse = np.mean((y - yhat)**2)
bce = -np.mean(y*np.log(yhat) + (1-y)*np.log(1-yhat))
print(f"MSE={mse:.3f}  BCE={bce:.3f}")

In practice

MSE for clean regression; Huber / smooth-L1 for robustness to outliers (object detection box loss, RL value targets). Bounded gradients also tame exploding updates — the same motivation as gradient clipping.

Training dynamics

Weight initialization

keeping signal alive through depth

The idea

Before any learning happens, the scale of the random weights decides whether a signal survives a deep stack of layers. Too large and activations explode; too small and they vanish to zero. The fix is to set the weight variance from the fan-in. Add layers and dial the gain to find the stable knife-edge.

Picture it

In the companion lab you can add layers, change the gain, switch activation. Watch the per-layer activation scale.

The math

Var(a_l) = fan_in · Var(W) · Var(a_{l−1}) · c_act. Pick Var(W) = g²/fan_in ⇒ Var(a_l) = g²·c_act·Var(a_{l−1}).

Derivation

Stable layers need g²·c_act = 1. ReLU zeroes half the inputs (c_act = ½) ⇒ g = √2 (He). Linear/tanh (c_act ≈ 1) ⇒ g = 1 (Xavier/Glorot).

Code

python

import numpy as np
# Initialization sets the scale of activations; Xavier/He keep variance ~1
fan_in, fan_out = 256, 256
xavier = np.random.randn(fan_in, fan_out) * np.sqrt(1/fan_in)
he     = np.random.randn(fan_in, fan_out) * np.sqrt(2/fan_in)   # for ReLU
x = np.random.randn(1000, fan_in)
print("act std (He):", round((x @ he).std(), 2))   # ~1, not exploding

In practice

Correct init is what made training very deep nets possible without normalization. Get it wrong and you see immediate NaNs (explode) or a dead network that never learns (vanish).

Vanishing & exploding gradients

the product that makes or breaks depth

The idea

Backprop sends the error backward by multiplying through every layer. Multiply many numbers below one and the gradient decays to nothing in early layers; many above one and it blows up. This is why sigmoid-deep nets barely trained — and why ReLU and residual connections changed everything. Deepen the net and switch the activation.

Picture it

In the companion lab you can set depth and weight scale; switch activation to compare gradient flow.

The math

∂L/∂h_l = ∏_{j>l} f′(z_j)·W_j. Magnitude ≈ (|w|·|f′|)^{depth−l}.

Derivation

Each backward hop multiplies by the layer Jacobian. Sigmoid has f′ ≤ ¼, so even with w=1 the factor ≤ 0.25 — raised to a high power → 0. ReLU has f′ = 1 on active units, so the factor stays near |w|.

Code

python

import numpy as np
# Vanishing gradients: stacking saturating activations shrinks the signal
g = 1.0
for layer in range(20):
    s = 1/(1+np.exp(-np.random.randn()))     # a sigmoid output ~0.5
    g *= s*(1-s)                              # multiply by sigmoid' (<=0.25)
print("gradient after 20 layers:", f"{g:.1e}")   # tiny -> vanished

In practice

The core reason for ReLU/GELU, residual connections (which add an identity path so the factor includes +1), normalization, careful init, gradient clipping, and LSTM/GRU gates in RNNs.

Normalization (BatchNorm / LayerNorm)

recenter, rescale, repeat

The idea

Normalization grabs a batch of activations, recenters them to mean 0 and variance 1, then lets the network rescale with two learned parameters γ and β. This keeps every layer’s inputs in a sane range no matter what the previous layers do. Shift and spread the raw activations and watch them snap back.

Picture it

In the companion lab you can mess up the raw activations; normalization fixes the spread, then γ,β reshape it on purpose.

The math

x̂ = (x − μ_B) / √(σ²_B + ε) y = γ·x̂ + β BN: μ,σ over the batch · LN: over the features.

Derivation

By construction x̂ has mean 0, variance 1. Then y has mean β and variance γ² — so the net can recover any scale/shift it wants (even undo the normalization), but starts from a stable point.

Code

python

import numpy as np
# Batch/Layer norm: standardize activations to mean 0, var 1, then scale+shift
def layernorm(x, gamma=1.0, beta=0.0, eps=1e-5):
    mu = x.mean(-1, keepdims=True); var = x.var(-1, keepdims=True)
    return gamma*(x - mu)/np.sqrt(var + eps) + beta
h = np.random.randn(4, 8)*5 + 3
print("after norm:", np.round(layernorm(h).mean(), 3), np.round(layernorm(h).std(), 3))

In practice

Faster, more stable training and tolerance of higher learning rates; smooths the loss landscape and reduces sensitivity to init. BatchNorm in CNNs, LayerNorm in Transformers (per-token, batch-independent).

SGD vs Momentum vs Adam

three ways down the same valley

The idea

On a stretched ravine, plain SGD bounces across the walls and inches along the floor. Momentum builds velocity down the long axis; Adam rescales each coordinate by its own gradient history so both axes move at a sensible pace. Run each optimizer from the same start and compare the paths.

Picture it

In the companion lab you can run each optimizer (same start, ill-conditioned bowl) and compare.

The math

SGD: θ ← θ − ηg Momentum: v ← βv − ηg; θ ← θ + v Adam: m,v = EMA(g), EMA(g²); θ ← θ − η·m̂/(√v̂ + ε)

Derivation

Momentum accumulates consistent gradients (the floor direction) while cancelling the oscillating ones — an exponential average of past gradients. Adam divides by √v̂, so a steep coordinate gets a small step and a flat one a large step: per-parameter learning rates.

Code

python

import numpy as np
# Adam: per-parameter adaptive step from 1st & 2nd moment estimates
w = np.array([5.0]); m = v = 0; b1,b2,lr,eps = 0.9,0.999,0.1,1e-8
grad = lambda w: 2*w
for t in range(1, 51):
    g = grad(w)
    m = b1*m + (1-b1)*g;  v = b2*v + (1-b2)*g**2
    mh = m/(1-b1**t);     vh = v/(1-b2**t)          # bias correction
    w -= lr*mh/(np.sqrt(vh)+eps)
print(np.round(w, 3))    # -> 0

In practice

Adam (and AdamW) is the default for Transformers and most modern training — robust to scale and quick to get going. SGD+momentum often generalizes better in vision and is preferred when tuned. Same idea fights the conditioning problem from the optimization group.

Linear algebra of layers

A layer as a matrix multiply

each output is a dot product

The idea

A dense layer is just h = Wx + b. Read it row by row: each output is the dot product of one weight row with the input — a little feature detector that fires when the input matches its row. Change the input and pick a row to see exactly which numbers get multiplied and summed.

Picture it

In the companion lab you can set the input; highlight a row to see its dot product.

The math

h = Wx + b, with W ∈ ℝ^{out×in}. h_i = Σ_j W_{ij} x_j + b_i = (row_i of W)·x + b_i.

Derivation

Matrix–vector product stacks dot products: the i-th entry pairs row i of W with x. Equivalently h = x₁·col₁ + x₂·col₂ + b — a linear combination of W’s columns.

Code

python

import numpy as np
# A dense layer is just  y = x W + b  (a batched matrix multiply)
x = np.random.randn(32, 128)             # batch of 32 vectors
W = np.random.randn(128, 64); b = np.random.randn(64)
y = x @ W + b
print(y.shape)                            # (32, 64)

In practice

The fundamental building block (nn.Linear / Dense). Rows of W are learned templates; b sets each detector’s threshold; an activation then gates the result. Params = out×in + out.

The Jacobian & backward pass

gradients flow back through Wᵀ

The idea

Forward, a linear layer multiplies by W. Backward, the gradient multiplies by Wᵀ. That transpose is the whole mechanism of backprop: the layer’s Jacobian sends gradients from outputs back to inputs. Drag the incoming gradient and watch the outgoing gradient it produces.

Picture it

In the companion lab you can drag the upstream gradient g_y; the downstream gradient is g_x = Wᵀ g_y.

The math

Forward: y = Wx ⇒ Jacobian ∂y/∂x = W. Backward: g_x = Jᵀ g_y = Wᵀ g_y. Activation a=φ(z): J = diag(φ′(z)).

Derivation

By the chain rule ∂L/∂x_j = Σ_i (∂L/∂y_i)(∂y_i/∂x_j) = Σ_i g_{y,i} W_{ij} = (Wᵀ g_y)_j. Reverse-mode never builds J — it computes this vector–Jacobian product directly.

Code

python

import numpy as np
# The Jacobian of a layer: how each output changes with each input
W = np.array([[2., 1.], [0., 3.]])
f = lambda x: W @ x                       # linear layer -> Jacobian is W
x = np.array([1., 1.]); h = 1e-5
J = np.array([[ (f(x+np.eye(2)[j]*h)[i]-f(x-np.eye(2)[j]*h)[i])/(2*h)
                for j in range(2)] for i in range(2)])
print(np.round(J, 2))                     # == W

In practice

Every backward() is a chain of VJPs: multiply the incoming gradient by each layer’s transposed Jacobian. For elementwise activations that’s a cheap diagonal scale by φ′; this is what autodiff frameworks implement per op.

Why nonlinearity?

depth needs a bend

The idea

Stack two linear layers and you get… one linear layer: W₂(W₁x) = (W₂W₁)x. No amount of depth helps until you put a nonlinearity between them. Toggle the activation to watch a flat grid either stay a boring parallelogram or fold into something a linear map could never make.

Picture it

In the companion lab you can toggle the activation between the two layers; randomize the weights.

The math

Linear∘linear: W₂(W₁x) = (W₂W₁)x — still one matrix. With φ: W₂·φ(W₁x) cannot be collapsed.

Derivation

Matrix products are associative, so any chain of linear maps equals a single matrix W_eff. The activation φ breaks associativity — it acts per-coordinate between the matmuls — so the composite becomes genuinely nonlinear.

Code

python

import numpy as np
# Stacking linear layers without nonlinearity collapses to one matrix
W1 = np.random.randn(4, 4); W2 = np.random.randn(4, 4)
x  = np.random.randn(4)
print(np.allclose(W2 @ (W1 @ x), (W2 @ W1) @ x))   # True -> need nonlinearity

In practice

This is why every layer has an activation. With nonlinearities a deep net is a universal approximator (it can fold space into arbitrary decision regions); without them, a 100-layer MLP is just linear regression.

Forward + backward: a live MLP

the whole training loop in miniature

The idea

Here is a complete tiny network learning by backprop. Two inputs feed a hidden ReLU layer, then a linear output, scored against a target. Step runs one forward pass (values flow right) and one backward pass (gradients flow left), then nudges every weight downhill. Watch the loss fall as it learns.

Picture it

In the companion lab you can step the network and watch loss drop. Change the target or learning rate.

The math

Forward: z₁ = W₁x + b₁, h = ReLU(z₁) ŷ = W₂h + b₂, L = (ŷ − y)² Backward: ∂L/∂ŷ = 2(ŷ−y); ∂L/∂W₂ = (∂L/∂ŷ)·hᵀ ∂L/∂h = W₂ᵀ(∂L/∂ŷ); ∂L/∂z₁ = ∂L/∂h ⊙ ReLU′(z₁) ∂L/∂W₁ = (∂L/∂z₁)·xᵀ; then W ← W − η·∂L/∂W

Code

python

import numpy as np
# Full forward pass of a 2-layer MLP
def mlp(x, W1, b1, W2, b2):
    h = np.maximum(0, x @ W1 + b1)        # ReLU hidden layer
    return h @ W2 + b2                     # linear output
x  = np.random.randn(8, 16)
W1, b1 = np.random.randn(16,32), np.zeros(32)
W2, b2 = np.random.randn(32,3),  np.zeros(3)
print(mlp(x, W1, b1, W2, b2).shape)       # (8, 3)

In practice

This is literally what every training step does, scaled up: one forward pass caches activations, one backward pass applies the chain rule to get all gradients, the optimizer updates. Everything else — depth, attention, batches — is more of this.

Sequence & attention

Attention weights

a soft, content-based lookup

The idea

Attention lets a query pull information from whichever items resemble it most. Score the query against each key by dot product, soften the scores into weights with softmax, then blend the values by those weights. Drag the query toward a key and watch nearly all the weight rush to it.

Picture it

In the companion lab you can drag the query. Each key’s weight is how much the output leans toward it.

The math

scoreⱼ = q·kⱼ / √d w = softmax(scores) output = Σⱼ wⱼ vⱼ (here vⱼ = kⱼ)

Derivation

The dot product measures alignment; softmax turns the scores into positive weights summing to 1, so the output is a convex combination of the values — a differentiable, content-addressed average that the network can learn to steer.

Code

python

import numpy as np
# Attention scores = how much each query attends to each key
Q = np.random.randn(3, 8)                 # 3 queries, dim 8
K = np.random.randn(5, 8)                 # 5 keys
scores = Q @ K.T                          # (3, 5) raw compatibility
print(scores.shape)

In practice

The atom of the Transformer. Stacking it lets every token gather context from any other in one step (unlike RNNs), which is why attention scales and parallelizes so well.

Self-attention: Q, K, V

every token attends to every token

The idea

In self-attention every token plays three roles: it emits a query (what am I looking for), a key (what do I offer), and a value (what I’ll pass on). The full attention matrix scores every query against every key; each row, softmaxed, says how one token distributes its attention across all tokens. Pick a token to see its row.

Picture it

In the companion lab you can select a query token to highlight its attention row; reshape the learned projections.

The math

Q = XW_Q, K = XW_K, V = XW_V. A = softmax(QKᵀ / √d) (row-wise) output = A·V

Derivation

Entry A_{ij} = softmax over j of (q_i·k_j)/√d: token i’s affinity for token j. Multiplying by V replaces each token with a weighted mixture of all tokens’ values — contextualization in one matmul.

Code

python

import numpy as np
def softmax(z): e=np.exp(z-z.max(-1,keepdims=True)); return e/e.sum(-1,keepdims=True)
# Self-attention: each token mixes in the values of relevant tokens
X = np.random.randn(4, 8)                 # 4 tokens
Wq,Wk,Wv = [np.random.randn(8,8) for _ in range(3)]
Q,K,V = X@Wq, X@Wk, X@Wv
A = softmax(Q @ K.T / np.sqrt(8))         # attention weights (rows sum to 1)
out = A @ V
print(out.shape, np.round(A.sum(1), 2))

In practice

The heart of every Transformer (BERT, GPT, ViT). A causal mask zeroes the upper triangle for autoregressive generation; the cost is O(N²) in sequence length, which is what long-context research keeps attacking.

Why divide by √d

keeping softmax from saturating

The idea

Dot products of high-dimensional vectors grow with the dimension, so raw attention scores get huge — softmax then collapses to a near one-hot spike and its gradient vanishes. Dividing by √d rescales the scores back to unit variance, keeping attention soft and trainable. Crank up d and compare the two rows.

Picture it

In the companion lab you can increase the key/query dimension d and compare unscaled vs scaled weights.

The math

If q,k have unit-variance entries, Var(q·k) = d. Dividing by √d ⇒ Var = 1, regardless of d.

Derivation

q·k = Σᵢ qᵢkᵢ is a sum of d zero-mean unit-variance terms, so its variance is d and its std is √d. Scaling by 1/√d standardizes it, so softmax sees comparable inputs at any depth.

Code

python

import numpy as np
# Why divide by sqrt(d): without it, dot products grow and softmax saturates
d = 64
q = np.random.randn(d); k = np.random.randn(d)
print("var of q.k :", round((q@k).var() if False else float(np.var([np.random.randn(d)@np.random.randn(d) for _ in range(1000)])),1))
# ~d, so scores ~sqrt(d); dividing by sqrt(d) keeps softmax in a useful range

In practice

The “scaled” in scaled-dot-product attention. Multi-head: split d into h heads of size d/h, attend in parallel, concatenate — each head can specialize (syntax, coreference, position) while keeping per-head dot products well-scaled.

Positional encoding

giving order to a set

The idea

Attention is permutation-invariant — shuffle the tokens and it can’t tell, because it has no built-in notion of order. Positional encodings fix this by adding a unique, smoothly varying signature to each position. The classic choice is a bank of sines and cosines at geometrically spaced frequencies. Slide the position to see its fingerprint.

Picture it

In the companion lab you can move the highlighted position and change the model width.

The math

PE(pos, 2i) = sin( pos / 10000^{2i/d} ) PE(pos, 2i+1) = cos( pos / 10000^{2i/d} ) Added to the token embedding before layer 1.

Derivation

Low dimensions oscillate fast, high dimensions slowly — like binary place-values in continuous form, giving every position a distinct vector. Because shifts act as rotations, PE(pos+k) is a linear function of PE(pos), so the model can learn relative offsets.

Code

python

import numpy as np
# Sinusoidal positional encodings give the model a sense of order
def pos_enc(seq_len, d):
    pos = np.arange(seq_len)[:, None]; i = np.arange(d)[None, :]
    ang = pos / np.power(10000, (2*(i//2))/d)
    pe = np.where(i % 2 == 0, np.sin(ang), np.cos(ang))
    return pe
print(pos_enc(10, 16).shape)              # (10, 16), added to token embeddings

In practice

Sinusoidal PE generalizes to unseen lengths; learned absolute embeddings are also common. Modern LLMs favor RoPE (rotary) and ALiBi, which inject relative position directly into the attention scores.

Convolution & vision

Convolution

a sliding weighted patch

The idea

A convolution slides a small kernel across an image; at each spot it multiplies the overlapping pixels by the kernel and sums them into one output value. A 3×3 kernel becomes an edge detector, a blur, or a sharpener depending on its nine numbers. Drag the window over the image and switch kernels to watch the feature map change.

Picture it

In the companion lab you can drag the 3×3 window over the image; switch the kernel.

The math

(I * K)[i,j] = Σ_{u,v} I[i+u, j+v] · K[u,v]. (Deep nets use cross-correlation — no kernel flip.)

Derivation

Each output pixel is a dot product of the kernel with the local patch — a tiny linear filter. The same kernel is reused at every location, so a feature (an edge, a corner) is detected wherever it appears.

Code

python

import numpy as np
# 2-D convolution: slide a kernel, take dot products (here: edge detector)
img = np.eye(5)
kernel = np.array([[-1,-1,-1],[0,0,0],[1,1,1]])    # horizontal edges
out = np.zeros((3,3))
for i in range(3):
    for j in range(3):
        out[i,j] = (img[i:i+3, j:j+3] * kernel).sum()
print(out)

In practice

The convolutional layer. Weight sharing + local connectivity give translation equivariance and far fewer parameters than a dense layer; stacks of conv layers build up from edges to textures to objects.

Stride, padding & output size

controlling spatial size

The idea

Three knobs decide how big a conv layer’s output is and how fast it shrinks: the kernel size, the stride (how far the window jumps), and the padding (zeros added around the border). Dial them and watch the window placements and the resulting output length update against the formula.

Picture it

In the companion lab you can adjust the four knobs; the windows and output length follow the formula.

The math

O = ⌊(N + 2p − k)/s⌋ + 1. “Same” padding: p = (k−1)/2 keeps O = N when s = 1.

Derivation

The first window starts at 0 and the last fits when its right edge ≤ N+2p−1; stepping by s gives (N+2p−k)/s + 1 start positions, floored. Larger stride downsamples; padding restores the size the border would otherwise lose.

Code

python

import numpy as np
# Output size = (W - K + 2P)/S + 1; stride/padding control downsampling
def out_size(W, K, P, S): return (W - K + 2*P)//S + 1
print(out_size(32, 3, 1, 1))      # 32 (same)
print(out_size(32, 3, 0, 2))      # 15 (downsampled)

In practice

Stride-2 convs (or pooling) halve resolution as channel count grows; “same” padding keeps shapes tidy through deep stacks. Stacking convs grows the receptive field each output pixel can see.

Pooling

summarize and shrink

The idea

Pooling shrinks a feature map by summarizing each little region with a single number — the maximum (keep the strongest response) or the average. It has no learnable weights; it just downsamples and grants a bit of translation invariance. Toggle between max and average and see which input each output came from.

Picture it

In the companion lab you can switch between max and average pooling (2×2, stride 2).

The math

Max: y = max over the window. Avg: y = mean over the window. 2×2 stride-2 → halves each spatial dimension.

Derivation

Pooling is a fixed (non-learned) reduction. Max-pool’s gradient flows only to the winning input; avg-pool spreads it evenly. Either way it discards precise location, buying small-shift invariance.

Code

python

import numpy as np
# Max pooling: downsample by keeping the strongest activation per window
x = np.arange(16).reshape(4,4).astype(float)
pooled = x.reshape(2,2,2,2).max(axis=(1,3))    # 2x2 max pool
print(pooled)

In practice

Classic CNNs alternate conv and pool to reduce resolution and computation. Many modern architectures replace pooling with strided convolutions or global average pooling before the classifier head.

Parameter sharing & feature maps

why CNNs fit in memory

The idea

Why convolutions instead of dense layers on images? A dense layer wires every pixel to every output — the parameter count explodes. A conv layer reuses one tiny kernel across the whole image, so its size depends only on the kernel and channel counts, not the image size. Compare the two and watch the gap.

Picture it

In the companion lab you can grow the image and the channels; compare conv vs dense parameters.

The math

Conv params = k²·C_in·C_out + C_out. Dense params ≈ (H·W·C_in)·(H·W·C_out).

Derivation

The conv count has no H,W term — the same k×k weights are shared across all H·W positions. The dense count multiplies the full flattened sizes, so it grows like the fourth power of the image side.

Code

python

# Parameter count of a conv layer vs a dense layer on the same input
# Conv: K*K*C_in*C_out (shared across all spatial positions)
K, Cin, Cout = 3, 64, 128
conv = K*K*Cin*Cout
dense = (32*32*Cin)*(32*32*Cout)               # if you used a dense layer
print(f"conv params={conv:,}  dense params={dense:,}")  # weight sharing wins

In practice

Weight sharing is what lets CNNs scale to real images and generalize across positions. Each of the C_out kernels produces one feature map (channel); stacking them builds a hierarchy of detectors.

Generative models

Autoencoder

compress, then reconstruct

The idea

An autoencoder squeezes input through a narrow bottleneck and reconstructs it. The squeeze forces it to learn the data’s low-dimensional structure — the manifold the points actually live on. Add noise and watch the 1-D bottleneck project messy points back onto the clean curve (denoising).

Picture it

In the companion lab you can add noise; the bottleneck reconstruction snaps points back to the learned curve.

The math

z = enc(x), x̂ = dec(z), minimize ‖x − x̂‖². If dim(z) < dim(x), the net must compress.

Derivation

With a bottleneck the best reconstruction lands on a learned dim(z)-manifold; a linear AE recovers PCA, a nonlinear one a curved manifold. Off-manifold noise is discarded — so the AE denoises by projection.

Code

python

import numpy as np
# Autoencoder: compress to a bottleneck z, then reconstruct
def encode(x, We): return np.maximum(0, x @ We)
def decode(z, Wd): return z @ Wd
x  = np.random.randn(10, 64)
We = np.random.randn(64, 8); Wd = np.random.randn(8, 64)     # bottleneck=8
recon = decode(encode(x, We), Wd)
print("reconstruction MSE:", round(np.mean((x-recon)**2), 3))

In practice

Representation learning, dimensionality reduction, anomaly detection (high reconstruction error = outlier), and denoising AEs. The latent code z is a compact feature you can reuse downstream.

Variational autoencoder

a smooth, samplable latent

The idea

A VAE makes the latent space probabilistic: the encoder outputs a Gaussian q(z|x), and a KL term pulls every such Gaussian toward the standard-normal prior. That keeps the latent space smooth and samplable, so you can draw new z from the prior and decode novel data. Tune q and watch the KL.

Picture it

In the companion lab you can set the encoder’s q(z|x). KL → 0 as q matches the prior N(0,1).

The math

ELBO = E_q[log p(x|z)] − KL(q(z|x)∥p(z)). Reparameterize: z = μ + σ·ε, ε~N(0,1). KL = ½(μ² + σ² − 1 − lnσ²).

Derivation

Maximizing the data likelihood is intractable, so we maximize its ELBO lower bound. The reparameterization trick writes z as a deterministic function of (μ,σ,ε), letting gradients flow through the sampling step.

Code

python

import numpy as np
# VAE: encoder outputs mu, logvar; sample z with the reparam trick
mu, logvar = np.zeros(8), np.zeros(8)
eps = np.random.randn(8)
z = mu + np.exp(0.5*logvar) * eps                  # differentiable sampling
# KL term pulls q(z|x) toward N(0,I)
kl = -0.5 * np.sum(1 + logvar - mu**2 - np.exp(logvar))
print("z:", np.round(z[:3],2), " KL:", round(kl,3))

In practice

VAEs generate by sampling z~N(0,1) and decoding. The KL regularizes the latent so it stays continuous (nearby z → similar outputs); β-VAE up-weights KL to encourage disentangled factors.

GAN: generator vs discriminator

a forger and a detective

The idea

A GAN pits two networks against each other: the generator invents fake samples, the discriminator tries to tell real from fake. They reach equilibrium only when the fakes are indistinguishable — at which point the discriminator is reduced to a coin flip everywhere. Move the generator toward the real data and watch D collapse to 0.5.

Picture it

In the companion lab you can steer the generator toward the real distribution; the optimal D flattens to 0.5.

The math

min_G max_D E_{x~data}[log D(x)] + E_{z}[log(1−D(G(z)))]. Optimal D*(x) = p_data(x) / (p_data(x) + p_g(x)).

Derivation

For fixed G the inner max is solved pointwise by D*. Substituting it back leaves the Jensen–Shannon divergence between p_data and p_g — minimized (to 0) exactly when p_g = p_data, where D* = ½ everywhere.

Code

python

import numpy as np
# GAN objective: generator fools a discriminator (minimax game)
def D(x, w): return 1/(1+np.exp(-(x@w)))            # discriminator
real = np.random.randn(100, 4) + 3
fake = np.random.randn(100, 4)                       # from generator
w = np.random.randn(4)
# D wants D(real)->1, D(fake)->0; G wants D(fake)->1
loss_D = -np.mean(np.log(D(real,w)) + np.log(1-D(fake,w)+1e-9))
print("discriminator loss:", round(loss_D, 3))

In practice

GANs produce sharp samples but training is delicate (mode collapse, instability), spawning WGAN, spectral norm, etc. The adversarial idea also powers domain adaptation and image-to-image translation.

Diffusion: noise & denoise

destroy structure, then learn to rebuild it

The idea

Diffusion models learn to reverse a gradual noising process. Forward, you blend a clean sample with more and more Gaussian noise until it’s pure static; a network then learns to undo one step at a time, so you can start from noise and generate. Slide the noise level to watch a shape dissolve and the signal schedule decay.

Picture it

In the companion lab you can increase the timestep t — the ring dissolves into Gaussian noise as the signal √ā fades.

The math

Forward: x_t = √ā_t·x₀ + √(1−ā_t)·ε, ε~N(0,I). ā_t decreases 1→0 (here a cosine schedule).

Derivation

Composing many small Gaussian steps stays Gaussian, giving the closed form above. Training fits ε_θ(x_t, t) to predict the added noise; sampling subtracts the predicted noise step by step from x_T ~ N(0,I).

Code

python

import numpy as np
# Forward diffusion: gradually add Gaussian noise over T steps
x0 = np.array([1.0, 0.5])
betas = np.linspace(1e-4, 0.2, 50)
alpha_bar = np.cumprod(1 - betas)
t = 30
xt = np.sqrt(alpha_bar[t])*x0 + np.sqrt(1-alpha_bar[t])*np.random.randn(2)
print("noised sample at t=30:", np.round(xt, 2))   # model learns to reverse this

In practice

DDPMs power modern image/audio/video generation. Predicting noise ≈ score matching (∇ log p). Few-step samplers (DDIM) and latent diffusion (operate in a VAE latent) make it fast enough for production.

Regularization & generalization

Overfitting & early stopping

stop at the validation minimum

The idea

Train long enough and the training loss keeps falling while the validation loss bottoms out and then climbs — the model is starting to memorize noise. The gap between the two is the generalization error. Slide the stop epoch to the validation minimum: that’s early stopping, the cheapest regularizer there is.

Picture it

In the companion lab you can choose where to stop training; watch the train–validation gap.

The math

Generalization gap = L_val − L_train. Pick t* = argmin_t L_val(t); stop there.

Derivation

Each gradient step lets weights grow and fit finer detail — signal first, then noise. Halting early bounds the effective parameter movement, so it acts like an implicit norm penalty (closely related to L2 for linear models).

Code

python

import numpy as np
# Early stopping: halt when validation loss stops improving (patience)
val = [1.0, 0.7, 0.55, 0.5, 0.52, 0.53, 0.55]
best, wait, patience, stop = np.inf, 0, 2, None
for epoch, v in enumerate(val):
    if v < best: best, wait = v, 0
    else:
        wait += 1
        if wait >= patience: stop = epoch; break
print("stopped at epoch", stop)             # before it overfits

In practice

Standard practice: monitor a held-out set, keep the best checkpoint, stop after “patience” epochs without improvement. Free regularization that also saves compute.

Dropout

train a different sub-network each step

The idea

Dropout randomly switches off a fraction of units on every forward pass, so the network can’t lean on any single neuron — it must spread the work. Each pass trains a different thinned sub-network; at test time you use them all, scaled. Resample the mask and change the drop rate.

Picture it

In the companion lab you can set the drop rate and resample the mask — each draw is a different sub-network.

The math

Train: h̃ = (m / (1−p)) ⊙ h, mᵢ ~ Bernoulli(1−p). Test: use full h (the 1/(1−p) keeps the scale matched — “inverted dropout”).

Derivation

Zeroing units samples one of 2ⁿ thinned networks; training averages over them, and the test-time full network approximates that ensemble’s average. The 1/(1−p) rescale keeps E[h̃] = h.

Code

python

import numpy as np
# Dropout: randomly zero activations at train time; scale to keep the mean
def dropout(h, p=0.5, train=True):
    if not train: return h
    mask = (np.random.rand(*h.shape) > p) / (1-p)   # inverted dropout
    return h * mask
h = np.ones((2, 6))
print(dropout(h, p=0.5))                    # ~half zeroed, rest scaled by 2

In practice

Cheap, effective regularization — it prevents co-adaptation of features. Common in MLPs and Transformer sublayers (p≈0.1); largely replaced by BatchNorm in conv nets. Disabled at inference.

Weight decay (L2)

shrink weights, smooth the fit

The idea

Weight decay adds a penalty on the size of the weights, nudging them toward zero each step. Big weights mean a wiggly, high-variance fit; shrinking them smooths the function. Turn the penalty λ up and watch an overfit curve relax toward a gentle trend.

Picture it

In the companion lab you can increase the L2 penalty λ and watch the fit smooth and the weights shrink.

The math

min ‖y − Xw‖² + λ‖w‖² ⇒ w = (XᵀX + λI)⁻¹Xᵀy. SGD form: w ← (1 − ηλ)·w − η∇L.

Derivation

The extra λ‖w‖² adds 2λw to the gradient, so every step first multiplies w by (1−ηλ) — a constant pull toward 0. It’s the MAP estimate under a Gaussian prior on the weights (= ridge regression).

Code

python

import numpy as np
# Weight decay = L2 penalty; the gradient pulls weights toward 0 each step
w, lr, lam = np.array([2.0, -3.0]), 0.1, 0.05
grad = np.array([0.0, 0.0])                  # suppose data gradient is 0
for _ in range(50):
    w -= lr*(grad + lam*w)                    # decay term shrinks w
print(np.round(w, 3))                         # decays toward 0

In practice

The most common explicit regularizer. AdamW decouples decay from the adaptive step for correct behavior. Larger λ → simpler model (higher bias, lower variance).

Label smoothing

targets that aren’t quite certain

The idea

Hard one-hot targets push the network to make a class infinitely more likely than the rest — driving logits to extremes and breeding overconfidence. Label smoothing replaces the 1 with a slightly-less-than-1 and spreads the rest evenly, giving the loss a finite optimum and better-calibrated probabilities. Slide ε.

Picture it

In the companion lab you can increase ε to soften the target; the optimal logit gap becomes finite.

The math

y_LS = (1−ε)·y_onehot + ε/K. Loss = cross-entropy(softmax(z), y_LS).

Derivation

With a hard target the CE keeps falling as the true logit → ∞ (never satisfied). A soft target is matched at a finite logit where p_true = 1−ε+ε/K, so the minimizer stops at sensible confidence rather than diverging.

Code

python

import numpy as np
# Label smoothing: soften one-hot targets to curb overconfidence
def smooth(onehot, eps=0.1):
    K = len(onehot)
    return (1-eps)*onehot + eps/K
print(np.round(smooth(np.array([0,1,0,0.])), 3))   # [.025 .925 .025 .025]

In practice

Standard in image classification and Transformers (ε≈0.1). Improves calibration and generalization and reduces overconfident errors, at a tiny cost in raw accuracy.

Sequence models

Recurrent neural networks

a hidden state that remembers

The idea

An RNN reads a sequence one step at a time, carrying a hidden state that summarizes everything seen so far. The same weights apply at every step, so it handles any length. Feed a short input pulse and tune the recurrence: a strong self-weight makes the memory linger long after the input is gone.

Picture it

In the companion lab you can a pulse enters at the start; the recurrent weight sets how long the memory persists.

The math

h_t = tanh(w_h·h_{t−1} + w_x·x_t + b). y_t = w_y·h_t. Same weights at every step.

Derivation

Unrolling the recurrence in time turns the RNN into a (very deep) feed-forward net with shared weights. The hidden state is a learned running summary; w_h near 1 lets information persist across many steps.

Code

python

import numpy as np
# Vanilla RNN: carry a hidden state across time steps
def rnn_step(x, h, Wx, Wh, b): return np.tanh(x@Wx + h@Wh + b)
Wx,Wh,b = np.random.randn(4,8), np.random.randn(8,8), np.zeros(8)
h = np.zeros(8)
for x in np.random.randn(6, 4):              # sequence of length 6
    h = rnn_step(x, h, Wx, Wh, b)
print(h.shape)                                # final hidden summary

In practice

Language modeling, speech, time series. Parameter sharing across time = translation in time. The same weights everywhere is what makes RNNs (and their gated successors) compact and length-agnostic.

Backprop through time

the gradient’s long walk back in time

The idea

Training an RNN means backpropagating the error through every timestep — and the gradient to an early step is a product of one factor per step in between. If that factor is below one the gradient vanishes over time (the net forgets long-range dependencies); above one it explodes. Set the recurrent weight and sequence length to see the decay.

Picture it

In the companion lab you can see how the gradient to step t shrinks with distance from the output.

The math

∂h_t/∂h_{t−k} = ∏ (w_h·tanh′(·)). Magnitude ≈ |w_h·tanh′|^k.

Derivation

The chain rule multiplies one Jacobian factor per timestep. A product of k factors each <1 decays geometrically to 0 (vanishing); each >1 blows up (exploding) — the same depth problem, now stretched along time.

Code

python

import numpy as np
# Backprop through time: gradient is a product of Jacobians -> can explode/vanish
Wh = np.array([[1.1, 0.],[0., 0.9]])         # recurrent weight
g = np.eye(2)
for t in range(20):
    g = g @ Wh                                # repeated multiplication
print("grad norms after 20 steps:", np.round(np.diag(g), 2))  # 1.1^20 vs 0.9^20

In practice

Why plain RNNs can’t learn long-range dependencies. Fixes: gradient clipping (explosion), and gated additive memory — LSTM/GRU — for vanishing. Transformers sidestep it entirely with direct attention.

LSTM gates

a gated memory highway

The idea

The LSTM fixes vanishing gradients with a protected cell state updated by addition, not repeated multiplication. A forget gate decides how much of the old memory to keep and an input gate how much new to write. With the forget gate near 1 a value rides the cell state for many steps untouched — a gradient highway. Tune the gates.

Picture it

In the companion lab you can a value is written early; the forget gate controls how long the cell state remembers it.

The math

f,i,o = σ(W·[h_{t−1},x_t]). c_t = f⊙c_{t−1} + i⊙g_t, h_t = o⊙tanh(c_t).

Derivation

∂c_t/∂c_{t−1} = f. With f ≈ 1 the gradient passes through nearly unchanged — the “constant error carousel” — so error signals survive across long gaps instead of decaying like w_h^k.

Code

python

import numpy as np
# LSTM cell: gates decide what to forget, write, and read from the cell state
sig = lambda z: 1/(1+np.exp(-z))
def lstm(x, h, c, W):
    z = np.concatenate([x, h]) @ W           # combined gate pre-activations
    f,i,o,g = np.split(z, 4)
    c = sig(f)*c + sig(i)*np.tanh(g)          # forget + input gates
    h = sig(o)*np.tanh(c)                      # output gate
    return h, c
print("gates route gradients past the vanishing problem")

In practice

LSTMs and the lighter GRU (which merges gates) powered pre-Transformer sequence models — translation, speech, captioning. The additive gated memory is the key idea; Transformers later replaced recurrence with attention.

Seq2seq & the need for attention

the bottleneck that attention fixed

The idea

A basic encoder–decoder crams an entire input sequence into one fixed-size context vector, then decodes from it. For long inputs that bottleneck loses information — each token gets a smaller share. Attention removes the bottleneck by letting the decoder look back at every encoder state. Grow the sequence and toggle attention.

Picture it

In the companion lab you can lengthen the input and toggle attention to remove the single-vector bottleneck.

The math

No attention: c = h_T (one vector for the whole input). Attention: c_t = Σ_s α_{t,s}·h_s, α = softmax(score(decoder_t, h_s)).

Derivation

A fixed vector has constant capacity, so information per input token ∼ 1/length — it degrades as inputs grow. Attention gives the decoder a content-weighted view of all encoder states, so capacity scales with length.

Code

python

import numpy as np
# Encoder-decoder: encode a sequence to a context, decode a new sequence
def encode(seq, h, step):
    for x in seq: h = step(x, h)
    return h                                   # context vector
# decoder then generates tokens conditioned on the context (+ attention)
print("context summarizes the input for the decoder")

In practice

This is exactly why attention was introduced for translation — and, taken to its limit (drop recurrence, keep only attention), it becomes the Transformer that dominates modern sequence modeling.

Transformer architecture

Residual stream & LayerNorm

an additive highway through depth

The idea

Transformers are deep because of a simple trick: every sublayer adds its output back to a running “residual stream” instead of replacing it. That additive highway keeps signal and gradient stable through dozens of layers. LayerNorm standardizes each token before the sublayer reads it. Toggle the skip connection to see depth blow up without it.

Picture it

In the companion lab you can stack more layers; toggle the residual skip to see the activation scale with and without it.

The math

Pre-LN block: x ← x + Sublayer(LayerNorm(x)). LayerNorm: (x − μ)/σ · γ + β (over the feature dim).

Derivation

∂(x+F(x))/∂x = I + ∂F/∂x, so the gradient always has a clean +1 path — it can’t vanish through the skip. Without it, signal scales like (gain)^depth and explodes or dies, exactly the depth problem.

Code

python

import numpy as np
# Residual connection: y = x + f(x); gradients flow through the identity path
def block(x, f): return x + f(x)
x = np.random.randn(4, 8)
y = block(x, lambda z: 0.01*z)               # even if f is tiny, x passes through
print(np.allclose(y, x, atol=0.1))           # signal preserved across depth

In practice

Residual connections (from ResNets) plus LayerNorm are what make 100-layer Transformers trainable. Pre-LN (norm inside the block) is the stable modern default; post-LN needs careful warmup. The residual stream is also how interpretability reads the model.

Position-wise feed-forward

the per-token compute block

The idea

Between attention blocks sits a small MLP applied to each token independently: project up to a wider hidden size, apply a nonlinearity (GELU), project back down. Attention mixes information across tokens; the FFN does the per-token “thinking” and holds most of the parameters. Resize it and watch the cost.

Picture it

In the companion lab you can set the model width and expansion factor; see the up-then-down projection and parameter count.

The math

FFN(x) = W₂·GELU(W₁x + b₁) + b₂. W₁: d → r·d, W₂: r·d → d (r usually 4). Params ≈ 2·r·d².

Derivation

The same two matrices act on every position (a 1×1 “conv” over the sequence) — hence “position-wise.” The widening to r·d gives capacity for nonlinear feature mixing before compressing back to the residual width.

Code

python

import numpy as np
# Position-wise feed-forward: expand, nonlinearity, project back (per token)
def ffn(x, W1, W2):
    return np.maximum(0, x @ W1) @ W2          # d -> 4d -> d
x = np.random.randn(10, 64)
W1, W2 = np.random.randn(64,256), np.random.randn(256,64)
print(ffn(x, W1, W2).shape)                    # (10, 64)

In practice

FFN layers hold roughly two-thirds of a Transformer’s parameters and are a prime target for efficiency work (mixture-of-experts routes tokens to a subset of FFNs). Attention moves info between tokens; the FFN transforms it.

Multi-head attention

several attentions in parallel

The idea

Rather than one big attention, a Transformer splits the model width into several heads, each attending in its own smaller subspace, then concatenates and projects the results. Different heads can specialize — one tracks the previous token, another long-range agreement. Change the head count and watch the width partition.

Picture it

In the companion lab you can change the number of heads; the model width d splits into h subspaces of size d/h.

The math

head_i = Attention(QWᵢ^Q, KWᵢ^K, VWᵢ^V), dim d/h. MHA = Concat(head₁..head_h)·W^O.

Derivation

Splitting d into h subspaces costs the same total compute as one head of size d, but lets the model attend to several patterns at once. Each head’s d/h dimensions keep its dot products well-scaled (the √(d/h) factor).

Code

python

import numpy as np
def softmax(z): e=np.exp(z-z.max(-1,keepdims=True)); return e/e.sum(-1,keepdims=True)
# Multi-head attention: run several attentions in parallel, concatenate
def head(X, Wq, Wk, Wv, dk):
    Q,K,V = X@Wq, X@Wk, X@Wv
    return softmax(Q@K.T/np.sqrt(dk)) @ V
X = np.random.randn(5, 64); H, dk = 8, 8
outs = [head(X, *[np.random.randn(64,dk) for _ in range(3)], dk) for _ in range(H)]
print(np.concatenate(outs, -1).shape)          # (5, 64)

In practice

Typical: d=512, h=8 → 64-dim heads. Heads learn distinct roles (positional, syntactic, “induction” heads that copy patterns). Variants like multi-query / grouped-query attention share K,V across heads to shrink the inference cache.

A Transformer block, assembled

attention + FFN + residual, ×N

The idea

Put it together: each block routes the residual stream through two sublayers — multi-head attention then the feed-forward network — each wrapped in LayerNorm and a skip connection. Stack N of these and you have GPT or BERT. Add a causal mask and it can only look left, which is what makes a decoder generate left-to-right.

Picture it

In the companion lab you can stack more blocks; toggle the causal mask (decoder) vs full attention (encoder).

The math

a = x + MHA(LN(x)) y = a + FFN(LN(a)) Repeat N times; final LN → output head.

Derivation

Two residual sublayers per block: attention mixes across tokens, FFN transforms each token. A causal mask sets attention scores to −∞ above the diagonal, so position t sees only ≤ t — enabling autoregressive next-token prediction.

Code

python

import numpy as np
# A transformer block = (attention + residual + norm) then (FFN + residual + norm)
def ln(x): return (x-x.mean(-1,keepdims=True))/(x.std(-1,keepdims=True)+1e-5)
def block(x, attn, ffn):
    x = ln(x + attn(x))      # pre/post-norm variants exist
    x = ln(x + ffn(x))
    return x
print("stack N of these to build the encoder/decoder")

In practice

Encoder-only (BERT) for understanding, decoder-only (GPT, causal mask) for generation, encoder–decoder (T5) for seq2seq. Depth N and width d set the parameter budget; everything scales from this one repeated block.

Self-supervised learning

Self-supervised pretext tasks

hide part of the input, predict it

The idea

Labels are scarce, but raw data is everywhere — so self-supervision invents its own labels by hiding part of the input and asking the model to predict it. Masked language modeling hides a token; the model must use the surrounding context to fill it in, learning rich representations for free. Slide the mask along the sentence.

Picture it

In the companion lab you can move the mask; the model predicts the hidden token from both-side context.

The math

Maximize Σ log P(x_masked | x_context). No human labels — the data supplies its own target.

Derivation

Predicting a hidden piece forces the model to capture how parts of the data relate. Pretraining on this objective over huge corpora yields features that transfer; a small labeled set then fine-tunes for the real task.

Code

python

import numpy as np
# Pretext task: hide part of the input, predict it (no human labels)
sentence = list("the cat sat on the mat")
mask_at = 4
target = sentence[mask_at]; sentence[mask_at] = "_"
print("".join(sentence), "-> predict:", repr(target))  # masked modeling

In practice

BERT (masked tokens), GPT (next token), MAE (masked image patches), and word2vec all use pretext tasks. This pretrain-then-fine-tune recipe is the backbone of modern foundation models.

Contrastive learning (InfoNCE)

pull positives close, push negatives away

The idea

Contrastive learning shapes an embedding space by example: pull an anchor toward its positive (two views of the same thing) and push it away from a batch of negatives. The InfoNCE loss is a softmax over similarities — essentially “classify which item is the true match.” Drag the anchor and set the temperature.

Picture it

In the companion lab you can drag the anchor; the loss drops as it nears its positive and leaves the negatives.

The math

L = −log [ exp(sim(a,p)/τ) / Σ_k exp(sim(a,k)/τ) ]. sim = cosine; sum runs over the positive + all negatives.

Derivation

It’s cross-entropy for a classifier whose “classes” are the candidates: the positive should win the softmax. Low τ sharpens focus on the hardest negatives; the gradient pulls a toward p and away from whichever negatives are closest.

Code

python

import numpy as np
# Contrastive (InfoNCE): pull two views of the same item together, push others apart
def info_nce(z, zp, temp=0.1):
    z = z/np.linalg.norm(z,axis=1,keepdims=True)
    zp = zp/np.linalg.norm(zp,axis=1,keepdims=True)
    logits = z @ zp.T / temp                   # similarity matrix
    labels = np.arange(len(z))                  # positives on the diagonal
    e = np.exp(logits - logits.max(1,keepdims=True))
    p = e[labels, labels]/e.sum(1)
    return -np.log(p).mean()
print(round(info_nce(np.random.randn(8,16), np.random.randn(8,16)), 2))

In practice

SimCLR, MoCo, CLIP (image–text pairs). Large batches supply many negatives; temperature is a key knob. Produces general-purpose embeddings used for retrieval, zero-shot classification, and as frozen features.

Augmentation & invariance

two views, one meaning

The idea

Where do positive pairs come from when there are no labels? From the same example, augmented two different ways — crop, flip, recolor. The model is trained to map both views to the same embedding, learning to ignore the nuisances augmentation changes while keeping what matters. Crank the augmentation strength.

Picture it

In the companion lab you can increase augmentation strength; the two views differ more, yet should embed to the same point.

The math

x → (t₁(x), t₂(x)) with random augmentations t. Objective: encoder f makes f(t₁(x)) ≈ f(t₂(x)).

Derivation

Treating two augmentations of one image as a positive pair bakes in invariance to those transforms. Choosing augmentations is choosing which factors to discard — the inductive bias that makes self-supervised features useful.

Code

python

import numpy as np
# Augmentations create two "views" the model must agree on
def augment(x):
    x = x + 0.1*np.random.randn(*x.shape)       # jitter
    if np.random.rand() < 0.5: x = x[..., ::-1] # flip
    return x
img = np.random.rand(8, 8)
v1, v2 = augment(img), augment(img)             # positive pair
print(v1.shape, v2.shape)

In practice

SimCLR showed augmentation choice (especially crop + color) is decisive. The same idea underlies BYOL, SwAV, and data2vec. Strong-but-not-destructive augmentation is the art.

Representation collapse

the trivial solution to avoid

The idea

There’s a cheap way to make all positive pairs match: map everything to the same point. That’s collapse — a useless constant encoder. Negatives prevent it by demanding that different inputs stay apart (uniformity). Toggle the negatives off and watch the embeddings implode to a dot.

Picture it

In the companion lab you can turn negatives off and the embeddings collapse to one point; on, they spread out.

The math

Alignment: positives close. Uniformity: features spread on the sphere. Collapse = all f(x) equal (alignment perfect, uniformity zero).

Derivation

With only an “attract positives” term, the global optimum is a constant function. Negatives (or alternatives below) add a repulsive term, forcing the encoder to spread inputs out — retaining information.

Code

python

import numpy as np
# Collapse: without negatives/stop-grad, all embeddings become identical
emb = np.random.randn(100, 16)
for _ in range(50):
    emb = 0.9*emb.mean(0) + 0.1*emb              # naive "pull together" only
print("embedding variance:", round(emb.var(), 4))  # -> ~0, collapsed

In practice

Negative-free methods avoid collapse by other means: BYOL/SimSiam use a stop-gradient + predictor, Barlow Twins decorrelates feature dimensions, VICReg adds variance/covariance terms. Avoiding collapse is the central design problem of SSL.

Reinforcement learning

Markov decision process & value

how good is it to be here?

The idea

Reinforcement learning casts a problem as states, actions, rewards, and transitions — a Markov decision process. The value of a state is the total (discounted) future reward you expect from it under good play. High value flows outward from the goal, and the greedy arrows trace the optimal path. Tune the discount.

Picture it

In the companion lab you can change the discount γ (how much the future counts). Values and the greedy policy update.

The math

Return G_t = Σ_{k≥0} γᵏ r_{t+k}. V*(s) = max_a [ r(s,a) + γ·V*(s′) ].

Derivation

A reward d steps away is worth γᵈ now, so γ<1 makes the sum finite and prefers sooner rewards. The optimal value satisfies a recursive consistency (the Bellman equation), and acting greedily w.r.t. V* is optimal.

Code

python

import numpy as np
# An MDP: states, actions, transition probs, rewards, discount gamma
states = [0, 1]; gamma = 0.9
P = {0:{ 'a':[(0.8,0,1.0),(0.2,1,0.0)] }, 1:{ 'a':[(1.0,1,0.0)] }}   # (prob, s', r)
# expected immediate reward of taking 'a' in state 0
exp_r = sum(p*r for p,_,r in P[0]['a'])
print("E[r | s=0, a] =", exp_r)

In practice

Deep RL replaces the value table with a network (DQN) for huge state spaces — Atari, Go, robotics. γ sets the planning horizon. Everything (Q-learning, actor-critic) builds on this MDP/value foundation.

The Bellman equation

value = reward + discounted next value

The idea

The Bellman equation is the recursion at the heart of RL: a state’s value equals the immediate reward plus the discounted value of the best next state. Value iteration just applies it as an update, sweep after sweep, and the values ripple outward from the goal until they stop changing. Step it and watch the wavefront spread.

Picture it

In the companion lab you can apply the Bellman backup one sweep at a time; values propagate from the goal.

The math

V_{k+1}(s) = max_a [ r(s,a) + γ·V_k(s′) ]. Iterate → converges to V* (a contraction, factor γ).

Derivation

The Bellman optimality operator is a γ-contraction in max-norm, so repeated application has a unique fixed point V* and converges geometrically. Each sweep lets information travel one more step back from the goal.

Code

python

import numpy as np
# Value iteration solves the Bellman optimality equation
R = np.array([0., 1.]); gamma = 0.9
P = np.array([[0.5,0.5],[0.,1.]])            # transition under the (only) policy
V = np.zeros(2)
for _ in range(100):
    V = R + gamma * P @ V                     # Bellman backup
print(np.round(V, 2))

In practice

Q-learning is the sampled, model-free version: Q(s,a) ← Q + α[r + γ max Q(s′,·) − Q]. DQN fits this target with a network; the same backup underlies temporal-difference learning broadly.

Policy gradient

make good actions more probable

The idea

Instead of valuing states, policy-gradient methods optimize the action policy directly. Sample actions, see which earned high reward, and nudge the policy to make those more likely — gradient ascent on expected return. Step it and watch a Gaussian policy drift onto the reward peak.

Picture it

In the companion lab you can run policy-gradient steps; the policy mean climbs toward the high-reward region.

The math

Maximize J(θ) = E_{a~π_θ}[R(a)]. ∇J = E[ ∇ log π_θ(a)·R(a) ] (REINFORCE).

Derivation

The log-derivative trick rewrites the gradient of an expectation as an expectation of ∇logπ·R, estimable from samples. Subtracting a baseline (e.g. mean return) cuts variance without bias — the actor-critic idea.

Code

python

import numpy as np
# REINFORCE: push up log-prob of actions weighted by the return
def softmax(z): e=np.exp(z-z.max()); return e/e.sum()
theta = np.zeros(2); lr = 0.1
for ep in range(2000):
    p = softmax(theta); a = np.random.choice(2, p=p)
    R = 1.0 if a == 1 else 0.0                 # reward action 1
    grad = -p; grad[a] += 1                     # d log pi / d theta
    theta += lr * R * grad
print(np.round(softmax(theta), 2))             # -> favors action 1

In practice

Policy nets handle continuous/high-dim actions where value tables fail; PPO/TRPO stabilize the updates and power robotics and RLHF (aligning LLMs from preference rewards). Higher variance than value methods, hence baselines and clipping.

Exploration vs exploitation

try new things, or bank the best?

The idea

An agent that always picks its current best guess may never discover a better option; one that always explores never cashes in. The multi-armed bandit is the cleanest version of this dilemma. ε-greedy explores at random a fraction ε of the time. Pull the arms and watch the estimates sharpen and regret accumulate.

Picture it

In the companion lab you can set ε and pull; estimates converge to the true arm values, regret grows slower with good ε.

The math

ε-greedy: best arm w.p. 1−ε, random w.p. ε. Update: Q(a) ← Q(a) + (1/N(a))·(r − Q(a)). Regret = Σ (μ* − μ_{chosen}).

Derivation

The running mean is an unbiased value estimate; a little forced exploration guarantees every arm is sampled enough to be ranked correctly. UCB/Thompson sampling explore smarter — by optimism or posterior sampling — for lower regret.

Code

python

import numpy as np
# Exploration vs exploitation: epsilon-greedy over action values
Q = np.array([0.1, 0.9, 0.3]); eps = 0.1
def act(Q, eps):
    return np.random.randint(len(Q)) if np.random.rand() < eps else Q.argmax()
choices = [act(Q, eps) for _ in range(1000)]
print("explored non-greedy", round(np.mean(np.array(choices)!=1), 2), "of the time")

In practice

The explore–exploit trade-off pervades RL (ε-greedy in DQN), recommender cold-start, A/B testing, and hyperparameter search. Too little exploration → stuck on a suboptimal arm; too much → wasted pulls.

Graph neural networks

Graphs & the adjacency matrix

structure as a matrix

The idea

A graph is nodes joined by edges — social networks, molecules, road maps. Its structure is captured by the adjacency matrix A, where A_{ij}=1 if i and j are linked. A node’s row lists its neighbors, and the degree is the row sum. Pick a node to light up its neighborhood and the matching matrix row.

Picture it

In the companion lab you can select a node to highlight its neighbors and its row/column in the adjacency matrix.

The math

A_{ij} = 1 if edge (i,j) else 0 (symmetric, undirected). degree d_i = Σ_j A_{ij} = row sum.

Derivation

The adjacency matrix turns graph structure into linear algebra: A·x spreads a signal x to neighbors, A² reaches two hops, and the eigenvectors of A (or the Laplacian D−A) reveal communities and the smooth modes GNNs operate on.

Code

python

import numpy as np
# A graph: node features X and an adjacency matrix A
A = np.array([[0,1,1,0],[1,0,1,0],[1,1,0,1],[0,0,1,0.]])   # 4 nodes
X = np.random.randn(4, 8)                                     # node features
deg = A.sum(1)                                                # node degrees
print("degrees:", deg)

In practice

GNNs operate on (A, X): the adjacency for structure, a node-feature matrix X for content. Real graphs are huge and sparse, so A is stored as an edge list — the substrate for every message-passing layer that follows.

Message passing

gather from neighbors, update, repeat

The idea

A graph neural net updates each node by gathering messages from its neighbors, aggregating them (sum or mean), and combining with its own state. One round mixes information one hop; repeat and a signal diffuses across the graph. Inject a signal at one node and step the rounds to watch it spread.

Picture it

In the companion lab you can node 0 starts with a signal; each round averages neighbors so it diffuses outward.

The math

m_i = AGG({ h_j : j ∈ N(i) }) (sum / mean / max) h_i′ = UPDATE(h_i, m_i).

Derivation

Stacking k message-passing layers gives each node a receptive field of its k-hop neighborhood — the graph analogue of a CNN’s growing field. Permutation-invariant aggregators (sum/mean) make the output independent of neighbor ordering.

Code

python

import numpy as np
# Message passing: each node aggregates features from its neighbors
A = np.array([[0,1,1,0],[1,0,1,0],[1,1,0,1],[0,0,1,0.]])
X = np.random.randn(4, 8)
msg = A @ X                                  # sum of neighbor features
agg = msg / A.sum(1, keepdims=True).clip(1)  # mean aggregation
print(agg.shape)                              # (4, 8) updated per node

In practice

The unifying framework (MPNN): GCN, GraphSAGE, GAT differ only in how they aggregate (mean, sampled, attention-weighted). Powers molecules, recommendation graphs, knowledge graphs, and physics simulations.

Graph convolution (GCN)

normalized neighbor mixing + a learned map

The idea

The graph convolutional layer is message passing with a normalized adjacency and a learnable weight: aggregate neighbors (scaled by degree so high-degree nodes don’t dominate), multiply by W, apply a nonlinearity. Stacking L layers grows each node’s receptive field to L hops. Slide the depth to see how far a node “sees.”

Picture it

In the companion lab you can set the number of GCN layers; node 0’s receptive field grows one hop per layer.

The math

H′ = σ( D̃^{−1/2} Ã D̃^{−1/2} H W ), Ã = A + I. Symmetric normalization tempers high-degree nodes.

Derivation

Adding self-loops (A+I) keeps a node’s own features; the D̃^{−1/2}·Ã·D̃^{−1/2} scaling is a first-order approximation of spectral graph convolution, keeping eigenvalues bounded so deep stacks stay stable. W is the only learned part.

Code

python

import numpy as np
# GCN layer: symmetric-normalized propagation  H = D^-.5 A~ D^-.5 X W
A = np.array([[0,1,1,0],[1,0,1,0],[1,1,0,1],[0,0,1,0.]])
A_tilde = A + np.eye(4)                       # add self-loops
d = A_tilde.sum(1); Dinv = np.diag(1/np.sqrt(d))
Ahat = Dinv @ A_tilde @ Dinv
X, W = np.random.randn(4,8), np.random.randn(8,16)
H = np.maximum(0, Ahat @ X @ W)
print(H.shape)                                # (4, 16)

In practice

Kipf & Welling’s GCN is the canonical baseline for node classification (citation graphs) and link prediction. Attention (GAT) learns the neighbor weights; the receptive field = number of layers, which is also why too many layers hurt.

Over-smoothing

why deeper GNNs can do worse

The idea

GNNs have a depth paradox: more layers reach farther, but each layer also averages neighbors — and too much averaging makes every node’s features converge to the same value. Discriminative signal washes out, so very deep GNNs often do worse. Increase the layers and watch the colors collapse toward uniform.

Picture it

In the companion lab you can add layers; repeated neighbor-averaging drives all node features toward the same value.

The math

Repeated H ← D⁻¹Ã H drives features toward the dominant eigenvector — a node-degree-weighted constant. feature variance across nodes → 0.

Derivation

The normalized propagation is a smoothing (low-pass) operator; iterating it is power iteration toward its top eigenvector, so differences between nodes shrink geometrically. Beyond a few layers, embeddings become nearly indistinguishable.

Code

python

import numpy as np
# Oversmoothing: too many GCN layers make all node features converge
A = np.array([[0,1,1,0],[1,0,1,0],[1,1,0,1],[0,0,1,0.]]) + np.eye(4)
d = A.sum(1); Ahat = np.diag(1/d) @ A
X = np.random.randn(4, 8)
for _ in range(50): X = Ahat @ X              # repeated smoothing
print("feature variance across nodes:", round(X.var(0).mean(), 5))  # ->0

In practice

The core obstacle to deep GNNs. Mitigations: residual/skip connections, jumping-knowledge, PairNorm, DropEdge, and initial-residual schemes (GCNII). Often 2–3 layers is the practical sweet spot.

Efficiency & deployment

Quantization

fewer bits per weight

The idea

Weights trained in 32-bit float can usually be stored in far fewer bits. Quantization snaps each value to the nearest level on a coarse grid — INT8 or even INT4 — shrinking the model 4–8× and speeding up integer math, at the cost of small rounding error. Drop the bit-width and watch the grid coarsen and error rise.

Picture it

In the companion lab you can reduce the bit-width; weights snap to a coarser grid — smaller model, larger rounding error.

The math

q = round(w / s)·s, s = range / (2ᵇ − 1). levels = 2ᵇ; rounding error ≈ s/√12.

Derivation

b bits give 2ᵇ levels spaced by step s, so each weight is off by at most s/2 (uniform error, std s/√12). Halving bits doubles the step — error grows geometrically, which is why 8-bit is usually safe and 4-bit needs care.

Code

python

import numpy as np
# Quantize float32 weights to int8 (scale + round); ~4x smaller
w = np.random.randn(1000).astype(np.float32)
scale = np.abs(w).max() / 127
q = np.round(w / scale).astype(np.int8)       # store int8 + one scale
w_recon = q.astype(np.float32) * scale
print("max error:", round(np.abs(w - w_recon).max(), 4))

In practice

Post-training quantization and quantization-aware training; INT8 inference is routine, 4-bit (GPTQ/AWQ) common for LLM serving. Outlier channels need per-channel scales or mixed precision. The big win: memory bandwidth and energy.

Pruning & sparsity

drop the weights that barely matter

The idea

Most weights in a trained network are near zero and contribute little. Pruning sets the smallest-magnitude ones to zero, producing a sparse model that stores and computes less. Accuracy holds up to surprisingly high sparsity, then falls off a cliff. Slide the pruning fraction to sparsify the matrix and trace the trade-off.

Picture it

In the companion lab you can prune the smallest-magnitude weights; accuracy holds, then drops past a point.

The math

Keep w if |w| > τ (the chosen percentile), else 0. sparsity = fraction of zeroed weights.

Derivation

Magnitude is a cheap proxy for importance: tiny weights change the output little, so zeroing them barely moves the loss — until you cut into load-bearing weights and accuracy collapses. Fine-tuning after pruning recovers much of the loss.

Code

python

import numpy as np
# Magnitude pruning: zero out the smallest-magnitude weights
w = np.random.randn(10, 10)
k = 0.5                                        # prune 50%
thresh = np.quantile(np.abs(w), k)
w_pruned = w * (np.abs(w) >= thresh)
print("sparsity:", round((w_pruned == 0).mean(), 2))

In practice

Unstructured pruning gives high sparsity but needs sparse kernels to speed up; structured pruning (whole channels/heads) yields real latency wins. The lottery-ticket hypothesis: sparse subnetworks can train to full accuracy.

Knowledge distillation

a small model learns from a big one’s hunches

The idea

A small “student” can be trained to mimic a big “teacher.” The trick is to learn from the teacher’s soft probabilities, not just the hard label: the relative scores across wrong classes (“dark knowledge”) reveal which classes look similar, a much richer signal. Raise the temperature to soften the targets.

Picture it

In the companion lab you can raise the softmax temperature T; the teacher’s targets soften, exposing inter-class similarity.

The math

p_i(T) = softmax(z_i / T). Loss = α·CE(hard) + (1−α)·T²·KL(student_T ∥ teacher_T).

Derivation

Higher T flattens the softmax, amplifying the small logits that encode “class 2 looks a bit like class 4.” Matching those soft targets transfers more information per example than a one-hot label; the T² factor keeps gradient scale constant.

Code

python

import numpy as np
# Knowledge distillation: student matches teacher's soft probabilities
def softmax_T(z, T): e=np.exp(z/T - (z/T).max()); return e/e.sum()
teacher_logits = np.array([4., 2., 1.])
soft = softmax_T(teacher_logits, T=2.0)        # softened targets carry "dark knowledge"
print(np.round(soft, 3))                        # student trains on these

In practice

Compresses large models into deployable ones (DistilBERT) and underlies self-distillation and label smoothing’s cousin. Common in LLM serving: distill a frontier model’s behavior into a smaller, cheaper one.

Mixture of experts

huge capacity, sparse compute

The idea

A mixture-of-experts layer holds many expert sub-networks but routes each token to only a few via a learned gate. Parameters scale with the number of experts, while compute scales only with how many you activate — so you grow capacity almost for free. Set the expert count and the top-k routing.

Picture it

In the companion lab you can set the number of experts and how many each token uses (top-k). Params grow; compute stays low.

The math

gate g(x) = softmax(W_g x); pick top-k experts. y = Σ_{i∈topk} g_i(x)·E_i(x). compute ∝ k, params ∝ N.

Derivation

Only k of N experts fire per token, so FLOPs depend on k while the parameter count depends on N — decoupling capacity from cost. A load-balancing loss keeps the gate from collapsing onto a few experts.

Code

python

import numpy as np
def softmax(z): e=np.exp(z-z.max()); return e/e.sum()
# Mixture of Experts: a router sends each token to its top-k experts
x = np.random.randn(8)                          # a token
gate = softmax(np.random.randn(4))              # scores over 4 experts
topk = gate.argsort()[::-1][:2]                 # use only 2 (sparse)
print("routed to experts:", topk, "  (rest skipped -> cheap)")

In practice

Sparse MoE scales LLMs to trillions of parameters (Switch Transformer, Mixtral) at the FLOPs of a much smaller dense model. Challenges: routing instability, load balancing, and the memory to hold all experts.

Interpretability

Saliency: input gradients

what the model is looking at

The idea

The simplest explanation of a prediction: which input pixels would change it most? That’s the gradient of the output with respect to the input — a saliency map. Raw gradients are noisy, so SmoothGrad averages the map over many slightly-noised copies of the input. Slide the smoothing and watch the edges sharpen out of the noise.

Picture it

In the companion lab you can increase SmoothGrad samples; averaging cancels gradient noise and the salient edges emerge.

The math

saliency(x) = |∂ score / ∂ x|. SmoothGrad: average |∂score/∂x| over x + noise, n times.

Derivation

The gradient says how a tiny pixel change moves the score — a first-order local explanation. Because that gradient is jagged, adding noise and averaging (SmoothGrad) keeps the consistent signal (edges) while the random part shrinks like 1/√n.

Code

python

import numpy as np
# Saliency: gradient of the output w.r.t. each input pixel (sensitivity)
def model(x, w): return 1/(1+np.exp(-(x@w)))
x, w = np.random.randn(16), np.random.randn(16)
p = model(x, w)
saliency = (p*(1-p))*w                          # d sigmoid / d x
print("most influential input:", np.abs(saliency).argmax())

In practice

Fast, model-agnostic attributions for vision. Caveats: raw saliency can be noisy and even input-insensitive; sanity-check against randomized-model baselines. Grad-CAM (gradients at conv layers) gives cleaner class-localized maps.

Integrated gradients

attribution that adds up

The idea

Plain gradients can be zero exactly where a feature already saturated the output — missing its importance. Integrated gradients fix this by accumulating gradients along a straight path from a neutral baseline to the actual input. The attributions then satisfy completeness: they sum to the change in output. Add path steps and watch the sum converge.

Picture it

In the companion lab you can increase the number of path steps; the Riemann sum converges and completeness holds.

The math

IG_i = (x_i − b_i)·∫₀¹ ∂f(b + α(x−b))/∂x_i dα. Completeness: Σ_i IG_i = f(x) − f(b).

Derivation

By the fundamental theorem of calculus, integrating the gradient along the baseline→input path recovers the total output change, split cleanly across features. A Riemann sum over k steps approximates the integral, converging as k grows.

Code

python

import numpy as np
# Integrated gradients: average the gradient along a path from a baseline
def grad(x, w): p=1/(1+np.exp(-(x@w))); return p*(1-p)*w
x, base, w = np.random.randn(8), np.zeros(8), np.random.randn(8)
steps = 50
ig = (x - base) * np.mean([grad(base + (k/steps)*(x-base), w) for k in range(steps)], 0)
print("attributions sum:", round(ig.sum(), 3))   # satisfies completeness

In practice

A principled, axiomatic attribution (sensitivity + implementation invariance) widely used for tabular, text, and vision models. The baseline choice matters — a black image, zero embedding, or average input encodes “absence.”

Probing & concept directions

reading concepts off the activations

The idea

Does a network actually encode a concept — say, sentiment or part-of-speech? Train a simple linear probe on its hidden activations: if a linear boundary separates the concept, the representation encodes it linearly. Probe accuracy tends to rise with depth as features grow more abstract. Slide the layer and watch the classes pull apart.

Picture it

In the companion lab you can move to deeper layers; the two concept classes become linearly separable and probe accuracy climbs.

The math

Fit w on frozen activations h: ŷ = sign(w·h + b). High probe accuracy ⇒ the concept is linearly decodable.

Derivation

A linear probe has tiny capacity, so it can only succeed if the concept already lies along a direction in activation space. Comparing probes across layers traces where and how a concept becomes explicit — the basis of the “linear representation” view.

Code

python

import numpy as np
from sklearn.linear_model import LogisticRegression
# Probing: train a simple classifier on frozen activations to test what they encode
acts = np.random.randn(500, 64)                  # hidden representations
concept = (acts[:, 0] > 0).astype(int)           # some property
acc = LogisticRegression().fit(acts, concept).score(acts, concept)
print("probe accuracy:", round(acc, 2))          # high -> info is linearly present

In practice

Standard for analyzing language and vision models (BERTology, concept directions, steering vectors). Caveat: probe success shows information is present and decodable, not that the model uses it — pair with causal interventions.

Feature visualization

what turns a neuron on

The idea

Another way in: find the input that maximally activates a given neuron, by optimizing the image rather than the weights. The results reveal a hierarchy — early units fire for oriented edges, middle ones for textures and curves, deep ones for object-like motifs. Move through the layers and switch neurons to see the preferred stimulus grow in complexity.

Picture it

In the companion lab you can move through layers (edges → textures → motifs) and switch neurons to see different preferred patterns.

The math

x* = argmax_x a_neuron(x) − λ·R(x). Optimize the input by gradient ascent; R regularizes for natural images.

Derivation

Ascending the neuron’s activation w.r.t. the input synthesizes its ideal stimulus. Regularizers (frequency penalties, jitter, transformations) keep the result interpretable rather than adversarial noise — revealing the feature the unit detects.

Code

python

import numpy as np
# Feature visualization: optimize an input to maximally activate a neuron
def activation(x, w): return x @ w
w = np.random.randn(16)
x = np.random.randn(16)
for _ in range(200):
    x += 0.1 * w                                 # ascend the activation
    x = np.clip(x, -3, 3)
print("found input that maximizes the neuron:", np.round(x[:4], 1))

In practice

The basis of the OpenAI/DeepDream “circuits” line of work: edges → textures → parts → objects across depth. Combined with attribution and probing, it builds mechanistic pictures of what networks compute.

Language models & alignment

Tokenization (BPE)

text into subword pieces

The idea

Language models don’t read characters or whole words — they read subword tokens. Byte-pair encoding starts from characters and repeatedly merges the most frequent adjacent pair into a new token, so common chunks become single units while rare words still decompose. Slide the merge count and watch words fuse from characters into subwords.

Picture it

In the companion lab you can increase BPE merges; frequent character pairs fuse into longer tokens (fewer, bigger pieces).

The math

Repeat: find the most frequent adjacent token pair, merge it into one new token, add to vocab. Stop at a target vocab size.

Derivation

Greedy frequency-based merging balances vocabulary size against sequence length: frequent words become single tokens (short sequences), rare/novel words fall back to pieces (no out-of-vocab). Byte-level BPE guarantees any string is representable.

Code

python

# Byte-pair encoding idea: merge the most frequent adjacent pair, repeat
from collections import Counter
tokens = list("low lower lowest")
for _ in range(3):
    pairs = Counter(zip(tokens[:-1], tokens[1:]))
    best = max(pairs, key=pairs.get)
    merged, i = [], 0
    while i < len(tokens):
        if i < len(tokens)-1 and (tokens[i],tokens[i+1])==best:
            merged.append(tokens[i]+tokens[i+1]); i+=2
        else: merged.append(tokens[i]); i+=1
    tokens = merged
print(tokens)

In practice

BPE / WordPiece / Unigram (SentencePiece) tokenizers underlie GPT, BERT, Llama. Token count drives context length and cost; quirks (digits, whitespace, code, non-English scripts) trace back to how the tokenizer was trained.

Decoding: temperature & top-p

turning probabilities into text

The idea

A language model outputs a probability over the next token; how you sample from it shapes the text. Temperature flattens or sharpens the distribution, and nucleus (top-p) sampling keeps only the smallest set of tokens whose probability sums to p. Tune both and watch the candidate set and shape change.

Picture it

In the companion lab you can raise temperature to flatten, lower top-p to truncate to the nucleus (kept = blue).

The math

p_i = softmax(logit_i / T). Top-p: keep the smallest set with Σ p_i ≥ p, renormalize.

Derivation

Dividing logits by T > 1 shrinks differences (more random, more diverse); T → 0 is greedy argmax. Nucleus sampling adapts the cutoff to the distribution’s shape — keeping few tokens when the model is confident, many when it’s unsure — unlike fixed top-k.

Code

python

import numpy as np
# Temperature + top-p sampling shape the next-token distribution
def sample(logits, T=1.0, p=0.9):
    z = logits / T; e = np.exp(z - z.max()); probs = e/e.sum()
    idx = np.argsort(probs)[::-1]
    cum = np.cumsum(probs[idx]); keep = idx[:np.searchsorted(cum, p)+1]
    pk = probs[keep]/probs[keep].sum()
    return np.random.choice(keep, p=pk)
print(sample(np.array([3.,2.,1.,0.5]), T=0.7))

In practice

The core generation knobs. Low T / small p → focused, repetitive; high T / p → creative, riskier. Add top-k, repetition penalties, and beam search for structured tasks. Sampling choice strongly shapes outputs at fixed model.

Scaling laws

predictable returns to scale

The idea

Across orders of magnitude, model loss falls as a clean power law in parameters, data, and compute — a straight line on a log-log plot, bottoming out at an irreducible floor. These laws let you predict a big model’s loss from small runs and budget compute optimally. Slide the compute and watch the loss slide down the line.

Picture it

In the companion lab you can increase the compute budget; loss follows a power law toward the irreducible floor (diminishing returns).

The math

L(C) ≈ E + a·C^{−b} (power law + irreducible loss E). Compute-optimal (Chinchilla): scale params N and data D together, N ∝ √C, D ∝ √C.

Derivation

Empirically loss is linear in log-log over many decades, so doubling compute cuts the reducible loss by a fixed fraction — until it hits the data/architecture floor E. Given a budget C ≈ 6·N·D, the optimum balances model size and tokens rather than over-growing one.

Code

python

import numpy as np
# Scaling laws: loss falls as a power law in params/data/compute
def loss(N, alpha=0.076, Nc=8.8e13): return (Nc/N)**alpha
for N in (1e8, 1e9, 1e10, 1e11):
    print(f"N={N:.0e}  loss~{loss(N):.3f}")     # diminishing but predictable

In practice

Scaling laws (Kaplan; Chinchilla) guide frontier model planning and the “tokens ≈ 20× params” rule of thumb. They also frame debates on data limits and where returns finally bend.

Alignment: RLHF

optimize for human preference, but stay grounded

The idea

A pretrained model predicts text; it isn’t trained to be helpful or safe. RLHF fixes that: collect human preferences, fit a reward model, then optimize the policy to earn reward — while a KL penalty tethers it to the original model so it doesn’t drift or exploit quirks. Lower the leash and watch the policy chase a spurious reward spike (reward hacking).

Picture it

In the companion lab you can change the KL penalty. Too weak → the policy over-optimizes a fake reward spike; too strong → it barely moves.

The math

max_π E_π[ r(x) ] − β·KL(π ∥ π_SFT). Pipeline: pretrain → SFT → reward model → PPO/DPO.

Derivation

The reward model is an imperfect proxy for human preference, so pure reward-maximization finds its exploits (reward hacking). The KL term penalizes drift from the trusted base policy, trading reward against staying in-distribution — β sets that balance.

Code

python

import numpy as np
# RLHF reward modeling: learn from human preferences (Bradley-Terry)
def pref_prob(r_win, r_lose): return 1/(1+np.exp(-(r_win - r_lose)))
# train reward so preferred response scores higher; then RL-optimize the policy
print("P(human prefers A) =", round(pref_prob(2.0, 1.0), 3))

In practice

RLHF (and DPO, which skips the explicit RL) aligns chat models for helpfulness and safety. Open problems: reward over-optimization (Goodhart), annotation bias, and sycophancy. The KL leash is the central safety knob.

Multimodal learning

Joint embedding space

pictures and words in one space

The idea

Multimodal models map different kinds of data — images and text — into one shared vector space, so a picture of a dog and the words “a dog” land near each other. Distance then means semantic similarity across modalities, enabling search by either. Pick a caption and watch it retrieve its nearest image.

Picture it

In the companion lab you can select a text query; the nearest image in the shared space is retrieved (its true match).

The math

image → f(img), text → g(txt), both in ℝᵈ. similarity = cos(f(img), g(txt)); retrieve argmax.

Derivation

Two separate encoders are trained so paired (image, text) embeddings align. Once the space is shared, cross-modal retrieval is just nearest-neighbor search — no task-specific head — and arithmetic/clustering work across modalities.

Code

python

import numpy as np
# Joint embedding: map images and text into one shared vector space
img_emb  = np.random.randn(10, 64)
text_emb = np.random.randn(10, 64)
def norm(x): return x/np.linalg.norm(x,axis=1,keepdims=True)
sim = norm(img_emb) @ norm(text_emb).T           # cross-modal similarity
print(sim.shape)                                  # (10, 10)

In practice

CLIP, ALIGN, and multimodal LLM encoders build these spaces. They power text→image search, captioning, and grounding vision models in language — the substrate for the contrastive training in the next card.

Contrastive image–text (CLIP)

match the diagonal, suppress the rest

The idea

CLIP learns the shared space from a batch of image–text pairs. Compute every image-vs-every-text similarity to form an N×N matrix; the correct pairs lie on the diagonal. Train to push the diagonal up and everything off-diagonal down — a symmetric contrastive loss. Train it and watch the diagonal light up.

Picture it

In the companion lab you can train the contrastive objective; the matching pairs (diagonal) brighten, mismatches darken.

The math

logits = (I · Tᵀ)/τ over the batch. Loss = ½[ CE(rows, diag) + CE(cols, diag) ] — symmetric InfoNCE.

Derivation

Each row is a classification: which of the N captions matches this image? The label is the diagonal. Minimizing cross-entropy both ways aligns paired embeddings and repels the N−1 negatives per example — contrastive learning at batch scale.

Code

python

import numpy as np
# CLIP: contrastive loss aligns matching image-text pairs on the diagonal
def clip_loss(I, T, temp=0.07):
    I = I/np.linalg.norm(I,axis=1,keepdims=True)
    T = T/np.linalg.norm(T,axis=1,keepdims=True)
    logits = I @ T.T / temp
    e = np.exp(logits - logits.max(1,keepdims=True))
    p = np.diag(e)/e.sum(1)
    return -np.log(p).mean()
print(round(clip_loss(np.random.randn(16,64), np.random.randn(16,64)), 2))

In practice

Trained on ~400M web image–text pairs, CLIP yields transferable features and zero-shot recognition. Big batches supply many negatives; τ is learned. The recipe generalizes to audio–text, video–text, and beyond.

Zero-shot classification

classify by comparing to label text

The idea

Because images and label-words share a space, you can classify without training on the classes: embed the image, embed each candidate label as a prompt like “a photo of a {class},” and pick the closest. New classes need only new text — no labeled images, no fine-tuning. Pick an image and read off the similarities.

Picture it

In the companion lab you can choose an image; its embedding is compared to class prompts — the closest wins (no training).

The math

pred = argmax_c cos(f(image), g("a photo of a "+c)). Softmax over class similarities → probabilities.

Derivation

The contrastive pretraining already placed images near their describing text, so an unseen class is recognizable as long as you can name it. Prompt wording matters (“prompt engineering”); ensembling templates improves accuracy.

Code

python

import numpy as np
# Zero-shot classification: compare an image to text prompts of each class
def norm(x): return x/np.linalg.norm(x)
img = np.random.randn(64)
classes = {c: np.random.randn(64) for c in ['cat','dog','car']}
scores = {c: norm(img) @ norm(e) for c, e in classes.items()}
print("prediction:", max(scores, key=scores.get))

In practice

CLIP’s headline ability: competitive ImageNet accuracy with zero training images. Enables open-vocabulary detection/segmentation and flexible filtering. Limits: fine-grained classes, counting, and concepts rare on the web.

The modality gap

aligned, yet not overlapping

The idea

A surprising wrinkle: even after contrastive training, image embeddings and text embeddings sit in two separate regions of the shared space — a persistent “modality gap.” Pairs are still nearest within their offset, so retrieval works, but the two clouds don’t overlap. Toggle the gap to see why it doesn’t break matching.

Picture it

In the companion lab you can toggle the modality gap. Even offset, each text’s nearest text-neighbor structure mirrors the images’.

The math

Embeddings live on a hypersphere; image and text occupy distinct cones separated by a roughly constant offset vector.

Derivation

Initialization places each encoder’s outputs in a narrow cone, and the contrastive loss with finite temperature preserves a gap rather than fully merging the clouds. Relative geometry within each modality still matches, so cross-modal similarity ranking is unaffected.

Code

python

import numpy as np
# Modality gap: image and text embeddings cluster in separate cones
img  = np.random.randn(100, 64) + 5            # shifted clusters
text = np.random.randn(100, 64) - 5
gap = np.linalg.norm(img.mean(0) - text.mean(0))
print("distance between modality centroids:", round(gap, 1))

In practice

Documented across CLIP-style models; it matters for tasks that mix modalities directly (e.g., embedding arithmetic, fusion). Closing or modeling the gap is an active research thread; for retrieval it’s usually benign.

Long-context & state-space models

The quadratic cost of attention

every token attends to every token

The idea

Self-attention compares every token with every other, so an n-token sequence builds an n×n score matrix — compute and memory grow like n². Double the context and the cost quadruples. That wall is why long-context is hard and why cheaper alternatives exist. Slide the sequence length and watch the matrix (and the cost) balloon.

Picture it

In the companion lab you can increase the sequence length; the attention matrix and its cost grow quadratically.

The math

scores = QKᵀ ∈ ℝ^{n×n}; softmax; ×V. Time and memory = O(n²·d).

Derivation

Forming all pairwise dot products is n² entries, each O(d) — and storing the matrix is n² memory. For sequences of thousands of tokens this dominates, capping context length and inflating training/inference cost.

Code

python

# Attention cost grows quadratically with sequence length n
for n in (1_000, 4_000, 16_000, 64_000):
    ops = n*n                                   # n^2 attention scores
    print(f"n={n:>6}: {ops:.1e} score computations")

In practice

The central bottleneck for long context. Mitigations: FlashAttention (same cost, far less memory via tiling), sparse/sliding-window attention, and the linear/SSM approaches in the next cards. Driving force behind long-context research.

The KV cache

remember the past to decode fast

The idea

When a model generates text token by token, it would be wasteful to recompute attention over the whole prefix each step. Instead it caches the keys and values of every past token and reuses them. That makes each new token cheap — but the cache grows linearly with context, and for long sequences it becomes the memory bottleneck. Slide the position and watch it fill.

Picture it

In the companion lab you can advance the decoding position; the K/V cache grows with every generated token.

The math

Store K,V for all past tokens. New token: q·K_cache, then ·V_cache. Cache memory = O(n·d·layers); per-step cost = O(n).

Derivation

Past keys/values don’t change as you append tokens, so caching them turns an O(n²) re-encode into O(n) per step. The price is memory that scales with context length × layers × heads — often gigabytes for long prompts.

Code

python

# KV cache: store past keys/values so each new token is O(n), not O(n^2)
seq, d = 0, 64
cache_k, cache_v = [], []
def step(k, v):
    cache_k.append(k); cache_v.append(v)        # append once
    return len(cache_k)                          # attend over all cached
for t in range(5): step([0]*d, [0]*d)
print("cached tokens:", len(cache_k), "(memory grows linearly)")

In practice

The KV cache dominates LLM-serving memory and limits batch size / context. Fixes: multi-query / grouped-query attention (share K,V across heads), paged attention (vLLM), quantized caches, and sliding-window attention.

Linear attention

reorder the matrix products

The idea

The quadratic cost comes from forming QKᵀ first. Replace softmax with a kernel feature map and you can reassociate the product: compute KᵀV once (a small d×d summary) and reuse it for every query — turning O(n²) into O(n). It’s an approximation, but it scales. Slide n and watch linear overtake quadratic.

Picture it

In the companion lab you can grow the sequence; the linear path (green) beats quadratic (pink) as n increases.

The math

softmax: (φ(Q)φ(K)ᵀ)V — O(n²). linear: φ(Q)(φ(K)ᵀV) — O(n·d²). Associativity is the trick.

Derivation

With a feature map φ replacing exp, attention becomes a sum of outer products φ(k)vᵀ — a d×d running state. Each query reads that state in O(d²), independent of n, and the state updates incrementally (a linear RNN).

Code

python

import numpy as np
# Linear attention: reorder (Q K^T) V as Q (K^T V) -> O(n) in sequence length
Q = np.random.randn(100, 16); K = np.random.randn(100, 16); V = np.random.randn(100, 16)
KV = K.T @ V                                     # (16,16) once
out = Q @ KV                                      # cheap, no n x n matrix
print(out.shape)                                  # (100, 16)

In practice

Performer, Linear Transformer, RWKV. Trades exact softmax for linear scaling and a recurrent form (great for streaming). Quality can lag full attention; hybrid and gated variants narrow the gap — and lead toward state-space models.

State-space models

a linear recurrence with long memory

The idea

State-space models process a sequence as a linear recurrence: a hidden state carries information forward, decayed and updated at each step. That’s O(n) and streamable like an RNN — yet it can be unrolled into a parallel scan for fast training. A decay parameter sets the memory horizon; selective SSMs (Mamba) make it input-dependent. Tune the decay.

Picture it

In the companion lab you can set the state decay; closer to 1 means the input pulse is remembered far longer.

The math

h_t = A·h_{t−1} + B·x_t, y_t = C·h_t. O(n) recurrence ≡ a convolution → parallel scan for training.

Derivation

Unrolling gives y_t = Σ_k C Aᵏ B x_{t−k} — a convolution with an exponentially decaying kernel; |A|→ 1 lengthens memory (horizon ≈ 1/(1−a)). Linearity lets an associative scan compute all states in parallel, unlike a nonlinear RNN.

Code

python

import numpy as np
# State-space model: a linear recurrence h_t = A h_{t-1} + B x_t
A, B, C = 0.9, 1.0, 1.0
h, ys = 0.0, []
for x in np.random.randn(20):
    h = A*h + B*x; ys.append(C*h)                # scan; O(n), long memory via A
print("output length:", len(ys))

In practice

S4 and Mamba reach O(n) with long-range memory rivaling attention. “Selective” SSMs make A,B,C input-dependent (content-aware gating). Strong on very long sequences (audio, genomics, long text) where attention’s n² is prohibitive.

Neural ODEs & continuous models

Residual blocks as ODE steps

depth is integration time

The idea

A residual block computes x + f(x) — which is exactly one Euler step of an ODE, dx/dt = f(x), with step size 1. Stacking blocks integrates that ODE forward, and a point’s journey through the network is a trajectory in a vector field. More blocks = finer integration. Add layers and watch the path smooth onto the true flow.

Picture it

In the companion lab you can add residual blocks; each is an Euler step, so more blocks trace the ODE trajectory more finely.

The math

x_{l+1} = x_l + f(x_l) ≈ one Euler step of dx/dt = f(x). L blocks integrate over time t = L·h.

Derivation

Rewriting the residual update as (x_{l+1}−x_l)/h = f(x_l) is the forward-Euler discretization of an ODE. As the number of blocks grows (step → 0), the discrete network limits to a continuous-depth model.

Code

python

import numpy as np
# A residual net is Euler's method: x_{t+1} = x_t + f(x_t)  (step size 1)
def f(x): return -0.5*x
x = np.array([2.0])
for _ in range(5): x = x + f(x)                  # discrete residual steps
print(np.round(x, 3))

In practice

This ODE view of ResNets motivated Neural ODEs and explained why very deep residual nets train stably. It also links architecture choices to numerical-integration stability.

Neural ODEs: adaptive solving

let the solver choose the depth

The idea

A Neural ODE replaces discrete layers with a learned vector field f(h,t) and hands the forward pass to an ODE solver. “Depth” becomes integration time, and an adaptive solver spends function evaluations where the dynamics are stiff. Coarse solving is inaccurate; refine it and the trajectory error collapses toward the true solution.

Picture it

In the companion lab you can increase solver steps (function evaluations); the trajectory error vs the true solution drops.

The math

h(T) = h(0) + ∫₀ᵀ f(h(t), t) dt, solved numerically. Adaptive solvers pick step sizes to hit an error tolerance.

Derivation

The model output is the ODE solution at time T. More function evaluations shrink the discretization error (Euler is O(h); Runge–Kutta higher order). Continuous depth means you can query any t and adjust accuracy at inference.

Code

python

import numpy as np
# Neural ODE: take the limit of infinitely many tiny residual steps
def f(x, t): return -0.5*x
x, t, dt = np.array([2.0]), 0.0, 0.01
for _ in range(500):
    x = x + dt*f(x, t); t += dt                   # ODE solver (Euler)
print(np.round(x, 3))                              # ~ 2 e^{-0.5 t}

In practice

Neural ODEs trade fixed layers for an accuracy–compute knob (number of function evals), handle irregularly-sampled time series naturally, and underpin continuous normalizing flows and some diffusion samplers.

The adjoint method

gradients at O(1) memory

The idea

Backprop normally stores every intermediate activation — memory that grows with depth. Neural ODEs instead solve a second ODE backward in time (the adjoint) to get gradients, reconstructing states on the fly. Memory becomes constant, independent of how many solver steps the forward pass took. Slide the depth and compare.

Picture it

In the companion lab you can increase integration depth; standard backprop memory grows, the adjoint stays flat.

The math

Define a(t) = ∂L/∂h(t). It obeys da/dt = −aᵀ ∂f/∂h, solved backward from T → 0. Memory: O(1) in the number of steps.

Derivation

The adjoint ODE propagates the loss sensitivity backward; integrating it (alongside the re-derived state) yields parameter gradients without storing the forward trajectory — trading recomputation for memory.

Code

python

# Adjoint method: get gradients by solving a second ODE backward in time
# -> constant memory regardless of the number of solver steps
# d a/dt = -a^T df/dx ;  integrate backward to accumulate parameter grads
print("memory cost: O(1) in solver steps, vs O(steps) for plain backprop")

In practice

Constant memory lets Neural ODEs go “arbitrarily deep.” Caveat: reverse reconstruction can be numerically delicate; checkpointing and interpolated adjoints improve stability and accuracy in practice.

Continuous normalizing flows

reshape a distribution with a flow

The idea

A continuous normalizing flow turns a simple density into a complex one by flowing samples along an ODE. Track how the log-density changes — it falls by the divergence of the velocity field — and you get exact likelihoods of a flexible distribution. Slide time and watch a Gaussian blob morph into a structured target.

Picture it

In the companion lab you can advance time t; samples flow from the base Gaussian to the target density.

The math

dz/dt = f(z,t); d log p(z)/dt = − tr(∂f/∂z) (instantaneous change of variables). log p(x) = log p(z₀) − ∫ tr(∂f/∂z) dt.

Derivation

The continuous change-of-variables formula replaces a flow’s log-det-Jacobian with a time integral of the trace (divergence) — cheap to estimate. Integrating it alongside z gives exact sample log-likelihoods.

Code

python

import numpy as np
# Continuous normalizing flow: track how log-density changes along the flow
# d log p / dt = -trace(df/dx)
def trace_jac(x): return -0.5*len(x)             # for f(x) = -0.5 x
logp = 0.0
for _ in range(100): logp += -trace_jac(np.zeros(2))*0.01
print("accumulated log-density change:", round(logp, 3))

In practice

CNFs (FFJORD) are exact-likelihood generative models; the modern descendant, flow matching, trains the velocity field directly and powers state-of-the-art image/audio generators alongside diffusion.

Score-based & diffusion (deep dive)

Score functions

the gradient that points to the data

The idea

Score-based models learn the gradient of the log-density, ∇ log p(x) — the score — a vector field that everywhere points toward higher data density. You never need the normalizing constant, and following the field walks you onto the data manifold. Drag the point and watch its score arrow swing toward the nearest mode.

Picture it

In the companion lab you can drag the point; its score vector ∇log p always points toward higher density (the modes).

The math

score(x) = ∇ₓ log p(x). Learned by denoising score matching — no normalizing constant needed.

Derivation

Taking the gradient of log p cancels the intractable partition function, leaving only the shape of the density. Training a network to denoise data at many noise levels provably estimates the score (Vincent’s identity).

Code

python

import numpy as np
# Score = gradient of log-density; points toward higher probability
mu, sigma = 0.0, 1.0
score = lambda x: -(x - mu)/sigma**2             # for a Gaussian
print([round(score(x), 2) for x in (-2, 0, 2)])  # pulls toward the mean

In practice

The score is the backbone of diffusion / score-based generative models: estimate it, then integrate a reverse SDE/ODE to turn noise into samples. The same object connects diffusion, energy-based models, and flow matching.

Langevin sampling

noise + score = samples

The idea

Given the score, you can sample without ever writing down p: repeatedly step a little way along the score and add a dash of noise. This Langevin dynamics settles into the data distribution — the noise keeps it from collapsing to a single mode. Step it and watch random dust drift into the density’s basins.

Picture it

In the companion lab you can run Langevin steps; samples climb the score and the noise spreads them across both modes.

The math

x ← x + ε·∇log p(x) + √(2ε)·η, η ~ N(0,I). As ε → 0, the stationary distribution is p.

Derivation

The drift ε·score pulls samples toward high density; the √(2ε) noise injects exploration so the chain matches p rather than just its peak. It’s gradient ascent on log p with calibrated noise — the sampling analogue of SGD.

Code

python

import numpy as np
# Langevin sampling: follow the score, add noise -> samples from p
score = lambda x: -x                             # target N(0,1)
x, eps, samples = 5.0, 0.1, []
for _ in range(5000):
    x = x + eps*score(x) + np.sqrt(2*eps)*np.random.randn()
    samples.append(x)
s = np.array(samples[500:])
print(round(s.mean(), 2), round(s.std(), 2))     # ~0, ~1

In practice

Annealed Langevin (decreasing noise levels) was the first score-based sampler. Modern diffusion replaces it with a reverse SDE/ODE solve, but the principle — follow the score, control the noise — is the same.

Classifier-free guidance

trade diversity for prompt fidelity

The idea

To make a conditional generator follow its prompt more strongly, mix two scores: the conditional and the unconditional. Push past the conditional (guidance scale w > 1) and samples lock onto the prompt — sharper and more on-target, but less diverse. Slide w and watch mass concentrate on the conditioned mode as variety shrinks.

Picture it

In the companion lab you can raise the guidance scale; the effective density sharpens onto the conditioned mode (fidelity up, diversity down).

The math

ẽ = ε_uncond + w·(ε_cond − ε_uncond). Effective density ∝ p_uncond(x)·[p(c|x)]^w.

Derivation

The guided score extrapolates from the unconditional toward the conditional; raising w sharpens the implied p(c|x)^w, concentrating probability where the condition is most satisfied. No separate classifier is needed — hence “classifier-free.”

Code

python

import numpy as np
# Classifier-free guidance: extrapolate from uncond toward cond score
def guided(eps_uncond, eps_cond, w): 
    return eps_uncond + w*(eps_cond - eps_uncond)
print(guided(0.2, 0.8, w=3.0))   # higher w -> stronger prompt adherence

In practice

The standard knob in text-to-image/audio diffusion: higher w → closer to the prompt but less varied and sometimes over-saturated. Tuning w (and its schedule) is central to sample quality.

Schedules & samplers

quality vs number of steps

The idea

Diffusion generates by reversing noise over many steps — but how many? More denoising steps mean higher quality and higher cost. Deterministic samplers (DDIM) get usable results in a few steps; stochastic ones (DDPM) need more but can be crisper. Slide the step budget and watch the reconstruction error fall.

Picture it

In the companion lab you can change the number of denoising steps; few steps are coarse, more steps refine the sample toward the target.

The math

Reverse process: x_{t−1} = denoise(x_t, t) over T steps. Error ≈ O(1/steps); DDIM = deterministic ODE, DDPM = stochastic SDE.

Derivation

Each step integrates a slice of the reverse SDE/ODE; finer discretization lowers truncation error, like an ODE solver. DDIM’s deterministic update lets large steps stay accurate, so far fewer evaluations suffice.

Code

python

import numpy as np
# Fewer denoising steps = faster but coarser; error ~ 1/steps
for steps in (5, 20, 100):
    err = 1/np.sqrt(steps)
    print(f"{steps:>3} steps  ->  approx error {err:.3f}")

In practice

Step count is the main quality–latency dial. Fast samplers (DDIM, DPM-Solver) and distillation (consistency models, few-step LCM) cut hundreds of steps to a handful for real-time generation.

Agents & tool use

Reasoning + acting (ReAct)

think, act, observe, repeat

The idea

An agent isn’t just a one-shot answer — it loops: think about what to do, take an action, observe the result, and repeat until it can answer. Interleaving reasoning with acting lets the model gather facts it doesn’t know and correct course. Step the loop and watch a question resolve through thought–action–observation.

Picture it

In the companion lab you can step the agent loop; each cycle adds a thought, an action, and an observation until it answers.

The math

Loop: thought → action → observation → (repeat) → answer. The transcript is fed back as context each step.

Derivation

Reasoning traces make the next action better; actions ground the reasoning in real observations, reducing hallucination. The growing transcript is the agent’s working memory — each step conditions on the whole history.

Code

python

# ReAct loop: interleave reasoning (thought) with acting (tool calls)
trace = []
def step(thought, action, observation):
    trace.append(("thought", thought)); trace.append(("action", action))
    trace.append(("observation", observation))
step("need the capital", 'search("capital of France")', "Paris")
print(trace[-1])      # the transcript is fed back as context each turn

In practice

ReAct and its descendants (function-calling agents, Reflexion) underpin tool-using assistants. Failure modes: loops, context-window overflow, and compounding errors — handled by limits, planning, and memory (later cards).

Tool calling

offload what the model can’t do

The idea

Language models are bad at arithmetic, fresh facts, and precise lookups — so let them call tools. The model emits a structured call (function name + arguments), an external system runs it, and the result is fed back to continue generation. Toggle the tool and watch a calculation flip from a confident guess to the correct answer.

Picture it

In the companion lab you can toggle tool access for a hard arithmetic query; the tool returns the exact result the model would otherwise approximate.

The math

Model → {name, args} (structured/JSON) → execute → result → model continues. Trained via function-calling fine-tuning.

Derivation

The model learns when and how to call a tool, not to compute the answer itself. This grounds outputs in reliable systems (calculators, search, code, databases) and extends capability beyond what weights memorized.

Code

python

# Tool calling: model emits a structured call; system runs it; result returns
import json
def calculator(expr): return eval(expr)            # the external tool
call = {"name": "calculator", "args": {"expr": "47291 * 8633"}}
result = calculator(call["args"]["expr"])
print(json.dumps({"tool_result": result}))         # fed back to the model

In practice

Function/tool calling powers assistants, retrieval, code execution, and API agents. Reliability hinges on schema adherence, argument correctness, and handling tool errors gracefully.

Planning & search

branch, evaluate, and prune

The idea

For multi-step problems, a single greedy chain of thought can wander off. Tree-of-thoughts lets the agent branch — explore several reasoning paths, score partial states, and pursue the most promising while pruning dead ends. Expand the tree and watch the search find a high-value path.

Picture it

In the companion lab you can expand the search tree; nodes are scored and the best path (green) is pursued while weak branches are pruned.

The math

Search over a tree of partial solutions; value each node v(s); expand argmax; backtrack on dead ends (BFS/DFS/beam, or MCTS).

Derivation

Branching converts one fragile chain into a search problem where a value/verifier guides exploration. Pruning low-value branches focuses compute, trading more inference for higher reliability on hard, multi-step tasks.

Code

python

import numpy as np
# Tree-of-thoughts: branch, score partial states, expand the best
frontier = [("start", 0.5)]
for depth in range(3):
    node, score = max(frontier, key=lambda x: x[1])      # best-first
    frontier = [(node+f"->{i}", score + np.random.rand()*0.3) for i in range(2)]
print("best path:", max(frontier, key=lambda x: x[1])[0])

In practice

Tree-of-Thoughts, self-consistency (sample many chains, vote), and MCTS-style search lift reasoning on math/planning benchmarks — at the cost of many more model calls per answer.

Memory & retrieval

recall the relevant, ignore the rest

The idea

Context windows are finite, so agents can’t keep everything in view. Instead they store information as embeddings and retrieve the few most relevant items for the current query — retrieval-augmented memory. Slide the query and watch the top-k relevant memories light up while the rest stay dormant.

Picture it

In the companion lab you can move the query; the k nearest memories (by embedding similarity) are retrieved into context.

The math

Retrieve top-k by similarity: argmax_k cos(embed(query), embed(memoryᵢ)). Append retrieved items to the prompt.

Derivation

Embedding memories into a vector space turns recall into nearest-neighbor search, so relevant context is injected on demand instead of held in the window. This scales memory beyond the context limit and keeps prompts focused.

Code

python

import numpy as np
# Retrieval-augmented memory: fetch the top-k relevant items by similarity
mem = np.random.randn(100, 64)                     # stored memories
query = np.random.randn(64)
sims = mem @ query / (np.linalg.norm(mem,axis=1)*np.linalg.norm(query))
topk = sims.argsort()[::-1][:3]
print("retrieved memory ids:", topk)               # injected into context

In practice

Retrieval-augmented generation (RAG) and long-term agent memory: vector databases, chunking, and re-ranking. Quality depends on embeddings, chunking, and retrieval precision — garbage retrieved is garbage generated.

World models

Learned dynamics models

predict the next state, then roll forward

The idea

A world model learns the environment’s dynamics: given the current state and an action, predict the next state. Chain those predictions and you can roll the future forward in your head — imagined trajectories, no real interaction needed. Slide the horizon and watch a rollout trace the predicted path.

Picture it

In the companion lab you can extend the rollout horizon; the model chains one-step predictions into an imagined trajectory.

The math

ŝ_{t+1} = fθ(s_t, a_t). Rollout: apply f repeatedly from s₀. Trained to match observed transitions.

Derivation

Fitting f to (state, action → next state) tuples gives a simulator. Multi-step rollouts compose single-step predictions, enabling planning and data-efficient RL — but errors accumulate (a later card).

Code

python

import numpy as np
# Learned dynamics: predict next state from state+action, then roll forward
def f(s, a): return np.array([s[0] + 0.1*s[1], s[1] - 0.1*s[0] + a])
s = np.array([1.0, 0.0]); traj = [s]
for t in range(20): s = f(s, a=0.0); traj.append(s)   # imagined rollout
print("rollout length:", len(traj))

In practice

Model-based RL (Dreamer, MuZero, world models for games/robotics) learns dynamics to plan or generate synthetic experience, often far more sample-efficient than model-free RL.

Latent world models

imagine in a compressed space

The idea

Pixels are high-dimensional and noisy; predicting them step by step is wasteful. Latent world models encode observations into a compact latent state, predict there, and decode only when needed. Planning happens in the small latent space. Step the latent rollout and watch the decoded prediction stay coherent.

Picture it

In the companion lab you can step the latent rollout; dynamics are predicted in the compact code, then decoded back to observations.

The math

encode: z_t = e(o_t). predict: ẑ_{t+1} = f(z_t, a_t). decode: ô = d(ẑ). Loss = reconstruction + latent-prediction + reward.

Derivation

A low-dimensional latent captures the task-relevant structure, so rollouts are cheap and stable; decoding is only needed for visualization or reward. This is the recipe behind PlaNet/Dreamer’s recurrent state-space model.

Code

python

import numpy as np
# Latent world model: encode obs -> z, predict in latent, decode only to view
encode = lambda o: o @ np.random.randn(64, 8)      # to compact latent
predict = lambda z: z @ np.random.randn(8, 8)       # cheap latent dynamics
z = encode(np.random.randn(64))
for _ in range(10): z = predict(z)                  # plan in latent space
print("latent dim:", z.shape)

In practice

Latent dynamics make pixel-based control tractable and sample-efficient; the agent learns and plans entirely in latent imagination, decoding only to act or inspect.

Planning in imagination

simulate futures, pick the best

The idea

With a world model in hand, an agent can plan by simulation: sample many candidate action sequences, roll each forward in the model, score the imagined outcomes, and execute the best plan’s first action — then replan. Step it and watch many imagined futures get scored and the best one chosen.

Picture it

In the companion lab you can sample candidate plans; each is rolled out in imagination and scored — the best (green) reaches the goal.

The math

For each candidate action sequence a_{1:H}: rollout in f, score Σ reward. Execute argmax’s first action (model-predictive control).

Derivation

The model turns control into search over imagined trajectories — no real-world trials needed. Re-planning each step (receding horizon) corrects for model error and changing conditions.

Code

python

import numpy as np
# Plan by imagination: sample action sequences, score rollouts, pick the best
def rollout_return(actions): return -np.sum((np.cumsum(actions) - 5)**2)
candidates = [np.random.randn(10) for _ in range(50)]
best = max(candidates, key=rollout_return)
print("execute first action of best plan:", round(best[0], 2))

In practice

MuZero plans with a learned model + MCTS; Dreamer learns a policy from imagined rollouts; CEM/MPC planners optimize action sequences. The model’s accuracy bounds the plan’s quality.

Compounding model error

small errors, long horizons, big drift

The idea

The catch with imagination: tiny one-step prediction errors compound. Each rollout step feeds the model its own slightly-wrong output, so the imagined trajectory drifts from reality — exponentially in the worst case. That caps how far ahead you can plan. Raise the per-step error and watch the prediction peel away from the truth.

Picture it

In the companion lab you can increase the per-step error; the predicted rollout diverges from the true trajectory as the horizon grows.

The math

If each step adds error ε and dynamics have Lipschitz constant L > 1, horizon-H error ≈ ε·(L^H − 1)/(L − 1) — grows fast.

Derivation

Feeding predictions back as inputs means errors propagate and amplify through the dynamics. Even a near-perfect one-step model can produce useless long rollouts — the core difficulty of model-based RL.

Code

python

import numpy as np
# Compounding error: per-step error grows over the rollout horizon
eps, L = 0.02, 1.1                                  # step error, expansion rate
for H in (5, 10, 20, 40):
    err = eps*(L**H - 1)/(L - 1)
    print(f"horizon {H:>2}: drift ~ {err:.2f}")     # grows fast

In practice

Mitigations: short rollout horizons, frequent replanning, ensembles to flag uncertainty, and training on the model’s own rollouts (to match the test distribution). Why model-based methods favor short imagined horizons.

Mixture-of-depths & adaptive compute

Early exit / adaptive depth

easy inputs leave early

The idea

Not every input needs the full network. With early-exit heads, an easy example can be classified after a few layers and skip the rest, while hard ones run deep. The model spends compute in proportion to difficulty. Slide the difficulty and watch the exit layer — and the FLOPs — move.

Picture it

In the companion lab you can change input difficulty; easy inputs exit at a shallow layer (less compute), hard ones run to the end.

The math

At layer ℓ, if confidence(ℓ) > τ, exit and return. Expected cost = Σ P(exit at ℓ)·cost(ℓ).

Derivation

Adding classifier heads at intermediate layers lets the model halt once it’s confident. Average compute drops because most inputs are easy — the long tail of hard cases still gets full depth.

Code

python

import numpy as np
# Early exit: stop once an intermediate head is confident enough
def confidence(layer, difficulty): return min(1.0, 0.3 + 0.1*layer - difficulty)
for difficulty in (0.1, 0.6):
    layer = next(l for l in range(1, 13) if confidence(l, difficulty) > 0.7)
    print(f"difficulty {difficulty}: exits at layer {layer}")

In practice

BranchyNet, DeeBERT, and CALM speed up inference by exiting early on easy tokens/inputs. Trade-off: extra heads and a confidence threshold to tune; accuracy can dip if exits are too aggressive.

Mixture-of-depths routing

spend compute on the tokens that matter

The idea

Mixture-of-depths gives each transformer block a token budget: a router picks which tokens actually pass through the heavy computation, and the rest take a cheap residual shortcut. Total compute is capped regardless of sequence length. Slide the capacity and watch tokens get routed in or skipped.

Picture it

In the companion lab you can set the per-block token capacity; only the top-scoring tokens are processed, the others skip the block.

The math

Router scores each token; top-k (k = capacity·n) go through the block, others use identity. FLOPs ∝ capacity, not n.

Derivation

A fixed capacity makes per-block cost independent of how many tokens are “interesting,” giving a static compute graph (good for hardware) while still letting the model allocate depth to the tokens the router deems important.

Code

python

import numpy as np
# Mixture-of-depths: a router sends only top-k tokens through the heavy block
scores = np.random.rand(24)                         # router score per token
capacity = 0.5
k = int(capacity*len(scores))
process = set(scores.argsort()[::-1][:k])            # rest skip the block
print(f"{k}/{len(scores)} tokens processed -> {capacity*100:.0f}% FLOPs")

In practice

Mixture-of-Depths (and the related Mixture-of-Experts) cut FLOPs at fixed quality by routing. MoD keeps a predictable compute budget; the router is trained jointly with the model.

Adaptive computation time

think longer when it’s hard

The idea

Some inputs need more thinking steps than others. A halting mechanism lets a network decide, per input, how many recurrent steps to take: it accumulates a halting probability and stops once it crosses a threshold, paying a small “ponder” penalty for extra steps. Slide difficulty and watch the step count adapt.

Picture it

In the companion lab you can change difficulty; the halting unit accumulates probability and stops later for harder inputs.

The math

Each step emits halting prob h_t; stop when Σ h_t ≥ 1−ε. Loss adds a ponder cost ∝ number of steps.

Derivation

Making the step count data-dependent lets the model allocate computation to difficulty, while the ponder penalty discourages thinking forever. It’s a differentiable approximation to “keep computing until confident.”

Code

python

import numpy as np
# Adaptive computation time: keep stepping until halting prob crosses 1
def steps_for(difficulty):
    cum, t = 0.0, 0
    while cum < 1.0 and t < 12:
        cum += max(0.05, 0.5 - difficulty); t += 1
    return t
print("easy:", steps_for(0.1), " hard:", steps_for(0.45))

In practice

Adaptive Computation Time (Graves), PonderNet, and Universal Transformers vary depth per token/input. Related in spirit to chain-of-thought length and test-time compute scaling.

Speculative decoding

draft fast, verify in parallel

The idea

Big models decode slowly, one token at a time. Speculative decoding uses a small, fast draft model to guess several tokens ahead, then the big model verifies them all in a single parallel pass — accepting the correct prefix and correcting the first mistake. Slide the acceptance rate and watch the speedup grow.

Picture it

In the companion lab you can change how often the draft’s guesses are accepted; higher acceptance → more tokens per verification → faster.

The math

Draft proposes k tokens; target verifies in one pass; accept the longest correct prefix. Expected tokens/step rises with acceptance — output identical to the target alone.

Derivation

Verification is parallel (one forward pass over k tokens), and a rejection-sampling correction guarantees the same distribution as the target model. So speed improves with no quality loss — the gain scales with draft accuracy.

Code

python

import numpy as np
# Speculative decoding: a draft model proposes k tokens, target verifies in 1 pass
k, accept_rate = 5, 0.7
accepted = np.random.binomial(k, accept_rate)        # longest correct prefix
tokens_per_pass = accepted + 1                        # +1 correction
print(f"{tokens_per_pass} tokens/verification (~{tokens_per_pass:.1f}x speedup, lossless)")

In practice

Speculative / Medusa / lookahead decoding give 2–3× lossless speedups in LLM serving. The draft can be a small model, the target’s own early layers, or learned heads.

Knowledge editing

Locating facts (causal tracing)

find where a fact lives

The idea

Where does a model “store” that Paris is in France? Causal tracing finds out: corrupt the input, then restore each hidden state one at a time and see which restoration recovers the correct answer. The layers that matter light up — usually mid-layer MLPs at the subject token. Step the trace and watch the responsible site emerge.

Picture it

In the companion lab you can step the trace; restoring each layer’s activation reveals which sites carry the fact (brighter = more causal effect).

The math

Effect(ℓ,t) = P(correct | restore state at layer ℓ, token t) − P(correct | corrupted). Peak = causal site.

Derivation

Corrupting the subject embedding breaks the prediction; restoring a single internal state and measuring recovery isolates which states cause the correct output — a causal, not merely correlational, localization of the fact.

Code

python

import numpy as np
# Causal tracing: restore each hidden state, see which recovers the answer
effect = np.zeros((10, 6))                            # layers x tokens
effect[5, 2] = 0.9                                     # a causal mid-layer site
l, t = np.unravel_index(effect.argmax(), effect.shape)
print(f"fact localized at layer {l}, token {t}")

In practice

ROME’s causal tracing showed factual associations concentrate in mid-layer MLPs at the last subject token — motivating targeted edits (next card). A pillar of mechanistic interpretability.

Editing a fact (ROME)

rewrite a weight, change a belief

The idea

Once you know where a fact lives, you can rewrite it with a targeted weight update — no retraining. ROME treats an MLP layer as a key–value store and applies a rank-one edit so the subject’s key now maps to a new value (“the Eiffel Tower is in Rome”). Toggle the edit and watch the target output flip while neighbors stay put.

Picture it

In the companion lab you can apply the rank-one edit; the target association changes while an unrelated fact is left intact.

The math

Treat MLP weight W as key→value memory. Solve for Δ (rank-one) s.t. (W+Δ)k* = v*, minimizing change elsewhere.

Derivation

A rank-one update Δ = (v* − Wk*)·(C⁻¹k*)ᵀ / (k*ᵀC⁻¹k*) sets the new association at key k* while least-squares-minimizing disruption to other keys (C is the covariance of keys). Surgical, closed-form, fast.

Code

python

import numpy as np
# ROME: rank-one update so a key maps to a new value, minimal collateral
W = np.random.randn(64, 64)                          # an MLP weight (key->value)
k_star = np.random.randn(64); v_star = np.random.randn(64)
delta = np.outer(v_star - W @ k_star, k_star) / (k_star @ k_star)
W_new = W + delta                                     # now W_new @ k_star == v_star
print("edit applied, error:", round(np.linalg.norm(W_new @ k_star - v_star), 6))

In practice

ROME and MEMIT edit thousands of facts directly in weights. Useful for fixing errors and updating stale knowledge, but raises specificity and safety questions (next cards).

Ripple effects & specificity

generalize to paraphrases, not to neighbors

The idea

A good edit must generalize to paraphrases of the same fact yet not leak into unrelated facts — and it should keep logically-entailed consequences consistent. Real edits are imperfect: they can miss paraphrases (poor generalization) or bleed into neighbors (poor specificity). Toggle edit quality and watch the ripple spread.

Picture it

In the companion lab you can switch between a clean and a leaky edit; see which related/unrelated facts change.

The math

Evaluate: efficacy (target), generalization (paraphrases), specificity (unrelated unchanged), and consistency (entailed facts).

Derivation

Because facts share representations, a weight edit can ripple to entangled associations. Quality means hitting the target and its paraphrases while leaving everything else — a tension that grows as you edit many facts at once.

Code

python

import numpy as np
# A good edit generalizes to paraphrases but not to unrelated facts
edited_fact   = "Eiffel Tower -> Rome"
checks = {"paraphrase": True, "entailed": True, "unrelated_1": False, "unrelated_2": False}
# specificity = fraction of unrelated facts left unchanged
specificity = np.mean([not v for k, v in checks.items() if "unrelated" in k])
print("specificity:", specificity)                    # want 1.0

In practice

Benchmarks (CounterFact, RippleEdits) measure these axes. Mass edits can degrade unrelated knowledge or fluency; specificity vs generalization is the core editing trade-off.

Forgetting & continual learning

learning the new without erasing the old

The idea

Train a network on a new task and it tends to overwrite the old one — catastrophic forgetting. Because the same weights serve all tasks, gradient steps for task B trample task A. Countermeasures (rehearsal, regularizing important weights) trade a little new-task speed for retained old-task accuracy. Toggle mitigation and watch the old task survive.

Picture it

In the companion lab you can toggle mitigation; without it, learning task B collapses task-A accuracy — with it, both are retained.

The math

Sequential tasks share θ. EWC adds Σᵢ Fᵢ(θᵢ − θ*ᵢ)² to anchor weights important for old tasks (F = Fisher information).

Derivation

Unconstrained SGD on task B moves θ freely, destroying task-A solutions. Penalizing changes to weights with high Fisher information (those that mattered for A), or replaying A’s data, preserves old performance while still fitting B.

Code

python

import numpy as np
# Catastrophic forgetting: EWC penalizes changing weights important for old tasks
theta_old = np.array([1.0, 2.0])                     # old-task solution
fisher    = np.array([10.0, 0.1])                     # importance per weight
def ewc_penalty(theta, lam=1.0):
    return lam * np.sum(fisher * (theta - theta_old)**2)
print("penalty for moving important weight:", round(ewc_penalty(theta_old + 0.5), 2))

In practice

Continual learning methods: EWC and synaptic intelligence (regularization), experience replay/rehearsal, and parameter isolation (adapters/LoRA per task). Central to lifelong learning and to safe model editing.