Chain rule

01.01 · Functions

Start with one prediction. For f(x)=2x+1, write f(4) before reading on. The input is 4, the fixed rule doubles it and adds 1, and the output is 9. A parameter would be a number inside the rule that training may change.

What is a function?

A function is a rule that maps each allowed input to exactly one output. The notation $f(x)$ means "the output of $f$ when the input is $x$." It does not mean $f$ multiplied by $x$.

Example: A simple rule

Consider $f(x) = 2x + 1$. This rule doubles the input and adds 1:

• $f(2) = 2(2) + 1 = 5$
• $f(3) = 2(3) + 1 = 7$
• $f(-1) = 2(-1) + 1 = -1$

Functions in neural networks

A neural network is a parameterised function $f_\theta: \mathcal{X} \to \mathcal{Y}$, where $\theta$ are the trainable parameters (weights and biases). A simple linear layer computes:

where $w$ (weight) and $b$ (bias) are parameters we adjust during training. In Python:

def predict(x, w=2, b=1):
    return w * x + b

print(predict(3))        # 7
print(predict(3, w=4))   # 13

Linear vs. affine

A linear map satisfies $f(x+y) = f(x) + f(y)$ and $f(cx) = cf(x)$. Example: $f(x) = 2x$.

An affine function is a linear map plus an offset: $f(x) = wx + b$. Neural network layers are typically affine transformations followed by nonlinear activations.

Composition

Neural networks compose many functions. If $f(x) = 2x$ and $g(x) = x^2$, then $(g \circ f)(x) = g(f(x)) = (2x)^2 = 4x^2$. Composition is associative but not commutative: $g \circ f \neq f \circ g$ in general.

01.02 · Derivatives

Predict a nearby value. If f(x)=x², then f(3)=9. The slope at 3 is 6, so moving to 3.01 predicts an increase near 6×0.01=0.06. The exact value is 9.0601, which shows both the usefulness and the local limit of the derivative.

What is a derivative?

The derivative measures the instantaneous rate of change of a function at a point. For $f(x)$, the derivative at $x$ is:

Geometrically, $f'(x)$ is the slope of the tangent line to the curve $y = f(x)$ at the point $(x, f(x))$.

Example: Derivative of $x^2$

Let $g(x) = x^2$. At $x = 3$, we compute:

The ratios approach 6, so $g'(3) = 6$.

Derivation

Using the limit definition:

So $g'(x) = 2x$, giving $g'(3) = 6$ and $g'(-2) = -4$.

Using the derivative to predict changes

If $x$ moves from 3 to 3.02 (change $h = 0.02$), the derivative predicts:

The exact change is $3.02^2 - 3^2 = 0.1204$. Close, but not exact — the approximation improves as $h$ shrinks.

Notation

01.03 · Vectors

Name the coordinates first. Let x=[2,3] mean two measured features. With weights w=[4,1], the dot product is 4×2+1×3=11. Swapping the vector entries changes the meaning and usually changes the answer.

What is a vector?

A vector is an ordered list of numbers. In AI, vectors represent features, parameters, gradients, and more:

Example: a study session might be represented as $\mathbf{x} = [2, 3]$, where $x_1 = 2$ hours reading and $x_2 = 3$ exercises attempted. Order matters: $[3, 2]$ is a different vector.

The dot product

The dot product combines two vectors into a scalar. It measures how much one vector "aligns" with another:

Example: with $\mathbf{w} = [4, 1]$ and $\mathbf{x} = [2, 3]$:

Vector operations in code

import numpy as np

x = np.array([2, 3])
w = np.array([4, 1])

# Dot product
score = np.dot(w, x)  # or w @ x
print(score)  # 11

# Addition
y = x + np.array([1, 2])  # [3, 5]

# Scalar multiplication
z = 2 * x  # [4, 6]

# Norm
length = np.linalg.norm(x)  # sqrt(13) ≈ 3.606

Geometric interpretation

A vector $\mathbf{x} = [x_1, x_2]$ can be drawn as an arrow from the origin $(0,0)$ to the point $(x_1, x_2)$. The dot product has a geometric formula:

