The Math Behind Modern AI, Derived and Then Coded in NumPy

Most people learn the maths of machine learning from one of two places, and both let them down in opposite ways.

The first is the academic paper or textbook. It gives you a tidy derivation with indices everywhere, a proof that runs over two pages, and not a single line of code. You can follow each step and still have no idea how to turn it into a vectorised program. The second is the beginner tutorial. It shows you model.fit(X, y), a loss curve, and a cheerful message that says "the maths happens under the hood." It does. You just never get to see it, so the first time a loss turns into nan you have nothing to reason with.

This article sits in the gap. For each idea I derive the result on paper, in plain English, and then write the NumPy that computes it and checks the derivation numerically. If the algebra and the code disagree, the code tells you. That habit, deriving somthing and then testing the derivation against a brute-force number, is the most useful debugging skill I know in this field, and it is the thread running through everything below.

This is Part 5 of the MLHub series, and it is meant to be the piece you read first if you are new. The earlier parts covered agents and MCP, LoRA and QLoRA fine-tuning, fast dataframes with Polars, and Vision Transformers. Each of those leans on the foundations here, and I will point back to them where it helps. LoRA is a low-rank matrix factorisation, which is the SVD section. The attention layer inside a Vision Transformer is softmax(QKᵀ/√d)V, which is the last section but one. Nothing in those articles is magic once you have this one.

What you need is Python, NumPy and the patience to read a few equations slowly. The output blocks show what the snippets print when run in order, and every number I quote is either produced by that code or labelled as an illustration.

Setting up

All the code below runs with NumPy 2.x on any laptop. No GPU, no framework. I use the modern random generator so results are repeatable, and I keep everything in float64 unless I say otherwise, because when you are checking derivations you want rounding error to be the last thing you worry about.

import numpy as np

np.set_printoptions(precision=4, suppress=True)
rng = np.random.default_rng(0)

One convention to fix now, because it saves confusion later. Data matrices have one row per example and one column per feature, so X has shape (N, D): N examples, D features. Weight matrices map input features to output features, so a layer that goes from D inputs to H outputs has W of shape (D, H). Whenever I derive something I will write the shape of every object next to it. A derivation whose shapes do not line up is wrong, and checking shapes is free.

Vectors, matrices and the dense layer

The dot product is a similarity score

A vector is a list of numbers, and in machine learning it usually means "a thing described by features". A house is (area, bedrooms, age). A song is a few hundred learned numbers. A word is a few thousand.

The dot product of two vectors multiplies matching entries and adds them up:

a · b = a₁b₁ + a₂b₂ + ... + a_D b_D = ‖a‖ ‖b‖ cos θ

The second form is the important one. It says the dot product is the product of the two lengths and the cosine of the angle between the vectors. If they point the same way, cos θ is 1 and the score is large and positive. If they are perpendicular it is 0. If they point in opposite directions it is negative. Divide out the lengths and you get cosine similarity, which ignores how long the vectors are and only asks whether they point the same way.

Here is a recommendation engine in one paragraph. Imagine a streaming service that has learned a short "taste" vector for each user and a matching "content" vector for each film, with the three coordinates meaning something like action, romance and documentary. The predicted affinity is simply the dot product. That is the whole scoring step of many recommenders: a matrix product between a user matrix and an item matrix.

def cosine(a, b):
    return a @ b / (np.linalg.norm(a) * np.linalg.norm(b))

u = np.array([1.0, 2.0, 3.0])
v = np.array([2.0, 4.0, 7.0])
w = np.array([-3.0, 1.0, 1.0])
print(u @ v, round(cosine(u, v), 4), round(cosine(u, w), 4))

# Two users, three films: every score at once is one matrix product.
users = np.array([[0.9, 0.1, 0.8],
                  [0.1, 0.9, 0.2]])
films = np.array([[0.8, 0.2, 0.9],
                  [0.1, 0.9, 0.1],
                  [0.5, 0.5, 0.5]])
print(users @ films.T)

31.0 0.9974 0.1612
[[1.46 0.26 0.9 ]
 [0.44 0.84 0.6 ]]

The first user, who leans towards the first coordinate, scores film 0 highest. The second user prefers film 1. Notice films.T: the transpose is what makes the shapes work, (2, 3) @ (3, 3) -> (2, 3). Every row is a user, every column a film.

Matrices are linear maps

A matrix A of shape (m, n) is a function that takes a vector with n numbers and returns a vector with m numbers, and it does so in a very restricted way: the output is linear in the input. Doubling the input doubles the output, and the output for a sum is the sum of the outputs. Rotations, stretches, projections and shears are all matrices. A neural network layer is one matrix followed by a nonlinearity, and the nonlinearity is there precisely because a stack of matrices with nothing in between collapses into a single matrix, since A(Bx) = (AB)x.

A dense layer is X @ W + b

A dense (fully connected) layer takes a batch of N examples with D features and produces H outputs per example. For one example x (a row), output unit j is

out_j = x₁W₁ⱼ + x₂W₂ⱼ + ... + x_D W_Dⱼ + b_j

That is a dot product between x and column j of W, plus a bias. Doing that for all N examples and all H units is exactly the matrix product plus a broadcast bias:

X : (N, D)     W : (D, H)     b : (H,)
X @ W : (N, H)          (X @ W) + b : (N, H)

Let me prove to myself that the matrix form and the loops give the same thing.

X = rng.normal(size=(4, 3))        # 4 examples, 3 features
W = rng.normal(size=(3, 2))        # 3 inputs -> 2 outputs
b = np.array([0.5, -1.0])

out = X @ W + b                    # b is broadcast across the 4 rows

manual = np.empty((4, 2))
for i in range(4):
    for j in range(2):
        manual[i, j] = sum(X[i, k] * W[k, j] for k in range(3)) + b[j]

print(out.shape, np.allclose(out, manual))

(4, 2) True

Two things happened there that deserve a name.

Broadcasting is NumPy's rule for combining arrays of different shapes. It lines the shapes up from the right, and a dimension of size 1 (or a missing dimension) is stretched to match the other one. Here b has shape (2,), which is treated as (1, 2) and copied down four rows. It is convenient, and it is also the source of some of the nastiest silent bugs in numerical code.

Cost. Multiplying an (n, k) matrix by a (k, m) matrix takes n·k·m multiply-and-add pairs, which is about 2nkm floating point operations. That formula explains most of the performance advice you hear. Matrix multiplication is associative, so (AB)x and A(Bx) are equal, but they are not equally cheap:

n = 1000   # A and B are n x n, x is a vector of length n
print("(A@B)@x flops:", 2*n*n*n + 2*n*n)
print("A@(B@x) flops:", 2*n*n + 2*n*n)

(A@B)@x flops: 2002000000
A@(B@x) flops: 4000000

That is roughly a 500 times difference from bracketing alone, and it is the same trick that makes LoRA cheap, which we will come back to.

Linear regression, derived two ways

Scenario: you want to predict a house price from its floor area, number of bedrooms and age. The data below is synthetic, generated from known weights plus noise, so we can see whether the methods recover the truth. The model is a weighted sum plus an offset:

ŷ = w₀ + w₁·area + w₂·bedrooms + w₃·age

Put a column of ones at the front of the feature matrix so the offset w₀ becomes just another weight, and the predictions for the whole dataset are ŷ = Xw, with X of shape (N, D) and w of shape (D,).

We need a way to say how wrong the predictions are. The standard choice is the mean squared error. I put a factor of one half in front, which changes nothing about where the minimum is but removes a 2 from the derivative:

J(w) = (1 / 2N) · ‖Xw − y‖²  =  (1 / 2N) · (Xw − y)ᵀ(Xw − y)

Way one: the normal equations

Set the gradient to zero. Using the rule that the gradient of ‖r‖²/2 with respect to w is Jᵣᵀ r for a residual r = Xw − y whose Jacobian with respect to w is X, we get