where $\theta$ is the angle between the vectors. This shows that:

• $\mathbf{u} \cdot \mathbf{v} > 0$ when $\theta < 90°$ (vectors point in similar directions)
• $\mathbf{u} \cdot \mathbf{v} = 0$ when $\theta = 90°$ (vectors are orthogonal/perpendicular)
• $\mathbf{u} \cdot \mathbf{v} < 0$ when $\theta > 90°$ (vectors point in opposite directions)

01.04 · Gradients

Use a loss instead of a hill. For L(w,b)=(w−1)²+(b+2)², the gradient at (3,−1) is [4,2]. Predict the descent direction: both coordinates should decrease, so we step opposite the gradient.

Partial derivatives

When a function depends on multiple variables, we can ask how it changes as one variable moves while the others stay fixed. The partial derivative with respect to $x_i$ is:

The gradient

The gradient collects all partial derivatives into a single vector:

For our example: $\nabla L(w,b) = [2(w-1),\; 2(b+2)]$. At $(3,-1)$: $\nabla L = [4, 2]$.

Gradient descent

To minimize a loss $L$, step opposite the gradient:

where $\eta > 0$ is the learning rate (step size).

Learning rate matters

For a quadratic $L(w) = (w-1)^2$ starting at $w_0 = 3$:

• $\eta = 0.1$: slow, steady convergence ✓
• $\eta = 0.5$: converges in one step ✓
• $\eta = 1.2$: overshoots, loss increases ✗

The gradient as a vector field

The gradient $\nabla f$ is a vector field: at each point in the input space, it assigns a vector. Plotting these vectors over level curves of $f$ reveals the geometry:

• The gradient is orthogonal to level curves (contour lines)
• It points uphill (steepest ascent)
• Its magnitude $\|\nabla f\|$ tells you how steep the slope is

01.05 · Matrices

Check the shape before multiplying. A matrix with shape (3,2) maps a two-value input to a three-value output. Each output row forms one weighted sum from the same two inputs. A common wrong answer reverses these dimensions.

What is a matrix?

A matrix is a rectangular array of numbers. It can represent a linear transformation, a dataset, or a set of parameters:

This $2 \times 2$ matrix has 2 rows and 2 columns. We describe its shape as (rows, columns).

Matrix-vector product

A matrix acts on a vector to produce a new vector. With $W$ as above and $\mathbf{x} = [2, 3]^\top$:

Each row of $W$ produces one output via a dot product with $\mathbf{x}$.

Shape compatibility

For $A \in \mathbb{R}^{m \times n}$ and $B \in \mathbb{R}^{p \times q}$, the product $AB$ is defined only if $n = p$ (inner dimensions match). The result has shape $(m, q)$:

Transpose

The transpose $W^\top$ swaps rows and columns:

Properties:

• $(A^\top)^\top = A$
• $(AB)^\top = B^\top A^\top$
• $(A + B)^\top = A^\top + B^\top$

Matrix operations in code

import numpy as np

W = np.array([[4, 1], [0, 2]])
x = np.array([2, 3])

# Matrix-vector product
y = W @ x  # or W.dot(x)
print(y)  # [11, 6]

# Transpose
Wt = W.T
print(Wt)  # [[4, 0], [1, 2]]

# Matrix-matrix product
X = np.array([[2, 5], [3, 3]])
Y = W @ X
print(Y)  # [[11, 23], [6, 6]]

01.06 · Derivation Rules & Examples

Choose the rule from the expression. A sum needs the sum rule. A multiplication of two changing functions needs the product rule. A function inside another function needs the chain rule. First mark the outermost operation; then work inward.

The toolbox of differentiation

Rather than returning to the limit definition every time, we use a small set of rules that compose to differentiate almost any function encountered in machine learning.

Exponentials and logarithms

These two functions are central to information theory and probability in ML:

Trigonometric functions

Activation functions in neural networks

Every activation function must be differentiable (almost everywhere) so gradients can flow backward:

Worked examples

Example 1: Mean squared error

$L(\theta) = \frac{1}{N}\sum_{i=1}^N (f_\theta(x_i) - y_i)^2$. Using the chain rule and power rule:

Example 2: Cross-entropy loss

$L = -\sum_i y_i \log p_i$. With respect to $p_j$:

Example 3: L2 regularisation

$R(\theta) = \frac{\lambda}{2}\|\theta\|^2 = \frac{\lambda}{2}\sum_i \theta_i^2$. Then:

01.07 · The Chain Rule

Trace one chain. Let u=2x and y=u². At x=3, dy/du=12 and du/dx=2. Their product gives dy/dx=24. Each factor measures one local link.

Single-variable chain rule

If $y = f(g(x))$, then:

Multi-step compositions

For a chain $y = f_3(f_2(f_1(x)))$:

Each factor is the derivative of one layer with respect to its input. This product structure is what makes backpropagation efficient.

Multivariable chain rule

When intermediate values are vectors, the derivatives become Jacobian matrices, and the chain rule becomes matrix multiplication:

where $J_i$ is the Jacobian of the $i$-th layer.

Leibniz notation makes the chain rule intuitive

The notation $\frac{dy}{dx} = \frac{dy}{du}\cdot\frac{du}{dx}$ looks like fractions cancelling. While this is technically an abuse of notation, it's a powerful mnemonic. For a chain $x \to a \to b \to y$:

Worked example: a tiny network

Consider a two-layer network with one neuron each:

To find $\frac{\partial L}{\partial W_1}$, apply the chain rule:

Computing each factor:

• $\frac{\partial L}{\partial y} = 2(y - t)$
• $\frac{\partial y}{\partial h} = W_2$
• $\frac{\partial h}{\partial a} = \sigma'(a) = h(1-h)$
• $\frac{\partial a}{\partial W_1} = x$

Multiplying: $\frac{\partial L}{\partial W_1} = 2(y-t) \cdot W_2 \cdot h(1-h) \cdot x$

01.08 · Backpropagation in Python

Predict signs before code. If a larger weight makes an already-too-large prediction even larger, its loss gradient should be positive. Gradient descent will subtract that value. If the program reports the opposite sign, inspect the error definition and local derivative.

From chain rule to algorithm

Backpropagation is the efficient algorithm for computing gradients in neural networks. It applies the chain rule right-to-left, reusing intermediate values from the forward pass.

Implementation from scratch

Here's a minimal backprop implementation for a two-layer network:

import numpy as np

def sigmoid(z): return 1 / (1 + np.exp(-z))
def sigmoid_grad(a): return a * (1 - a)

# Forward pass
def forward(x, W1, b1, W2, b2):
    z1 = W1 @ x + b1
    a1 = sigmoid(z1)
    z2 = W2 @ a1 + b2
    a2 = sigmoid(z2)
    cache = (x, z1, a1, z2, a2)
    return a2, cache

# Backward pass
def backward(y_true, cache, W2):
    x, z1, a1, z2, a2 = cache
    m = x.shape[1]  # batch size
    
    # Output layer error
    # Squared error, summed over outputs and averaged over the batch.
    dz2 = 2 * (a2 - y_true) * sigmoid_grad(a2)  # shape: (n_out, m)
    dW2 = (1/m) * dz2 @ a1.T                   # shape: (n_out, n_hid)
    db2 = (1/m) * np.sum(dz2, axis=1, keepdims=True)
    
    # Hidden layer error
    da1 = W2.T @ dz2                            # shape: (n_hid, m)
    dz1 = da1 * sigmoid_grad(a1)               # element-wise
    dW1 = (1/m) * dz1 @ x.T                    # shape: (n_hid, n_in)
    db1 = (1/m) * np.sum(dz1, axis=1, keepdims=True)
    
    return dW1, db1, dW2, db2

Step-by-step walkthrough

Computational graphs

Modern frameworks (PyTorch, JAX) build a computational graph during the forward pass, recording every operation. Backward traversal of this graph applies the chain rule automatically:

import torch

x = torch.tensor([1.0, 0.0], requires_grad=True)
W1 = torch.randn(4, 2, requires_grad=True)
b1 = torch.randn(4, requires_grad=True)
W2 = torch.randn(1, 4, requires_grad=True)
b2 = torch.randn(1, requires_grad=True)

# Forward — PyTorch records the graph
h = torch.sigmoid(W1 @ x + b1)
y = torch.sigmoid(W2 @ h + b2)
loss = (y - 1.0)**2

# Backward — gradients computed automatically
loss.backward()
print(W1.grad)  # ∂L/∂W1
print(W2.grad)  # ∂L/∂W2

Memory and efficiency

Backprop requires storing all intermediate activations $a^{(\ell)}$ from the forward pass. This is why deep networks are memory-intensive. Techniques to reduce memory:

• Gradient checkpointing: recompute some activations during backward pass
• Mixed precision: use FP16 for activations, FP32 for gradients
• Activation recomputation: trade compute for memory

01.09 · The Jacobian Matrix

Count rows and columns. If a function takes two inputs and returns three outputs, its Jacobian has shape (3,2). Row i describes output i. Column j describes sensitivity to input j.

When outputs are vectors

The gradient applies to scalar-valued functions. When $f: \mathbb{R}^n \to \mathbb{R}^m$ produces a vector output, we need the Jacobian matrix:

Example: a linear map

For $f(\mathbf{x}) = A\mathbf{x}$ where $A \in \mathbb{R}^{m \times n}$, the Jacobian is simply $A$ itself:

This makes sense: the linear map is its own best linear approximation everywhere.

Example: element-wise nonlinearity

For $f(\mathbf{x}) = \sigma(\mathbf{x})$ applied element-wise (e.g., sigmoid, ReLU), the Jacobian is diagonal:

Each output $f_i$ depends only on input $x_i$, so all off-diagonal entries are zero.

The Jacobian in the chain rule

For a composition $\mathbf{y} = f(g(\mathbf{x}))$, the chain rule becomes:

This is just matrix multiplication. For a neural network with $K$ layers:

Jacobian-vector products (JVPs)

In practice, we rarely compute the full Jacobian (which is $m \times n$ and can be huge). Instead, we compute:

• JVP (forward-mode): $J\mathbf{v}$ — how the output changes when the input moves in direction $\mathbf{v}$
• VJP (reverse-mode): $\mathbf{u}^\top J$ — the gradient of $\mathbf{u}^\top f(\mathbf{x})$ with respect to $\mathbf{x}$

Backpropagation computes VJPs. PyTorch's torch.autograd.grad computes VJPs; torch.autograd.functional.jvp computes JVPs.

01.10 · Hadamard Product (Element-wise Operations)

Keep two products separate. The Hadamard product of [2,3] and [4,5] is [8,15]. Their dot product is the single number 23. Same inputs, different output shape, different purpose.

Element-wise multiplication

When two arrays have the same shape, you can multiply them position-by-position. This is the Hadamard product, written $\odot$:

The top-right entry is $2 \times 6 = 12$. It does not mix numbers from other positions.

Hadamard vs. matrix multiplication

Using the same two matrices:

The top-right entry here is $1 \times 6 + 2 \times 8 = 22$ — a dot product of row 1 and column 2. Very different from the Hadamard result of 12.

Masks and gating

The Hadamard product is how you apply masks — zeroing out some entries while keeping others:

The mask $[1, 0, 1]$ hides the middle entry. This pattern appears everywhere in ML:

• Dropout: multiply activations by a random binary mask
• Attention: mask out padding tokens before softmax
• Gating (LSTM, GRU): element-wise multiply gates with hidden states

In code

import numpy as np

a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6], [7, 8]])

# Element-wise (Hadamard)
print((a * b).tolist())   # [[5, 12], [21, 32]]

# Matrix multiplication
print((a @ b).tolist())   # [[19, 22], [43, 50]]

Broadcasting