∇J(w) = (1 / N) · Xᵀ(Xw − y)

Here is why, spelled out. Write the residual as r_n = Σ_d X_nd w_d − y_n. Then J = (1/2N) Σ_n r_n², so ∂J/∂w_d = (1/N) Σ_n r_n · ∂r_n/∂w_d = (1/N) Σ_n r_n X_nd. That is the d-th entry of Xᵀr / N. Setting it to zero:

Xᵀ(Xw − y) = 0
XᵀX w = Xᵀy                     (the normal equations)
w = (XᵀX)⁻¹ Xᵀ y                (if XᵀX is invertible)

The name "normal" comes from geometry. The condition Xᵀ(Xw − y) = 0 says the residual is perpendicular (normal) to every column of X, which is what least squares means: you project y onto the space spanned by the columns of X.

A practical note. Nobody should compute the explicit inverse. Solving the linear system XᵀX w = Xᵀy with np.linalg.solve is faster and more accurate, and np.linalg.lstsq is better still when the columns are nearly dependent.

Way two: gradient descent

The gradient points uphill, so step the opposite way. With a learning rate η:

w ← w − η · (1/N) · Xᵀ(Xw − y)

Why does this converge, and how big may η be? The loss is a quadratic bowl. Its second derivative (the Hessian) is H = XᵀX / N. For any vector v, vᵀHv = ‖Xv‖² / N ≥ 0, so H is positive semi-definite, and that is exactly the statement that J is convex: it has no local minima other than the global one, and a bowl can only curve upwards.

Let w* be the minimiser, so H w* = Xᵀy / N. Subtract w* from both sides of the update and the gradient becomes H(w − w*):

w_{t+1} − w* = (I − ηH)(w_t − w*)

Now write the error in the eigenvector basis of H, with eigenvalues λ₁ ≥ ... ≥ λ_D ≥ 0. Each component of the error gets multiplied by (1 − ηλᵢ) every step. It shrinks if and only if |1 − ηλᵢ| < 1, that is 0 < ηλᵢ < 2. The most demanding direction is the steepest one, λ₁, which is the largest curvature and usually called L:

0 < η < 2 / L        where L = largest eigenvalue of XᵀX / N

Above 2/L, the steepest direction overshoots by more than it corrects and the loss blows up. Below it, the speed is set by the flattest direction: the ratio λ₁/λ_D, the condition number, decides how many steps you need. That is why standardising features (so the bowl is round rather than a long thin valley) makes gradient descent faster.

Let us check both claims in code, on data with known true weights.

n = 200
area = rng.uniform(50, 200, n)
beds = rng.integers(1, 6, n).astype(float)
age = rng.uniform(0, 50, n)
F = np.column_stack([area, beds, age])
Xs = (F - F.mean(0)) / F.std(0)               # standardise the features
X = np.column_stack([np.ones(n), Xs])         # prepend the bias column
true_w = np.array([300.0, 90.0, 20.0, -30.0])
y = X @ true_w + rng.normal(0, 15, n)         # price in thousands, plus noise

# Way one: normal equations
w_ne = np.linalg.solve(X.T @ X, X.T @ y)

# Way two: gradient descent with a step size chosen from L
H = X.T @ X / n
L = np.linalg.eigvalsh(H).max()
lr = 1.0 / L
w = np.zeros(X.shape[1])
for step in range(2000):
    w -= lr * X.T @ (X @ w - y) / n

print("L =", round(L, 4), " lr =", round(lr, 4))
print("normal eq:", w_ne)
print("grad desc:", w)
print("max abs difference:", np.abs(w - w_ne).max())

L = 1.0759  lr = 0.9294
normal eq: [299.9306  91.7348  20.403  -30.9107]
grad desc: [299.9306  91.7348  20.403  -30.9107]
max abs difference: 7.105427357601002e-14

The two answers agree to numerical precision, and both land close to the true weights we planted (300, 90, 20, −30), off by a little because of the noise. Now the 2/L threshold. I will run 50 steps at several multiples of it and print the final loss:

def final_loss(lr, steps=50):
    w = np.zeros(X.shape[1])
    for _ in range(steps):
        w -= lr * X.T @ (X @ w - y) / n
    return 0.5 * np.mean((X @ w - y) ** 2)

print(f"start (w = 0): loss = {0.5 * np.mean(y ** 2):,.1f}")
for factor in (0.5, 0.99, 1.01, 1.1):
    print(f"lr = {factor} x (2/L): loss = {final_loss(factor * 2 / L):,.1f}")

start (w = 0): loss = 50,221.0
lr = 0.5 x (2/L): loss = 101.0
lr = 0.99 x (2/L): loss = 678.0
lr = 1.01 x (2/L): loss = 31,620.4
lr = 1.1 x (2/L): loss = 360,319,704,568.5

Just under the threshold the loss still falls from its starting value, but slowly, because at 0.99 · 2/L the steepest direction flips sign every step while shrinking by only 2% of its size. Just over it, that direction grows instead, and the loss explodes. The threshold behaves exactly as the argument predicted. For non-quadratic losses (neural networks), L is only a local quantity, but the same intuition holds: too large a step size and you diverge, which is why loss suddenly going to inf in the first few steps is usually a learning-rate problem.

Calculus for learning, and a gradient checker you can reuse

Training a model means finding parameters that make a loss small, and calculus is the tool that tells you which way to move. Four ideas are enough.

Derivative. For a function of one number, f'(x) is the slope: how much f changes per unit change in x, for small changes.

Partial derivative. For a function of several numbers, ∂f/∂x_i is the slope along coordinate i, holding all the others fixed.

Gradient. Stack all the partial derivatives into a vector: ∇f = (∂f/∂x₁, ..., ∂f/∂x_D). Its key property is that it is the direction of steepest ascent. Here is the short argument. Move a tiny distance in a unit direction u. The change in f is the directional derivative, ∇f · u, which equals ‖∇f‖ cos θ with θ the angle between u and the gradient. That is largest when θ = 0, so the best direction to climb is the gradient, and the best direction to descend is its negative. Gradient descent is nothing more than "keep walking downhill in the steepest direction".

Chain rule. If y = f(g(x)), then dy/dx = f'(g(x)) · g'(x). For several variables, sensitivities multiply along each path and add across paths. This one rule is the whole of backpropagation. A network is a long chain of simple functions, and the chain rule lets you compute the derivative of the final loss with respect to every parameter by walking backwards, reusing partial results as you go.

Never trust a hand-derived gradient until you have checked it

The most valuable trick in this article is small. A derivative is, by definition, a limit of finite differences, so you can estimate it with no calculus at all. The central difference is more accurate than the one-sided version:

∂f/∂x_i ≈ ( f(x + ε·eᵢ) − f(x − ε·eᵢ) ) / (2ε)

Here eᵢ is the vector with a 1 in position i and 0 elsewhere. The error shrinks like ε², so with ε = 10⁻⁶ in float64 you typically agree to eight or more digits. It is far too slow to train with (two loss evaluations per parameter), but it is a perfect referee for your analytic gradients. Here is the reusable version:

def numerical_grad(f, x, eps=1e-6):
    """Central finite differences of a scalar function f at array x."""
    x = x.astype(np.float64).copy()
    grad = np.zeros_like(x)
    it = np.nditer(x, flags=["multi_index"])
    for _ in it:
        idx = it.multi_index
        old = x[idx]
        x[idx] = old + eps
        f_plus = f(x)
        x[idx] = old - eps
        f_minus = f(x)
        x[idx] = old
        grad[idx] = (f_plus - f_minus) / (2 * eps)
    return grad

def rel_error(a, b):
    """Largest relative difference between two arrays."""
    return np.max(np.abs(a - b) / np.maximum(1e-8, np.abs(a) + np.abs(b)))