NumPy extends element-wise operations to compatible but unequal shapes via broadcasting. For example, adding a vector to every row of a matrix:

X = np.array([[1, 2, 3], [4, 5, 6]])  # shape (2, 3)
b = np.array([10, 20, 30])             # shape (3,)
print(X + b)  # [[11, 22, 33], [14, 25, 36]]

01.11 · Entropy & Information Theory

Predict the uncertain case. A fair coin gives two equally likely outcomes, so one observation carries 1 bit of uncertainty. A coin that always lands heads has 0 bits: the outcome adds no surprise under that distribution.

Measuring uncertainty

Entropy quantifies the uncertainty or "surprise" in a probability distribution. It's the foundation of cross-entropy loss, the standard loss for classification.

Intuition

Entropy is the expected number of bits needed to encode an outcome drawn from $p$. Equivalently, it's the expected "surprise" — where surprise of event $i$ is $-\log_2 p_i$:

Properties of entropy

• $H(p) \geq 0$, with equality iff $p$ is a point mass (deterministic)
• $H(p) \leq \log_2 n$, with equality iff $p$ is uniform
• $H$ is concave: mixing distributions increases entropy

Cross-entropy

When we use a model $q$ to encode data actually drawn from the true distribution $p$, the expected code length is the cross-entropy:

By Gibbs' inequality, $H(p, q) \geq H(p)$, with equality iff $q = p$. The excess $H(p,q) - H(p)$ is the KL divergence.

Cross-entropy loss in classification

In classification, the true label is a one-hot vector $y$ (the true distribution $p$), and the model outputs probabilities $\hat{y}$ (the model distribution $q$). The cross-entropy loss is:

where $y^*$ is the true class. Minimising cross-entropy pushes the model to assign high probability to the correct class.

Perplexity

In language modelling, perplexity is the exponentiated cross-entropy:

It measures the effective branching factor of the model. Lower perplexity = better model. A perplexity of 10 means the model is as uncertain as choosing uniformly among 10 options.

01.12 · KL Divergence

Fix the direction. D_KL(p||q) asks how costly it is to use model distribution q when data follows p. Reversing the arguments asks a different question and can produce a different number.

Distance between distributions

The Kullback-Leibler (KL) divergence measures how one probability distribution differs from another. It's not a true distance (it's asymmetric), but it's the fundamental measure of dissimilarity in information theory.

Properties

• $D_{KL}(p \| q) \geq 0$ (Gibbs' inequality), with equality iff $p = q$
• $D_{KL}(p \| q) \neq D_{KL}(q \| p)$ in general (asymmetric)
• Does not satisfy the triangle inequality (not a metric)

Relationship to entropy and cross-entropy

KL divergence decomposes neatly:

It's the extra bits needed when using $q$ to encode data from $p$, beyond the optimal code length $H(p)$.

Forward vs. reverse KL

The asymmetry matters in practice:

• Forward KL $D_{KL}(p \| q)$: $q$ must cover all modes of $p$ (zero-avoiding). Used in maximum likelihood.
• Reverse KL $D_{KL}(q \| p)$: $q$ can focus on one mode of $p$ (zero-forcing). Used in variational inference.

Applications in ML

Variational inference (VAEs)

In a VAE, we minimise $D_{KL}(q_\phi(z|x) \| p(z))$ to regularise the encoder's posterior toward the prior.

Knowledge distillation

Distill a large teacher model into a small student by minimising $D_{KL}(p_{\text{teacher}} \| p_{\text{student}})$ on soft predictions.

Policy gradient (RL)

Trust-region methods (TRPO, PPO) constrain $D_{KL}(\pi_{\text{old}} \| \pi_{\text{new}})$ to prevent destructive policy updates.

01.13 · Singular Value Decomposition (SVD)

Predict what compression removes. Keeping the largest singular values preserves the strongest matrix directions. Dropping a small singular value removes its direction's contribution. This is a controlled approximation, not automatic removal of “unimportant meaning.”

The most useful matrix factorisation

The SVD decomposes any matrix into three interpretable pieces. It reveals the fundamental structure of a linear map and is used throughout ML for dimensionality reduction, compression, and analysis.

Geometric interpretation

Any linear map $A$ can be decomposed into three steps:

1. Rotate or reflect the input space ($V^\top$)
2. Scale each axis independently ($\Sigma$)
3. Rotate or reflect into the output space ($U$)

The singular values $\sigma_i$ tell you how much $A$ stretches space along each principal direction.

Truncated SVD and low-rank approximation

Keeping only the top $k$ singular values gives the best rank-$k$ approximation:

The Eckart-Young theorem guarantees this minimises $\|A - B\|_F$ over all rank-$k$ matrices $B$.

Applications in ML

Principal Component Analysis (PCA)

PCA on a centred data matrix $X$ is the SVD of $X/\sqrt{n-1}$. The right singular vectors are the principal components; the singular values give the standard deviations along each component.

LoRA (Low-Rank Adaptation)

LoRA fine-tunes large language models by decomposing weight updates as low-rank matrices:

where $B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times k}$ with $r \ll \min(d,k)$. This is an implicit SVD-like decomposition of the update $\Delta W = BA$.