Let us test it on a function whose gradient we can derive by hand with the product rule. Take f(x) = Σ xᵢ² sin(xᵢ). Each term depends only on its own coordinate, so ∂f/∂xᵢ = 2xᵢ sin(xᵢ) + xᵢ² cos(xᵢ). I will also check the steepest-ascent claim by trying ten thousand random unit directions and seeing whether any of them beats the gradient direction.

f = lambda x: np.sum(x**2 * np.sin(x))
x = np.array([0.5, -1.2, 2.0])

analytic = 2 * x * np.sin(x) + x**2 * np.cos(x)
numeric = numerical_grad(f, x)
print("analytic:", analytic)
print("numeric: ", numeric)
print("relative error:", rel_error(analytic, numeric))

best = -np.inf
for _ in range(10_000):
    d = rng.normal(size=3)
    d /= np.linalg.norm(d)                 # a random unit direction
    best = max(best, d @ analytic)         # directional derivative along d
print("gradient norm:", np.linalg.norm(analytic), " best random direction:", best)

analytic: [0.6988 2.7587 1.9726]
numeric:  [0.6988 2.7587 1.9726]
relative error: 7.19552314947556e-11
gradient norm: 3.462640028052271  best random direction: 3.462327669344936

The analytic and numeric gradients match to about ten digits, and the best of ten thousand random directions gets close to the gradient's norm without ever exceeding it. Keep numerical_grad around; we will use it on a full neural network in a moment.

Softmax, cross-entropy and the gradient p − y

Classification needs a way to turn raw scores into probabilities. Suppose a spam filter has three classes (spam, promotion, normal) and the model produces one score per class, called logits, z = (z₁, z₂, z₃). Softmax turns them into a probability distribution:

p_k = exp(z_k) / Σ_j exp(z_j)

Every p_k is positive and they sum to 1. The exponential makes the largest score dominate, and it keeps everything positive.

To train, we need a loss. If the true class is c, the natural loss is the negative log of the probability the model gave to the truth. Write the label as a one-hot vector y (a 1 at position c, zeros elsewhere), and the cross-entropy loss for one example is

L = − Σ_k y_k log p_k  =  − log p_c

Deriving the gradient step by step

We want ∂L/∂z_i for every logit. The trick that makes this painless is to simplify the loss first. Substitute the softmax into log p_k:

log p_k = z_k − log Σ_j exp(z_j)

L = − Σ_k y_k ( z_k − log Σ_j exp(z_j) )
  = − Σ_k y_k z_k  +  ( Σ_k y_k ) · log Σ_j exp(z_j)
  = − Σ_k y_k z_k  +  log Σ_j exp(z_j)          because Σ_k y_k = 1

Now differentiate with respect to z_i. The first term contributes −y_i. The second is a log of a sum, so by the chain rule its derivative is exp(z_i) / Σ_j exp(z_j), which is exactly p_i. Hence

∂L/∂z_i = p_i − y_i        so       ∇_z L = p − y

The gradient of the loss with respect to the logits is the predicted probability vector minus the one-hot target. It is bounded between −1 and 1, it is zero exactly when the prediction is perfect, and it is the reason softmax plus cross-entropy is such a well-behaved pairing. As a cross-check, here is the long way through the softmax Jacobian, ∂p_k/∂z_i = p_k(δ_ki − p_i), where δ is 1 when k = i and 0 otherwise:

∂L/∂z_i = − Σ_k (y_k / p_k) · p_k (δ_ki − p_i)
        = − Σ_k y_k (δ_ki − p_i)
        = − y_i + p_i Σ_k y_k
        = p_i − y_i

Same answer. With two classes and a sigmoid (logistic regression) the identical p − y appears, so logistic regression, softmax regression and the last layer of a classifier network all share one gradient.

For a batch of N examples, the loss is the average of the per-example losses, so each row of the gradient is divided by N. That division is the "forgetting to average" bug I warn about later.

def softmax(z):
    z = z - z.max(axis=1, keepdims=True)     # stability trick, explained below
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)

def cross_entropy(z, y):
    """z: (N, K) logits, y: (N,) integer class labels. Mean loss over the batch."""
    p = softmax(z)
    return -np.mean(np.log(p[np.arange(len(y)), y]))

def cross_entropy_grad(z, y):
    p = softmax(z)
    p[np.arange(len(y)), y] -= 1.0           # p - onehot(y)
    return p / len(y)                        # average over the batch

z = rng.normal(size=(5, 4))
y = rng.integers(0, 4, size=5)
g_analytic = cross_entropy_grad(z, y)
g_numeric = numerical_grad(lambda zz: cross_entropy(zz, y), z)
print("relative error:", rel_error(g_analytic, g_numeric))

relative error: 8.232183701310522e-09

An error near 10⁻⁹ means the derivation and the code agree.

Why cross-entropy, and not something else

There is a principled reason for this loss, and it is maximum likelihood. Suppose the model claims that, given input x, the label is class k with probability p_k(x; θ). If the training examples are independent, the probability the model assigns to the whole dataset is the product of the probabilities it assigned to each true label. Maximum likelihood says choose the parameters that make the observed data as probable as possible. Products are awkward and underflow, so take a log, and flip the sign to get something to minimise:

maximise  Π_n p(y_n | x_n; θ)
⇔ maximise  Σ_n log p(y_n | x_n; θ)
⇔ minimise  − (1/N) Σ_n log p(y_n | x_n; θ)   =   mean cross-entropy

So cross-entropy is the negative log-likelihood of the labels under the model. The same recipe applied to a regression model whose errors are Gaussian gives squared error: −log N(y; ŷ, σ²) is (y − ŷ)²/(2σ²) plus constants. Squared error for regression and cross-entropy for classification are the same principle wearing two different noise assumptions.

Backpropagation for a two-layer network

Now the main event. We will derive backpropagation for a small network, being fussy about shapes, then implement it and prove it right with the gradient checker.

The network has one hidden layer with a ReLU activation, relu(t) = max(0, t), and a softmax output over K classes. For a batch:

X  : (N, D)   inputs
Z1 = X @ W1 + b1        W1: (D, H)   b1: (H,)     Z1: (N, H)
A1 = relu(Z1)                                     A1: (N, H)
Z2 = A1 @ W2 + b2       W2: (H, K)   b2: (K,)     Z2: (N, K)  the logits
P  = softmax(Z2)                                  P : (N, K)
L  = −(1/N) Σ_n log P[n, y_n]                     a single number

Backpropagation runs the chain rule from the loss back to the inputs. At each step we already hold the gradient with respect to a layer's output, and we compute the gradient with respect to its parameters and its input. A rule that saves a lot of pain: the gradient of a scalar loss with respect to any array has the same shape as that array. Use that as a running check.

Step 1, the logits. From the previous section, and including the batch average:

dZ2 = (P − Y) / N              shape (N, K)      (Y is the one-hot labels)

Step 2, the second layer's parameters. Element by element, Z2[n,k] = Σ_h A1[n,h]·W2[h,k] + b2[k]. So ∂L/∂W2[h,k] = Σ_n A1[n,h] · dZ2[n,k]. Read that sum: it is a matrix product with the first index summed over n.

dW2 = A1ᵀ @ dZ2               (H, N) @ (N, K) = (H, K)   matches W2
db2 = Σ_n dZ2[n, :]           sum over the batch axis    (K,)   matches b2

The bias appears once per example, so its gradient adds up over the batch.

Step 3, pass the gradient back through the layer. A1[n,h] affects the loss through every logit k, so ∂L/∂A1[n,h] = Σ_k dZ2[n,k] · W2[h,k]:

dA1 = dZ2 @ W2ᵀ               (N, K) @ (K, H) = (N, H)   matches A1

Step 4, through the ReLU. ReLU acts on each entry separately, and its derivative is 1 where the input was positive and 0 elsewhere. So the gradient is masked, entry by entry:

dZ1 = dA1 ⊙ 1[Z1 > 0]         elementwise product        (N, H)

Step 5, the first layer's parameters. Exactly the same shape reasoning as step 2, with X in place of A1:

dW1 = Xᵀ @ dZ1                (D, N) @ (N, H) = (D, H)   matches W1
db1 = Σ_n dZ1[n, :]                                     (H,)   matches b1

That is the whole algorithm. Notice the pattern: every layer needs its stored input to compute its weight gradient (A1 and X), and its weights to pass the gradient further back (W2). That is why the forward pass has to keep its activations in memory, and why memory, not compute, is often what limits training.

The code

For data I will generate the classic two-spirals problem: two interleaved arms that no straight line can separate, which is the reason we need a hidden layer at all.

def make_spirals(n_per_class=150, noise=0.15, seed=0):
    r = np.random.default_rng(seed)
    t = np.sqrt(r.uniform(0.05, 1.0, n_per_class)) * 3 * np.pi
    X, y = [], []
    for k in range(2):
        angle = t + k * np.pi                     # second arm is rotated by 180 degrees
        radius = t / (3 * np.pi)
        pts = np.column_stack([radius * np.cos(angle), radius * np.sin(angle)])
        X.append(pts + r.normal(0, noise * 0.3, pts.shape))
        y.append(np.full(n_per_class, k))
    return np.vstack(X), np.concatenate(y)


class MLP:
    def __init__(self, n_in, n_hidden, n_out, seed=0):
        r = np.random.default_rng(seed)
        self.p = {
            "W1": r.normal(0, np.sqrt(2.0 / n_in), (n_in, n_hidden)),       # He init
            "b1": np.zeros(n_hidden),
            "W2": r.normal(0, np.sqrt(2.0 / n_hidden), (n_hidden, n_out)),
            "b2": np.zeros(n_out),
        }

    def forward(self, X, y):
        p = self.p
        Z1 = X @ p["W1"] + p["b1"]
        A1 = np.maximum(0.0, Z1)
        Z2 = A1 @ p["W2"] + p["b2"]
        P = softmax(Z2)
        loss = -np.mean(np.log(P[np.arange(len(y)), y] + 1e-12))
        return loss, (X, Z1, A1, P)

    def backward(self, y, cache):
        X, Z1, A1, P = cache
        N = len(y)
        dZ2 = P.copy()
        dZ2[np.arange(N), y] -= 1.0
        dZ2 /= N                                   # (N, K)
        grads = {}
        grads["W2"] = A1.T @ dZ2                   # (H, K)
        grads["b2"] = dZ2.sum(axis=0)              # (K,)
        dA1 = dZ2 @ self.p["W2"].T                 # (N, H)
        dZ1 = dA1 * (Z1 > 0)                       # (N, H)
        grads["W1"] = X.T @ dZ1                    # (D, H)
        grads["b1"] = dZ1.sum(axis=0)              # (H,)
        return grads

    def predict(self, X):
        A1 = np.maximum(0.0, X @ self.p["W1"] + self.p["b1"])
        return (A1 @ self.p["W2"] + self.p["b2"]).argmax(axis=1)

Before training anything, verify the backward pass against finite differences on a small network and a small dataset. I nudge the weights away from their initial values first so that no ReLU input sits exactly at zero, where the derivative is undefined and finite differences would disagree with the subgradient we chose.

Xs_small, ys_small = make_spirals(20, seed=3)
net = MLP(2, 8, 2, seed=1)
jitter = np.random.default_rng(5)
for k in net.p:
    net.p[k] = net.p[k] + 0.1 * jitter.normal(size=net.p[k].shape)

loss, cache = net.forward(Xs_small, ys_small)
grads = net.backward(ys_small, cache)

for name in net.p:
    def loss_of(v, name=name):
        old = net.p[name]
        net.p[name] = v
        l, _ = net.forward(Xs_small, ys_small)
        net.p[name] = old
        return l
    numeric = numerical_grad(loss_of, net.p[name])
    print(f"{name}: relative error {rel_error(grads[name], numeric):.2e}")

W1: relative error 1.16e-07
b1: relative error 7.53e-08
W2: relative error 3.12e-09
b2: relative error 1.15e-09

Every parameter matches to about eight digits or better. Now we can train with full-batch gradient descent, and evaluate on a fresh spiral drawn with a different seed so we are not just admiring the training set.

Xtr, ytr = make_spirals(150, seed=0)
Xte, yte = make_spirals(150, seed=1)

net = MLP(2, 64, 2, seed=0)
for step in range(3001):
    loss, cache = net.forward(Xtr, ytr)
    grads = net.backward(ytr, cache)
    for k in net.p:
        net.p[k] -= 0.5 * grads[k]                 # plain gradient descent, lr = 0.5
    if step % 500 == 0:
        acc = (net.predict(Xtr) == ytr).mean()
        print(f"step {step:4d}  loss {loss:.4f}  train acc {acc:.3f}")

print("test accuracy:", (net.predict(Xte) == yte).mean())

step    0  loss 0.7568  train acc 0.640
step  500  loss 0.1634  train acc 0.970
step 1000  loss 0.0666  train acc 0.993
step 1500  loss 0.0374  train acc 0.993
step 2000  loss 0.0247  train acc 1.000
step 2500  loss 0.0182  train acc 1.000
step 3000  loss 0.0140  train acc 1.000
test accuracy: 1.0

Forty lines of NumPy learn a problem a linear model cannot touch. Frameworks such as PyTorch automate exactly this bookkeeping (that is what "autograd" means), but the algorithm they run is the one above.

Vanishing and exploding gradients, and why initialisation matters

Look at the backward pass again: the gradient at layer ℓ is the gradient at the output multiplied by a Wᵀ and an activation mask for every layer in between. In a deep network that is a long product of matrices. If each multiplication tends to shrink the signal a little, the gradient reaching the early layers is essentially zero (vanishing), and if each grows it a little, it becomes enormous (exploding). The forward pass has the same issue with the activations themselves.

The fix starts with a variance calculation. Take one unit that sums D inputs, y = Σᵢ wᵢ xᵢ, where the weights and inputs are independent with zero mean. Then

Var(y) = Σᵢ Var(wᵢ xᵢ) = D · Var(w) · E[x²]

To keep the signal at the same scale from layer to layer we want Var(y) ≈ E[x²], so we need Var(w) = 1/D. Glorot (Xavier) initialisation averages the forward and backward requirements and uses Var(w) = 2/(D_in + D_out). ReLU sets negative values to zero, which throws away about half the second moment, so a ReLU layer needs twice as much variance to compensate: Var(w) = 2/D_in, called He initialisation. That is the np.sqrt(2.0 / n_in) in the constructor. Here is the effect over twenty layers of width 256, with three choices of weight scale:

x0 = rng.normal(size=(1000, 256))
scales = {"too small (0.01)": lambda n: 0.01,
          "too big (0.2)": lambda n: 0.2,
          "He (sqrt(2/n))": lambda n: np.sqrt(2.0 / n)}

for name, scale in scales.items():
    h = x0
    stds = []
    for layer in range(20):
        Wl = rng.normal(0, scale(256), (256, 256))
        h = np.maximum(0, h @ Wl)
        stds.append(h.std())
    print(f"{name:18s} std after layers 1, 5, 10, 20:",
          ["%.1e" % stds[i] for i in (0, 4, 9, 19)])

too small (0.01)   std after layers 1, 5, 10, 20: ['9.4e-02', '1.5e-05', '2.3e-10', '6.2e-20']
too big (0.2)      std after layers 1, 5, 10, 20: ['1.9e+00', '6.1e+01', '3.1e+03', '1.1e+07']
He (sqrt(2/n))     std after layers 1, 5, 10, 20: ['8.3e-01', '7.3e-01', '6.9e-01', '6.5e-01']