Condition number

The condition number of $A$ is $\kappa(A) = \sigma_{\max}/\sigma_{\min}$. Large $\kappa$ means the matrix is ill-conditioned — small input changes produce large output changes. This affects numerical stability of linear solvers and gradient flow.

Computing SVD

import numpy as np

A = np.array([[1, 2], [3, 4], [5, 6]])
U, s, Vt = np.linalg.svd(A, full_matrices=False)

print("Singular values:", s)      # [9.525, 0.514]
print("Rank-1 approx:", s[0] * np.outer(U[:,0], Vt[0,:]))

# Truncated to top-k
k = 1
A_k = U[:, :k] @ np.diag(s[:k]) @ Vt[:k, :]

01.14 · Moving Averages (EMA)

Calculate one update. With old average 10, new value 14, and decay 0.75, the exponential moving average becomes 0.75×10 + 0.25×14 = 11. Larger decay keeps more history and reacts more slowly.

Smoothing noisy sequences

During training, loss values, gradients, and metrics fluctuate. Moving averages smooth these sequences to reveal underlying trends.

Simple moving average (SMA)

Average the last $k$ values:

When a new value arrives, drop the oldest and add the newest.

Exponential moving average (EMA)

The EMA keeps a single running estimate, blending the old estimate with the new observation:

Effective window size

An EMA with decay $\beta$ behaves roughly like a simple average over the last $\frac{1}{1-\beta}$ observations:

• $\beta = 0.9 \implies$ effective window $\approx 10$
• $\beta = 0.99 \implies$ effective window $\approx 100$
• $\beta = 0.999 \implies$ effective window $\approx 1000$

Bias correction

When initialised at $v_0 = 0$, the first few estimates are biased toward zero. The correction divides by $1 - \beta^t$:

At $t=1$ with $\beta=0.9$: $\hat{v}_1 = v_1 / 0.1 = 10 v_1$. As $t \to \infty$, $\beta^t \to 0$ and the correction vanishes.

EMA in optimisers

EMAs are the core mechanism of modern optimisers:

EMA for model weights

Averaging the model weights themselves over training often improves generalisation:

EMA gives recent weights more influence. It differs from standard Stochastic Weight Averaging (SWA), which equally averages selected training checkpoints. Evaluate the averaged weights before choosing them for inference; averaging is not a guarantee of better accuracy.

Implementation

class EMA:
    def __init__(self, beta=0.999):
        self.beta = beta
        self.v = None
        self.t = 0
    
    def update(self, x):
        self.t += 1
        if self.v is None:
            self.v = x  # initialise with first observation
        else:
            self.v = self.beta * self.v + (1 - self.beta) * x
        # Bias correction (only if zero-initialised)
        return self.v  # or self.v / (1 - self.beta**self.t)

# Usage: smoothing training loss
ema = EMA(beta=0.99)
for loss in training_losses:
    smoothed = ema.update(loss)