With too-small weights the activations collapse to around 1e-19; with too-big weights they climb past 1e7; with He scaling they stay between roughly 0.6 and 0.8. The exact values depend on the random draw, but the pattern does not. Activation choice matters too: a sigmoid has a maximum slope of 0.25, so its factor in the chain is at best a quarter per layer, a large part of why deep sigmoid networks were hard to train and why ReLU-family activations, residual connections and normalisation layers took over.

Probability and information

Models are uncertain, and probability is the language for saying how uncertain. A few definitions, then the two ideas that matter most: Bayes' rule and cross-entropy as a distance between distributions.

The expectation of a random variable is its probability-weighted average, E[X] = Σ x·p(x). The variance measures spread around that average, Var(X) = E[(X − E[X])²] = E[X²] − E[X]², and its square root is the standard deviation. The Gaussian (normal) distribution with mean μ and variance σ² has density

N(x; μ, σ²) = 1 / sqrt(2πσ²) · exp( −(x − μ)² / (2σ²) )

It shows up everywhere becuase sums of many small independent effects tend towards it (the central limit theorem), and because its log is a plain quadratic, which is the reason maximum likelihood with Gaussian noise turns into least squares. Maximising the Gaussian likelihood over μ gives the sample mean, and over σ² gives the average squared deviation from it.

Bayes' rule with a real number attached

Bayes' rule tells you how to update a belief when you see evidence:

P(A | B) = P(B | A) · P(A) / P(B)

The classic trap is ignoring the base rate, so here is a medical-test scenario with numbers chosen to make the point (they are illustrative, not from a real test). A condition affects 1% of a population. A screening test catches 95% of people who have it (sensitivity) and correctly clears 90% of people who do not (specificity). You test positive. What is the chance you have it?

P(disease) = 0.01           P(positive | disease) = 0.95
P(positive | healthy) = 0.10

P(positive) = 0.95 · 0.01 + 0.10 · 0.99 = 0.0095 + 0.0990 = 0.1085
P(disease | positive) = 0.0095 / 0.1085 ≈ 0.0876

About 9%, not 95%. Because healthy people vastly outnumber sick ones, most positives are false alarms. The same reasoning drives a spam filter. Suppose 40% of mail is spam, the word "winner" appears in 30% of spam and 2% of legitimate mail. Then P(spam | "winner") = 0.30·0.40 / (0.30·0.40 + 0.02·0.60) = 0.12 / 0.132 ≈ 0.909. Naive Bayes multiplies evidence like this across many words, assuming they are independent given the class. Let me confirm the medical number by simulation instead of trusting the algebra:

prior, sensitivity, specificity = 0.01, 0.95, 0.90
p_pos = sensitivity * prior + (1 - specificity) * (1 - prior)
print("analytic P(disease | positive):", round(sensitivity * prior / p_pos, 4))

n = 1_000_000
sick = rng.random(n) < prior
positive = np.where(sick, rng.random(n) < sensitivity, rng.random(n) < (1 - specificity))
print("simulated:                    ", round(sick[positive].mean(), 4))

analytic P(disease | positive): 0.0876
simulated:                     0.0872

Entropy, cross-entropy and KL divergence

Now the information-theory trio. For a discrete distribution p:

Entropy            H(p)    = − Σ p(x) log p(x)          the average surprise; the best possible
                                                        average code length (in nats for natural log)
Cross-entropy      H(p, q) = − Σ p(x) log q(x)          the average surprise when you believe q
                                                        but reality is p
KL divergence      KL(p‖q) = Σ p(x) log( p(x) / q(x) )  the extra surprise from believing q

Splitting the log of the ratio gives the relationship that explains the loss function we have been using:

KL(p‖q) = Σ p log p − Σ p log q = H(p, q) − H(p)

KL is never negative (a consequence of Jensen's inequality) and is zero only when q = p. It is not symmetric, so it is a divergence and not a distance: KL(p‖q) ≠ KL(q‖p) in general.

Here is the punchline. In supervised learning, p is the true label distribution and q is the model's output. H(p) is fixed by the data and does not depend on the model's parameters, so minimising cross-entropy over the parameters is the same as minimising KL divergence to the data distribution. For one-hot labels the entropy of the data is exactly zero, so cross-entropy and KL are literally the same number. That is a third way of arriving at the same loss, next to maximum likelihood and the p − y gradient. Three viewpoints, one function.

def entropy(p):
    p = p[p > 0]
    return -np.sum(p * np.log(p))

def cross_ent(p, q):
    m = p > 0
    return -np.sum(p[m] * np.log(q[m]))

def kl(p, q):
    m = p > 0
    return np.sum(p[m] * np.log(p[m] / q[m]))

p = np.array([0.7, 0.2, 0.1])       # "reality"
q = np.array([0.4, 0.4, 0.2])       # a model's belief
print("H(p)          =", round(entropy(p), 4))
print("H(p, q)       =", round(cross_ent(p, q), 4))
print("KL(p || q)    =", round(kl(p, q), 4))
print("H(p) + KL     =", round(entropy(p) + kl(p, q), 4), "(equals H(p, q))")
print("KL(q || p)    =", round(kl(q, p), 4), "(not symmetric)")
print("uniform, 3 outcomes:", round(entropy(np.ones(3) / 3), 4), "= log 3 =", round(np.log(3), 4))

H(p)          = 0.8018
H(p, q)       = 0.9856
KL(p || q)    = 0.1838
H(p) + KL     = 0.9856 (equals H(p, q))
KL(q || p)    = 0.192 (not symmetric)
uniform, 3 outcomes: 1.0986 = log 3 = 1.0986

Note the last line: uniform uncertainty over three outcomes has the maximum possible entropy, log 3. It is also why a randomly initialised 3-class classifier should start with a loss near log 3 ≈ 1.10. If your first loss is wildly different from log(K), something is off with your setup, and this is one of the cheapest sanity checks there is.

Eigenvectors, SVD, PCA and low-rank structure

Linear algebra has one more idea that keeps coming back: decomposing a matrix into simple pieces.

An eigenvector of a square matrix S is a direction that S only stretches: Sv = λv. For a symmetric matrix (a covariance matrix, say) the eigenvectors are perpendicular and the matrix can be written S = QΛQᵀ, with the eigenvectors in the columns of Q and the eigenvalues on the diagonal of Λ.

Most matrices in machine learning are not square or symmetric, and there the workhorse is the singular value decomposition. Any (m, n) matrix can be written

A = U Σ Vᵀ       U: (m, m)   Σ: (m, n) diagonal, σ₁ ≥ σ₂ ≥ ... ≥ 0   Vᵀ: (n, n)

with U and V having perpendicular unit-length columns. In words: rotate, stretch along the axes by the singular values, rotate again. It is connected to eigenvectors, because AᵀA = V Σ² Vᵀ: the right singular vectors are the eigenvectors of AᵀA, and the squared singular values are its eigenvalues.

PCA as variance maximisation

Principal component analysis answers the question "if I could keep only one direction, which one preserves the most information?" Take centred data Xc of shape (N, D) (each column has mean zero), with covariance C = XcᵀXc / (N − 1), shape (D, D). Project onto a unit vector v. The projected values are Xc v, and their variance is

Var(Xc v) = vᵀ C v

We want to maximise vᵀCv subject to vᵀv = 1. Add a Lagrange multiplier for the constraint, ℒ = vᵀCv − λ(vᵀv − 1), and set the gradient with respect to v to zero:

2Cv − 2λv = 0     ⟹     Cv = λv

So the best direction is an eigenvector of C, and plugging back in, the variance achieved is vᵀCv = λvᵀv = λ. To capture the most variance, choose the eigenvector with the largest eigenvalue. The second component is the best remaining direction perpendicular to the first, and so on. So PCA is the eigendecomposition of the covariance matrix, and by the relation above it is also the SVD of the centred data, with λᵢ = σᵢ² / (N − 1). In practice people use the SVD route because it is numerically kinder.

Low-rank approximation, and the link to LoRA

Write the SVD as a sum of rank-one pieces, A = Σᵢ σᵢ uᵢ vᵢᵀ. If you keep only the first k terms you get a matrix A_k of rank k. The Eckart–Young theorem (stated here, not proven) says this truncation is the best possible rank-k approximation to A, in both the Frobenius and spectral norm, and that the Frobenius error is exactly the size of what you dropped:

min over rank-k B of ‖A − B‖_F  =  ‖A − A_k‖_F  =  sqrt( σ_{k+1}² + σ_{k+2}² + ... )

Here is a numerical check. I build a 500 × 50 matrix with five real directions of structure plus a little noise, decompose it, and reconstruct at several ranks.

n_rows, n_cols, true_rank = 500, 50, 5
A = rng.normal(size=(n_rows, true_rank)) @ rng.normal(size=(true_rank, n_cols))
A += 0.1 * rng.normal(size=(n_rows, n_cols))

U, S, Vt = np.linalg.svd(A, full_matrices=False)
print("top singular values:", S[:8])

for k in (1, 3, 5, 10):
    A_k = (U[:, :k] * S[:k]) @ Vt[:k]
    err = np.linalg.norm(A - A_k)                  # Frobenius norm by default
    predicted = np.sqrt(np.sum(S[k:] ** 2))        # Eckart-Young
    stored = k * (n_rows + n_cols) / (n_rows * n_cols)
    print(f"rank {k:2d}: error {err:8.4f}   sqrt(sum dropped sigma^2) {predicted:8.4f}"
          f"   storage {stored:.1%}")

# PCA two ways: eigendecomposition of the covariance vs SVD of the centred data
Xc = A - A.mean(axis=0)
evals, evecs = np.linalg.eigh(Xc.T @ Xc / (n_rows - 1))
evals, evecs = evals[::-1], evecs[:, ::-1]         # eigh sorts ascending
_, S2, Vt2 = np.linalg.svd(Xc, full_matrices=False)
print("eigenvalues == sigma^2/(N-1):", np.allclose(evals, S2**2 / (n_rows - 1)))
print("first principal axis agrees (up to sign):", np.allclose(np.abs(evecs[:, 0]), np.abs(Vt2[0])))
ratio = evals / evals.sum()
print("explained variance ratio:", ratio[:6], " top 5 total:", round(ratio[:5].sum(), 4))

top singular values: [168.1764 154.1708 149.6911 136.465  122.9326   2.8255   2.7771   2.7302]
rank  1: error 283.0787   sqrt(sum dropped sigma^2) 283.0787   storage 2.2%
rank  3: error 184.2756   sqrt(sum dropped sigma^2) 184.2756   storage 6.6%
rank  5: error  14.9124   sqrt(sum dropped sigma^2)  14.9124   storage 11.0%
rank 10: error  13.5882   sqrt(sum dropped sigma^2)  13.5882   storage 22.0%
eigenvalues == sigma^2/(N-1): True
first principal axis agrees (up to sign): True
explained variance ratio: [0.2609 0.2196 0.2063 0.1716 0.1396 0.0001]  top 5 total: 0.9979

Five singular values stand well above the rest, matching the five directions we planted, and once k reaches 5 the error collapses to the noise floor. The error at every rank equals the Eckart–Young prediction to the printed digits. And the PCA two ways agree.

This is the idea behind LoRA, which we covered in the fine-tuning part of this series. A weight matrix W₀ of shape (d_out, d_in) has d_out · d_in numbers. LoRA freezes it and learns an update in factored form:

W = W₀ + (α / r) · B A        B: (d_out, r)     A: (r, d_in)     r much smaller than d_in

The bet, which the LoRA paper argues from experiments, is that the update needed to adapt a pretrained model has low intrinsic rank, so a rank-r factorisation loses little. The saving is the same arithmetic as the storage column above, and the same bracketing trick as our earlier cost example: compute B(Ax) rather than forming BA first.

d_out, d_in, r = 512, 512, 8
print("full update:", d_out * d_in, "  LoRA rank 8:", r * (d_out + d_in),
      f"  ({r * (d_out + d_in) / (d_out * d_in):.1%})")

full update: 262144   LoRA rank 8: 8192   (3.1%)

One caveat: Eckart–Young gives the best low-rank approximation of a given matrix. LoRA does not decompose anything; it learns B and A directly by gradient descent, trusting that a low-rank update exists.

Optimisers: SGD, momentum and Adam

Plain gradient descent uses the whole dataset for every step. Stochastic gradient descent (SGD) uses a random mini-batch, which gives a noisy but unbiased estimate of the gradient at a fraction of the cost. The three update rules you need to know, with g the current gradient and η the learning rate:

SGD:        w ← w − η g

Momentum:   v ← β v + g
            w ← w − η v                     (β is typically 0.9)

Adam:       m ← β₁ m + (1 − β₁) g           running mean of gradients
            v ← β₂ v + (1 − β₂) g²          running mean of squared gradients
            m̂ = m / (1 − β₁ᵗ)               bias-corrected
            v̂ = v / (1 − β₂ᵗ)
            w ← w − η m̂ / (√v̂ + ε)          (β₁ = 0.9, β₂ = 0.999, ε = 1e-8)

Momentum keeps a running velocity, so consistent directions build up speed while directions that keep flipping sign cancel out. Adam (Kingma and Ba) additionally divides each coordinate by the root of its recent squared gradient, so every parameter gets its own step size: coordinates with big, noisy gradients get small steps, quiet ones get large steps.

Why the bias correction? Both m and v start at zero. At step 1, m = (1 − β₁) g = 0.1 g, which badly underestimates the gradient. In general, if the gradient were a constant g, then m_t = (1 − β₁ᵗ) g. So the running mean is biased towards zero by exactly the factor (1 − β₁ᵗ) and dividing by it removes the bias. The effect fades as t grows, because β₁ᵗ → 0. For v the correction matters even more since β₂ = 0.999 means the average takes about a thousand steps to warm up, and without it the first steps would divide by an underestimate of the scale and take huge, unstable steps. A tiny check:

g, m = 1.0, 0.0
for t in range(1, 4):
    m = 0.9 * m + 0.1 * g
    print(f"t={t}: raw m = {m:.3f}   corrected = {m / (1 - 0.9**t):.3f}")

t=1: raw m = 0.100   corrected = 1.000
t=2: raw m = 0.190   corrected = 1.000
t=3: raw m = 0.271   corrected = 1.000

Now the three optimisers on a deliberately awkward function: a long thin bowl f(w) = ½(1·w₀² + 50·w₁²), curved 50 times more steeply in one direction. The 2/L rule from earlier says plain SGD needs η < 2/50 = 0.04, and at that safe step size it crawls along the flat direction.

curv = np.array([1.0, 50.0])
def f(w):    return 0.5 * np.sum(curv * w ** 2)
def grad(w): return curv * w

def sgd(w, g, state, lr=0.03):
    return w - lr * g

def momentum(w, g, state, lr=0.02, beta=0.9):
    v = beta * state.get("v", np.zeros_like(w)) + g
    state["v"] = v
    return w - lr * v

def adam(w, g, state, lr=0.1, b1=0.9, b2=0.999, eps=1e-8):
    t = state.get("t", 0) + 1
    m = b1 * state.get("m", np.zeros_like(w)) + (1 - b1) * g
    v = b2 * state.get("v", np.zeros_like(w)) + (1 - b2) * g ** 2
    state.update(t=t, m=m, v=v)
    m_hat = m / (1 - b1 ** t)
    v_hat = v / (1 - b2 ** t)
    return w - lr * m_hat / (np.sqrt(v_hat) + eps)

for name, opt in [("sgd", sgd), ("momentum", momentum), ("adam", adam)]:
    w, state, losses = np.array([5.0, 1.0]), {}, []
    for t in range(1, 101):
        w = opt(w, grad(w), state)
        if t in (10, 50, 100):
            losses.append(round(float(f(w)), 4))
    print(f"{name:9s} loss after 10, 50, 100 steps: {losses}")

sgd       loss after 10, 50, 100 steps: [6.7975, 0.5944, 0.0283]
momentum  loss after 10, 50, 100 steps: [1.9474, 0.13, 0.0003]
adam      loss after 10, 50, 100 steps: [8.1758, 0.4066, 0.001]

Read this as an illustration, not a benchmark, since the learning rates are hand-picked and untuned. Two things are visible. Momentum accelerates along the flat direction where plain SGD is held back by the steep one. Adam lags at step 10 because it moves each coordinate by roughly η = 0.1 per step whatever the gradient size, so crossing the 5 units of the flat direction takes about 50 steps, but it is ahead of plain SGD by step 50. Which optimiser wins on a real network depends on the problem and the tuning.

Numerical stability: the log-sum-exp trick

Real floating point numbers overflow. exp(1000) is far beyond what float64 can hold, and a naive softmax computes exactly that. The fix relies on a fact you can verify in one line: softmax is unchanged when you subtract the same constant from every logit,

exp(z_i − c) / Σ_j exp(z_j − c) = [exp(−c) · exp(z_i)] / [exp(−c) · Σ_j exp(z_j)] = exp(z_i) / Σ_j exp(z_j)

Choose c = max(z) and the largest exponent becomes exp(0) = 1, so nothing overflows, and at least one term is exactly 1, so the denominator is never zero. The same idea gives the log-sum-exp function, needed whenever you want log-probabilities:

log Σ exp(z_i) = m + log Σ exp(z_i − m)      with m = max(z)
log softmax(z)_i = z_i − logsumexp(z)

Computing log(softmax(z)) as two separate steps is a classic mistake, because the softmax can underflow to exactly 0 and the log of that is −inf. Working in log space throughout avoids it. Here is the failure, and the repair:

def softmax_naive(z):
    e = np.exp(z)
    return e / e.sum()

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

def logsumexp(z, axis=-1):
    m = z.max(axis=axis, keepdims=True)
    return (m + np.log(np.exp(z - m).sum(axis=axis, keepdims=True))).squeeze(axis)

def log_softmax(z, axis=-1):
    return z - np.expand_dims(logsumexp(z, axis), axis)

big = np.array([1000.0, 1001.0, 1002.0])
with np.errstate(all="ignore"):                  # silence the overflow warnings
    print("naive softmax:   ", softmax_naive(big))
    print("naive log(softmax) of [0, 1000]:", np.log(softmax_naive(np.array([0.0, 1000.0]))))
print("stable softmax:  ", softmax_stable(big))
print("logsumexp:       ", logsumexp(big))
print("log_softmax:     ", log_softmax(big))
print("log_softmax of [0, 1000]:", log_softmax(np.array([0.0, 1000.0])))

naive softmax:    [nan nan nan]
naive log(softmax) of [0, 1000]: [-inf  nan]
stable softmax:   [0.09   0.2447 0.6652]
logsumexp:        1002.4076059644444
log_softmax:      [-2.4076 -1.4076 -0.4076]
log_softmax of [0, 1000]: [-1000.     0.]

The naive version returns nan (infinity divided by infinity) and the log of a zero probability gives -inf; the stable versions return correct, finite answers. Every serious framework's cross-entropy function takes raw logits and does this internally, which is why you should pass logits, not probabilities, to those functions.

Attention is a dot product with a normaliser

Everything so far assembles into the most important formula in current deep learning. Given queries Q of shape (T, d_k), keys K of shape (T, d_k) and values V of shape (T, d_v), from Attention Is All You Need:

Attention(Q, K, V) = softmax( Q Kᵀ / √d_k ) V

Read it with what you now know. Q Kᵀ has shape (T, T): entry (i, j) is the dot product of query i with key j, the similarity score from the very first section. Softmax along each row turns those scores into weights that sum to 1. Multiplying by V takes, for each position, a weighted average of the value vectors. So attention is "similarity search, then a soft lookup". The Vision Transformer article in this series applied it to image patches, and the mechanism is identical: each patch asks every other patch how relevant it is and blends their values.

Why divide by √d_k? Suppose the entries of a query and a key are independent with mean 0 and variance 1. The score is q·k = Σᵢ qᵢkᵢ. Each term has mean 0 and variance E[qᵢ²]E[kᵢ²] = 1, and the terms are independent, so the variances add: Var(q·k) = d_k. Scores grow in spread as √d_k. A softmax fed numbers with a standard deviation of 22 (for d_k = 512) is almost one-hot, and a softmax that is almost one-hot has a gradient close to zero: its Jacobian, diag(p) − ppᵀ, vanishes when p is a one-hot vector. Dividing by √d_k restores unit variance and keeps the softmax in the range where it can learn. The causal mask used in language models sets the scores of future positions to a large negative number before the softmax, so their weights become zero.

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.swapaxes(-1, -2) / np.sqrt(d_k)     # (..., T, T)
    if mask is not None:
        scores = np.where(mask, scores, -1e9)          # blocked positions get ~zero weight
    weights = softmax_stable(scores, axis=-1)          # normalise over the keys
    return weights @ V, weights

T, d = 4, 8
Q, K, V = (rng.normal(size=(T, d)) for _ in range(3))
causal = np.tril(np.ones((T, T), dtype=bool))
out, w = attention(Q, K, V, causal)
print("output shape:", out.shape, " row sums:", w.sum(axis=-1))
print(w)

# The variance argument, checked: 100,000 random query/key pairs at d_k = 64
qs, ks = rng.normal(size=(100_000, 64)), rng.normal(size=(100_000, 64))
s = (qs * ks).sum(axis=1)
print("variance of q.k:", round(s.var(), 2), " after / sqrt(d_k):", round((s / np.sqrt(64)).var(), 4))

# And what saturation looks like at d_k = 512
q1, K1 = rng.normal(size=512), rng.normal(size=(6, 512))
print("unscaled: ", softmax_stable(K1 @ q1))
print("scaled:   ", softmax_stable(K1 @ q1 / np.sqrt(512)))

output shape: (4, 8)  row sums: [1. 1. 1. 1.]
[[1.     0.     0.     0.    ]
 [0.4545 0.5455 0.     0.    ]
 [0.4129 0.2784 0.3087 0.    ]
 [0.2366 0.2506 0.2162 0.2965]]
variance of q.k: 63.92  after / sqrt(d_k): 0.9987
unscaled:  [0.0003 0.     0.9185 0.     0.0813 0.    ]
scaled:    [0.2282 0.0402 0.3277 0.0647 0.2944 0.0447]

The variance is close to 64 before scaling and close to 1 after it, the causal weights are lower triangular with each row summing to 1, and the unscaled softmax at d_k = 512 puts about 92% of its weight on a single key of six while the scaled one spreads it out. The -1e9 is a common stand-in for -inf: it avoids nan if a whole row is ever masked.

Production reality: floating point, reproducibility and bugs

Everything above ran in float64, the forgiving default. Real training does not, and the differences bite.

The number formats

Floating point trades range against precision using a fixed number of bits. The common formats:

The trade is stark. float16 keeps more precision than bfloat16 but overflows at 65,504 and its smallest normal number is about 6.1e-5, so small gradients underflow to zero. bfloat16 has the same range as float32 but only about two to three decimal digits of precision. NumPy has no bfloat16 type, so I show float32 and float16 (np.finfo reports their limits) and the failure modes that follow from them:

with np.errstate(over="ignore"):
    print("float16 of 70000:", np.float16(70000.0))
print("float16 of 1e-8: ", np.float16(1e-8), "  (underflow to zero)")
print("float16: 2048 + 1 =", np.float16(2048) + np.float16(1), " (the +1 is lost)")

total = np.float16(0)
for _ in range(5000):
    total += np.float16(0.01)
print("float16 sum of 5000 x 0.01:", total, " (true value: 50)")
print("0.1 + 0.2 == 0.3 in float64:", 0.1 + 0.2 == 0.3)

float16 of 70000: inf
float16 of 1e-8:  0.0   (underflow to zero)
float16: 2048 + 1 = 2048.0  (the +1 is lost)
float16 sum of 5000 x 0.01: 32.0  (true value: 50)
0.1 + 0.2 == 0.3 in float64: False

The accumulation result is worth staring at. Adding 0.01 to a float16 running total stalls at 32, because at that magnitude the gap between neighbouring float16 numbers (0.031) is larger than twice the increment, so every addition rounds back to the same value. This is precisely why mixed-precision training keeps a float32 master copy of the weights and accumulates sums in float32: tiny updates would otherwise vanish.

Why fp16 needs loss scaling

Gradients in a deep network are often very small, smaller than the smallest float16 can represent. Loss scaling fixes it with the chain rule. Multiply the loss by a constant S before backpropagation, and every gradient is multiplied by S too, since the whole computation is linear in the loss. The gradients now sit comfortably in float16's range. Before the optimiser step you divide by S in float32 to recover the true values. (Mixed Precision Training describes the recipe.) Dynamic loss scaling adjusts S automatically: lower it when an overflow (inf or nan) appears in the gradients, raise it slowly while things stay finite. bfloat16 rarely needs any of this, since its range matches float32.

grad_true = np.float32(1e-8)
S = np.float32(65536)                          # 2**16

print("without scaling:", np.float16(grad_true))
scaled = np.float16(grad_true * S)
print("with scaling:   ", scaled, "-> unscaled in float32:", np.float32(scaled) / S)

without scaling: 0.0
with scaling:    0.000655 -> unscaled in float32: 9.997166e-09

Without scaling the gradient is flushed to zero. With scaling it survives, and dividing back gives the original value to about three or four significant digits, which is the precision float16 has to offer.

Reproducibility

Floating point addition is not associative: (a + b) + c and a + (b + c) can differ in the last bits. A parallel sum on a GPU or a multi-threaded BLAS may add in a different order from run to run, so two runs with the same seed can differ slightly. Practical advice: seed every generator you use (np.random.default_rng(seed) in NumPy; the framework and data loader seeds in PyTorch), record library versions, use the framework's deterministic mode when you need bit-exact runs and accept the speed cost, and compare results with np.allclose instead of ==. If a result flips between success and failure across seeds, you have a stability problem to understand, not a seed to pick.

Test numerical code like any other code

The tools in this article make good unit tests: a gradient check for anything with a hand-written backward pass, and property tests that assert things which must hold for any input. A minimal suite for our softmax and cross-entropy, using only the standard library and NumPy:

import unittest

class TestSoftmax(unittest.TestCase):
    def test_rows_sum_to_one(self):
        z = np.random.default_rng(0).normal(size=(6, 5)) * 50   # large logits on purpose
        np.testing.assert_allclose(softmax(z).sum(axis=1), 1.0)

    def test_shift_invariance(self):
        z = np.random.default_rng(1).normal(size=(3, 4))
        np.testing.assert_allclose(softmax(z), softmax(z + 123.0))

    def test_gradient_matches_finite_differences(self):
        r = np.random.default_rng(2)
        z, y = r.normal(size=(4, 3)), r.integers(0, 3, size=4)
        numeric = numerical_grad(lambda a: cross_entropy(a, y), z)
        self.assertLess(rel_error(cross_entropy_grad(z, y), numeric), 1e-6)

    def test_loss_is_a_mean_not_a_sum(self):
        r = np.random.default_rng(3)
        z, y = r.normal(size=(4, 3)), r.integers(0, 3, size=4)
        self.assertAlmostEqual(cross_entropy(z, y),
                               cross_entropy(np.tile(z, (2, 1)), np.tile(y, 2)))

unittest.main(argv=["numerics"], exit=False)  # or run with `python -m unittest`

....
----------------------------------------------------------------------
Ran 4 tests in 0.020s

OK

The last test catches a whole family of bugs: if you sum the loss instead of averaging it, duplicating the batch doubles the loss, the test fails, and you find out before a training run quietly depends on batch size.

The bugs I would check first

Most numerical bugs are not exotic. They are these:

1. Softmax over the wrong axis. Your logits are (N, K) and the normalisation must run over K, which is axis=1 (or -1). Using axis=0 normalises across the batch instead. Nothing crashes: the numbers just don't mean anything.
2. Silent broadcasting. A (N, 1) array minus a (N,) array does not raise an error. It gives an (N, N) matrix. Then the mean of that matrix is a plausible number that is not your loss.
3. Forgetting to average the loss over the batch, so the gradient scales with the batch size and the right learning rate silently changes whenever you change it.
4. Feeding probabilities where logits are expected (or the reverse), so softmax is applied twice.
5. Not checking shapes. Print them. Assert them.

Two of these reproduce in a few lines:

z = rng.normal(size=(3, 4))
print("softmax over axis 0, row sums:", softmax_stable(z, axis=0).sum(axis=1))
print("softmax over axis 1, row sums:", softmax_stable(z, axis=1).sum(axis=1))

pred = rng.normal(size=(8, 1))          # e.g. a model output left with a trailing axis
target = rng.normal(size=(8,))
print("pred - target shape:", (pred - target).shape)
print("wrong loss:  ", round(float(np.mean((pred - target) ** 2)), 4))
print("right loss:  ", round(float(np.mean((pred.ravel() - target) ** 2)), 4))

softmax over axis 0, row sums: [1.5339 1.6081 0.858 ]
softmax over axis 1, row sums: [1. 1. 1.]
pred - target shape: (8, 8)
wrong loss:   2.2106
right loss:   1.4501

The first shows row sums that are not 1, the tell-tale for a wrong axis. The second shows an (8, 8) result and two different losses; neither raises an error, and only one of them is right.

Keeping the ideas straight, and where to go next

If you want a reading order, the one I would follow is:

1. Linear algebra, geometrically. The 3Blue1Brown "Essence of linear algebra" videos build the intuition for matrices as maps and for eigenvectors, and Gilbert Strang's MIT lectures cover the SVD.
2. Calculus and the chain rule, by hand. Derive backprop for the two-layer network yourself, then break the code deliberately and watch the gradient check fail.
3. Probability with an eye on likelihood. Know maximum likelihood, Bayes' rule and KL divergence cold, because they explain why every loss function looks the way it does.
4. Optimisation. Change the learning rates on the toy problems here and watch the 2/L behaviour appear.
5. Then frameworks. PyTorch's autograd will read as bookkeeping rather than magic. The free textbook Deep Learning by Goodfellow, Bengio and Courville covers all these foundations at greater depth.

If you only remember three things: check shapes on every line, check every gradient numerically, and keep the numerics in mind when you drop below float32.

The five parts, briefly

This series started with the practical end, and this article is the ground under it. The earlier parts covered building agents and MCP integrations, LoRA and QLoRA for cheap fine-tuning (the low-rank section above is its mathematical core), Polars for fast data preparation (the X matrices come from somewhere), and Vision Transformers (whose attention layer is the attention section). Each is a tool, and this one is what explains why the tools behave as they do. None of them needs all of this maths to use, but when something goes wrong, the derivations are what you debug with.

Take the code, run it, change the numbers, and try to break it. A derivation you have checked against a brute-force number is one you understand; the rest is memorised.